diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 9bbcbe12e..483a82796 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -71,6 +71,7 @@ pub struct McpConnectionManager { clients: HashMap, server_origins: HashMap, elicitation_requests: ElicitationRequestManager, + startup_cancellation_token: CancellationToken, } impl McpConnectionManager { @@ -85,6 +86,7 @@ impl McpConnectionManager { approval_policy.value(), permission_profile.get().clone(), ), + startup_cancellation_token: CancellationToken::new(), } } @@ -92,6 +94,24 @@ impl McpConnectionManager { !self.clients.is_empty() } + /// Drain all MCP clients from this manager and return a future that stops + /// them and terminates their stdio server processes. + pub fn begin_shutdown(&mut self) -> impl std::future::Future + Send + 'static { + self.startup_cancellation_token.cancel(); + let clients = std::mem::take(&mut self.clients); + self.server_origins.clear(); + async move { + for client in clients.into_values() { + client.shutdown().await; + } + } + } + + /// Stop all MCP clients owned by this manager and terminate stdio server processes. + pub async fn shutdown(&mut self) { + self.begin_shutdown().await; + } + pub fn server_origin(&self, server_name: &str) -> Option<&str> { self.server_origins.get(server_name).map(String::as_str) } @@ -221,6 +241,7 @@ impl McpConnectionManager { clients, server_origins, elicitation_requests: elicitation_requests.clone(), + startup_cancellation_token: cancel_token.clone(), }; tokio::spawn(async move { let outcomes = join_set.join_all().await; @@ -609,6 +630,13 @@ impl McpConnectionManager { } } +impl Drop for McpConnectionManager { + fn drop(&mut self) { + self.startup_cancellation_token.cancel(); + self.clients.clear(); + } +} + async fn emit_update( submit_id: &str, tx_event: &Sender, diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 02c4bcc73..3fcef0c06 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -662,6 +662,7 @@ async fn list_all_tools_uses_startup_snapshot_while_client_is_pending() { startup_snapshot: Some(startup_tools), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), }, ); @@ -690,6 +691,7 @@ async fn resolve_tool_info_accepts_canonical_namespaced_tool_names() { startup_snapshot: Some(startup_tools), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), }, ); @@ -726,6 +728,7 @@ async fn list_all_tools_blocks_while_client_is_pending_without_startup_snapshot( startup_snapshot: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), }, ); @@ -750,6 +753,7 @@ async fn list_all_tools_does_not_block_when_startup_snapshot_cache_hit_is_empty( startup_snapshot: Some(Vec::new()), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), }, ); @@ -784,6 +788,7 @@ async fn list_all_tools_uses_startup_snapshot_when_client_startup_fails() { startup_snapshot: Some(startup_tools), startup_complete, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), }, ); diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 074e57c88..bdfea8932 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -60,6 +60,7 @@ use rmcp::model::Implementation; use rmcp::model::InitializeRequestParams; use rmcp::model::ProtocolVersion; use tokio_util::sync::CancellationToken; +use tracing::warn; /// MCP server capability indicating that Codex should include [`SandboxState`] /// in tool-call request `_meta` under this key. @@ -115,6 +116,7 @@ pub(crate) struct AsyncManagedClient { pub(crate) startup_snapshot: Option>, pub(crate) startup_complete: Arc, pub(crate) tool_plugin_provenance: Arc, + pub(crate) cancel_token: CancellationToken, } impl AsyncManagedClient { @@ -142,8 +144,9 @@ impl AsyncManagedClient { let startup_tool_filter = tool_filter; let startup_complete = Arc::new(AtomicBool::new(false)); let startup_complete_for_fut = Arc::clone(&startup_complete); + let cancel_token_for_fut = cancel_token.clone(); let fut = async move { - let outcome = async { + let outcome = match async { if let Err(error) = validate_mcp_server_name(&server_name) { return Err(error.into()); } @@ -158,7 +161,7 @@ impl AsyncManagedClient { ) .await?, ); - match start_server_task( + start_server_task( server_name, client, StartServerTaskParams { @@ -172,14 +175,14 @@ impl AsyncManagedClient { codex_apps_tools_cache_context, }, ) - .or_cancel(&cancel_token) .await - { - Ok(result) => result, - Err(CancelErr::Cancelled) => Err(StartupOutcomeError::Cancelled), - } } - .await; + .or_cancel(&cancel_token_for_fut) + .await + { + Ok(result) => result, + Err(CancelErr::Cancelled) => Err(StartupOutcomeError::Cancelled), + }; startup_complete_for_fut.store(true, Ordering::Release); outcome @@ -197,6 +200,7 @@ impl AsyncManagedClient { startup_snapshot, startup_complete, tool_plugin_provenance, + cancel_token, } } @@ -204,6 +208,17 @@ impl AsyncManagedClient { self.client.clone().await } + pub(crate) async fn shutdown(&self) { + self.cancel_token.cancel(); + match self.client().await { + Ok(client) => client.client.shutdown().await, + Err(StartupOutcomeError::Cancelled) => {} + Err(error) => { + warn!("failed to initialize MCP client during shutdown: {error:#}"); + } + } + } + fn startup_snapshot_while_initializing(&self) -> Option> { if !self.startup_complete.load(Ordering::Acquire) { return self.startup_snapshot.clone(); diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index be4f56097..d06ea2460 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -267,7 +267,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager( .default_environment() .unwrap_or_else(|| environment_manager.local_environment()); - let (mcp_connection_manager, cancel_token) = McpConnectionManager::new( + let (mut mcp_connection_manager, cancel_token) = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, auth_status_entries, @@ -346,6 +346,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager( } let accessible_connectors = with_app_plugin_sources(accessible_connectors, &tool_plugin_provenance); + mcp_connection_manager.shutdown().await; Ok(AccessibleConnectorsStatus { connectors: accessible_connectors, codex_apps_ready, diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 8f7255ec4..4965cdb38 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -883,6 +883,11 @@ pub async fn shutdown(sess: &Arc, sub_id: String) -> bool { .unified_exec_manager .terminate_all_processes() .await; + let mcp_shutdown = { + let mut manager = sess.services.mcp_connection_manager.write().await; + manager.begin_shutdown() + }; + mcp_shutdown.await; sess.guardian_review_session.shutdown().await; info!("Shutting down Codex instance"); let history = sess.clone_history().await; @@ -1166,8 +1171,19 @@ pub(super) async fn submission_loop( break; } } - // Also drain cached guardian state if the submission loop exits because - // the channel closed without receiving an explicit shutdown op. + // If the submission loop exits because the channel closed without an + // explicit shutdown op, still run process teardown for child processes + // owned by this session. + sess.services + .unified_exec_manager + .terminate_all_processes() + .await; + let mcp_shutdown = { + let mut manager = sess.services.mcp_connection_manager.write().await; + manager.begin_shutdown() + }; + mcp_shutdown.await; + // Also drain cached guardian state on this implicit shutdown path. sess.guardian_review_session.shutdown().await; debug!("Agent loop exited"); } diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index eae4608c4..18cc19a72 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -255,8 +255,11 @@ impl Session { *guard = cancel_token; } - let mut manager = self.services.mcp_connection_manager.write().await; - *manager = refreshed_manager; + let mut old_manager = { + let mut manager = self.services.mcp_connection_manager.write().await; + std::mem::replace(&mut *manager, refreshed_manager) + }; + old_manager.shutdown().await; } pub(crate) async fn refresh_mcp_servers_if_requested(&self, turn_context: &TurnContext) { diff --git a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs index 4d3a2ba77..cb83da5e6 100644 --- a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs @@ -755,6 +755,9 @@ fn parse_data_url(url: &str) -> Option<(String, String)> { #[tokio::main] async fn main() -> Result<(), Box> { eprintln!("starting rmcp test server"); + if let Ok(pid_file) = std::env::var("MCP_TEST_PID_FILE") { + std::fs::write(pid_file, std::process::id().to_string())?; + } // Run the server with STDIO transport. If the client disconnects we simply // bubble up the error so the process exits. let service = TestToolServer::new(); diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 5cdb1d441..6f38acadd 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -67,6 +67,7 @@ use crate::oauth::OAuthPersistor; use crate::oauth::StoredOAuthTokens; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; +use crate::stdio_server_launcher::StdioServerProcessHandle; use crate::stdio_server_launcher::StdioServerTransport; use crate::utils::apply_default_headers; use crate::utils::build_default_headers; @@ -93,6 +94,7 @@ enum ClientState { service: Arc>, oauth: Option, }, + Closed, } #[derive(Clone)] @@ -265,6 +267,7 @@ pub struct ListToolsWithConnectorIdResult { /// https://github.com/modelcontextprotocol/rust-sdk pub struct RmcpClient { state: Mutex, + stdio_process: Option, transport_recipe: TransportRecipe, initialize_context: Mutex>, session_recovery_lock: Semaphore, @@ -287,11 +290,17 @@ impl RmcpClient { let transport = Self::create_pending_transport(&transport_recipe) .await .map_err(io::Error::other)?; + let stdio_process = match &transport { + PendingTransport::Stdio { transport } => Some(transport.process_handle()), + PendingTransport::StreamableHttp { .. } + | PendingTransport::StreamableHttpWithOAuth { .. } => None, + }; Ok(Self { state: Mutex::new(ClientState::Connecting { transport: Some(transport), }), + stdio_process, transport_recipe, initialize_context: Mutex::new(None), session_recovery_lock: Semaphore::new(/*permits*/ 1), @@ -325,6 +334,7 @@ impl RmcpClient { state: Mutex::new(ClientState::Connecting { transport: Some(transport), }), + stdio_process: None, transport_recipe, initialize_context: Mutex::new(None), session_recovery_lock: Semaphore::new(/*permits*/ 1), @@ -353,6 +363,7 @@ impl RmcpClient { None => return Err(anyhow!("client already initializing")), }, ClientState::Ready { .. } => return Err(anyhow!("client already initialized")), + ClientState::Closed => return Err(anyhow!("MCP client is shut down")), } }; @@ -376,6 +387,9 @@ impl RmcpClient { { let mut guard = self.state.lock().await; + if matches!(*guard, ClientState::Closed) { + return Err(anyhow!("MCP client is shut down")); + } *guard = ClientState::Ready { service, oauth: oauth_persistor.clone(), @@ -623,6 +637,7 @@ impl RmcpClient { match &*guard { ClientState::Ready { service, .. } => Ok(Arc::clone(service)), ClientState::Connecting { .. } => Err(anyhow!("MCP client not initialized")), + ClientState::Closed => Err(anyhow!("MCP client is shut down")), } } @@ -637,6 +652,22 @@ impl RmcpClient { } } + /// Stop the MCP transport and any stdio server process owned by this client. + pub async fn shutdown(&self) { + let previous_state = { + let mut guard = self.state.lock().await; + std::mem::replace(&mut *guard, ClientState::Closed) + }; + + if let Some(process) = &self.stdio_process + && let Err(error) = process.terminate().await + { + warn!("failed to terminate MCP stdio server process: {error}"); + } + + drop(previous_state); + } + /// This should be called after every tool call so that if a given tool call triggered /// a refresh of the OAuth tokens, they are persisted. async fn persist_oauth_tokens(&self) { @@ -900,6 +931,9 @@ impl RmcpClient { ClientState::Connecting { .. } => { return Err(anyhow!("MCP client not initialized")); } + ClientState::Closed => { + return Err(anyhow!("MCP client is shut down")); + } } } @@ -919,6 +953,9 @@ impl RmcpClient { { let mut guard = self.state.lock().await; + if matches!(*guard, ClientState::Closed) { + return Err(anyhow!("MCP client is shut down")); + } *guard = ClientState::Ready { service, oauth: oauth_persistor.clone(), diff --git a/codex-rs/rmcp-client/src/stdio_server_launcher.rs b/codex-rs/rmcp-client/src/stdio_server_launcher.rs index ced594b78..fb3dab525 100644 --- a/codex-rs/rmcp-client/src/stdio_server_launcher.rs +++ b/codex-rs/rmcp-client/src/stdio_server_launcher.rs @@ -18,6 +18,8 @@ use std::io; use std::path::PathBuf; use std::process::Stdio; use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; #[cfg(unix)] use std::thread::sleep; #[cfg(unix)] @@ -31,6 +33,7 @@ use codex_config::types::McpServerEnvVar; use codex_exec_server::ExecBackend; use codex_exec_server::ExecEnvPolicy; use codex_exec_server::ExecParams; +use codex_exec_server::ExecProcess; use codex_protocol::config_types::ShellEnvironmentPolicyInherit; #[cfg(unix)] use codex_utils_pty::process_group::kill_process_group; @@ -88,11 +91,7 @@ pub struct StdioServerCommand { /// directly to `rmcp::service::serve_client`. pub struct StdioServerTransport { inner: StdioServerTransportInner, - // Local child processes can leave subprocesses behind, so the local - // variant keeps a process-group guard with the transport. Executor-backed - // processes are owned and cleaned up by the executor, so that variant uses - // `None`. - _process_group_guard: Option, + process: StdioServerProcessHandle, } enum StdioServerTransportInner { @@ -127,6 +126,7 @@ impl Transport for StdioServerTransport { } async fn close(&mut self) -> std::result::Result<(), Self::Error> { + self.process.terminate().await?; match &mut self.inner { StdioServerTransportInner::Local(transport) => transport.close().await, StdioServerTransportInner::Executor(transport) => transport.close().await, @@ -134,6 +134,12 @@ impl Transport for StdioServerTransport { } } +impl StdioServerTransport { + pub(crate) fn process_handle(&self) -> StdioServerProcessHandle { + self.process.clone() + } +} + impl StdioServerCommand { /// Build the stdio process parameters before choosing where the process /// runs. @@ -192,12 +198,33 @@ impl StdioServerLauncher for LocalStdioServerLauncher { const PROCESS_GROUP_TERM_GRACE_PERIOD: Duration = Duration::from_secs(2); #[cfg(unix)] -struct ProcessGroupGuard { +struct LocalProcessTerminator { process_group_id: u32, } -#[cfg(not(unix))] -struct ProcessGroupGuard; +#[cfg(windows)] +struct LocalProcessTerminator { + pid: u32, +} + +#[cfg(not(any(unix, windows)))] +struct LocalProcessTerminator; + +#[derive(Clone)] +pub(crate) struct StdioServerProcessHandle { + inner: Arc, +} + +struct StdioServerProcessHandleInner { + program_name: String, + kind: StdioServerProcessKind, + terminated: AtomicBool, +} + +enum StdioServerProcessKind { + Local(Option), + Executor(Arc), +} mod private { pub trait Sealed {} @@ -238,7 +265,10 @@ impl LocalStdioServerLauncher { let (transport, stderr) = TokioChildProcess::builder(command) .stderr(Stdio::piped()) .spawn()?; - let process_group_guard = transport.id().map(ProcessGroupGuard::new); + let process = StdioServerProcessHandle::local( + program_name.clone(), + transport.id().map(LocalProcessTerminator::new), + ); if let Some(stderr) = stderr { tokio::spawn(async move { @@ -260,18 +290,24 @@ impl LocalStdioServerLauncher { Ok(StdioServerTransport { inner: StdioServerTransportInner::Local(transport), - _process_group_guard: process_group_guard, + process, }) } } -impl ProcessGroupGuard { +impl LocalProcessTerminator { fn new(process_group_id: u32) -> Self { #[cfg(unix)] { Self { process_group_id } } - #[cfg(not(unix))] + #[cfg(windows)] + { + Self { + pid: process_group_id, + } + } + #[cfg(not(any(unix, windows)))] { let _ = process_group_id; Self @@ -279,7 +315,7 @@ impl ProcessGroupGuard { } #[cfg(unix)] - fn maybe_terminate_process_group(&self) { + fn terminate(&self) { let process_group_id = self.process_group_id; let should_escalate = match terminate_process_group(process_group_id) { Ok(exists) => exists, @@ -298,14 +334,93 @@ impl ProcessGroupGuard { } } - #[cfg(not(unix))] - fn maybe_terminate_process_group(&self) {} + #[cfg(windows)] + fn terminate(&self) { + let _ = std::process::Command::new("taskkill") + .arg("/PID") + .arg(self.pid.to_string()) + .arg("/T") + .arg("/F") + .status(); + } + + #[cfg(not(any(unix, windows)))] + fn terminate(&self) {} } -impl Drop for ProcessGroupGuard { +impl StdioServerProcessHandle { + fn local(program_name: String, terminator: Option) -> Self { + Self { + inner: Arc::new(StdioServerProcessHandleInner { + program_name, + kind: StdioServerProcessKind::Local(terminator), + terminated: AtomicBool::new(false), + }), + } + } + + pub(crate) fn executor(program_name: String, process: Arc) -> Self { + Self { + inner: Arc::new(StdioServerProcessHandleInner { + program_name, + kind: StdioServerProcessKind::Executor(process), + terminated: AtomicBool::new(false), + }), + } + } + + pub(crate) async fn terminate(&self) -> io::Result<()> { + if self.inner.terminated.swap(true, Ordering::AcqRel) { + return Ok(()); + } + + match &self.inner.kind { + StdioServerProcessKind::Local(Some(terminator)) => { + terminator.terminate(); + Ok(()) + } + StdioServerProcessKind::Local(None) => Ok(()), + StdioServerProcessKind::Executor(process) => match process.terminate().await { + Ok(()) => Ok(()), + Err(error) => { + self.inner.terminated.store(false, Ordering::Release); + Err(io::Error::other(error)) + } + }, + } + } +} + +impl Drop for StdioServerProcessHandleInner { fn drop(&mut self) { - if cfg!(unix) { - self.maybe_terminate_process_group(); + if self.terminated.swap(true, Ordering::AcqRel) { + return; + } + + match &self.kind { + StdioServerProcessKind::Local(Some(terminator)) => { + terminator.terminate(); + } + StdioServerProcessKind::Local(None) => {} + StdioServerProcessKind::Executor(process) => { + let process = Arc::clone(process); + let program_name = self.program_name.clone(); + let Ok(handle) = tokio::runtime::Handle::try_current() else { + warn!( + "Could not schedule remote MCP server process termination on drop ({}): no Tokio runtime is available", + self.program_name + ); + return; + }; + + std::mem::drop(handle.spawn(async move { + if let Err(error) = process.terminate().await { + warn!( + "Failed to terminate remote MCP server process on drop ({program_name}): {error}" + ); + } + })); + } } } } @@ -392,12 +507,14 @@ impl ExecutorStdioServerLauncher { .await .map_err(io::Error::other)?; + let process = + StdioServerProcessHandle::executor(program_name.clone(), Arc::clone(&started.process)); Ok(StdioServerTransport { inner: StdioServerTransportInner::Executor(ExecutorProcessTransport::new( started.process, program_name, )), - _process_group_guard: None, + process, }) } diff --git a/codex-rs/rmcp-client/tests/process_group_cleanup.rs b/codex-rs/rmcp-client/tests/process_group_cleanup.rs index 5f742f981..10d1e8ec7 100644 --- a/codex-rs/rmcp-client/tests/process_group_cleanup.rs +++ b/codex-rs/rmcp-client/tests/process_group_cleanup.rs @@ -9,8 +9,43 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; +use codex_rmcp_client::ElicitationAction; +use codex_rmcp_client::ElicitationResponse; use codex_rmcp_client::LocalStdioServerLauncher; use codex_rmcp_client::RmcpClient; +use futures::FutureExt as _; +use rmcp::model::ClientCapabilities; +use rmcp::model::Implementation; +use rmcp::model::InitializeRequestParams; +use rmcp::model::ProtocolVersion; +use serde_json::json; + +fn stdio_server_bin() -> Result { + codex_utils_cargo_bin::cargo_bin("test_stdio_server").map_err(Into::into) +} + +fn init_params() -> InitializeRequestParams { + InitializeRequestParams { + meta: None, + capabilities: ClientCapabilities { + experimental: None, + extensions: None, + roots: None, + sampling: None, + elicitation: None, + tasks: None, + }, + client_info: Implementation { + name: "codex-test".into(), + version: "0.0.0-test".into(), + title: Some("Codex rmcp shutdown test".into()), + description: None, + icons: None, + website_url: None, + }, + protocol_version: ProtocolVersion::V_2025_06_18, + } +} fn process_exists(pid: u32) -> bool { std::process::Command::new("kill") @@ -94,3 +129,67 @@ async fn drop_kills_wrapper_process_group() -> Result<()> { wait_for_process_exit(grandchild_pid).await } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shutdown_kills_initialized_stdio_server_with_in_flight_operation() -> Result<()> { + let temp_dir = tempfile::tempdir()?; + let server_pid_file = temp_dir.path().join("server.pid"); + let server_pid_file_str = server_pid_file.to_string_lossy().into_owned(); + + let client = Arc::new( + RmcpClient::new_stdio_client( + stdio_server_bin()?.into(), + Vec::::new(), + Some(HashMap::from([( + OsString::from("MCP_TEST_PID_FILE"), + OsString::from(server_pid_file_str), + )])), + &[], + /*cwd*/ None, + Arc::new(LocalStdioServerLauncher::new(std::env::current_dir()?)), + ) + .await?, + ); + + client + .initialize( + init_params(), + Some(Duration::from_secs(5)), + Box::new(|_, _| { + async { + Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({})), + meta: None, + }) + } + .boxed() + }), + ) + .await?; + + let server_pid = wait_for_pid_file(&server_pid_file).await?; + assert!( + process_exists(server_pid), + "expected MCP server process {server_pid} to be running before shutdown" + ); + + let call_client = Arc::clone(&client); + let call_task = tokio::spawn(async move { + call_client + .call_tool( + "sync".to_string(), + Some(json!({ "sleep_after_ms": 300_000 })), + /*meta*/ None, + Some(Duration::from_secs(300)), + ) + .await + }); + tokio::time::sleep(Duration::from_millis(200)).await; + + client.shutdown().await; + + wait_for_process_exit(server_pid).await?; + let _ = tokio::time::timeout(Duration::from_secs(5), call_task).await?; + Ok(()) +}