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,