Prewarm MCP runtime updates in the background (#35144)

## Why

Changes to MCP inputs can require rebuilding the runtime before the next model
step. Start that work eagerly so the step is less likely to wait for it.

## What changed

- Add a per-session worker that prewarms dirty MCP state after configuration,
  authentication, or explicit refresh changes.
- Coalesce queued refreshes and always rebuild from the latest session state,
  while keeping exact model-step refreshes as the correctness path.
- Stop and join the worker during session shutdown.

## Testing

- Verify rapid refreshes skip superseded configuration and publish the latest
  state.
- Verify MCP-relevant policy changes mark the runtime for refresh.

GitOrigin-RevId: 045ed406caff97dcbb483bc183fd16eae0c7d7ad

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

@@ -422,7 +422,7 @@ pub async fn dynamic_tool_response(sess: &Arc<Session>, id: String, response: Dy
 }
 
 pub fn refresh_mcp_servers(sess: &Session) {
-    sess.mark_mcp_runtime_dirty();
+    sess.request_mcp_runtime_refresh();
 }
 
 pub async fn reload_user_config(sess: &Arc<Session>) {
@@ -585,9 +585,12 @@ 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}");
     }
-    let _refresh = sess.mcp_refresh_lock.acquire().await;
-    sess.mcp_refresh_lock.close();
-    sess.services.mcp_runtime.shutdown().await;
+    sess.stop_mcp_prewarm_worker().await;
+    {
+        let _refresh = sess.mcp_refresh_lock.acquire().await;
+        sess.mcp_refresh_lock.close();
+        sess.services.mcp_runtime.shutdown().await;
+    }
     sess.guardian_review_session.shutdown().await;
 
     crate::hook_runtime::run_session_end_hooks(sess).await;

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

@@ -502,7 +502,7 @@ impl Session {
         self.services
             .supports_openai_form_elicitation
             .store(supported, std::sync::atomic::Ordering::Relaxed);
-        self.mark_mcp_runtime_dirty();
+        self.request_mcp_runtime_refresh();
         Ok(())
     }
 

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

@@ -0,0 +1,75 @@
+//! Best-effort MCP prewarming.
+//!
+//! A bounded channel coalesces refresh requests. The worker only prepares the
+//! newest thread state; exact model steps remain the correctness path.
+
+use super::*;
+
+impl Session {
+    pub(super) fn request_mcp_runtime_refresh(&self) {
+        self.mark_mcp_runtime_dirty();
+        self.schedule_mcp_prewarm();
+    }
+
+    pub(super) fn start_mcp_prewarm_worker(
+        self: &Arc<Self>,
+        requests: async_channel::Receiver<()>,
+        mut auth_changes: tokio::sync::watch::Receiver<u64>,
+    ) {
+        let session = Arc::downgrade(self);
+        let shutdown = self.mcp_prewarm_shutdown.clone();
+        let worker = self.services.runtime_handle.spawn(async move {
+            loop {
+                let auth_changed = tokio::select! {
+                    biased;
+                    _ = shutdown.cancelled() => break,
+                    request = requests.recv() => {
+                        if request.is_err() {
+                            break;
+                        }
+                        false
+                    },
+                    auth_change = auth_changes.changed() => {
+                        if auth_change.is_err() {
+                            break;
+                        }
+                        true
+                    },
+                };
+                let Some(session) = session.upgrade() else {
+                    break;
+                };
+                if auth_changed {
+                    session.mark_mcp_runtime_dirty();
+                }
+                tokio::select! {
+                    biased;
+                    _ = shutdown.cancelled() => break,
+                    _ = session.refresh_mcp_if_dirty() => {},
+                }
+            }
+        });
+        *self
+            .mcp_prewarm_task
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(worker);
+    }
+
+    pub(super) fn schedule_mcp_prewarm(&self) {
+        let _ = self.mcp_prewarm_tx.try_send(());
+    }
+
+    pub(super) async fn stop_mcp_prewarm_worker(&self) {
+        self.mcp_prewarm_shutdown.cancel();
+        let worker = self
+            .mcp_prewarm_task
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .take();
+        if let Some(worker) = worker
+            && let Err(error) = worker.await
+        {
+            warn!(%error, "MCP prewarm worker stopped unexpectedly");
+        }
+    }
+}

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

@@ -211,6 +211,7 @@ mod handlers;
 mod inject;
 mod input_queue;
 mod mcp;
+mod mcp_prewarm;
 mod mcp_runtime;
 pub(crate) mod multi_agents;
 mod review;
@@ -1483,7 +1484,7 @@ impl Session {
         updates: SessionSettingsUpdate,
     ) -> ConstraintResult<()> {
         let notify_config_contributors = !self.services.extensions.config_contributors().is_empty();
-        let (previous_config, new_config, permission_profile_changed) = {
+        let (previous_config, new_config, permission_profile_changed, mcp_inputs_changed) = {
             let mut state = self.state.lock().await;
             let updated = match state.session_configuration.apply(&updates) {
                 Ok(updated) => updated,
@@ -1507,17 +1508,25 @@ impl Session {
                     .turn_environments
                     .update_selections(updated.environment_selections());
             }
+            state.session_configuration = updated;
             if mcp_inputs_changed {
                 self.mark_mcp_runtime_dirty();
             }
-            state.session_configuration = updated;
-            (previous_config, new_config, permission_profile_changed)
+            (
+                previous_config,
+                new_config,
+                permission_profile_changed,
+                mcp_inputs_changed,
+            )
         };
         self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref());
         if permission_profile_changed {
             self.refresh_managed_network_proxy_for_current_permission_profile()
                 .await;
         }
+        if mcp_inputs_changed {
+            self.schedule_mcp_prewarm();
+        }
         Ok(())
     }
 
@@ -1633,6 +1642,7 @@ impl Session {
             (previous_config, new_config, config)
         };
         self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref());
+        self.schedule_mcp_prewarm();
         let environments = self.services.turn_environments.snapshot().await;
         let hooks = build_hooks_for_config(
             config.as_ref(),
@@ -1698,6 +1708,8 @@ impl Session {
         }
         state.session_configuration.original_config_do_not_use = Arc::new(config);
         self.mark_mcp_runtime_dirty();
+        drop(state);
+        self.schedule_mcp_prewarm();
     }
 
     fn emit_config_changed_contributors(

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

@@ -45,6 +45,9 @@ pub(crate) struct Session {
     pub(super) mcp_refresh_lock: Semaphore,
     pub(super) mcp_elicitation_reviewer_handle: OnceLock<codex_mcp::ElicitationReviewerHandle>,
     pub(super) mcp_elicitation_lifecycle_handle: OnceLock<codex_mcp::ElicitationLifecycle>,
+    pub(super) mcp_prewarm_tx: async_channel::Sender<()>,
+    pub(super) mcp_prewarm_shutdown: CancellationToken,
+    pub(super) mcp_prewarm_task: std::sync::Mutex<Option<JoinHandle<()>>>,
     pub(crate) conversation: Arc<RealtimeConversationManager>,
     pub(crate) active_turn: Mutex<Option<ActiveTurn>>,
     pub(crate) input_queue: InputQueue,
@@ -1145,6 +1148,7 @@ impl Session {
                 tool_search_handler_cache: Default::default(),
                 turn_environments: Arc::clone(&turn_environments),
             };
+            let (mcp_prewarm_tx, mcp_prewarm_rx) = async_channel::bounded(1);
             let sess = Arc::new(Session {
                 thread_id,
                 installation_id,
@@ -1159,6 +1163,9 @@ impl Session {
                 mcp_refresh_lock: Semaphore::new(/*permits*/ 1),
                 mcp_elicitation_reviewer_handle: OnceLock::new(),
                 mcp_elicitation_lifecycle_handle: OnceLock::new(),
+                mcp_prewarm_tx,
+                mcp_prewarm_shutdown: CancellationToken::new(),
+                mcp_prewarm_task: std::sync::Mutex::new(None),
                 conversation: Arc::new(RealtimeConversationManager::new()),
                 active_turn: Mutex::new(None),
                 input_queue: InputQueue::new(),
@@ -1209,6 +1216,7 @@ impl Session {
             }
             turn_environments.start_connection_event_forwarding(tx_event.clone());
 
+            let mcp_auth_changes = sess.services.auth_manager.auth_change_receiver();
             sess.install_initial_mcp_runtime(
                 &session_configuration,
                 mcp_projection,
@@ -1216,6 +1224,7 @@ impl Session {
                 mcp_runtime_cwd,
             )
             .await?;
+            sess.start_mcp_prewarm_worker(mcp_prewarm_rx, mcp_auth_changes);
             sess.schedule_startup_prewarm(session_configuration.base_instructions.clone())
                 .await;
             let session_start_source = match &initial_history {

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

@@ -5548,6 +5548,9 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
         mcp_refresh_lock: Semaphore::new(/*permits*/ 1),
         mcp_elicitation_reviewer_handle: OnceLock::new(),
         mcp_elicitation_lifecycle_handle: OnceLock::new(),
+        mcp_prewarm_tx: async_channel::bounded(1).0,
+        mcp_prewarm_shutdown: CancellationToken::new(),
+        mcp_prewarm_task: std::sync::Mutex::new(None),
         conversation: Arc::new(RealtimeConversationManager::new()),
         active_turn: Mutex::new(None),
         input_queue: super::input_queue::InputQueue::new(),
@@ -7715,6 +7718,9 @@ where
         mcp_refresh_lock: Semaphore::new(/*permits*/ 1),
         mcp_elicitation_reviewer_handle: OnceLock::new(),
         mcp_elicitation_lifecycle_handle: OnceLock::new(),
+        mcp_prewarm_tx: async_channel::bounded(1).0,
+        mcp_prewarm_shutdown: CancellationToken::new(),
+        mcp_prewarm_task: std::sync::Mutex::new(None),
         conversation: Arc::new(RealtimeConversationManager::new()),
         active_turn: Mutex::new(None),
         input_queue: super::input_queue::InputQueue::new(),
@@ -7907,6 +7913,29 @@ async fn mcp_elicitation_reviewer_is_reused_across_runtime_refreshes() {
     assert!(Arc::ptr_eq(&previous, &session.mcp_elicitation_reviewer()));
 }
 
+#[tokio::test]
+async fn mcp_policy_changes_schedule_runtime_refresh() {
+    let (session, _turn_context) = make_session_and_context().await;
+    let session = Arc::new(session);
+
+    session
+        .new_turn_with_sub_id(
+            "policy-change".to_string(),
+            SessionSettingsUpdate {
+                approval_policy: Some(AskForApproval::Never),
+                ..Default::default()
+            },
+        )
+        .await
+        .expect("approval policy update should succeed");
+
+    assert!(
+        session
+            .mcp_refresh_pending
+            .load(std::sync::atomic::Ordering::Acquire)
+    );
+}
+
 struct PendingNoiseConnectProvider;
 
 impl codex_exec_server::NoiseRendezvousConnectProvider for PendingNoiseConnectProvider {

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

@@ -638,6 +638,7 @@ impl Session {
                     state.session_configuration = next.clone();
                     Ok((
                         next,
+                        mcp_inputs_changed,
                         permission_profile_changed,
                         previous_config,
                         new_config,
@@ -647,23 +648,31 @@ impl Session {
             }
         };
 
-        let (session_configuration, permission_profile_changed, previous_config, new_config) =
-            match update_result {
-                Ok(update) => update,
-                Err(err) => {
-                    let message = err.to_string();
-                    self.send_event_raw(Event {
-                        id: sub_id.clone(),
-                        msg: EventMsg::Error(ErrorEvent {
-                            message: message.clone(),
-                            codex_error_info: Some(CodexErrorInfo::BadRequest),
-                        }),
-                    })
-                    .await;
-                    return Err(CodexErr::InvalidRequest(message));
-                }
-            };
+        let (
+            session_configuration,
+            mcp_inputs_changed,
+            permission_profile_changed,
+            previous_config,
+            new_config,
+        ) = match update_result {
+            Ok(update) => update,
+            Err(err) => {
+                let message = err.to_string();
+                self.send_event_raw(Event {
+                    id: sub_id.clone(),
+                    msg: EventMsg::Error(ErrorEvent {
+                        message: message.clone(),
+                        codex_error_info: Some(CodexErrorInfo::BadRequest),
+                    }),
+                })
+                .await;
+                return Err(CodexErr::InvalidRequest(message));
+            }
+        };
         self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref());
+        if mcp_inputs_changed {
+            self.schedule_mcp_prewarm();
+        }
 
         if permission_profile_changed {
             self.refresh_managed_network_proxy_for_current_permission_profile()
@@ -727,7 +736,6 @@ impl Session {
             .and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok())
             .unwrap_or_else(|| session_configuration.cwd().clone());
         let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone());
-
         let model_info = self
             .services
             .models_manager

codex-rs/core/src/tools/spec_plan_tests.rs

@@ -1119,6 +1119,7 @@ async fn tool_search_cache_rebuilds_when_deferred_world_state_changes() {
                 )],
                 tool_suggest_candidates: None,
                 extension_tool_executors: Vec::new(),
+                wait_for_environment_tool_config: None,
                 dynamic_tools: &[],
             },
             &cache,

codex-rs/core/tests/suite/mcp_tool_exposure.rs

@@ -3,6 +3,9 @@ use codex_config::Constrained;
 use codex_core::config::Config;
 use codex_extension_api::ExtensionFuture;
 use codex_extension_api::ExtensionRegistryBuilder;
+use codex_extension_api::McpServerContribution;
+use codex_extension_api::McpServerContributionContext;
+use codex_extension_api::McpServerContributor;
 use codex_extension_api::ThreadLifecycleContributor;
 use codex_extension_api::ThreadStartInput;
 use codex_features::Feature;
@@ -43,12 +46,70 @@ use serde_json::json;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::sync::Mutex;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
 use std::time::Duration;
+use tokio::sync::Semaphore;
 
 struct McpResourceClientCapture {
     client: Arc<Mutex<Option<McpResourceClient>>>,
 }
 
+struct CoalescingMcpContributor {
+    block_next: AtomicBool,
+    entered: Semaphore,
+    release: Semaphore,
+    observed_markers: Mutex<Vec<String>>,
+}
+
+impl CoalescingMcpContributor {
+    fn new() -> Self {
+        Self {
+            block_next: AtomicBool::new(false),
+            entered: Semaphore::new(0),
+            release: Semaphore::new(0),
+            observed_markers: Mutex::new(Vec::new()),
+        }
+    }
+}
+
+impl McpServerContributor<Config> for CoalescingMcpContributor {
+    fn id(&self) -> &'static str {
+        "coalescing_mcp_refresh_test"
+    }
+
+    fn contribute<'a>(
+        &'a self,
+        context: McpServerContributionContext<'a, Config>,
+    ) -> ExtensionFuture<'a, Vec<McpServerContribution>> {
+        Box::pin(async move {
+            let marker = context
+                .config()
+                .mcp_servers
+                .get()
+                .keys()
+                .next()
+                .cloned()
+                .unwrap_or_else(|| "initial".to_string());
+            self.observed_markers
+                .lock()
+                .expect("observed markers lock should not be poisoned")
+                .push(marker.clone());
+            if marker != "initial" {
+                self.entered.add_permits(1);
+            }
+            if self.block_next.swap(false, Ordering::SeqCst) {
+                self.release
+                    .acquire()
+                    .await
+                    .expect("release semaphore should remain open")
+                    .forget();
+            }
+            Vec::new()
+        })
+    }
+}
+
 fn format_labeled_requests_snapshot(
     scenario: &str,
     sections: &[(&str, &ResponsesRequest)],
@@ -108,6 +169,89 @@ impl ThreadLifecycleContributor<Config> for McpResourceClientCapture {
     }
 }
 
+fn config_with_mcp_marker(base: &Config, marker: &str) -> Config {
+    let mut config = base.clone();
+    let server = serde_json::from_value(json!({
+        "url": "http://127.0.0.1:1/mcp",
+        "enabled": false,
+    }))
+    .expect("test MCP server config");
+    config
+        .mcp_servers
+        .set(HashMap::from([(marker.to_string(), server)]))
+        .expect("test config should allow MCP servers");
+    config
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn rapid_mcp_refreshes_coalesce_to_the_latest_config() -> 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 contributor = Arc::new(CoalescingMcpContributor::new());
+    let mut extensions = ExtensionRegistryBuilder::<Config>::new();
+    extensions.mcp_server_contributor(contributor.clone());
+    let test = core_test_support::test_codex::test_codex()
+        .with_extensions(Arc::new(extensions.build()))
+        .build(&server)
+        .await?;
+
+    contributor.block_next.store(true, Ordering::SeqCst);
+    test.codex
+        .refresh_runtime_config(config_with_mcp_marker(&test.config, "config-a"))
+        .await;
+    tokio::time::timeout(Duration::from_secs(5), contributor.entered.acquire())
+        .await
+        .expect("configuration A should enter MCP projection")
+        .expect("entered semaphore should remain open")
+        .forget();
+
+    test.codex
+        .refresh_runtime_config(config_with_mcp_marker(&test.config, "config-b"))
+        .await;
+    test.codex
+        .refresh_runtime_config(config_with_mcp_marker(&test.config, "config-c"))
+        .await;
+    contributor.release.add_permits(1);
+
+    tokio::time::timeout(Duration::from_secs(5), contributor.entered.acquire())
+        .await
+        .expect("the coalesced refresh should project the latest configuration")
+        .expect("entered semaphore should remain open")
+        .forget();
+    let observed = contributor
+        .observed_markers
+        .lock()
+        .expect("observed markers lock should not be poisoned")
+        .iter()
+        .filter(|marker| marker.as_str() != "initial")
+        .cloned()
+        .collect::<Vec<_>>();
+    assert_eq!(
+        observed,
+        vec!["config-a".to_string(), "config-c".to_string()]
+    );
+
+    test.submit_turn("bind the latest MCP state").await?;
+    assert!(
+        !contributor
+            .observed_markers
+            .lock()
+            .expect("observed markers lock should not be poisoned")
+            .iter()
+            .any(|marker| marker == "config-b")
+    );
+    response.single_request();
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn out_of_band_resource_read_reconciles_the_published_mcp_runtime() -> Result<()> {
     let server = responses::start_mock_server().await;