From 8751fd3fcb8031d42b62670a6131872074635c9b Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 23 Jun 2026 18:29:27 +0100 Subject: [PATCH] Shut down superseded MCP managers on refresh (#29608) ## Summary MCP refresh replaced the published connection manager without shutting down the manager it superseded. If another task retained that old manager, its stdio MCP processes stayed alive and accumulated across refreshes. Atomically swap in the refreshed manager, then explicitly shut down the exact manager returned by the swap. Add a process-level regression test that retains the old manager during refresh and verifies its stdio process exits while the replacement remains available. ## Context Explicit cleanup was lost when manager publication moved to `ArcSwap`. Dropping the old manager is not a reliable shutdown boundary because active callers can retain its `Arc` and underlying client process handles. --- codex-rs/codex-mcp/src/connection_manager.rs | 11 +- .../codex-mcp/src/connection_manager_tests.rs | 53 ++++++++ codex-rs/core/src/session/mcp.rs | 6 +- .../core/tests/suite/mcp_refresh_cleanup.rs | 127 ++++++++++++++++++ codex-rs/core/tests/suite/mod.rs | 2 + 5 files changed, 195 insertions(+), 4 deletions(-) create mode 100644 codex-rs/core/tests/suite/mcp_refresh_cleanup.rs diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 53e9efc6d..33631393c 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -345,8 +345,15 @@ impl McpConnectionManager { /// Stop all MCP clients owned by this manager and terminate stdio server processes. pub async fn shutdown(&self) { self.startup_cancellation_token.cancel(); - for client in self.clients.values() { - client.shutdown().await; + let clients = self.clients.values().cloned().collect::>(); + // Keep cleanup alive if an interrupt cancels the refresh that requested it. + let shutdown_task = tokio::spawn(async move { + for client in clients { + client.shutdown().await; + } + }); + if let Err(error) = shutdown_task.await { + warn!("MCP client shutdown task failed: {error}"); } } diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 887a9b3d4..469318c69 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1020,6 +1020,59 @@ async fn shutdown_cancels_pending_tool_listing() { assert!(tools.is_empty()); } +#[tokio::test] +async fn shutdown_continues_after_caller_is_aborted() { + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (completed_tx, completed_rx) = tokio::sync::oneshot::channel(); + let release = Arc::new(tokio::sync::Notify::new()); + let release_for_client = Arc::clone(&release); + let blocking_client = async move { + let _ = started_tx.send(()); + release_for_client.notified().await; + let _ = completed_tx.send(()); + Err(StartupOutcomeError::Cancelled) + } + .boxed() + .shared(); + let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionManager::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + manager.clients.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + AsyncManagedClient { + client: blocking_client, + is_codex_apps_mcp_server: true, + cached_tool_info_snapshot: None, + cached_server_info: None, + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + }, + ); + let manager = Arc::new(manager); + let shutdown_task = tokio::spawn({ + let manager = Arc::clone(&manager); + async move { manager.shutdown().await } + }); + + started_rx.await.expect("client shutdown should start"); + shutdown_task.abort(); + let shutdown_error = shutdown_task + .await + .expect_err("caller shutdown task should be aborted"); + assert!(shutdown_error.is_cancelled()); + release.notify_one(); + + tokio::time::timeout(Duration::from_secs(1), completed_rx) + .await + .expect("client shutdown should survive caller cancellation") + .expect("client shutdown completion sender should stay alive"); +} + #[tokio::test] async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() { let pending_client = futures::future::pending::>() diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 42c274fdb..737335c59 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -371,9 +371,11 @@ impl Session { let current_manager = self.services.mcp_connection_manager.load_full(); refreshed_manager.set_elicitations_auto_deny(current_manager.elicitations_auto_deny()); } - self.services + let superseded_manager = self + .services .mcp_connection_manager - .store(Arc::new(refreshed_manager)); + .swap(Arc::new(refreshed_manager)); + superseded_manager.shutdown().await; } pub(crate) async fn refresh_mcp_servers_if_requested( diff --git a/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs b/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs new file mode 100644 index 000000000..a8e06bb35 --- /dev/null +++ b/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs @@ -0,0 +1,127 @@ +use std::collections::HashMap; +use std::fs; +use std::sync::Arc; +use std::time::Duration; + +use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; +use codex_config::types::McpServerConfig; +use codex_config::types::McpServerTransportConfig; +use core_test_support::process::process_is_alive; +use core_test_support::process::wait_for_pid_file; +use core_test_support::process::wait_for_process_exit; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::stdio_server_bin; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_mcp_server; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let temp_dir = tempfile::tempdir()?; + let pid_file = temp_dir.path().join("mcp.pid"); + let pid_file_for_config = pid_file.clone(); + let command = stdio_server_bin()?; + let fixture = test_codex() + .with_config(move |config| { + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + "refresh_cleanup".to_string(), + McpServerConfig { + transport: McpServerTransportConfig::Stdio { + command, + args: Vec::new(), + env: Some(HashMap::from([( + "MCP_TEST_PID_FILE".to_string(), + pid_file_for_config.to_string_lossy().into_owned(), + )])), + env_vars: Vec::new(), + cwd: None, + }, + environment_id: DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: Some(Duration::from_secs(10)), + tool_timeout_sec: None, + default_tools_approval_mode: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth: None, + oauth_resource: None, + tools: HashMap::new(), + }, + ); + config + .mcp_servers + .set(servers) + .expect("test MCP servers should accept any configuration"); + }) + .build(&server) + .await?; + wait_for_mcp_server(&fixture.codex, "refresh_cleanup").await?; + + let superseded_pid = wait_for_pid_file(&pid_file).await?; + assert!(process_is_alive(&superseded_pid)?); + + let barrier = serde_json::json!({ + "id": "mcp-refresh-cleanup", + "participants": 2, + "timeout_ms": 1_000 + }); + let long_call = tokio::spawn({ + let codex = Arc::clone(&fixture.codex); + let barrier = barrier.clone(); + async move { + codex + .call_mcp_tool( + "refresh_cleanup", + "sync", + Some(serde_json::json!({ + "barrier": barrier, + "sleep_after_ms": 30_000 + })), + /*meta*/ None, + ) + .await + } + }); + fixture + .codex + .call_mcp_tool( + "refresh_cleanup", + "sync", + Some(serde_json::json!({ "barrier": barrier })), + /*meta*/ None, + ) + .await?; + fs::remove_file(&pid_file)?; + + responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_assistant_message("msg-1", "done"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + fixture + .codex + .set_openai_form_elicitation_support(/*supported*/ true) + .await?; + fixture.submit_turn("refresh MCP servers").await?; + + let replacement_pid = wait_for_pid_file(&pid_file).await?; + assert_ne!(replacement_pid, superseded_pid); + wait_for_process_exit(&superseded_pid).await?; + assert!(process_is_alive(&replacement_pid)?); + assert!(long_call.await?.is_err()); + + fixture.codex.shutdown_and_wait().await?; + wait_for_process_exit(&replacement_pid).await +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index a9298ec6d..3d4a191f1 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -65,6 +65,8 @@ mod image_rollout; mod items; mod json_result; mod live_cli; +#[cfg(unix)] +mod mcp_refresh_cleanup; mod mcp_tool_exposure; mod mcp_turn_metadata; mod model_overrides;