diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs index f66e08858b..fc48427399 100644 --- a/codex-rs/codex-mcp/src/binding.rs +++ b/codex-rs/codex-mcp/src/binding.rs @@ -1,4 +1,4 @@ -//! Immutable MCP state bound to one model sampling request. +//! Immutable MCP catalog and execution handles. use std::collections::HashMap; use std::fmt; @@ -26,7 +26,7 @@ use crate::rmcp_client::ManagedClient; use crate::server::McpServerMetadata; use crate::tools::ToolInfo; -/// The exact tool catalog and execution handles for one model sampling request. +/// The exact tool catalog and execution handles shared by compatible sampling steps. pub struct McpBinding { connections: Arc, clients: Arc, @@ -75,7 +75,7 @@ impl McpBinding { self.plugins_available } - /// Returns the frozen catalog advertised for this sampling request. + /// Returns the frozen catalog captured for this binding. pub fn tools(&self) -> &[ToolInfo] { &self.tools } diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 173e94053a..b2c8a01745 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -598,7 +598,7 @@ impl McpConnectionSet { let Some(view) = self.servers.get(server_name) else { return false; }; - view.connection.client().await.is_ok() + view.connection.client.ready_transport().is_some() || view.connection.client().await.is_ok() } /// Stop all MCP clients owned by this manager and terminate stdio server processes. 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 b5c27c903e..341b5d7100 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -58,6 +58,32 @@ pub fn tool_is_model_visible(tool: &ToolInfo) -> bool { } impl McpConnectionSet { + pub(crate) async fn stable_catalog_revision(&self) -> Option { + for (server_name, view) in &self.servers { + if !view + .connection + .client + .startup_complete + .load(Ordering::Acquire) + { + return None; + } + let Some(client) = view.connection.client.ready_transport() else { + if !view.connection.client.is_codex_apps_mcp_server + && self.required_servers.binary_search(server_name).is_err() + && matches!(view.connection.client.client.peek(), Some(Err(_))) + { + continue; + } + return None; + }; + if client.is_closed().await { + return None; + } + } + Some(*self.tool_catalog_revision.read().await) + } + /// Returns all tools with model-visible names normalized. #[instrument(level = "trace", skip_all, fields(mcp_server_count = self.servers.len()))] pub async fn list_all_tools(&self) -> Vec { diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index b31dd8db4b..2197468b60 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1911,6 +1911,7 @@ async fn capture_binding_skips_pending_optional_servers_after_one_shared_startup } let manager = Arc::new(manager); + assert_eq!(manager.stable_catalog_revision().await, None); let binding = tokio::time::timeout(Duration::from_millis(1500), capture_binding(&manager)) .await .expect("all optional servers should share a single startup grace"); @@ -1953,6 +1954,39 @@ async fn capture_binding_skips_pending_optional_servers_after_one_shared_startup assert!(binding.is_err(), "explicitly requested servers must wait"); } +#[tokio::test] +async fn stable_catalog_revision_ignores_terminal_optional_server_failures() { + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionSet::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + let ready = create_ready_async_managed_client(vec![create_test_tool("ready", "echo")]).await; + assert!(ready.client().await.is_ok()); + let mut failed = ready.clone(); + manager.insert_test_client("ready", ready); + failed.client = futures::future::ready::>(Err( + StartupOutcomeError::Failed { + error: "optional startup failed".to_string(), + is_authentication_required: false, + }, + )) + .boxed() + .shared(); + assert!(failed.client().await.is_err()); + manager.insert_test_client("failed", failed); + + assert_eq!(manager.stable_catalog_revision().await, Some(0)); + manager.required_servers.push("failed".to_string()); + assert_eq!(manager.stable_catalog_revision().await, None); + manager.required_servers.clear(); + + let binding = capture_binding(&Arc::new(manager)).await; + assert!(binding.prepare_call("ready", "echo").is_some()); +} + #[tokio::test(start_paused = true)] async fn capture_binding_shares_optional_startup_grace_across_connection_sets() { let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); @@ -2519,6 +2553,7 @@ async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client() }; let manager = Arc::new(manager); + assert_eq!(manager.stable_catalog_revision().await, None); let reconnect_finished_wait = reconnect_finished.notified(); let tools = manager.list_all_tools().await; assert!(tools.is_empty()); @@ -2533,6 +2568,7 @@ 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_eq!(manager.stable_catalog_revision().await, Some(0)); let step = capture_binding(&manager).await; let prepared = step diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index d5e25977c3..1c326bffc7 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -492,6 +492,24 @@ impl AsyncManagedClient { self.client.clone().await } + pub(crate) fn ready_transport(&self) -> Option> { + let recovered = self.startup_reconnect.as_ref().and_then(|reconnect| { + reconnect + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .current_client + .as_ref() + .map(|client| Arc::clone(&client.client)) + }); + recovered.or_else(|| { + self.client + .peek() + .and_then(|result| result.as_ref().ok()) + .map(|client| Arc::clone(&client.client)) + }) + } + pub(crate) async fn reconnect_failed_startup(&self) { let Some(startup_reconnect) = self.startup_reconnect.as_ref() else { return; diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 33debec79d..e9feec69ec 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -8,6 +8,7 @@ use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; +use std::sync::Mutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; @@ -84,6 +85,12 @@ struct PublishedMcpRuntime { auth_token: Option, plugins_available: bool, ready_selected_capability_roots: Vec, + cached_binding: Mutex>, +} + +struct CachedMcpBinding { + catalog_revision: u64, + binding: Arc, } struct McpReconnectGuard<'a> { @@ -148,6 +155,7 @@ impl McpRuntime { auth_token: None, plugins_available: false, ready_selected_capability_roots: Vec::new(), + cached_binding: Mutex::new(None), }), reconnect_pending: AtomicBool::new(false), elicitation_router: ElicitationRequestRouter::default(), @@ -204,6 +212,7 @@ impl McpRuntime { auth_token, plugins_available, ready_selected_capability_roots, + cached_binding: Mutex::new(None), })); let _ = publish.send(true); } @@ -223,14 +232,51 @@ impl McpRuntime { &self, required_servers: &[String], ) -> Option> { - let current = self.current.load_full(); + Self::binding_from_published_runtime(self.current.load_full(), required_servers).await + } + + async fn binding_from_published_runtime( + current: Arc, + required_servers: &[String], + ) -> Option> { let config = Arc::clone(current.config.as_ref()?); - Some(Arc::new( + let stable_catalog_revision = current.connections.stable_catalog_revision().await; + if let Some(catalog_revision) = stable_catalog_revision { + let cached = current + .cached_binding + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(cached) = cached.as_ref() + && cached.catalog_revision == catalog_revision + { + return Some(Arc::clone(&cached.binding)); + } + } + + let binding = Arc::new( current .connections .capture_binding_with_metadata(config, current.plugins_available, required_servers) .await, - )) + ); + if let Some(catalog_revision) = stable_catalog_revision + && current.connections.stable_catalog_revision().await == Some(catalog_revision) + { + let mut cached = current + .cached_binding + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(cached) = cached.as_ref() + && cached.catalog_revision == catalog_revision + { + return Some(Arc::clone(&cached.binding)); + } + *cached = Some(CachedMcpBinding { + catalog_revision, + binding: Arc::clone(&binding), + }); + } + Some(binding) } /// Returns whether the published snapshot still belongs to the current credentials. @@ -261,16 +307,11 @@ impl McpRuntime { /// Captures the current runtime after its selected server has finished startup. pub async fn current_binding_for_call(&self, server: &str) -> Option> { let current = self.current.load_full(); - let config = Arc::clone(current.config.as_ref()?); + current.config.as_ref()?; if !current.connections.wait_for_server_startup(server).await { return None; } - Some(Arc::new( - current - .connections - .capture_binding_with_metadata(config, current.plugins_available, &[]) - .await, - )) + Self::binding_from_published_runtime(current, /*required_servers*/ &[]).await } /// Returns the latest published configuration without waiting for clients. @@ -513,6 +554,45 @@ mod tests { assert!(!gate.wait().await); } + #[tokio::test] + async fn cached_bindings_are_scoped_to_the_published_runtime() { + let published = Arc::new(PublishedMcpRuntime { + connections: Arc::new(McpConnectionSet::empty(/*prefix_mcp_tool_names*/ true)), + config: Some(Arc::new(crate::mcp::tests::test_mcp_config( + std::env::temp_dir(), + ))), + auth: None, + auth_token: None, + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + cached_binding: Mutex::new(None), + }); + let first = McpRuntime::binding_from_published_runtime( + Arc::clone(&published), + /*required_servers*/ &[], + ) + .await + .expect("first binding"); + let repeated = McpRuntime::binding_from_published_runtime( + Arc::clone(&published), + /*required_servers*/ &[], + ) + .await + .expect("repeated binding"); + assert!(Arc::ptr_eq(&first, &repeated)); + + let previous = Arc::into_inner(published).expect("published runtime has no other owners"); + let republished = Arc::new(PublishedMcpRuntime { + cached_binding: Mutex::new(None), + ..previous + }); + let refreshed = + McpRuntime::binding_from_published_runtime(republished, /*required_servers*/ &[]) + .await + .expect("republished binding"); + assert!(!Arc::ptr_eq(&first, &refreshed)); + } + fn http_server(environment_id: &str) -> McpServerConfig { McpServerConfig { auth: Default::default(),