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 84716448f4..36ad04e173 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -248,7 +248,7 @@ impl McpConnectionSet { .startup_complete .load(Ordering::Acquire) { - (None, cached_tools.or_else(|| view.connection.client.cached_tools())?) + (None, view.connection.client.cached_tools_or(cached_tools)?) } else { view.connection.client.reconnect_failed_startup().await; let Ok(mut client) = view.connection.client().await else { diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 7bf924b7f9..ee20c57598 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -3017,6 +3017,62 @@ async fn capture_binding_retains_cached_tools_that_expire_while_waiting_for_anot assert!(manager.servers["server_a"].connection.startup_is_dormant()); } +#[tokio::test] +async fn capture_binding_omits_cache_disabled_while_waiting_for_another_server() { + let runtime_context = McpRuntimeContext::new( + Arc::new(environment_manager_without_environments()), + std::env::temp_dir(), + ); + let server_config: McpServerConfig = + serde_json::from_value(serde_json::json!({ "command": "cached-mcp" })) + .expect("server configuration"); + let cache_context = McpToolCatalogCache::default() + .context( + "cached", + &server_config, + &runtime_context, + /*resolved_environment*/ None, + ( + &ElicitationCapability::default(), + &ClientMcpExtensions::default(), + ), + /*connection_identity*/ None, + ) + .expect("shared catalog"); + cache_context.publish_if_newest( + cache_context.begin_fetch(), + &[create_test_tool("cached", "cached_tool")], + ); + let (mut cached, _started, _release) = create_gated_async_managed_client( + create_test_managed_client(vec![create_test_tool("cached", "live_tool")]).await, + ); + cached.tool_catalog_cache_context = Some(cache_context.clone()); + let (waiting, started, release) = create_gated_async_managed_client( + create_test_managed_client(vec![create_test_tool("waiting", "ready_tool")]).await, + ); + let mut manager = McpConnectionSet::new_uninitialized( + &Constrained::allow_any(AskForApproval::OnRequest), + &Constrained::allow_any(PermissionProfile::default()), + /*prefix_mcp_tool_names*/ true, + ); + manager.insert_test_client("cached", cached); + manager.insert_test_client("waiting", waiting); + manager.required_servers = vec!["waiting".to_string()]; + let manager = Arc::new(manager); + + let binding = capture_binding(&manager); + tokio::pin!(binding); + assert!(futures::poll!(&mut binding).is_pending()); + started.await.expect("required server startup began"); + cache_context.disable(); + release.send(()).expect("release required server startup"); + + assert_eq!( + model_tool_names(binding.await.tools()), + HashSet::from([ToolName::namespaced("mcp__waiting", "ready_tool")]), + ); +} + #[tokio::test] async fn capture_binding_resolves_concurrently_and_rechecks_cached_clients() { let codex_home = tempdir().expect("tempdir"); diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index a4c1cd4380..563d6bedf1 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -553,13 +553,17 @@ impl AsyncManagedClient { } pub(crate) fn cached_tools(&self) -> Option> { + self.cached_tools_or(/*fallback*/ None) + } + + pub(crate) fn cached_tools_or(&self, fallback: Option>) -> Option> { self.codex_apps_tools_cache_context .as_ref() .and_then(ConnectorRuntimeContext::current_tools) .or_else(|| { self.tool_catalog_cache_context .as_ref() - .and_then(McpToolCatalogCacheContext::current_tools) + .and_then(|cache| cache.current_tools_or(fallback)) }) } diff --git a/codex-rs/codex-mcp/src/tool_catalog_cache.rs b/codex-rs/codex-mcp/src/tool_catalog_cache.rs index 77fb60e804..c0165f4ae9 100644 --- a/codex-rs/codex-mcp/src/tool_catalog_cache.rs +++ b/codex-rs/codex-mcp/src/tool_catalog_cache.rs @@ -150,11 +150,24 @@ impl McpToolCatalogCacheContext { } pub(crate) fn current_tools(&self) -> Option> { - lock_unpoisoned(&self.entry.state) + self.current_tools_or(/*fallback*/ None) + } + + /// Prefers the current catalog, retaining a capture's fallback across expiry but not opt-out. + pub(crate) fn current_tools_or( + &self, + fallback: Option>, + ) -> Option> { + let state = lock_unpoisoned(&self.entry.state); + if state.disabled_by_server { + return None; + } + state .snapshot .as_ref() .filter(|snapshot| snapshot.published_at.elapsed() <= TOOL_CATALOG_CACHE_TTL) .map(|snapshot| snapshot.tools.clone()) + .or(fallback) } pub(crate) fn begin_fetch(&self) -> McpToolCatalogFetchTicket { diff --git a/codex-rs/core/tests/suite/mcp_tool_cache.rs b/codex-rs/core/tests/suite/mcp_tool_cache.rs index e727f97e84..dae587c23d 100644 --- a/codex-rs/core/tests/suite/mcp_tool_cache.rs +++ b/codex-rs/core/tests/suite/mcp_tool_cache.rs @@ -1,4 +1,6 @@ use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; use anyhow::Context; @@ -8,8 +10,15 @@ use codex_config::types::McpServerTransportConfig; use codex_core::NewThread; use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; +use codex_core::config::Config; use codex_exec_server::ExecutorFileSystem; use codex_exec_server::RemoveOptions; +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_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_protocol::mcp::McpServerConnectionStatus; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; @@ -23,6 +32,9 @@ use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::user_input::UserInput; use codex_utils_path_uri::PathUri; use core_test_support::apps_test_server::AppsTestServer; +use core_test_support::apps_test_server::SEARCH_CALENDAR_CREATE_TOOL; +use core_test_support::apps_test_server::SEARCH_CALENDAR_NAMESPACE; +use core_test_support::apps_test_server::apps_enabled_builder; use core_test_support::is_remote_test_environment; use core_test_support::responses; use core_test_support::responses::ResponseMock; @@ -250,6 +262,164 @@ async fn mcp_calls_stay_bound_to_each_thread() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn apps_cache_filled_during_binding_capture_reaches_the_model() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + struct ChildAppsEndpoint(String); + + impl McpServerContributor for ChildAppsEndpoint { + fn id(&self) -> &'static str { + "child_apps_cache_test" + } + + fn contribute<'a>( + &'a self, + context: McpServerContributionContext<'a, Config>, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + if !matches!(context.session_source(), Some(SessionSource::SubAgent(_))) { + return Vec::new(); + } + vec![McpServerContribution::HostedApps { + config: Box::new( + serde_json::from_value(json!({ "url": self.0 })) + .expect("child Apps MCP config"), + ), + }] + }) + } + } + + tokio::time::timeout(Duration::from_secs(30), async { + let server = responses::start_mock_server().await; + let tools_available = Arc::new(AtomicBool::new(false)); + let apps = + AppsTestServer::mount_with_tools_available_when(&server, Arc::clone(&tools_available)) + .await?; + let waiting_server = responses::start_mock_server().await; + let (waiting, waiting_startup) = + AppsTestServer::mount_with_startup_control(&waiting_server).await?; + let pending_server = responses::start_mock_server().await; + let (pending_apps, pending_startup) = + AppsTestServer::mount_with_startup_control(&pending_server).await?; + // Gate only the child's MCP endpoint, leaving backend directory requests unblocked. + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.mcp_server_contributor(Arc::new(ChildAppsEndpoint(format!( + "{}/api/codex/ps/mcp", + pending_apps.chatgpt_base_url + )))); + let test = apps_enabled_builder(apps.chatgpt_base_url) + .with_model_info_override("gpt-5.5", |model| model.supports_search_tool = false) + .with_extensions(Arc::new(extensions.build())) + .with_config(move |config| { + // Keep the gated HTTP request on the app host. Wine's executor serializes RPCs, + // so blocking it here would also block the peer startup that releases this gate. + config + .mcp_servers + .set(std::collections::HashMap::from([( + SERVER_NAME.to_string(), + serde_json::from_value(json!({ + "url": format!("{}/api/codex/ps/mcp", waiting.chatgpt_base_url), + "http_headers": { "Authorization": "Bearer cache-test-token" }, + "enabled_tools": ["calendar_list_events"], + })) + .expect("cacheable MCP config"), + )])) + .expect("test MCP config"); + }) + .build_with_auto_env(&server) + .await?; + // Startup emits one summary for both servers, not one event per server. + let EventMsg::McpStartupComplete(startup) = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::McpStartupComplete(_)) + }) + .await + else { + unreachable!("event predicate guarantees the startup summary"); + }; + assert_eq!( + startup + .ready + .into_iter() + .collect::>(), + std::collections::BTreeSet::from([ + CODEX_APPS_MCP_SERVER_NAME.to_string(), + SERVER_NAME.to_string(), + ]) + ); + + let release_apps = pending_startup.hold_next_successful_initialize(); + let release_waiting = waiting_startup.hold_next_successful_initialize(); + let NewThread { thread: child, .. } = test + .thread_manager + .start_thread(StartThreadOptions { + session_source: Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: test.session_configured.thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + })), + ..StartThreadOptions::new(test.config.clone()) + }) + .await?; + assert_eq!(waiting_startup.initialize_attempts(), 1); + let response = mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("cache-refresh"), + responses::ev_assistant_message("done", "done"), + responses::ev_completed("cache-refresh"), + ]), + ) + .await; + child + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Mention { + name: SERVER_NAME.to_string(), + path: format!("mcp://{SERVER_NAME}"), + }])) + .await?; + wait_for_event(&child, |event| { + matches!(event, EventMsg::McpStartupUpdate(update) + if update.server == SERVER_NAME && matches!(update.status, McpStartupStatus::Starting)) + }) + .await; + + // A peer fills the initially empty shared Apps cache while capture waits for the named server. + tools_available.store(true, Ordering::SeqCst); + let mut peer_config = test.config.clone(); + peer_config + .mcp_servers + .set(std::collections::HashMap::new())?; + let NewThread { thread: peer, .. } = test + .thread_manager + .start_thread(StartThreadOptions::new(peer_config)) + .await?; + wait_for_mcp_server(&peer, CODEX_APPS_MCP_SERVER_NAME).await?; + release_waiting.send(())?; + wait_for_event(&child, |event| matches!(event, EventMsg::TurnComplete(_))).await; + // The child's own Apps client must still be pending when its request reaches inference. + release_apps.send(())?; + let body = response.single_request().body_json(); + assert!( + responses::namespace_child_tool( + &body, + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_CREATE_TOOL, + ) + .is_some(), + "the first model request must include the peer's Apps tools: {body}" + ); + + child.shutdown_and_wait().await?; + peer.shutdown_and_wait().await?; + Ok(()) + }) + .await + .context("timed out exercising an Apps cache fill during binding capture")? +} + #[test_case(false, false, 1; "optional server uses cache")] #[test_case(true, false, 1; "required server uses cache")] #[test_case(false, true, 2; "headers helper bypasses cache")]