feat(shell): support high-resolution duration if available (#1484)

* feat(bash): measure duration in microsecond resolution with ble.sh

* feat(zsh): measure duration in nanosecond resolution with zsh/datetime

* refactor(history): use Option<i64> for command-line arg duration

Co-authored-by: Ellie Huxtable <ellie@elliehuxtable.com>

* style(history): apply suggestion by `cargo fmt`

* fix(history): use Option<u64> for arg duration

---------

Co-authored-by: Ellie Huxtable <ellie@elliehuxtable.com>

4 files changed, 33 insertions, 6 deletions

diff --git a/atuin/src/command/client/history.rs b/atuin/src/command/client/history.rs
index 8cf35d5b..85ca69ff 100644
--- a/atuin/src/command/client/history.rs
+++ b/atuin/src/command/client/history.rs
@@ -34,6 +34,8 @@ pub enum Cmd {
         id: String,
         #[arg(long, short)]
         exit: i64,
+        #[arg(long, short)]
+        duration: Option<u64>,
     },
 
     /// List all items in history
@@ -272,6 +274,7 @@ impl Cmd {
         settings: &Settings,
         id: &str,
         exit: i64,
+        duration: Option<u64>,
     ) -> Result<()> {
         if id.trim() == "" {
             return Ok(());
@@ -290,8 +293,11 @@ impl Cmd {
         }
 
         h.exit = exit;
-        h.duration = i64::try_from((OffsetDateTime::now_utc() - h.timestamp).whole_nanoseconds())
-            .context("command took over 292 years")?;
+        h.duration = match duration {
+            Some(value) => i64::try_from(value).context("command took over 292 years")?,
+            None => i64::try_from((OffsetDateTime::now_utc() - h.timestamp).whole_nanoseconds())
+                .context("command took over 292 years")?,
+        };
 
         db.update(&h).await?;
 
@@ -366,7 +372,9 @@ impl Cmd {
 
         match self {
             Self::Start { command } => Self::handle_start(db, settings, &command).await,
-            Self::End { id, exit } => Self::handle_end(db, settings, &id, exit).await,
+            Self::End { id, exit, duration } => {
+                Self::handle_end(db, settings, &id, exit, duration).await
+            }
             Self::List {
                 session,
                 cwd,
diff --git a/atuin/src/command/client/search/duration.rs b/atuin/src/command/client/search/duration.rs
index 08dadb95..dfa9426b 100644
--- a/atuin/src/command/client/search/duration.rs
+++ b/atuin/src/command/client/search/duration.rs
@@ -29,6 +29,7 @@ pub fn format_duration_into(dur: Duration, f: &mut fmt::Formatter<'_>) -> fmt::R
         let seconds = day_secs % 60;
 
         let millis = nanos / 1_000_000;
+        let micros = nanos / 1_000;
 
         // a difference from our impl than the original is that
         // we only care about the most-significant segment of the duration.
@@ -41,6 +42,8 @@ pub fn format_duration_into(dur: Duration, f: &mut fmt::Formatter<'_>) -> fmt::R
         item("m", minutes)?;
         item("s", seconds)?;
         item("ms", u64::from(millis))?;
+        item("us", u64::from(micros))?;
+        item("ns", u64::from(nanos))?;
         ControlFlow::Continue(())
     }
 
diff --git a/atuin/src/shell/atuin.bash b/atuin/src/shell/atuin.bash
index d05859d2..1a5a6383 100644
--- a/atuin/src/shell/atuin.bash
+++ b/atuin/src/shell/atuin.bash
@@ -13,7 +13,16 @@ __atuin_precmd() {
 
     [[ -z "${ATUIN_HISTORY_ID}" ]] && return
 
-    (ATUIN_LOG=error atuin history end --exit "${EXIT}" -- "${ATUIN_HISTORY_ID}" &) >/dev/null 2>&1
+    local duration=""
+    # shellcheck disable=SC2154,SC2309
+    if [[ -n "${BLE_ATTACHED-}" && _ble_bash -ge 50000 && -n "${_ble_exec_time_ata-}" ]]; then
+      # We use the high-resolution duration based on EPOCHREALTIME (bash >=
+      # 5.0) that is recorded by ble.sh. The shell variable
+      # `_ble_exec_time_ata` contains the execution time in microseconds.
+      duration=${_ble_exec_time_ata}000
+    fi
+
+    (ATUIN_LOG=error atuin history end --exit "${EXIT}" ${duration:+--duration "$duration"} -- "${ATUIN_HISTORY_ID}" &) >/dev/null 2>&1
     export ATUIN_HISTORY_ID=""
 }
 
diff --git a/atuin/src/shell/atuin.zsh b/atuin/src/shell/atuin.zsh
index 491ff9ac..16c001c4 100644
--- a/atuin/src/shell/atuin.zsh
+++ b/atuin/src/shell/atuin.zsh
@@ -9,6 +9,7 @@
 # Source this in your ~/.zshrc
 autoload -U add-zsh-hook
 
+zmodload zsh/datetime 2>/dev/null
 
 # If zsh-autosuggestions is installed, configure it to use Atuin's search. If
 # you'd like to override this, then add your config after the $(atuin init zsh)
@@ -24,14 +25,20 @@ _atuin_preexec() {
     local id
     id=$(atuin history start -- "$1")
     export ATUIN_HISTORY_ID="$id"
+    __atuin_preexec_time=${EPOCHREALTIME-}
 }
 
 _atuin_precmd() {
-    local EXIT="$?"
+    local EXIT="$?" __atuin_precmd_time=${EPOCHREALTIME-}
 
     [[ -z "${ATUIN_HISTORY_ID:-}" ]] && return
 
-    (ATUIN_LOG=error atuin history end --exit $EXIT -- $ATUIN_HISTORY_ID &) >/dev/null 2>&1
+    local duration=""
+    if [[ -n $__atuin_preexec_time && -n $__atuin_precmd_time ]]; then
+        printf -v duration %.0f $(((__atuin_precmd_time - __atuin_preexec_time) * 1000000000))
+    fi
+
+    (ATUIN_LOG=error atuin history end --exit $EXIT ${=duration:+--duration $duration} -- $ATUIN_HISTORY_ID &) >/dev/null 2>&1
     export ATUIN_HISTORY_ID=""
 }