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()); +}