diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs index 4fbe33a433..447cbd62e6 100644 --- a/codex-rs/codex-mcp/src/binding.rs +++ b/codex-rs/codex-mcp/src/binding.rs @@ -19,7 +19,6 @@ use rmcp::model::ReadResourceResult; use rmcp::model::Resource; use rmcp::model::ResourceTemplate; use serde_json::Value as JsonValue; -use tokio::sync::RwLock; use crate::McpConfig; use crate::binding_clients::McpBindingClients; @@ -174,7 +173,6 @@ pub struct PreparedMcpCall { client: Arc, config: Arc, catalog_revision: u64, - catalog_revision_source: Arc>, tool_info: ToolInfo, server_name: String, server_metadata: McpServerMetadata, @@ -192,7 +190,6 @@ impl PreparedMcpCall { client: Arc, config: Arc, catalog_revision: u64, - catalog_revision_source: Arc>, tool_info: ToolInfo, server_metadata: McpServerMetadata, plugin_id: Option, @@ -205,7 +202,6 @@ impl PreparedMcpCall { client, config, catalog_revision, - catalog_revision_source, tool_info, server_name, server_metadata, @@ -305,10 +301,6 @@ 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. /// A caller-supplied timeout can further restrict the server's configured timeout. - #[expect( - clippy::await_holding_invalid_type, - reason = "catalog replacement must remain serialized with call preparation and execution" - )] pub async fn call_with_preparation( &self, requested_timeout: Option, @@ -325,52 +317,52 @@ impl PreparedMcpCall { (server_timeout, requested_timeout) => server_timeout.or(requested_timeout), }; let tool_name = self.tool_info.tool.name.to_string(); - let current_revision = self.catalog_revision_source.read().await; - if *current_revision != self.catalog_revision { - return Err(anyhow::anyhow!( + self.client + .tool_catalog + .run_with_revision(self.catalog_revision, || async { + let (arguments, meta) = prepare().await?; + let timeout_deadline = + effective_timeout.map(|timeout| tokio::time::Instant::now() + timeout); + let add_trusted_access_context = self.connections.add_trusted_access_context( + &self.tool_info, + &self.server_metadata, + arguments.as_ref(), + meta, + ); + let meta = match effective_timeout.zip(timeout_deadline) { + Some((timeout, deadline)) => { + tokio::time::timeout_at(deadline, add_trusted_access_context) + .await + .map_err(|_| { + anyhow::anyhow!("timed out awaiting tools/call after {timeout:.0?}") + })? + } + None => add_trusted_access_context.await, + }; + let remaining_timeout = match effective_timeout.zip(timeout_deadline) { + Some((timeout, deadline)) => { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return Err(anyhow::anyhow!( + "timed out awaiting tools/call after {timeout:.0?}" + )); + } + Some(remaining) + } + None => None, + }; + self.client + .client + .call_tool(tool_name.clone(), arguments, meta, remaining_timeout) + .await + .with_context(|| format!("tool call failed for `{}/{tool_name}`", self.server_name)) + }) + .await + .ok_or_else(|| anyhow::anyhow!( "tool call rejected because the catalog changed after `{}/{tool_name}` was prepared", self.server_name - )); - } - let (arguments, meta) = prepare().await?; - let timeout_deadline = - effective_timeout.map(|timeout| tokio::time::Instant::now() + timeout); - let add_trusted_access_context = self.connections.add_trusted_access_context( - &self.tool_info, - &self.server_metadata, - arguments.as_ref(), - meta, - ); - let meta = match effective_timeout.zip(timeout_deadline) { - Some((timeout, deadline)) => { - tokio::time::timeout_at(deadline, add_trusted_access_context) - .await - .map_err(|_| { - anyhow::anyhow!("timed out awaiting tools/call after {timeout:.0?}") - })? - } - None => add_trusted_access_context.await, - }; - let remaining_timeout = match effective_timeout.zip(timeout_deadline) { - Some((timeout, deadline)) => { - let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); - if remaining.is_zero() { - return Err(anyhow::anyhow!( - "timed out awaiting tools/call after {timeout:.0?}" - )); - } - Some(remaining) - } - None => None, - }; - let result = self - .client - .client - .call_tool(tool_name.clone(), arguments, meta, remaining_timeout) - .await - .with_context(|| format!("tool call failed for `{}/{tool_name}`", self.server_name))?; - drop(current_revision); - Ok(call_tool_result_from_rmcp(result)) + ))? + .map(call_tool_result_from_rmcp) } } diff --git a/codex-rs/codex-mcp/src/binding_tests.rs b/codex-rs/codex-mcp/src/binding_tests.rs index b42f75c53a..9ba610159f 100644 --- a/codex-rs/codex-mcp/src/binding_tests.rs +++ b/codex-rs/codex-mcp/src/binding_tests.rs @@ -22,6 +22,7 @@ use tokio::sync::Notify; use super::McpBinding; use super::PreparedMcpCall; use crate::binding_clients::McpBindingClients; +use crate::client_tool_catalog::ClientToolCatalog; use crate::connection_manager::McpConnectionSet; use crate::rmcp_client::ManagedClient; use crate::server::McpServerMetadata; @@ -46,7 +47,7 @@ impl InProcessTransportFactory for TestInProcessTransportFactory { struct TestStep { step: Arc, client: Arc, - tool_catalog_revision: Arc>, + tool_catalog: Arc, } async fn test_step( @@ -76,6 +77,7 @@ async fn test_step( .await .expect("create in-process MCP client"), ); + let tool_catalog = Arc::new(ClientToolCatalog::new(vec![tool.clone()])); let managed_client = Arc::new(ManagedClient { client: Arc::clone(&client), server_info: McpServerInfo { @@ -86,7 +88,7 @@ async fn test_step( icons: None, website_url: None, }, - tools: vec![tool.clone()], + tool_catalog: Arc::clone(&tool_catalog), tool_timeout: None, server_instructions: None, server_supports_sandbox_state_meta_capability: supports_sandbox_state_meta, @@ -97,7 +99,6 @@ async fn test_step( Arc::clone(&managed_client), )]))); let connections = Arc::new(McpConnectionSet::empty(/*prefix_mcp_tool_names*/ true)); - let tool_catalog_revision = Arc::new(tokio::sync::RwLock::new(0)); let mut config = crate::mcp::tests::test_mcp_config(std::env::temp_dir()); if label == "old" { config.approval_policy = Constrained::allow_any(AskForApproval::Never); @@ -114,7 +115,6 @@ async fn test_step( managed_client, Arc::clone(&config), /*catalog_revision*/ 0, - Arc::clone(&tool_catalog_revision), tool.clone(), McpServerMetadata { environment_id: format!("{label}-environment"), @@ -142,7 +142,7 @@ async fn test_step( calls, )), client, - tool_catalog_revision, + tool_catalog, } } @@ -296,7 +296,13 @@ async fn prepared_call_is_rejected_after_catalog_refresh() { .prepare_call(SERVER_NAME, TOOL_NAME) .expect("step should prepare the advertised tool"); - *step.tool_catalog_revision.write().await += 1; + step.tool_catalog + .refresh( + || async { Ok((step.step.tools().to_vec(), ())) }, + |_, ()| {}, + ) + .await + .expect("refresh tool catalog"); let error = prepared .call( @@ -324,7 +330,13 @@ async fn stale_prepared_call_does_not_run_preparation() { .step .prepare_call(SERVER_NAME, TOOL_NAME) .expect("step should prepare the advertised tool"); - *step.tool_catalog_revision.write().await += 1; + step.tool_catalog + .refresh( + || async { Ok((step.step.tools().to_vec(), ())) }, + |_, ()| {}, + ) + .await + .expect("refresh tool catalog"); let prepared_side_effect_ran = Arc::new(AtomicBool::new(false)); let marker = Arc::clone(&prepared_side_effect_ran); @@ -366,13 +378,20 @@ async fn preparation_holds_catalog_authority_until_it_finishes() { }); preparation_started.notified().await; + let refresh = step.tool_catalog.refresh( + || async { Ok((step.step.tools().to_vec(), ())) }, + |_, ()| {}, + ); + tokio::pin!(refresh); assert!( - step.tool_catalog_revision.try_write().is_err(), + futures::poll!(&mut refresh).is_pending(), "catalog replacement must wait for irreversible call preparation" ); finish_preparation.notify_one(); call.await .expect("call task should finish") .expect_err("the test preparation should stop the call"); - assert!(step.tool_catalog_revision.try_write().is_ok()); + refresh + .await + .expect("catalog refresh should finish after preparation"); } diff --git a/codex-rs/codex-mcp/src/client_tool_catalog.rs b/codex-rs/codex-mcp/src/client_tool_catalog.rs new file mode 100644 index 0000000000..8103e19180 --- /dev/null +++ b/codex-rs/codex-mcp/src/client_tool_catalog.rs @@ -0,0 +1,110 @@ +//! Tool catalog state and synchronization owned by one live MCP client. +//! +//! 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. + +use std::collections::HashSet; +use std::future::Future; +use std::sync::Arc; + +use anyhow::Result; +use tokio::sync::Mutex; +use tokio::sync::RwLock; + +use crate::tools::ToolInfo; + +/// The exact Apps catalog returned by an awaited refresh of one published runtime. +pub struct CodexAppsToolSnapshot { + /// Raw installed tools, including tools hidden or disabled for the model. + pub tools: Vec, + /// Raw MCP tool names allowed by the same runtime's generic MCP policy. + /// App-specific policy is applied by the caller. + pub model_visible_tool_names: HashSet, +} + +pub(crate) struct ClientToolCatalog { + current: RwLock, + /// Serialize fetches without blocking calls against the current catalog. + refresh_lock: Mutex<()>, +} + +pub(crate) struct ToolCatalogSnapshot { + /// Zero is the startup catalog; successful explicit refreshes advance it. + pub(crate) revision: u64, + pub(crate) tools: Vec, +} + +impl ClientToolCatalog { + pub(crate) fn new(tools: Vec) -> Self { + Self { + current: RwLock::new(ToolCatalogSnapshot { revision: 0, tools }), + refresh_lock: Mutex::new(()), + } + } + + pub(crate) async fn read(&self, read: impl FnOnce(&ToolCatalogSnapshot) -> R) -> R { + let current = self.current.read().await; + read(¤t) + } + + /// Serialize fetching and publication, leaving the current catalog usable during the fetch. + /// The publication callback runs alongside the exact-client update under the write lock. + #[expect( + clippy::await_holding_invalid_type, + reason = "refreshes must remain serialized through fetching and catalog publication" + )] + pub(crate) async fn refresh(&self, fetch: F, publish: P) -> Result + where + F: FnOnce() -> Fut, + Fut: Future, C)>>, + P: FnOnce(&[ToolInfo], C) -> R, + { + let _refresh = self.refresh_lock.lock().await; + let (tools, context) = fetch().await?; + let mut current = self.current.write().await; + let result = publish(&tools, context); + current.tools = tools; + current.revision += 1; + Ok(result) + } + + /// Reject stale calls before preparation and hold catalog authority until execution finishes. + #[expect( + clippy::await_holding_invalid_type, + reason = "catalog publication must wait for call preparation and execution" + )] + pub(crate) async fn run_with_revision( + &self, + expected_revision: u64, + run: F, + ) -> Option + where + F: FnOnce() -> Fut, + Fut: Future, + { + let current = self.current.read().await; + if current.revision != expected_revision { + return None; + } + let result = run().await; + drop(current); + Some(result) + } +} + +/// A binding cache key includes client identity, since new clients start at zero. +#[derive(Clone)] +pub(crate) struct ClientToolCatalogRevision { + pub(crate) catalog: Arc, + pub(crate) revision: u64, +} + +impl PartialEq for ClientToolCatalogRevision { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.catalog, &other.catalog) && self.revision == other.revision + } +} + +impl Eq for ClientToolCatalogRevision {} diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index b211631aff..adbac51863 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -73,8 +73,6 @@ use codex_protocol::protocol::McpStartupFailureReason; use codex_protocol::protocol::McpStartupStatus; use codex_protocol::protocol::McpStartupUpdateEvent; use codex_rmcp_client::determine_streamable_http_auth_status_from_credentials; -use tokio::sync::Mutex; -use tokio::sync::RwLock; use tokio::sync::watch; use tokio::task::JoinSet; use tracing::warn; @@ -190,9 +188,6 @@ pub(crate) struct McpConnectionSet { protocol_mode: crate::McpProtocolMode, required_servers: Vec, optional_startup_deadline: OnceLock, - tool_catalog_revision: Arc>, - codex_apps_tools_override: RwLock>>, - codex_apps_refresh_lock: Mutex<()>, tool_plugin_provenance: Arc, prefix_mcp_tool_names: bool, non_prefixed_mcp_tool_servers: Vec, @@ -711,9 +706,6 @@ impl McpConnectionSet { protocol_mode, required_servers, optional_startup_deadline: OnceLock::new(), - tool_catalog_revision: Arc::new(RwLock::new(0)), - codex_apps_tools_override: RwLock::new(None), - codex_apps_refresh_lock: Mutex::new(()), tool_plugin_provenance, prefix_mcp_tool_names, non_prefixed_mcp_tool_servers, @@ -773,9 +765,6 @@ impl McpConnectionSet { protocol_mode: crate::McpProtocolMode::Legacy, required_servers: Vec::new(), optional_startup_deadline: OnceLock::new(), - tool_catalog_revision: Arc::new(RwLock::new(0)), - codex_apps_tools_override: RwLock::new(None), - codex_apps_refresh_lock: Mutex::new(()), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), prefix_mcp_tool_names, non_prefixed_mcp_tool_servers: Vec::new(), @@ -993,4 +982,4 @@ impl McpConnectionSet { #[cfg(test)] #[path = "connection_manager_tests.rs"] -mod tests; +pub(crate) mod tests; 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 d944eec810..5385787b57 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -19,6 +19,8 @@ use super::McpServerMetadata; use crate::binding::McpBinding; use crate::binding::PreparedMcpCall; use crate::binding_clients::McpBindingClients; +use crate::client_tool_catalog::ClientToolCatalogRevision; +use crate::client_tool_catalog::CodexAppsToolSnapshot; 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; @@ -57,7 +59,10 @@ pub fn tool_is_model_visible(tool: &ToolInfo) -> bool { } impl McpConnectionSet { - pub(crate) async fn stable_catalog_revision(&self) -> Option { + pub(crate) async fn stable_catalog_revisions( + &self, + ) -> Option> { + let mut revisions = HashMap::new(); for (server_name, view) in &self.servers { if !view .connection @@ -67,7 +72,7 @@ impl McpConnectionSet { { return None; } - let Some(client) = view.connection.client.ready_transport() else { + let Some(client) = view.connection.client.ready_client() else { if !view.connection.client.is_codex_apps_mcp_server && self.required_servers.binary_search(server_name).is_err() && matches!(view.connection.client.client.peek(), Some(Err(_))) @@ -76,11 +81,19 @@ impl McpConnectionSet { } return None; }; - if client.is_closed().await { + if client.client.is_closed().await { return None; } + let revision = client.tool_catalog.read(|catalog| catalog.revision).await; + revisions.insert( + server_name.clone(), + ClientToolCatalogRevision { + catalog: Arc::clone(&client.tool_catalog), + revision, + }, + ); } - Some(*self.tool_catalog_revision.read().await) + Some(revisions) } /// Returns all tools with model-visible names normalized. @@ -102,30 +115,15 @@ impl McpConnectionSet { .client .startup_complete .load(Ordering::Acquire); - let catalog_override = if server_name == CODEX_APPS_MCP_SERVER_NAME { - self.codex_apps_tools_override.read().await.clone() - } else { - None - }; - let server_tools = async { - match catalog_override { - Some(tools) => { - let tools = filter_tools(tools, &view.tool_filter); - Ok(prepare_codex_apps_tools_for_model( - tools, - &self.tool_plugin_provenance, - )) - } - None => view.listed_tools(&self.tool_plugin_provenance).await, - } - } - .instrument(trace_span!( - "list_tools_for_server", - server_name = %server_name, - has_cached_tools, - startup_complete - )) - .await; + let server_tools = view + .listed_tools(&self.tool_plugin_provenance) + .instrument(trace_span!( + "list_tools_for_server", + server_name = %server_name, + has_cached_tools, + startup_complete + )) + .await; let result = match server_tools { Ok(server_tools) => Ok(server_tools .into_iter() @@ -170,10 +168,6 @@ impl McpConnectionSet { (tools, errors) } - #[expect( - clippy::await_holding_invalid_type, - reason = "catalog capture must remain serialized with catalog replacement" - )] pub(crate) async fn capture_binding_with_metadata( self: &Arc, config: Arc, @@ -181,9 +175,8 @@ impl McpConnectionSet { required_servers: &[String], required_plugins: &HashSet, ) -> McpBinding { - let revision = self.tool_catalog_revision.read().await; let mut listed_tools = Vec::new(); - let mut clients = std::collections::HashMap::new(); + let mut clients = HashMap::new(); let optional_mcp_startup_grace = config.optional_mcp_startup_grace; let server_snapshots = join_all(self.servers.iter().map(|(server_name, view)| async move { if !view @@ -269,13 +262,11 @@ impl McpConnectionSet { return None; }; client.tool_timeout = view.tool_timeout; - let catalog_override = if server_name == CODEX_APPS_MCP_SERVER_NAME { - self.codex_apps_tools_override.read().await.clone() - } else { - None - }; - let server_tools = catalog_override.unwrap_or_else(|| client.tools.clone()); - (Some(Arc::new(client)), server_tools) + let (revision, server_tools) = client + .tool_catalog + .read(|catalog| (catalog.revision, catalog.tools.clone())) + .await; + (Some((Arc::new(client), revision)), server_tools) }; let server_tools = filter_tools(server_tools, &view.tool_filter); let server_tools = if server_name == CODEX_APPS_MCP_SERVER_NAME { @@ -301,12 +292,11 @@ impl McpConnectionSet { })) .await; for (server_name, client, server_tools) in server_results.into_iter().flatten() { - if let Some(client) = client { - clients.insert(server_name, client); + if let Some((client, revision)) = client { + clients.insert(server_name, (client, revision)); } listed_tools.extend(server_tools); } - let clients = Arc::new(McpBindingClients::new(clients)); let listed_tools = normalize_tools_for_model_with_prefix( listed_tools, self.prefix_mcp_tool_names, @@ -316,14 +306,18 @@ 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) = clients.client(&tool_info.server_name) else { + let Some((client, revision)) = clients.get(&tool_info.server_name) else { if model_visible { tools.push(tool_info); } continue; }; - let Some(call) = self.prepare_call(&tool_info, client, Arc::clone(&config), *revision) - else { + let Some(call) = self.prepare_call( + &tool_info, + Arc::clone(client), + Arc::clone(&config), + *revision, + ) else { trace!( server_name = %tool_info.server_name, tool_name = %tool_info.tool.name, @@ -342,6 +336,12 @@ impl McpConnectionSet { tools.push(tool_info); } } + let clients = Arc::new(McpBindingClients::new( + clients + .into_iter() + .map(|(server_name, (client, _))| (server_name, client)) + .collect(), + )); McpBinding::new( Arc::clone(self), clients, @@ -366,7 +366,6 @@ impl McpConnectionSet { client, config, tool_catalog_revision, - Arc::clone(&self.tool_catalog_revision), tool_info.clone(), view.metadata.clone(), self.plugin_id_for_mcp_server_name(server_name) @@ -375,67 +374,48 @@ impl McpConnectionSet { ) } - /// Force-refresh Codex Apps tools and publish one new exact catalog revision. - #[expect( - clippy::await_holding_invalid_type, - reason = "catalog publication must remain serialized with captured tool calls" - )] - pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { - let _refresh = self.codex_apps_refresh_lock.lock().await; + /// Refreshes one exact Apps catalog, preserving the raw inventory for app policy. + pub(crate) async fn refresh_codex_apps_client_catalog( + &self, + config: &crate::McpConfig, + ) -> Result { let refresh_start = Instant::now(); let view = self .servers .get(CODEX_APPS_MCP_SERVER_NAME) .ok_or_else(|| anyhow!("unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"))?; - let managed_client = view - .connection - .client() - .await - .context("failed to get client")?; - - let list_start = Instant::now(); - let fetch_ticket = - managed_client - .codex_apps_tools_cache_context - .as_ref() - .map(|cache_context| { - cache_context.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh) - }); - let client_tools = list_tools_for_client_uncached( - CODEX_APPS_MCP_SERVER_NAME, - /*is_codex_apps_mcp_server*/ true, - /*codex_apps_refresh_trigger*/ "explicit", - &managed_client.client, - view.tool_timeout, - view.catalog_item_limit, - managed_client.server_instructions.as_deref(), - ) - .await - .with_context(|| { - format!("failed to refresh tools for MCP server '{CODEX_APPS_MCP_SERVER_NAME}'") - })?; - - let mut tool_catalog_revision = self.tool_catalog_revision.write().await; - let tools = match ( - managed_client.codex_apps_tools_cache_context.as_ref(), - fetch_ticket, - ) { - (Some(cache_context), Some(fetch_ticket)) => cache_context.publish_if_newest_accepted( - fetch_ticket, - &managed_client.server_info, - client_tools.clone(), - ), - (None, None) => client_tools.clone(), - _ => unreachable!("Codex Apps fetch ticket requires cache context"), - }; - *self.codex_apps_tools_override.write().await = Some(client_tools); - *tool_catalog_revision += 1; - drop(tool_catalog_revision); + let (tools, _) = self.refresh_codex_apps_tool_catalog().await?; + let server_has_permission = config + .permission_profile_for_server(CODEX_APPS_MCP_SERVER_NAME) + .is_some(); + let model_visible_tool_names = tools + .iter() + .filter(|tool| { + server_has_permission + && view.tool_filter.allows(&tool.tool.name) + && tool_is_model_visible(tool) + }) + .map(|tool| tool.tool.name.to_string()) + .collect(); emit_duration( - MCP_TOOLS_LIST_DURATION_METRIC, - list_start.elapsed(), - &[("cache", "miss")], + CODEX_APPS_REFRESH_DURATION_METRIC, + refresh_start.elapsed(), + &[("path", "legacy"), ("trigger", "explicit")], ); + Ok(CodexAppsToolSnapshot { + tools, + model_visible_tool_names, + }) + } + + /// Refreshes Apps tools and returns the prepared shared-cache winner for discovery. + pub async fn refresh_codex_apps_tools_for_discovery(&self) -> Result> { + let refresh_start = Instant::now(); + let view = self + .servers + .get(CODEX_APPS_MCP_SERVER_NAME) + .ok_or_else(|| anyhow!("unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"))?; + let (_, tools) = self.refresh_codex_apps_tool_catalog().await?; let tools = prepare_codex_apps_tools_for_model( filter_tools(tools, &view.tool_filter), &self.tool_plugin_provenance, @@ -455,6 +435,72 @@ impl McpConnectionSet { Ok(tools) } + /// Publishes the exact client catalog and returns both raw inventories. + async fn refresh_codex_apps_tool_catalog(&self) -> Result<(Vec, Vec)> { + let view = self + .servers + .get(CODEX_APPS_MCP_SERVER_NAME) + .ok_or_else(|| anyhow!("unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"))?; + let managed_client = view + .connection + .client() + .await + .context("failed to get client")?; + let (client_tools, tools, list_start) = managed_client + .tool_catalog + .refresh( + || async { + let list_start = Instant::now(); + let fetch_ticket = managed_client.codex_apps_tools_cache_context.as_ref().map( + |cache_context| { + cache_context.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh) + }, + ); + let client_tools = list_tools_for_client_uncached( + CODEX_APPS_MCP_SERVER_NAME, + /*is_codex_apps_mcp_server*/ true, + /*codex_apps_refresh_trigger*/ "explicit", + &managed_client.client, + view.tool_timeout, + view.catalog_item_limit, + managed_client.server_instructions.as_deref(), + ) + .await + .with_context(|| { + format!( + "failed to refresh tools for MCP server '{CODEX_APPS_MCP_SERVER_NAME}'" + ) + })?; + Ok((client_tools, (fetch_ticket, list_start))) + }, + |client_tools, (fetch_ticket, list_start)| { + // Discovery may accept another client's newer fetch. The catalog + // retains the tools fetched through this exact connection. + let tools = match ( + managed_client.codex_apps_tools_cache_context.as_ref(), + fetch_ticket, + ) { + (Some(cache_context), Some(fetch_ticket)) => cache_context + .publish_if_newest_accepted( + fetch_ticket, + &managed_client.server_info, + client_tools.to_vec(), + ), + (None, None) => client_tools.to_vec(), + _ => unreachable!("Codex Apps fetch ticket requires cache context"), + }; + (client_tools.to_vec(), tools, list_start) + }, + ) + .await?; + emit_duration( + MCP_TOOLS_LIST_DURATION_METRIC, + list_start.elapsed(), + &[("cache", "miss")], + ); + Ok((client_tools, tools)) + } + fn with_server_metadata(mut tool: ToolInfo, metadata: &McpServerMetadata) -> ToolInfo { tool.supports_parallel_tool_calls = metadata.supports_parallel_tool_calls; tool.server_origin = metadata diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 0c0e57306e..bcbfb1dd4a 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1,5 +1,6 @@ use super::*; use crate::McpBinding; +use crate::client_tool_catalog::ClientToolCatalog; use crate::elicitation::ElicitationLifecycle; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationRequestRouter; @@ -108,9 +109,6 @@ impl McpConnectionSet { protocol_mode: crate::McpProtocolMode::Legacy, required_servers: Vec::new(), optional_startup_deadline: OnceLock::new(), - tool_catalog_revision: Arc::new(RwLock::new(0)), - codex_apps_tools_override: RwLock::new(None), - codex_apps_refresh_lock: Mutex::new(()), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), prefix_mcp_tool_names, non_prefixed_mcp_tool_servers: Vec::new(), @@ -459,7 +457,7 @@ async fn create_test_managed_client(tools: Vec) -> ManagedClient { .expect("create in-process RMCP client"), ), server_info: create_test_server_info("Ready"), - tools, + tool_catalog: Arc::new(ClientToolCatalog::new(tools)), tool_timeout: None, server_instructions: None, server_supports_sandbox_state_meta_capability: false, @@ -520,7 +518,6 @@ async fn prepared_call_timeout_includes_trusted_access_lookup() { Arc::new(create_test_managed_client(vec![tool.clone()]).await), Arc::new(config), /*catalog_revision*/ 0, - Arc::new(RwLock::new(0)), tool, server_metadata, Some("docs@test".to_string()), @@ -727,7 +724,7 @@ fn create_gated_async_managed_client( ) } -async fn create_test_manager_with_ready_apps_client( +pub(crate) async fn create_test_manager_with_ready_apps_client( cache_context: ConnectorRuntimeContext, tool_name: &str, list_started: Option>, @@ -759,7 +756,7 @@ async fn create_test_manager_with_ready_apps_client( let managed_client = ManagedClient { client, server_info: create_test_server_info("Codex Apps"), - tools: vec![tool], + tool_catalog: Arc::new(ClientToolCatalog::new(vec![tool])), tool_timeout: Some(Duration::from_secs(5)), server_instructions: None, server_supports_sandbox_state_meta_capability: false, @@ -2168,8 +2165,8 @@ async fn capture_binding_uses_the_ready_clients_own_tools() { } #[tokio::test] -async fn hard_refresh_keeps_binding_override_local_when_shared_cache_loses_race() --> anyhow::Result<()> { +async fn hard_refresh_keeps_client_catalog_local_when_shared_cache_loses_race() -> anyhow::Result<()> +{ let codex_home = tempdir()?; let shared_cache = ConnectorRuntimeManager::::default(); let cache_key = ConnectorRuntimeContextKey::personal( @@ -2187,7 +2184,7 @@ async fn hard_refresh_keeps_binding_override_local_when_shared_cache_loses_race( Some(Arc::clone(&release_list)), ) .await?; - let manager_b = create_test_manager_with_ready_apps_client( + let mut manager_b = create_test_manager_with_ready_apps_client( cache_context_b, "b_only", /*list_started*/ None, @@ -2198,11 +2195,11 @@ async fn hard_refresh_keeps_binding_override_local_when_shared_cache_loses_race( let manager_a_for_refresh = Arc::clone(&manager_a); let refresh_a = tokio::spawn(async move { manager_a_for_refresh - .hard_refresh_codex_apps_tools_cache() + .refresh_codex_apps_tools_for_discovery() .await }); list_started.notified().await; - let tools_b = manager_b.hard_refresh_codex_apps_tools_cache().await?; + let tools_b = manager_b.refresh_codex_apps_tools_for_discovery().await?; release_list.notify_one(); let tools_a = refresh_a.await??; @@ -2247,6 +2244,64 @@ async fn hard_refresh_keeps_binding_override_local_when_shared_cache_loses_race( .collect::>(), vec!["b_only"] ); + + let mut config = crate::mcp::tests::test_mcp_config(codex_home.path().to_path_buf()); + config.server_permission_profiles.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + PermissionProfile::default(), + ); + let manager_a_for_refresh = Arc::clone(&manager_a); + let config_for_refresh = config.clone(); + let refresh_a = tokio::spawn(async move { + manager_a_for_refresh + .refresh_codex_apps_client_catalog(&config_for_refresh) + .await + }); + list_started.notified().await; + manager_b.refresh_codex_apps_tools_for_discovery().await?; + release_list.notify_one(); + let snapshot_a = refresh_a.await??; + assert_eq!( + snapshot_a + .tools + .iter() + .map(|tool| tool.tool.name.as_ref()) + .collect::>(), + vec!["a_only"] + ); + assert_eq!( + snapshot_a.model_visible_tool_names, + HashSet::from(["a_only".to_string()]) + ); + assert_eq!( + cache_context_a + .current_tools() + .expect("shared cache tools") + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["b_only"] + ); + Arc::get_mut(&mut manager_b) + .expect("unshared manager") + .servers + .get_mut(CODEX_APPS_MCP_SERVER_NAME) + .expect("Apps server") + .tool_filter + .disabled + .insert("b_only".to_string()); + let snapshot_b = manager_b.refresh_codex_apps_client_catalog(&config).await?; + assert_eq!( + ( + snapshot_b + .tools + .iter() + .map(|tool| tool.tool.name.as_ref()) + .collect::>(), + snapshot_b.model_visible_tool_names, + ), + (vec!["b_only"], HashSet::new()) + ); Ok(()) } @@ -2365,7 +2420,7 @@ fn tool_catalog_cache_bypasses_http_headers_helpers() { let cache = McpToolCatalogCache::default(); let runtime_context = reusable_server_runtime_context(); let mut config = reusable_server_config("https://example.com/mcp"); - let identity = reusable_server_identity(&config, &runtime_context); + let identity = reusable_server_identity("docs", &config, &runtime_context); let context = |config: &McpServerConfig, identity: &McpServerConnectionIdentity| { cache.context( "docs", @@ -2393,7 +2448,7 @@ fn tool_catalog_cache_bypasses_http_headers_helpers() { unreachable!("expected HTTP transport"); }; *http_headers_helper = Some("auth-cli headers".to_string()); - let identity = reusable_server_identity(&config, &runtime_context); + let identity = reusable_server_identity("docs", &config, &runtime_context); assert!(context(&config, &identity).is_none()); } @@ -2634,7 +2689,7 @@ async fn capture_binding_skips_pending_optional_servers_after_configured_shared_ required_manager.required_servers = vec!["pending-selected".to_string()]; let manager = Arc::new(manager); - assert_eq!(manager.stable_catalog_revision().await, None); + assert!(manager.stable_catalog_revisions().await.is_none()); let started = tokio::time::Instant::now(); let binding = tokio::time::timeout( Duration::from_millis(500), @@ -2767,7 +2822,7 @@ async fn capture_binding_waits_for_optional_startup_when_shared_grace_is_disable } #[tokio::test] -async fn stable_catalog_revision_ignores_terminal_optional_server_failures() { +async fn stable_catalog_revisions_ignore_terminal_optional_server_failures() { let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionSet::new_uninitialized( @@ -2790,9 +2845,9 @@ async fn stable_catalog_revision_ignores_terminal_optional_server_failures() { assert!(failed.client().await.is_err()); manager.insert_test_client("failed", failed); - assert_eq!(manager.stable_catalog_revision().await, Some(0)); + assert!(manager.stable_catalog_revisions().await.is_some()); manager.required_servers.push("failed".to_string()); - assert_eq!(manager.stable_catalog_revision().await, None); + assert!(manager.stable_catalog_revisions().await.is_none()); manager.required_servers.clear(); let binding = capture_binding(&Arc::new(manager)).await; @@ -3496,7 +3551,10 @@ async fn cancelling_startup_does_not_disable_a_ready_client() { .await .expect("startup cancellation should not disable a ready client"); assert_eq!( - model_tool_names(&managed.tools), + managed + .tool_catalog + .read(|catalog| model_tool_names(&catalog.tools)) + .await, HashSet::from([ToolName::namespaced("ready", "search")]) ); } @@ -3740,7 +3798,7 @@ async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client() }; let manager = Arc::new(manager); - assert_eq!(manager.stable_catalog_revision().await, None); + assert!(manager.stable_catalog_revisions().await.is_none()); let reconnect_finished_wait = reconnect_finished.notified(); let tools = manager.list_all_tools().await; assert!(tools.is_empty()); @@ -3755,7 +3813,7 @@ async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client() vec!["drive_search"] ); assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); - assert_eq!(manager.stable_catalog_revision().await, Some(0)); + assert!(manager.stable_catalog_revisions().await.is_some()); let step = capture_binding(&manager).await; let prepared = step @@ -4727,13 +4785,14 @@ fn reusable_server_runtime_context() -> McpRuntimeContext { } fn reusable_server_identity( + server_name: &str, config: &McpServerConfig, runtime_context: &McpRuntimeContext, ) -> McpServerConnectionIdentity { let server = EffectiveMcpServer::configured(config.clone()); - let resolved_environment = runtime_context.resolve_server_environment("docs", config); + let resolved_environment = runtime_context.resolve_server_environment(server_name, config); McpServerConnectionIdentity::new( - "docs", + server_name, &server, /*host_plugin_root*/ None, OAuthCredentialsStoreMode::default(), @@ -4767,7 +4826,7 @@ async fn manager_with_reusable_ready_server( "docs".to_string(), McpServerView { connection: Arc::new(McpServerConnection { - identity: Some(reusable_server_identity(config, runtime_context)), + identity: Some(reusable_server_identity("docs", config, runtime_context)), client: create_ready_async_managed_client(tools).await, startup_timeout: config .startup_timeout_sec @@ -4792,6 +4851,7 @@ async fn reconcile_reusable_server( let codex_home = tempdir().expect("tempdir"); reconcile_reusable_server_with_mcp_config( previous, + "docs", config, runtime_context, crate::mcp::tests::test_mcp_config(codex_home.path().to_path_buf()), @@ -4801,6 +4861,7 @@ async fn reconcile_reusable_server( async fn reconcile_reusable_server_with_mcp_config( previous: &McpConnectionSet, + server_name: &str, config: McpServerConfig, runtime_context: McpRuntimeContext, mcp_config: crate::McpConfig, @@ -4815,7 +4876,7 @@ async fn reconcile_reusable_server_with_mcp_config( plugins_available: false, ready_selected_capability_roots: Vec::new(), mcp_servers: HashMap::from([( - "docs".to_string(), + server_name.to_string(), EffectiveMcpServer::configured(config), )]), submit_id: "refresh".to_string(), @@ -4838,6 +4899,97 @@ async fn reconcile_reusable_server_with_mcp_config( .await } +#[tokio::test] +async fn refreshed_catalog_follows_reused_client_without_mutating_old_bindings() +-> anyhow::Result<()> { + let codex_home = tempdir()?; + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + /*account_id*/ None, + /*chatgpt_user_id*/ None, + ); + let runtime_context = reusable_server_runtime_context(); + let config = reusable_server_config("http://127.0.0.1:1"); + let mut previous = create_test_manager_with_ready_apps_client( + cache_context, + "refreshed", + /*list_started*/ None, + /*release_list*/ None, + ) + .await?; + let manager = Arc::get_mut(&mut previous).expect("unshared manager"); + manager.insert_test_client( + "docs", + create_ready_async_managed_client(vec![create_test_tool("docs", "unrelated")]).await, + ); + let connection = Arc::get_mut( + &mut manager + .servers + .get_mut(CODEX_APPS_MCP_SERVER_NAME) + .expect("Apps server") + .connection, + ) + .expect("unshared Apps connection"); + connection.identity = Some(reusable_server_identity( + CODEX_APPS_MCP_SERVER_NAME, + &config, + &runtime_context, + )); + let client = connection.client().await?; + // Seed an earlier catalog before capturing the binding under test. + let startup_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "startup")]; + client + .tool_catalog + .refresh(|| async { Ok((startup_tools, ())) }, |_, ()| {}) + .await?; + let old_binding = capture_binding(&previous).await; + let old_tools = serde_json::to_value(old_binding.tools())?; + let old_call = old_binding + .prepare_call(CODEX_APPS_MCP_SERVER_NAME, "startup") + .expect("startup call"); + let unrelated_call = old_binding + .prepare_call("docs", "unrelated") + .expect("unrelated call"); + + previous.refresh_codex_apps_tools_for_discovery().await?; + let republished = Arc::new( + reconcile_reusable_server_with_mcp_config( + &previous, + CODEX_APPS_MCP_SERVER_NAME, + config, + runtime_context, + crate::mcp::tests::test_mcp_config(codex_home.path().to_path_buf()), + ) + .await, + ); + assert!(previous.shares_test_connection_with(&republished, CODEX_APPS_MCP_SERVER_NAME)); + assert_eq!( + model_tool_names(capture_binding(&republished).await.tools()), + HashSet::from([ToolName::namespaced("mcp__codex_apps", "refreshed")]), + ); + assert_eq!(serde_json::to_value(old_binding.tools())?, old_tools); + + let error = old_call + .call_with_preparation(/*requested_timeout*/ None, || async { + panic!("stale call preparation must not run"); + }) + .await + .expect_err("a call from the old catalog must be rejected"); + assert!(error.to_string().contains("catalog changed")); + let error = unrelated_call + .call_with_preparation(/*requested_timeout*/ None, || async { + Err(anyhow!("unrelated preparation reached")) + }) + .await + .expect_err("stop before the unrelated tool executes"); + assert!( + error.to_string().contains("unrelated preparation reached"), + "Apps refresh must not invalidate another client's calls" + ); + + Ok(()) +} + #[tokio::test] async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> anyhow::Result<()> { let tools = Arc::new(tokio::sync::RwLock::new(vec![Tool::new( @@ -4888,7 +5040,7 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a let managed_client = ManagedClient { client, server_info: create_test_server_info("Mutable tools"), - tools: initial_tools, + tool_catalog: Arc::new(ClientToolCatalog::new(initial_tools)), tool_timeout: None, server_instructions: initialize.instructions, server_supports_sandbox_state_meta_capability: false, @@ -4908,7 +5060,7 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a "docs".to_string(), McpServerView { connection: Arc::new(McpServerConnection { - identity: Some(reusable_server_identity(&config, &runtime_context)), + identity: Some(reusable_server_identity("docs", &config, &runtime_context)), client: AsyncManagedClient { client: futures::future::ready(Ok(managed_client)).boxed().shared(), is_codex_apps_mcp_server: false, @@ -5471,7 +5623,7 @@ async fn reconciliation_replaces_closed_connections() -> anyhow::Result<()> { let mut connected_client = view.connection.client().await?; connected_client.client = Arc::clone(&client); view.connection = Arc::new(McpServerConnection { - identity: Some(reusable_server_identity(&config, &runtime_context)), + identity: Some(reusable_server_identity("docs", &config, &runtime_context)), client: AsyncManagedClient { client: futures::future::ready(Ok(connected_client)) .boxed() @@ -5586,6 +5738,7 @@ async fn reconciliation_reconnects_when_host_plugin_root_changes() { let unchanged = reconcile_reusable_server_with_mcp_config( &previous, + "docs", server_config.clone(), runtime_context.clone(), config_for_root(original_root), @@ -5596,6 +5749,7 @@ async fn reconciliation_reconnects_when_host_plugin_root_changes() { let replacement_config = config_for_root(replacement_root); let replacement = reconcile_reusable_server_with_mcp_config( &unchanged, + "docs", server_config, runtime_context, replacement_config, diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 32e80e099f..bbf52e638e 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -1,6 +1,7 @@ pub use binding::McpBinding; pub use binding::PreparedMcpCall; pub use client_capabilities::client_mcp_extensions; +pub use client_tool_catalog::CodexAppsToolSnapshot; pub use codex_rmcp_client::McpProtocolMode; pub use connection_manager::tool_is_model_visible; pub use elicitation::ElicitationLifecycle; @@ -100,6 +101,7 @@ mod binding; pub(crate) mod binding_clients; mod catalog; mod client_capabilities; +mod client_tool_catalog; pub(crate) mod codex_apps; pub(crate) mod connection_manager; pub(crate) mod elicitation; diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 9a1cb073ba..8f12ad8de9 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -21,6 +21,7 @@ use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; +use crate::client_tool_catalog::ClientToolCatalog; use crate::codex_apps::normalize_codex_apps_callable_name; use crate::codex_apps::normalize_codex_apps_callable_namespace; use crate::codex_apps::normalize_codex_apps_tool_title; @@ -113,7 +114,7 @@ const UNTRUSTED_CONNECTOR_META_KEYS: &[&str] = &[ pub(crate) struct ManagedClient { pub(crate) client: Arc, pub(crate) server_info: McpServerInfo, - pub(crate) tools: Vec, + pub(crate) tool_catalog: Arc, pub(crate) tool_timeout: Option, pub(crate) server_instructions: Option, pub(crate) server_supports_sandbox_state_meta_capability: bool, @@ -121,30 +122,28 @@ pub(crate) struct ManagedClient { } impl ManagedClient { - pub(crate) fn listed_tools(&self) -> Vec { + pub(crate) async fn listed_tools(&self) -> Vec { let total_start = Instant::now(); - if let Some(tools) = self - .codex_apps_tools_cache_context - .as_ref() - .and_then(ConnectorRuntimeContext::current_tools) - { - emit_duration( - MCP_TOOLS_LIST_DURATION_METRIC, - total_start.elapsed(), - &[("cache", "hit")], - ); - return tools; - } - - if self.codex_apps_tools_cache_context.is_some() { - emit_duration( - MCP_TOOLS_LIST_DURATION_METRIC, - total_start.elapsed(), - &[("cache", "miss")], - ); - } - - self.tools.clone() + self.tool_catalog + .read(|catalog| { + // Discovery may use the shared cache until this client is refreshed. + // Executable bindings always capture this client's own catalog. + if catalog.revision == 0 + && let Some(cache_context) = &self.codex_apps_tools_cache_context + { + let tools = cache_context.current_tools(); + emit_duration( + MCP_TOOLS_LIST_DURATION_METRIC, + total_start.elapsed(), + &[("cache", if tools.is_some() { "hit" } else { "miss" })], + ); + if let Some(tools) = tools { + return tools; + } + } + catalog.tools.clone() + }) + .await } } @@ -585,7 +584,7 @@ impl AsyncManagedClient { Ok(startup_tools) } else { match self.client().await { - Ok(client) => Ok(client.listed_tools()), + Ok(client) => Ok(client.listed_tools().await), Err(error) if self.is_codex_apps_mcp_server => self.cached_tools().ok_or(error), Err(error) => Err(error), } @@ -976,7 +975,7 @@ async fn start_server_task( let managed = ManagedClient { client: Arc::clone(&client), server_info, - tools: client_tools, + tool_catalog: Arc::new(ClientToolCatalog::new(client_tools)), tool_timeout: None, server_instructions: initialize_result.instructions, server_supports_sandbox_state_meta_capability, diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 7a55d2cc52..45f2e97237 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -46,6 +46,8 @@ use tokio_util::sync::CancellationToken; use crate::McpConfig; use crate::binding::McpBinding; +use crate::client_tool_catalog::ClientToolCatalogRevision; +use crate::client_tool_catalog::CodexAppsToolSnapshot; use crate::connection_manager::McpConnectionSet; use crate::elicitation::ElicitationLifecycle; use crate::elicitation::ElicitationRequestRouter; @@ -117,7 +119,7 @@ struct PublishedMcpRuntime { } struct CachedMcpBinding { - catalog_revision: u64, + catalog_revisions: HashMap, binding: Arc, } @@ -361,14 +363,14 @@ impl McpRuntime { required_plugins: &HashSet, ) -> Option> { let config = Arc::clone(current.config.as_ref()?); - let stable_catalog_revision = current.connections.stable_catalog_revision().await; - if let Some(catalog_revision) = stable_catalog_revision { + let stable_catalog_revisions = current.connections.stable_catalog_revisions().await; + if let Some(catalog_revisions) = &stable_catalog_revisions { let cached = current .cached_binding .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if let Some(cached) = cached.as_ref() - && cached.catalog_revision == catalog_revision + && &cached.catalog_revisions == catalog_revisions { return Some(Arc::clone(&cached.binding)); } @@ -385,20 +387,25 @@ impl McpRuntime { ) .await, ); - if let Some(catalog_revision) = stable_catalog_revision - && current.connections.stable_catalog_revision().await == Some(catalog_revision) + if let Some(catalog_revisions) = stable_catalog_revisions + && current + .connections + .stable_catalog_revisions() + .await + .as_ref() + == Some(&catalog_revisions) { let mut cached = current .cached_binding .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if let Some(cached) = cached.as_ref() - && cached.catalog_revision == catalog_revision + && cached.catalog_revisions == catalog_revisions { return Some(Arc::clone(&cached.binding)); } *cached = Some(CachedMcpBinding { - catalog_revision, + catalog_revisions, binding: Arc::clone(&binding), }); } @@ -520,7 +527,20 @@ impl McpRuntime { &self, ) -> anyhow::Result> { self.latest_connections() - .hard_refresh_codex_apps_tools_cache() + .refresh_codex_apps_tools_for_discovery() + .await + } + + /// Refreshes the published Apps client and returns its exact inventory and MCP eligibility. + pub async fn refresh_codex_apps_tools(&self) -> anyhow::Result { + let current = self.current.load_full(); + let config = current + .config + .as_deref() + .ok_or_else(|| anyhow::anyhow!("MCP runtime is not configured"))?; + current + .connections + .refresh_codex_apps_client_catalog(config) .await } @@ -880,6 +900,91 @@ mod tests { assert!(!gate.wait().await); } + #[tokio::test] + async fn cached_bindings_follow_the_clients_catalog_revision() -> anyhow::Result<()> { + let codex_home = tempfile::tempdir()?; + let cache_context = ConnectorRuntimeManager::::default().context( + codex_home.path().to_path_buf(), + ConnectorRuntimeContextKey::personal( + /*account_id*/ None, /*chatgpt_user_id*/ None, + ), + ); + let connections = + crate::connection_manager::tests::create_test_manager_with_ready_apps_client( + cache_context, + "search", + /*list_started*/ None, + /*release_list*/ None, + ) + .await?; + // Complete the fixture's shared startup future before testing stable reuse. + connections.list_all_tools().await; + let mut config = crate::mcp::tests::test_mcp_config(codex_home.path().to_path_buf()); + config.server_permission_profiles.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + PermissionProfile::default(), + ); + let published = Arc::new(PublishedMcpRuntime { + connections: Arc::clone(&connections), + config: Some(Arc::new(config)), + auth: None, + auth_token: None, + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + selected_environments: HashMap::new(), + cached_binding: Mutex::new(None), + }); + let before = McpRuntime::binding_from_published_runtime( + Arc::clone(&published), + /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), + ) + .await + .expect("initial binding"); + let repeated = McpRuntime::binding_from_published_runtime( + Arc::clone(&published), + /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), + ) + .await + .expect("cached initial binding"); + assert!(Arc::ptr_eq(&before, &repeated)); + + connections.refresh_codex_apps_tools_for_discovery().await?; + + let refreshed = McpRuntime::binding_from_published_runtime( + Arc::clone(&published), + /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), + ) + .await + .expect("refreshed binding"); + assert!(!Arc::ptr_eq(&before, &refreshed)); + let call = refreshed + .prepare_call(CODEX_APPS_MCP_SERVER_NAME, "search") + .expect("refreshed call"); + let error = call + .call_with_preparation(/*requested_timeout*/ None, || async { + Err(anyhow::anyhow!("reached refreshed call preparation")) + }) + .await + .expect_err("stop before dispatch"); + assert!( + error + .to_string() + .contains("reached refreshed call preparation") + ); + let repeated = McpRuntime::binding_from_published_runtime( + published, + /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), + ) + .await + .expect("cached refreshed binding"); + assert!(Arc::ptr_eq(&refreshed, &repeated)); + Ok(()) + } + #[tokio::test] async fn cached_bindings_are_scoped_to_the_published_runtime() { let published = Arc::new(PublishedMcpRuntime { diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 7772ce985e..ae02620448 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -839,6 +839,13 @@ impl CodexThread { self.session.refresh_mcp_config(next_config).await; } + /// Refreshes this thread's Apps tools before returning their runtime state. + pub async fn refresh_codex_apps_tools( + &self, + ) -> anyhow::Result { + self.session.refresh_codex_apps_tools().await + } + pub async fn environment_selections(&self) -> Vec { self.session.services.turn_environments.selections() } diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index fb79de7d7b..bac8c641cb 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -255,6 +255,21 @@ impl Session { } } + /// Refreshes Apps tools on the published thread runtime and returns that client's snapshot. + pub(crate) async fn refresh_codex_apps_tools( + self: &Arc, + ) -> anyhow::Result { + // Reconcile unchanged config so failed or closed clients can be replaced. + self.mark_mcp_runtime_dirty(); + self.refresh_mcp_if_dirty().await; + let _refresh = self + .mcp_refresh + .acquire() + .await + .map_err(|_| anyhow::anyhow!("MCP runtime refresh semaphore closed"))?; + self.services.mcp_runtime.refresh_codex_apps_tools().await + } + /// Reconnects the runtime so refreshed Apps tools belong to their new exact client. pub(crate) async fn hard_refresh_latest_codex_apps_tools( self: &Arc,