fix: shell tool preview stuck as Running after completion (#3436)

1 files changed, 42 insertions, 7 deletions

diff --git a/crates/atuin-ai/src/fsm/mod.rs b/crates/atuin-ai/src/fsm/mod.rs
index 92be1cd8..d32d6d7b 100644
--- a/crates/atuin-ai/src/fsm/mod.rs
+++ b/crates/atuin-ai/src/fsm/mod.rs
@@ -447,11 +447,24 @@ impl AgentFsm {
                 },
             ) => {
                 if let Some(tracked) = self.ctx.tools.get_mut(&tool_id) {
-                    tracked.preview = Some(tools::ToolPreviewData::Shell {
-                        lines,
-                        exit_code,
-                        interrupted: false,
-                    });
+                    if tracked.is_resolved() {
+                        // Tool already completed — a late preview update raced with
+                        // ToolExecutionDone. Update lines (they may carry the final
+                        // screen) but preserve the finalized exit_code/interrupted.
+                        if let Some(tools::ToolPreviewData::Shell {
+                            lines: existing_lines,
+                            ..
+                        }) = &mut tracked.preview
+                        {
+                            *existing_lines = lines;
+                        }
+                    } else {
+                        tracked.preview = Some(tools::ToolPreviewData::Shell {
+                            lines,
+                            exit_code,
+                            interrupted: false,
+                        });
+                    }
                 }
                 vec![]
             }
@@ -799,8 +812,30 @@ impl AgentFsm {
         }
 
         tracked.state = ToolState::Completed;
-        if preview.is_some() {
-            tracked.preview = preview;
+
+        // Merge shell preview: the final ToolExecutionDone carries exit_code/interrupted
+        // but has empty lines (the live lines were accumulated via ToolPreviewUpdate).
+        // Preserve the accumulated lines and fold in the terminal metadata.
+        match (&mut tracked.preview, preview) {
+            (
+                Some(tools::ToolPreviewData::Shell {
+                    exit_code,
+                    interrupted,
+                    ..
+                }),
+                Some(tools::ToolPreviewData::Shell {
+                    exit_code: final_exit,
+                    interrupted: final_interrupted,
+                    ..
+                }),
+            ) => {
+                *exit_code = final_exit;
+                *interrupted = final_interrupted;
+            }
+            (_, Some(p)) => {
+                tracked.preview = Some(p);
+            }
+            _ => {}
         }
 
         let content = outcome.format_for_llm();