fix: cancel and drain MCP-owned shell work on disconnect

Track shell admission and completion at the MCP connection owner. On input EOF or error, cancel before response draining; wait for shell cleanup and join output-reader tasks before returning from server teardown. Test real separately grouped shell descendants and preservation of an unrelated peer. This is cooperative MCP teardown, not an arbitrary descendant containment or Desktop Stopped certificate.

Signed-off-by: Logan Johnson <loganj@squareup.com>

crates/buzz-dev-mcp/src/lib.rs

@@ -3,9 +3,7 @@
 use rmcp::{
     handler::server::{router::tool::ToolRouter, wrapper::Parameters},
     model::{CallToolResult, ServerCapabilities, ServerInfo},
-    tool, tool_handler, tool_router,
-    transport::stdio,
-    ErrorData, ServerHandler, ServiceExt,
+    tool, tool_handler, tool_router, ErrorData, ServerHandler, ServiceExt,
 };
 use std::path::Path;
 use std::sync::Arc;
@@ -19,6 +17,7 @@ mod str_replace;
 mod todo;
 mod tree;
 mod view_image;
+mod workloads;
 
 #[derive(Clone)]
 struct DevMcp {
@@ -180,9 +179,34 @@ async fn async_main(cmd: String) -> Result<(), Box<dyn std::error::Error>> {
     let shim = shim::Shim::install()?;
     let state = Arc::new(shell::SharedState::new(cwd, shim)?);
 
-    let service = DevMcp::new(state).serve(stdio()).await?;
-    service.waiting().await?;
-    Ok(())
+    serve_connection(state, tokio::io::stdin(), tokio::io::stdout())
+        .await
+        .map_err(|e| e as Box<dyn std::error::Error>)
+}
+
+async fn serve_connection<R, W>(
+    state: Arc<shell::SharedState>,
+    reader: R,
+    writer: W,
+) -> Result<(), Box<dyn std::error::Error + Send + Sync>>
+where
+    R: tokio::io::AsyncRead + Send + Unpin + 'static,
+    W: tokio::io::AsyncWrite + Send + Unpin + 'static,
+{
+    let input = workloads::CancelOnEof {
+        reader,
+        state: state.clone(),
+    };
+    let result = async {
+        let service = DevMcp::new(state.clone()).serve((input, writer)).await?;
+        service.waiting().await?;
+        Ok::<_, Box<dyn std::error::Error + Send + Sync>>(())
+    }
+    .await;
+    // Also covers failed initialization and broken stdout. Never leave runtime
+    // destruction to race the independently grouped shell cleanup futures.
+    state.workloads.drain().await?;
+    result
 }
 
 /// Suppress the console window that Windows otherwise allocates for every

crates/buzz-dev-mcp/src/shell.rs

@@ -24,6 +24,7 @@ const ARTIFACT_RING_SIZE: usize = 8;
 const READ_CHUNK: usize = 16 * 1024;
 
 pub struct SharedState {
+    pub(crate) workloads: crate::workloads::Workloads,
     pub cwd: PathBuf,
     pub shim: Shim,
     pub session_dir: TempDir,
@@ -52,6 +53,7 @@ impl SharedState {
         };
         let bootstrap_instructions = build_bootstrap(&cwd, shell_hint);
         Ok(Self {
+            workloads: Default::default(),
             cwd,
             shim,
             session_dir,
@@ -132,6 +134,13 @@ pub async fn run(
     p: ShellParams,
     ct: CancellationToken,
 ) -> Result<CallToolResult, ErrorData> {
+    let _work = state
+        .workloads
+        .enter()
+        .ok_or_else(|| ErrorData::internal_error("shell connection is shutting down", None))?;
+    if state.workloads.cancel.is_cancelled() {
+        return Ok(CallToolResult::error(vec![Content::text("cancelled")]));
+    }
     if p.command.len() > MAX_COMMAND_BYTES {
         return Err(ErrorData::invalid_params(
             format!("command exceeds {MAX_COMMAND_BYTES} byte limit"),
@@ -217,7 +226,12 @@ pub async fn run(
     let mut notes: Vec<String> = Vec::new();
     let (status, timed_out) = tokio::select! {
         biased;
-        _ = ct.cancelled() => {
+        _ = async {
+            tokio::select! {
+                _ = ct.cancelled() => {}
+                _ = state.workloads.cancel.cancelled() => {}
+            }
+        } => {
             // Kill process group, reap child, abort reader tasks.
             kill_group.kill_immediate();
             // Bounded reap so we don't leak zombies. If reap times out,
@@ -234,6 +248,8 @@ pub async fn run(
             }
             stdout_handle.abort();
             stderr_handle.abort();
+            let _ = stdout_handle.await;
+            let _ = stderr_handle.await;
             return Ok(CallToolResult::error(vec![Content::text("cancelled")]));
         }
         r = tokio::time::timeout(timeout_dur, child.wait()) => match r {

crates/buzz-dev-mcp/src/workloads.rs

@@ -0,0 +1,76 @@
+//! Connection-owned shell work. EOF is a cancellation boundary, not permission
+//! to drop the runtime while independently grouped shell children still run.
+use std::io;
+use std::pin::Pin;
+use std::sync::Mutex;
+use std::task::{Context, Poll};
+use tokio::io::{AsyncRead, ReadBuf};
+use tokio_util::sync::CancellationToken;
+use tokio_util::task::task_tracker::TaskTrackerToken;
+use tokio_util::task::TaskTracker;
+
+#[derive(Default)]
+pub(crate) struct Workloads {
+    closed: Mutex<bool>,
+    tasks: TaskTracker,
+    pub(crate) cancel: CancellationToken,
+}
+
+impl Workloads {
+    pub(crate) fn enter(&self) -> Option<TaskTrackerToken> {
+        let closed = self.closed.lock().ok()?;
+        if *closed {
+            None
+        } else {
+            // Admission and close are serialized: wait cannot miss a shell
+            // whose handler was scheduled just as the transport ended.
+            Some(self.tasks.token())
+        }
+    }
+
+    pub(crate) fn close(&self) {
+        if let Ok(mut closed) = self.closed.lock() {
+            *closed = true;
+        }
+        self.cancel.cancel();
+        self.tasks.close();
+    }
+
+    pub(crate) async fn drain(&self) -> io::Result<()> {
+        self.close();
+        tokio::time::timeout(std::time::Duration::from_secs(3), self.tasks.wait())
+            .await
+            .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "shell cleanup did not finish"))
+    }
+}
+
+/// Observe input EOF before rmcp's response-drain timeout. That gives the shell
+/// owner time to kill its process group, wait its child and join output readers.
+pub(crate) struct CancelOnEof<R> {
+    pub(crate) reader: R,
+    pub(crate) state: std::sync::Arc<super::shell::SharedState>,
+}
+
+impl<R: AsyncRead + Unpin> AsyncRead for CancelOnEof<R> {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        let this = self.get_mut();
+        let before = buf.filled().len();
+        let remaining = buf.remaining();
+        let result = Pin::new(&mut this.reader).poll_read(cx, buf);
+        if matches!(result, Poll::Ready(Err(_)))
+            || (matches!(result, Poll::Ready(Ok(())))
+                && remaining > 0
+                && buf.filled().len() == before)
+        {
+            this.state.workloads.close();
+        }
+        result
+    }
+}
+
+#[cfg(test)]
+mod tests;

crates/buzz-dev-mcp/src/workloads/tests.rs

@@ -0,0 +1,124 @@
+use super::*;
+
+#[tokio::test]
+async fn closed_owner_rejects_late_admission_and_waits_for_existing_work() {
+    let owner = Workloads::default();
+    let work = owner.enter().unwrap();
+    owner.close();
+    assert!(owner.enter().is_none());
+    assert!(owner.cancel.is_cancelled());
+    assert!(
+        tokio::time::timeout(std::time::Duration::from_millis(20), owner.drain())
+            .await
+            .is_err()
+    );
+    drop(work);
+    owner.drain().await.unwrap();
+    owner.drain().await.unwrap();
+}
+
+#[cfg(unix)]
+#[tokio::test]
+async fn mcp_eof_reaps_separate_shell_group_without_stopping_peer() {
+    use nix::sys::signal::{kill, Signal};
+    use nix::unistd::Pid;
+    use std::time::Duration;
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+    let dir = tempfile::tempdir().unwrap();
+    let marker = dir.path().join("children");
+    let state = std::sync::Arc::new(
+        crate::shell::SharedState::new(dir.path().into(), crate::shim::Shim::install().unwrap())
+            .unwrap(),
+    );
+    let (client, server) = tokio::io::duplex(64 * 1024);
+    let (reader, writer) = tokio::io::split(server);
+    let owned = state.clone();
+    let server_task = tokio::spawn(async move {
+        crate::serve_connection(owned, reader, writer)
+            .await
+            .map_err(|e| e.to_string())
+    });
+    let (mut reader, mut writer) = tokio::io::split(client);
+    let output = tokio::spawn(async move {
+        let mut bytes = vec![];
+        reader.read_to_end(&mut bytes).await.unwrap();
+        bytes
+    });
+    for message in [
+        serde_json::json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{
+            "protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"teardown-test","version":"1"}
+        }}),
+        serde_json::json!({"jsonrpc":"2.0","method":"notifications/initialized"}),
+        serde_json::json!({"jsonrpc":"2.0","id":2,"method":"tools/call","params":{
+            "name":"shell","arguments":{"command":format!("sleep 30 & echo \"$$ $!\" > '{}'; wait", marker.display()),"timeout_ms":10000}
+        }}),
+    ] {
+        writer
+            .write_all(format!("{message}\n").as_bytes())
+            .await
+            .unwrap();
+    }
+    let mut peer = tokio::process::Command::new("/bin/sleep")
+        .arg("30")
+        .process_group(0)
+        .kill_on_drop(true)
+        .spawn()
+        .unwrap();
+    let pids = tokio::time::timeout(Duration::from_secs(5), async {
+        loop {
+            if let Ok(text) = std::fs::read_to_string(&marker) {
+                let pids: Vec<i32> = text
+                    .split_whitespace()
+                    .filter_map(|p| p.parse().ok())
+                    .collect();
+                if pids.len() == 2 {
+                    break pids;
+                }
+            }
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
+    })
+    .await;
+    let group = pids.as_ref().ok().map(|pids| {
+        std::process::Command::new("/bin/ps")
+            .args(["-o", "pgid=", "-p", &pids[0].to_string()])
+            .output()
+            .unwrap()
+    });
+    // Always trigger cleanup before asserting fixture readiness.
+    writer.shutdown().await.unwrap();
+    drop(writer);
+    // Finish before rmcp's five-second EOF response-drain timeout: cleanup
+    // starts at EOF, not only after the transport has already given up.
+    let result = tokio::time::timeout(Duration::from_secs(4), server_task).await;
+    state.workloads.close();
+    let peer_alive = peer.try_wait().unwrap().is_none();
+    peer.kill().await.unwrap();
+    peer.wait().await.unwrap();
+    let pids = pids.unwrap();
+    let group = group.unwrap();
+    assert!(group.status.success());
+    assert_eq!(
+        String::from_utf8_lossy(&group.stdout).trim(),
+        pids[0].to_string()
+    );
+    result.unwrap().unwrap().unwrap();
+    output.await.unwrap();
+    assert!(peer_alive, "a separate placement is not a teardown target");
+    // A real shell and its child occupied their own group, not the MCP/root
+    // group. EOF must make the natural shell owner cancel and reap them.
+    for pid in pids {
+        let gone = tokio::time::timeout(Duration::from_secs(3), async {
+            while kill(Pid::from_raw(pid), None).is_ok() {
+                tokio::time::sleep(Duration::from_millis(10)).await;
+            }
+        })
+        .await;
+        if gone.is_err() {
+            let _ = kill(Pid::from_raw(pid), Signal::SIGKILL);
+        }
+        assert!(gone.is_ok(), "managed child {pid} survived MCP EOF");
+    }
+    assert!(state.workloads.enter().is_none());
+}