Share optional MCP startup grace across connection sets (#36011)

## Why

Each connection set could restart the one-second startup grace for the same
cacheable optional MCP server, repeatedly delaying catalog capture while that
server was still pending.

## What changed

- Store the optional startup deadline in the shared tool catalog cache entry.
- Reuse that deadline until a catalog is published, and allow a fresh grace
  period after the cached catalog expires.
- Preserve per-connection-set startup grace for servers that disable catalog
  caching.

## Testing

Added a paused-time test covering shared deadlines, immediate cached-tool reuse,
cache expiration, and disabled caching.

GitOrigin-RevId: 3494969ee7002c70ce9e4ffd370e46a22d38cc11
This commit is contained in:
jif
2026-07-29 18:53:15 +00:00
committed by copyberry
parent 09cf609218
commit 85c082cccc
3 changed files with 142 additions and 6 deletions
@@ -173,12 +173,16 @@ impl McpConnectionSet {
return;
}
if !must_wait_for_startup {
if tokio::time::timeout_at(
optional_startup_deadline,
view.connection.client.client(),
)
.await
.is_err()
let startup_deadline = view
.connection
.client
.tool_catalog_cache_context
.as_ref()
.map(|cache| cache.optional_startup_deadline(optional_startup_deadline))
.unwrap_or(optional_startup_deadline);
if tokio::time::timeout_at(startup_deadline, view.connection.client.client())
.await
.is_err()
{
trace!(server_name = %server_name, "omitting pending optional MCP server");
}
@@ -1873,6 +1873,121 @@ async fn capture_binding_skips_pending_optional_servers_after_one_shared_startup
assert!(binding.is_err(), "explicitly requested servers must wait");
}
#[tokio::test(start_paused = true)]
async fn capture_binding_shares_optional_startup_grace_across_connection_sets() {
let approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
let permission_profile = Constrained::allow_any(PermissionProfile::default());
let cache = McpToolCatalogCache::default();
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("pending MCP server configuration");
let cache_context = cache
.context(
"pending",
&server_config,
&runtime_context,
/*resolved_environment*/ None,
&ElicitationCapability::default(),
/*supports_openai_form_elicitation*/ false,
)
.expect("shared pending MCP catalog");
let create_connection_set = || {
let mut manager = McpConnectionSet::new_uninitialized(
&approval_policy,
&permission_profile,
/*prefix_mcp_tool_names*/ true,
);
manager.insert_test_client(
"pending",
AsyncManagedClient {
client: futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared(),
is_codex_apps_mcp_server: false,
cached_server_info: None,
codex_apps_tools_cache_context: None,
tool_catalog_cache_context: Some(cache_context.clone()),
startup_complete: Arc::new(AtomicBool::new(false)),
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
);
Arc::new(manager)
};
let first_started = tokio::time::Instant::now();
let first = tokio::time::timeout(
Duration::from_millis(1500),
capture_binding(&create_connection_set()),
)
.await
.expect("the first thread should receive the optional startup grace");
assert!(first.tools().is_empty());
assert_eq!(first_started.elapsed(), Duration::from_secs(1));
let second = tokio::time::timeout(
Duration::from_millis(1),
capture_binding(&create_connection_set()),
)
.await
.expect("the next thread must not restart the same server's startup grace");
assert!(second.tools().is_empty());
cache_context.publish_if_newest(
cache_context.begin_fetch(),
&[create_test_tool("pending", "cached_tool")],
);
let deadline_after_publication = tokio::time::Instant::now() + Duration::from_secs(1);
assert_eq!(
cache_context.optional_startup_deadline(deadline_after_publication),
deadline_after_publication,
"publishing a catalog must not install a stale startup deadline"
);
let cached = tokio::time::timeout(
Duration::from_millis(1),
capture_binding(&create_connection_set()),
)
.await
.expect("cached tools should be immediately available to later threads");
assert_eq!(
cached
.tools()
.iter()
.map(|tool| tool.callable_name.as_str())
.collect::<Vec<_>>(),
vec!["cached_tool"]
);
tokio::time::advance(Duration::from_secs(30 * 60 + 1)).await;
assert!(
tokio::time::timeout(
Duration::from_millis(1),
capture_binding(&create_connection_set()),
)
.await
.is_err(),
"an expired catalog should receive a fresh startup grace"
);
cache_context.disable();
for _ in 0..2 {
let started = tokio::time::Instant::now();
let binding = tokio::time::timeout(
Duration::from_millis(1500),
capture_binding(&create_connection_set()),
)
.await
.expect("non-cacheable servers should keep their per-thread startup grace");
assert!(binding.tools().is_empty());
assert_eq!(started.elapsed(), Duration::from_secs(1));
}
}
#[tokio::test]
async fn capture_binding_resolves_concurrently_and_rechecks_cached_clients() {
let codex_home = tempdir().expect("tempdir");
@@ -51,6 +51,7 @@ struct ToolCatalogCacheEntry {
#[derive(Default)]
struct ToolCatalogCacheState {
snapshot: Option<ToolCatalogSnapshot>,
optional_startup_deadline: Option<Instant>,
last_accepted_generation: u64,
disabled_by_server: bool,
}
@@ -108,6 +109,21 @@ impl McpToolCatalogCacheContext {
self.current_tools().is_some()
}
pub(crate) fn optional_startup_deadline(&self, default_deadline: Instant) -> Instant {
let mut state = lock_unpoisoned(&self.entry.state);
if state.disabled_by_server
|| state
.snapshot
.as_ref()
.is_some_and(|snapshot| snapshot.published_at.elapsed() <= TOOL_CATALOG_CACHE_TTL)
{
return default_deadline;
}
*state
.optional_startup_deadline
.get_or_insert(default_deadline)
}
pub(crate) fn current_tools(&self) -> Option<Vec<ToolInfo>> {
lock_unpoisoned(&self.entry.state)
.snapshot
@@ -147,6 +163,7 @@ impl McpToolCatalogCacheContext {
tool.tool.annotations = None;
}
state.last_accepted_generation = ticket.generation;
state.optional_startup_deadline = None;
state.snapshot = Some(ToolCatalogSnapshot {
tools,
published_at: Instant::now(),