diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index c12d23e8da..b44e549507 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -21,6 +21,7 @@ use startup::emit_update; use startup::mcp_init_error_display; use startup::mcp_startup_failure_reason; use startup::should_share_codex_apps_tools_cache; +pub(crate) use tool_catalog::BindingCatalogRevision; pub use tool_catalog::tool_is_model_visible; use std::collections::HashMap; 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 e6148cbb4d..2e08cf004d 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -58,10 +58,19 @@ pub fn tool_is_model_visible(tool: &ToolInfo) -> bool { .any(|target| target.as_str() == Some(MCP_UI_MODEL_VISIBILITY)) } +/// Catalog identity within one published connection set, including dormant servers. +#[derive(PartialEq, Eq)] +pub(crate) enum BindingCatalogRevision { + Ready(ClientToolCatalogRevision), + Dormant(u64), +} + impl McpConnectionSet { pub(crate) async fn stable_catalog_revisions( &self, - ) -> Option> { + required_servers: &[String], + required_plugins: &HashSet, + ) -> Option> { let mut revisions = HashMap::new(); for (server_name, view) in &self.servers { if !view @@ -70,7 +79,31 @@ impl McpConnectionSet { .startup_complete .load(Ordering::Acquire) { - return None; + // A cached catalog can remain stable without starting its server. + // Explicit requirements must still start the server during binding capture. + if !view.connection.startup_is_dormant() + || view.connection.client.cancel_token.is_cancelled() + || required_servers + .iter() + .any(|required| required == server_name) + || (self.is_selected_plugin_mcp_server(server_name) + && self + .plugin_id_for_mcp_server_name(server_name) + .is_some_and(|plugin_id| required_plugins.contains(plugin_id))) + { + return None; + } + let revision = view + .connection + .client + .tool_catalog_cache_context + .as_ref()? + .current_revision()?; + revisions.insert( + server_name.clone(), + BindingCatalogRevision::Dormant(revision), + ); + continue; } let Some(client) = view.connection.client.ready_client() else { if !view.connection.client.is_codex_apps_mcp_server @@ -87,10 +120,10 @@ impl McpConnectionSet { let revision = client.tool_catalog.read(|catalog| catalog.revision).await; revisions.insert( server_name.clone(), - ClientToolCatalogRevision { + BindingCatalogRevision::Ready(ClientToolCatalogRevision { catalog: Arc::clone(&client.tool_catalog), revision, - }, + }), ); } Some(revisions) @@ -168,6 +201,7 @@ impl McpConnectionSet { (tools, errors) } + #[instrument(level = "trace", skip_all)] pub(crate) async fn capture_binding_with_metadata( self: &Arc, config: Arc, diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 9d3f3d5837..f6bc828115 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -2731,7 +2731,15 @@ async fn capture_binding_skips_pending_optional_servers_after_configured_shared_ required_manager.required_servers = vec!["pending-selected".to_string()]; let manager = Arc::new(manager); - assert!(manager.stable_catalog_revisions().await.is_none()); + assert!( + manager + .stable_catalog_revisions( + /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new() + ) + .await + .is_none() + ); let started = tokio::time::Instant::now(); let binding = tokio::time::timeout( Duration::from_millis(500), @@ -2887,9 +2895,25 @@ async fn stable_catalog_revisions_ignore_terminal_optional_server_failures() { assert!(failed.client().await.is_err()); manager.insert_test_client("failed", failed); - assert!(manager.stable_catalog_revisions().await.is_some()); + assert!( + manager + .stable_catalog_revisions( + /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new() + ) + .await + .is_some() + ); manager.required_servers.push("failed".to_string()); - assert!(manager.stable_catalog_revisions().await.is_none()); + assert!( + manager + .stable_catalog_revisions( + /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new() + ) + .await + .is_none() + ); manager.required_servers.clear(); let binding = capture_binding(&Arc::new(manager)).await; @@ -3840,7 +3864,15 @@ async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client() }; let manager = Arc::new(manager); - assert!(manager.stable_catalog_revisions().await.is_none()); + assert!( + manager + .stable_catalog_revisions( + /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new() + ) + .await + .is_none() + ); let reconnect_finished_wait = reconnect_finished.notified(); let tools = manager.list_all_tools().await; assert!(tools.is_empty()); @@ -3855,7 +3887,15 @@ async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client() vec!["drive_search"] ); assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); - assert!(manager.stable_catalog_revisions().await.is_some()); + assert!( + manager + .stable_catalog_revisions( + /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new() + ) + .await + .is_some() + ); let step = capture_binding(&manager).await; let prepared = step diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 45f2e97237..2ad549d5e7 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -46,8 +46,8 @@ use tokio_util::sync::CancellationToken; use crate::McpConfig; use crate::binding::McpBinding; -use crate::client_tool_catalog::ClientToolCatalogRevision; use crate::client_tool_catalog::CodexAppsToolSnapshot; +use crate::connection_manager::BindingCatalogRevision; use crate::connection_manager::McpConnectionSet; use crate::elicitation::ElicitationLifecycle; use crate::elicitation::ElicitationRequestRouter; @@ -119,7 +119,7 @@ struct PublishedMcpRuntime { } struct CachedMcpBinding { - catalog_revisions: HashMap, + catalog_revisions: HashMap, binding: Arc, } @@ -363,7 +363,10 @@ impl McpRuntime { required_plugins: &HashSet, ) -> Option> { let config = Arc::clone(current.config.as_ref()?); - let stable_catalog_revisions = current.connections.stable_catalog_revisions().await; + let stable_catalog_revisions = current + .connections + .stable_catalog_revisions(required_servers, required_plugins) + .await; if let Some(catalog_revisions) = &stable_catalog_revisions { let cached = current .cached_binding @@ -390,7 +393,7 @@ impl McpRuntime { if let Some(catalog_revisions) = stable_catalog_revisions && current .connections - .stable_catalog_revisions() + .stable_catalog_revisions(required_servers, required_plugins) .await .as_ref() == Some(&catalog_revisions) diff --git a/codex-rs/codex-mcp/src/tool_catalog_cache.rs b/codex-rs/codex-mcp/src/tool_catalog_cache.rs index c0165f4ae9..72d723b981 100644 --- a/codex-rs/codex-mcp/src/tool_catalog_cache.rs +++ b/codex-rs/codex-mcp/src/tool_catalog_cache.rs @@ -116,7 +116,18 @@ impl Default for ToolCatalogCacheEntry { impl McpToolCatalogCacheContext { pub(crate) fn has_tools(&self) -> bool { - self.current_tools().is_some_and(|tools| !tools.is_empty()) + self.current_revision().is_some() + } + + /// Identifies the current usable catalog without cloning its tool definitions. + /// The entry is fixed for a connection; accepted publications advance its revision. + pub(crate) fn current_revision(&self) -> Option { + let state = lock_unpoisoned(&self.entry.state); + let snapshot = state.snapshot.as_ref()?; + (!state.disabled_by_server + && !snapshot.tools.is_empty() + && snapshot.published_at.elapsed() <= TOOL_CATALOG_CACHE_TTL) + .then_some(state.last_accepted_generation) } pub(crate) fn optional_startup_deadline( diff --git a/codex-rs/core/tests/suite/mcp_tool_cache.rs b/codex-rs/core/tests/suite/mcp_tool_cache.rs index 9602d53248..d4549c4cf3 100644 --- a/codex-rs/core/tests/suite/mcp_tool_cache.rs +++ b/codex-rs/core/tests/suite/mcp_tool_cache.rs @@ -1,5 +1,6 @@ use std::sync::Arc; use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; @@ -49,6 +50,13 @@ use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; use test_case::test_case; +use tracing::Subscriber; +use tracing::span::Attributes; +use tracing::span::Id; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::Context as LayerContext; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::util::SubscriberInitExt; use super::rmcp_client::remote_aware_environment_id; use super::rmcp_client::remote_aware_stdio_server_bin; @@ -56,6 +64,18 @@ use super::rmcp_client::remote_aware_stdio_server_bin; const SERVER_NAME: &str = "cached_rmcp"; const NAMESPACE: &str = "mcp__cached_rmcp"; +struct McpBindingCaptureCounter(Arc); + +impl Layer for McpBindingCaptureCounter { + fn on_new_span(&self, attributes: &Attributes<'_>, _id: &Id, _context: LayerContext<'_, S>) { + if attributes.metadata().target() == "codex_mcp::connection_manager::tool_catalog" + && attributes.metadata().name() == "capture_binding_with_metadata" + { + self.0.fetch_add(1, Ordering::SeqCst); + } + } +} + fn user_turn(prompt: &str) -> TurnInputRequest { TurnInputRequest::user_input(vec![UserInput::Text { text: prompt.to_string(), @@ -558,7 +578,8 @@ async fn cached_http_mcp_starts_lazily_for_subagents( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +// Keep spawned tasks on the thread with the scoped tracing subscriber. +#[tokio::test(flavor = "current_thread")] async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow::Result<()> { skip_if_wine_exec!( Ok(()), @@ -566,6 +587,11 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow ); skip_if_no_network!(Ok(())); + let binding_captures = Arc::new(AtomicUsize::new(0)); + let _tracing = tracing_subscriber::registry() + .with(McpBindingCaptureCounter(Arc::clone(&binding_captures))) + .set_default(); + let responses_server = responses::start_mock_server().await; let command = remote_aware_stdio_server_bin()?; let environment_id = remote_aware_environment_id(); @@ -691,15 +717,29 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow .await?; second_thread.submit(Op::Interrupt).await?; + binding_captures.store(0, Ordering::SeqCst); let unused_response = mount_sse_once( &responses_server, responses::sse(vec![ responses::ev_response_created("unused"), - responses::ev_assistant_message("unused-message", "done"), + responses::ev_function_call( + "unused-plan", + "update_plan", + r#"{"plan":[{"step":"Work without MCP tools","status":"in_progress"}]}"#, + ), responses::ev_completed("unused"), ]), ) .await; + let unused_done_response = mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("unused-done"), + responses::ev_assistant_message("unused-message", "done"), + responses::ev_completed("unused-done"), + ]), + ) + .await; second_thread .start_or_steer_turn(user_turn("Do not call any MCP tools.")) .await?; @@ -718,10 +758,17 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow !reported_ready_before_startup, "a dormant MCP server must not be reported as ready" ); - assert_definition( - &unused_response, - &format!("Use the tools from {cached_process}."), - &format!("Echo from {cached_process}."), + for response in [&unused_response, &unused_done_response] { + assert_definition( + response, + &format!("Use the tools from {cached_process}."), + &format!("Echo from {cached_process}."), + ); + } + assert_eq!( + binding_captures.load(Ordering::SeqCst), + 1, + "both model steps must share one binding while the cached server stays dormant" ); let (mcp_config, _) = second_thread.current_mcp_config_and_runtime_context().await; assert_eq!( @@ -847,6 +894,10 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow .await?; let expected_error = format!("MCP tool `{SERVER_NAME}/cwd` is not available to the model"); assert_eq!(cached_turn.await??, second_process); + assert!( + binding_captures.load(Ordering::SeqCst) > 1, + "starting the cached server must invalidate the dormant binding" + ); assert_definition( &cached_done_response, &format!("Use the tools from {second_process}."),