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

@@ -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;

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<ResourceContent>,
 }
 
-/// 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<ArcSwap<McpConnectionManager>>,
+    runtime: Arc<McpRuntime>,
 }
 
-/// 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<McpConnectionManager>);
 
@@ -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<ArcSwap<McpConnectionManager>>) -> Self {
-        Self { manager }
+    /// Creates a resource client backed by the thread's MCP runtime.
+    pub fn new(runtime: Arc<McpRuntime>) -> 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<McpResourceReadResult> {
         let result = self
-            .manager
-            .load_full()
+            .runtime
+            .snapshot()
             .read_resource(server, ReadResourceRequestParams::new(uri.to_string()))
             .await?;
         let contents = result

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<McpConnectionManager>,
+}
+
+impl McpRuntime {
+    pub fn new(connections: Arc<McpConnectionManager>) -> Self {
+        Self {
+            connections: ArcSwap::from(connections),
+        }
+    }
+
+    pub fn snapshot(&self) -> Arc<McpConnectionManager> {
+        self.connections.load_full()
+    }
+
+    pub fn replace(&self, connections: McpConnectionManager) -> Arc<McpConnectionManager> {
+        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 {

codex-rs/core/src/session/handlers.rs

@@ -607,11 +607,7 @@ async fn shutdown_session_runtime(sess: &Arc<Session>) {
     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;
 }
 

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(

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;

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,

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(

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<ArcSwap<McpConnectionManager>>,
-    /// The latest atomically published MCP config and manager pair.
-    pub(crate) mcp_runtime: ArcSwapOption<McpRuntimeSnapshot>,
+    /// The single owner of live MCP connections for this thread.
+    pub(crate) mcp_runtime: Arc<McpRuntime>,
+    /// The latest atomically published MCP config and connection snapshot.
+    pub(crate) mcp_runtime_snapshot: ArcSwapOption<McpRuntimeSnapshot>,
     /// Serializes environment-driven runtime rebuilds.
     pub(crate) mcp_projection_lock: Mutex<()>,
     pub(crate) mcp_startup_cancellation_token: Mutex<CancellationToken>,
@@ -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<McpConfig>,
         plugins_available: bool,
         runtime_context: McpRuntimeContext,
         ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
-        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<SelectedCapabilityRoot>,
-        manager: McpConnectionManager,
+        connections: McpConnectionManager,
     ) -> Arc<McpRuntimeSnapshot> {
-        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<McpRuntimeSnapshot> {
-        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

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<Mutex<Option<McpResourceClient>>>,
+}
+
+impl ThreadLifecycleContributor<Config> 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::<McpResourceClient>()
+                .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::<Config>::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(()));