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 c09c1b9090..84716448f4 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -175,7 +175,7 @@ impl McpConnectionSet { let mut listed_tools = Vec::new(); let mut clients = std::collections::HashMap::new(); let optional_mcp_startup_grace = config.optional_mcp_startup_grace; - join_all(self.servers.iter().map(|(server_name, view)| async move { + let server_snapshots = join_all(self.servers.iter().map(|(server_name, view)| async move { if !view .connection .client @@ -183,7 +183,11 @@ impl McpConnectionSet { .load(Ordering::Acquire) { let required = self.required_servers.binary_search(server_name).is_ok(); - let has_cached_tools = view.connection.client.has_cached_tools(); + // Keep the catalog that lets us skip startup even if it expires during the wait. + let cached_tools = view.connection.client.cached_tools().filter(|tools| { + view.connection.client.is_codex_apps_mcp_server || !tools.is_empty() + }); + let has_cached_tools = cached_tools.is_some(); let must_wait_for_startup = (required && (!view.connection.startup_is_dormant() || !has_cached_tools)) || self.is_selected_plugin_mcp_server(server_name) @@ -192,7 +196,7 @@ impl McpConnectionSet { .any(|required| required == server_name) || (server_name == CODEX_APPS_MCP_SERVER_NAME && !has_cached_tools); if !must_wait_for_startup && has_cached_tools { - return; + return (server_name, view, cached_tools); } if !must_wait_for_startup && optional_mcp_startup_grace.is_zero() { if let Some(cache) = view.connection.client.tool_catalog_cache_context.as_ref() @@ -229,46 +233,37 @@ impl McpConnectionSet { { trace!(server_name = %server_name, "omitting pending optional MCP server"); } - return; + return (server_name, view, cached_tools); } let _ = view.connection.client().await; + return (server_name, view, cached_tools); } + (server_name, view, None) })) .await; - let server_results = join_all(self.servers.iter().map(|(server_name, view)| async move { - if !view + let server_results = join_all(server_snapshots.into_iter().map(|(server_name, view, cached_tools)| async move { + let (client, server_tools) = if !view .connection .client .startup_complete .load(Ordering::Acquire) { - if !view.connection.client.has_cached_tools() { - return None; - } - let server_tools = view.listed_tools(&self.tool_plugin_provenance).await?; - let server_tools = server_tools - .into_iter() - .map(|mut tool| { - if let Some(annotations) = tool.tool.annotations.as_mut() { - annotations.read_only_hint = None; - } - Self::with_server_metadata(tool, &view.metadata) - }) - .collect::>(); - return Some((server_name.clone(), None, server_tools)); - } - view.connection.client.reconnect_failed_startup().await; - let Ok(mut client) = view.connection.client().await else { - trace!(server_name = %server_name, "omitting MCP server without an exact ready client"); - return None; - }; - client.tool_timeout = view.tool_timeout; - let catalog_override = if server_name == CODEX_APPS_MCP_SERVER_NAME { - self.codex_apps_tools_override.read().await.clone() + (None, cached_tools.or_else(|| view.connection.client.cached_tools())?) } else { - None + view.connection.client.reconnect_failed_startup().await; + let Ok(mut client) = view.connection.client().await else { + trace!(server_name = %server_name, "omitting MCP server without an exact ready client"); + return None; + }; + client.tool_timeout = view.tool_timeout; + let catalog_override = if server_name == CODEX_APPS_MCP_SERVER_NAME { + self.codex_apps_tools_override.read().await.clone() + } else { + None + }; + let server_tools = catalog_override.unwrap_or_else(|| client.tools.clone()); + (Some(Arc::new(client)), server_tools) }; - let server_tools = catalog_override.unwrap_or_else(|| client.tools.clone()); let server_tools = filter_tools(server_tools, &view.tool_filter); let server_tools = if server_name == CODEX_APPS_MCP_SERVER_NAME { prepare_codex_apps_tools_for_model(server_tools, &self.tool_plugin_provenance) @@ -280,9 +275,16 @@ impl McpConnectionSet { }; let server_tools = server_tools .into_iter() - .map(|tool| Self::with_server_metadata(tool, &view.metadata)) + .map(|mut tool| { + if client.is_none() + && let Some(annotations) = tool.tool.annotations.as_mut() + { + annotations.read_only_hint = None; + } + Self::with_server_metadata(tool, &view.metadata) + }) .collect::>(); - Some((server_name.clone(), Some(Arc::new(client)), server_tools)) + Some((server_name.clone(), client, server_tools)) })) .await; for (server_name, client, server_tools) in server_results.into_iter().flatten() { diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index e49845b876..7bf924b7f9 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -2885,6 +2885,138 @@ async fn capture_binding_shares_optional_startup_grace_across_connection_sets() } } +#[tokio::test(start_paused = true)] +async fn capture_binding_uses_cache_published_during_optional_startup() { + 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": "pending-mcp" })) + .expect("server configuration"); + let cache_context = McpToolCatalogCache::default() + .context( + "pending", + &server_config, + &runtime_context, + /*resolved_environment*/ None, + ( + &ElicitationCapability::default(), + &ClientMcpExtensions::default(), + ), + /*connection_identity*/ None, + ) + .expect("shared catalog"); + let (mut client, started, _release) = create_gated_async_managed_client( + create_test_managed_client(vec![create_test_tool("pending", "live_tool")]).await, + ); + client.tool_catalog_cache_context = Some(cache_context.clone()); + 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("pending", client); + let manager = Arc::new(manager); + + let binding = capture_binding(&manager); + tokio::pin!(binding); + assert!(futures::poll!(&mut binding).is_pending()); + started + .await + .expect("optional startup began without a cache"); + cache_context.publish_if_newest( + cache_context.begin_fetch(), + &[create_test_tool("pending", "peer_tool")], + ); + tokio::time::advance(Duration::from_secs(/*secs*/ 2)).await; + + let binding = binding.await; + assert!( + !manager.servers["pending"] + .connection + .client + .startup_complete + .load(Ordering::Acquire) + ); + assert_eq!( + model_tool_names(binding.tools()), + HashSet::from([ToolName::namespaced("mcp__pending", "peer_tool")]), + ); +} + +#[tokio::test(start_paused = true)] +async fn capture_binding_retains_cached_tools_that_expire_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": "server-a-mcp" })) + .expect("MCP server configuration"); + let cache_context = McpToolCatalogCache::default() + .context( + "server_a", + &server_config, + &runtime_context, + /*resolved_environment*/ None, + ( + &ElicitationCapability::default(), + &ClientMcpExtensions::default(), + ), + /*connection_identity*/ None, + ) + .expect("server A cache context"); + let (mut client_a, _started_a, _release_a) = create_gated_async_managed_client( + create_test_managed_client(vec![create_test_tool("server_a", "live_tool")]).await, + ); + client_a.tool_catalog_cache_context = Some(cache_context.clone()); + let (client_b, started_b, release_b) = create_gated_async_managed_client( + create_test_managed_client(vec![create_test_tool("server_b", "tool_b")]).await, + ); + cache_context.publish_if_newest( + cache_context.begin_fetch(), + &[create_test_tool("server_a", "cached_tool")], + ); + tokio::time::advance(Duration::from_secs(30 * 60 - 1)).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("server_a", client_a); + manager.insert_test_client("server_b", client_b); + Arc::get_mut(&mut manager.servers.get_mut("server_a").unwrap().connection) + .expect("unique server A connection") + .startup_trigger = Some(watch::channel(/*init*/ false).0); + manager.required_servers = vec!["server_a".to_string(), "server_b".to_string()]; + let manager = Arc::new(manager); + + let binding = capture_binding(&manager); + tokio::pin!(binding); + assert!(futures::poll!(&mut binding).is_pending()); + started_b.await.expect("server B startup should begin"); + assert!(manager.servers["server_a"].connection.startup_is_dormant()); + + // Expire A's catalog while capture is waiting for B's required startup. + tokio::time::advance(Duration::from_secs(/*secs*/ 2)).await; + assert!(cache_context.current_tools().is_none()); + release_b.send(()).expect("release server B startup"); + + let binding = tokio::time::timeout(Duration::from_secs(/*secs*/ 1), binding) + .await + .expect("cached server A should not need startup"); + assert_eq!( + model_tool_names(binding.tools()), + HashSet::from([ + ToolName::namespaced("mcp__server_a", "cached_tool"), + ToolName::namespaced("mcp__server_b", "tool_b"), + ]) + ); + assert!(manager.servers["server_a"].connection.startup_is_dormant()); +} + #[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 ff1e0c900a..a4c1cd4380 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -552,7 +552,7 @@ impl AsyncManagedClient { .is_some_and(McpToolCatalogCacheContext::has_tools) } - fn cached_tools(&self) -> Option> { + pub(crate) fn cached_tools(&self) -> Option> { self.codex_apps_tools_cache_context .as_ref() .and_then(ConnectorRuntimeContext::current_tools)