fix(buzz-agent-core): drain on cancel instead of dropping the stream

The new cancel test caught two real bugs, both from breaking out of the select
loop the moment the token fired.

Dropping `stream` drops the futures goose is awaiting, so
`mcp_client.rs:688` never reaches its `cancel_token.cancelled()` arm and never
sends `notifications/cancelled`. Consequences: the MCP child keeps running its
tool after the turn is over, and any announced `tool_call` never reaches a
terminal state -- the desktop renders that as a spinner forever, which is the
invariant buzz-agent held at agent.rs:470-477.

Cancellation is cooperative, so treat it that way: keep polling the stream and
let goose unwind (emit tool responses, send the MCP cancellations, end the
stream), bounded by CANCEL_DRAIN_TIMEOUT = 5s. Track announced-minus-resolved
tool call ids and synthesise terminal updates for any stragglers if the drain
times out -- a wrong status beats a stuck spinner.

3 new tests: cancel mid-tool-call returns stopReason=cancelled with every tool
call resolved and activeRunId cleared; notifications/cancelled actually reaches
the MCP server (asserted via FAKE_MCP_CANCEL_LOG); cancel for an unknown
session doesn't kill the process.

28 tests green, fmt + clippy -D warnings clean.

Signed-off-by: Michael Neale <michael.neale@gmail.com>

buzz-agent-core/src/agent.rs

@@ -60,6 +60,15 @@ const ERROR_REFLECTION: &str =
 /// conversation.
 const MAX_REFLECTIONS: usize = 8;
 
+/// How long to let goose unwind after `session/cancel` before giving up.
+///
+/// Cancellation is cooperative: goose has to notice the token, send
+/// `notifications/cancelled` to each in-flight MCP request
+/// (`mcp_client.rs:687-690`), emit the resulting tool responses, and end the
+/// stream. Dropping the stream instead skips all of that. This bounds the
+/// wait so a wedged MCP child cannot hold the turn open indefinitely.
+const CANCEL_DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
+
 /// Per-turn token accounting, mirroring buzz-agent's contract.
 #[derive(Debug, Default, Clone, Copy)]
 pub struct TurnTokens {
@@ -223,6 +232,14 @@ async fn drive_stream(
     let mut compacted = false;
     let mut reflections = 0usize;
 
+    // Set once cancellation is observed; bounds how long we let goose unwind.
+    let mut drain: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = None;
+
+    // Tool calls announced to the harness, minus those that reached a terminal
+    // state. Anything still here when the stream ends gets a synthetic
+    // terminal update — see the cancel arm below.
+    let mut open_tool_calls: Vec<String> = Vec::new();
+
     loop {
         tokio::select! {
             // Keep the harness idle clock alive during provider silence.
@@ -234,8 +251,28 @@ async fn drive_stream(
                 .await;
             }
 
-            _ = cancel.cancelled() => {
+            // Cancellation is a cooperative DRAIN, not an abort.
+            //
+            // Breaking here drops `stream`, which drops the futures goose is
+            // awaiting — so `mcp_client.rs:688` never reaches its
+            // `cancel_token.cancelled()` arm and never sends
+            // `notifications/cancelled`. The MCP child keeps running its tool
+            // after the turn is over, and any announced `tool_call` never
+            // reaches a terminal state, leaving the desktop UI spinning
+            // (the invariant buzz-agent held at `agent.rs:470-477`).
+            //
+            // So we keep polling the stream and let goose unwind: it emits the
+            // tool responses, sends the MCP cancellations, and ends the stream
+            // itself. `drain` only bounds how long we are willing to wait.
+            _ = cancel.cancelled(), if drain.is_none() => {
                 stop = StopReason::Cancelled;
+                drain = Some(Box::pin(tokio::time::sleep(CANCEL_DRAIN_TIMEOUT)));
+            }
+
+            // Goose did not unwind in time. Give up and synthesise terminal
+            // states, because a stuck spinner is worse than a wrong status.
+            _ = async { drain.as_mut().unwrap().await }, if drain.is_some() => {
+                tracing::warn!("cancel drain timed out; forcing terminal tool states");
                 break;
             }
 
@@ -251,6 +288,7 @@ async fn drive_stream(
                             tokens,
                             &mut compacted,
                             &mut reflections,
+                            &mut open_tool_calls,
                         )
                         .await
                         {
@@ -276,6 +314,26 @@ async fn drive_stream(
         }
     }
 
+    // Anything announced but never resolved gets a synthetic terminal update.
+    // buzz-agent guaranteed this invariant (`agent.rs:470-477`) because the
+    // desktop renders an unresolved `tool_call` as a spinner forever. Normally
+    // the drain above means this list is already empty; it only fires when
+    // goose failed to unwind inside CANCEL_DRAIN_TIMEOUT.
+    for tool_call_id in open_tool_calls.drain(..) {
+        wire::send(
+            wire_tx,
+            wire::session_update(
+                session_id,
+                json!({
+                    "sessionUpdate": "tool_call_update",
+                    "toolCallId": tool_call_id,
+                    "status": "failed",
+                }),
+            ),
+        )
+        .await;
+    }
+
     // Goose compacted history mid-turn. Re-inject the todo list so it survives
     // the truncation — this is what buzz-agent's `_PostCompact` existed for
     // (`handoff.rs:73-81`). Steering is the right channel: goose drains it at
@@ -318,13 +376,20 @@ async fn handle_event(
     tokens: &mut TurnTokens,
     compacted: &mut bool,
     reflections: &mut usize,
+    open_tool_calls: &mut Vec<String>,
 ) -> Option<StopReason> {
     match event {
         AgentEvent::Message(msg) => {
             for content in &msg.content {
                 emit_content(content, session_id, wire_tx).await;
 
+                if let MessageContent::ToolRequest(req) = content {
+                    open_tool_calls.push(req.id.clone());
+                }
+
                 if let MessageContent::ToolResponse(resp) = content {
+                    open_tool_calls.retain(|id| id != &resp.id);
+
                     let failed = match &resp.tool_result {
                         Ok(r) => r.is_error.unwrap_or(false),
                         Err(_) => true,

buzz-agent-core/tests/cancel.rs

@@ -0,0 +1,355 @@
+//! `session/cancel` — the last ACP method without coverage.
+//!
+//! buzz-acp sends this as a *notification* (no id, no response expected) and
+//! then waits for the in-flight `session/prompt` to return. Two things must
+//! hold or the harness wedges:
+//!
+//! 1. The prompt response arrives, with `stopReason: "cancelled"`.
+//! 2. Every announced tool call reaches a terminal state — buzz-agent
+//!    guaranteed this (`agent.rs:470-477`) because the desktop UI otherwise
+//!    shows a spinner forever.
+//!
+//! Goose plumbs the cancellation token into MCP tool calls and emits
+//! `notifications/cancelled` per in-flight request (`mcp_client.rs:687-690`),
+//! so this is a cooperative drain rather than a hard abort.
+
+use std::io::{BufRead, BufReader, Write};
+use std::net::TcpListener;
+use std::process::{Child, Command, Stdio};
+use std::time::{Duration, Instant};
+
+use serde_json::{json, Value};
+
+/// Provider whose first generation calls a tool that will hang.
+fn spawn_hanging_tool_provider() -> String {
+    let listener = TcpListener::bind("127.0.0.1:0").expect("bind");
+    let addr = listener.local_addr().expect("addr");
+
+    std::thread::spawn(move || {
+        for stream in listener.incoming() {
+            let Ok(mut stream) = stream else { continue };
+            let mut reader = BufReader::new(stream.try_clone().expect("clone"));
+
+            let mut request_line = String::new();
+            if reader.read_line(&mut request_line).is_err() {
+                continue;
+            }
+            let mut content_length = 0usize;
+            loop {
+                let mut line = String::new();
+                if reader.read_line(&mut line).is_err() {
+                    break;
+                }
+                let t = line.trim_end();
+                if t.is_empty() {
+                    break;
+                }
+                if let Some(v) = t
+                    .strip_prefix("content-length: ")
+                    .or_else(|| t.strip_prefix("Content-Length: "))
+                {
+                    content_length = v.parse().unwrap_or(0);
+                }
+            }
+            let mut body = vec![0u8; content_length];
+            use std::io::Read;
+            let _ = reader.read_exact(&mut body);
+
+            let respond = |stream: &mut std::net::TcpStream, ct: &str, payload: String| {
+                let _ = stream.write_all(
+                    format!(
+                        "HTTP/1.1 200 OK\r\ncontent-type: {ct}\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
+                        payload.len(),
+                        payload
+                    )
+                    .as_bytes(),
+                );
+                let _ = stream.flush();
+            };
+
+            if request_line.contains("/models") {
+                respond(
+                    &mut stream,
+                    "application/json",
+                    json!({ "object": "list", "data": [] }).to_string(),
+                );
+                continue;
+            }
+
+            let call = json!({
+                "id": "c", "object": "chat.completion.chunk", "created": 1,
+                "model": "fake-model",
+                "choices": [{
+                    "index": 0,
+                    "delta": {
+                        "role": "assistant",
+                        "tool_calls": [{
+                            "index": 0,
+                            "id": "call_1",
+                            "type": "function",
+                            "function": { "name": "devtools__tool_0", "arguments": "{}" },
+                        }],
+                    },
+                    "finish_reason": null,
+                }],
+            });
+            let done = json!({
+                "id": "c", "object": "chat.completion.chunk", "created": 1,
+                "model": "fake-model",
+                "choices": [{ "index": 0, "delta": {}, "finish_reason": "tool_calls" }],
+                "usage": { "prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7 },
+            });
+            respond(
+                &mut stream,
+                "text/event-stream",
+                format!("data: {call}\n\ndata: {done}\n\ndata: [DONE]\n\n"),
+            );
+        }
+    });
+
+    format!("http://{addr}")
+}
+
+struct Harness {
+    child: Child,
+    stdin: std::process::ChildStdin,
+    stdout: BufReader<std::process::ChildStdout>,
+    next_id: i64,
+}
+
+impl Harness {
+    fn start(base_url: &str, home: &std::path::Path) -> Self {
+        let mut child = Command::new(env!("CARGO_BIN_EXE_buzz-agent-core"))
+            .env("BUZZ_AGENT_PROVIDER", "openai-compat")
+            .env("BUZZ_AGENT_MODEL", "fake-model")
+            .env("OPENAI_COMPAT_API_KEY", "k")
+            .env("OPENAI_COMPAT_BASE_URL", base_url)
+            .env("HOME", home)
+            .env("XDG_CONFIG_HOME", home.join("config"))
+            .env("XDG_DATA_HOME", home.join("data"))
+            .env("GOOSE_DISABLE_KEYRING", "1")
+            .env("RUST_LOG", "warn")
+            .stdin(Stdio::piped())
+            .stdout(Stdio::piped())
+            .stderr(Stdio::inherit())
+            .spawn()
+            .expect("spawn");
+        let stdin = child.stdin.take().expect("stdin");
+        let stdout = BufReader::new(child.stdout.take().expect("stdout"));
+        Self {
+            child,
+            stdin,
+            stdout,
+            next_id: 0,
+        }
+    }
+
+    fn request(&mut self, method: &str, params: Value) -> i64 {
+        self.next_id += 1;
+        let id = self.next_id;
+        writeln!(
+            self.stdin,
+            "{}",
+            json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params })
+        )
+        .expect("write");
+        self.stdin.flush().expect("flush");
+        id
+    }
+
+    /// Fire-and-forget, exactly as buzz-acp sends `session/cancel`.
+    fn notify(&mut self, method: &str, params: Value) {
+        writeln!(
+            self.stdin,
+            "{}",
+            json!({ "jsonrpc": "2.0", "method": method, "params": params })
+        )
+        .expect("write");
+        self.stdin.flush().expect("flush");
+    }
+
+    fn await_response(&mut self, id: i64) -> (Value, Vec<Value>) {
+        let mut notifications = Vec::new();
+        loop {
+            let mut line = String::new();
+            let n = self.stdout.read_line(&mut line).expect("read");
+            assert_ne!(n, 0, "agent closed stdout awaiting id={id}");
+            let Ok(msg) = serde_json::from_str::<Value>(&line) else {
+                continue;
+            };
+            if msg.get("id").and_then(Value::as_i64) == Some(id) {
+                return (msg, notifications);
+            }
+            if msg.get("method").is_some() {
+                notifications.push(msg);
+            }
+        }
+    }
+}
+
+impl Drop for Harness {
+    fn drop(&mut self) {
+        let _ = self.child.kill();
+        let _ = self.child.wait();
+    }
+}
+
+fn new_session(h: &mut Harness, cwd: &std::path::Path, cancel_log: &std::path::Path) -> String {
+    let id = h.request("initialize", json!({ "protocolVersion": 2 }));
+    let _ = h.await_response(id);
+
+    let id = h.request(
+        "session/new",
+        json!({
+            "cwd": cwd.to_str().unwrap(),
+            "mcpServers": [{
+                "name": "devtools",
+                "command": env!("CARGO_BIN_EXE_fake-mcp"),
+                "args": [],
+                "env": [
+                    // Tool call blocks ~forever, so cancel lands mid-flight.
+                    { "name": "FAKE_MCP_TOOL_DELAY", "value": "999" },
+                    { "name": "FAKE_MCP_CANCEL_LOG", "value": cancel_log.to_str().unwrap() },
+                ],
+            }],
+        }),
+    );
+    let (resp, _) = h.await_response(id);
+    resp["result"]["sessionId"]
+        .as_str()
+        .unwrap_or_else(|| panic!("session/new failed: {resp}"))
+        .to_string()
+}
+
+#[test]
+fn cancel_mid_tool_call_returns_cancelled() {
+    let base_url = spawn_hanging_tool_provider();
+    let home = tempfile::tempdir().expect("home");
+    let cwd = tempfile::tempdir().expect("cwd");
+    let logs = tempfile::tempdir().expect("logs");
+    let cancel_log = logs.path().join("cancelled.jsonl");
+
+    let mut h = Harness::start(&base_url, home.path());
+    let session_id = new_session(&mut h, cwd.path(), &cancel_log);
+
+    let id = h.request(
+        "session/prompt",
+        json!({
+            "sessionId": session_id,
+            "prompt": [{ "type": "text", "text": "run the slow tool" }],
+        }),
+    );
+
+    // Let the turn reach the hanging tool, then cancel.
+    std::thread::sleep(Duration::from_secs(2));
+    h.notify("session/cancel", json!({ "sessionId": session_id }));
+
+    let started = Instant::now();
+    let (resp, notifications) = h.await_response(id);
+    let elapsed = started.elapsed();
+
+    assert!(
+        elapsed < Duration::from_secs(30),
+        "cancel took {elapsed:?}; the tool sleeps 999s, so this did not unblock"
+    );
+    assert_eq!(
+        resp["result"]["stopReason"], "cancelled",
+        "expected stopReason=cancelled: {resp}"
+    );
+
+    // Every announced tool call must reach a terminal state, or the desktop
+    // spins forever (buzz-agent agent.rs:470-477).
+    let announced: Vec<String> = notifications
+        .iter()
+        .filter(|n| n["params"]["update"]["sessionUpdate"] == "tool_call")
+        .filter_map(|n| {
+            n["params"]["update"]["toolCallId"]
+                .as_str()
+                .map(String::from)
+        })
+        .collect();
+    let terminal: Vec<String> = notifications
+        .iter()
+        .filter(|n| n["params"]["update"]["sessionUpdate"] == "tool_call_update")
+        .filter_map(|n| {
+            n["params"]["update"]["toolCallId"]
+                .as_str()
+                .map(String::from)
+        })
+        .collect();
+
+    for id in &announced {
+        assert!(
+            terminal.contains(id),
+            "tool call {id} was announced but never resolved — UI would hang.\n\
+             announced={announced:?} terminal={terminal:?}"
+        );
+    }
+
+    // activeRunId must be cleared, or a later steer targets a dead run.
+    let cleared = notifications.iter().any(|n| {
+        n["params"]["update"]["sessionUpdate"] == "session_info_update"
+            && n["params"]["update"]["_meta"]["goose"]["activeRunId"].is_null()
+    });
+    assert!(cleared, "activeRunId was never cleared after cancel");
+}
+
+#[test]
+fn cancel_propagates_notifications_cancelled_to_mcp() {
+    let base_url = spawn_hanging_tool_provider();
+    let home = tempfile::tempdir().expect("home");
+    let cwd = tempfile::tempdir().expect("cwd");
+    let logs = tempfile::tempdir().expect("logs");
+    let cancel_log = logs.path().join("cancelled.jsonl");
+
+    let mut h = Harness::start(&base_url, home.path());
+    let session_id = new_session(&mut h, cwd.path(), &cancel_log);
+
+    let id = h.request(
+        "session/prompt",
+        json!({
+            "sessionId": session_id,
+            "prompt": [{ "type": "text", "text": "run the slow tool" }],
+        }),
+    );
+
+    std::thread::sleep(Duration::from_secs(2));
+    h.notify("session/cancel", json!({ "sessionId": session_id }));
+    let _ = h.await_response(id);
+
+    // Goose should tell the MCP server to stop (mcp_client.rs:687-690).
+    // Give it a moment to flush.
+    let mut saw = false;
+    for _ in 0..20 {
+        if let Ok(contents) = std::fs::read_to_string(&cancel_log) {
+            if contents.contains("notifications/cancelled") {
+                saw = true;
+                break;
+            }
+        }
+        std::thread::sleep(Duration::from_millis(100));
+    }
+    assert!(
+        saw,
+        "MCP server never received notifications/cancelled — a hung tool would \
+         keep running after the turn ended"
+    );
+}
+
+#[test]
+fn cancel_for_unknown_session_is_ignored() {
+    let base_url = spawn_hanging_tool_provider();
+    let home = tempfile::tempdir().expect("home");
+    let mut h = Harness::start(&base_url, home.path());
+
+    let id = h.request("initialize", json!({ "protocolVersion": 2 }));
+    let _ = h.await_response(id);
+
+    // Must not crash the agent: it is a notification, so there is nothing to
+    // reply to, and buzz-acp would hang if the process died here.
+    h.notify("session/cancel", json!({ "sessionId": "ses_nope" }));
+
+    let id = h.request("initialize", json!({ "protocolVersion": 2 }));
+    let (resp, _) = h.await_response(id);
+    assert_eq!(resp["result"]["agentInfo"]["name"], "buzz-agent");
+}