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/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(()));