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.
This commit is contained in:
jif
2026-06-23 18:29:27 +01:00
committed by GitHub
parent 9fe689783d
commit 8751fd3fcb
5 changed files with 195 additions and 4 deletions
+9 -2
View File
@@ -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::<Vec<_>>();
// 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}");
}
}
@@ -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::<Result<ManagedClient, StartupOutcomeError>>()