From ef2d3edb959a75a90665ade37d554f3ca65fe880 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 24 Jul 2026 09:35:03 +0000 Subject: [PATCH] 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 | 11 +- codex-rs/core/src/session/mcp.rs | 2 +- codex-rs/core/src/session/mcp_prewarm.rs | 75 +++++++++ codex-rs/core/src/session/mod.rs | 18 ++- codex-rs/core/src/session/session.rs | 9 ++ codex-rs/core/src/session/tests.rs | 29 ++++ codex-rs/core/src/session/turn_context.rs | 42 ++--- codex-rs/core/src/tools/spec_plan_tests.rs | 1 + .../core/tests/suite/mcp_tool_exposure.rs | 144 ++++++++++++++++++ 9 files changed, 306 insertions(+), 25 deletions(-) create mode 100644 codex-rs/core/src/session/mcp_prewarm.rs diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index ee2a0d3624..485769eba8 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -422,7 +422,7 @@ pub async fn dynamic_tool_response(sess: &Arc, 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) { @@ -585,9 +585,12 @@ 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}"); } - 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; diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 76f83306c0..0a467a94fd 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/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(()) } diff --git a/codex-rs/core/src/session/mcp_prewarm.rs b/codex-rs/core/src/session/mcp_prewarm.rs new file mode 100644 index 0000000000..e923b6b86f --- /dev/null +++ b/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, + requests: async_channel::Receiver<()>, + mut auth_changes: tokio::sync::watch::Receiver, + ) { + 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"); + } + } +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 5af8a34a3d..ec1ad75175 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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( diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 41c24df898..ac3e357f0b 100644 --- a/codex-rs/core/src/session/session.rs +++ b/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, pub(super) mcp_elicitation_lifecycle_handle: OnceLock, + pub(super) mcp_prewarm_tx: async_channel::Sender<()>, + pub(super) mcp_prewarm_shutdown: CancellationToken, + pub(super) mcp_prewarm_task: std::sync::Mutex>>, pub(crate) conversation: Arc, pub(crate) active_turn: Mutex>, 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 { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 308f39de98..edc39358a3 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/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 { diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 124842e8f5..49896774a0 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/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 diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 3957f42202..87be4d9275 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/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, diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index b1b6cf8266..92bc905779 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/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>>, } +struct CoalescingMcpContributor { + block_next: AtomicBool, + entered: Semaphore, + release: Semaphore, + observed_markers: Mutex>, +} + +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 for CoalescingMcpContributor { + fn id(&self) -> &'static str { + "coalescing_mcp_refresh_test" + } + + fn contribute<'a>( + &'a self, + context: McpServerContributionContext<'a, Config>, + ) -> ExtensionFuture<'a, Vec> { + 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 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::::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::>(); + 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;