From e1340f2ea8d4b6a8577e734fe940fdcf46710514 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 13 Nov 2026 00:46:23 +0000 Subject: [PATCH] Deduplicate overlapping Apps startup tool listings (#61220) ## Why Concurrent Apps startups with equivalent discovery inputs each issue a `tools/list` request even though they share a live tool catalog. ## What changed Share successful tool listings across overlapping startups in the same discovery scope while keeping each client's MCP connection independent. Later startups and explicit refreshes still fetch from the server. Include time spent waiting for a shared listing in the configured startup timeout. If a listing fails or is cancelled, let a waiting startup retry using its own connection. Preserve separate catalogs for different scopes and honor servers that opt out of catalog sharing. ## Testing Add tests for shared startup listings with independent connections, waiting timeouts, cancellation recovery, and catalog isolation across discovery scopes. GitOrigin-RevId: cb626f6cb9c38e27b3d75a8da55bda584aaef6f4 --- .../codex-mcp/src/connection_manager_tests.rs | 34 +++- codex-rs/codex-mcp/src/rmcp_client.rs | 50 +++-- .../src/rmcp_client/apps_startup_tests.rs | 171 ++++++++++++++++++ .../connectors/src/connector_runtime/mod.rs | 48 ++++- .../connectors/src/connector_runtime/tests.rs | 43 +++++ 5 files changed, 306 insertions(+), 40 deletions(-) create mode 100644 codex-rs/codex-mcp/src/rmcp_client/apps_startup_tests.rs diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 0961577236..dd9705c53a 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -433,6 +433,23 @@ impl InProcessTransportFactory for MutableToolsTransportFactory { } } +pub(crate) async fn create_test_apps_client( + tool: Tool, + list_started: Option>, + release_list: Option>, + list_requests: Arc, +) -> anyhow::Result> { + let client = RmcpClient::new_in_process_client(Arc::new(RefreshTestTransportFactory { + tool, + list_started, + release_list, + next_cursor: None, + list_requests, + })) + .await?; + Ok(Arc::new(client)) +} + struct DisconnectingToolsTransportFactory { server: MutableToolsServer, disconnect: CancellationToken, @@ -856,16 +873,13 @@ pub(crate) async fn create_test_manager_with_ready_apps_client( release_list: Option>, ) -> anyhow::Result> { let tool = create_test_tool(CODEX_APPS_MCP_SERVER_NAME, tool_name); - let client = Arc::new( - RmcpClient::new_in_process_client(Arc::new(RefreshTestTransportFactory { - tool: tool.tool.clone(), - list_started, - release_list, - next_cursor: None, - list_requests: Arc::new(AtomicUsize::new(0)), - })) - .await?, - ); + let client = create_test_apps_client( + tool.tool.clone(), + list_started, + release_list, + Arc::default(), + ) + .await?; client .initialize( InitializeRequestParams::new( diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 13a7a420d8..f2b1139d2e 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -14,6 +14,9 @@ mod managed_client; #[path = "rmcp_client/status.rs"] mod status; +#[cfg(test)] +#[path = "rmcp_client/apps_startup_tests.rs"] +mod apps_startup_tests; #[cfg(test)] #[path = "rmcp_client/provenance_tests.rs"] mod provenance_tests; @@ -69,7 +72,6 @@ use codex_config::McpServerTransportConfig; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_connectors::ConnectorRuntimeContext; -use codex_connectors::ConnectorRuntimeFetchSource; use codex_exec_server::Environment; use codex_login::AuthChangeState; use codex_otel::auth_storage::AuthStorageOriginator; @@ -386,7 +388,7 @@ impl ManagedClientStartup { sharing_declined, publish_apps, is_codex_apps_mcp_server, - startup_timeout: Some(startup_timeout), + startup_timeout, tx_event, elicitation_requests, codex_apps_tools_cache_context, @@ -1182,7 +1184,7 @@ async fn start_server_task( let requested_capabilities = params.capabilities.clone(); let started_at = Instant::now(); let initialize_result = client - .initialize(params, startup_timeout, send_elicitation) + .initialize(params, Some(startup_timeout), send_elicitation) .await; record_protocol_discovery_metrics( client.protocol_mode(), @@ -1255,40 +1257,32 @@ async fn start_server_task( let list_start = Instant::now(); let server_info = mcp_server_info_from_implementation(&server_name, initialize_result.server_info); - let fetch_ticket = codex_apps_tools_cache_context - .as_ref() - .map(|context| context.begin_fetch(ConnectorRuntimeFetchSource::Startup)); if let Some(tx) = &publish_apps { tx.send_replace(Some(Arc::downgrade(&client))); } - let (tools, meta) = list_tools_for_client_uncached( + let list = list_tools_for_client_uncached( &server_name, is_codex_apps_mcp_server, /*codex_apps_refresh_trigger*/ "initial", &client, - startup_timeout, + Some(startup_timeout), catalog_item_limit, initialize_result.instructions.as_deref(), - ) - .await - .map_err(StartupOutcomeError::from)?; - let client_tools: Arc<[ToolInfo]> = - match (codex_apps_tools_cache_context.as_ref(), fetch_ticket) { - (Some(context), Some(ticket)) if server_disables_tool_catalog_cache => { - context.publish_runtime_if_newest_accepted( - ticket, - &server_info, - tools.clone(), - meta.clone(), - ); - tools.into() + ); + let (client_tools, meta) = match &codex_apps_tools_cache_context { + Some(context) => tokio::select! { + // If the lock and deadline become ready together, skip the expired request. + biased; + () = tokio::time::sleep(startup_timeout) => { + return Err(anyhow!("tools/list startup timed out after {startup_timeout:?}").into()); } - (Some(context), Some(ticket)) => context - .publish_runtime_if_newest_accepted(ticket, &server_info, tools, meta.clone()) - .shared_tools(), - (None, None) => tools.into(), - _ => unreachable!("Codex Apps fetch ticket requires cache context"), - }; + result = context.fetch_startup_tools(&server_info, list) => result?, + }, + None => { + let (tools, meta) = list.await?; + (tools.into(), meta) + } + }; if let (Some(cache_context), Some(mut fetch_ticket)) = ( tool_catalog_cache_context.as_ref(), tool_catalog_fetch_ticket, @@ -1420,7 +1414,7 @@ struct StartServerTaskParams { sharing_declined: Option, server_capabilities: Arc>>, is_codex_apps_mcp_server: bool, - startup_timeout: Option, // TODO: cancel_token should handle this. + startup_timeout: Duration, tx_event: Option>, elicitation_requests: ElicitationRequestManager, codex_apps_tools_cache_context: Option>, diff --git a/codex-rs/codex-mcp/src/rmcp_client/apps_startup_tests.rs b/codex-rs/codex-mcp/src/rmcp_client/apps_startup_tests.rs new file mode 100644 index 0000000000..6a0268d6f1 --- /dev/null +++ b/codex-rs/codex-mcp/src/rmcp_client/apps_startup_tests.rs @@ -0,0 +1,171 @@ +//! Apps tool discovery across concurrent startups on independent MCP connections. +#![cfg(test)] + +use super::*; +use crate::client_tool_catalog::ToolCatalogSnapshot; +use crate::connection_manager::tests::create_test_apps_client; +use crate::connection_manager::tests::create_test_tool; +use crate::mcp::tests::test_mcp_config; +use codex_connectors::ConnectorRuntimeContextKey; +use codex_connectors::ConnectorRuntimeManager; +use pretty_assertions::assert_eq; +use std::path::Path; +use std::sync::atomic::AtomicUsize; +use tokio::sync::Notify; + +fn tool_names(catalog: ToolCatalogSnapshot) -> Vec { + catalog + .tools + .iter() + .map(|tool| tool.tool.name.to_string()) + .collect() +} + +struct AppsStartup { + task: tokio::task::JoinHandle>, + list_started: Arc, + release_list: Arc, + list_requests: Arc, +} + +async fn start_overlapping_apps_clients( + home: &Path, + startup_timeout: Duration, +) -> anyhow::Result<(AppsStartup, AppsStartup)> { + let cache = ConnectorRuntimeManager::new_without_cache().context( + home.to_path_buf(), + ConnectorRuntimeContextKey::personal(Some("account".into()), Some("user".into())), + ); + let elicitation_requests = ElicitationRequestManager::new( + Arc::new(test_mcp_config(home.to_path_buf())), + /*reviewer*/ None, + /*lifecycle*/ None, + Default::default(), + ); + let start = |tool_name: &str| { + let (publish, mut initialized) = watch::channel(None); + let params = StartServerTaskParams { + server_key: McpServerKey::unscoped(CODEX_APPS_MCP_SERVER_NAME), + publish_apps: Some(publish), + sharing_declined: None, + server_capabilities: Arc::default(), + is_codex_apps_mcp_server: true, + startup_timeout, + tx_event: None, + elicitation_requests: elicitation_requests.clone(), + codex_apps_tools_cache_context: Some(cache.clone()), + tool_catalog_cache_context: None, + tool_catalog_fetch_ticket: None, + client_elicitation_capability: ElicitationCapability::default(), + client_mcp_extensions: ClientMcpExtensions::default(), + auth_changes: None, + catalog_item_limit: crate::pagination::MAX_CODEX_APPS_TOOL_CATALOG_ITEMS, + }; + let tool = create_test_tool(CODEX_APPS_MCP_SERVER_NAME, tool_name).tool; + async move { + let list_started = Arc::new(Notify::new()); + let release_list = Arc::new(Notify::new()); + let list_requests = Arc::new(AtomicUsize::new(0)); + let client = create_test_apps_client( + tool, + Some(Arc::clone(&list_started)), + Some(Arc::clone(&release_list)), + Arc::clone(&list_requests), + ) + .await?; + let mut startup = Box::pin(start_server_task( + CODEX_APPS_MCP_SERVER_NAME.into(), + client, + params, + )); + tokio::select! { + biased; + _ = &mut startup => anyhow::bail!("startup finished before the listing was released"), + ready = initialized.changed() => ready?, + } + Ok::<_, anyhow::Error>(AppsStartup { + task: tokio::spawn(startup), + list_started, + release_list, + list_requests, + }) + } + }; + let first = start("lookup").await?; + tokio::time::timeout(startup_timeout, first.list_started.notified()).await?; + let second = start("peer_lookup").await?; + Ok((first, second)) +} + +#[tokio::test(start_paused = true)] +async fn overlapping_apps_startups_share_tools_and_keep_their_own_connections() -> anyhow::Result<()> +{ + let home = tempfile::tempdir()?; + let (first, second) = + start_overlapping_apps_clients(home.path(), DEFAULT_STARTUP_TIMEOUT).await?; + first.release_list.notify_one(); + second.release_list.notify_one(); + let first_client = first.task.await??; + let second_client = second.task.await??; + assert_eq!(first.list_requests.load(Ordering::SeqCst), 1); + assert_eq!(second.list_requests.load(Ordering::SeqCst), 0); + for client in [&first_client, &second_client] { + assert_eq!(client.tool_catalog.read(tool_names).await, ["lookup"]); + } + + first_client.client.shutdown().await; + let reply = second_client + .call_tool( + "lookup".into(), + /*arguments*/ None, + |_| Ok(None), + Some(DEFAULT_TOOL_TIMEOUT), + ) + .await?; + assert!(!reply.is_error.unwrap_or(false)); + second_client.client.shutdown().await; + Ok(()) +} + +#[tokio::test(start_paused = true)] +async fn waiting_for_apps_tools_uses_the_configured_startup_timeout() -> anyhow::Result<()> { + let home = tempfile::tempdir()?; + let (first, second) = + start_overlapping_apps_clients(home.path(), Duration::from_secs(5)).await?; + // The second server could reply immediately if its client issued another listing. + second.release_list.notify_one(); + let (first_result, second_result) = tokio::try_join!(first.task, second.task)?; + for result in [first_result, second_result] { + assert!( + matches!(result, Err(StartupOutcomeError::Failed { error, .. }) if error.contains("timed out")) + ); + } + assert_eq!(second.list_requests.load(Ordering::SeqCst), 0); + first.release_list.notify_one(); + Ok(()) +} + +#[tokio::test(start_paused = true)] +async fn cancelling_an_apps_startup_releases_the_waiting_client() -> anyhow::Result<()> { + let home = tempfile::tempdir()?; + let (first, second) = + start_overlapping_apps_clients(home.path(), DEFAULT_STARTUP_TIMEOUT).await?; + let interval = DEFAULT_STARTUP_TIMEOUT / 3; + tokio::select! { + biased; + () = second.list_started.notified() => anyhow::bail!("peer listed before cancellation"), + () = tokio::time::sleep(interval) => {} + } + first.task.abort(); + // The peer must start its own request before the abandoned listing would have timed out. + tokio::time::timeout(interval, second.list_started.notified()).await?; + first.release_list.notify_one(); + second.release_list.notify_one(); + let recovered = second.task.await??; + assert_eq!( + recovered.tool_catalog.read(tool_names).await, + ["peer_lookup"] + ); + recovered.client.shutdown().await; + Ok(()) +} diff --git a/codex-rs/connectors/src/connector_runtime/mod.rs b/codex-rs/connectors/src/connector_runtime/mod.rs index bf0096a464..e9988404d7 100644 --- a/codex-rs/connectors/src/connector_runtime/mod.rs +++ b/codex-rs/connectors/src/connector_runtime/mod.rs @@ -6,8 +6,8 @@ //! owned by the connector metadata store, not by this module. //! Live catalog subscriptions publish only successful fetches from a matching //! discovery scope, never disk snapshots or another scope's discovery winner. -//! Equivalent live contexts share one current catalog. Each fetch still contacts -//! the server; equal definitions retain their established storage and ordering. +//! Equivalent live contexts share one current catalog and any overlapping startup +//! fetch. Later startups and explicit refreshes fetch from the server again. //! Providers expire with their last context. use std::collections::HashMap; @@ -22,12 +22,14 @@ use std::time::Duration; use std::time::Instant; use std::time::SystemTime; +use anyhow::Context; use arc_swap::ArcSwapOption; use codex_login::CodexAuth; use codex_protocol::mcp::McpServerInfo; use serde::Deserialize; use serde::Serialize; use serde::de::DeserializeOwned; +use tokio::sync::Mutex as AsyncMutex; use tokio::sync::watch; use self::persistence::load_cached_codex_apps_server_info; @@ -41,6 +43,7 @@ const MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC: &str = "codex.mcp.tools.cache_pub /// The current immutable tools for matching discovery inputs. struct CatalogProvider { scope: Vec, + startup: AsyncMutex<()>, updates: watch::Sender>>>, } @@ -260,6 +263,7 @@ impl ConnectorRuntimeContext { self.live_catalog = Some(catalog.upgrade().unwrap_or_else(|| { let provider = Arc::new(CatalogProvider { scope: scope_parts, + startup: AsyncMutex::default(), updates: watch::channel(None).0, }); *catalog = Arc::downgrade(&provider); @@ -275,6 +279,46 @@ impl ConnectorRuntimeContext { self } + /// Lists and publishes startup tools. Only overlapping calls in the same live scope + /// can join a fetch. If the listing fails or is cancelled, a waiter can use its own + /// connection to retry. Opted-out clients keep the exact tools their server returned. + #[expect( + clippy::await_holding_invalid_type, + reason = "serialize startup listings through publication within one discovery scope" + )] + pub async fn fetch_startup_tools( + &self, + server_info: &McpServerInfo, + list: impl Future, Option>)>>, + ) -> anyhow::Result<(Arc<[T]>, Option>)> { + let Some(catalog) = &self.live_catalog else { + let ticket = self.begin_fetch(ConnectorRuntimeFetchSource::Startup); + let (tools, meta) = list.await?; + self.publish_runtime_if_newest_accepted( + ticket, + server_info, + tools.clone(), + meta.clone(), + ); + return Ok((tools.into(), meta)); + }; + + let updates = catalog.updates.subscribe(); + let _guard = catalog.startup.lock().await; + // A peer may have published while this startup waited for the lock. + if !updates.borrow().has_changed() { + let ticket = self.begin_fetch(ConnectorRuntimeFetchSource::Startup); + let (tools, meta) = list.await?; + self.publish_runtime_if_newest_accepted(ticket, server_info, tools, meta); + } + // Account-wide discovery may have accepted a different scope's newer fetch. + let current = catalog.updates.borrow(); + let snapshot = current + .as_ref() + .context("Apps tool catalog is unavailable after startup")?; + Ok((snapshot.shared_tools(), snapshot.shared_tool_list_meta())) + } + /// Subscribes to accepted live tools without refetching or waiting for other clients. pub fn subscribe(&self) -> Option>>>> { self.live_catalog diff --git a/codex-rs/connectors/src/connector_runtime/tests.rs b/codex-rs/connectors/src/connector_runtime/tests.rs index 183c5262c5..a5f154ecae 100644 --- a/codex-rs/connectors/src/connector_runtime/tests.rs +++ b/codex-rs/connectors/src/connector_runtime/tests.rs @@ -514,6 +514,49 @@ fn codex_apps_tools_cache_scopes_non_utf8_home_disk_paths() { assert_ne!(cache_paths[0], cache_paths[1]); } +#[tokio::test] +async fn startup_returns_its_own_catalog_when_another_scope_fetches_first() -> anyhow::Result<()> { + let home = tempdir()?; + let context = create_codex_apps_tools_cache_context( + home.path().into(), + Some("account"), + /*chatgpt_user_id*/ None, + ); + let first = context.clone().with_live_scope("one".into()); + let other = context.with_live_scope("two".into()); + let server = create_test_server_info("Apps"); + let tools = vec![create_test_tool("apps", "one")]; + let meta = Some(Arc::new(serde_json::json!({"version": "one"}))); + let (release, wait) = tokio::sync::oneshot::channel(); + let mut list = std::pin::pin!(first.fetch_startup_tools(&server, async { + wait.await?; + Ok((tools.clone(), meta.clone())) + })); + let mut cx = std::task::Context::from_waker(std::task::Waker::noop()); + assert!(list.as_mut().poll(&mut cx).is_pending()); + let other_reply = tokio::time::timeout( + Duration::from_secs(5), + other.fetch_startup_tools(&server, async { + Ok((vec![create_test_tool("apps", "other")], None)) + }), + ) + .await??; + assert_eq!(other_reply.0.as_ref(), [create_test_tool("apps", "other")]); + release.send(()).unwrap(); + let (listed, listed_meta) = list.await?; + assert_eq!( + (listed.as_ref(), listed_meta.as_deref()), + (tools.as_slice(), meta.as_deref()) + ); + + let updated = vec![create_test_tool("apps", "updated")]; + let (fresh, _) = first + .fetch_startup_tools(&server, async { Ok((updated.clone(), None)) }) + .await?; + assert_eq!(fresh.as_ref(), updated); + Ok(()) +} + #[test] fn live_catalogs_isolate_scopes_and_reject_older_refreshes() { let codex_home = tempdir().expect("tempdir");