From 5e3496e0ad6a85d7f57cbe3b6d576db3b317766c Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sat, 24 Oct 2026 23:12:22 +0000 Subject: [PATCH] Share eligible local stdio MCP connections across chats (#57211) ## What changed Extend `mcp_connection_sharing` to pool compatible local stdio connections that negotiate MCP `2026-07-28`. Keep legacy servers and servers that request input during startup private. Preserve each chat's tool catalog and route input requests to the calling chat. Release connections independently for each chat, shutting down the shared server after the final release, and retire pooled connections on forced refresh without interrupting other chats. ## Testing Add pool, runtime, app-server, and core scenario tests covering protocol gating, startup cancellation, refresh, final-release shutdown, and separate tool visibility and simultaneous input responses across chats sharing one process. GitOrigin-RevId: 896b107267a20ec1bffccdefca6c40fe2945a022 --- .../tests/suite/v2/mcp_connection_sharing.rs | 172 +++++ codex-rs/app-server/tests/suite/v2/mod.rs | 1 + codex-rs/codex-mcp/src/binding.rs | 104 +-- codex-rs/codex-mcp/src/binding_clients.rs | 7 +- codex-rs/codex-mcp/src/binding_tests.rs | 2 + codex-rs/codex-mcp/src/client_tool_catalog.rs | 1 + codex-rs/codex-mcp/src/connection_manager.rs | 85 ++- .../src/connection_manager/resources.rs | 13 +- .../codex-mcp/src/connection_manager_tests.rs | 27 +- codex-rs/codex-mcp/src/connection_pool.rs | 445 ++++++++++++ .../codex-mcp/src/connection_pool_tests.rs | 657 ++++++++++++++++++ codex-rs/codex-mcp/src/lib.rs | 2 + codex-rs/codex-mcp/src/rmcp_client.rs | 146 ++-- .../src/rmcp_client/managed_client.rs | 150 ++++ codex-rs/codex-mcp/src/runtime.rs | 54 +- .../src/runtime_shared_connection_tests.rs | 172 +++++ codex-rs/core/src/mcp.rs | 6 + codex-rs/core/src/session/mcp.rs | 3 +- codex-rs/core/src/session/mcp_runtime.rs | 10 +- codex-rs/core/tests/suite/scenarios.rs | 3 + .../core/tests/suite/scenarios_mcp_sharing.rs | 179 +++++ ...ls_and_model_requested_input_per_chat.snap | 68 ++ codex-rs/features/src/lib.rs | 2 +- .../rmcp-client/src/bin/test_stdio_server.rs | 30 + 24 files changed, 2197 insertions(+), 142 deletions(-) create mode 100644 codex-rs/app-server/tests/suite/v2/mcp_connection_sharing.rs create mode 100644 codex-rs/codex-mcp/src/connection_pool.rs create mode 100644 codex-rs/codex-mcp/src/connection_pool_tests.rs create mode 100644 codex-rs/codex-mcp/src/rmcp_client/managed_client.rs create mode 100644 codex-rs/codex-mcp/src/runtime_shared_connection_tests.rs create mode 100644 codex-rs/core/tests/suite/scenarios_mcp_sharing.rs create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mcp_sharing__shared_mcp_keeps_tools_and_model_requested_input_per_chat.snap diff --git a/codex-rs/app-server/tests/suite/v2/mcp_connection_sharing.rs b/codex-rs/app-server/tests/suite/v2/mcp_connection_sharing.rs new file mode 100644 index 0000000000..6558b34c55 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/mcp_connection_sharing.rs @@ -0,0 +1,172 @@ +//! Exercises protocol-gated stdio sharing through chat startup, tool dispatch, and shutdown. + +use std::collections::BTreeMap; +use std::collections::HashMap; +use std::time::Duration; + +use anyhow::Result; +use app_test_support::MockResponsesConfig; +use app_test_support::TestAppServer; +use app_test_support::create_mock_responses_server_sequence_unchecked; +use codex_app_server_protocol::ClientRequest; +use codex_app_server_protocol::ExperimentalFeatureEnablementSetParams; +use codex_app_server_protocol::ExperimentalFeatureEnablementSetResponse; +use codex_app_server_protocol::McpServerToolCallParams; +use codex_app_server_protocol::McpServerToolCallResponse; +use codex_app_server_protocol::ThreadClosedNotification; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadUnsubscribeParams; +use codex_app_server_protocol::ThreadUnsubscribeResponse; +use codex_app_server_protocol::ThreadUnsubscribeStatus; +use core_test_support::skip_if_remote; +use core_test_support::stdio_server_bin; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::TempDir; +use tokio::time::timeout; + +#[tokio::test] +async fn stdio_sharing_follows_protocol_and_enablement_for_loaded_chats() -> Result<()> { + skip_if_remote!(Ok(()), "the local stdio executable is a host-local fixture"); + for (modern, modern_without_sharing) in [(false, false), (true, false), (true, true)] { + let responses = create_mock_responses_server_sequence_unchecked(Vec::new()).await; + let codex_home = TempDir::new()?; + let command = toml::Value::String(stdio_server_bin()?); + // Exercise both pool-driven protocol selection and modern connections without pooling. + let mut config = format!("[features]\nmcp_2026_07_28 = {modern_without_sharing}\n"); + for name in ["alpha", "beta"] { + config.push_str(&format!( + "[mcp_servers.{name}]\ncommand = {command}\nenabled_tools = [\"echo\"]\n\ + [mcp_servers.{name}.env]\nMCP_TEST_DYNAMIC_SERVER_METADATA = \"1\"\n" + )); + if name == "beta" && modern { + config.push_str("CODEX_MCP_PROTOCOL_VERSION = \"2026-07-28\"\n"); + } else if name == "beta" { + config.push_str("CODEX_MCP_PROTOCOL_VERSION = \"2025-06-18\"\n"); + } + } + MockResponsesConfig::new(&responses.uri()) + .with_root_config("thread_unload_delay_secs = 0") + .with_extra_config(&config) + .write(codex_home.path())?; + let mut app = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .build_initialized() + .await?; + set_sharing(&mut app, /*enabled*/ true).await?; + let first = app.start_thread(ThreadStartParams::default()).await?.thread; + let mut first_processes = call_servers(&mut app, &first.id).await?; + assert_ne!(first_processes["alpha"], first_processes["beta"]); + let second = app.start_thread(ThreadStartParams::default()).await?.thread; + let mut second_processes = call_servers(&mut app, &second.id).await?; + for name in ["alpha", "beta"] { + assert_eq!( + first_processes[name] == second_processes[name], + name == "alpha" || modern + ); + } + if modern_without_sharing { + for sharing in [false, true] { + set_sharing(&mut app, sharing).await?; + first_processes = call_servers(&mut app, &first.id).await?; + second_processes = call_servers(&mut app, &second.id).await?; + for name in ["alpha", "beta"] { + assert_eq!(first_processes[name] == second_processes[name], sharing); + } + } + } + unload_thread(&mut app, &first.id).await?; + assert_eq!(call_servers(&mut app, &second.id).await?, second_processes); + unload_thread(&mut app, &second.id).await?; + #[cfg(unix)] + for process in first_processes.values().chain(second_processes.values()) { + core_test_support::process::wait_for_process_exit( + process + .strip_prefix("rmcp-test-process-") + .expect("the fixture identifies its PID"), + ) + .await?; + } + #[cfg(not(unix))] + { + let third = app.start_thread(ThreadStartParams::default()).await?.thread; + let third_processes = call_servers(&mut app, &third.id).await?; + for name in ["alpha", "beta"] { + assert_ne!(third_processes[name], first_processes[name]); + } + } + } + Ok(()) +} + +async fn set_sharing(app: &mut TestAppServer, enabled: bool) -> Result<()> { + let enablement = BTreeMap::from([("mcp_connection_sharing".to_string(), enabled)]); + let response: ExperimentalFeatureEnablementSetResponse = app + .request( + |request_id| ClientRequest::ExperimentalFeatureEnablementSet { + request_id, + params: ExperimentalFeatureEnablementSetParams { + enablement: enablement.clone(), + }, + }, + ) + .await?; + assert_eq!(response.enablement, enablement); + Ok(()) +} + +async fn call_servers( + app: &mut TestAppServer, + thread_id: &str, +) -> Result> { + let mut processes = HashMap::new(); + for name in ["alpha", "beta"] { + let response: McpServerToolCallResponse = app + .request(|request_id| ClientRequest::McpServerToolCall { + request_id, + params: McpServerToolCallParams { + thread_id: thread_id.to_string(), + server: name.to_string(), + tool: "echo".to_string(), + arguments: Some(json!({ "message": "ping" })), + meta: None, + }, + }) + .await?; + let result = response.structured_content.expect("echo result"); + let process = result["echo"] + .as_str() + .expect("the fixture identifies its process"); + processes.insert(name, process.to_owned()); + } + Ok(processes) +} + +async fn unload_thread(app: &mut TestAppServer, thread_id: &str) -> Result<()> { + let response: ThreadUnsubscribeResponse = app + .request(|request_id| ClientRequest::ThreadUnsubscribe { + request_id, + params: ThreadUnsubscribeParams { + thread_id: thread_id.to_string(), + }, + }) + .await?; + assert_eq!( + response, + ThreadUnsubscribeResponse { + status: ThreadUnsubscribeStatus::Unsubscribed, + } + ); + let closed: ThreadClosedNotification = timeout( + Duration::from_secs(/*secs*/ 10), + app.read_notification("thread/closed"), + ) + .await??; + assert_eq!( + closed, + ThreadClosedNotification { + thread_id: thread_id.to_string(), + } + ); + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 7532c62a64..42c6ecb14f 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -87,6 +87,7 @@ mod luna_reserve; mod marketplace_add; mod marketplace_remove; mod marketplace_upgrade; +mod mcp_connection_sharing; mod mcp_event_stream; mod mcp_protocol_default; mod mcp_resource; diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs index a5c17924dc..fb9c00f8b5 100644 --- a/codex-rs/codex-mcp/src/binding.rs +++ b/codex-rs/codex-mcp/src/binding.rs @@ -398,55 +398,63 @@ impl PreparedMcpCall { (server_timeout, requested_timeout) => server_timeout.or(requested_timeout), }; let tool_name = self.tool_info().tool.name.to_string(); - self.client - .tool_catalog - .run_with_snapshot(&self.catalog_snapshot, || async { - let (arguments, meta) = prepare().await?; - let timeout_deadline = - effective_timeout.map(|timeout| tokio::time::Instant::now() + timeout); - let add_trusted_access_context = self.connections.add_trusted_access_context( - self.tool_info(), - &self.server_metadata, - arguments.as_ref(), - meta, - ); - let meta = match effective_timeout.zip(timeout_deadline) { - Some((timeout, deadline)) => { - tokio::time::timeout_at(deadline, add_trusted_access_context) - .await - .map_err(|_| { - codex_rmcp_client::request_timeout_error("tools/call", timeout) - })? - } - None => add_trusted_access_context.await, - }; - let remaining_timeout = match effective_timeout.zip(timeout_deadline) { - Some((timeout, deadline)) => { - let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); - if remaining.is_zero() { - return Err(codex_rmcp_client::request_timeout_error("tools/call", timeout)); + let operation = + self.client + .tool_catalog + .run_with_snapshot(&self.catalog_snapshot, || async { + let (arguments, meta) = prepare().await?; + let timeout_deadline = + effective_timeout.map(|timeout| tokio::time::Instant::now() + timeout); + let add_trusted_access_context = self.connections.add_trusted_access_context( + self.tool_info(), + &self.server_metadata, + arguments.as_ref(), + meta, + ); + let meta = match effective_timeout.zip(timeout_deadline) { + Some((timeout, deadline)) => { + tokio::time::timeout_at(deadline, add_trusted_access_context) + .await + .map_err(|_| { + codex_rmcp_client::request_timeout_error("tools/call", timeout) + })? } - Some(remaining) - } - None => None, - }; - self.client - .client - .call_tool( - tool_name.clone(), - arguments, - self.connections.network_policy_meta_for_call( - &self.client, - Some(self.permission_profile()), - self.server_environment_id(), - meta, - ), - remaining_timeout, - ) - .await - .with_context(|| format!("tool call failed for `{}/{tool_name}`", self.server_name)) - }) - .await + None => add_trusted_access_context.await, + }; + let remaining_timeout = match effective_timeout.zip(timeout_deadline) { + Some((timeout, deadline)) => { + let remaining = + deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return Err(codex_rmcp_client::request_timeout_error( + "tools/call", + timeout, + )); + } + Some(remaining) + } + None => None, + }; + self.client + .call_tool( + tool_name.clone(), + arguments, + self.connections.network_policy_meta_for_call( + &self.client, + Some(self.permission_profile()), + self.server_environment_id(), + meta, + ), + remaining_timeout, + ) + .await + .with_context(|| { + format!("tool call failed for `{}/{tool_name}`", self.server_name) + }) + }); + self.client + .with_active_shared_lease(operation) + .await? .ok_or_else(|| anyhow::anyhow!( "tool call rejected because the catalog changed after `{}/{tool_name}` was prepared", self.server_name diff --git a/codex-rs/codex-mcp/src/binding_clients.rs b/codex-rs/codex-mcp/src/binding_clients.rs index b2b1c174b8..5ef9c4baf9 100644 --- a/codex-rs/codex-mcp/src/binding_clients.rs +++ b/codex-rs/codex-mcp/src/binding_clients.rs @@ -40,7 +40,6 @@ impl McpBindingClients { .client(server) .ok_or_else(|| anyhow!("MCP server '{server}' was not ready for this step"))?; managed - .client .list_resources(params, managed.tool_timeout) .await .with_context(|| format!("resources/list failed for `{server}`")) @@ -55,7 +54,6 @@ impl McpBindingClients { .client(server) .ok_or_else(|| anyhow!("MCP server '{server}' was not ready for this step"))?; managed - .client .list_resource_templates(params, managed.tool_timeout) .await .with_context(|| format!("resources/templates/list failed for `{server}`")) @@ -71,7 +69,6 @@ impl McpBindingClients { .ok_or_else(|| anyhow!("MCP server '{server}' was not ready for this step"))?; let uri = params.uri.clone(); managed - .client .read_resource(params, managed.tool_timeout) .await .with_context(|| format!("resources/read failed for `{server}` ({uri})")) @@ -88,7 +85,7 @@ impl McpBindingClients { .filter(|(server_name, _)| include_server(server_name)) { let server_name = server_name.clone(); - let client = Arc::clone(&managed.client); + let client = Arc::clone(managed); let timeout = managed.tool_timeout; join_set.spawn(async move { let resources = collect_paginated("resources/list", timeout, |params| { @@ -116,7 +113,7 @@ impl McpBindingClients { .filter(|(server_name, _)| include_server(server_name)) { let server_name = server_name.clone(); - let client = Arc::clone(&managed.client); + let client = Arc::clone(managed); let timeout = managed.tool_timeout; join_set.spawn(async move { let templates = collect_paginated("resources/templates/list", timeout, |params| { diff --git a/codex-rs/codex-mcp/src/binding_tests.rs b/codex-rs/codex-mcp/src/binding_tests.rs index d816befd1b..60d33c8a19 100644 --- a/codex-rs/codex-mcp/src/binding_tests.rs +++ b/codex-rs/codex-mcp/src/binding_tests.rs @@ -84,6 +84,8 @@ async fn test_step( /*updates*/ None, )); let managed_client = Arc::new(ManagedClient { + connection_sharing: Default::default(), + shared_elicitation: None, _auth_change_notifications: None, client: Arc::clone(&client), server_info: McpServerInfo { diff --git a/codex-rs/codex-mcp/src/client_tool_catalog.rs b/codex-rs/codex-mcp/src/client_tool_catalog.rs index 1bb85a8213..cec29a9190 100644 --- a/codex-rs/codex-mcp/src/client_tool_catalog.rs +++ b/codex-rs/codex-mcp/src/client_tool_catalog.rs @@ -69,6 +69,7 @@ enum CatalogSource { }, } +#[derive(Clone)] pub(crate) struct ToolCatalogSnapshot { pub(crate) revision: u64, pub(crate) tools: Arc<[ToolInfo]>, diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 2ad43785d2..06c835e180 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -41,6 +41,8 @@ use std::time::Duration; use crate::McpProtocolMode; use crate::binding::call_tool_result_from_rmcp; use crate::catalog::McpServerSource; +use crate::connection_pool::ConnectionKey; +use crate::connection_pool::ConnectionSharing; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationRequestRouter; use crate::event_stream::EventStreamConnectionSettings; @@ -279,8 +281,8 @@ impl McpServerView { let timeout = self.refresh_timeout; let result = client .tool_catalog - .refresh_after_notifications(&client.client.tool_list_revision(), timeout, || { - crate::rmcp_client::list_tools_for_client_uncached( + .refresh_after_notifications(&client.client.tool_list_revision(), timeout, || async { + let relist = crate::rmcp_client::list_tools_for_client_uncached( server_name, /*is_codex_apps_mcp_server*/ false, "list_changed", @@ -288,7 +290,8 @@ impl McpServerView { Some(timeout), self.catalog_item_limit, client.server_instructions.as_deref(), - ) + ); + client.with_active_shared_lease(relist.boxed()).await? }) .await; if result.is_err() { @@ -675,6 +678,12 @@ impl McpConnectionSet { expected_protocol_mode }; let protocol_mode = expected_protocol_mode.unwrap_or(protocol_mode); + let connection_pool = runtime_context.connection_pool.as_ref().filter(|_| { + // Sharing a signed-in helper requires tracking account changes. + discover_local_stdio + && protocol_mode != McpProtocolMode::Legacy + && (auth.is_none() || auth_manager.is_some()) + }); if matches!(configured_config.auth, McpServerAuth::EmaAuth) && existing_ema.is_some() { if let Some(previous_view) = existing_ema .filter(|_| elicitation_previous.is_some()) @@ -743,6 +752,10 @@ impl McpConnectionSet { && (same_owner || !previous_view.connection.may_use_legacy_fallback()) { let connection = Arc::clone(&previous_view.connection); + // A failed pool acquisition can leave this chat on its private helper. + let retained_legacy = connection.retained_legacy().is_some(); + let can_reuse = retained_legacy + || connection.client.connection_slot.is_some() == connection_pool.is_some(); let reusable_pending_startup = connection.identity.as_ref() == Some(&connection_identity) && !connection.client.startup_complete.load(Ordering::Acquire) @@ -770,15 +783,15 @@ impl McpConnectionSet { .reusable_client(&connection_identity) .await .filter(|_| previous_view.catalog_item_limit == catalog_item_limit); - if reusable_pending_startup - || unchanged_auth_failure.is_some() - || ready.as_ref().is_some_and(|client| { - expected_protocol_mode.is_some_and(|expected| { - client.client.protocol_mode() == expected - || (previous_view.protocol_mode == expected - && connection.retained_legacy().is_some()) - }) - }) + if can_reuse + && (reusable_pending_startup + || unchanged_auth_failure.is_some() + || ready.as_ref().is_some_and(|client| { + expected_protocol_mode.is_some_and(|expected| { + client.client.protocol_mode() == expected + || (previous_view.protocol_mode == expected && retained_legacy) + }) + })) { let pending_client = reusable_pending_startup.then(|| connection.client.clone()); @@ -853,6 +866,7 @@ impl McpConnectionSet { && !connection.client.cancel_token.is_cancelled() && ready.is_some_and(|client| { client.client.protocol_mode() == McpProtocolMode::Legacy + && matches!(client.connection_sharing, ConnectionSharing::Private) }) { legacy_to_upgrade = Some(connection); @@ -861,6 +875,24 @@ impl McpConnectionSet { let cancel_token = startup_cancellation_token.child_token(); let tool_catalog_cache_context = tool_catalog_cache_context.take(); let has_runtime_auth = runtime_auth_provider.is_some(); + let connection_pool = connection_pool.map(|pool| { + ( + pool.clone(), + ConnectionKey { + server_name: server_name.clone(), + identity: connection_identity.clone(), + account: codex_apps_tools_cache_key.clone(), + auth_state: auth_manager.as_ref().map(|manager| { + ( + Arc::clone(manager), + *manager.auth_change_state_receiver().borrow(), + ) + }), + startup_timeout, + catalog_item_limit, + }, + ) + }); let mut async_managed_client = AsyncManagedClient::new( ema_auth.clone(), server_name.clone(), @@ -890,6 +922,7 @@ impl McpConnectionSet { .map(|manager| manager.auth_change_state_receiver()), protocol_mode, catalog_item_limit, + connection_pool, ); // A forced reconnect must also wait for the previously published helper to exit. if let Some(previous) = existing_ema.and_then(|set| set.servers.get(&server_name)) { @@ -1192,14 +1225,29 @@ impl McpConnectionSet { || view.connection.client_for_call().await.is_ok() } - /// Stop all MCP clients owned by this manager and terminate stdio server processes. + /// Release this manager's clients, terminating shared stdio servers after their final release. pub async fn shutdown(&self) { let connections = self .servers .values() .map(|view| Arc::clone(&view.connection)) .collect::>(); - // Keep cleanup alive if an interrupt cancels the refresh that requested it. + Self::shutdown_connections(connections).await; + } + + /// Releases pool leases retained by older bindings. + pub(crate) async fn shutdown_shared_connections(&self) { + let connections = self + .servers + .values() + .filter(|view| view.connection.client.connection_slot.is_some()) + .map(|view| Arc::clone(&view.connection)) + .collect(); + Self::shutdown_connections(connections).await; + } + + /// Keeps cleanup alive if an interrupt cancels the refresh that requested it. + async fn shutdown_connections(connections: Vec>) { let shutdown_task = tokio::spawn(async move { join_all(connections.iter().map(|connection| connection.shutdown())).await; }); @@ -1214,6 +1262,14 @@ impl McpConnectionSet { } } + pub(crate) fn retire_shared_connections(&self) { + for view in self.servers.values() { + if let Some(slot) = &view.connection.client.connection_slot { + slot.retire(); + } + } + } + pub fn plugin_id_for_mcp_server_name(&self, server_name: &str) -> Option<&str> { self.tool_plugin_context .plugin_id_for_mcp_server_name(server_name) @@ -1304,7 +1360,6 @@ impl McpConnectionSet { meta.remove(ENTITLEMENT_CONTEXT_KEY); } let result: rmcp::model::CallToolResult = client - .client .call_tool( tool.to_string(), arguments, diff --git a/codex-rs/codex-mcp/src/connection_manager/resources.rs b/codex-rs/codex-mcp/src/connection_manager/resources.rs index b42f09da89..a0b10b1eff 100644 --- a/codex-rs/codex-mcp/src/connection_manager/resources.rs +++ b/codex-rs/codex-mcp/src/connection_manager/resources.rs @@ -46,14 +46,13 @@ impl McpConnectionSet { .filter(|(server_name, _)| include_server(server_name)) { let server_name = server_name.clone(); - let Ok(managed_client) = view.connection.client().await else { + let Ok(client) = view.connection.client().await else { continue; }; let timeout = view.tool_timeout; - let client = managed_client.client; join_set.spawn(async move { let resources = collect_paginated("resources/list", timeout, |params| { - let client = Arc::clone(&client); + let client = client.clone(); async move { let response = client.list_resources(params, timeout).await?; Ok((response.resources, response.next_cursor)) @@ -93,14 +92,13 @@ impl McpConnectionSet { .filter(|(server_name, _)| include_server(server_name)) { let server_name = server_name.clone(); - let Ok(managed_client) = view.connection.client().await else { + let Ok(client) = view.connection.client().await else { continue; }; let timeout = view.tool_timeout; - let client = managed_client.client; join_set.spawn(async move { let templates = collect_paginated("resources/templates/list", timeout, |params| { - let client = Arc::clone(&client); + let client = client.clone(); async move { let response = client.list_resource_templates(params, timeout).await?; Ok((response.resource_templates, response.next_cursor)) @@ -137,7 +135,6 @@ impl McpConnectionSet { ) -> Result { let (managed, timeout) = self.client_by_name(server).await?; managed - .client .list_resources(params, timeout) .await .with_context(|| format!("resources/list failed for `{server}`")) @@ -150,7 +147,6 @@ impl McpConnectionSet { ) -> Result { let (managed, timeout) = self.client_by_name(server).await?; managed - .client .list_resource_templates(params, timeout) .await .with_context(|| format!("resources/templates/list failed for `{server}`")) @@ -164,7 +160,6 @@ impl McpConnectionSet { let (managed, timeout) = self.client_by_name(server).await?; let uri = params.uri.clone(); managed - .client .read_resource(params, timeout) .await .with_context(|| format!("resources/read failed for `{server}` ({uri})")) diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 08323befbd..1eaa51d8ed 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -187,7 +187,7 @@ impl McpConnectionSet { } } -fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { +pub(crate) fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { ToolInfo { server_name: server_name.to_string(), supports_parallel_tool_calls: false, @@ -485,6 +485,8 @@ async fn legacy_tool_catalog_rejects_repeated_pagination_cursor() -> anyhow::Res async fn create_test_managed_client(tools: Vec) -> ManagedClient { ManagedClient { + connection_sharing: Default::default(), + shared_elicitation: None, _auth_change_notifications: None, client: Arc::new( RmcpClient::new_in_process_client(Arc::new(TestInProcessTransportFactory)) @@ -597,6 +599,7 @@ pub(crate) async fn create_ready_async_managed_client(tools: Vec) -> A codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), } @@ -770,6 +773,7 @@ fn create_gated_async_managed_client( codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, startup_complete, + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -808,6 +812,8 @@ pub(crate) async fn create_test_manager_with_ready_apps_client( .await?; let managed_client = ManagedClient { + connection_sharing: Default::default(), + shared_elicitation: None, _auth_change_notifications: None, client, server_info: create_test_server_info("Codex Apps"), @@ -844,6 +850,7 @@ pub(crate) async fn create_test_manager_with_ready_apps_client( codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -899,6 +906,7 @@ fn create_test_manager_with_failed_apps_startup( codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + connection_slot: None, startup_reconnect: Some(Arc::new(CodexAppsStartupReconnect::new(reconnect_factory))), cancel_token: CancellationToken::new(), }, @@ -2351,6 +2359,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -2424,6 +2433,7 @@ async fn capture_binding_preserves_ready_client_catalog_and_shares_prepared_meta codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -2821,6 +2831,7 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -2919,6 +2930,7 @@ async fn capture_binding_exposes_cached_tools_before_startup() { codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, startup_complete, + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -3022,6 +3034,7 @@ async fn capture_binding_skips_pending_optional_servers_after_configured_shared_ codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, startup_complete: Arc::new(AtomicBool::new(false)), + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -3277,6 +3290,7 @@ async fn capture_binding_shares_optional_startup_grace_across_connection_sets() codex_apps_tools_cache_context: None, tool_catalog_cache_context: Some(cache_context.clone()), startup_complete: Arc::new(AtomicBool::new(false)), + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -3967,6 +3981,7 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -4026,6 +4041,7 @@ async fn shutdown_cancels_pending_tool_listing() { codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + connection_slot: None, startup_reconnect: None, cancel_token, }, @@ -4074,6 +4090,7 @@ async fn shutdown_continues_after_caller_is_aborted() { codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -4128,6 +4145,7 @@ async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() { codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -4184,6 +4202,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, startup_complete, + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -5357,7 +5376,7 @@ fn reusable_server_runtime_context() -> McpRuntimeContext { ) } -fn reusable_server_identity( +pub(crate) fn reusable_server_identity( server_name: &str, config: &McpServerConfig, runtime_context: &McpRuntimeContext, @@ -5815,6 +5834,8 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a ) .await?; let managed_client = ManagedClient { + connection_sharing: Default::default(), + shared_elicitation: None, _auth_change_notifications: None, client: Arc::clone(&client), server_info: create_test_server_info("Mutable tools"), @@ -5857,6 +5878,7 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, @@ -6693,6 +6715,7 @@ async fn reconciliation_replaces_closed_connections() -> anyhow::Result<()> { codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + connection_slot: None, startup_reconnect: None, cancel_token: CancellationToken::new(), }, diff --git a/codex-rs/codex-mcp/src/connection_pool.rs b/codex-rs/codex-mcp/src/connection_pool.rs new file mode 100644 index 0000000000..a44aba1184 --- /dev/null +++ b/codex-rs/codex-mcp/src/connection_pool.rs @@ -0,0 +1,445 @@ +//! Shares MCP 2026-07-28 connections while keeping each chat's ownership separate. +//! +//! The protocol requires request-scoped context and explicit handles for application state. +//! Request continuations use the calling chat's elicitation authority. Legacy connections and +//! servers that request input during startup remain private. + +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::OnceLock; +use std::sync::Weak; +use std::time::Duration; + +use codex_async_utils::OrCancelExt; +use codex_connectors::ConnectorRuntimeContextKey; +use codex_login::AuthChangeState; +use codex_login::AuthManager; +use codex_rmcp_client::ElicitationResponse; +use codex_rmcp_client::SendElicitation; +use futures::FutureExt; +use rmcp::model::ElicitationAction; +use rmcp::model::ProtocolVersion; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +use crate::client_tool_catalog::ClientToolCatalog; +use crate::client_tool_catalog::ToolCatalogSnapshot; +use crate::rmcp_client::ManagedClient; +use crate::rmcp_client::ManagedClientFuture; +use crate::rmcp_client::StartupOutcomeError; +use crate::rmcp_client::list_tools_for_client_uncached; +use crate::server::McpServerConnectionIdentity; +use crate::tool_catalog_cache::McpToolCatalogCacheContext; + +/// Process-owned registry for compatible local stdio servers using MCP 2026-07-28. +#[derive(Clone, Default)] +pub struct McpConnectionPool { + entries: Arc>>>, +} + +#[derive(Clone)] +pub(crate) struct ConnectionKey { + pub(crate) server_name: String, + pub(crate) identity: McpServerConnectionIdentity, + pub(crate) account: ConnectorRuntimeContextKey, + pub(crate) auth_state: Option<(Arc, AuthChangeState)>, + pub(crate) startup_timeout: Duration, + pub(crate) catalog_item_limit: usize, +} + +impl PartialEq for ConnectionKey { + fn eq(&self, other: &Self) -> bool { + self.server_name == other.server_name + && self.identity == other.identity + && self.account == other.account + && self.startup_timeout == other.startup_timeout + && self.catalog_item_limit == other.catalog_item_limit + && match (&self.auth_state, &other.auth_state) { + (Some((left, left_state)), Some((right, right_state))) => { + Arc::ptr_eq(left, right) + && left_state.owner_generation == right_state.owner_generation + && left_state.source_generation == right_state.source_generation + && left_state.login_generation == right_state.login_generation + } + (None, None) => true, + (Some(_), None) | (None, Some(_)) => false, + } + } +} + +struct Entry { + key: ConnectionKey, + client: ManagedClientFuture, + capabilities: Arc>>, + // Share immutable definitions; every chat owns its catalog revisions and refresh locks. + catalog: Mutex>, + cancellation: CancellationToken, + /// Wakes joiners without cancelling the creator's startup. + sharing_declined: CancellationToken, + owners: Weak, + pool: Weak>>>, +} + +/// One strong reference per active chat lease, independent of retained client bindings. +/// Every strong reference must be consumed by `ConnectionLease::release`. +struct LeaseOwners; + +#[derive(Clone, Copy, PartialEq, Eq)] +enum AcquisitionRole { + Creator, + Joiner, +} + +impl Entry { + fn remove_from_pool(self: &Arc) { + if let Some(pool) = self.pool.upgrade() { + pool.lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .retain(|entry| !entry.ptr_eq(&Arc::downgrade(self))); + } + } + + async fn tool_catalog_for_chat( + &self, + client: &ManagedClient, + role: AcquisitionRole, + cache_context: Option<&McpToolCatalogCacheContext>, + ) -> Result, StartupOutcomeError> { + // Listing must not update another chat's catalog or wait for its calls. + let listed = if role == AcquisitionRole::Creator { + client.tool_catalog.read(|snapshot| snapshot).await + } else { + let fetch_ticket = cache_context.map(|context| (context, context.begin_fetch())); + let (tools, meta) = list_tools_for_client_uncached( + &self.key.server_name, + self.key.server_name == crate::CODEX_APPS_MCP_SERVER_NAME, + "shared_connection", + &client.client, + Some(self.key.startup_timeout), + self.key.catalog_item_limit, + client.server_instructions.as_deref(), + ) + .await?; + if let Some((context, mut ticket)) = fetch_ticket { + if context.track_tool_list_changes { + ticket.tool_list_revision = Some(client.client.tool_list_revision()); + } + context.publish_if_newest(ticket, &tools, meta.clone()); + } + ToolCatalogSnapshot { + revision: 0, + tools: tools.into(), + meta, + } + }; + let mut catalog = self + .catalog + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let catalog = catalog.get_or_insert_with(|| listed.clone()); + if catalog.tools != listed.tools || catalog.meta != listed.meta { + *catalog = listed; + } + Ok(Arc::new(ClientToolCatalog::new( + Arc::clone(&catalog.tools), + catalog.meta.clone(), + /*updates*/ None, + ))) + } +} + +/// A chat releases its lease explicitly on shutdown, even while bindings retain its client. +pub(crate) struct ConnectionLease { + entry: Arc, + owner: Mutex>>, + cancellation: CancellationToken, + slot: Arc, +} + +/// Tracks an acquisition without retaining its lease, including before startup completes. +#[derive(Default)] +pub(crate) struct ConnectionSlot { + state: Mutex, + shutdown_tasks: OnceLock, +} + +#[derive(Default)] +enum SlotState { + #[default] + Available, + Acquired(Weak), + Retired, +} + +impl ConnectionSlot { + pub(crate) fn retire(&self) { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let SlotState::Acquired(entry) = std::mem::replace(&mut *state, SlotState::Retired) + && let Some(entry) = entry.upgrade() + { + entry.remove_from_pool(); + } + } + + pub(crate) async fn wait_for_shutdown(&self) { + if let Some(tasks) = self.shutdown_tasks.get() { + tasks.close(); + tasks.wait().await; + } + } +} + +#[derive(Clone, Default)] +pub(crate) enum ConnectionSharing { + #[default] + Private, + Negotiated, + /// The server declined sharing, but its creator still owns the startup entry. + Owned(Arc), + Shared(Arc), +} + +impl McpConnectionPool { + pub(crate) async fn connect( + &self, + slot: Arc, + key: ConnectionKey, + capabilities: Arc>>, + tool_catalog_cache_context: Option, + cancellation: CancellationToken, + start: impl Fn(CancellationToken, Option) -> ManagedClientFuture + Sync, + ) -> Result { + let (lease, role) = self + .acquire(key, &slot, Arc::clone(&capabilities), &cancellation, &start) + .or_cancel(&cancellation) + .await + .unwrap_or(Err(StartupOutcomeError::Cancelled))?; + let outcome = tokio::select! { + outcome = lease.entry.client.clone() => Some(outcome), + () = lease.entry.sharing_declined.cancelled(), if role == AcquisitionRole::Joiner => None, + () = cancellation.cancelled() => return Err(StartupOutcomeError::Cancelled), + }; + match outcome { + Some(Ok(mut client)) + if role == AcquisitionRole::Creator + || matches!(client.connection_sharing, ConnectionSharing::Negotiated) => + { + let shared_capabilities = lease + .entry + .capabilities + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + *capabilities + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = shared_capabilities; + client.connection_sharing = + if matches!(client.connection_sharing, ConnectionSharing::Negotiated) { + client.client.enable_request_cancellation(); + client.tool_catalog = lease + .entry + .tool_catalog_for_chat( + &client, + role, + tool_catalog_cache_context.as_ref(), + ) + .or_cancel(&cancellation) + .await + .unwrap_or(Err(StartupOutcomeError::Cancelled))?; + ConnectionSharing::Shared(lease) + } else { + lease.entry.remove_from_pool(); + ConnectionSharing::Owned(lease) + }; + Ok(client) + } + Some(Err(error)) if role == AcquisitionRole::Creator => Err(error), + _ => { + // Older or failed servers leave other chats to initialize privately + // with their own elicitation authority. + drop(lease); + start(cancellation, /*sharing_declined*/ None).await + } + } + } + + async fn acquire( + &self, + key: ConnectionKey, + slot: &Arc, + capabilities: Arc>>, + cancellation: &CancellationToken, + start: impl Fn(CancellationToken, Option) -> ManagedClientFuture, + ) -> Result<(Arc, AcquisitionRole), StartupOutcomeError> { + loop { + let (lease, role) = { + let mut state = slot + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if matches!(&*state, SlotState::Retired) { + return Err(StartupOutcomeError::Cancelled); + } + let mut entries = self + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + entries.retain(|entry| entry.strong_count() > 0); + let lease_cancellation = cancellation.child_token(); + let existing = entries + .iter() + .filter_map(Weak::upgrade) + .filter(|entry| { + entry.key == key + && !entry.sharing_declined.is_cancelled() + && entry.client.peek().is_none_or(|result| { + result.as_ref().is_ok_and(|client| { + matches!( + client.connection_sharing, + ConnectionSharing::Negotiated + ) + }) + }) + }) + .find_map(|entry| entry.owners.upgrade().map(|owner| (entry, owner))); + let (entry, owner, role) = if let Some((entry, owner)) = existing { + (entry, owner, AcquisitionRole::Joiner) + } else { + let startup_cancellation = CancellationToken::new(); + let sharing_declined = CancellationToken::new(); + let client = + start(startup_cancellation.clone(), Some(sharing_declined.clone())); + // Startup belongs to the connection, including if its first chat is cancelled. + tokio::spawn(client.clone()); + let owner = Arc::new(LeaseOwners); + let entry = Arc::new(Entry { + key: key.clone(), + client, + capabilities: Arc::clone(&capabilities), + catalog: Mutex::new(None), + cancellation: startup_cancellation, + sharing_declined, + owners: Arc::downgrade(&owner), + pool: Arc::downgrade(&self.entries), + }); + entries.push(Arc::downgrade(&entry)); + (entry, owner, AcquisitionRole::Creator) + }; + let lease = Arc::new(ConnectionLease { + entry, + owner: Mutex::new(Some(owner)), + cancellation: lease_cancellation, + slot: Arc::clone(slot), + }); + *state = SlotState::Acquired(Arc::downgrade(&lease.entry)); + (lease, role) + }; + if let Some(Ok(client)) = lease.entry.client.peek() + && client.client.is_closed().await + { + lease.entry.remove_from_pool(); + lease.release(); + continue; + } + return Ok((lease, role)); + } + } +} + +impl ConnectionLease { + pub(crate) fn cancellation(&self) -> &CancellationToken { + &self.cancellation + } + + /// Returns the slot tracking shutdown when this chat releases the last lease. + pub(crate) fn release(&self) -> Option<&ConnectionSlot> { + self.cancellation.cancel(); + let mut owner = self + .owner + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Arc::into_inner(owner.take()?)?; + self.entry.cancellation.cancel(); + self.entry.remove_from_pool(); + let client = self.entry.client.clone(); + let runtime = tokio::runtime::Handle::try_current().ok()?; + // The chat must be able to await its final release even when cancellation drops the lease. + let tasks = self.slot.shutdown_tasks.get_or_init(TaskTracker::new); + drop(tasks.spawn_on( + async move { + match client.await { + Ok(client) => client.client.shutdown().await, + Err(error) => error.wait_for_shutdown().await, + } + }, + &runtime, + )); + Some(&self.slot) + } +} + +impl Drop for ConnectionLease { + fn drop(&mut self) { + let _ = self.release(); + } +} + +/// Releases waiting chats when startup becomes private and detaches callbacks after sharing. +pub(crate) struct SharingNegotiation { + callback: Arc>>, + sharing_declined: CancellationToken, +} + +impl SharingNegotiation { + pub(crate) fn new( + sender: SendElicitation, + sharing_declined: CancellationToken, + ) -> (Self, SendElicitation) { + let callback = Arc::new(Mutex::new(Some(sender))); + let negotiation = Self { + callback: Arc::clone(&callback), + sharing_declined: sharing_declined.clone(), + }; + let sender: SendElicitation = Box::new(move |id, request| { + let callback = callback + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(callback) = callback.as_ref() { + // Client input binds this startup to its original chat. + sharing_declined.cancel(); + callback(id, request) + } else { + futures::future::ready(Ok(ElicitationResponse { + action: ElicitationAction::Cancel, + content: None, + meta: None, + })) + .boxed() + } + }); + (negotiation, sender) + } + + pub(crate) fn finish(self, protocol_version: &ProtocolVersion) -> ConnectionSharing { + let mut callback = self + .callback + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if *protocol_version == ProtocolVersion::V_2026_07_28 + && !self.sharing_declined.is_cancelled() + { + callback.take(); + ConnectionSharing::Negotiated + } else { + self.sharing_declined.cancel(); + ConnectionSharing::Private + } + } +} + +#[cfg(test)] +#[path = "connection_pool_tests.rs"] +pub(crate) mod tests; diff --git a/codex-rs/codex-mcp/src/connection_pool_tests.rs b/codex-rs/codex-mcp/src/connection_pool_tests.rs new file mode 100644 index 0000000000..9f330aa35c --- /dev/null +++ b/codex-rs/codex-mcp/src/connection_pool_tests.rs @@ -0,0 +1,657 @@ +//! Tests for shared MCP connection ownership and protocol negotiation. + +use super::*; +use crate::binding_clients::McpBindingClients; +use crate::connection_manager::tests::create_ready_async_managed_client; +use crate::connection_manager::tests::create_test_tool; +use crate::connection_manager::tests::reusable_server_identity; +use crate::runtime::McpRuntimeContext; +use codex_config::McpServerConfig; +use codex_exec_server::EnvironmentManager; +use codex_login::CodexAuth; +use codex_rmcp_client::Elicitation; +use codex_rmcp_client::InProcessTransportFactory; +use codex_rmcp_client::RmcpClient; +use futures::future::BoxFuture; +use pretty_assertions::assert_eq; +use rmcp::ServerHandler; +use rmcp::ServiceExt; +use rmcp::model::ClientCapabilities; +use rmcp::model::Implementation; +use rmcp::model::InitializeRequestParams; +use rmcp::model::ListResourceTemplatesResult; +use rmcp::model::ListResourcesResult; +use rmcp::model::ListToolsResult; +use rmcp::model::PaginatedRequestParams; +use rmcp::model::ProtocolVersion; +use rmcp::model::ServerCapabilities; +use rmcp::model::ServerConfig; +use rmcp::model::Tool; +use rmcp::service::RequestContext; +use rmcp::service::RoleServer; +use serde_json::json; +use std::collections::HashMap; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use tokio::sync::Notify; +use tokio::sync::Semaphore; + +#[derive(Clone)] +pub(crate) struct TestServer { + tools: Arc>>, + block_listing: Arc, + started_listings: Arc, + cancelled_listings: Arc, +} + +impl Default for TestServer { + fn default() -> Self { + Self { + tools: Arc::new(Mutex::new(vec![ + create_test_tool("test_server", "test").tool, + ])), + block_listing: Arc::default(), + started_listings: Arc::new(Semaphore::new(/*permits*/ 0)), + cancelled_listings: Arc::new(Semaphore::new(/*permits*/ 0)), + } + } +} + +impl ServerHandler for TestServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_resources() + .build(), + ) + } + + async fn call_tool( + &self, + _request: rmcp::model::CallToolRequestParams, + _context: RequestContext, + ) -> Result { + Ok(rmcp::model::CallToolResult::success(Vec::new()).into()) + } + + async fn list_tools( + &self, + _request: Option, + context: RequestContext, + ) -> Result { + if self.block_listing.load(Ordering::Acquire) { + context.ct.cancelled().await; + self.cancelled_listings.add_permits(1); + return Err(rmcp::ErrorData::internal_error( + "listing cancelled", + /*data*/ None, + )); + } + Ok(ListToolsResult::with_all_items( + self.tools.lock().unwrap().clone(), + )) + } + + async fn list_resources( + &self, + _request: Option, + context: RequestContext, + ) -> Result { + if self.block_listing.load(Ordering::Acquire) { + self.started_listings.add_permits(1); + context.ct.cancelled().await; + self.cancelled_listings.add_permits(1); + } + Ok(ListResourcesResult::with_all_items(Vec::new())) + } + + async fn list_resource_templates( + &self, + _request: Option, + context: RequestContext, + ) -> Result { + if self.block_listing.load(Ordering::Acquire) { + self.started_listings.add_permits(1); + context.ct.cancelled().await; + self.cancelled_listings.add_permits(1); + } + Ok(ListResourceTemplatesResult::with_all_items(Vec::new())) + } +} + +impl InProcessTransportFactory for TestServer { + fn open(&self) -> BoxFuture<'static, std::io::Result> { + let server = self.clone(); + async move { + let (client, transport) = tokio::io::duplex(4096); + tokio::spawn(async move { + server + .serve(transport) + .await + .unwrap() + .waiting() + .await + .unwrap(); + }); + Ok(client) + } + .boxed() + } +} + +pub(crate) fn key(context: &McpRuntimeContext, config: &McpServerConfig) -> ConnectionKey { + ConnectionKey { + server_name: "test_server".to_string(), + identity: reusable_server_identity("test_server", config, context), + account: ConnectorRuntimeContextKey::personal( + /*account_id*/ None, /*chatgpt_user_id*/ None, + ), + auth_state: None, + startup_timeout: Duration::from_secs(30), + catalog_item_limit: 100, + } +} + +pub(crate) fn context() -> McpRuntimeContext { + McpRuntimeContext::new( + Arc::new(EnvironmentManager::default_for_tests()), + std::env::temp_dir(), + ) +} + +pub(crate) fn config() -> McpServerConfig { + serde_json::from_value(json!({"command": "test-server", "args": ["server.mjs"]})).unwrap() +} + +pub(crate) fn startup( + sharing: ConnectionSharing, + starts: Arc, + first_started: Arc, + release_first: Arc, + server: TestServer, +) -> impl Fn(CancellationToken, Option) -> ManagedClientFuture { + move |cancellation, _sharing_declined| { + let sharing = sharing.clone(); + let starts = Arc::clone(&starts); + let first_started = Arc::clone(&first_started); + let release_first = Arc::clone(&release_first); + let server = server.clone(); + async move { + async { + if starts.fetch_add(1, Ordering::SeqCst) == 0 { + first_started.notify_one(); + release_first.notified().await; + } + let mut client = create_ready_async_managed_client(Vec::new()) + .await + .client() + .await?; + client.client = Arc::new( + RmcpClient::new_in_process_client(Arc::new(server)) + .await + .expect("create test MCP client"), + ); + client + .client + .initialize( + InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("test", "1"), + ) + .with_protocol_version(ProtocolVersion::V_2025_06_18), + Some(Duration::from_secs(30)), + Box::new(|_, _| { + futures::future::ready(Ok(ElicitationResponse { + action: ElicitationAction::Decline, + content: None, + meta: None, + })) + .boxed() + }), + ) + .await?; + let (tools, meta) = list_tools_for_client_uncached( + "test_server", + /*is_codex_apps_mcp_server*/ false, + "test", + &client.client, + /*timeout*/ None, + /*catalog_item_limit*/ 100, + /*server_instructions*/ None, + ) + .await?; + client.tool_catalog = + Arc::new(ClientToolCatalog::new(tools, meta, /*updates*/ None)); + client.connection_sharing = sharing; + Ok(client) + } + .or_cancel(&cancellation) + .await + .unwrap_or(Err(StartupOutcomeError::Cancelled)) + } + .boxed() + .shared() + } +} + +#[tokio::test] +async fn shared_startup_survives_first_chat_cancellation_and_rejects_released_bindings() { + let pool = McpConnectionPool::default(); + let key = key(&context(), &config()); + let server = TestServer::default(); + let starts = Arc::new(AtomicUsize::new(0)); + let started = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let launch = startup( + ConnectionSharing::Negotiated, + Arc::clone(&starts), + Arc::clone(&started), + Arc::clone(&release), + server.clone(), + ); + let first_cancel = CancellationToken::new(); + let first = pool.connect( + Arc::default(), + key.clone(), + Arc::default(), + /*tool_catalog_cache_context*/ None, + first_cancel.clone(), + &launch, + ); + let second = pool.connect( + Arc::default(), + key.clone(), + Arc::default(), + /*tool_catalog_cache_context*/ None, + CancellationToken::new(), + &launch, + ); + tokio::pin!(first, second); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut second).is_pending()); + started.notified().await; + first_cancel.cancel(); + assert!(matches!(first.await, Err(StartupOutcomeError::Cancelled))); + release.notify_one(); + let client = second.await.unwrap(); + assert_eq!(starts.load(Ordering::SeqCst), 1); + assert!(!client.client.is_closed().await); + + let other = pool + .connect( + Arc::default(), + key, + Arc::default(), + /*tool_catalog_cache_context*/ None, + CancellationToken::new(), + &launch, + ) + .await + .unwrap(); + let mut config = crate::mcp::tests::test_mcp_config(std::env::temp_dir()); + config + .server_permission_profiles + .insert("test_server".to_string(), config.permission_profile.clone()); + let prepared = crate::binding::PreparedMcpCall::new( + Arc::new(crate::connection_manager::McpConnectionSet::empty( + /*prefix_mcp_tool_names*/ true, + )), + Arc::new(client.clone()), + Arc::new(config), + client.tool_catalog.read(Arc::new).await, + create_test_tool("test_server", "test"), + crate::server::McpServerMetadata { + environment_id: "local".to_string(), + pollutes_memory: false, + origin: None, + supports_parallel_tool_calls: true, + default_tools_approval_mode: None, + tool_approval_modes: Default::default(), + }, + /*plugin_id*/ None, + /*selected_plugin_server*/ false, + ) + .unwrap(); + let bindings = McpBindingClients::new(HashMap::from([( + "test_server".to_string(), + Arc::new(client.clone()), + )])); + server.block_listing.store(true, Ordering::Release); + let waiting = async { + tokio::join!( + bindings.list_resources("test_server", /*params*/ None), + bindings.list_resource_templates("test_server", /*params*/ None), + ) + }; + tokio::pin!(waiting); + tokio::select! { + _ = &mut waiting => panic!("resource listings must wait for cancellation"), + permits = server.started_listings.acquire_many(2) => permits.unwrap().forget(), + } + let mut owned = create_ready_async_managed_client(Vec::new()).await; + owned.client = futures::future::ready(Ok(client.clone())).boxed().shared(); + owned.shutdown().await; + let (resources, templates) = waiting.await; + assert!(resources.is_err()); + assert!(templates.is_err()); + tokio::time::timeout( + Duration::from_secs(1), + server.cancelled_listings.acquire_many(2), + ) + .await + .expect("released resource listings must be cancelled") + .unwrap() + .forget(); + let prepared_error = prepared + .call_with_preparation(/*requested_timeout*/ None, || async { + panic!("released bindings must not perform irreversible preparation") + }) + .await + .unwrap_err(); + assert!( + prepared_error + .to_string() + .contains("lease has been released") + ); + other + .call_tool( + "test".to_string(), + /*arguments*/ None, + |_| Ok(None), + /*timeout*/ None, + ) + .await + .unwrap(); + owned.client = futures::future::ready(Ok(other)).boxed().shared(); + owned.shutdown().await; +} + +#[tokio::test] +async fn callback_during_startup_releases_waiting_chats_without_sharing_authority() { + let pool = McpConnectionPool::default(); + let key = key(&context(), &config()); + let starts = Arc::new(AtomicUsize::new(0)); + let started = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let second_callback = Arc::new(Notify::new()); + let release_second = Arc::new(Notify::new()); + let callbacks = Arc::new(Mutex::new(Vec::new())); + let launch = |name| { + let launch = startup( + ConnectionSharing::Private, + Arc::clone(&starts), + Arc::clone(&started), + Arc::clone(&release), + TestServer::default(), + ); + let callbacks = Arc::clone(&callbacks); + let second_callback = Arc::clone(&second_callback); + let release_second = Arc::clone(&release_second); + move |cancellation, sharing_declined: Option| { + let startup = launch(cancellation, sharing_declined.clone()); + let callbacks = Arc::clone(&callbacks); + let second_callback = Arc::clone(&second_callback); + let release_second = Arc::clone(&release_second); + async move { + let mut client = startup.await?; + let (negotiation, callback) = SharingNegotiation::new( + Box::new(move |_, _| { + callbacks.lock().unwrap().push(name); + futures::future::ready(Ok(ElicitationResponse { + action: ElicitationAction::Decline, + content: None, + meta: None, + })) + .boxed() + }), + sharing_declined.unwrap_or_default(), + ); + callback( + rmcp::model::RequestId::Number(0), + Elicitation::OpenAiForm { + meta: None, + message: "startup prompt".to_string(), + requested_schema: json!({}), + }, + ) + .await + .unwrap(); + if name == "first" { + second_callback.notified().await; + } else { + second_callback.notify_one(); + release_second.notified().await; + } + client.connection_sharing = negotiation.finish(&ProtocolVersion::V_2026_07_28); + Ok(client) + } + .boxed() + .shared() + } + }; + let first = pool.connect( + Arc::default(), + key.clone(), + Arc::default(), + /*tool_catalog_cache_context*/ None, + CancellationToken::new(), + launch("first"), + ); + let second = pool.connect( + Arc::default(), + key, + Arc::default(), + /*tool_catalog_cache_context*/ None, + CancellationToken::new(), + launch("second"), + ); + tokio::pin!(first, second); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut second).is_pending()); + started.notified().await; + release.notify_one(); + let first = tokio::time::timeout(Duration::from_secs(10), async { + tokio::select! { + first = &mut first => first, + _ = &mut second => panic!("the joining chat should still be starting"), + } + }) + .await + .expect("the joining chat must start before the first finishes") + .unwrap(); + first + .call_tool( + "test".to_string(), + /*arguments*/ None, + |_| Ok(None), + /*timeout*/ None, + ) + .await + .unwrap(); + let mut owner = create_ready_async_managed_client(Vec::new()).await; + owner.client = futures::future::ready(Ok(first.clone())).boxed().shared(); + tokio::time::timeout(Duration::from_secs(10), owner.shutdown()) + .await + .expect("private owner shutdown must not wait for the joining chat"); + assert!(first.client.is_closed().await); + + release_second.notify_one(); + let second = second.await.unwrap(); + assert_eq!(starts.load(Ordering::SeqCst), 2); + assert_eq!(*callbacks.lock().unwrap(), ["first", "second"]); + second + .call_tool( + "test".to_string(), + /*arguments*/ None, + |_| Ok(None), + /*timeout*/ None, + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn connection_key_changes_with_the_owner_but_survives_credential_refresh() { + let manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let mut original = key(&context(), &config()); + original.auth_state = Some((manager, AuthChangeState::default())); + let mut refreshed = original.clone(); + refreshed.auth_state.as_mut().unwrap().1.generation += 1; + assert!(original == refreshed); + + refreshed.auth_state.as_mut().unwrap().1.owner_generation += 1; + assert!(original != refreshed); +} + +#[tokio::test(start_paused = true)] +async fn new_chat_relists_tools_without_invalidating_or_waiting_for_existing_calls() { + let pool = McpConnectionPool::default(); + let key = key(&context(), &config()); + let server = TestServer::default(); + let starts = Arc::new(AtomicUsize::new(0)); + let release = Arc::new(Notify::new()); + release.notify_one(); + let launch = startup( + ConnectionSharing::Negotiated, + Arc::clone(&starts), + Arc::default(), + release, + server.clone(), + ); + let first = pool + .connect( + Arc::default(), + key.clone(), + Arc::default(), + /*tool_catalog_cache_context*/ None, + CancellationToken::new(), + &launch, + ) + .await + .unwrap(); + let before = first.tool_catalog.read(|snapshot| snapshot).await; + let release_call = Notify::new(); + let call = first + .tool_catalog + .run_with_snapshot(&before, || release_call.notified()); + tokio::pin!(call); + assert!(futures::poll!(&mut call).is_pending()); + for tool_name in ["test", "new"] { + *server.tools.lock().unwrap() = vec![create_test_tool("test_server", tool_name).tool]; + let second = tokio::time::timeout( + Duration::from_secs(1), + pool.connect( + Arc::default(), + key.clone(), + Arc::default(), + /*tool_catalog_cache_context*/ None, + CancellationToken::new(), + &launch, + ), + ) + .await + .expect("new chats must not wait for existing calls") + .unwrap(); + assert_eq!(starts.load(Ordering::SeqCst), 1); + assert_eq!( + second + .tool_catalog + .read(|snapshot| snapshot.tools[0].tool.name.to_string()) + .await, + tool_name + ); + tokio::time::timeout( + Duration::from_secs(1), + second.tool_catalog.refresh( + || async { Ok((vec![create_test_tool("test_server", "refreshed")], None, ())) }, + |_, _, ()| {}, + ), + ) + .await + .expect("one chat's refresh must not wait for another chat's calls") + .unwrap(); + assert_eq!( + first + .tool_catalog + .read(|snapshot| snapshot.tools[0].tool.name.to_string()) + .await, + "test" + ); + } + release_call.notify_one(); + assert_eq!(call.await, Some(())); + + server.block_listing.store(true, Ordering::Release); + let failed = pool + .connect( + Arc::default(), + key, + Arc::default(), + /*tool_catalog_cache_context*/ None, + CancellationToken::new(), + &launch, + ) + .await; + assert!(failed.is_err()); + let _cancelled = + tokio::time::timeout(Duration::from_secs(1), server.cancelled_listings.acquire()) + .await + .expect("timed-out listings must be cancelled") + .unwrap(); + first + .call_tool( + "test".to_string(), + /*arguments*/ None, + |_| Ok(None), + /*timeout*/ None, + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn negotiation_requires_the_stateless_protocol_and_detaches_startup_callbacks() { + for (protocol_version, expected_sharing, expected_action) in [ + ( + ProtocolVersion::V_2025_06_18, + false, + ElicitationAction::Decline, + ), + ( + ProtocolVersion::V_2026_07_28, + true, + ElicitationAction::Cancel, + ), + ] { + let (negotiation, callback) = SharingNegotiation::new( + Box::new(|_, _| { + futures::future::ready(Ok(ElicitationResponse { + action: ElicitationAction::Decline, + content: None, + meta: None, + })) + .boxed() + }), + CancellationToken::new(), + ); + let sharing = negotiation.finish(&protocol_version); + assert_eq!( + matches!(sharing, ConnectionSharing::Negotiated), + expected_sharing + ); + let response = callback( + rmcp::model::RequestId::Number(0), + Elicitation::OpenAiForm { + meta: None, + message: "unexpected".to_string(), + requested_schema: json!({}), + }, + ) + .await + .unwrap(); + assert_eq!(response.action, expected_action); + } +} diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 71db3cb5a7..a3e2549171 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -6,6 +6,7 @@ pub use codex_rmcp_client::McpProtocolMode; pub use connection_manager::NetworkPolicyMetadata; pub use connection_manager::NetworkPolicyMetadataProvider; pub use connection_manager::tool_is_model_visible; +pub use connection_pool::McpConnectionPool; pub use elicitation::ElicitationLifecycle; pub use elicitation::ElicitationReviewRequest; pub use elicitation::ElicitationReviewer; @@ -118,6 +119,7 @@ mod client_capabilities; mod client_tool_catalog; pub(crate) mod codex_apps; pub(crate) mod connection_manager; +mod connection_pool; pub(crate) mod elicitation; mod ema_account; mod ema_auth; diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 8e66a9ecfc..018da8ae70 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -9,6 +9,8 @@ //! Higher-level aggregation and resource/tool APIs live in //! [`crate::connection_manager`]. +#[path = "rmcp_client/managed_client.rs"] +mod managed_client; #[path = "rmcp_client/status.rs"] mod status; @@ -16,6 +18,8 @@ mod status; #[path = "rmcp_client/provenance_tests.rs"] mod provenance_tests; +pub(crate) use managed_client::ManagedClient; + use std::borrow::Cow; use std::collections::BTreeMap; use std::collections::HashMap; @@ -33,6 +37,10 @@ use crate::codex_apps::normalize_codex_apps_callable_name; use crate::codex_apps::normalize_codex_apps_callable_namespace; use crate::codex_apps::normalize_codex_apps_tool_title; use crate::codex_apps::prepare_openai_file_params_for_model; +use crate::connection_pool::ConnectionKey; +use crate::connection_pool::ConnectionSharing; +use crate::connection_pool::ConnectionSlot; +use crate::connection_pool::SharingNegotiation; use crate::elicitation::ElicitationRequestManager; use crate::executor_environment_http_client::ExecutorEnvironmentHttpClient; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; @@ -94,7 +102,6 @@ use rmcp::model::Tool as RmcpTool; use tokio::sync::watch; use tokio::time::Instant as TokioInstant; use tokio_util::sync::CancellationToken; -use tokio_util::task::AbortOnDropHandle; use tracing::Instrument; use tracing::instrument; use tracing::warn; @@ -125,45 +132,6 @@ const UNTRUSTED_CONNECTOR_META_KEYS: &[&str] = &[ "connectorDescription", ]; -#[derive(Clone)] -pub(crate) struct ManagedClient { - pub(crate) _auth_change_notifications: Option>>, - pub(crate) client: Arc, - pub(crate) server_info: McpServerInfo, - pub(crate) tool_catalog: Arc, - pub(crate) tool_timeout: Option, - pub(crate) server_instructions: Option, - pub(crate) server_supports_sandbox_state_meta_capability: bool, - pub(crate) client_supports_network_policy_meta_capability: bool, - pub(crate) codex_apps_tools_cache_context: Option>, -} - -impl ManagedClient { - pub(crate) async fn listed_catalog(&self) -> (Vec, Option>) { - let total_start = Instant::now(); - self.tool_catalog - .read(|catalog| { - // Discovery may use the shared cache until this client is refreshed. - // Executable bindings always capture this client's own catalog. - if catalog.revision == 0 - && let Some(cache_context) = &self.codex_apps_tools_cache_context - { - let tools = cache_context.current_catalog_snapshot(); - emit_duration( - MCP_TOOLS_LIST_DURATION_METRIC, - total_start.elapsed(), - &[("cache", if tools.is_some() { "hit" } else { "miss" })], - ); - if let Some(tools) = tools { - return (tools.tools().to_vec(), tools.shared_tool_list_meta()); - } - } - (catalog.tools.to_vec(), catalog.meta) - }) - .await - } -} - pub(crate) type ManagedClientFuture = Shared>>; type StartupShutdown = Shared>>; @@ -328,7 +296,7 @@ struct ManagedClientStartup { } impl ManagedClientStartup { - fn start(&self) -> ManagedClientFuture { + fn start(&self, sharing_declined: Option) -> ManagedClientFuture { // A new attempt must not expose capabilities from an earlier connection. *self .server_capabilities @@ -408,6 +376,7 @@ impl ManagedClientStartup { server_name.clone(), client, StartServerTaskParams { + sharing_declined, is_codex_apps_mcp_server, startup_timeout: Some(startup_timeout), tx_event, @@ -491,6 +460,7 @@ pub(crate) struct AsyncManagedClient { pub(crate) tool_catalog_cache_context: Option, pub(crate) startup_complete: Arc, pub(crate) startup_reconnect: Option>, + pub(crate) connection_slot: Option>, pub(crate) cancel_token: CancellationToken, } @@ -520,6 +490,7 @@ impl AsyncManagedClient { auth_changes: Option>, protocol_mode: McpProtocolMode, catalog_item_limit: usize, + connection_pool: Option<(crate::McpConnectionPool, ConnectionKey)>, ) -> Self { let is_codex_apps_mcp_server = server_name == CODEX_APPS_MCP_SERVER_NAME; let reconnect_server_name = server_name.clone(); @@ -569,19 +540,64 @@ impl AsyncManagedClient { startup_complete: Arc::clone(&startup_complete), server_capabilities: Arc::clone(&server_capabilities), }); - // Startup can publish a legacy fallback. Readers wait for the negotiated catalog. + // A discovery attempt may publish a legacy catalog. It cannot read that catalog + // before negotiation or publish a modern joiner's tool list into it. let tool_catalog_cache_context = tool_catalog_cache_context .filter(|_| protocol_mode != McpProtocolMode::DiscoverLocalStdio); - let client = startup.start(); + let (client, connection_slot) = if let Some((pool, key)) = connection_pool { + let slot = Arc::new(ConnectionSlot::default()); + let connection_slot = Arc::clone(&slot); + let startup = Arc::clone(&startup); + let cancel_token = cancel_token.clone(); + let startup_complete = Arc::clone(&startup_complete); + let server_capabilities = Arc::clone(&server_capabilities); + let tool_catalog_cache_context = tool_catalog_cache_context.clone(); + let client = async move { + let outcome = pool + .connect( + slot, + key, + server_capabilities, + tool_catalog_cache_context, + cancel_token, + |cancellation, sharing_declined| { + let mut shared_startup = (*startup).clone(); + shared_startup.cancel_token = cancellation; + shared_startup.start(sharing_declined) + }, + ) + .await + .map(|mut client| { + if matches!(client.connection_sharing, ConnectionSharing::Shared(_)) { + client.shared_elicitation = + Some(Arc::new(startup.elicitation_requests.make_sender( + startup.server_name.clone(), + startup.tx_event.clone(), + &startup.client_mcp_extensions, + ))); + } + client + }); + startup_complete.store(true, Ordering::Release); + outcome + } + .boxed() + .shared(); + (client, Some(connection_slot)) + } else { + (startup.start(/*sharing_declined*/ None), None) + }; let startup_reconnect = is_codex_apps_mcp_server.then(|| { let startup = Arc::clone(&startup); Arc::new( - CodexAppsStartupReconnect::new(Arc::new(move || startup.start())) - .with_startup_status_context( - startup_submit_id, - reconnect_server_name, - reconnect_tx_event, - ), + CodexAppsStartupReconnect::new(Arc::new(move || { + startup.start(/*sharing_declined*/ None) + })) + .with_startup_status_context( + startup_submit_id, + reconnect_server_name, + reconnect_tx_event, + ), ) }); Self { @@ -594,6 +610,7 @@ impl AsyncManagedClient { tool_catalog_cache_context, startup_complete, startup_reconnect, + connection_slot, cancel_token, } } @@ -669,13 +686,29 @@ impl AsyncManagedClient { pub(crate) async fn shutdown(&self) { self.cancel_token.cancel(); match self.client().await { - Ok(client) => client.client.shutdown().await, + Ok(client) => match &client.connection_sharing { + ConnectionSharing::Owned(lease) | ConnectionSharing::Shared(lease) => { + if let Some(slot) = lease.release() { + return slot.wait_for_shutdown().await; + } + if matches!(&client.connection_sharing, ConnectionSharing::Owned(_)) { + // Waiters from startup may still hold leases even though sharing was declined. + client.client.shutdown().await; + } + } + ConnectionSharing::Private | ConnectionSharing::Negotiated => { + client.client.shutdown().await; + } + }, Err(StartupOutcomeError::Cancelled) => {} Err(error) => { warn!("failed to initialize MCP client during shutdown: {error:#}"); error.wait_for_shutdown().await; } } + if let Some(slot) = &self.connection_slot { + slot.wait_for_shutdown().await; + } } pub(crate) fn has_cached_tools(&self) -> bool { @@ -1081,6 +1114,7 @@ async fn start_server_task( params: StartServerTaskParams, ) -> Result { let StartServerTaskParams { + sharing_declined, is_codex_apps_mcp_server, startup_timeout, tx_event, @@ -1098,6 +1132,12 @@ async fn start_server_task( elicitation_requests.make_sender(server_name.clone(), tx_event, &client_mcp_extensions); let client_supports_network_policy_meta = client_mcp_extensions.contains(MCP_NETWORK_POLICY_META_CAPABILITY); + let (sharing_negotiation, send_elicitation) = if let Some(sharing_declined) = sharing_declined { + let (negotiation, sender) = SharingNegotiation::new(send_elicitation, sharing_declined); + (Some(negotiation), sender) + } else { + (None, send_elicitation) + }; let mut params = mcp_initialize_request_params(client_elicitation_capability, client_mcp_extensions); if auth_changes.is_some() { @@ -1130,6 +1170,9 @@ async fn start_server_task( { context.invalidate_if_newest(ticket); } + let connection_sharing = sharing_negotiation + .map(|negotiation| negotiation.finish(&initialize_result.protocol_version)) + .unwrap_or_default(); *server_capabilities .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) = @@ -1234,6 +1277,8 @@ async fn start_server_task( ); } let managed = ManagedClient { + connection_sharing, + shared_elicitation: None, _auth_change_notifications: auth_change_notifications, client: Arc::clone(&client), server_info, @@ -1338,6 +1383,7 @@ fn mcp_server_info_from_implementation( } struct StartServerTaskParams { + sharing_declined: Option, server_capabilities: Arc>>, is_codex_apps_mcp_server: bool, startup_timeout: Option, // TODO: cancel_token should handle this. diff --git a/codex-rs/codex-mcp/src/rmcp_client/managed_client.rs b/codex-rs/codex-mcp/src/rmcp_client/managed_client.rs new file mode 100644 index 0000000000..1749bf853d --- /dev/null +++ b/codex-rs/codex-mcp/src/rmcp_client/managed_client.rs @@ -0,0 +1,150 @@ +//! A chat's access to an initialized MCP client and its tool catalog. +//! A pooled connection honors this chat's lease and routes input requests back to it. + +use std::future::Future; +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; + +use super::MCP_TOOLS_LIST_DURATION_METRIC; +use crate::client_tool_catalog::ClientToolCatalog; +use crate::connection_pool::ConnectionSharing; +use crate::runtime::emit_duration; +use crate::tools::ToolInfo; +use anyhow::Result; +use anyhow::anyhow; +use codex_connectors::ConnectorRuntimeContext; +use codex_protocol::mcp::McpServerInfo; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::SendElicitation; +use rmcp::model::ListResourceTemplatesResult; +use rmcp::model::ListResourcesResult; +use rmcp::model::PaginatedRequestParams; +use rmcp::model::ServerCapabilities; +use serde_json::Value; +use tokio_util::task::AbortOnDropHandle; + +#[derive(Clone)] +pub(crate) struct ManagedClient { + pub(crate) connection_sharing: ConnectionSharing, + /// The calling chat's input handler, separate from the shared transport's startup handler. + pub(crate) shared_elicitation: Option>, + pub(crate) _auth_change_notifications: Option>>, + pub(crate) client: Arc, + pub(crate) server_info: McpServerInfo, + pub(crate) tool_catalog: Arc, + pub(crate) tool_timeout: Option, + pub(crate) server_instructions: Option, + pub(crate) server_supports_sandbox_state_meta_capability: bool, + pub(crate) client_supports_network_policy_meta_capability: bool, + pub(crate) codex_apps_tools_cache_context: Option>, +} + +impl ManagedClient { + pub(crate) async fn with_active_shared_lease( + &self, + operation: impl Future, + ) -> Result { + match &self.connection_sharing { + ConnectionSharing::Shared(lease) => tokio::select! { + biased; + () = lease.cancellation().cancelled() => { + Err(anyhow!("MCP connection lease has been released")) + } + result = operation => Ok(result), + }, + ConnectionSharing::Private + | ConnectionSharing::Negotiated + | ConnectionSharing::Owned(_) => Ok(operation.await), + } + } + + pub(crate) async fn call_tool( + &self, + name: String, + arguments: Option, + meta_for_session: impl Fn(&ServerCapabilities) -> Result> + Send + Sync, + timeout: Option, + ) -> Result { + self.with_active_shared_lease(async { + if let Some(send_elicitation) = &self.shared_elicitation { + self.client + .call_tool_with_elicitation( + name, + arguments, + meta_for_session, + timeout, + Arc::clone(send_elicitation), + ) + .await + } else { + self.client + .call_tool(name, arguments, meta_for_session, timeout) + .await + } + }) + .await? + } + + /// Lists resources within this chat's lease lifetime. + pub(crate) async fn list_resources( + &self, + params: Option, + timeout: Option, + ) -> Result { + self.with_active_shared_lease(self.client.list_resources(params, timeout)) + .await? + } + + /// Lists resource templates within this chat's lease lifetime. + pub(crate) async fn list_resource_templates( + &self, + params: Option, + timeout: Option, + ) -> Result { + self.with_active_shared_lease(self.client.list_resource_templates(params, timeout)) + .await? + } + + /// Reads a resource using this chat's input authority and lease lifetime. + pub(crate) async fn read_resource( + &self, + params: rmcp::model::ReadResourceRequestParams, + timeout: Option, + ) -> Result { + self.with_active_shared_lease(async { + if let Some(send_elicitation) = &self.shared_elicitation { + self.client + .read_resource_with_elicitation(params, timeout, Arc::clone(send_elicitation)) + .await + } else { + self.client.read_resource(params, timeout).await + } + }) + .await? + } + + pub(crate) async fn listed_catalog(&self) -> (Vec, Option>) { + let total_start = Instant::now(); + self.tool_catalog + .read(|catalog| { + // Discovery may use the shared cache until this client is refreshed. + // Executable bindings always capture this client's own catalog. + if catalog.revision == 0 + && let Some(cache_context) = &self.codex_apps_tools_cache_context + { + let tools = cache_context.current_catalog_snapshot(); + emit_duration( + MCP_TOOLS_LIST_DURATION_METRIC, + total_start.elapsed(), + &[("cache", if tools.is_some() { "hit" } else { "miss" })], + ); + if let Some(tools) = tools { + return (tools.tools().to_vec(), tools.shared_tool_list_meta()); + } + } + (catalog.tools.to_vec(), catalog.meta) + }) + .await + } +} diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index e994fb0f38..a9f1f09c62 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -8,6 +8,10 @@ mod status; +#[cfg(test)] +#[path = "runtime_shared_connection_tests.rs"] +mod shared_connection_tests; + use std::collections::HashMap; use std::collections::HashSet; use std::path::PathBuf; @@ -46,6 +50,7 @@ use rmcp::model::ReadResourceResult; use rmcp::model::RequestId; use tokio::sync::watch; use tokio_util::sync::CancellationToken; +use tracing::warn; use crate::McpConfig; use crate::binding::McpBinding; @@ -54,6 +59,7 @@ use crate::client_tool_catalog::CodexAppsToolSnapshot; use crate::connection_manager::BindingCatalogRevision; use crate::connection_manager::McpConnectionSet; use crate::connection_manager::PreviousMcpConnections; +use crate::connection_pool::McpConnectionPool; use crate::elicitation::ElicitationLifecycle; use crate::elicitation::ElicitationRequestRouter; use crate::elicitation::ElicitationReviewerHandle; @@ -112,6 +118,8 @@ pub struct McpRuntimeInput { /// their exact connections and configuration for as long as they are needed. pub struct McpRuntime { current: ArcSwap, + // Old bindings keep leases alive after publication; cancellation alone does not release them. + superseded_connections: Mutex>>, event_stream_cancellation: Mutex, reconnect_pending: AtomicBool, resource_cache_generation: AtomicU64, @@ -269,6 +277,7 @@ impl McpRuntime { ready_environments: HashMap::new(), cached_binding: Mutex::new(None), }), + superseded_connections: Mutex::default(), event_stream_cancellation: Mutex::new(EventStreamCancellation { event_server_available: false, cancel_event_streams_on_server_removal: watch::channel(()).0, @@ -398,6 +407,9 @@ impl McpRuntime { let auth = input.auth.clone(); let auth_manager = input.auth_manager.clone(); let auth_token = auth.as_ref().and_then(|auth| auth.get_token().ok()); + if !reuse_previous { + current.connections.retire_shared_connections(); + } let auth_matches = current.auth_matches(auth.as_ref()); let auth_owner = AuthOwner::capture(auth.as_ref(), auth_manager.as_deref()); let same_owner = match (current.auth_owner, auth_owner) { @@ -421,8 +433,9 @@ impl McpRuntime { let ready_selected_capability_roots = input.ready_selected_capability_roots.clone(); let environment_selections = Arc::clone(&input.runtime_context.environment_selections); let ready_environments = input.runtime_context.ready_environments.clone(); + // Keep MCP startup state out of the caller's nested session-initialization future. let connections = Arc::new( - McpConnectionSet::new( + Box::pin(McpConnectionSet::new( ema_auth, retain_ema.then_some(current.connections.as_ref()), reuse_previous.then_some(PreviousMcpConnections { @@ -432,7 +445,7 @@ impl McpRuntime { publication_gate, input, self.elicitation_router.clone(), - ) + )) .await, ); let hosted_event_server_retained = connections.contains_server(CODEX_APPS_MCP_SERVER_NAME) @@ -448,6 +461,14 @@ impl McpRuntime { .event_stream_cancellation .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); + { + let mut superseded = self + .superseded_connections + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + superseded.retain(|connections| connections.strong_count() > 0); + superseded.push(Arc::downgrade(¤t.connections)); + } self.current.store(Arc::new(PublishedMcpRuntime { connections, config: Some(config), @@ -923,7 +944,26 @@ impl McpRuntime { pub async fn shutdown(&self) { self.elicitation_router.close().await; - self.latest_connections().shutdown().await; + let superseded = std::mem::take( + &mut *self + .superseded_connections + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ); + let current = self.latest_connections(); + // Keep every publication's cleanup alive if the caller is interrupted. + let shutdown = tokio::spawn(async move { + for connections in superseded + .into_iter() + .filter_map(|connections| connections.upgrade()) + { + connections.shutdown_shared_connections().await; + } + current.shutdown().await; + }); + if let Err(error) = shutdown.await { + warn!("MCP runtime shutdown task failed: {error}"); + } } } @@ -933,6 +973,7 @@ impl McpRuntime { /// environment registry, captured thread selections, and the host-local cwd used by MCP processes. #[derive(Clone)] pub struct McpRuntimeContext { + pub(crate) connection_pool: Option, environment_manager: Arc, environment_selections: Arc<[TurnEnvironmentSelection]>, ready_environments: HashMap>, @@ -981,6 +1022,7 @@ impl McpRuntimeContext { .with_tls_backend_fallback(), ); Self { + connection_pool: None, environment_manager, environment_selections: Arc::default(), ready_environments: HashMap::new(), @@ -989,6 +1031,12 @@ impl McpRuntimeContext { } } + /// Uses this process's pool for compatible local stdio servers using MCP 2026-07-28. + pub fn with_connection_pool(mut self, pool: McpConnectionPool) -> Self { + self.connection_pool = Some(pool); + self + } + /// Pins all captured selections and the concrete handles for environments already ready. pub fn with_selected_environments( mut self, diff --git a/codex-rs/codex-mcp/src/runtime_shared_connection_tests.rs b/codex-rs/codex-mcp/src/runtime_shared_connection_tests.rs new file mode 100644 index 0000000000..812bbcc491 --- /dev/null +++ b/codex-rs/codex-mcp/src/runtime_shared_connection_tests.rs @@ -0,0 +1,172 @@ +use super::*; +use crate::connection_manager::tests::create_ready_async_managed_client; +use crate::connection_pool::ConnectionSharing; +use crate::connection_pool::ConnectionSlot; +use crate::connection_pool::McpConnectionPool; +use crate::connection_pool::tests as fixture; +use futures::FutureExt; +use pretty_assertions::assert_eq; +use std::sync::atomic::AtomicUsize; +use tokio::sync::Notify; + +#[tokio::test] +async fn forced_refresh_replaces_pending_connection_without_interrupting_other_chats() { + let pool = McpConnectionPool::default(); + let key = fixture::key(&fixture::context(), &fixture::config()); + let starts = Arc::new(AtomicUsize::new(0)); + let started = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let launch = Arc::new(fixture::startup( + ConnectionSharing::Negotiated, + Arc::clone(&starts), + Arc::clone(&started), + Arc::clone(&release), + fixture::TestServer::default(), + )); + let slot = Arc::new(ConnectionSlot::default()); + let first = { + let pool = pool.clone(); + let key = key.clone(); + let slot = Arc::clone(&slot); + let launch = Arc::clone(&launch); + async move { + pool.connect( + slot, + key, + Arc::default(), + /*tool_catalog_cache_context*/ None, + CancellationToken::new(), + &*launch, + ) + .await + } + .boxed() + .shared() + }; + let mut poll_first = first.clone(); + assert!(futures::poll!(&mut poll_first).is_pending()); + let other = pool.connect( + Arc::default(), + key.clone(), + Arc::default(), + /*tool_catalog_cache_context*/ None, + CancellationToken::new(), + &*launch, + ); + tokio::pin!(other); + assert!(futures::poll!(&mut other).is_pending()); + started.notified().await; + let mut owner = create_ready_async_managed_client(Vec::new()).await; + owner.client = first; + owner.connection_slot = Some(slot); + let mut connections = McpConnectionSet::empty(/*prefix_mcp_tool_names*/ true); + connections.insert_test_client("test_server", owner); + let connections = Arc::new(connections); + let config = Arc::new(crate::mcp::tests::test_mcp_config(std::env::temp_dir())); + let runtime = McpRuntime::empty(/*prefix_mcp_tool_names*/ true); + runtime.current.store(Arc::new(PublishedMcpRuntime { + connections: Arc::clone(&connections), + config: Some(Arc::clone(&config)), + auth: None, + auth_manager: None, + auth_token: None, + auth_owner: None, + auth_generation: Arc::new(()), + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + environment_selections: Arc::default(), + ready_environments: HashMap::new(), + cached_binding: Mutex::new(None), + })); + let input = McpRuntimeInput { + network_policy_metadata: None, + startup_policy: McpStartupPolicy::Eager, + config, + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + mcp_servers: HashMap::new(), + submit_id: "refresh".to_string(), + tx_event: None, + startup_cancellation_token: CancellationToken::new(), + runtime_context: fixture::context(), + codex_apps_tools_cache: ConnectorRuntimeManager::default(), + tool_catalog_cache: McpToolCatalogCache::default(), + codex_apps_tools_cache_key: ConnectorRuntimeContextKey::personal( + /*account_id*/ None, /*chatgpt_user_id*/ None, + ), + client_mcp_extensions: ClientMcpExtensions::default(), + auth: None, + auth_manager: None, + elicitation_reviewer: None, + elicitation_lifecycle: None, + }; + let binding = runtime.current_binding_for_call("test_server"); + tokio::pin!(binding); + assert!(futures::poll!(&mut binding).is_pending()); + runtime.reconnect_on_next_refresh(); + runtime.replace(input).await; + let _replacement = tokio::time::timeout( + Duration::from_secs(1), + pool.connect( + Arc::default(), + key.clone(), + Arc::default(), + /*tool_catalog_cache_context*/ None, + CancellationToken::new(), + &*launch, + ), + ) + .await + .expect("forced refresh must not wait for the old startup") + .unwrap(); + release.notify_one(); + let other = other.await.unwrap(); + other + .call_tool( + "test".to_string(), + /*arguments*/ None, + |_| Ok(None), + /*timeout*/ None, + ) + .await + .unwrap(); + + // A second refresh of the old owner must not evict the newer replacement. + connections.retire_shared_connections(); + let _joined = pool + .connect( + Arc::default(), + key, + Arc::default(), + /*tool_catalog_cache_context*/ None, + CancellationToken::new(), + &*launch, + ) + .await + .unwrap(); + assert_eq!(starts.load(Ordering::SeqCst), 2); + + let binding = binding.await.expect("binding from the old publication"); + runtime.shutdown().await; + assert!( + binding + .list_resources("test_server", /*params*/ None) + .await + .is_err() + ); + other + .call_tool( + "test".to_string(), + /*arguments*/ None, + |_| Ok(None), + /*timeout*/ None, + ) + .await + .unwrap(); + let transport = Arc::clone(&other.client); + let mut other_owner = create_ready_async_managed_client(Vec::new()).await; + other_owner.client = futures::future::ready(Ok(other)).boxed().shared(); + other_owner.shutdown().await; + assert!(transport.is_closed().await); + drop(binding); +} diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index f4170cb011..3bfa3b412c 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -111,6 +111,7 @@ enum OrderedMcpOverlay { #[derive(Clone)] pub struct McpManager { + connection_pool: codex_mcp::McpConnectionPool, plugins_manager: Arc, extensions: Arc>, codex_apps_tools_cache: ConnectorRuntimeManager, @@ -135,6 +136,7 @@ impl McpManager { codex_apps_tools_cache: ConnectorRuntimeManager, ) -> Self { Self { + connection_pool: codex_mcp::McpConnectionPool::default(), plugins_manager, extensions, codex_apps_tools_cache, @@ -152,6 +154,10 @@ impl McpManager { self.tool_catalog_cache.clone() } + pub(crate) fn connection_pool(&self) -> codex_mcp::McpConnectionPool { + self.connection_pool.clone() + } + pub(crate) fn shared_definitions(&self) -> &SharedMcpDefinitions { &self.shared_definitions } diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 895770147c..b1d8e2a9b4 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -175,7 +175,7 @@ impl Session { self.mcp_runtime_context(&environments, &host_fallback_cwd) } - fn mcp_runtime_context( + pub(super) fn mcp_runtime_context( &self, environments: &TurnEnvironmentSnapshot, host_fallback_cwd: &std::path::Path, @@ -188,6 +188,7 @@ impl Session { self.services.turn_environments.environment_manager(), local_process_cwd, ) + .with_connection_pool(self.services.mcp_manager.connection_pool()) .with_selected_environments( environments.all_selections().into(), environments.ready_environment_handles(), diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index d00bb6af41..ac0bc4a983 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -358,14 +358,8 @@ impl Session { }), ); let mcp_config = Arc::new(config); - let runtime_context = McpRuntimeContext::new( - self.services.turn_environments.environment_manager(), - desired.local_process_cwd.clone(), - ) - .with_selected_environments( - desired.environments.all_selections().into(), - desired.environments.ready_environment_handles(), - ); + let runtime_context = + self.mcp_runtime_context(&desired.environments, &desired.local_process_cwd); McpRuntimeInput { network_policy_metadata: None, startup_policy: if matches!(desired.session_source, SessionSource::SubAgent(_)) { diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index afe1c89c16..a3a7a47e3e 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -152,6 +152,9 @@ mod mcp_resource_messages; #[path = "scenarios_mcp_ema.rs"] mod mcp_ema; +#[path = "scenarios_mcp_sharing.rs"] +mod mcp_sharing; + #[path = "scenarios_guardian_agent_messages_tests.rs"] mod guardian_agent_messages; diff --git a/codex-rs/core/tests/suite/scenarios_mcp_sharing.rs b/codex-rs/core/tests/suite/scenarios_mcp_sharing.rs new file mode 100644 index 0000000000..54b7aebe95 --- /dev/null +++ b/codex-rs/core/tests/suite/scenarios_mcp_sharing.rs @@ -0,0 +1,179 @@ +//! Shared MCP transports preserve each chat's model-visible tools and input responses. + +use super::*; +use codex_protocol::approvals::ElicitationAction; +use codex_protocol::approvals::ElicitationRequest; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::ThreadSettingsOverrides; +use core_test_support::responses; +use pretty_assertions::assert_eq; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shared_mcp_keeps_tools_and_model_requested_input_per_chat() -> Result<()> { + skip_if_no_network!(Ok(())); + core_test_support::skip_if_remote!(Ok(()), "the shared stdio fixture runs on the host"); + skip_if_wine_exec!(Ok(()), "requires a Windows test_stdio_server binary"); + + let first_server = start_mock_server().await; + let second_server = start_mock_server().await; + let command = stdio_server_bin()?; + let test = test_codex() + .with_config(move |config| { + super::configure_scenario_catalog(config); + config.approvals_reviewer = ApprovalsReviewer::User; + config + .features + .enable(Feature::McpConnectionSharing) + .expect("enable sharing"); + config + .features + .disable(Feature::Mcp20260728) + .expect("let sharing select the protocol"); + config + .mcp_servers + .set( + serde_json::from_value(json!({ + "shared": { + "command": command, + "enabled_tools": ["echo"], + "default_tools_approval_mode": "approve", + "env": { + "MCP_TEST_ECHO_INPUT": "1", + }, + }, + })) + .expect("shared MCP config"), + ) + .expect("set MCP config"); + }) + .with_model_info_override("gpt-5.5", |model| { + model.tool_mode = Some(ToolMode::Direct); + model.supports_search_tool = false; + model.use_responses_lite = false; + }) + .build_with_auto_env(&first_server) + .await?; + wait_for_mcp_server(&test.codex, "shared").await?; + + let mut options = test.start_thread_options().await; + options.config.model_provider.base_url = Some(format!("{}/v1", second_server.uri())); + let mut servers = options.config.mcp_servers.get().clone(); + servers + .get_mut("shared") + .expect("shared server") + .enabled_tools = Some(vec!["echo-tool".to_string()]); + options.config.mcp_servers.set(servers)?; + let second = test.thread_manager.start_thread(options).await?.thread; + wait_for_mcp_server(&second, "shared").await?; + + let chats = [ + ( + &test.codex, + &first_server, + "echo", + "First task owner?", + "Mira", + ), + ( + &second, + &second_server, + "echo_tool", + "Second task owner?", + "Sora", + ), + ]; + let mut mocks = Vec::new(); + for (thread, server, tool, question, _) in chats { + mocks.push( + mount_sse_sequence( + server, + vec![ + sse(vec![ + ev_response_created("ask-owner"), + ev_function_call_with_namespace( + "ask-owner", + "mcp__shared", + tool, + &json!({ "message": question }).to_string(), + ), + ev_completed("ask-owner"), + ]), + responses::sse_completed("done"), + ], + ) + .await, + ); + thread + .start_or_steer_turn( + TurnInputRequest::user_input(vec![UserInput::Text { + text: format!("Ask the user: {question}"), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + approval_policy: Some(AskForApproval::OnRequest), + permission_profile: Some(PermissionProfile::read_only()), + ..Default::default() + }), + ) + .await?; + } + + let mut inputs = Vec::new(); + let mut processes = Vec::new(); + for (thread, _, _, question, _) in chats { + let input = wait_for_event_match(thread, |event| match event { + EventMsg::ElicitationRequest(input) => Some(input.clone()), + _ => None, + }) + .await; + let ElicitationRequest::Form { meta, message, .. } = &input.request else { + panic!("expected a form input"); + }; + assert_eq!(message, question); + processes.push(meta.as_ref().expect("fixture metadata")["fixture/processId"].clone()); + inputs.push(input); + } + assert!(processes[0].is_u64()); + assert_eq!( + processes[0], processes[1], + "the chats must share one process" + ); + + // Both calls use the same server input key; resolve the joiner's input first. + for ((thread, _, _, _, answer), input) in chats.into_iter().zip(inputs).rev() { + thread + .submit(Op::ResolveElicitation { + server_name: input.server_name, + request_id: input.id, + decision: ElicitationAction::Accept, + content: Some(json!({ "answer": answer })), + meta: None, + }) + .await?; + } + + let mut histories = Vec::new(); + for ((thread, ..), mock) in chats.into_iter().zip(mocks) { + wait_for_event(thread, |event| matches!(event, EventMsg::TurnComplete(_))).await; + histories.push(mock.requests()); + } + let entries = histories + .iter() + .zip(["creator", "joiner"]) + .flat_map(|(requests, label)| { + requests + .iter() + .map(move |request| SnapshotEntry::captured(request).labeled(label)) + }) + .collect::>(); + insta::assert_snapshot!(context_snapshot::format_context_snapshot( + "Two chats share one MCP process, expose different tools, and answer simultaneous model-requested questions in reverse order.", + &entries, + &ContextSnapshotOptions::default() + .rewrite_known_segments() + .include_request_settings(), + )); + Ok(()) +} diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mcp_sharing__shared_mcp_keeps_tools_and_model_requested_input_per_chat.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mcp_sharing__shared_mcp_keeps_tools_and_model_requested_input_per_chat.snap new file mode 100644 index 0000000000..b007c48c77 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mcp_sharing__shared_mcp_keeps_tools_and_model_requested_input_per_chat.snap @@ -0,0 +1,68 @@ +--- +source: core/tests/suite/scenarios_mcp_sharing.rs +expression: "context_snapshot::format_context_snapshot(\"Two chats share one MCP process, expose different tools, and answer simultaneous model-requested questions in reverse order.\",\n&entries,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)" +--- +Scenario: Two chats share one MCP process, expose different tools, and answer simultaneous model-requested questions in reverse order. + +## Window 1 +Settings: + include: ["reasoning.encrypted_content"] + model: "gpt-5.5" + parallel_tool_calls: true + prompt_cache_key: "" + reasoning: {"effort":"medium"} + store: false + stream: true + text: {"verbosity":"low"} + tool_choice: "auto" + tools (11; hash=e4ad6947d6c88cac): + - function/exec_command: Runs a command in a PTY, returning output or a session ID for ongoing interaction.; args=[cmd, justification, login, max_output_tokens, prefix_rule, sandbox_permissions, shell, tty, workdir, yield_time_ms]; hash=c04108e51a8548c5 + - function/write_stdin: Writes characters to an existing unified exec session and returns recent output.; args=[chars, max_output_tokens, session_id, yield_time_ms]; hash=7fb62d8a3e2df4e2 + - function/list_mcp_resources: Lists resources provided by MCP servers. Resources allow servers to share data that provides cont...; args=[cursor, server]; hash=6792edaff13561bd + - function/list_mcp_resource_templates: Lists resource templates provided by MCP servers. Parameterized resource templates allow servers ...; args=[cursor, server]; hash=60f17c90ef7a4ebb + - function/read_mcp_resource: Read a specific resource from an MCP server given the server name and resource URI.; args=[server, uri]; hash=24246896b39008ae + - function/request_user_input: Request user input for one to three short questions and wait for the response. This tool is only ...; args=[questions]; hash=64a729b25c424fc6 + - custom/apply_patch: The `apply_patch` tool can be used to edit files. This is a FREEFORM tool, so do not wrap the pat...; hash=dea9f12b0111a4ba + - function/view_image: View a local image file from the filesystem when visual inspection is needed. Use this for images...; args=[detail, path]; hash=3b766d5bdec5d50e + - namespace/multi_agent_v1: Tools for spawning and managing sub-agents.; hash=2327142bda177c50 + - function/close_agent + - function/resume_agent + - function/send_input + - function/spawn_agent + - function/wait_agent + - namespace/mcp__shared: Use these tools to exercise the rmcp test server.; hash=c2583dfa504bbe3f + - function/echo + - web_search; hash=f103e327c6870a46 +-- request 1 (turn; creator) -- +00:message/developer: + +01:message/developer: + +02:message/user: + +03:message/user: + Ask the user: First task owner? +-- request 2 (turn; creator) -- +04:function_call/mcp__shared.echo:{"message":"First task owner?"} +05:function_call_output:Wall time: seconds + Output: | Mira + +## Window 2 (after request 2: input diverged at item 03) +Settings: relative to window 1 + prompt_cache_key: "" + tools (11; hash=b0f271ad4e160a41): + ~ namespace/mcp__shared: Use these tools to exercise the rmcp test server.; hash=a9e155b340098fac + - function/echo_tool +-- request 3 (turn; joiner) -- +00:message/developer: + +01:message/developer: + +02:message/user: + +03:message/user: + Ask the user: Second task owner? +-- request 4 (turn; joiner) -- +04:function_call/mcp__shared.echo_tool:{"message":"Second task owner?"} +05:function_call_output:Wall time: seconds + Output: | Sora diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 5fd61d1445..83daf42efc 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -269,7 +269,7 @@ pub enum Feature { Mcp20260728, /// Honor ordinary MCP server tool-list notifications within a turn. McpToolListChanged, - /// Discover the MCP protocol for eligible local stdio servers. + /// Discover and share eligible local stdio connections across chats. McpConnectionSharing, /// Enable MCP protocol version 2026-07-28 for the host-owned Codex Apps server. CodexAppsMcp20260728, 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 8ea0150268..83780e2a40 100644 --- a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs @@ -13,9 +13,15 @@ use rmcp::ServiceExt; use rmcp::handler::server::ServerHandler; use rmcp::model::CallToolRequestParams; use rmcp::model::CallToolResult; +use rmcp::model::ContentBlock; +use rmcp::model::ElicitRequest; +use rmcp::model::ElicitRequestParams; +use rmcp::model::ElicitationSchema; use rmcp::model::Implementation; use rmcp::model::InitializeRequestParams; use rmcp::model::InitializeResult; +use rmcp::model::InputRequest; +use rmcp::model::InputRequiredResult; use rmcp::model::JsonObject; use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; @@ -781,6 +787,30 @@ impl ServerHandler for TestToolServer { } }; + if std::env::var_os("MCP_TEST_ECHO_INPUT").is_some() { + if let Some(responses) = request.input_responses { + let answer = responses["confirmation"]["content"]["answer"] + .as_str() + .ok_or_else(|| McpError::invalid_params("missing answer text", None))?; + return Ok(CallToolResult::success(vec![ContentBlock::text(answer)]).into()); + } + let mut meta = MetaObject::new(); + meta.insert("fixture/processId".into(), json!(std::process::id())); + let form = ElicitRequestParams::FormElicitationParams { + meta: Some(meta.into()), + message: args.message, + requested_schema: ElicitationSchema::builder() + .required_string_property("answer", |s| s) + .build() + .map_err(|error| McpError::internal_error(error.to_string(), None))?, + }; + let input = InputRequest::Elicitation(ElicitRequest::new(form)); + return Ok(InputRequiredResult::from_input_requests(BTreeMap::from([ + ("confirmation".into(), input), + ])) + .into()); + } + let env_snapshot: HashMap = std::env::vars().collect(); let env_name = args.env_var.as_deref().unwrap_or("MCP_TEST_VALUE"); let echo = dynamic_server_process_label()