From 921d79ef389d4918d4524247ddbf3a893eaa6c14 Mon Sep 17 00:00:00 2001 From: Logan Johnson Date: Mon, 31 Aug 2026 14:55:41 -0400 Subject: [PATCH] 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 --- crates/buzz-dev-mcp/src/lib.rs | 36 +++++- crates/buzz-dev-mcp/src/shell.rs | 18 ++- crates/buzz-dev-mcp/src/workloads.rs | 76 +++++++++++++ crates/buzz-dev-mcp/src/workloads/tests.rs | 124 +++++++++++++++++++++ 4 files changed, 247 insertions(+), 7 deletions(-) create mode 100644 crates/buzz-dev-mcp/src/workloads.rs create mode 100644 crates/buzz-dev-mcp/src/workloads/tests.rs diff --git a/crates/buzz-dev-mcp/src/lib.rs b/crates/buzz-dev-mcp/src/lib.rs index 87c3a11931..0ef12992b7 100644 --- a/crates/buzz-dev-mcp/src/lib.rs +++ b/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> { 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) +} + +async fn serve_connection( + state: Arc, + reader: R, + writer: W, +) -> Result<(), Box> +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>(()) + } + .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 diff --git a/crates/buzz-dev-mcp/src/shell.rs b/crates/buzz-dev-mcp/src/shell.rs index 7aa95b1d87..e8df67751c 100644 --- a/crates/buzz-dev-mcp/src/shell.rs +++ b/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 { + 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 = 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 { diff --git a/crates/buzz-dev-mcp/src/workloads.rs b/crates/buzz-dev-mcp/src/workloads.rs new file mode 100644 index 0000000000..f59e81f5f8 --- /dev/null +++ b/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, + tasks: TaskTracker, + pub(crate) cancel: CancellationToken, +} + +impl Workloads { + pub(crate) fn enter(&self) -> Option { + 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 { + pub(crate) reader: R, + pub(crate) state: std::sync::Arc, +} + +impl AsyncRead for CancelOnEof { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + 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; diff --git a/crates/buzz-dev-mcp/src/workloads/tests.rs b/crates/buzz-dev-mcp/src/workloads/tests.rs new file mode 100644 index 0000000000..5af122e2c6 --- /dev/null +++ b/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 = 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()); +}