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 authored and copyberry committed 2026-07-29 18:53:15 +00:00
1 parent 09cf609218
commit 85c082cccc
3 files changed
+142 -6

No files matched your search

@@ -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(),