Preserve cached MCP tools during binding capture (#41336)

## Why Binding capture can wait for multiple MCP servers concurrently. During that wait, a tool catalog may be published by another connection or expire after it has already allowed a server to remain dormant. Re-reading only the current cache could omit those tools or trigger unnecessary startup. ## What changed - Retain the cached catalog that permits lazy startup for the duration of binding capture. - Recheck the cache after the optional startup grace so catalogs published while waiting are included. - Add regression coverage for catalogs published during startup and catalogs that expire while another server is resolving. GitOrigin-RevId: b5c66831db8f9ebecc79ff25c5b4b0182ad68a2f
2026-10-21 05:22:21 +00:00 · 2026-08-28 13:42:37 +00:00
parent 430d26b543
commit 7625343977
3 changed files with 168 additions and 34 deletions
@@ -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::<Vec<_>>();
-                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::<Vec<_>>();
-            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() {
@@ -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");
@@ -552,7 +552,7 @@ impl AsyncManagedClient {
                .is_some_and(McpToolCatalogCacheContext::has_tools)
    }

-    fn cached_tools(&self) -> Option<Vec<ToolInfo>> {
+    pub(crate) fn cached_tools(&self) -> Option<Vec<ToolInfo>> {
        self.codex_apps_tools_cache_context
            .as_ref()
            .and_then(ConnectorRuntimeContext::current_tools)