Resolve MCP tool catalogs concurrently (#35777)

## What changed - Resolve per-server tool catalogs concurrently in `list_all_tools` and binding capture. - After the shared startup wait, recheck servers that previously exposed cached tools so a newly ready client contributes its current catalog and callable client. ## Testing - Add gated-client tests that verify multiple server catalogs start before any one server is released. - Verify binding capture replaces cached Codex Apps tools with the ready client's tools when startup completes during capture. GitOrigin-RevId: a30d1a535ed1588f2bba21f2741e8bbd9342e339
2026-11-02 19:43:24 +00:00 · 2026-07-28 14:13:03 +00:00
parent 03748ad5e1
commit 84ccb2938b
2 changed files with 246 additions and 43 deletions
@@ -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::<Vec<_>>(),
+                ),
+                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::<Vec<_>>();
+                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::<Vec<_>>();
+            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(
@@ -429,6 +429,42 @@ async fn create_ready_async_managed_client(tools: Vec<ToolInfo>) -> 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<ToolInfo>,
    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<_>>(),
+        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::<Result<ManagedClient, StartupOutcomeError>>()