Reuse MCP bindings while cached servers remain dormant (#44121)

## Why

Dormant MCP servers with usable cached tool catalogs prevented binding reuse, causing each model step to capture a new binding even when the catalog was unchanged.

## What changed

Track dormant catalog revisions alongside ready client revisions so bindings can be reused until the catalog changes or the server starts. Explicit server and plugin requirements still trigger binding capture and startup. Check cached catalog availability without cloning tool definitions.

## Testing

Extend the lazy-startup integration test to verify that two model steps share one binding while the server stays dormant, and that starting the server invalidates that binding.

GitOrigin-RevId: ef0469a85000f84a928946c24350e3bda1522fb5
This commit is contained in:
jif authored and copyberry committed 2026-09-09 10:51:24 +00:00
1 parent 9e868bd9dc
commit 20f109eadb
6 files changed
+160 -20

No files matched your search

@@ -21,6 +21,7 @@ use startup::emit_update;
use startup::mcp_init_error_display;
use startup::mcp_startup_failure_reason;
use startup::should_share_codex_apps_tools_cache;
pub(crate) use tool_catalog::BindingCatalogRevision;
pub use tool_catalog::tool_is_model_visible;
use std::collections::HashMap;
@@ -58,10 +58,19 @@ pub fn tool_is_model_visible(tool: &ToolInfo) -> bool {
.any(|target| target.as_str() == Some(MCP_UI_MODEL_VISIBILITY))
}
/// Catalog identity within one published connection set, including dormant servers.
#[derive(PartialEq, Eq)]
pub(crate) enum BindingCatalogRevision {
Ready(ClientToolCatalogRevision),
Dormant(u64),
}
impl McpConnectionSet {
pub(crate) async fn stable_catalog_revisions(
&self,
) -> Option<HashMap<String, ClientToolCatalogRevision>> {
required_servers: &[String],
required_plugins: &HashSet<String>,
) -> Option<HashMap<String, BindingCatalogRevision>> {
let mut revisions = HashMap::new();
for (server_name, view) in &self.servers {
if !view
@@ -70,7 +79,31 @@ impl McpConnectionSet {
.startup_complete
.load(Ordering::Acquire)
{
return None;
// A cached catalog can remain stable without starting its server.
// Explicit requirements must still start the server during binding capture.
if !view.connection.startup_is_dormant()
|| view.connection.client.cancel_token.is_cancelled()
|| required_servers
.iter()
.any(|required| required == server_name)
|| (self.is_selected_plugin_mcp_server(server_name)
&& self
.plugin_id_for_mcp_server_name(server_name)
.is_some_and(|plugin_id| required_plugins.contains(plugin_id)))
{
return None;
}
let revision = view
.connection
.client
.tool_catalog_cache_context
.as_ref()?
.current_revision()?;
revisions.insert(
server_name.clone(),
BindingCatalogRevision::Dormant(revision),
);
continue;
}
let Some(client) = view.connection.client.ready_client() else {
if !view.connection.client.is_codex_apps_mcp_server
@@ -87,10 +120,10 @@ impl McpConnectionSet {
let revision = client.tool_catalog.read(|catalog| catalog.revision).await;
revisions.insert(
server_name.clone(),
ClientToolCatalogRevision {
BindingCatalogRevision::Ready(ClientToolCatalogRevision {
catalog: Arc::clone(&client.tool_catalog),
revision,
},
}),
);
}
Some(revisions)
@@ -168,6 +201,7 @@ impl McpConnectionSet {
(tools, errors)
}
#[instrument(level = "trace", skip_all)]
pub(crate) async fn capture_binding_with_metadata(
self: &Arc<Self>,
config: Arc<crate::McpConfig>,
@@ -2731,7 +2731,15 @@ async fn capture_binding_skips_pending_optional_servers_after_configured_shared_
required_manager.required_servers = vec!["pending-selected".to_string()];
let manager = Arc::new(manager);
assert!(manager.stable_catalog_revisions().await.is_none());
assert!(
manager
.stable_catalog_revisions(
/*required_servers*/ &[],
/*required_plugins*/ &HashSet::new()
)
.await
.is_none()
);
let started = tokio::time::Instant::now();
let binding = tokio::time::timeout(
Duration::from_millis(500),
@@ -2887,9 +2895,25 @@ async fn stable_catalog_revisions_ignore_terminal_optional_server_failures() {
assert!(failed.client().await.is_err());
manager.insert_test_client("failed", failed);
assert!(manager.stable_catalog_revisions().await.is_some());
assert!(
manager
.stable_catalog_revisions(
/*required_servers*/ &[],
/*required_plugins*/ &HashSet::new()
)
.await
.is_some()
);
manager.required_servers.push("failed".to_string());
assert!(manager.stable_catalog_revisions().await.is_none());
assert!(
manager
.stable_catalog_revisions(
/*required_servers*/ &[],
/*required_plugins*/ &HashSet::new()
)
.await
.is_none()
);
manager.required_servers.clear();
let binding = capture_binding(&Arc::new(manager)).await;
@@ -3840,7 +3864,15 @@ async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client()
};
let manager = Arc::new(manager);
assert!(manager.stable_catalog_revisions().await.is_none());
assert!(
manager
.stable_catalog_revisions(
/*required_servers*/ &[],
/*required_plugins*/ &HashSet::new()
)
.await
.is_none()
);
let reconnect_finished_wait = reconnect_finished.notified();
let tools = manager.list_all_tools().await;
assert!(tools.is_empty());
@@ -3855,7 +3887,15 @@ 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!(manager.stable_catalog_revisions().await.is_some());
assert!(
manager
.stable_catalog_revisions(
/*required_servers*/ &[],
/*required_plugins*/ &HashSet::new()
)
.await
.is_some()
);
let step = capture_binding(&manager).await;
let prepared = step
+7 -4
View File
@@ -46,8 +46,8 @@ use tokio_util::sync::CancellationToken;
use crate::McpConfig;
use crate::binding::McpBinding;
use crate::client_tool_catalog::ClientToolCatalogRevision;
use crate::client_tool_catalog::CodexAppsToolSnapshot;
use crate::connection_manager::BindingCatalogRevision;
use crate::connection_manager::McpConnectionSet;
use crate::elicitation::ElicitationLifecycle;
use crate::elicitation::ElicitationRequestRouter;
@@ -119,7 +119,7 @@ struct PublishedMcpRuntime {
}
struct CachedMcpBinding {
catalog_revisions: HashMap<String, ClientToolCatalogRevision>,
catalog_revisions: HashMap<String, BindingCatalogRevision>,
binding: Arc<McpBinding>,
}
@@ -363,7 +363,10 @@ impl McpRuntime {
required_plugins: &HashSet<String>,
) -> Option<Arc<McpBinding>> {
let config = Arc::clone(current.config.as_ref()?);
let stable_catalog_revisions = current.connections.stable_catalog_revisions().await;
let stable_catalog_revisions = current
.connections
.stable_catalog_revisions(required_servers, required_plugins)
.await;
if let Some(catalog_revisions) = &stable_catalog_revisions {
let cached = current
.cached_binding
@@ -390,7 +393,7 @@ impl McpRuntime {
if let Some(catalog_revisions) = stable_catalog_revisions
&& current
.connections
.stable_catalog_revisions()
.stable_catalog_revisions(required_servers, required_plugins)
.await
.as_ref()
== Some(&catalog_revisions)
+12 -1
View File
@@ -116,7 +116,18 @@ impl Default for ToolCatalogCacheEntry {
impl McpToolCatalogCacheContext {
pub(crate) fn has_tools(&self) -> bool {
self.current_tools().is_some_and(|tools| !tools.is_empty())
self.current_revision().is_some()
}
/// Identifies the current usable catalog without cloning its tool definitions.
/// The entry is fixed for a connection; accepted publications advance its revision.
pub(crate) fn current_revision(&self) -> Option<u64> {
let state = lock_unpoisoned(&self.entry.state);
let snapshot = state.snapshot.as_ref()?;
(!state.disabled_by_server
&& !snapshot.tools.is_empty()
&& snapshot.published_at.elapsed() <= TOOL_CATALOG_CACHE_TTL)
.then_some(state.last_accepted_generation)
}
pub(crate) fn optional_startup_deadline(