diff --git a/codex-rs/codex-mcp/src/binding_tests.rs b/codex-rs/codex-mcp/src/binding_tests.rs index 58ec1aeabc..2ffb3e86bd 100644 --- a/codex-rs/codex-mcp/src/binding_tests.rs +++ b/codex-rs/codex-mcp/src/binding_tests.rs @@ -356,6 +356,7 @@ async fn prepared_call_is_rejected_after_catalog_refresh() { step.tool_catalog .refresh( || async { Ok((step.step.tools().to_vec(), None, ())) }, + || Some(()), |_, _, ()| {}, ) .await @@ -390,6 +391,7 @@ async fn stale_prepared_call_does_not_run_preparation() { step.tool_catalog .refresh( || async { Ok((step.step.tools().to_vec(), None, ())) }, + || Some(()), |_, _, ()| {}, ) .await @@ -435,8 +437,13 @@ async fn preparation_holds_catalog_authority_until_it_finishes() { }); preparation_started.notified().await; + let validation_reached = AtomicBool::new(false); let refresh = step.tool_catalog.refresh( || async { Ok((step.step.tools().to_vec(), None, ())) }, + || { + validation_reached.store(true, Ordering::SeqCst); + Some(()) + }, |_, _, ()| {}, ); tokio::pin!(refresh); @@ -444,6 +451,10 @@ async fn preparation_holds_catalog_authority_until_it_finishes() { futures::poll!(&mut refresh).is_pending(), "catalog replacement must wait for irreversible call preparation" ); + assert!( + !validation_reached.load(Ordering::SeqCst), + "publication validation must not hold its guard while calls drain" + ); finish_preparation.notify_one(); call.await .expect("call task should finish") @@ -451,6 +462,7 @@ async fn preparation_holds_catalog_authority_until_it_finishes() { refresh .await .expect("catalog refresh should finish after preparation"); + assert!(validation_reached.load(Ordering::SeqCst)); } #[tokio::test] diff --git a/codex-rs/codex-mcp/src/client_tool_catalog.rs b/codex-rs/codex-mcp/src/client_tool_catalog.rs index cec29a9190..38de69879e 100644 --- a/codex-rs/codex-mcp/src/client_tool_catalog.rs +++ b/codex-rs/codex-mcp/src/client_tool_catalog.rs @@ -232,22 +232,28 @@ impl ClientToolCatalog { result } - /// Serialize fetching and publication, leaving the current catalog usable during the fetch. - /// The publication callback runs alongside the exact-client update under the write lock. + /// Validate under the caller's publication guard after fetching and draining calls. + /// The guard covers both shared-provider publication and this client's revision update. #[expect( clippy::await_holding_invalid_type, reason = "refreshes must remain serialized through fetching and catalog publication" )] - pub(crate) async fn refresh(&self, fetch: F, publish: P) -> Result + pub(crate) async fn refresh( + &self, + fetch: impl FnOnce() -> Fut, + validate: impl FnOnce() -> Option, + publish: impl FnOnce(&[ToolInfo], Option>, C) -> R, + ) -> Result> where - F: FnOnce() -> Fut, Fut: Future, Option>, C)>>, - P: FnOnce(&[ToolInfo], Option>, C) -> R, { let _refresh = self.refresh_lock.lock().await; let (tools, meta, context) = fetch().await?; let _calls = self.calls.write().await; let mut current = self.current.write().await; + let Some(_validity) = validate() else { + return Ok(None); + }; let result = publish(&tools, meta.clone(), context); match &mut current.source { CatalogSource::Local { @@ -269,7 +275,7 @@ impl ClientToolCatalog { } current.revision += 1; current.last_refresh_revision = current.revision; - Ok(result) + Ok(Some(result)) } /// Reject stale calls before preparation. Explicit refreshes wait for execution; diff --git a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs index 58d8f0dc88..9ac9a8932e 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -660,16 +660,19 @@ impl McpConnectionSet { } /// Refreshes one exact Apps catalog, preserving the raw inventory for app policy. - pub(crate) async fn refresh_codex_apps_client_catalog( + pub(crate) async fn refresh_codex_apps_client_catalog( &self, config: &crate::McpConfig, - ) -> Result { + validate: impl FnOnce() -> Option, + ) -> Result> { let refresh_start = Instant::now(); let view = self .servers .get(CODEX_APPS_MCP_SERVER_NAME) .ok_or_else(|| anyhow!("unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"))?; - let (tools, _) = self.refresh_codex_apps_tool_catalog().await?; + let Some((tools, _)) = self.refresh_codex_apps_tool_catalog(validate).await? else { + return Ok(None); + }; let server_has_permission = config .permission_profile_for_server(CODEX_APPS_MCP_SERVER_NAME) .is_some(); @@ -690,10 +693,10 @@ impl McpConnectionSet { refresh_start.elapsed(), &[("path", "legacy"), ("trigger", "explicit")], ); - Ok(CodexAppsToolSnapshot { + Ok(Some(CodexAppsToolSnapshot { tools, model_visible_tool_names, - }) + })) } /// Refreshes Apps tools and returns the prepared shared-cache winner for discovery. @@ -703,7 +706,10 @@ impl McpConnectionSet { .servers .get(CODEX_APPS_MCP_SERVER_NAME) .ok_or_else(|| anyhow!("unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"))?; - let (_, tools) = self.refresh_codex_apps_tool_catalog().await?; + let (_, tools) = self + .refresh_codex_apps_tool_catalog(|| Some(())) + .await? + .context("Apps discovery refresh skipped publication")?; let tools = prepare_codex_apps_tools_for_model( filter_tools(tools, &view.tool_filter), &self.tool_plugin_context, @@ -744,7 +750,10 @@ impl McpConnectionSet { } /// Publishes the exact client catalog and returns both raw inventories. - async fn refresh_codex_apps_tool_catalog(&self) -> Result<(Vec, Vec)> { + async fn refresh_codex_apps_tool_catalog( + &self, + validate: impl FnOnce() -> Option, + ) -> Result, Vec)>> { let view = self .servers .get(CODEX_APPS_MCP_SERVER_NAME) @@ -754,7 +763,7 @@ impl McpConnectionSet { .client() .await .context("failed to get client")?; - let (tools, list_start) = managed_client + let Some((tools, list_start)) = managed_client .tool_catalog .refresh( || async { @@ -781,6 +790,7 @@ impl McpConnectionSet { })?; Ok((client_tools, meta, (fetch_ticket, list_start))) }, + validate, |client_tools, meta, (fetch_ticket, list_start)| { // Discovery can accept another scope's winner; executable catalogs // receive only the latest successful fetch from their own scope. @@ -803,7 +813,10 @@ impl McpConnectionSet { (tools, list_start) }, ) - .await?; + .await? + else { + return Ok(None); + }; let client_tools = managed_client .tool_catalog .read(|catalog| catalog.tools.to_vec()) @@ -813,7 +826,7 @@ impl McpConnectionSet { list_start.elapsed(), &[("cache", "miss")], ); - Ok((client_tools, tools)) + Ok(Some((client_tools, tools))) } fn with_server_metadata(mut tool: ToolInfo, metadata: &McpServerMetadata) -> ToolInfo { diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index ca4edee537..2216fecdda 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -2701,13 +2701,13 @@ async fn hard_refresh_keeps_client_catalog_local_when_shared_cache_loses_race() let config_for_refresh = config.clone(); let refresh_a = tokio::spawn(async move { manager_a_for_refresh - .refresh_codex_apps_client_catalog(&config_for_refresh) + .refresh_codex_apps_client_catalog(&config_for_refresh, || Some(())) .await }); list_started.notified().await; manager_b.refresh_codex_apps_tools_for_discovery().await?; release_list.notify_one(); - let snapshot_a = refresh_a.await??; + let snapshot_a = refresh_a.await??.expect("Apps catalog was published"); assert_eq!( snapshot_a .tools @@ -2737,7 +2737,10 @@ async fn hard_refresh_keeps_client_catalog_local_when_shared_cache_loses_race() .tool_filter .disabled .insert("b_only".to_string()); - let snapshot_b = manager_b.refresh_codex_apps_client_catalog(&config).await?; + let snapshot_b = manager_b + .refresh_codex_apps_client_catalog(&config, || Some(())) + .await? + .expect("Apps catalog was published"); assert_eq!( ( snapshot_b @@ -5743,6 +5746,7 @@ async fn apps_catalog_broadcast_restores_equivalent_calls() -> anyhow::Result<() catalog .refresh( || async { Ok((snapshot.tools.to_vec(), None, ())) }, + || Some(()), |tools, _, ()| store_current_tools(&context, tools.to_vec()), ) .await?; @@ -5801,6 +5805,7 @@ async fn apps_catalog_broadcast_survives_an_older_local_refresh() -> anyhow::Res catalog .refresh( || async { Ok((Vec::new(), None, older_ticket)) }, + || Some(()), |tools, _, ticket| { context .publish_runtime_if_newest_accepted( @@ -5859,7 +5864,11 @@ async fn refreshed_catalog_follows_reused_client_without_mutating_old_bindings() let startup_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "startup")]; client .tool_catalog - .refresh(|| async { Ok((startup_tools, None, ())) }, |_, _, ()| {}) + .refresh( + || async { Ok((startup_tools, None, ())) }, + || Some(()), + |_, _, ()| {}, + ) .await?; let old_binding = capture_binding(&previous).await; let old_tools = serde_json::to_value(old_binding.tools())?; diff --git a/codex-rs/codex-mcp/src/connection_pool_tests.rs b/codex-rs/codex-mcp/src/connection_pool_tests.rs index 932f39d61a..64f2491096 100644 --- a/codex-rs/codex-mcp/src/connection_pool_tests.rs +++ b/codex-rs/codex-mcp/src/connection_pool_tests.rs @@ -767,6 +767,7 @@ async fn new_chat_relists_tools_without_invalidating_or_waiting_for_existing_cal Duration::from_secs(1), second.tool_catalog.refresh( || async { Ok((vec![create_test_tool("test_server", "refreshed")], None, ())) }, + || Some(()), |_, _, ()| {}, ), ) diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 166f675a66..35ac0f980c 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -6,6 +6,9 @@ //! [`crate::connection_manager`]. Registry removal revokes EMA authority before //! the thread becomes untracked. +#[cfg(test)] +#[path = "runtime/apps_tool_refresh_tests.rs"] +mod apps_tool_refresh_tests; mod status; #[cfg(test)] @@ -124,6 +127,8 @@ pub struct McpRuntimeInput { /// their exact connections and configuration for as long as they are needed. pub struct McpRuntime { current: ArcSwap, + // This mutex is held only through synchronous publication, never network I/O or call drain. + catalog_publication: Mutex, // Old bindings keep leases alive after publication; cancellation alone does not release them. superseded_connections: Mutex>>, event_stream_cancellation: Mutex, @@ -263,6 +268,15 @@ impl McpPublicationGate { } impl McpRuntime { + /// Supersede pending fetches as soon as Core requests a configuration change. + pub fn invalidate_apps_tool_refresh(&self) { + let mut generation = self + .catalog_publication + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + *generation = generation.wrapping_add(1); + } + /// Creates a runtime with no configured servers. /// /// This is useful while constructing a thread that must publish a stable @@ -283,6 +297,7 @@ impl McpRuntime { ready_environments: HashMap::new(), cached_binding: Mutex::new(None), }), + catalog_publication: Mutex::new(0), superseded_connections: Mutex::default(), event_stream_cancellation: Mutex::new(EventStreamCancellation { event_server_available: false, @@ -389,6 +404,8 @@ impl McpRuntime { } async fn publish(&self, input: McpRuntimeInput, reuse_previous: bool) { + // Also cover callers which replace the runtime without Core's dirty-state machinery. + self.invalidate_apps_tool_refresh(); let current = self.current.load_full(); let retain_ema = current.config.is_some(); let ema_auth = (input.config.xaa_enabled @@ -475,6 +492,10 @@ impl McpRuntime { superseded.retain(|connections| connections.strong_count() > 0); superseded.push(Arc::downgrade(¤t.connections)); } + let publication = self + .catalog_publication + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); self.current.store(Arc::new(PublishedMcpRuntime { connections, config: Some(config), @@ -489,6 +510,7 @@ impl McpRuntime { ready_environments, cached_binding: Mutex::new(None), })); + drop(publication); let _ = publish.send(true); cancellation.event_server_available = hosted_event_server_retained; if !hosted_event_server_retained { @@ -732,17 +754,51 @@ impl McpRuntime { .await } - /// Refreshes the published Apps client and returns its exact inventory and MCP eligibility. - pub async fn refresh_codex_apps_tools(&self) -> anyhow::Result { - let current = self.current.load_full(); - let config = current - .config - .as_deref() - .ok_or_else(|| anyhow::anyhow!("MCP runtime is not configured"))?; - current - .connections - .refresh_codex_apps_client_catalog(config) - .await + /// Captures the current Apps client under Core's session permit. + /// Await the returned fetch after releasing the permit. A stale result is skipped. + pub fn capture_apps_tool_refresh( + &self, + ) -> impl Future>> + Send { + let (generation, current) = { + let generation = self + .catalog_publication + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + (*generation, self.current.load_full()) + }; + let auth_changes = current.auth_manager.as_ref().map(|manager| { + let mut changes = manager.auth_change_state_receiver(); + // Subscribe before reading auth to catch changes during capture. + if !current.auth_matches(manager.auth_cached().as_ref()) { + changes.mark_changed(); + } + changes + }); + async move { + let validate = || { + let publication = self + .catalog_publication + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // Hold the watch read guard through publication. Auth writers take cache then watch. + let auth = auth_changes.as_ref().map(watch::Receiver::borrow); + (*publication == generation + && Arc::ptr_eq(&self.current.load(), ¤t) + && auth_changes + .as_ref() + .is_none_or(|changes| matches!(changes.has_changed(), Ok(false)))) + .then_some((publication, auth)) + }; + let config = current + .config + .as_deref() + .ok_or_else(|| anyhow::anyhow!("MCP runtime is not configured"))?; + let snapshot = current + .connections + .refresh_codex_apps_client_catalog(config, validate) + .await?; + Ok(snapshot.filter(|_| validate().is_some())) + } } /// Returns tools from this thread’s ready, unshared Apps MCP client. diff --git a/codex-rs/codex-mcp/src/runtime/apps_tool_refresh_tests.rs b/codex-rs/codex-mcp/src/runtime/apps_tool_refresh_tests.rs new file mode 100644 index 0000000000..0009faaecd --- /dev/null +++ b/codex-rs/codex-mcp/src/runtime/apps_tool_refresh_tests.rs @@ -0,0 +1,91 @@ +#![cfg(test)] +//! Verifies that an in-flight Apps refresh cannot publish after invalidation. + +use super::*; +use tokio::sync::Notify; + +#[tokio::test] +async fn stale_apps_refresh_does_not_publish() -> anyhow::Result<()> { + let home = tempfile::tempdir()?; + let auth = CodexAuth::from_external_chatgpt_tokens( + "header.e30.token", + "account", + /*chatgpt_plan_type*/ None, + )?; + let auth_manager = + AuthManager::from_auth_for_testing_with_home(auth.clone(), home.path().to_path_buf()); + let cache = ConnectorRuntimeManager::::default().context( + home.path().to_path_buf(), + ConnectorRuntimeContextKey::personal( + /*account_id*/ None, /*chatgpt_user_id*/ None, + ), + ); + let started = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let connections = crate::connection_manager::tests::create_test_manager_with_ready_apps_client( + cache.clone(), + "search", + Some(started.clone()), + Some(release.clone()), + ) + .await?; + connections.list_all_tools().await; + let revision = connections + .stable_catalog_revisions(/*required_servers*/ &[], &HashSet::new()) + .await + .expect("Apps catalog is ready"); + let runtime = McpRuntime::empty(/*prefix_mcp_tool_names*/ true); + runtime.current.store(Arc::new(PublishedMcpRuntime { + connections: connections.clone(), + config: Some(Arc::new(crate::mcp::tests::test_mcp_config( + home.path().to_path_buf(), + ))), + auth_token: Some(auth.get_token()?), + auth: Some(auth), + auth_manager: Some(auth_manager.clone()), + auth_owner: None, + auth_generation: Arc::new(()), + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + environment_selections: Arc::default(), + ready_environments: HashMap::new(), + cached_binding: Mutex::new(None), + })); + + let refresh = runtime.capture_apps_tool_refresh(); + tokio::pin!(refresh); + tokio::select! { + _ = started.notified() => {} + _ = &mut refresh => panic!("refresh completed before tools/list was released"), + } + runtime.invalidate_apps_tool_refresh(); + release.notify_one(); + assert!( + refresh.await?.is_none(), + "runtime invalidation supersedes the fetch" + ); + + let refresh = runtime.capture_apps_tool_refresh(); + tokio::pin!(refresh); + tokio::select! { + _ = started.notified() => {} + _ = &mut refresh => panic!("refresh completed before tools/list was released"), + } + auth_manager.logout().await?; + release.notify_one(); + assert!(refresh.await?.is_none(), "logout supersedes the fetch"); + + assert!( + !cache.has_current_tools(), + "stale tools reached the shared cache" + ); + assert!( + connections + .stable_catalog_revisions(/*required_servers*/ &[], &HashSet::new()) + .await + .as_ref() + == Some(&revision), + "a stale fetch invalidated the existing client catalog" + ); + Ok(()) +} diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 1a4424c89e..f055cdba19 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -358,8 +358,9 @@ pub(super) async fn shutdown_session_runtime(sess: &Arc) { { let _refresh = sess.mcp_refresh.acquire().await; sess.mcp_refresh.close(); - sess.services.mcp_runtime.shutdown().await; } + sess.wait_for_apps_tool_refresh().await; + sess.services.mcp_runtime.shutdown().await; sess.drain_code_mode_messages().await; diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 4865b61d38..b2ec905920 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -3,6 +3,7 @@ use crate::McpRefreshReason; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::environment_selection::combine_selected_capability_roots; use crate::mcp_tool_call::McpToolApprovalContext; +use anyhow::Context; use codex_exec_server::DiscoverV2CapabilitiesRequest; use codex_exec_server::ExecutorCapabilityDiscoveryCache; use codex_exec_server::ExecutorCapabilityDiscoverySnapshot; @@ -38,6 +39,7 @@ use codex_protocol::openai_models::GuardianScope; use codex_rmcp_client::Elicitation; use rmcp::model::ElicitationAction; use serde_json::Map; +use tokio::sync::SemaphorePermit; const TOOL_SUGGESTION_ACTION_INSTALL: &str = "install"; const TOOL_SUGGESTION_ACTION_KEY: &str = "suggest_type"; @@ -194,6 +196,11 @@ impl Session { error!("MCP runtime refresh semaphore closed"); return; }; + self.reconcile_mcp_runtime().await; + } + + /// Reconcile pending state under the session permit. + async fn reconcile_mcp_runtime(self: &Arc) { loop { // Compare and rebuild from current environments, not choices saved for a future turn. let environments = self.services.turn_environments.snapshot().await; @@ -280,15 +287,47 @@ impl Session { pub(crate) async fn refresh_codex_apps_tools( self: &Arc, ) -> anyhow::Result { - // Reconcile unchanged config so failed or closed clients can be replaced. - self.mark_mcp_runtime_dirty(McpRefreshReason::CodexAppsToolRefresh); - self.refresh_mcp_if_dirty().await; - let _refresh = self + let permit = self .mcp_refresh .acquire() .await - .map_err(|_| anyhow::anyhow!("MCP runtime refresh semaphore closed"))?; - self.services.mcp_runtime.refresh_codex_apps_tools().await + .context("MCP runtime refresh semaphore closed")?; + let _in_flight = self.mcp_refresh.apps_refreshes.token(); + // Reconcile once per batch so a failed or closed client can be replaced. + if self.mcp_refresh.apps_refreshes.len() == 1 { + self.mark_mcp_runtime_dirty(McpRefreshReason::CodexAppsToolRefresh); + } + if let Some(snapshot) = self.fetch_apps_catalog(permit).await? { + return Ok(snapshot); + } + let permit = self + .mcp_refresh + .acquire() + .await + .context("MCP runtime refresh semaphore closed")?; + self.fetch_apps_catalog(permit) + .await? + .context("MCP runtime changed during both Apps tool refresh attempts") + } + + async fn fetch_apps_catalog( + self: &Arc, + permit: SemaphorePermit<'_>, + ) -> anyhow::Result> { + self.reconcile_mcp_runtime().await; + let refresh = self.services.mcp_runtime.capture_apps_tool_refresh(); + // An invalidation may arrive after reconciliation but before capture. + if self.mcp_refresh.is_pending() { + return Ok(None); + } + drop(permit); + refresh.await + } + + /// Wait outside the session permit so a superseded refresh can retry. + #[tracing::instrument(name = "mcp.runtime.apps_catalog_wait", skip_all)] + pub(crate) async fn wait_for_apps_tool_refresh(&self) { + self.mcp_refresh.apps_refreshes.wait().await; } /// Reconnects the runtime so refreshed Apps tools belong to their new exact client. @@ -379,6 +418,7 @@ impl Session { pub(super) fn mark_mcp_runtime_dirty(&self, reason: McpRefreshReason) { self.mcp_refresh.invalidate(reason); + self.services.mcp_runtime.invalidate_apps_tool_refresh(); } #[tracing::instrument( @@ -457,6 +497,7 @@ impl Session { .chain(&recovered_oauth_servers) .cloned() .collect::>(); + self.wait_for_apps_tool_refresh().await; if let Some(binding) = self .services .mcp_runtime diff --git a/codex-rs/core/src/session/mcp_refresh.rs b/codex-rs/core/src/session/mcp_refresh.rs index 6e33c577be..d2a94f32d2 100644 --- a/codex-rs/core/src/session/mcp_refresh.rs +++ b/codex-rs/core/src/session/mcp_refresh.rs @@ -5,18 +5,24 @@ use std::sync::Mutex; use tokio::sync::AcquireError; use tokio::sync::Semaphore; use tokio::sync::SemaphorePermit; +use tokio_util::task::TaskTracker; /// Owns MCP invalidation and the single gate used to publish runtime updates. pub(super) struct McpRefresh { pending: Mutex>, gate: Semaphore, + pub(super) apps_refreshes: TaskTracker, } impl McpRefresh { pub(super) fn new() -> Self { + let apps_refreshes = TaskTracker::new(); + // Closed trackers still accept tokens and wake waiters when empty. + apps_refreshes.close(); Self { pending: Mutex::default(), gate: Semaphore::new(/*permits*/ 1), + apps_refreshes, } } @@ -27,7 +33,6 @@ impl McpRefresh { .insert(reason); } - #[cfg(test)] pub(super) fn is_pending(&self) -> bool { !self .pending diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index e257f19318..7b26b2712c 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -62,6 +62,11 @@ impl Session { advertised_tool: &ToolInfo, ) -> anyhow::Result { self.refresh_mcp_if_dirty().await; + // Apps refreshes fetch tools outside the MCP runtime lock. Wait for any + // in-flight refresh before capturing the catalog for this call. + if advertised_tool.server_name == CODEX_APPS_MCP_SERVER_NAME { + self.wait_for_apps_tool_refresh().await; + } self.services .mcp_runtime .prepare_call(advertised_tool) diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index a693630786..95bcb95102 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -83,14 +83,20 @@ impl SessionTask for RegularTask { == TurnStartPhase::RegularTaskStart && contributor.requires_mcp_runtime(&sess.services.thread_extension_data) }); - let preparation = sess - .emit_turn_start_lifecycle( + let preparation = async { + sess.emit_turn_start_lifecycle( &ctx, /*token_usage_at_turn_start*/ None, TurnStartPhase::RegularTaskStart, ) - .or_cancel(&cancellation_token) .await; + if prepares_mcp { + // Resource discovery can overlap the fetch. Wait before reprojection. + sess.wait_for_apps_tool_refresh().await; + } + } + .or_cancel(&cancellation_token) + .await; // Even cancelled discovery may have cleared the previous account's catalog. if prepares_mcp { sess.request_mcp_runtime_reprojection( diff --git a/codex-rs/core/tests/suite/codex_apps_protocol.rs b/codex-rs/core/tests/suite/codex_apps_protocol.rs index 561cc9b8aa..687cd25c9b 100644 --- a/codex-rs/core/tests/suite/codex_apps_protocol.rs +++ b/codex-rs/core/tests/suite/codex_apps_protocol.rs @@ -3,12 +3,14 @@ use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; +use std::time::Duration; use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; use codex_features::Feature; use codex_login::CodexAuth; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_mcp::McpSnapshotDetail; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; @@ -28,10 +30,12 @@ use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use core_test_support::wait_for_mcp_server; +use core_test_support::wait_for_successful_turn; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; use test_case::test_case; +use tokio::sync::Notify; #[derive(Clone, Copy)] enum HostedProtocolSetting { @@ -197,3 +201,116 @@ async fn apps_call_survives_catalog_restoration_while_awaiting_approval() -> any peer.shutdown_and_wait().await?; Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn next_model_step_uses_the_in_flight_apps_refresh() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let model = responses::start_mock_server().await; + let apps_server = responses::start_mock_server().await; + let apps = AppsTestServer::mount(&apps_server).await?; + let fixture = apps_enabled_builder(apps.chatgpt_base_url) + .with_config(|config| config.update_plan_enabled = true) + .with_model_info_override("gpt-5.5", |model| model.supports_search_tool = false) + .build_with_auto_env(&model) + .await?; + wait_for_mcp_server(&fixture.codex, CODEX_APPS_MCP_SERVER_NAME).await?; + + let first_requested = Arc::new(Notify::new()); + let requested = Arc::clone(&first_requested); + let (release_first, first_gate) = std::sync::mpsc::channel(); + let first_gate = std::sync::Mutex::new(Some(first_gate)); + let first = responses::sse_response(responses::sse(vec![ + responses::ev_response_created("first"), + responses::ev_function_call("plan", "update_plan", r#"{"plan":[]}"#), + responses::ev_completed("first"), + ])); + let response = responses::mount_response_handler(&model, move |_| { + if let Some(gate) = first_gate.lock().unwrap().take() { + requested.notify_one(); + let _ = gate.recv(); + first.clone() + } else { + responses::sse_response(responses::sse_completed("second")) + } + }) + .await; + + // Wait for the first model request to capture the original tools before overriding tools/list. + fixture + .codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Plan the next steps.".into(), + text_elements: Vec::new(), + }])) + .await?; + tokio::time::timeout(Duration::from_secs(10), first_requested.notified()).await?; + let listed = Arc::new(Notify::new()); + let listing = Arc::clone(&listed); + let (release_list, list_gate) = std::sync::mpsc::channel(); + let list_gate = std::sync::Mutex::new(Some(list_gate)); + wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::body_partial_json( + json!({"method":"tools/list"}), + )) + .respond_with(move |request: &wiremock::Request| { + let body: Value = serde_json::from_slice(&request.body).unwrap(); + if let Some(gate) = list_gate.lock().unwrap().take() { + listing.notify_one(); + let _ = gate.recv(); + } + wiremock::ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc":"2.0", "id":body["id"], "result":{"tools":[]} + })) + }) + .with_priority(1) + .mount(&apps_server) + .await; + + let refresh = fixture.codex.refresh_codex_apps_tools(); + tokio::pin!(refresh); + tokio::select! { + _ = listed.notified() => {} + result = &mut refresh => { + result?; + panic!("refresh ended before tools/list"); + } + } + release_first.send(())?; + wait_for_event(&fixture.codex, |event| { + matches!(event, EventMsg::PlanUpdate(_)) + }) + .await; + tokio::time::timeout( + Duration::from_secs(10), + fixture.codex.mcp_server_status_snapshot( + CODEX_APPS_MCP_SERVER_NAME, + McpSnapshotDetail::ToolsAndAuthOnly, + /*expected_registration*/ None, + ), + ) + .await??; + assert_eq!( + response.requests().len(), + 1, + "the next model step must wait for Apps" + ); + release_list.send(())?; + refresh.await?; + wait_for_successful_turn(&fixture.codex).await; + + let calendar_advertised = response + .requests() + .iter() + .map(|request| { + responses::namespace_child_tool( + &request.body_json(), + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_CREATE_TOOL, + ) + .is_some() + }) + .collect::>(); + assert_eq!(calendar_advertised, [true, false]); + Ok(()) +}