diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs index 447cbd62e6..ade3747729 100644 --- a/codex-rs/codex-mcp/src/binding.rs +++ b/codex-rs/codex-mcp/src/binding.rs @@ -22,6 +22,7 @@ use serde_json::Value as JsonValue; use crate::McpConfig; use crate::binding_clients::McpBindingClients; +use crate::client_tool_catalog::ToolCatalogSnapshot; use crate::connection_manager::McpConnectionSet; use crate::rmcp_client::ManagedClient; use crate::server::McpServerMetadata; @@ -172,7 +173,7 @@ pub struct PreparedMcpCall { connections: Arc, client: Arc, config: Arc, - catalog_revision: u64, + catalog_snapshot: Arc, tool_info: ToolInfo, server_name: String, server_metadata: McpServerMetadata, @@ -189,7 +190,7 @@ impl PreparedMcpCall { connections: Arc, client: Arc, config: Arc, - catalog_revision: u64, + catalog_snapshot: Arc, tool_info: ToolInfo, server_metadata: McpServerMetadata, plugin_id: Option, @@ -201,7 +202,7 @@ impl PreparedMcpCall { connections, client, config, - catalog_revision, + catalog_snapshot, tool_info, server_name, server_metadata, @@ -299,7 +300,7 @@ impl PreparedMcpCall { } /// Runs irreversible call preparation and execution under the authority of - /// this call's exact catalog revision and the extensions owned by the Codex session. + /// this call's captured catalog and the extensions owned by the Codex session. /// A caller-supplied timeout can further restrict the server's configured timeout. pub async fn call_with_preparation( &self, @@ -319,7 +320,7 @@ impl PreparedMcpCall { let tool_name = self.tool_info.tool.name.to_string(); self.client .tool_catalog - .run_with_revision(self.catalog_revision, || async { + .run_with_snapshot(&self.catalog_snapshot, || async { let (arguments, meta) = prepare().await?; let timeout_deadline = effective_timeout.map(|timeout| tokio::time::Instant::now() + timeout); diff --git a/codex-rs/codex-mcp/src/binding_tests.rs b/codex-rs/codex-mcp/src/binding_tests.rs index edfef4a72b..447a92364d 100644 --- a/codex-rs/codex-mcp/src/binding_tests.rs +++ b/codex-rs/codex-mcp/src/binding_tests.rs @@ -118,7 +118,7 @@ async fn test_step( Arc::clone(&connections), managed_client, Arc::clone(&config), - /*catalog_revision*/ 0, + tool_catalog.read(Arc::new).await, tool.clone(), McpServerMetadata { environment_id: format!("{label}-environment"), diff --git a/codex-rs/codex-mcp/src/client_tool_catalog.rs b/codex-rs/codex-mcp/src/client_tool_catalog.rs index c95a8d44e4..0501f1c782 100644 --- a/codex-rs/codex-mcp/src/client_tool_catalog.rs +++ b/codex-rs/codex-mcp/src/client_tool_catalog.rs @@ -1,10 +1,9 @@ -//! Tool catalog state and synchronization owned by one live MCP client. +//! Client-local catalog revisions and synchronization for MCP calls. //! -//! Reusing the client preserves this state; a replacement client starts a new -//! catalog. Snapshot reads and revision-checked calls keep the locks private; -//! successful refreshes publish after calls using the current catalog finish. -//! Refresh results retain raw tools and eligibility from the same published runtime. -//! Optional live updates are adopted before reading or starting another call. +//! Apps clients read the provider's current catalog without retaining old arrays. +//! Only active readers, calls, and frozen model bindings pin earlier snapshots. +//! Explicit refreshes publish after calls using this client's current revision finish. +//! Equivalent shared publications preserve calls, but explicit refreshes invalidate them. use std::collections::HashSet; use std::future::Future; @@ -31,67 +30,111 @@ pub struct CodexAppsToolSnapshot { } pub(crate) struct ClientToolCatalog { - current: RwLock, + current: RwLock, /// Serialize fetches without blocking calls against the current catalog. refresh_lock: Mutex<()>, } +struct CatalogState { + revision: u64, + /// The catalog revision installed by this client's last explicit refresh. + last_refresh_revision: u64, + source: CatalogSource, +} + +/// Ordinary MCP clients own their tools; Apps clients observe the shared provider. +enum CatalogSource { + Local(Arc<[ToolInfo]>), + Live { + updates: ToolCatalogUpdates, + tools_version: u64, + }, +} + pub(crate) struct ToolCatalogSnapshot { - /// Advances on explicit refresh or adoption of changed live tools. pub(crate) revision: u64, - pub(crate) tools: Vec, - updates: Option, + pub(crate) tools: Arc<[ToolInfo]>, } impl ClientToolCatalog { - pub(crate) fn new(tools: Vec, mut updates: Option) -> Self { - let tools = updates - .as_mut() - .and_then(|updates| { - updates + /// Live clients read the provider; callers publish startup tools before subscribing. + pub(crate) fn new( + tools: impl Into>, + updates: Option, + ) -> Self { + let source = match updates { + Some(mut updates) => { + let tools_version = updates .borrow_and_update() .as_ref() - .map(|snapshot| snapshot.tools().to_vec()) - }) - .unwrap_or(tools); + .map_or(0, |snapshot| snapshot.tools_version()); + CatalogSource::Live { + updates, + tools_version, + } + } + None => CatalogSource::Local(tools.into()), + }; Self { - current: RwLock::new(ToolCatalogSnapshot { + current: RwLock::new(CatalogState { revision: 0, - tools, - updates, + last_refresh_revision: 0, + source, }), refresh_lock: Mutex::new(()), } } - pub(crate) async fn read(&self, read: impl FnOnce(&ToolCatalogSnapshot) -> R) -> R { - let current = self.read_current().await; - read(¤t) + pub(crate) async fn read(&self, read: impl FnOnce(ToolCatalogSnapshot) -> R) -> R { + let (_current, snapshot) = self.read_current().await; + read(snapshot) } - async fn read_current(&self) -> RwLockReadGuard<'_, ToolCatalogSnapshot> { + /// Captures tools and their client-local revision together, without retaining them in the client. + async fn read_current(&self) -> (RwLockReadGuard<'_, CatalogState>, ToolCatalogSnapshot) { loop { { let current = self.current.read().await; - if !current - .updates - .as_ref() - .is_some_and(|updates| updates.has_changed().unwrap_or(false)) - { - return current; + let tools = match ¤t.source { + CatalogSource::Local(tools) => Some(Arc::clone(tools)), + CatalogSource::Live { updates, .. } => { + // Hold the watch borrow while checking its version so the tools and + // revision come from the same publication. + let snapshot = updates.borrow(); + (!updates.has_changed().unwrap_or(false)).then(|| { + snapshot + .as_ref() + .map(|snapshot| snapshot.shared_tools()) + .unwrap_or_default() + }) + } + }; + if let Some(tools) = tools { + let snapshot = ToolCatalogSnapshot { + revision: current.revision, + tools, + }; + return (current, snapshot); } } let mut current = self.current.write().await; - if let Some(updates) = current.updates.as_mut() - && updates.has_changed().unwrap_or(false) - { - let snapshot = updates.borrow_and_update().clone(); - if let Some(snapshot) = snapshot - && current.tools != snapshot.tools() - { - current.tools = snapshot.tools().to_vec(); - current.revision += 1; + let changed = match &mut current.source { + CatalogSource::Live { + updates, + tools_version, + } => { + let version = updates + .borrow_and_update() + .as_ref() + .map_or(0, |snapshot| snapshot.tools_version()); + let changed = *tools_version != version; + *tools_version = version; + changed } + CatalogSource::Local(_) => false, + }; + if changed { + current.revision += 1; } } } @@ -112,17 +155,20 @@ impl ClientToolCatalog { let (tools, context) = fetch().await?; let mut current = self.current.write().await; let result = publish(&tools, context); - current.tools = current - .updates - .as_mut() - .and_then(|updates| { - updates + match &mut current.source { + CatalogSource::Local(current_tools) => *current_tools = tools.into(), + CatalogSource::Live { + updates, + tools_version, + } => { + *tools_version = updates .borrow_and_update() .as_ref() - .map(|snapshot| snapshot.tools().to_vec()) - }) - .unwrap_or(tools); + .map_or(0, |snapshot| snapshot.tools_version()); + } + } current.revision += 1; + current.last_refresh_revision = current.revision; Ok(result) } @@ -131,17 +177,20 @@ impl ClientToolCatalog { clippy::await_holding_invalid_type, reason = "catalog publication must wait for call preparation and execution" )] - pub(crate) async fn run_with_revision( + pub(crate) async fn run_with_snapshot( &self, - expected_revision: u64, + expected: &ToolCatalogSnapshot, run: F, ) -> Option where F: FnOnce() -> Fut, Fut: Future, { - let current = self.read_current().await; - if current.revision != expected_revision { + let (current, snapshot) = self.read_current().await; + if current.last_refresh_revision > expected.revision + || (snapshot.revision != expected.revision + && !catalogs_match(&snapshot.tools, &expected.tools)) + { return None; } let result = run().await; @@ -150,6 +199,32 @@ impl ClientToolCatalog { } } +/// Compares complete definitions independently of tool-list order. Stable sorting keeps +/// conflicting duplicate identities in their original order, since deduplication can pick the +/// first definition. This slow path runs only after the client's revision has changed. +fn catalogs_match(left: &[ToolInfo], right: &[ToolInfo]) -> bool { + if left == right { + return true; + } + if left.len() != right.len() { + return false; + } + let mut left = left.iter().collect::>(); + let mut right = right.iter().collect::>(); + for tools in [&mut left, &mut right] { + tools.sort_by_key(|tool| { + ( + tool.server_name.as_str(), + tool.tool.name.as_ref(), + tool.connector_id.as_deref(), + tool.callable_namespace.as_str(), + tool.callable_name.as_str(), + ) + }); + } + left == right +} + /// A binding cache key includes client identity, since new clients start at zero. #[derive(Clone)] pub(crate) struct ClientToolCatalogRevision { diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 933fd60ad1..6022f1c119 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -352,8 +352,7 @@ impl McpConnectionSet { let shares_codex_apps_tools_cache = is_host_owned_codex_apps && should_share_codex_apps_tools_cache(&server_name, uses_env_bearer_token); let codex_apps_tools_cache_context = shares_codex_apps_tools_cache.then(|| { - // Tools/list has no thread selection or UI capabilities. Only equivalent - // transport/auth and listing settings may share executable Apps tools. + // Only equivalent discovery inputs may share executable Apps tools. let mut transport = configured_config.transport.clone(); if let McpServerTransportConfig::StreamableHttp { http_headers: Some(headers), 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 8b750ba326..58b810366c 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -21,6 +21,7 @@ use crate::binding::PreparedMcpCall; use crate::binding_clients::McpBindingClients; use crate::client_tool_catalog::ClientToolCatalogRevision; use crate::client_tool_catalog::CodexAppsToolSnapshot; +use crate::client_tool_catalog::ToolCatalogSnapshot; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::rmcp_client::CODEX_APPS_REFRESH_DURATION_METRIC; use crate::rmcp_client::MCP_TOOLS_LIST_DURATION_METRIC; @@ -296,11 +297,9 @@ impl McpConnectionSet { return None; }; client.tool_timeout = view.tool_timeout; - let (revision, server_tools) = client - .tool_catalog - .read(|catalog| (catalog.revision, catalog.tools.clone())) - .await; - (Some((Arc::new(client), revision)), server_tools) + let snapshot = client.tool_catalog.read(Arc::new).await; + let server_tools = snapshot.tools.to_vec(); + (Some((Arc::new(client), snapshot)), server_tools) }; let server_tools = filter_tools(server_tools, &view.tool_filter); let server_tools = if server_name == CODEX_APPS_MCP_SERVER_NAME { @@ -326,8 +325,8 @@ impl McpConnectionSet { })) .await; for (server_name, client, server_tools) in server_results.into_iter().flatten() { - if let Some((client, revision)) = client { - clients.insert(server_name, (client, revision)); + if let Some((client, snapshot)) = client { + clients.insert(server_name, (client, snapshot)); } listed_tools.extend(server_tools); } @@ -340,7 +339,7 @@ impl McpConnectionSet { let mut calls = std::collections::HashMap::with_capacity(listed_tools.len()); for tool_info in listed_tools { let model_visible = crate::tool_is_model_visible(&tool_info); - let Some((client, revision)) = clients.get(&tool_info.server_name) else { + let Some((client, snapshot)) = clients.get(&tool_info.server_name) else { if model_visible { tools.push(tool_info); } @@ -350,7 +349,7 @@ impl McpConnectionSet { &tool_info, Arc::clone(client), Arc::clone(&config), - *revision, + Arc::clone(snapshot), ) else { trace!( server_name = %tool_info.server_name, @@ -391,7 +390,7 @@ impl McpConnectionSet { tool_info: &ToolInfo, client: Arc, config: Arc, - tool_catalog_revision: u64, + tool_catalog_snapshot: Arc, ) -> Option { let server_name = &tool_info.server_name; let view = self.servers.get(server_name)?; @@ -399,7 +398,7 @@ impl McpConnectionSet { Arc::clone(self), client, config, - tool_catalog_revision, + tool_catalog_snapshot, tool_info.clone(), view.metadata.clone(), self.plugin_id_for_mcp_server_name(server_name) @@ -532,7 +531,7 @@ impl McpConnectionSet { .await?; let client_tools = managed_client .tool_catalog - .read(|catalog| catalog.tools.clone()) + .read(|catalog| catalog.tools.to_vec()) .await; emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 721865c457..0d10a93e70 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -518,11 +518,13 @@ async fn prepared_call_timeout_includes_trusted_access_lookup() { default_tools_approval_mode: None, tool_approval_modes: HashMap::new(), }; + let client = Arc::new(create_test_managed_client(vec![tool.clone()]).await); + let catalog_snapshot = client.tool_catalog.read(Arc::new).await; let prepared = crate::PreparedMcpCall::new( manager, - Arc::new(create_test_managed_client(vec![tool.clone()]).await), + client, Arc::new(config), - /*catalog_revision*/ 0, + catalog_snapshot, tool, server_metadata, Some("docs@test".to_string()), @@ -5126,7 +5128,9 @@ async fn apps_catalog_broadcast_preserves_running_calls_and_rejects_stale_calls( .with_live_scope("apps".to_string()); let original = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "original")]; let updated = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "updated")]; + store_current_tools(&context, original.clone()); let catalog = ClientToolCatalog::new(original.clone(), context.subscribe()); + let snapshot = catalog.read(Arc::new).await; store_current_tools(&context, original.clone()); assert_eq!( @@ -5136,17 +5140,11 @@ async fn apps_catalog_broadcast_preserves_running_calls_and_rejects_stale_calls( ); let (release, released) = tokio::sync::oneshot::channel::<()>(); - let running = catalog.run_with_revision( - /*expected_revision*/ 0, - || async { released.await.unwrap() }, - ); + let running = catalog.run_with_snapshot(&snapshot, || async { released.await.unwrap() }); tokio::pin!(running); assert!(futures::poll!(&mut running).is_pending()); store_current_tools(&context, updated.clone()); - let stale = catalog.run_with_revision( - /*expected_revision*/ 0, - || async { panic!("stale preparation") }, - ); + let stale = catalog.run_with_snapshot(&snapshot, || async { panic!("stale preparation") }); tokio::pin!(stale); assert!( futures::poll!(&mut stale).is_pending(), @@ -5157,17 +5155,96 @@ async fn apps_catalog_broadcast_preserves_running_calls_and_rejects_stale_calls( assert_eq!(stale.await, None::<()>); assert_eq!( catalog - .read(|catalog| (catalog.revision, catalog.tools.clone())) + .read(|catalog| (catalog.revision, catalog.tools.to_vec())) .await, (1, updated.clone()) ); // A client whose startup finishes late adopts the already-published result. let late = ClientToolCatalog::new(original, context.subscribe()); - assert_eq!(late.read(|catalog| catalog.tools.clone()).await, updated); + assert_eq!(late.read(|catalog| catalog.tools.to_vec()).await, updated); Ok(()) } +#[tokio::test] +async fn apps_catalog_broadcast_restores_equivalent_calls() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + /*account_id*/ None, + /*chatgpt_user_id*/ None, + ) + .with_live_scope("apps".to_string()); + let original = vec![ + create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "first"), + create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "second"), + ]; + store_current_tools(&context, original.clone()); + let catalog = ClientToolCatalog::new(original.clone(), context.subscribe()); + let snapshot = catalog.read(Arc::new).await; + + let mut changed = original.clone(); + changed[0].tool.description = Some("Changed definition".into()); + store_current_tools(&context, changed); + assert_eq!( + catalog + .run_with_snapshot(&snapshot, || async { panic!("changed preparation") }) + .await, + None::<()> + ); + + let mut restored = original; + restored.reverse(); + store_current_tools(&context, restored); + assert_eq!( + catalog + .run_with_snapshot(&snapshot, || async { "prepared" }) + .await, + Some("prepared") + ); + catalog + .refresh( + || async { Ok((snapshot.tools.to_vec(), ())) }, + |tools, ()| store_current_tools(&context, tools.to_vec()), + ) + .await?; + assert_eq!( + catalog + .run_with_snapshot(&snapshot, || async { panic!("refreshed preparation") }) + .await, + None::<()> + ); + Ok(()) +} + +#[test] +fn idle_apps_clients_do_not_retain_replaced_tools() { + let context = ConnectorRuntimeManager::::new_without_cache() + .context( + PathBuf::from("unused"), + ConnectorRuntimeContextKey::personal( + /*account_id*/ None, /*chatgpt_user_id*/ None, + ), + ) + .with_live_scope("apps".into()); + let tool = create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "original"); + let schema = Arc::downgrade(&tool.tool.input_schema); + store_current_tools(&context, vec![tool]); + let clients = [ + ClientToolCatalog::new(Vec::new(), context.subscribe()), + ClientToolCatalog::new(Vec::new(), context.subscribe()), + ]; + store_current_tools( + &context, + vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "updated")], + ); + assert!( + schema.upgrade().is_none(), + "idle clients must not own the old schema" + ); + drop(clients); +} + #[tokio::test] async fn apps_catalog_broadcast_survives_an_older_local_refresh() -> anyhow::Result<()> { let codex_home = tempdir()?; @@ -5177,11 +5254,12 @@ async fn apps_catalog_broadcast_survives_an_older_local_refresh() -> anyhow::Res /*chatgpt_user_id*/ None, ) .with_live_scope("apps".to_string()); + store_current_tools(&context, Vec::new()); let catalog = ClientToolCatalog::new(Vec::new(), context.subscribe()); let older_ticket = context.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh); let newer = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "newer")]; store_current_tools(&context, newer.clone()); - assert_eq!(catalog.read(|catalog| catalog.tools.clone()).await, newer); + assert_eq!(catalog.read(|catalog| catalog.tools.to_vec()).await, newer); catalog .refresh( || async { Ok((Vec::new(), older_ticket)) }, @@ -5194,7 +5272,7 @@ async fn apps_catalog_broadcast_survives_an_older_local_refresh() -> anyhow::Res }, ) .await?; - assert_eq!(catalog.read(|catalog| catalog.tools.clone()).await, newer); + assert_eq!(catalog.read(|catalog| catalog.tools.to_vec()).await, newer); Ok(()) } diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index c0ac4b3a2e..06ea54e861 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -146,7 +146,7 @@ impl ManagedClient { return tools; } } - catalog.tools.clone() + catalog.tools.to_vec() }) .await } @@ -938,6 +938,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) @@ -981,11 +982,30 @@ async fn start_server_task( .as_ref() .and_then(|exp| exp.get(MCP_SANDBOX_STATE_META_CAPABILITY)) .is_some(); + let codex_apps_tools_cache_context = codex_apps_tools_cache_context.map(|context| { + if server_disables_tool_catalog_cache { + context.without_live_scope() + } else { + // Converted tools can inherit these instructions. Server capabilities and the + // negotiated protocol can also differ between otherwise identical connections. + let mut scope = serde_json::json!([ + requested_capabilities, + initialize_result.protocol_version, + initialize_result.capabilities, + initialize_result.instructions, + initialize_result.server_info, + ]); + scope.sort_all_objects(); + context.with_live_scope(scope.to_string()) + } + }); 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(|cache_context| cache_context.begin_fetch(ConnectorRuntimeFetchSource::Startup)); - let client_tools = list_tools_for_client_uncached( + .map(|context| context.begin_fetch(ConnectorRuntimeFetchSource::Startup)); + let tools = list_tools_for_client_uncached( &server_name, is_codex_apps_mcp_server, /*codex_apps_refresh_trigger*/ "initial", @@ -996,25 +1016,25 @@ async fn start_server_task( ) .await .map_err(StartupOutcomeError::from)?; - let server_info = - mcp_server_info_from_implementation(&server_name, initialize_result.server_info); - let shared_tools = match (codex_apps_tools_cache_context.as_ref(), fetch_ticket) { - (Some(cache_context), Some(fetch_ticket)) => cache_context.publish_if_newest_accepted( - fetch_ticket, - &server_info, - client_tools.clone(), - ), - (None, None) => client_tools.clone(), - _ => unreachable!("Codex Apps fetch ticket requires cache context"), - }; - let has_shared_tool_catalog = is_codex_apps_mcp_server || tool_catalog_cache_context.is_some(); + 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()); + tools.into() + } + (Some(context), Some(ticket)) => context + .publish_runtime_if_newest_accepted(ticket, &server_info, tools) + .shared_tools(), + (None, None) => tools.into(), + _ => unreachable!("Codex Apps fetch ticket requires cache context"), + }; if let (Some(cache_context), Some(fetch_ticket)) = ( tool_catalog_cache_context.as_ref(), tool_catalog_fetch_ticket, ) { - cache_context.publish_if_newest(fetch_ticket, &shared_tools); + cache_context.publish_if_newest(fetch_ticket, &client_tools); } - if has_shared_tool_catalog { + if is_codex_apps_mcp_server || tool_catalog_cache_context.is_some() { emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, list_start.elapsed(), diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 09064ebc9a..85cefccf0d 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -10,6 +10,7 @@ use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex; +use std::sync::Weak; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; @@ -139,7 +140,8 @@ fn ensure_host_owned_apps_registration( struct CachedMcpBinding { catalog_revisions: HashMap, - binding: Arc, + // Reuse a frozen binding while a model step or caller still needs it. + binding: Weak, } struct McpReconnectGuard<'a> { @@ -393,8 +395,9 @@ impl McpRuntime { .unwrap_or_else(std::sync::PoisonError::into_inner); if let Some(cached) = cached.as_ref() && &cached.catalog_revisions == catalog_revisions + && let Some(binding) = cached.binding.upgrade() { - return Some(Arc::clone(&cached.binding)); + return Some(binding); } } @@ -423,12 +426,13 @@ impl McpRuntime { .unwrap_or_else(std::sync::PoisonError::into_inner); if let Some(cached) = cached.as_ref() && cached.catalog_revisions == catalog_revisions + && let Some(binding) = cached.binding.upgrade() { - return Some(Arc::clone(&cached.binding)); + return Some(binding); } *cached = Some(CachedMcpBinding { catalog_revisions, - binding: Arc::clone(&binding), + binding: Arc::downgrade(&binding), }); } Some(binding) @@ -995,13 +999,20 @@ mod tests { .contains("reached refreshed call preparation") ); let repeated = McpRuntime::binding_from_published_runtime( - published, + Arc::clone(&published), /*required_servers*/ &[], /*required_plugins*/ &HashSet::new(), ) .await .expect("cached refreshed binding"); assert!(Arc::ptr_eq(&refreshed, &repeated)); + let released = Arc::downgrade(&refreshed); + drop(refreshed); + drop(repeated); + assert!( + released.upgrade().is_none(), + "the runtime cache must not pin an unused binding" + ); Ok(()) } diff --git a/codex-rs/connectors/src/connector_runtime/mod.rs b/codex-rs/connectors/src/connector_runtime/mod.rs index d14b21554c..23858e82d3 100644 --- a/codex-rs/connectors/src/connector_runtime/mod.rs +++ b/codex-rs/connectors/src/connector_runtime/mod.rs @@ -6,12 +6,16 @@ //! 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. +//! Providers expire with their last context. use std::collections::HashMap; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex; +use std::sync::Weak; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use std::time::Duration; @@ -34,15 +38,21 @@ use self::persistence::tools_cache_path; const MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC: &str = "codex.mcp.tools.cache_publish.duration_ms"; -type LiveCatalog = watch::Sender>>>; +/// The current immutable tools for matching discovery inputs. +struct CatalogProvider { + scope: Vec, + updates: watch::Sender>>>, +} /// Values stored in the connector runtime's persisted tool snapshot. /// /// The runtime uses the connector-owned Codex Apps cache layout for every -/// serializable, cloneable payload. -pub trait ConnectorRuntimePayload: Clone + Serialize + DeserializeOwned {} +/// serializable, cloneable payload. Equality determines whether fresh results can +/// retain the previous storage and tool version, so it must include all metadata +/// that affects readers or prepared calls. +pub trait ConnectorRuntimePayload: Clone + PartialEq + Serialize + DeserializeOwned {} -impl ConnectorRuntimePayload for T where T: Clone + Serialize + DeserializeOwned {} +impl ConnectorRuntimePayload for T where T: Clone + PartialEq + Serialize + DeserializeOwned {} /// The account and workspace identity of a connector runtime catalog. #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -96,9 +106,10 @@ pub fn connector_runtime_cache_path(codex_home: &Path, auth: Option<&CodexAuth>) /// intentionally applied by readers rather than persisted in this snapshot. #[derive(Debug, Clone)] pub struct ConnectorRuntimeSnapshot { - tools: Vec, + tools: Arc<[T]>, refreshed_at: SystemTime, generation: u64, + tools_version: u64, } impl ConnectorRuntimeSnapshot { @@ -106,6 +117,15 @@ impl ConnectorRuntimeSnapshot { &self.tools } + pub fn shared_tools(&self) -> Arc<[T]> { + Arc::clone(&self.tools) + } + + /// Advances only when tool definitions change, so equal refreshes preserve prepared calls. + pub fn tools_version(&self) -> u64 { + self.tools_version + } + pub fn refreshed_at(&self) -> SystemTime { self.refreshed_at } @@ -182,7 +202,7 @@ impl ConnectorRuntimeManager { /// Handle to one shared account/workspace connector runtime. pub struct ConnectorRuntimeContext { entry: Arc>, - live_catalog: Option>, + live_catalog: Option>>, } impl Clone for ConnectorRuntimeContext { @@ -196,21 +216,42 @@ impl Clone for ConnectorRuntimeContext { impl ConnectorRuntimeContext { /// Groups executable catalogs by equivalent discovery inputs within this account/home. - /// The caller must include the endpoint, auth configuration and listing protocol in `scope`. + /// Include transport/auth, requested capabilities, and initialization results. + /// Calling this again refines the existing scope; it does not replace it. /// Account-wide discovery reads are unaffected. pub fn with_live_scope(mut self, scope: String) -> Self { - self.live_catalog = Some( - lock_unpoisoned(&self.entry.live_catalogs) - .entry(scope) - .or_insert_with(|| watch::channel(None).0) - .clone(), - ); + let mut scope_parts = self + .live_catalog + .as_ref() + .map(|catalog| catalog.scope.clone()) + .unwrap_or_default(); + scope_parts.push(scope); + let mut catalogs = lock_unpoisoned(&self.entry.live_catalogs); + catalogs.retain(|_, catalog| catalog.strong_count() > 0); + let catalog = catalogs.entry(scope_parts.clone()).or_default(); + self.live_catalog = Some(catalog.upgrade().unwrap_or_else(|| { + let provider = Arc::new(CatalogProvider { + scope: scope_parts, + updates: watch::channel(None).0, + }); + *catalog = Arc::downgrade(&provider); + provider + })); + drop(catalogs); + self + } + + /// Detaches a server that opts out of live catalog sharing, including publication. + pub fn without_live_scope(mut self) -> Self { + self.live_catalog = None; self } /// Subscribes to accepted live tools without refetching or waiting for other clients. pub fn subscribe(&self) -> Option>>>> { - self.live_catalog.as_ref().map(watch::Sender::subscribe) + self.live_catalog + .as_ref() + .map(|catalog| catalog.updates.subscribe()) } pub fn current_snapshot(&self) -> Option>> { @@ -249,7 +290,7 @@ impl ConnectorRuntimeContext { pub fn current_tools(&self) -> Option> { self.current_snapshot() - .map(|snapshot| snapshot.tools.clone()) + .map(|snapshot| snapshot.tools.to_vec()) } pub fn publish_runtime_if_newest_accepted( @@ -283,13 +324,26 @@ impl ConnectorRuntimeContext { ) -> Arc> { let publish_start = Instant::now(); let mut last_accepted_generation = lock_unpoisoned(&self.entry.last_accepted_generation); + let prior = self + .live_catalog + .as_ref() + .and_then(|catalog| catalog.updates.borrow().clone()); + let (tools, tools_version) = match prior { + Some(prior) + if prior.generation < ticket.generation && prior.tools() == tools.as_slice() => + { + (prior.shared_tools(), prior.tools_version) + } + _ => (tools.into(), ticket.generation), + }; let snapshot = Arc::new(ConnectorRuntimeSnapshot { tools, + tools_version, refreshed_at: SystemTime::now(), generation: ticket.generation, }); if let Some(live_catalog) = &self.live_catalog { - live_catalog.send_if_modified(|current| { + live_catalog.updates.send_if_modified(|current| { if current .as_ref() .is_some_and(|current| current.generation >= ticket.generation) @@ -336,7 +390,7 @@ impl ConnectorRuntimeContext { ) -> Vec { self.publish_runtime_if_newest_accepted(ticket, server_info, tools) .tools - .clone() + .to_vec() } } @@ -367,7 +421,7 @@ struct ConnectorRuntimeEntry { current_snapshot: ArcSwapOption>, next_fetch_generation: AtomicU64, last_accepted_generation: Mutex, - live_catalogs: Mutex>>, + live_catalogs: Mutex, Weak>>>, } impl ConnectorRuntimeEntry { diff --git a/codex-rs/connectors/src/connector_runtime/persistence.rs b/codex-rs/connectors/src/connector_runtime/persistence.rs index 34662f7f31..8abed04a74 100644 --- a/codex-rs/connectors/src/connector_runtime/persistence.rs +++ b/codex-rs/connectors/src/connector_runtime/persistence.rs @@ -5,7 +5,6 @@ use std::io::Read; use std::io::Write; use std::path::Path; use std::path::PathBuf; -#[cfg(test)] use std::sync::Arc; use std::time::Instant; use std::time::SystemTime; @@ -65,6 +64,7 @@ pub(crate) fn load_cached_connector_runtime_for_identity { schema_version: u8, - tools: Vec, + tools: Arc<[T]>, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -232,9 +232,10 @@ pub(crate) fn write_cached_codex_apps_tools_for_test( T: ConnectorRuntimePayload, { let snapshot = ConnectorRuntimeSnapshot { - tools: tools.to_vec(), + tools: tools.into(), refreshed_at: SystemTime::now(), generation: 0, + tools_version: 0, }; cache_context .entry @@ -251,7 +252,7 @@ where T: ConnectorRuntimePayload, { load_cached_connector_runtime_for_identity(&cache_context.entry.identity) - .map(|snapshot| snapshot.tools) + .map(|snapshot| snapshot.tools.to_vec()) } #[cfg(test)] @@ -263,9 +264,10 @@ where T: ConnectorRuntimePayload, { let snapshot = ConnectorRuntimeSnapshot { - tools: tools.to_vec(), + tools: tools.into(), refreshed_at: SystemTime::now(), generation: 0, + tools_version: 0, }; write_cached_connector_runtime(cache_context, &snapshot) } diff --git a/codex-rs/connectors/src/connector_runtime/tests.rs b/codex-rs/connectors/src/connector_runtime/tests.rs index 743daa9aac..e4a4996dc3 100644 --- a/codex-rs/connectors/src/connector_runtime/tests.rs +++ b/codex-rs/connectors/src/connector_runtime/tests.rs @@ -630,6 +630,56 @@ fn contexts_for_different_identities_keep_isolated_snapshots() { ); } +#[test] +fn live_provider_reuses_equal_tools_only_within_its_scope() { + let manager = ConnectorRuntimeManager::::new_without_cache(); + let context = manager.context( + PathBuf::from("unused"), + ConnectorRuntimeContextKey::personal( + /*account_id*/ None, /*chatgpt_user_id*/ None, + ), + ); + let first = context + .clone() + .with_live_scope("endpoint".into()) + .with_live_scope("capabilities".into()); + let second = context + .clone() + .with_live_scope("endpoint".into()) + .with_live_scope("capabilities".into()); + let other = context + .clone() + .with_live_scope("other-endpoint".into()) + .with_live_scope("capabilities".into()); + let tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "search")]; + let publish = |context: &ConnectorRuntimeContext| { + context.publish_runtime_if_newest_accepted( + context.begin_fetch(ConnectorRuntimeFetchSource::Startup), + &create_test_server_info("Apps"), + tools.clone(), + ) + }; + let original = publish(&first); + let repeated = publish(&second); + let separate = publish(&other); + assert!(Arc::ptr_eq( + &original.shared_tools(), + &repeated.shared_tools() + )); + assert_eq!(original.tools_version(), repeated.tools_version()); + assert!(!Arc::ptr_eq( + &original.shared_tools(), + &separate.shared_tools() + )); + drop(first); + drop(second); + let restarted = context + .with_live_scope("endpoint".into()) + .with_live_scope("capabilities".into()); + assert!(restarted.current_snapshot().is_some()); + assert!(restarted.subscribe().unwrap().borrow().is_none()); +} + #[test] fn oversized_tools_cache_is_ignored_during_initial_load() { let codex_home = tempdir().expect("tempdir"); diff --git a/codex-rs/core/tests/suite/codex_apps_protocol.rs b/codex-rs/core/tests/suite/codex_apps_protocol.rs index fadc529bae..561cc9b8aa 100644 --- a/codex-rs/core/tests/suite/codex_apps_protocol.rs +++ b/codex-rs/core/tests/suite/codex_apps_protocol.rs @@ -1,15 +1,36 @@ -//! Verifies the hosted Apps protocol choice in standalone Codex sessions. +//! Verifies hosted Apps discovery and prepared-call authority across sessions. +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use codex_core::StartThreadOptions; +use codex_core::TurnInputRequest; use codex_features::Feature; use codex_login::CodexAuth; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::ElicitationAction; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_protocol::user_input::UserInput; +use core_test_support::PathExt; use core_test_support::apps_test_server::AppsTestServer; +use core_test_support::apps_test_server::SEARCH_CALENDAR_CREATE_TOOL; +use core_test_support::apps_test_server::SEARCH_CALENDAR_NAMESPACE; +use core_test_support::apps_test_server::apps_enabled_builder; +use core_test_support::apps_test_server::recorded_apps_tool_call_by_call_id; use core_test_support::responses; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; use core_test_support::wait_for_mcp_server; use pretty_assertions::assert_eq; use serde_json::Value; +use serde_json::json; use test_case::test_case; #[derive(Clone, Copy)] @@ -62,3 +83,117 @@ async fn standalone_codex_apps_respects_protocol_setting( assert_eq!(methods, expected_methods); Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn apps_call_survives_catalog_restoration_while_awaiting_approval() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let tools_available = Arc::new(AtomicBool::new(true)); + let apps = + AppsTestServer::mount_with_tools_available_when(&server, Arc::clone(&tools_available)) + .await?; + let fixture = apps_enabled_builder(apps.chatgpt_base_url) + .with_model_info_override("gpt-5.5", |model| model.supports_search_tool = false) + .with_config(|config| { + config.approvals_reviewer = ApprovalsReviewer::User; + config.config_layer_stack = config + .config_layer_stack + .with_user_config( + &config.codex_home.join("config.toml").abs(), + toml::from_str("[apps.calendar]\ndefault_tools_approval_mode = 'prompt'") + .expect("Apps approval config"), + ) + .expect("test config should allow Apps approvals"); + config + .features + .enable(Feature::ToolCallMcpElicitation) + .expect("test config should allow MCP approvals"); + }) + .build_with_auto_env(&server) + .await?; + wait_for_mcp_server(&fixture.codex, CODEX_APPS_MCP_SERVER_NAME).await?; + let call = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("calendar-call"), + responses::ev_function_call_with_namespace( + "calendar-call", + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_CREATE_TOOL, + &json!({"title": "Lunch", "starts_at": "2026-03-10T12:00:00Z"}).to_string(), + ), + responses::ev_completed("calendar-call"), + ]), + ) + .await; + let completion = responses::mount_sse_once(&server, responses::sse_completed("done")).await; + fixture + .codex + .start_or_steer_turn( + TurnInputRequest::user_input(vec![UserInput::Text { + text: "Use [$calendar](app://calendar) to create a calendar event.".into(), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + approval_policy: Some(AskForApproval::OnRequest), + permission_profile: Some(PermissionProfile::Disabled), + ..Default::default() + }), + ) + .await?; + let EventMsg::ElicitationRequest(approval) = wait_for_event(&fixture.codex, |event| { + matches!( + event, + EventMsg::ElicitationRequest(_) | EventMsg::TurnComplete(_) + ) + }) + .await + else { + anyhow::bail!("the calendar call should wait for approval"); + }; + assert!( + responses::namespace_child_tool( + &call.single_request().body_json(), + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_CREATE_TOOL, + ) + .is_some() + ); + + // A peer removes the tools, then restores them without refreshing the waiting client. + tools_available.store(false, Ordering::SeqCst); + let peer = fixture + .thread_manager + .start_thread(StartThreadOptions::new(fixture.config.clone())) + .await? + .thread; + wait_for_mcp_server(&peer, CODEX_APPS_MCP_SERVER_NAME).await?; + tools_available.store(true, Ordering::SeqCst); + assert!(!peer.refresh_codex_apps_tools().await?.tools.is_empty()); + + fixture + .codex + .submit(Op::ResolveElicitation { + server_name: approval.server_name, + request_id: approval.id, + decision: ElicitationAction::Accept, + content: None, + meta: None, + }) + .await?; + wait_for_event(&fixture.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let executed = recorded_apps_tool_call_by_call_id(&server, "calendar-call").await; + assert_eq!(executed["params"]["name"], "calendar_create_event"); + assert!( + completion + .single_request() + .function_call_output_text("calendar-call") + .is_some() + ); + peer.shutdown_and_wait().await?; + Ok(()) +}