diff --git a/crates/mcp/src/stdio.rs b/crates/mcp/src/stdio.rs index a3eb426b3..174afd401 100644 --- a/crates/mcp/src/stdio.rs +++ b/crates/mcp/src/stdio.rs @@ -617,7 +617,8 @@ impl fmt::Display for McpRequestStage { } } -/// A running initialized stdio MCP client. +/// An owned stdio MCP process, initialized by `connect` or explicitly after +/// `spawn_uninitialized` when startup cleanup must retain it before any await. pub struct McpStdioClient { server_name: String, limits: McpStdioLimits, @@ -636,6 +637,9 @@ pub struct McpStdioClient { impl McpStdioClient { /// Spawn, initialize, negotiate capabilities, and send notifications/initialized. + /// Initialization errors attempt shutdown, but this convenience API cannot + /// return an unreaped client on failure/cancellation. Its error is not proof + /// of cleanup; use `spawn_uninitialized` with retained ownership when needed. pub async fn connect( spec: McpStdioServerSpec, limits: McpStdioLimits, @@ -667,6 +671,21 @@ impl McpStdioClient { async fn spawn( spec: McpStdioServerSpec, limits: McpStdioLimits, + ) -> Result { + Self::spawn_uninitialized(spec, limits) + } + + /// Spawn without awaiting the handshake, so a cleanup owner can retain the + /// client before initialization can fail or be cancelled. Call + /// `initialize_with_timeout` before discovery; on error/cancellation, shut + /// down this same client rather than spawning a replacement. + /// + /// Once the process spawns, this returns its owner even if an expected pipe + /// is missing. Initialization then fails, but cleanup can still reap the + /// actual child rather than relying on kill-on-drop for a post-spawn error. + pub fn spawn_uninitialized( + spec: McpStdioServerSpec, + limits: McpStdioLimits, ) -> Result { let redactor = spec.redactor(); let mut command = Command::new(&spec.command); @@ -688,29 +707,23 @@ impl McpStdioClient { McpErrorKind::Io(redactor.redact(&err.to_string())), ) })?; - let stdin = child.stdin.take().ok_or_else(|| { - McpClientError::new( - &spec.name, - McpPhase::Spawn, - McpErrorKind::Protocol("child stdin was not piped".into()), - ) - })?; - let stdout = child.stdout.take().ok_or_else(|| { - McpClientError::new( - &spec.name, - McpPhase::Spawn, - McpErrorKind::Protocol("child stdout was not piped".into()), - ) - })?; - let stderr = child.stderr.take().ok_or_else(|| { - McpClientError::new( - &spec.name, - McpPhase::Spawn, - McpErrorKind::Protocol("child stderr was not piped".into()), - ) - })?; + let stdin = child.stdin.take(); + let stdout = child.stdout.take(); + let stderr = child.stderr.take(); + let missing_pipe = if stdin.is_none() { + Some("stdin") + } else if stdout.is_none() { + Some("stdout") + } else if stderr.is_none() { + Some("stderr") + } else { + None + }; - let stdin: SharedStdin = Arc::new(Mutex::new(Some(Box::new(stdin)))); + let stdin: SharedStdin = + Arc::new(Mutex::new(stdin.map(|stdin| { + Box::new(stdin) as Box + }))); let diagnostics = Arc::new(Mutex::new(BoundedDiagnostics::new( spec.name.clone(), limits.max_diagnostic_lines, @@ -718,16 +731,32 @@ impl McpStdioClient { ))); let list_changes = Arc::new(Mutex::new(BoundedListChanged::new(spec.name.clone()))); let (tx, rx) = mpsc::channel(16); - let reader_task = spawn_stdout_reader( - spec.name.clone(), - stdout, - stdin.clone(), - tx, - limits.clone(), - redactor.clone(), - list_changes.clone(), - ); - let stderr_task = spawn_stderr_reader(stderr, diagnostics.clone(), limits.clone()); + // Do not return an error that discards an already spawned child. Queue + // setup failure before reader events so initialization cannot succeed. + if let Some(pipe) = missing_pipe { + tx.try_send(ReaderEvent::Error(McpClientError::new( + &spec.name, + McpPhase::Spawn, + McpErrorKind::Protocol(format!("child {pipe} was not piped")), + ))) + .expect("new response channel has capacity"); + } + let reader_task = match stdout { + Some(stdout) => spawn_stdout_reader( + spec.name.clone(), + stdout, + stdin.clone(), + tx, + limits.clone(), + redactor.clone(), + list_changes.clone(), + ), + None => tokio::spawn(async {}), + }; + let stderr_task = match stderr { + Some(stderr) => spawn_stderr_reader(stderr, diagnostics.clone(), limits.clone()), + None => tokio::spawn(async {}), + }; Ok(Self { server_name: spec.name, @@ -746,6 +775,25 @@ impl McpStdioClient { }) } + /// Perform the bounded handshake on an already retained client. Failure or + /// cancellation does not drop or shut down the client: its owner must do so. + /// Do not retry initialization on a partially initialized connection. + pub async fn initialize_with_timeout(&mut self) -> Result<(), McpClientError> { + let started = Instant::now(); + match timeout(self.limits.startup_timeout, self.initialize()).await { + Ok(result) => result, + Err(_) => Err(McpClientError::new( + &self.server_name, + McpPhase::Initialize, + McpErrorKind::Timeout { + operation: "startup".to_string(), + elapsed: started.elapsed(), + }, + ) + .with_diagnostics(self.snapshot_diagnostics().await)), + } + } + async fn initialize(&mut self) -> Result<(), McpClientError> { let result: InitializeResult = self .request( @@ -1080,6 +1128,8 @@ impl McpStdioClient { } /// Close stdin and wait for process exit, falling back to terminate and kill. + /// Errors and cancellation retain the child for another attempt on this + /// client. Only confirmed reaping releases the process handle. pub async fn shutdown(&mut self) -> Result { self.shutdown_started = true; let deadline = Instant::now() + self.limits.shutdown_timeout; @@ -1091,64 +1141,55 @@ impl McpStdioClient { .await .is_ok(); - let mut child = match self.child.take() { - Some(child) => child, - None => return Ok(ShutdownReport::already_finished()), - }; + // Keep the actual process handle in the client across every await and + // error. Taking it before wait loses retry/reaping authority on cancel. + if self.child.is_none() { + return Ok(ShutdownReport::already_finished()); + } if !closed { self.reader_task.abort(); - return self.terminate_then_kill(child).await; + return self.terminate_then_kill().await; } + let child = self.child.as_mut().expect("checked child"); if let Ok(Some(status)) = child.try_wait() { - self.reader_task.abort(); - self.stderr_task.abort(); - return Ok(ShutdownReport { - exit_status: Some(status), - terminated: false, - killed: false, - }); + return Ok(self.finish_shutdown(status, false, false)); } match timeout_at(deadline, child.wait()).await { - Ok(Ok(status)) => { - self.reader_task.abort(); - self.stderr_task.abort(); - Ok(ShutdownReport { - exit_status: Some(status), - terminated: false, - killed: false, - }) - } + Ok(Ok(status)) => Ok(self.finish_shutdown(status, false, false)), Ok(Err(err)) => Err(McpClientError::new( &self.server_name, McpPhase::Shutdown, McpErrorKind::Io(self.redactor.redact(&err.to_string())), ) .with_diagnostics(self.snapshot_diagnostics().await)), - Err(_) => self.terminate_then_kill(child).await, + Err(_) => self.terminate_then_kill().await, } } - async fn terminate_then_kill( + fn finish_shutdown( &mut self, - mut child: Child, - ) -> Result { - let mut terminated = false; - let mut killed = false; - if send_terminate(&mut child).is_ok() { - terminated = true; + status: ExitStatus, + terminated: bool, + killed: bool, + ) -> ShutdownReport { + // No await separates confirmed reaping from releasing the handle. + drop(self.child.take()); + self.reader_task.abort(); + self.stderr_task.abort(); + ShutdownReport { + exit_status: Some(status), + terminated, + killed, } + } + + async fn terminate_then_kill(&mut self) -> Result { + let child = self.child.as_mut().expect("shutdown retains child"); + let terminated = send_terminate(child).is_ok(); match timeout(self.limits.kill_timeout, child.wait()).await { - Ok(Ok(status)) => { - self.reader_task.abort(); - self.stderr_task.abort(); - Ok(ShutdownReport { - exit_status: Some(status), - terminated, - killed, - }) - } + Ok(Ok(status)) => Ok(self.finish_shutdown(status, terminated, false)), Ok(Err(err)) => Err(McpClientError::new( &self.server_name, McpPhase::Shutdown, @@ -1163,7 +1204,6 @@ impl McpStdioClient { McpErrorKind::Io(self.redactor.redact(&err.to_string())), ) })?; - killed = true; let status = timeout(self.limits.kill_timeout, child.wait()) .await .map_err(|_| { @@ -1183,13 +1223,7 @@ impl McpStdioClient { McpErrorKind::Io(self.redactor.redact(&err.to_string())), ) })?; - self.reader_task.abort(); - self.stderr_task.abort(); - Ok(ShutdownReport { - exit_status: Some(status), - terminated, - killed, - }) + Ok(self.finish_shutdown(status, terminated, true)) } } } @@ -1197,10 +1231,10 @@ impl McpStdioClient { impl Drop for McpStdioClient { fn drop(&mut self) { - if !self.shutdown_started { - if let Some(child) = &mut self.child { - let _ = child.start_kill(); - } + // An attempted shutdown is not proof of exit. Errors and cancellation + // leave the child here so Drop must still make a best-effort kill. + if let Some(child) = &mut self.child { + let _ = child.start_kill(); } self.reader_task.abort(); self.stderr_task.abort(); @@ -1784,3 +1818,70 @@ fn send_terminate(child: &mut Child) -> Result<(), ()> { #[cfg(test)] mod deadline_tests; + +#[cfg(all(test, unix))] +mod shutdown_retention_tests { + //! The real client and OS process are used; only the external server is a + //! local shell that acknowledges initialization and ignores TERM/EOF. + use super::*; + use std::future::{Future, poll_fn}; + use std::task::Poll; + + #[tokio::test] + async fn cancelled_graceful_or_terminate_wait_retries_reaping_the_same_child() { + for cancel_in_fallback in [false, true] { + let script = r#"trap '' TERM +read init || exit 1 +printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2025-11-25","capabilities":{},"serverInfo":{"name":"retained","version":"1"}}}' +while :; do :; done +"#; + let mut client = McpStdioClient::connect( + McpStdioServerSpec::new("retained", "/bin/sh").args(["-c", script]), + McpStdioLimits { + shutdown_timeout: if cancel_in_fallback { + Duration::ZERO + } else { + Duration::from_secs(30) + }, + kill_timeout: Duration::from_secs(30), + ..Default::default() + }, + ) + .await + .expect("initialize after installing TERM handler"); + let pid = client.child.as_ref().unwrap().id(); + let mut shutdown = Box::pin(client.shutdown()); + poll_fn(|cx| { + assert!(shutdown.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(shutdown); + assert_eq!(client.child.as_ref().and_then(Child::id), pid); + assert!(client.shutdown_started); + + client.limits.shutdown_timeout = Duration::from_millis(20); + client.limits.kill_timeout = Duration::from_millis(50); + let report = timeout(Duration::from_secs(2), client.shutdown()) + .await + .expect("retry bounded") + .expect("retry reaps original child"); + assert!( + report.exit_status.is_some(), + "not an already-finished receipt" + ); + assert!(report.killed, "TERM was ignored"); + assert!(client.child.is_none()); + let mut status = 0; + assert_eq!( + unsafe { libc::waitpid(pid.unwrap() as libc::pid_t, &mut status, libc::WNOHANG) }, + -1 + ); + assert_eq!( + std::io::Error::last_os_error().raw_os_error(), + Some(libc::ECHILD), + "shutdown, not a later observer, reaped the original process" + ); + } + } +}