mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
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:
@@ -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(),
|
||||
|
||||
Reference in New Issue
Block a user