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");