From 85c082ccccf6b5ac4d6c31d14f960057348b78f4 Mon Sep 17 00:00:00 2001 From: jif Date: Wed, 29 Jul 2026 18:48:00 +0000 Subject: [PATCH] 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 --- .../src/connection_manager/tool_catalog.rs | 16 ++- .../codex-mcp/src/connection_manager_tests.rs | 115 ++++++++++++++++++ codex-rs/codex-mcp/src/tool_catalog_cache.rs | 17 +++ 3 files changed, 142 insertions(+), 6 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs index 0a3447e56c..3e4f7721d6 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -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"); } diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 63a9cc3675..e78e05668d 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -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::>() + .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!["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"); diff --git a/codex-rs/codex-mcp/src/tool_catalog_cache.rs b/codex-rs/codex-mcp/src/tool_catalog_cache.rs index 15c3ae3d59..76062e388d 100644 --- a/codex-rs/codex-mcp/src/tool_catalog_cache.rs +++ b/codex-rs/codex-mcp/src/tool_catalog_cache.rs @@ -51,6 +51,7 @@ struct ToolCatalogCacheEntry { #[derive(Default)] struct ToolCatalogCacheState { snapshot: Option, + optional_startup_deadline: Option, 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> { 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(),