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 47afb2242d..0a3447e56c 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -7,6 +7,7 @@ use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use codex_connectors::ConnectorRuntimeFetchSource; +use futures::future::join_all; use tracing::Instrument; use tracing::instrument; use tracing::trace; @@ -63,7 +64,7 @@ impl McpConnectionSet { let mut tools = Vec::new(); let mut available_server_count = 0; let mut unavailable_server_count = 0; - for (server_name, view) in &self.servers { + let server_results = join_all(self.servers.iter().map(|(server_name, view)| async move { view.connection.client.reconnect_failed_startup().await; let has_cached_tools = view.connection.client.has_cached_tools(); let startup_complete = view @@ -76,7 +77,7 @@ impl McpConnectionSet { } else { None }; - let Some(server_tools) = async { + let server_tools = async { match catalog_override { Some(tools) => { let tools = filter_tools(tools, &view.tool_filter); @@ -94,23 +95,34 @@ impl McpConnectionSet { has_cached_tools, startup_complete )) - .await - else { - unavailable_server_count += 1; - trace!( - server_name = %server_name, - has_cached_tools, - startup_complete, - "MCP server tools unavailable while building tool list" - ); - continue; - }; - available_server_count += 1; - tools.extend( - server_tools - .into_iter() - .map(|tool| Self::with_server_metadata(tool, &view.metadata)), - ); + .await; + match server_tools { + Some(server_tools) => Some( + server_tools + .into_iter() + .map(|tool| Self::with_server_metadata(tool, &view.metadata)) + .collect::>(), + ), + None => { + trace!( + server_name = %server_name, + has_cached_tools, + startup_complete, + "MCP server tools unavailable while building tool list" + ); + None + } + } + })) + .await; + for server_tools in server_results { + match server_tools { + Some(server_tools) => { + available_server_count += 1; + tools.extend(server_tools); + } + None => unavailable_server_count += 1, + } } let tools = normalize_tools_for_model_with_prefix( tools, @@ -142,7 +154,7 @@ impl McpConnectionSet { let optional_startup_deadline = *self .optional_startup_deadline .get_or_init(|| tokio::time::Instant::now() + OPTIONAL_MCP_STARTUP_GRACE); - for (server_name, view) in &self.servers { + join_all(self.servers.iter().map(|(server_name, view)| async move { if !view .connection .client @@ -158,35 +170,50 @@ 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 { - if let Some(server_tools) = - view.listed_tools(&self.tool_plugin_provenance).await - { - listed_tools.extend(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) - })); - } - continue; + return; } - if !must_wait_for_startup - && tokio::time::timeout_at( + if !must_wait_for_startup { + if tokio::time::timeout_at( optional_startup_deadline, view.connection.client.client(), ) .await .is_err() - { - trace!(server_name = %server_name, "omitting pending optional MCP server"); - continue; + { + trace!(server_name = %server_name, "omitting pending optional MCP server"); + } + return; } let _ = view.connection.client.client().await; } + })) + .await; + let server_results = join_all(self.servers.iter().map(|(server_name, view)| async move { + 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.client().await else { trace!(server_name = %server_name, "omitting MCP server without an exact ready client"); - continue; + return None; }; client.tool_timeout = view.tool_timeout; let catalog_override = if server_name == CODEX_APPS_MCP_SERVER_NAME { @@ -204,12 +231,18 @@ impl McpConnectionSet { &self.tool_plugin_provenance, ) }; - clients.insert(server_name.clone(), Arc::new(client)); - listed_tools.extend( - server_tools - .into_iter() - .map(|tool| Self::with_server_metadata(tool, &view.metadata)), - ); + let server_tools = server_tools + .into_iter() + .map(|tool| Self::with_server_metadata(tool, &view.metadata)) + .collect::>(); + Some((server_name.clone(), Some(Arc::new(client)), server_tools)) + })) + .await; + for (server_name, client, server_tools) in server_results.into_iter().flatten() { + if let Some(client) = client { + clients.insert(server_name, client); + } + listed_tools.extend(server_tools); } let clients = Arc::new(McpBindingClients::new(clients)); let listed_tools = normalize_tools_for_model_with_prefix( diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index ed8f226e1d..8230ff0334 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -429,6 +429,42 @@ async fn create_ready_async_managed_client(tools: Vec) -> AsyncManaged } } +fn create_gated_async_managed_client( + client: ManagedClient, +) -> ( + AsyncManagedClient, + tokio::sync::oneshot::Receiver<()>, + tokio::sync::oneshot::Sender<()>, +) { + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let startup_complete = Arc::new(AtomicBool::new(false)); + let startup_complete_for_client = Arc::clone(&startup_complete); + let client = async move { + started_tx.send(()).expect("signal client startup"); + release_rx.await.expect("release client startup"); + startup_complete_for_client.store(true, std::sync::atomic::Ordering::Release); + Ok(client) + } + .boxed() + .shared(); + + ( + AsyncManagedClient { + client, + is_codex_apps_mcp_server: false, + cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_catalog_cache_context: None, + startup_complete, + startup_reconnect: None, + cancel_token: CancellationToken::new(), + }, + started_rx, + release_tx, + ) +} + async fn create_test_manager_with_ready_apps_client( cache_context: ConnectorRuntimeContext, tool_name: &str, @@ -1749,6 +1785,99 @@ 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 capture_binding_resolves_concurrently_and_rechecks_cached_clients() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + store_current_tools( + &cache_context, + vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "shared_cached_tool", + )], + ); + let ready_apps_client = create_test_managed_client(vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "client_local_tool", + )]) + .await; + let (mut apps_client, apps_started, release_apps) = + create_gated_async_managed_client(ready_apps_client); + apps_client.is_codex_apps_mcp_server = true; + apps_client.codex_apps_tools_cache_context = Some(cache_context); + let first_client = + create_test_managed_client(vec![create_test_tool("first", "first_tool")]).await; + let second_client = + create_test_managed_client(vec![create_test_tool("second", "second_tool")]).await; + let (first_client, first_started, release_first) = + create_gated_async_managed_client(first_client); + let (second_client, second_started, release_second) = + create_gated_async_managed_client(second_client); + + 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, + ); + manager.insert_test_client(CODEX_APPS_MCP_SERVER_NAME, apps_client); + manager.insert_test_client("first", first_client); + manager.insert_test_client("second", second_client); + let manager = Arc::new(manager); + + let manager_for_startup = Arc::clone(&manager); + let startup = tokio::spawn(async move { + manager_for_startup + .wait_for_server_startup(CODEX_APPS_MCP_SERVER_NAME) + .await + }); + tokio::time::timeout(Duration::from_secs(1), apps_started) + .await + .expect("Codex Apps startup should begin") + .expect("signal Codex Apps startup"); + + let manager_for_binding = Arc::clone(&manager); + let binding = tokio::spawn(async move { capture_binding(&manager_for_binding).await }); + tokio::time::timeout(Duration::from_secs(1), async { + first_started.await.expect("first server startup"); + second_started.await.expect("second server startup"); + }) + .await + .expect("both uncached servers should start before either is released"); + + release_apps.send(()).expect("release Codex Apps startup"); + assert!(startup.await.expect("Codex Apps startup task")); + release_first.send(()).expect("release first server"); + release_second.send(()).expect("release second server"); + + let binding = binding.await.expect("binding capture should complete"); + assert_eq!( + binding + .tools() + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + HashSet::from(["client_local_tool", "first_tool", "second_tool"]) + ); + assert!( + binding + .prepare_call(CODEX_APPS_MCP_SERVER_NAME, "client_local_tool") + .is_some() + ); + assert!( + binding + .prepare_call(CODEX_APPS_MCP_SERVER_NAME, "shared_cached_tool") + .is_none() + ); + assert!(binding.prepare_call("first", "first_tool").is_some()); + assert!(binding.prepare_call("second", "second_tool").is_some()); +} + #[tokio::test] async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { let managed_client = @@ -1780,6 +1909,47 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { ); } +#[tokio::test] +async fn list_all_tools_resolves_server_catalogs_concurrently() { + let first_client = create_test_managed_client(vec![create_test_tool("first", "search")]).await; + let second_client = + create_test_managed_client(vec![create_test_tool("second", "lookup")]).await; + let (first_client, first_started, release_first) = + create_gated_async_managed_client(first_client); + let (second_client, second_started, release_second) = + create_gated_async_managed_client(second_client); + 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, + ); + manager.insert_test_client("first", first_client); + manager.insert_test_client("second", second_client); + let manager = Arc::new(manager); + let manager_for_listing = Arc::clone(&manager); + let listing = tokio::spawn(async move { manager_for_listing.list_all_tools().await }); + + tokio::time::timeout(Duration::from_secs(1), async { + first_started.await.expect("first server startup"); + second_started.await.expect("second server startup"); + }) + .await + .expect("both server catalogs should start before either is released"); + release_first.send(()).expect("release first server"); + release_second.send(()).expect("release second server"); + + let tools = listing.await.expect("tool listing should complete"); + assert_eq!( + model_tool_names(&tools), + HashSet::from([ + ToolName::namespaced("mcp__first", "search"), + ToolName::namespaced("mcp__second", "lookup"), + ]) + ); +} + #[tokio::test] async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { let pending_client = futures::future::pending::>()