feat(ai): surface live AI generation status with smooth timing

crates/dbx-core/src/agent_events.rs

@@ -22,7 +22,12 @@ pub enum AgentEvent {
     ToolCallEnd { tool_call_id: String, tool_name: String, result: serde_json::Value, is_error: bool },
     /// A turn in the agent loop has ended.
     TurnEnd { turn: u32 },
-    /// The agent loop has finished.
+    /// The reply stream is fully consumed, but the run is NOT yet confirmed
+    /// successful (the CLI process may still exit non-zero or hang after closing
+    /// stdout). Non-terminal: the frontend may stop the reply animation on it,
+    /// but must keep listening for the real `AgentEnd` / `Error`.
+    ResponseComplete,
+    /// The agent loop has finished successfully.
     AgentEnd { input_tokens: Option<u32>, output_tokens: Option<u32> },
     /// Context was compacted to stay within context window limits.
     ContextCompacted {

crates/dbx-core/src/ai_cli_agent.rs

@@ -1000,6 +1000,23 @@ fn grok_usage_tokens(usage: &Value) -> Option<TokenUsage> {
     (input > 0 || output > 0).then_some(TokenUsage { input_tokens: input, output_tokens: output })
 }
 
+/// Stream a CLI agent's JSONL output to `on_event`, then return the accumulated
+/// final text.
+///
+/// `AgentEnd` is emitted ONLY after the process is confirmed successful: the CLI
+/// is reaped via `child.wait()` and the run must have no terminal error, no
+/// cancellation, and a zero exit status. Both the CLI's own `end` marker and the
+/// synthesized `AgentEnd` (used when the CLI emits none) are gated behind that
+/// confirmed exit — a CLI that closes stdout then hangs or exits non-zero must
+/// never be shown to the frontend as a completed reply.
+///
+/// To still let the frontend stop the reply animation promptly, a NON-terminal
+/// `ResponseComplete` is dispatched as soon as stdout is fully consumed (before
+/// `child.wait()` / the stderr drain, which can take seconds on some CLIs). The
+/// frontend may hide the status line on `ResponseComplete` but must keep
+/// listening for the real `AgentEnd` (success) / `Error` (failure). Cancellation
+/// / read-error paths skip `ResponseComplete` so the frontend cancel flow keeps
+/// its `cancelling` phase until `finally`.
 pub async fn run_cli_jsonl_agent(
     spec: CliAgentProcessSpec,
     cancelled: &Notify,
@@ -1039,7 +1056,7 @@ pub async fn run_cli_jsonl_agent(
     });
 
     let mut final_text = String::new();
-    let mut saw_agent_end = false;
+    let mut pending_agent_end: Option<AgentEvent> = None;
     let mut total_usage = TokenUsage::default();
     let mut terminal_error: Option<String> = None;
 
@@ -1063,8 +1080,17 @@ pub async fn run_cli_jsonl_agent(
                     total_usage.add(&usage);
                 }
                 for event in parsed.events {
+                    // A CLI-emitted `end` is a stream marker, not a confirmed
+                    // success: the process may still exit non-zero (or hang)
+                    // after closing stdout. Buffer it and re-emit only after
+                    // child.wait() confirms the exit status (see below).
+                    // Dialects that emit one `end` per tool result (e.g. Cursor)
+                    // collapse into the single confirmed-success AgentEnd here —
+                    // forwarding them inline would show the run as complete at
+                    // the first tool result.
                     if matches!(event, AgentEvent::AgentEnd { .. }) {
-                        saw_agent_end = true;
+                        pending_agent_end = Some(event);
+                        continue;
                     }
                     on_event(event);
                 }
@@ -1082,8 +1108,34 @@ pub async fn run_cli_jsonl_agent(
         }
     }
 
-    let status = child.wait().await.map_err(|e| format!("CLI agent wait failed: {e}"))?;
-    let stderr = stderr_task.await.unwrap_or_default();
+    // stdout is fully consumed. Emit a NON-terminal marker so the frontend can
+    // stop the reply animation (hide the status line) without treating the run
+    // as finished — the process is not confirmed successful until child.wait()
+    // returns a zero exit status below. Skipped on error/cancel paths so the
+    // frontend keeps its `cancelling`/`error` handling until `finally`.
+    if terminal_error.is_none() {
+        on_event(AgentEvent::ResponseComplete);
+    }
+
+    // The read loop above is gone, so this wait is the last chance to observe a
+    // cancel. Race the (potentially slow) post-EOF reap and stderr drain against
+    // `cancelled` so Stop still kills and reaps a CLI that closes stdout and
+    // then hangs. On cancel we kill() then wait() (kill-and-reap) so the child
+    // leaves no zombie behind, then return the cancel error. `stderr_task` is
+    // dropped on that branch — the pipe closes once the child is terminated, so
+    // the reader completes on its own (no leak, no hang).
+    let (status, stderr) = tokio::select! {
+        status = child.wait() => {
+            let status = status.map_err(|e| format!("CLI agent wait failed: {e}"))?;
+            let stderr = stderr_task.await.unwrap_or_default();
+            (status, stderr)
+        }
+        _ = cancelled.notified() => {
+            let _ = child.kill().await;
+            let _ = child.wait().await; // reap; kill() then wait() prevents a zombie
+            return Err("Agent loop cancelled".to_string());
+        }
+    };
 
     if let Some(error) = terminal_error {
         return Err(error);
@@ -1093,7 +1145,19 @@ pub async fn run_cli_jsonl_agent(
         return Err((spec.classify_run_error)(&stderr));
     }
 
-    if !saw_agent_end {
+    // Confirmed success. Forward the CLI's own `end` (backfilling usage from the
+    // accumulated totals when it carried none), else synthesize an `AgentEnd`.
+    if let Some(agent_end) = pending_agent_end {
+        if let AgentEvent::AgentEnd { input_tokens, output_tokens } = agent_end {
+            let agent_end = AgentEvent::AgentEnd {
+                input_tokens: input_tokens
+                    .or_else(|| (total_usage.input_tokens > 0).then_some(total_usage.input_tokens)),
+                output_tokens: output_tokens
+                    .or_else(|| (total_usage.output_tokens > 0).then_some(total_usage.output_tokens)),
+            };
+            on_event(agent_end);
+        }
+    } else {
         on_event(AgentEvent::AgentEnd {
             input_tokens: (total_usage.input_tokens > 0).then_some(total_usage.input_tokens),
             output_tokens: (total_usage.output_tokens > 0).then_some(total_usage.output_tokens),
@@ -1286,6 +1350,166 @@ mod tests {
             .any(|event| { matches!(event, AgentEvent::AgentEnd { input_tokens: Some(13), output_tokens: Some(6) }) }));
     }
 
+    #[tokio::test]
+    async fn jsonl_agent_never_emits_agent_end_when_process_fails_after_stdout_eof() {
+        // The CLI emits its final JSONL line, closes stdout (so the read loop
+        // sees EOF), lingers for a second, then exits with status 1. AgentEnd
+        // must NOT be emitted before the failure is determined — the frontend
+        // must not be shown a completed reply for a failed run. `ResponseComplete`
+        // (non-terminal) is still dispatched early so the reply animation stops.
+        let spec = CliAgentProcessSpec {
+            command: CliAgentCommandSpec {
+                program: "sh".to_string(),
+                args: vec![
+                    "-c".to_string(),
+                    concat!(
+                        "printf '%s\\n' ",
+                        "'{\"type\":\"step_finish\",\"part\":{\"tokens\":{\"input\":10,\"output\":2}}}' ",
+                        "'{\"type\":\"text\",\"part\":{\"text\":\"hello\"}}' ",
+                        "'{\"type\":\"step_finish\",\"part\":{\"tokens\":{\"input\":3,\"output\":4}}}'; ",
+                        "exec 1>&-; sleep 1; exit 1",
+                    )
+                    .to_string(),
+                ],
+            },
+            env: Vec::new(),
+            env_remove: Vec::new(),
+            current_dir: None,
+            stdin: None,
+            dialect: CliAgentJsonlDialect::OpenCodeRun,
+            classify_spawn_error,
+            classify_run_error,
+        };
+        let events = Arc::new(Mutex::new(Vec::new()));
+        let captured = Arc::clone(&events);
+        let (probe_tx, mut probe_rx) = tokio::sync::watch::channel(0u8);
+        let cancelled = Arc::new(Notify::new());
+        let runner_cancelled = Arc::clone(&cancelled);
+
+        let runner = tokio::spawn(async move {
+            run_cli_jsonl_agent(spec, runner_cancelled.as_ref(), move |event| {
+                if matches!(event, AgentEvent::ResponseComplete) {
+                    let _ = probe_tx.send(1);
+                }
+                if matches!(event, AgentEvent::AgentEnd { .. }) {
+                    let _ = probe_tx.send(2);
+                }
+                captured.lock().unwrap().push(event);
+            })
+            .await
+        });
+
+        // ResponseComplete must arrive as soon as stdout is consumed — well
+        // before the lingering process finishes its 1s shutdown.
+        timeout(Duration::from_secs(1), probe_rx.wait_for(|seen| *seen >= 1))
+            .await
+            .expect("ResponseComplete should be emitted while the process still runs")
+            .expect("probe watch channel closed before ResponseComplete was emitted");
+
+        // While the process is still alive (sleeping before its exit 1), AgentEnd
+        // must NOT have been emitted yet — the failure is not yet confirmed.
+        assert_eq!(*probe_rx.borrow(), 1, "AgentEnd must not be emitted before the exit status is known");
+
+        let result = timeout(Duration::from_secs(5), runner)
+            .await
+            .expect("runner should return once the failing process exits")
+            .expect("runner task should not panic");
+        assert!(result.is_err(), "run must fail for a non-zero exit");
+
+        let events = events.lock().unwrap();
+        assert_eq!(
+            events.iter().filter(|event| matches!(event, AgentEvent::AgentEnd { .. })).count(),
+            0,
+            "AgentEnd must never be emitted for a failed run"
+        );
+        assert!(
+            events.iter().any(|event| matches!(event, AgentEvent::ResponseComplete)),
+            "ResponseComplete should still be emitted for the failed run"
+        );
+    }
+
+    #[tokio::test]
+    async fn jsonl_agent_emits_response_complete_early_and_agent_end_only_after_confirmed_exit() {
+        // The CLI emits its final JSONL line, closes stdout (so the read loop
+        // sees EOF), lingers for a second, then exits with status 0. The
+        // non-terminal ResponseComplete must be dispatched as soon as stdout is
+        // consumed (before the process reap), while AgentEnd must wait for the
+        // confirmed-success exit and be the final event.
+        let spec = CliAgentProcessSpec {
+            command: CliAgentCommandSpec {
+                program: "sh".to_string(),
+                args: vec![
+                    "-c".to_string(),
+                    concat!(
+                        "printf '%s\\n' ",
+                        "'{\"type\":\"step_finish\",\"part\":{\"tokens\":{\"input\":10,\"output\":2}}}' ",
+                        "'{\"type\":\"text\",\"part\":{\"text\":\"hello\"}}' ",
+                        "'{\"type\":\"step_finish\",\"part\":{\"tokens\":{\"input\":3,\"output\":4}}}'; ",
+                        "exec 1>&-; sleep 1",
+                    )
+                    .to_string(),
+                ],
+            },
+            env: Vec::new(),
+            env_remove: Vec::new(),
+            current_dir: None,
+            stdin: None,
+            dialect: CliAgentJsonlDialect::OpenCodeRun,
+            classify_spawn_error,
+            classify_run_error,
+        };
+        let events = Arc::new(Mutex::new(Vec::new()));
+        let captured = Arc::clone(&events);
+        let (probe_tx, mut probe_rx) = tokio::sync::watch::channel(0u8);
+        let cancelled = Arc::new(Notify::new());
+        let runner_cancelled = Arc::clone(&cancelled);
+
+        let runner = tokio::spawn(async move {
+            run_cli_jsonl_agent(spec, runner_cancelled.as_ref(), move |event| {
+                if matches!(event, AgentEvent::ResponseComplete) {
+                    let _ = probe_tx.send(1);
+                }
+                if matches!(event, AgentEvent::AgentEnd { .. }) {
+                    let _ = probe_tx.send(2);
+                }
+                captured.lock().unwrap().push(event);
+            })
+            .await
+        });
+
+        // ResponseComplete must arrive as soon as the CLI output is consumed —
+        // well before the lingering process finishes its 1s shutdown.
+        timeout(Duration::from_secs(1), probe_rx.wait_for(|seen| *seen >= 1))
+            .await
+            .expect("ResponseComplete should be emitted while the process still runs")
+            .expect("probe watch channel closed before ResponseComplete was emitted");
+
+        // While the process is still alive, AgentEnd must not have been emitted yet.
+        assert_eq!(*probe_rx.borrow(), 1, "AgentEnd must not be emitted before the exit status is known");
+
+        let result = timeout(Duration::from_secs(5), runner)
+            .await
+            .expect("runner should return once the lingering process exits")
+            .expect("runner task should not panic");
+        assert_eq!(result.unwrap(), "hello");
+
+        let events = events.lock().unwrap();
+        assert_eq!(events.iter().filter(|event| matches!(event, AgentEvent::AgentEnd { .. })).count(), 1);
+        assert!(events
+            .iter()
+            .any(|event| { matches!(event, AgentEvent::AgentEnd { input_tokens: Some(13), output_tokens: Some(6) }) }));
+        let response_complete_at = events
+            .iter()
+            .position(|event| matches!(event, AgentEvent::ResponseComplete))
+            .expect("ResponseComplete should have been emitted");
+        let agent_end_at = events
+            .iter()
+            .position(|event| matches!(event, AgentEvent::AgentEnd { .. }))
+            .expect("AgentEnd should have been emitted after confirmed success");
+        assert!(response_complete_at < agent_end_at, "ResponseComplete must precede AgentEnd");
+        assert!(matches!(events.last(), Some(AgentEvent::AgentEnd { .. })), "AgentEnd should be the terminal event");
+    }
+
     #[tokio::test]
     async fn jsonl_error_kills_and_waits_for_child() {
         let pid_file = std::env::temp_dir().join(format!(
@@ -1362,4 +1586,81 @@ mod tests {
         assert!(!process_is_alive(pid.trim()));
         let _ = std::fs::remove_file(pid_file);
     }
+
+    #[tokio::test]
+    async fn jsonl_eof_cancel_kills_and_reaps_hanging_child() {
+        // The CLI emits a valid JSONL line, closes stdout (EOF -> ResponseComplete),
+        // then hangs before exiting. A cancel fired after EOF must race the
+        // post-EOF child.wait(): the child is killed AND reaped (no zombie) and
+        // the run returns promptly with the cancel error. No AgentEnd may be
+        // emitted on the cancel path.
+        let pid_file = std::env::temp_dir().join(format!(
+            "dbx-cli-agent-eof-cancel-{}-{}",
+            std::process::id(),
+            SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos()
+        ));
+        let script = format!(
+            "echo $$ > {}; printf '%s\\n' '{{\"type\":\"text\",\"part\":{{\"text\":\"hello\"}}}}'; exec 1>&-; exec sleep 30",
+            pid_file.display()
+        );
+
+        let spec = CliAgentProcessSpec {
+            command: CliAgentCommandSpec { program: "sh".to_string(), args: vec!["-c".to_string(), script] },
+            env: Vec::new(),
+            env_remove: Vec::new(),
+            current_dir: None,
+            stdin: None,
+            dialect: CliAgentJsonlDialect::OpenCodeRun,
+            classify_spawn_error,
+            classify_run_error,
+        };
+        let events = Arc::new(Mutex::new(Vec::new()));
+        let captured = Arc::clone(&events);
+        let (probe_tx, mut probe_rx) = tokio::sync::watch::channel(0u8);
+        let cancelled = Arc::new(Notify::new());
+        let runner_cancelled = Arc::clone(&cancelled);
+
+        let runner = tokio::spawn(async move {
+            run_cli_jsonl_agent(spec, runner_cancelled.as_ref(), move |event| {
+                if matches!(event, AgentEvent::ResponseComplete) {
+                    let _ = probe_tx.send(1);
+                }
+                if matches!(event, AgentEvent::AgentEnd { .. }) {
+                    let _ = probe_tx.send(2);
+                }
+                captured.lock().unwrap().push(event);
+            })
+            .await
+        });
+
+        // ResponseComplete must arrive as soon as stdout is consumed — the child
+        // is still hanging at that point, before the post-EOF wait.
+        timeout(Duration::from_secs(1), probe_rx.wait_for(|seen| *seen >= 1))
+            .await
+            .expect("ResponseComplete should be emitted while the process still runs")
+            .expect("probe watch channel closed before ResponseComplete was emitted");
+
+        assert_eq!(*probe_rx.borrow(), 1, "AgentEnd must not be emitted before the exit status is known");
+        let pid = std::fs::read_to_string(&pid_file).expect("child pid should be captured");
+        assert!(process_is_alive(pid.trim()), "hanging child should still be alive after ResponseComplete");
+
+        cancelled.notify_waiters();
+
+        let result = timeout(Duration::from_secs(3), runner)
+            .await
+            .expect("runner should return promptly after post-EOF cancellation")
+            .expect("runner task should not panic");
+        assert_eq!(result.unwrap_err(), "Agent loop cancelled");
+
+        sleep(Duration::from_millis(100)).await;
+        assert!(!process_is_alive(pid.trim()), "hanging child should be killed and reaped");
+        let _ = std::fs::remove_file(pid_file);
+
+        let events = events.lock().unwrap();
+        assert_eq!(
+            events.iter().filter(|event| matches!(event, AgentEvent::AgentEnd { .. })).count(),
+            0,
+            "AgentEnd must never be emitted on the cancel path"
+        );
+    }
 }