From f24e6954702aa5f5705c65ee57748c5817bdb507 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 17 Jul 2026 18:38:44 +0000 Subject: [PATCH] Centralize thread MCP connections in `McpRuntime` (#33889) ## What changed - Add a thread-owned `McpRuntime` that atomically publishes connection-set replacements and owns shutdown. - Route `McpResourceClient` snapshots and cache keys through that runtime so extension clients observe refreshed MCP servers. - Keep model-step configuration snapshots separate from ownership of live MCP connections. ## Testing - Add coverage that captures an extension's session resource client, refreshes the MCP server configuration, and verifies that the existing client sees the newly published server. GitOrigin-RevId: 351b2e25594b8eb6b91e50ff8d06b7d49768ae1c --- codex-rs/codex-mcp/src/lib.rs | 1 + codex-rs/codex-mcp/src/resource_client.rs | 33 +++--- codex-rs/codex-mcp/src/runtime.rs | 41 ++++++- codex-rs/core/src/session/handlers.rs | 6 +- codex-rs/core/src/session/mcp.rs | 4 +- codex-rs/core/src/session/mod.rs | 1 + codex-rs/core/src/session/session.rs | 25 ++-- codex-rs/core/src/session/tests.rs | 18 ++- codex-rs/core/src/state/service.rs | 32 +++-- .../core/tests/suite/mcp_tool_exposure.rs | 109 ++++++++++++++++++ 10 files changed, 203 insertions(+), 67 deletions(-) diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 949ef4f7af..499bcfc0b7 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -10,6 +10,7 @@ pub use resource_client::McpResourceClientCacheKey; pub use resource_client::McpResourcePage; pub use resource_client::McpResourceReadResult; pub use rmcp_client::MCP_SANDBOX_STATE_META_CAPABILITY; +pub use runtime::McpRuntime; pub use runtime::McpRuntimeContext; pub use runtime::SandboxState; pub use tool_catalog_cache::McpToolCatalogCache; diff --git a/codex-rs/codex-mcp/src/resource_client.rs b/codex-rs/codex-mcp/src/resource_client.rs index 4406b81a53..84d64f441f 100644 --- a/codex-rs/codex-mcp/src/resource_client.rs +++ b/codex-rs/codex-mcp/src/resource_client.rs @@ -3,13 +3,13 @@ use std::sync::Weak; use anyhow::Context; use anyhow::Result; -use arc_swap::ArcSwap; use codex_protocol::mcp::Resource; use codex_protocol::mcp::ResourceContent; use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceRequestParams; use crate::McpConnectionManager; +use crate::McpRuntime; /// One page of resources returned by an MCP server. #[derive(Clone, Debug, PartialEq)] @@ -27,16 +27,15 @@ pub struct McpResourceReadResult { pub contents: Vec, } -/// Session-scoped access to MCP resources through the currently installed manager. +/// Session-scoped access to MCP resources through the thread runtime. /// -/// The client retains the manager's shared publication handle rather than a manager -/// snapshot, so calls automatically use replacements installed during startup and refresh. +/// Calls automatically use the connection set most recently published by the runtime. #[derive(Clone)] pub struct McpResourceClient { - manager: Arc>, + runtime: Arc, } -/// Opaque identity for the manager currently used by an MCP resource client. +/// Opaque identity for the connection set currently used by an MCP resource client. #[derive(Clone)] pub struct McpResourceClientCacheKey(Weak); @@ -57,21 +56,21 @@ impl std::fmt::Debug for McpResourceClient { } impl McpResourceClient { - /// Creates a resource client backed by the session's replaceable MCP manager. - pub fn new(manager: Arc>) -> Self { - Self { manager } + /// Creates a resource client backed by the thread's MCP runtime. + pub fn new(runtime: Arc) -> Self { + Self { runtime } } - /// Returns an identity that changes whenever the published manager changes. + /// Returns an identity that changes whenever the published connection set changes. pub fn cache_key(&self) -> McpResourceClientCacheKey { - McpResourceClientCacheKey(Arc::downgrade(&self.manager.load_full())) + McpResourceClientCacheKey(Arc::downgrade(&self.runtime.snapshot())) } - /// Returns whether the current manager contains the named server. + /// Returns whether the current connection set contains the named server. /// /// This does not wait for server startup or imply that startup succeeded. pub async fn has_server(&self, server: &str) -> bool { - self.manager.load_full().contains_server(server) + self.runtime.snapshot().contains_server(server) } /// Lists one resource page from the named server. @@ -83,8 +82,8 @@ impl McpResourceClient { let params = cursor.map(|cursor| PaginatedRequestParams::default().with_cursor(Some(cursor))); let result = self - .manager - .load_full() + .runtime + .snapshot() .list_resources(server, params) .await?; let resources = result @@ -101,8 +100,8 @@ impl McpResourceClient { /// Reads one resource from the named server. pub async fn read_resource(&self, server: &str, uri: &str) -> Result { let result = self - .manager - .load_full() + .runtime + .snapshot() .read_resource(server, ReadResourceRequestParams::new(uri.to_string())) .await?; let contents = result diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index b9b769ffb4..93caad8b25 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -1,14 +1,15 @@ //! Runtime support for Model Context Protocol (MCP) servers. //! -//! This module contains data that describes the runtime environment in which MCP -//! servers execute, plus the sandbox state payload sent to capable servers and a -//! tiny shared metrics helper. Transport startup and orchestration live in -//! [`crate::rmcp_client`] and [`crate::connection_manager`]. +//! This module contains the thread-owned MCP runtime and data that describes the +//! environment in which MCP servers execute. Transport startup lives in +//! [`crate::rmcp_client`] and connection-set behavior lives in +//! [`crate::connection_manager`]. use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; +use arc_swap::ArcSwap; use codex_exec_server::Environment; use codex_exec_server::EnvironmentManager; use codex_exec_server::HttpClient; @@ -18,6 +19,38 @@ use codex_utils_path_uri::PathUri; use serde::Deserialize; use serde::Serialize; +use crate::McpConnectionManager; + +/// Owns the currently published MCP connection set for one Codex thread. +/// +/// Replacements are published atomically. Callers that already hold a snapshot +/// keep the previous connection set alive until their work completes. +pub struct McpRuntime { + connections: ArcSwap, +} + +impl McpRuntime { + pub fn new(connections: Arc) -> Self { + Self { + connections: ArcSwap::from(connections), + } + } + + pub fn snapshot(&self) -> Arc { + self.connections.load_full() + } + + pub fn replace(&self, connections: McpConnectionManager) -> Arc { + let connections = Arc::new(connections); + self.connections.store(Arc::clone(&connections)); + connections + } + + pub async fn shutdown(&self) { + self.snapshot().shutdown().await; + } +} + #[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct SandboxState { diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 552c03803d..8a70249968 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -607,11 +607,7 @@ async fn shutdown_session_runtime(sess: &Arc) { if let Err(err) = sess.services.code_mode_service.shutdown().await { warn!("failed to shutdown code mode session: {err}"); } - sess.services - .latest_mcp_runtime() - .manager_arc() - .shutdown() - .await; + sess.services.mcp_runtime.shutdown().await; sess.guardian_review_session.shutdown().await; } diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index ee7b7d17c6..f4bbc7d650 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -185,7 +185,9 @@ impl Session { current.runtime_context().clone(), ready_selected_capability_roots, )); - self.services.mcp_runtime.store(Some(Arc::clone(&runtime))); + self.services + .mcp_runtime_snapshot + .store(Some(Arc::clone(&runtime))); return runtime; } self.refresh_mcp_servers_inner( diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index c9cc8a95b1..26b6689946 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -73,6 +73,7 @@ use codex_login::CodexAuth; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_mcp::McpConnectionManager; use codex_mcp::McpResourceClient; +use codex_mcp::McpRuntime; use codex_mcp::McpRuntimeContext; use codex_models_manager::manager::RefreshStrategy; use codex_models_manager::manager::SharedModelsManager; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 0759789ab2..04aaba64b6 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1036,20 +1036,16 @@ impl Session { config.analytics_enabled, ) }); - // Keep one stable manager handle for the session so extension resource clients - // automatically observe the manager installed at startup and on later refreshes. - let mcp_connection_manager = Arc::new(arc_swap::ArcSwap::from_pointee( + let mcp_runtime = Arc::new(McpRuntime::new(Arc::new( McpConnectionManager::new_uninitialized_with_permission_profile( &config.permissions.approval_policy, config.permissions.permission_profile(), config.prefix_mcp_tool_names(), ), - )); + ))); let session_extension_data = codex_extension_api::ExtensionData::new(session_id.to_string()); - session_extension_data.insert(McpResourceClient::new(Arc::clone( - &mcp_connection_manager, - ))); + session_extension_data.insert(McpResourceClient::new(Arc::clone(&mcp_runtime))); for contributor in extensions.thread_lifecycle_contributors() { contributor.on_thread_start(codex_extension_api::ThreadStartInput { config: config.as_ref(), @@ -1062,15 +1058,10 @@ impl Session { } let services = SessionServices { - // Initialize the MCP connection manager with an uninitialized - // instance. It will be replaced with one created via - // McpConnectionManager::new() once all its constructor args are - // available. This also ensures `SessionConfigured` is emitted - // before any MCP-related events. It is reasonable to consider - // changing this to use Option or OnceCell, though the current - // setup is straightforward enough and performs well. - mcp_connection_manager, - mcp_runtime: arc_swap::ArcSwapOption::empty(), + // Start with an empty connection set. The initialized set is + // published after SessionConfigured so MCP events follow it. + mcp_runtime, + mcp_runtime_snapshot: arc_swap::ArcSwapOption::empty(), mcp_projection_lock: Mutex::new(()), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( @@ -1257,7 +1248,7 @@ impl Session { )) .await; sess.services - .install_mcp_connection_manager( + .install_mcp_runtime( Arc::new(mcp_projection.config), mcp_projection.plugins_available, mcp_runtime_context, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index bc180b4629..95b5adefa9 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5305,11 +5305,14 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); - let mcp_runtime = + let mcp_runtime_snapshot = crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref()); + let mcp_runtime = Arc::new(codex_mcp::McpRuntime::new( + mcp_runtime_snapshot.manager_arc(), + )); let services = SessionServices { - mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), - mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), + mcp_runtime, + mcp_runtime_snapshot: arc_swap::ArcSwapOption::from(Some(mcp_runtime_snapshot)), mcp_projection_lock: Mutex::new(()), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( @@ -7460,11 +7463,14 @@ where /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); - let mcp_runtime = + let mcp_runtime_snapshot = crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref()); + let mcp_runtime = Arc::new(codex_mcp::McpRuntime::new( + mcp_runtime_snapshot.manager_arc(), + )); let services = SessionServices { - mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), - mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), + mcp_runtime, + mcp_runtime_snapshot: arc_swap::ArcSwapOption::from(Some(mcp_runtime_snapshot)), mcp_projection_lock: Mutex::new(()), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 64aee2d497..33c0a1b087 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -34,6 +34,7 @@ use codex_hooks::Hooks; use codex_login::AuthManager; use codex_mcp::McpConfig; use codex_mcp::McpConnectionManager; +use codex_mcp::McpRuntime; use codex_mcp::McpRuntimeContext; use codex_models_manager::manager::SharedModelsManager; use codex_otel::SessionTelemetry; @@ -48,10 +49,10 @@ use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; pub(crate) struct SessionServices { - /// Mirror of the latest manager for extension resource clients that predate runtime snapshots. - pub(crate) mcp_connection_manager: Arc>, - /// The latest atomically published MCP config and manager pair. - pub(crate) mcp_runtime: ArcSwapOption, + /// The single owner of live MCP connections for this thread. + pub(crate) mcp_runtime: Arc, + /// The latest atomically published MCP config and connection snapshot. + pub(crate) mcp_runtime_snapshot: ArcSwapOption, /// Serializes environment-driven runtime rebuilds. pub(crate) mcp_projection_lock: Mutex<()>, pub(crate) mcp_startup_cancellation_token: Mutex, @@ -104,22 +105,22 @@ pub(crate) struct SessionServices { } impl SessionServices { - /// Installs the manager before validating required servers so startup-time elicitation can - /// resolve through the session's manager while validation waits. - pub(crate) async fn install_mcp_connection_manager( + /// Publishes the initial connections before validating required servers so startup-time + /// elicitation can resolve through the thread runtime while validation waits. + pub(crate) async fn install_mcp_runtime( &self, config: Arc, plugins_available: bool, runtime_context: McpRuntimeContext, ready_selected_capability_roots: Vec, - manager: McpConnectionManager, + connections: McpConnectionManager, ) -> Result<()> { let runtime = self.publish_mcp_runtime( config, plugins_available, runtime_context, ready_selected_capability_roots, - manager, + connections, ); runtime.manager().validate_required_servers().await } @@ -130,25 +131,22 @@ impl SessionServices { plugins_available: bool, runtime_context: McpRuntimeContext, ready_selected_capability_roots: Vec, - manager: McpConnectionManager, + connections: McpConnectionManager, ) -> Arc { - let manager = Arc::new(manager); - // Publish the manager for legacy resource clients first. Once the paired snapshot is - // visible, every model-scoped consumer observes this exact manager. - self.mcp_connection_manager.store(Arc::clone(&manager)); + let connections = self.mcp_runtime.replace(connections); let runtime = Arc::new(McpRuntimeSnapshot::new( config, plugins_available, - manager, + connections, runtime_context, ready_selected_capability_roots, )); - self.mcp_runtime.store(Some(Arc::clone(&runtime))); + self.mcp_runtime_snapshot.store(Some(Arc::clone(&runtime))); runtime } pub(crate) fn latest_mcp_runtime(&self) -> Arc { - let Some(runtime) = self.mcp_runtime.load_full() else { + let Some(runtime) = self.mcp_runtime_snapshot.load_full() else { unreachable!("MCP runtime must be installed before handling requests"); }; runtime diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index 954c590c7b..7fae5eb633 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -1,6 +1,15 @@ use anyhow::Result; +use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; +use codex_config::types::McpServerConfig; +use codex_config::types::McpServerTransportConfig; +use codex_core::config::Config; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadStartInput; use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_mcp::McpResourceClient; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::McpServerRefreshConfig; use codex_protocol::protocol::Op; @@ -27,8 +36,108 @@ use core_test_support::wait_for_mcp_server; use serde_json::Value; use serde_json::json; use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; +struct McpResourceClientCapture { + client: Arc>>, +} + +impl ThreadLifecycleContributor for McpResourceClientCapture { + fn on_thread_start<'a>( + &'a self, + input: ThreadStartInput<'a, Config>, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let client = input + .session_store + .get::() + .expect("session store should contain an MCP resource client"); + *self + .client + .lock() + .expect("capture lock should not be poisoned") = Some(client.as_ref().clone()); + }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn session_resource_client_follows_published_mcp_runtime() -> Result<()> { + let server = responses::start_mock_server().await; + let response = responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-1"), + ]), + ) + .await; + + let captured_client = Arc::new(Mutex::new(None)); + let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.thread_lifecycle_contributor(Arc::new(McpResourceClientCapture { + client: Arc::clone(&captured_client), + })); + let test = core_test_support::test_codex::test_codex() + .with_extensions(Arc::new(extensions.build())) + .build(&server) + .await?; + let resource_client = captured_client + .lock() + .expect("capture lock should not be poisoned") + .clone() + .expect("thread start should capture the MCP resource client"); + assert!(!resource_client.has_server("refreshed").await); + + let refreshed_server = McpServerConfig { + transport: McpServerTransportConfig::StreamableHttp { + url: format!("{}/mcp", server.uri()), + bearer_token_env_var: None, + http_headers: None, + env_http_headers: None, + }, + auth: Default::default(), + environment_id: DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: Some(Duration::from_millis(100)), + tool_timeout_sec: None, + default_tools_approval_mode: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth: None, + oauth_resource: None, + tools: HashMap::new(), + }; + test.codex + .submit(Op::RefreshMcpServers { + config: McpServerRefreshConfig { + mcp_servers: serde_json::to_value(HashMap::from([( + "refreshed".to_string(), + refreshed_server, + )]))?, + mcp_oauth_credentials_store_mode: serde_json::to_value( + test.config.mcp_oauth_credentials_store_mode, + )?, + auth_keyring_backend_kind: serde_json::to_value( + test.config.auth_keyring_backend_kind(), + )?, + }, + }) + .await?; + test.submit_turn("observe the refreshed MCP runtime") + .await?; + + assert!(resource_client.has_server("refreshed").await); + response.single_request(); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_only_exposes_direct_model_only_mcp_namespaces() -> Result<()> { skip_if_no_network!(Ok(()));