From ae12a9d645774982fe9deda590dc1f237d249d1d Mon Sep 17 00:00:00 2001 From: Krish Chainani Date: Fri, 13 Nov 2026 17:11:52 +0000 Subject: [PATCH] Add deferred MCP reprojection for the next regular turn (#61356) ## What changed Add `CodexThread::reproject_mcp_on_next_turn()` to request an MCP reprojection after `TurnStarted` and before tools are bound. Repeated requests coalesce, and idle reconciliation or prewarming leaves the request pending until the next regular turn. ## Testing Add coverage for deferred requests surviving prewarm and idle reconciliation, coalescing into one reprojection, and exposing refreshed tools in the turn's first model request. Verify that discovery starts after `TurnStarted`. GitOrigin-RevId: 260b2381013291a47eb9a7f736c129de76c5897f --- codex-rs/core/src/codex_thread.rs | 6 + codex-rs/core/src/session/mcp_prewarm.rs | 10 ++ codex-rs/core/src/session/mcp_refresh.rs | 13 ++ .../core/src/session/mcp_refresh_tests.rs | 15 ++ .../core/src/session/turn_start_mcp_tests.rs | 2 +- codex-rs/core/src/tasks/regular.rs | 1 + .../core/tests/suite/mcp_tool_exposure.rs | 143 +++++++++++++++++- 7 files changed, 186 insertions(+), 4 deletions(-) diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 6540c74665..0e8bde3dde 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -1276,6 +1276,12 @@ impl CodexThread { self.session.hooks().list_with_environment_hooks(sources) } + /// Rebuilds the MCP projection on the next regular turn, after turn start and before + /// tools are bound. Repeated requests coalesce without starting a background refresh. + pub fn reproject_mcp_on_next_turn(&self) { + self.session.reproject_mcp_on_next_turn(); + } + /// Observes this thread's published MCP connections that match the requested config. pub async fn mcp_connection_statuses( &self, diff --git a/codex-rs/core/src/session/mcp_prewarm.rs b/codex-rs/core/src/session/mcp_prewarm.rs index 49e828c863..4df2ec1674 100644 --- a/codex-rs/core/src/session/mcp_prewarm.rs +++ b/codex-rs/core/src/session/mcp_prewarm.rs @@ -7,6 +7,16 @@ use super::*; use crate::McpRefreshReason; impl Session { + pub(crate) fn reproject_mcp_on_next_turn(&self) { + self.mcp_refresh.request_for_next_turn(); + } + + pub(crate) fn start_deferred_mcp_reprojection(&self) { + if self.mcp_refresh.claim_for_turn_start() { + self.request_mcp_runtime_reprojection(McpRefreshReason::ExplicitInvalidation); + } + } + pub(crate) fn request_mcp_runtime_refresh(&self, reason: McpRefreshReason) { // Plugin changes can reuse connections but still change their skill resources. self.services.mcp_runtime.invalidate_resource_caches(); diff --git a/codex-rs/core/src/session/mcp_refresh.rs b/codex-rs/core/src/session/mcp_refresh.rs index d2a94f32d2..5cc6b69110 100644 --- a/codex-rs/core/src/session/mcp_refresh.rs +++ b/codex-rs/core/src/session/mcp_refresh.rs @@ -2,6 +2,8 @@ use crate::McpRefreshReason; use codex_protocol::ThreadId; use std::collections::BTreeSet; use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use tokio::sync::AcquireError; use tokio::sync::Semaphore; use tokio::sync::SemaphorePermit; @@ -10,6 +12,8 @@ use tokio_util::task::TaskTracker; /// Owns MCP invalidation and the single gate used to publish runtime updates. pub(super) struct McpRefresh { pending: Mutex>, + // Kept separate from pending so prewarm cannot claim work reserved for a real turn. + on_next_turn: AtomicBool, gate: Semaphore, pub(super) apps_refreshes: TaskTracker, } @@ -21,11 +25,20 @@ impl McpRefresh { apps_refreshes.close(); Self { pending: Mutex::default(), + on_next_turn: AtomicBool::new(false), gate: Semaphore::new(/*permits*/ 1), apps_refreshes, } } + pub(super) fn request_for_next_turn(&self) { + self.on_next_turn.store(true, Ordering::Release); + } + + pub(super) fn claim_for_turn_start(&self) -> bool { + self.on_next_turn.swap(false, Ordering::AcqRel) + } + pub(super) fn invalidate(&self, reason: McpRefreshReason) { self.pending .lock() diff --git a/codex-rs/core/src/session/mcp_refresh_tests.rs b/codex-rs/core/src/session/mcp_refresh_tests.rs index cb2d402b96..0f6026c3f7 100644 --- a/codex-rs/core/src/session/mcp_refresh_tests.rs +++ b/codex-rs/core/src/session/mcp_refresh_tests.rs @@ -3,6 +3,21 @@ use super::*; use pretty_assertions::assert_eq; +#[test] +fn next_turn_reprojection_is_not_claimed_by_prewarm() { + let refresh = McpRefresh::new(); + let thread_id = ThreadId::new(); + refresh.request_for_next_turn(); + refresh.request_for_next_turn(); + assert!(refresh.claim(thread_id).is_none()); + refresh.invalidate(McpRefreshReason::McpConfigReloaded); + let mut prewarm = refresh.claim(thread_id).expect("unrelated refresh"); + prewarm.published = true; + drop(prewarm); + assert!(refresh.claim_for_turn_start()); + assert!(!refresh.claim_for_turn_start()); +} + #[test] fn refresh_keeps_concurrent_invalidations_and_restores_unpublished_causes() { let thread_id = ThreadId::new(); diff --git a/codex-rs/core/src/session/turn_start_mcp_tests.rs b/codex-rs/core/src/session/turn_start_mcp_tests.rs index 65eae0ba8c..3bf66d8955 100644 --- a/codex-rs/core/src/session/turn_start_mcp_tests.rs +++ b/codex-rs/core/src/session/turn_start_mcp_tests.rs @@ -74,7 +74,7 @@ async fn turn_start_refreshes_dirty_mcp_reuses_clean_runtime_and_cancels_discove .requires_mcp .store(turn > 0, std::sync::atomic::Ordering::Relaxed); if turn == 0 { - session.mark_mcp_runtime_dirty(McpRefreshReason::ExplicitInvalidation); + session.reproject_mcp_on_next_turn(); } else if turn == 3 { session.services.mcp_runtime.reconnect_on_next_refresh(); session.request_mcp_runtime_refresh(McpRefreshReason::ExplicitInvalidation); diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 1923d1fbe8..dcb07bfd48 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -73,6 +73,7 @@ impl SessionTask for RegularTask { // not wait on startup prewarm resolution. let prewarmed_client_session = async { sess.emit_turn_started(&ctx, TaskKind::Regular).await; + sess.start_deferred_mcp_reprojection(); // Regular-start contributors run once, after the task is visible and interruptible. let prepares_mcp = sess .services diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index 59d01135fa..3408e536fc 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -1,5 +1,6 @@ use anyhow::Result; use codex_config::Constrained; +use codex_core::CodexThread; use codex_core::EnvironmentConfig; use codex_core::TurnInputRequest; use codex_core::TurnStartOptions; @@ -17,6 +18,7 @@ use codex_history::RolloutItem; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::CodexAppsResourceListParams; use codex_mcp::McpResourceClient; +use codex_mcp::McpSnapshotDetail; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::WindowsSandboxLevel; @@ -60,6 +62,7 @@ use serde_json::json; use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; +use std::sync::Weak; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -110,6 +113,9 @@ struct CoalescingMcpContributor { entered: Semaphore, release: Semaphore, observed_markers: Mutex>, + contributions: Mutex>, + discovery_thread: Mutex>>, + saw_turn_started: AtomicBool, } struct AppsMcpServerContributor { @@ -129,6 +135,9 @@ impl CoalescingMcpContributor { entered: Semaphore::new(0), release: Semaphore::new(0), observed_markers: Mutex::new(Vec::new()), + contributions: Mutex::new(Vec::new()), + discovery_thread: Mutex::new(None), + saw_turn_started: AtomicBool::new(false), } } } @@ -155,17 +164,39 @@ impl McpServerContributor for CoalescingMcpContributor { .lock() .expect("observed markers lock should not be poisoned") .push(marker.clone()); - if marker != "initial" { + let discovery_thread = self + .discovery_thread + .lock() + .expect("discovery thread lock should not be poisoned") + .take(); + if let Some(thread) = discovery_thread.and_then(|thread| thread.upgrade()) { + // The test transfers its sole event reader here until discovery is released. + // Snapshot the queue so a later TurnStarted cannot satisfy the ordering check. + for _ in 0..thread.queued_event_count() { + let event = thread + .next_event() + .await + .expect("thread should remain open"); + if matches!(event.msg, EventMsg::TurnStarted(_)) { + self.saw_turn_started.store(true, Ordering::SeqCst); + } + } + } + let block_next = self.block_next.swap(false, Ordering::SeqCst); + if marker != "initial" || block_next { self.entered.add_permits(1); } - if self.block_next.swap(false, Ordering::SeqCst) { + if block_next { self.release .acquire() .await .expect("release semaphore should remain open") .forget(); } - Vec::new() + self.contributions + .lock() + .expect("contributions lock should not be poisoned") + .clone() }) } } @@ -475,6 +506,112 @@ async fn rapid_mcp_refreshes_coalesce_to_the_latest_config() -> Result<()> { Ok(()) } +/// Deferred requests survive idle reconciliation and coalesce into a refresh after TurnStarted, +/// with the refreshed tools available to the next turn's first model request. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn deferred_mcp_reprojection_waits_for_turn_start_and_binds_refreshed_tools() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let apps_server = AppsTestServer::mount(&server).await?; + let mcp_config = serde_json::from_value(json!({ + "url": format!("{}/api/codex/ps/mcp", apps_server.chatgpt_base_url), + "required": true, + }))?; + 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 mut builder = core_test_support::test_codex::test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| { + config + .mcp_servers + .set(HashMap::from([("initial".to_string(), mcp_config)])) + .expect("test config should allow MCP servers"); + config + .features + .enable(Feature::CodeModeOnly) + .expect("test config should allow feature update"); + config.code_mode.direct_only_tool_namespaces = vec!["mcp__refreshed".to_string()]; + }); + let test = builder.build_with_auto_env(&server).await?; + wait_for_mcp_server(&test.codex, "initial").await?; + test.codex + .mcp_server_status_snapshot( + "initial", + /*environment_id*/ None, + McpSnapshotDetail::ToolsAndAuthOnly, + /*expected_registration*/ None, + ) + .await?; + let initial_projections = contributor.observed_markers.lock().unwrap().len(); + + *contributor.contributions.lock().unwrap() = vec![McpServerContribution::Set { + name: "refreshed".to_string(), + config: Box::new(test.config.mcp_servers.get()["initial"].clone()), + }]; + test.codex.reproject_mcp_on_next_turn(); + test.codex.reproject_mcp_on_next_turn(); + + // Exercise ordinary reconciliation while idle; it must leave deferred work for the turn. + test.codex + .mcp_server_status_snapshot( + "initial", + /*environment_id*/ None, + McpSnapshotDetail::ToolsAndAuthOnly, + /*expected_registration*/ None, + ) + .await?; + assert_eq!( + contributor.observed_markers.lock().unwrap().len(), + initial_projections + ); + + *contributor.discovery_thread.lock().unwrap() = Some(Arc::downgrade(&test.codex)); + contributor.block_next.store(true, Ordering::SeqCst); + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "use the refreshed tools".to_string(), + text_elements: Vec::new(), + }])) + .await?; + tokio::time::timeout(Duration::from_secs(5), contributor.entered.acquire()) + .await + .expect("the regular turn should start deferred discovery")? + .forget(); + assert!( + contributor.saw_turn_started.load(Ordering::SeqCst), + "TurnStarted must already be queued when discovery begins" + ); + assert!(response.requests().is_empty()); + contributor.release.add_permits(1); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + assert_eq!( + contributor.observed_markers.lock().unwrap().len(), + initial_projections + 1, + "repeated requests should cause only one deferred projection" + ); + let body = response.single_request().body_json(); + assert!( + namespace_child_tool(&body, "mcp__refreshed", "calendar_list_events").is_some(), + "the first request must include the refreshed tool: {body}" + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn root_reconciliation_reuses_pending_apps_startup() -> Result<()> { skip_if_no_network!(Ok(()));