From 65f8bf68533332628b7fc213eade2a91d18d36ee Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 21 Jul 2026 18:09:06 +0000 Subject: [PATCH] Bind MCP calls to captured catalog revisions (#34588) ## Why An MCP connection or tool catalog can change after a model step captures its available tools. Calls from that step must not be rerouted to a replacement client or run against a catalog revision the model did not see. ## What changed - Add `McpBinding` to capture the ready clients, visible tools, resources, and server metadata for one sampling request. - Prepare calls against the captured client and reject them if the tool catalog changes before call preparation and execution finish. - Keep resource access and cache identity tied to the binding's exact connection set, and omit servers that do not have a ready client. - Preserve each client's unfiltered tool catalog so bindings can apply the current filter consistently. ## Testing Add coverage for connection replacement, stale catalog rejection, serialized call preparation and refresh, ready-client capture, and binding-scoped resource identity. GitOrigin-RevId: a25da4d2c7f5ed06c41caa0c6d15c464c1698314 --- codex-rs/codex-mcp/src/binding.rs | 294 +++++++++++ codex-rs/codex-mcp/src/binding_clients.rs | 92 +++- codex-rs/codex-mcp/src/binding_tests.rs | 351 +++++++++++++ codex-rs/codex-mcp/src/connection_manager.rs | 19 + .../src/connection_manager/tool_catalog.rs | 128 ++++- .../codex-mcp/src/connection_manager_tests.rs | 490 +++++++++++++++++- codex-rs/codex-mcp/src/lib.rs | 3 + codex-rs/codex-mcp/src/mcp/mod.rs | 2 +- codex-rs/codex-mcp/src/mcp/mod_tests.rs | 2 +- codex-rs/codex-mcp/src/resource_client.rs | 99 +++- codex-rs/codex-mcp/src/rmcp_client.rs | 66 ++- 11 files changed, 1495 insertions(+), 51 deletions(-) create mode 100644 codex-rs/codex-mcp/src/binding.rs create mode 100644 codex-rs/codex-mcp/src/binding_tests.rs diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs new file mode 100644 index 0000000000..da94379ce6 --- /dev/null +++ b/codex-rs/codex-mcp/src/binding.rs @@ -0,0 +1,294 @@ +//! Immutable MCP state bound to one model sampling request. + +use std::collections::HashMap; +use std::fmt; +use std::future::Future; +use std::sync::Arc; + +use anyhow::Context; +use anyhow::Result; +use codex_config::AppToolApproval; +use codex_protocol::mcp::CallToolResult; +use rmcp::model::ListResourceTemplatesResult; +use rmcp::model::ListResourcesResult; +use rmcp::model::PaginatedRequestParams; +use rmcp::model::ReadResourceRequestParams; +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; +use crate::connection_manager::McpConnectionManager; +use crate::resource_client::McpResourceClient; +use crate::rmcp_client::ManagedClient; +use crate::server::McpServerMetadata; +use crate::tools::ToolInfo; + +/// The exact tool catalog and execution handles for one model sampling request. +pub struct McpBinding { + connections: Arc, + clients: Arc, + config: Arc, + plugins_available: bool, + tools: Vec, + calls: HashMap<(String, String), PreparedMcpCall>, +} + +impl McpBinding { + /// Creates an empty binding for tests and callers without a materialized runtime. + pub fn empty(config: Arc) -> Self { + Self::new( + Arc::new(McpConnectionManager::empty(config.prefix_mcp_tool_names)), + Arc::new(McpBindingClients::new(HashMap::new())), + config, + /*plugins_available*/ false, + Vec::new(), + HashMap::new(), + ) + } + + pub(crate) fn new( + connections: Arc, + clients: Arc, + config: Arc, + plugins_available: bool, + tools: Vec, + calls: HashMap<(String, String), PreparedMcpCall>, + ) -> Self { + Self { + connections, + clients, + config, + plugins_available, + tools, + calls, + } + } + + pub fn config(&self) -> &Arc { + &self.config + } + + pub fn plugins_available(&self) -> bool { + self.plugins_available + } + + /// Returns the frozen catalog advertised for this sampling request. + pub fn tools(&self) -> &[ToolInfo] { + &self.tools + } + + /// Binds a call to the exact client and metadata advertised by this binding. + pub fn prepare_call(&self, server: &str, tool: &str) -> Option { + self.calls + .get(&(server.to_string(), tool.to_string())) + .cloned() + } + + pub fn has_servers(&self) -> bool { + self.connections.has_servers() + } + + /// Returns resource access bound to this binding's exact connection set. + pub fn resource_client(&self) -> McpResourceClient { + McpResourceClient::for_binding(Arc::clone(&self.clients)) + } + + pub async fn list_resources( + &self, + server: &str, + params: Option, + ) -> Result { + self.clients.list_resources(server, params).await + } + + pub async fn list_all_resources( + &self, + include_server: impl Fn(&str) -> bool, + ) -> HashMap> { + self.clients.list_all_resources(include_server).await + } + + pub async fn list_resource_templates( + &self, + server: &str, + params: Option, + ) -> Result { + self.clients.list_resource_templates(server, params).await + } + + pub async fn list_all_resource_templates( + &self, + include_server: impl Fn(&str) -> bool, + ) -> HashMap> { + self.clients + .list_all_resource_templates(include_server) + .await + } + + pub async fn read_resource( + &self, + server: &str, + params: ReadResourceRequestParams, + ) -> Result { + self.clients.read_resource(server, params).await + } +} + +impl fmt::Debug for McpBinding { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("McpBinding") + .field("tools", &self.tools) + .field("prepared_call_count", &self.calls.len()) + .finish_non_exhaustive() + } +} + +/// A call bound to the exact client, tool, timeout, and server metadata seen by +/// one [`McpBinding`]. +#[derive(Clone)] +pub struct PreparedMcpCall { + _connections: Arc, + client: Arc, + catalog_revision: u64, + catalog_revision_source: Arc>, + tool_info: ToolInfo, + server_name: String, + server_metadata: McpServerMetadata, + plugin_id: Option, + selected_plugin_server: bool, +} + +impl PreparedMcpCall { + #[expect( + clippy::too_many_arguments, + reason = "the exact call authority stays together" + )] + pub(crate) fn new( + connections: Arc, + client: Arc, + catalog_revision: u64, + catalog_revision_source: Arc>, + tool_info: ToolInfo, + server_metadata: McpServerMetadata, + plugin_id: Option, + selected_plugin_server: bool, + ) -> Self { + let server_name = tool_info.server_name.clone(); + Self { + _connections: connections, + client, + catalog_revision, + catalog_revision_source, + tool_info, + server_name, + server_metadata, + plugin_id, + selected_plugin_server, + } + } + + pub fn tool_info(&self) -> &ToolInfo { + &self.tool_info + } + + pub fn server_name(&self) -> &str { + &self.server_name + } + + pub fn server_origin(&self) -> Option<&str> { + self.server_metadata + .origin + .as_ref() + .map(super::server::McpServerOrigin::as_str) + } + + pub fn server_environment_id(&self) -> &str { + &self.server_metadata.environment_id + } + + pub fn server_pollutes_memory(&self) -> bool { + self.server_metadata.pollutes_memory + } + + pub fn tool_approval_mode(&self) -> AppToolApproval { + self.server_metadata + .tool_approval_mode(&self.tool_info.tool.name) + } + + pub fn plugin_id(&self) -> Option<&str> { + self.plugin_id.as_deref() + } + + pub fn is_selected_plugin_server(&self) -> bool { + self.selected_plugin_server + } + + pub async fn server_supports_sandbox_state_meta_capability(&self) -> Result { + Ok(self.client.server_supports_sandbox_state_meta_capability) + } + + pub async fn call( + &self, + arguments: Option, + meta: Option, + ) -> Result { + self.call_with_preparation(|| async move { Ok((arguments, meta)) }) + .await + } + + /// Runs irreversible call preparation and execution under the authority of + /// this call's exact catalog revision. + #[expect( + clippy::await_holding_invalid_type, + reason = "catalog replacement must remain serialized with call preparation and execution" + )] + pub async fn call_with_preparation(&self, prepare: F) -> Result + where + F: FnOnce() -> Fut, + Fut: Future, Option)>>, + { + 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!( + "tool call rejected because the catalog changed after `{}/{tool_name}` was prepared", + self.server_name + )); + } + let (arguments, meta) = prepare().await?; + let result = self + .client + .client + .call_tool(tool_name.clone(), arguments, meta, self.client.tool_timeout) + .await + .with_context(|| format!("tool call failed for `{}/{tool_name}`", self.server_name))?; + drop(current_revision); + Ok(call_tool_result_from_rmcp(result)) + } +} + +fn call_tool_result_from_rmcp(result: rmcp::model::CallToolResult) -> CallToolResult { + let content = result + .content + .into_iter() + .map(|content| { + serde_json::to_value(content) + .unwrap_or_else(|_| JsonValue::String("".to_string())) + }) + .collect(); + CallToolResult { + content, + structured_content: result.structured_content, + is_error: result.is_error, + meta: result.meta.and_then(|meta| serde_json::to_value(meta).ok()), + } +} + +#[cfg(test)] +#[path = "binding_tests.rs"] +mod tests; diff --git a/codex-rs/codex-mcp/src/binding_clients.rs b/codex-rs/codex-mcp/src/binding_clients.rs index 08e5b1c465..69ccd66b2b 100644 --- a/codex-rs/codex-mcp/src/binding_clients.rs +++ b/codex-rs/codex-mcp/src/binding_clients.rs @@ -1,9 +1,16 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::Weak; +use anyhow::Context; use anyhow::Result; use anyhow::anyhow; +use codex_rmcp_client::RmcpClient; +use rmcp::model::ListResourceTemplatesResult; +use rmcp::model::ListResourcesResult; use rmcp::model::PaginatedRequestParams; +use rmcp::model::ReadResourceRequestParams; +use rmcp::model::ReadResourceResult; use rmcp::model::Resource; use rmcp::model::ResourceTemplate; use tokio::task::JoinSet; @@ -14,11 +21,94 @@ use crate::rmcp_client::ManagedClient; /// The ready clients captured for one model step. pub(crate) struct McpBindingClients { clients: HashMap>, + identity: McpBindingClientIdentity, } +#[derive(Clone)] +pub(crate) struct McpBindingClientIdentity(Vec<(String, Weak)>); + +impl PartialEq for McpBindingClientIdentity { + fn eq(&self, other: &Self) -> bool { + self.0.len() == other.0.len() + && self.0.iter().zip(&other.0).all( + |((server, client), (other_server, other_client))| { + server == other_server && client.ptr_eq(other_client) + }, + ) + } +} + +impl Eq for McpBindingClientIdentity {} + impl McpBindingClients { pub(crate) fn new(clients: HashMap>) -> Self { - Self { clients } + let mut identity = clients + .iter() + .map(|(server, client)| (server.clone(), Arc::downgrade(&client.client))) + .collect::>(); + identity.sort_by(|left, right| left.0.cmp(&right.0)); + Self { + clients, + identity: McpBindingClientIdentity(identity), + } + } + + pub(crate) fn client(&self, server: &str) -> Option> { + self.clients.get(server).cloned() + } + + pub(crate) fn contains_server(&self, server: &str) -> bool { + self.clients.contains_key(server) + } + + pub(crate) fn identity(&self) -> McpBindingClientIdentity { + self.identity.clone() + } + + pub(crate) async fn list_resources( + &self, + server: &str, + params: Option, + ) -> Result { + let managed = self + .client(server) + .ok_or_else(|| anyhow!("MCP server '{server}' was not ready for this step"))?; + managed + .client + .list_resources(params, managed.tool_timeout) + .await + .with_context(|| format!("resources/list failed for `{server}`")) + } + + pub(crate) async fn list_resource_templates( + &self, + server: &str, + params: Option, + ) -> Result { + let managed = self + .client(server) + .ok_or_else(|| anyhow!("MCP server '{server}' was not ready for this step"))?; + managed + .client + .list_resource_templates(params, managed.tool_timeout) + .await + .with_context(|| format!("resources/templates/list failed for `{server}`")) + } + + pub(crate) async fn read_resource( + &self, + server: &str, + params: ReadResourceRequestParams, + ) -> Result { + let managed = self + .client(server) + .ok_or_else(|| anyhow!("MCP server '{server}' was not ready for this step"))?; + let uri = params.uri.clone(); + managed + .client + .read_resource(params, managed.tool_timeout) + .await + .with_context(|| format!("resources/read failed for `{server}` ({uri})")) } pub(crate) async fn list_all_resources( diff --git a/codex-rs/codex-mcp/src/binding_tests.rs b/codex-rs/codex-mcp/src/binding_tests.rs new file mode 100644 index 0000000000..0533de7d79 --- /dev/null +++ b/codex-rs/codex-mcp/src/binding_tests.rs @@ -0,0 +1,351 @@ +use std::collections::HashMap; +use std::io; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use codex_config::AppToolApproval; +use codex_config::Constrained; +use codex_protocol::mcp::McpServerInfo; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; +use codex_rmcp_client::InProcessTransportFactory; +use codex_rmcp_client::RmcpClient; +use futures::FutureExt; +use pretty_assertions::assert_eq; +use rmcp::model::JsonObject; +use rmcp::model::Tool; +use tokio::io::DuplexStream; +use tokio::sync::Notify; + +use super::McpBinding; +use super::PreparedMcpCall; +use crate::binding_clients::McpBindingClients; +use crate::connection_manager::McpConnectionManager; +use crate::rmcp_client::ManagedClient; +use crate::server::McpServerMetadata; +use crate::server::McpServerOrigin; +use crate::tools::ToolFilter; +use crate::tools::ToolInfo; + +const SERVER_NAME: &str = "docs"; +const TOOL_NAME: &str = "search"; + +struct TestInProcessTransportFactory; + +impl InProcessTransportFactory for TestInProcessTransportFactory { + fn open(&self) -> futures::future::BoxFuture<'static, io::Result> { + async { + let (client_stream, _server_stream) = tokio::io::duplex(1); + Ok(client_stream) + } + .boxed() + } +} + +struct TestStep { + step: Arc, + client: Arc, + tool_catalog_revision: Arc>, +} + +async fn test_step( + label: &str, + approval_mode: AppToolApproval, + supports_sandbox_state_meta: bool, +) -> TestStep { + let tool = ToolInfo { + server_name: SERVER_NAME.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, + callable_name: TOOL_NAME.to_string(), + callable_namespace: SERVER_NAME.to_string(), + namespace_description: None, + tool: Tool::new( + TOOL_NAME.to_string(), + format!("{label} catalog"), + Arc::new(JsonObject::default()), + ), + openai_file_input_optional_fields: Default::default(), + connector_id: None, + connector_name: None, + plugin_display_names: Vec::new(), + }; + let client = Arc::new( + RmcpClient::new_in_process_client(Arc::new(TestInProcessTransportFactory)) + .await + .expect("create in-process MCP client"), + ); + let managed_client = Arc::new(ManagedClient { + client: Arc::clone(&client), + server_info: McpServerInfo { + name: label.to_string(), + title: Some(format!("{label} server")), + version: "1.0.0".to_string(), + description: None, + icons: None, + website_url: None, + }, + tools: vec![tool.clone()], + tool_filter: ToolFilter::default(), + tool_timeout: None, + server_instructions: None, + server_supports_sandbox_state_meta_capability: supports_sandbox_state_meta, + codex_apps_tools_cache_context: None, + }); + let clients = Arc::new(McpBindingClients::new(HashMap::from([( + SERVER_NAME.to_string(), + Arc::clone(&managed_client), + )]))); + let connections = Arc::new( + McpConnectionManager::new_uninitialized_with_permission_profile( + &Constrained::allow_any(AskForApproval::OnRequest), + &PermissionProfile::default(), + /*prefix_mcp_tool_names*/ true, + ), + ); + let tool_catalog_revision = Arc::new(tokio::sync::RwLock::new(0)); + let prepared = PreparedMcpCall::new( + Arc::clone(&connections), + managed_client, + /*catalog_revision*/ 0, + Arc::clone(&tool_catalog_revision), + tool.clone(), + McpServerMetadata { + environment_id: format!("{label}-environment"), + pollutes_memory: label == "old", + origin: Some(McpServerOrigin::StreamableHttp(format!( + "https://{label}.example" + ))), + supports_parallel_tool_calls: false, + default_tools_approval_mode: Some(approval_mode), + tool_approval_modes: HashMap::new(), + }, + Some(format!("{label}-plugin")), + label == "old", + ); + let calls = HashMap::from([((SERVER_NAME.to_string(), TOOL_NAME.to_string()), prepared)]); + + TestStep { + step: Arc::new(McpBinding::new( + connections, + clients, + Arc::new(crate::mcp::tests::test_mcp_config(std::env::temp_dir())), + /*plugins_available*/ false, + vec![tool], + calls, + )), + client, + tool_catalog_revision, + } +} + +#[tokio::test] +async fn prepared_call_keeps_captured_connection_and_authority_after_refresh() -> anyhow::Result<()> +{ + let old = test_step( + "old", + AppToolApproval::Prompt, + /*supports_sandbox_state_meta*/ true, + ) + .await; + let old_call = old + .step + .prepare_call(SERVER_NAME, TOOL_NAME) + .expect("old step should prepare the advertised tool"); + let old_connections = Arc::downgrade(&old.step.connections); + + let new = test_step( + "new", + AppToolApproval::Approve, + /*supports_sandbox_state_meta*/ false, + ) + .await; + let new_call = new + .step + .prepare_call(SERVER_NAME, TOOL_NAME) + .expect("new step should prepare the advertised tool"); + + assert_eq!( + ( + old.step.tools()[0].tool.description.as_deref(), + old_call.tool_info().tool.description.as_deref(), + old_call.server_origin(), + old_call.server_environment_id(), + old_call.server_pollutes_memory(), + old_call.tool_approval_mode(), + old_call.plugin_id(), + old_call.is_selected_plugin_server(), + old_call + .server_supports_sandbox_state_meta_capability() + .await?, + ), + ( + Some("old catalog"), + Some("old catalog"), + Some("https://old.example"), + "old-environment", + true, + AppToolApproval::Prompt, + Some("old-plugin"), + true, + true, + ) + ); + assert_eq!( + ( + new.step.tools()[0].tool.description.as_deref(), + new_call.tool_info().tool.description.as_deref(), + new_call.server_environment_id(), + new_call.tool_approval_mode(), + ), + ( + Some("new catalog"), + Some("new catalog"), + "new-environment", + AppToolApproval::Approve, + ) + ); + assert!(Arc::ptr_eq(&old_call.client.client, &old.client)); + assert!(!Arc::ptr_eq(&old.client, &new.client)); + + drop(old.step); + assert!( + old_connections.upgrade().is_some(), + "the prepared call should keep its captured connection set alive" + ); + drop(old_call); + assert!( + old_connections.upgrade().is_none(), + "the captured connection set should be released with the prepared call" + ); + Ok(()) +} + +#[tokio::test] +async fn prepared_call_does_not_reroute_after_captured_connection_closes() { + let old = test_step( + "old", + AppToolApproval::Prompt, + /*supports_sandbox_state_meta*/ true, + ) + .await; + let old_call = old + .step + .prepare_call(SERVER_NAME, TOOL_NAME) + .expect("old step should prepare the advertised tool"); + let new = test_step( + "new", + AppToolApproval::Approve, + /*supports_sandbox_state_meta*/ false, + ) + .await; + assert!(!Arc::ptr_eq(&old.client, &new.client)); + + old.client.shutdown().await; + + let error = old_call + .call( + Some(serde_json::json!({"query": "codex"})), + /*meta*/ None, + ) + .await + .expect_err("a call bound to a closed connection must fail"); + assert!( + format!("{error:#}").contains("MCP client is shut down"), + "the prepared call should fail on its captured client: {error:#}" + ); +} + +#[tokio::test] +async fn prepared_call_is_rejected_after_catalog_refresh() { + let step = test_step( + "old", + AppToolApproval::Prompt, + /*supports_sandbox_state_meta*/ true, + ) + .await; + let prepared = step + .step + .prepare_call(SERVER_NAME, TOOL_NAME) + .expect("step should prepare the advertised tool"); + + *step.tool_catalog_revision.write().await += 1; + + let error = prepared + .call( + Some(serde_json::json!({"query": "codex"})), + /*meta*/ None, + ) + .await + .expect_err("a call from an older catalog must be rejected"); + assert!( + format!("{error:#}").contains("catalog changed"), + "unexpected error: {error:#}" + ); +} + +#[tokio::test] +async fn stale_prepared_call_does_not_run_preparation() { + let step = test_step( + "old", + AppToolApproval::Prompt, + /*supports_sandbox_state_meta*/ true, + ) + .await; + let prepared = step + .step + .prepare_call(SERVER_NAME, TOOL_NAME) + .expect("step should prepare the advertised tool"); + *step.tool_catalog_revision.write().await += 1; + let prepared_side_effect_ran = Arc::new(AtomicBool::new(false)); + let marker = Arc::clone(&prepared_side_effect_ran); + + prepared + .call_with_preparation(|| async move { + marker.store(true, Ordering::SeqCst); + Ok((None, None)) + }) + .await + .expect_err("a call from an older catalog must be rejected"); + + assert!(!prepared_side_effect_ran.load(Ordering::SeqCst)); +} + +#[tokio::test] +async fn preparation_holds_catalog_authority_until_it_finishes() { + let step = test_step( + "old", + AppToolApproval::Prompt, + /*supports_sandbox_state_meta*/ true, + ) + .await; + let prepared = step + .step + .prepare_call(SERVER_NAME, TOOL_NAME) + .expect("step should prepare the advertised tool"); + let preparation_started = Arc::new(Notify::new()); + let finish_preparation = Arc::new(Notify::new()); + let started = Arc::clone(&preparation_started); + let finish = Arc::clone(&finish_preparation); + let call = tokio::spawn(async move { + prepared + .call_with_preparation(|| async move { + started.notify_one(); + finish.notified().await; + Err(anyhow::anyhow!("stop after preparation")) + }) + .await + }); + + preparation_started.notified().await; + assert!( + step.tool_catalog_revision.try_write().is_err(), + "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()); +} diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 1cc3f9fe9b..66cda0239b 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -72,6 +72,8 @@ use rmcp::model::RequestId; use rmcp::model::Resource; use rmcp::model::ResourceTemplate; use serde_json::Value as JsonValue; +use tokio::sync::Mutex; +use tokio::sync::RwLock; use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; use tracing::warn; @@ -109,6 +111,9 @@ pub struct McpConnectionManager { clients: HashMap, server_metadata: HashMap, required_servers: Vec, + tool_catalog_revision: Arc>, + codex_apps_tools_override: RwLock>>, + codex_apps_refresh_lock: Mutex<()>, tool_plugin_provenance: Arc, prefix_mcp_tool_names: bool, elicitation_requests: ElicitationRequestManager, @@ -352,6 +357,9 @@ impl McpConnectionManager { clients, server_metadata, required_servers, + 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, elicitation_requests: elicitation_requests.clone(), @@ -393,6 +401,9 @@ impl McpConnectionManager { clients: HashMap::new(), server_metadata: HashMap::new(), required_servers: Vec::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, elicitation_requests: ElicitationRequestManager::new( @@ -406,6 +417,14 @@ impl McpConnectionManager { } } + pub fn empty(prefix_mcp_tool_names: bool) -> Self { + Self::new_uninitialized_with_permission_profile( + &Constrained::allow_any(AskForApproval::OnRequest), + &PermissionProfile::default(), + prefix_mcp_tool_names, + ) + } + pub fn has_servers(&self) -> bool { !self.clients.is_empty() } 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 7d79765ad0..8cc61cecb8 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -1,3 +1,4 @@ +use std::sync::Arc; use std::time::Instant; use anyhow::Context; @@ -10,10 +11,14 @@ use tracing::trace; use tracing::trace_span; use super::McpConnectionManager; +use crate::binding::McpBinding; +use crate::binding::PreparedMcpCall; +use crate::binding_clients::McpBindingClients; use crate::codex_apps::prepare_openai_file_params_for_model; 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; +use crate::rmcp_client::ManagedClient; use crate::rmcp_client::list_tools_for_client_uncached; use crate::runtime::emit_duration; use crate::tools::ToolInfo; @@ -80,12 +85,109 @@ impl McpConnectionManager { Some(self.with_server_metadata(tool)) } + #[expect( + clippy::await_holding_invalid_type, + reason = "catalog capture must remain serialized with catalog replacement" + )] + /// Captures the ready clients, their exact tools, and the supplied runtime metadata. + pub async fn capture_binding_with_metadata( + self: &Arc, + config: Arc, + plugins_available: bool, + ) -> McpBinding { + let revision = self.tool_catalog_revision.read().await; + let mut listed_tools = Vec::new(); + let mut clients = std::collections::HashMap::new(); + for (server_name, managed_client) in &self.clients { + let catalog_override = if server_name == CODEX_APPS_MCP_SERVER_NAME { + self.codex_apps_tools_override.read().await.clone() + } else { + None + }; + let Some((client, server_tools)) = managed_client + .capture_ready_client_and_tools(catalog_override) + .await + else { + trace!( + server_name = %server_name, + "omitting MCP server without an exact ready client" + ); + continue; + }; + clients.insert(server_name.clone(), client); + listed_tools.extend( + server_tools + .into_iter() + .map(|tool| self.with_server_metadata(tool)), + ); + } + let clients = Arc::new(McpBindingClients::new(clients)); + let listed_tools = + normalize_tools_for_model_with_prefix(listed_tools, self.prefix_mcp_tool_names); + let mut tools = Vec::with_capacity(listed_tools.len()); + let mut calls = std::collections::HashMap::with_capacity(listed_tools.len()); + for tool_info in listed_tools { + let Some(client) = clients.client(&tool_info.server_name) else { + continue; + }; + let Some(call) = self.prepare_call(&tool_info, client, *revision) else { + trace!( + server_name = %tool_info.server_name, + tool_name = %tool_info.tool.name, + "omitting MCP tool without an exact ready client" + ); + continue; + }; + calls.insert( + ( + tool_info.server_name.clone(), + tool_info.tool.name.to_string(), + ), + call, + ); + tools.push(tool_info); + } + McpBinding::new( + Arc::clone(self), + clients, + config, + plugins_available, + tools, + calls, + ) + } + + fn prepare_call( + self: &Arc, + tool_info: &ToolInfo, + client: Arc, + tool_catalog_revision: u64, + ) -> Option { + let server_name = &tool_info.server_name; + Some(PreparedMcpCall::new( + Arc::clone(self), + client, + tool_catalog_revision, + Arc::clone(&self.tool_catalog_revision), + tool_info.clone(), + self.server_metadata.get(server_name)?.clone(), + self.plugin_id_for_mcp_server_name(server_name) + .map(str::to_string), + self.is_selected_plugin_mcp_server(server_name), + )) + } + /// Force-refresh codex apps tools by bypassing the in-process cache. /// /// On success, the refreshed tools replace shared cache contents when the /// cache is enabled and the latest filtered tools are returned directly to /// the caller. On failure, existing shared cache contents remain unchanged. + #[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; let refresh_start = Instant::now(); let managed_client = self .clients @@ -103,7 +205,7 @@ impl McpConnectionManager { .map(|cache_context| { cache_context.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh) }); - let tools = list_tools_for_client_uncached( + let client_tools = list_tools_for_client_uncached( CODEX_APPS_MCP_SERVER_NAME, /*is_codex_apps_mcp_server*/ true, /*codex_apps_refresh_trigger*/ "explicit", @@ -116,16 +218,22 @@ impl McpConnectionManager { format!("failed to refresh tools for MCP server '{CODEX_APPS_MCP_SERVER_NAME}'") })?; - let tools = - match ( - managed_client.codex_apps_tools_cache_context.as_ref(), + 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, - ) { - (Some(cache_context), Some(fetch_ticket)) => cache_context - .publish_if_newest_accepted(fetch_ticket, &managed_client.server_info, tools), - (None, None) => tools, - _ => unreachable!("Codex Apps fetch ticket requires cache context"), - }; + &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); emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, list_start.elapsed(), diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 88297a5592..473cee6ada 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::required::startup_outcome_error_message; use super::*; +use crate::McpBinding; use crate::elicitation::ElicitationLifecycle; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationRequestRouter; @@ -38,11 +39,22 @@ use codex_rmcp_client::RmcpClient; use futures::FutureExt; use futures::future::BoxFuture; use pretty_assertions::assert_eq; +use rmcp::ErrorData as McpError; +use rmcp::ServerHandler; +use rmcp::ServiceExt; +use rmcp::model::ClientCapabilities; use rmcp::model::CreateElicitationRequestParams; use rmcp::model::ElicitationAction; use rmcp::model::ElicitationCapability; +use rmcp::model::Implementation; +use rmcp::model::InitializeRequestParams; use rmcp::model::JsonObject; +use rmcp::model::ListToolsResult; use rmcp::model::NumberOrString; +use rmcp::model::PaginatedRequestParams; +use rmcp::model::ProtocolVersion; +use rmcp::model::ServerCapabilities; +use rmcp::model::ServerInfo; use rmcp::model::Tool; use std::collections::HashMap; use std::collections::HashSet; @@ -51,6 +63,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicUsize; use tempfile::tempdir; use tokio::io::DuplexStream; +use tokio::sync::Notify; fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { ToolInfo { @@ -94,6 +107,15 @@ fn store_current_tools(cache_context: &ConnectorRuntimeContext, tools: ); } +async fn capture_binding(manager: &Arc) -> McpBinding { + manager + .capture_binding_with_metadata( + Arc::new(crate::mcp::tests::test_mcp_config(std::env::temp_dir())), + /*plugins_available*/ false, + ) + .await +} + fn create_test_server_info(title: &str) -> McpServerInfo { McpServerInfo { name: "codex-apps".to_string(), @@ -117,6 +139,55 @@ impl InProcessTransportFactory for TestInProcessTransportFactory { } } +#[derive(Clone)] +struct RefreshTestTransportFactory { + tool: Tool, + list_started: Option>, + release_list: Option>, +} + +impl ServerHandler for RefreshTestTransportFactory { + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + } + + async fn list_tools( + &self, + _request: Option, + _context: rmcp::service::RequestContext, + ) -> Result { + if let Some(list_started) = &self.list_started { + list_started.notify_one(); + } + if let Some(release_list) = &self.release_list { + release_list.notified().await; + } + Ok(ListToolsResult { + tools: vec![self.tool.clone()], + next_cursor: None, + meta: None, + }) + } +} + +impl InProcessTransportFactory for RefreshTestTransportFactory { + fn open(&self) -> BoxFuture<'static, io::Result> { + let server = self.clone(); + async move { + let (client_stream, server_stream) = tokio::io::duplex(4096); + tokio::spawn(async move { + let server = server + .serve(server_stream) + .await + .expect("serve test MCP server"); + server.waiting().await.expect("wait for test MCP server"); + }); + Ok(client_stream) + } + .boxed() + } +} + async fn create_test_managed_client(tools: Vec) -> ManagedClient { ManagedClient { client: Arc::new( @@ -153,6 +224,83 @@ async fn create_ready_async_managed_client(tools: Vec) -> AsyncManaged } } +async fn create_test_manager_with_ready_apps_client( + cache_context: ConnectorRuntimeContext, + tool_name: &str, + list_started: Option>, + 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, + })) + .await?, + ); + client + .initialize( + InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("codex-test", "0.0.0-test"), + ) + .with_protocol_version(ProtocolVersion::V_2025_06_18), + Some(Duration::from_secs(5)), + Box::new(|_, _| async { Err(anyhow!("unexpected elicitation")) }.boxed()), + ) + .await?; + + let managed_client = ManagedClient { + client, + server_info: create_test_server_info("Codex Apps"), + tools: vec![tool], + tool_filter: ToolFilter::default(), + tool_timeout: Some(Duration::from_secs(5)), + server_instructions: None, + server_supports_sandbox_state_meta_capability: false, + codex_apps_tools_cache_context: Some(cache_context.clone()), + }; + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionManager::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + manager.clients.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + AsyncManagedClient { + client: futures::future::ready::>(Ok( + managed_client, + )) + .boxed() + .shared(), + is_codex_apps_mcp_server: true, + cached_server_info: Some(create_test_server_info("Codex Apps")), + codex_apps_tools_cache_context: Some(cache_context), + tool_catalog_cache_context: None, + tool_filter: ToolFilter::default(), + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + startup_reconnect: None, + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + }, + ); + manager.server_metadata.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + McpServerMetadata { + environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + pollutes_memory: false, + origin: None, + supports_parallel_tool_calls: false, + default_tools_approval_mode: None, + tool_approval_modes: HashMap::new(), + }, + ); + Ok(Arc::new(manager)) +} + fn create_test_manager_with_failed_apps_startup( cached_tools: Vec, reconnect_factory: Arc ManagedClientFuture + Send + Sync>, @@ -747,6 +895,181 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { assert_eq!(tool.callable_name, "calendar_create_event"); } +#[tokio::test] +async fn capture_binding_uses_the_ready_clients_own_tools() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + store_current_tools( + &cache_context, + vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "shared_cached_tool", + )], + ); + let mut ready_client = create_test_managed_client(vec![ + create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "client_local_tool"), + create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "client_local_blocked"), + ]) + .await; + let tool_filter = ToolFilter { + enabled: None, + disabled: HashSet::from(["client_local_blocked".to_string()]), + }; + ready_client.tool_filter = tool_filter.clone(); + ready_client.codex_apps_tools_cache_context = Some(cache_context.clone()); + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionManager::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + manager.clients.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + AsyncManagedClient { + client: futures::future::ready(Ok(ready_client)).boxed().shared(), + is_codex_apps_mcp_server: true, + cached_server_info: None, + codex_apps_tools_cache_context: Some(cache_context), + tool_catalog_cache_context: None, + tool_filter, + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + startup_reconnect: None, + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + }, + ); + manager.server_metadata.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + McpServerMetadata { + environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + pollutes_memory: false, + origin: None, + supports_parallel_tool_calls: false, + default_tools_approval_mode: None, + tool_approval_modes: HashMap::new(), + }, + ); + let manager = Arc::new(manager); + + assert_eq!( + manager + .list_all_tools() + .await + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["shared_cached_tool"] + ); + let step = capture_binding(&manager).await; + assert_eq!( + step.tools() + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["client_local_tool"] + ); + assert!( + step.prepare_call(CODEX_APPS_MCP_SERVER_NAME, "client_local_tool") + .is_some() + ); + assert!( + step.prepare_call(CODEX_APPS_MCP_SERVER_NAME, "shared_cached_tool") + .is_none() + ); + assert!( + step.prepare_call(CODEX_APPS_MCP_SERVER_NAME, "client_local_blocked") + .is_none() + ); +} + +#[tokio::test] +async fn hard_refresh_keeps_binding_override_local_when_shared_cache_loses_race() +-> anyhow::Result<()> { + let codex_home = tempdir()?; + let shared_cache = ConnectorRuntimeManager::::default(); + let cache_key = ConnectorRuntimeContextKey::personal( + Some("shared-account".to_string()), + Some("shared-user".to_string()), + ); + let cache_context_a = shared_cache.context(codex_home.path().to_path_buf(), cache_key.clone()); + let cache_context_b = shared_cache.context(codex_home.path().to_path_buf(), cache_key); + let list_started = Arc::new(Notify::new()); + let release_list = Arc::new(Notify::new()); + let manager_a = create_test_manager_with_ready_apps_client( + cache_context_a.clone(), + "a_only", + Some(Arc::clone(&list_started)), + Some(Arc::clone(&release_list)), + ) + .await?; + let manager_b = create_test_manager_with_ready_apps_client( + cache_context_b, + "b_only", + /*list_started*/ None, + /*release_list*/ None, + ) + .await?; + + 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() + .await + }); + list_started.notified().await; + let tools_b = manager_b.hard_refresh_codex_apps_tools_cache().await?; + release_list.notify_one(); + let tools_a = refresh_a.await??; + + assert_eq!( + tools_b + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["b_only"] + ); + assert_eq!( + tools_a + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["b_only"] + ); + assert_eq!( + cache_context_a + .current_tools() + .expect("shared cache tools") + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["b_only"] + ); + assert_eq!( + capture_binding(&manager_a) + .await + .tools() + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["a_only"] + ); + assert_eq!( + capture_binding(&manager_b) + .await + .tools() + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["b_only"] + ); + Ok(()) +} + #[tokio::test(start_paused = true)] async fn tool_catalog_cache_sanitizes_tools_and_tracks_environment_generation() { let cache = McpToolCatalogCache::default(); @@ -920,6 +1243,93 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() { ); } +#[tokio::test] +async fn capture_binding_waits_for_fresh_startup_even_with_cached_tools() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + store_current_tools( + &cache_context, + vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "shared_cached_tool", + )], + ); + let startup_complete = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let startup_complete_for_client = Arc::clone(&startup_complete); + let (startup_started, wait_for_startup) = tokio::sync::oneshot::channel(); + let (release_startup, startup_released) = tokio::sync::oneshot::channel(); + let pending_client = async move { + startup_started.send(()).expect("signal client startup"); + startup_released.await.expect("release client startup"); + startup_complete_for_client.store(true, std::sync::atomic::Ordering::Release); + Ok(create_test_managed_client(vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "client_local_tool", + )]) + .await) + } + .boxed() + .shared(); + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionManager::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + manager.clients.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + AsyncManagedClient { + client: pending_client, + is_codex_apps_mcp_server: true, + cached_server_info: None, + codex_apps_tools_cache_context: Some(cache_context), + tool_catalog_cache_context: None, + tool_filter: ToolFilter::default(), + startup_complete, + startup_reconnect: None, + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + }, + ); + manager.server_metadata.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + McpServerMetadata { + environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + pollutes_memory: false, + origin: None, + supports_parallel_tool_calls: false, + default_tools_approval_mode: None, + tool_approval_modes: HashMap::new(), + }, + ); + let manager = Arc::new(manager); + let manager_for_capture = Arc::clone(&manager); + let capture = tokio::spawn(async move { capture_binding(&manager_for_capture).await }); + + wait_for_startup.await.expect("client startup should begin"); + assert!(!capture.is_finished()); + release_startup.send(()).expect("release client startup"); + + let step = capture.await.expect("capture task"); + assert_eq!( + step.tools() + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["client_local_tool"] + ); + assert!( + step.resource_client() + .has_server(CODEX_APPS_MCP_SERVER_NAME) + .await + ); +} + #[tokio::test] async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { let managed_client = @@ -1232,7 +1642,19 @@ async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client() .boxed() .shared() }); - let manager = create_test_manager_with_failed_apps_startup(Vec::new(), reconnect_factory); + let mut manager = create_test_manager_with_failed_apps_startup(Vec::new(), reconnect_factory); + manager.server_metadata.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + McpServerMetadata { + environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + pollutes_memory: false, + origin: None, + supports_parallel_tool_calls: false, + default_tools_approval_mode: None, + tool_approval_modes: HashMap::new(), + }, + ); + let manager = Arc::new(manager); let reconnect_finished_wait = reconnect_finished.notified(); let tools = manager.list_all_tools().await; @@ -1249,6 +1671,17 @@ async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client() ); assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + let step = capture_binding(&manager).await; + let prepared = step + .prepare_call(CODEX_APPS_MCP_SERVER_NAME, "drive_search") + .expect("recovered tool should have a prepared call"); + assert!( + !prepared + .server_supports_sandbox_state_meta_capability() + .await + .expect("prepared call should use the recovered client") + ); + let tools = manager.list_all_tools().await; assert_eq!( tools @@ -1394,13 +1827,25 @@ async fn tool_lists_do_not_block_and_share_codex_apps_startup_reconnect() { .boxed() .shared() }); - let manager = Arc::new(create_test_manager_with_failed_apps_startup( + let mut manager = create_test_manager_with_failed_apps_startup( vec![create_test_tool( CODEX_APPS_MCP_SERVER_NAME, "cached_drive_search", )], reconnect_factory, - )); + ); + manager.server_metadata.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + McpServerMetadata { + environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + pollutes_memory: false, + origin: None, + supports_parallel_tool_calls: false, + default_tools_approval_mode: None, + tool_approval_modes: HashMap::new(), + }, + ); + let manager = Arc::new(manager); let reconnect_started_wait = reconnect_started.notified(); let first_tools = tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()) .await @@ -1425,6 +1870,19 @@ async fn tool_lists_do_not_block_and_share_codex_apps_startup_reconnect() { .collect::>(), vec!["cached_drive_search"] ); + let pending_step = tokio::time::timeout(Duration::from_millis(10), capture_binding(&manager)) + .await + .expect("step capture should not wait for reconnect"); + assert!( + pending_step.tools().is_empty(), + "a model step must not advertise cached tools without an exact ready client" + ); + let pending_resources = pending_step.resource_client(); + assert!( + !pending_resources + .has_server(CODEX_APPS_MCP_SERVER_NAME) + .await + ); release_reconnect.notify_one(); tokio::task::yield_now().await; @@ -1436,6 +1894,32 @@ async fn tool_lists_do_not_block_and_share_codex_apps_startup_reconnect() { .collect::>(), vec!["drive_search"] ); + let recovered_step = capture_binding(&manager).await; + assert_eq!( + recovered_step + .tools() + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["drive_search"] + ); + assert!( + recovered_step + .prepare_call(CODEX_APPS_MCP_SERVER_NAME, "drive_search") + .is_some() + ); + let recovered_resources = recovered_step.resource_client(); + assert!( + recovered_resources + .has_server(CODEX_APPS_MCP_SERVER_NAME) + .await + ); + assert!( + !pending_resources + .has_server(CODEX_APPS_MCP_SERVER_NAME) + .await + ); + assert!(pending_resources.cache_key() != recovered_resources.cache_key()); assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); } diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 499fe43cf1..557c92b005 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -1,3 +1,5 @@ +pub use binding::McpBinding; +pub use binding::PreparedMcpCall; pub use connection_manager::McpConnectionManager; pub use connection_manager::tool_is_model_visible; pub use elicitation::ElicitationLifecycle; @@ -82,6 +84,7 @@ pub use mcp::mcp_permission_prompt_is_auto_approved; pub use mcp::qualified_mcp_tool_name_prefix; pub(crate) mod auth_elicitation; +mod binding; pub(crate) mod binding_clients; mod catalog; pub(crate) mod codex_apps; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 97f5f7061f..b0d7f557a5 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -696,4 +696,4 @@ async fn collect_mcp_server_status_snapshot_from_manager( #[cfg(test)] #[path = "mod_tests.rs"] -mod tests; +pub(crate) mod tests; diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index fb84700b15..6c48a9f75b 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -17,7 +17,7 @@ use std::collections::HashMap; use std::collections::HashSet; use std::path::PathBuf; -fn test_mcp_config(codex_home: PathBuf) -> McpConfig { +pub(crate) fn test_mcp_config(codex_home: PathBuf) -> McpConfig { McpConfig { chatgpt_base_url: "https://chatgpt.com".to_string(), apps_mcp_product_sku: None, diff --git a/codex-rs/codex-mcp/src/resource_client.rs b/codex-rs/codex-mcp/src/resource_client.rs index 84d64f441f..c248d889a1 100644 --- a/codex-rs/codex-mcp/src/resource_client.rs +++ b/codex-rs/codex-mcp/src/resource_client.rs @@ -8,8 +8,10 @@ use codex_protocol::mcp::ResourceContent; use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceRequestParams; -use crate::McpConnectionManager; use crate::McpRuntime; +use crate::binding_clients::McpBindingClientIdentity; +use crate::binding_clients::McpBindingClients; +use crate::connection_manager::McpConnectionManager; /// One page of resources returned by an MCP server. #[derive(Clone, Debug, PartialEq)] @@ -27,21 +29,48 @@ pub struct McpResourceReadResult { pub contents: Vec, } -/// Session-scoped access to MCP resources through the thread runtime. -/// -/// Calls automatically use the connection set most recently published by the runtime. +/// Access to MCP resources through either the latest runtime or one exact step. #[derive(Clone)] pub struct McpResourceClient { - runtime: Arc, + source: McpResourceSource, } /// Opaque identity for the connection set currently used by an MCP resource client. #[derive(Clone)] -pub struct McpResourceClientCacheKey(Weak); +pub struct McpResourceClientCacheKey(McpResourceClientCacheKeyInner); + +#[derive(Clone)] +enum McpResourceClientCacheKeyInner { + Latest(Weak), + Exact(McpBindingClientIdentity), +} + +#[derive(Clone)] +enum McpResourceSource { + Latest(Arc), + Exact(Arc), +} impl PartialEq for McpResourceClientCacheKey { fn eq(&self, other: &Self) -> bool { - self.0.ptr_eq(&other.0) + match (&self.0, &other.0) { + ( + McpResourceClientCacheKeyInner::Latest(left), + McpResourceClientCacheKeyInner::Latest(right), + ) => left.ptr_eq(right), + ( + McpResourceClientCacheKeyInner::Exact(left), + McpResourceClientCacheKeyInner::Exact(right), + ) => left == right, + ( + McpResourceClientCacheKeyInner::Latest(_), + McpResourceClientCacheKeyInner::Exact(_), + ) + | ( + McpResourceClientCacheKeyInner::Exact(_), + McpResourceClientCacheKeyInner::Latest(_), + ) => false, + } } } @@ -56,21 +85,40 @@ impl std::fmt::Debug for McpResourceClient { } impl McpResourceClient { - /// Creates a resource client backed by the thread's MCP runtime. + /// Creates a resource client that follows the thread's latest published runtime. pub fn new(runtime: Arc) -> Self { - Self { runtime } + Self { + source: McpResourceSource::Latest(runtime), + } } - /// Returns an identity that changes whenever the published connection set changes. + pub(crate) fn for_binding(clients: Arc) -> Self { + Self { + source: McpResourceSource::Exact(clients), + } + } + + /// Returns the identity of the connection set used by this client. pub fn cache_key(&self) -> McpResourceClientCacheKey { - McpResourceClientCacheKey(Arc::downgrade(&self.runtime.snapshot())) + let key = match &self.source { + McpResourceSource::Latest(runtime) => { + McpResourceClientCacheKeyInner::Latest(Arc::downgrade(&runtime.snapshot())) + } + McpResourceSource::Exact(clients) => { + McpResourceClientCacheKeyInner::Exact(clients.identity()) + } + }; + McpResourceClientCacheKey(key) } - /// Returns whether the current connection set contains the named server. + /// Returns whether this client can address the named server. /// - /// This does not wait for server startup or imply that startup succeeded. + /// This does not wait for server startup. pub async fn has_server(&self, server: &str) -> bool { - self.runtime.snapshot().contains_server(server) + match &self.source { + McpResourceSource::Latest(runtime) => runtime.snapshot().contains_server(server), + McpResourceSource::Exact(clients) => clients.contains_server(server), + } } /// Lists one resource page from the named server. @@ -81,11 +129,12 @@ impl McpResourceClient { ) -> Result { let params = cursor.map(|cursor| PaginatedRequestParams::default().with_cursor(Some(cursor))); - let result = self - .runtime - .snapshot() - .list_resources(server, params) - .await?; + let result = match &self.source { + McpResourceSource::Latest(runtime) => { + runtime.snapshot().list_resources(server, params).await + } + McpResourceSource::Exact(clients) => clients.list_resources(server, params).await, + }?; let resources = result .resources .into_iter() @@ -99,11 +148,13 @@ impl McpResourceClient { /// Reads one resource from the named server. pub async fn read_resource(&self, server: &str, uri: &str) -> Result { - let result = self - .runtime - .snapshot() - .read_resource(server, ReadResourceRequestParams::new(uri.to_string())) - .await?; + let params = ReadResourceRequestParams::new(uri.to_string()); + let result = match &self.source { + McpResourceSource::Latest(runtime) => { + runtime.snapshot().read_resource(server, params).await + } + McpResourceSource::Exact(clients) => clients.read_resource(server, params).await, + }?; let contents = result .contents .into_iter() diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index aa34b47cf6..912a5375b1 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -140,7 +140,7 @@ impl ManagedClient { ); } - self.tools.clone() + filter_tools(self.tools.clone(), &self.tool_filter) } } @@ -516,6 +516,50 @@ impl AsyncManagedClient { self.client.clone().await } + /// Captures the ready client revision that is current now. + /// + /// A recovered Codex Apps connection replaces the failed startup future for + /// future steps, but cannot reroute a call that was already prepared. + fn ready_client_snapshot(&self) -> Option { + if let Some(client) = self + .startup_reconnect + .as_ref() + .and_then(|reconnect| reconnect.current_client()) + { + return Some(futures::future::ready(Ok(client)).boxed().shared()); + } + match self.client.peek() { + Some(Ok(_)) => Some(self.client.clone()), + Some(Err(_)) | None => None, + } + } + + /// Captures one ready client and derives its model-visible tools. + /// + /// A fresh client waits for initial startup even when metadata is cached. + /// After a failed startup, cached tools remain available to metadata-only + /// callers while recovery runs, but a model step still requires one exact + /// ready client. + pub(crate) async fn capture_ready_client_and_tools( + &self, + catalog_override: Option>, + ) -> Option<(Arc, Vec)> { + if !self.startup_complete.load(Ordering::Acquire) { + let _ = self.client().await; + } + self.reconnect_failed_startup().await; + let client = if self.has_cached_tools() { + self.ready_client_snapshot()? + } else { + self.client().await.ok()?; + self.ready_client_snapshot()? + }; + let managed_client = Arc::new(client.await.ok()?); + let tools = catalog_override.unwrap_or_else(|| managed_client.tools.clone()); + let tools = filter_tools(tools, &managed_client.tool_filter); + Some((Arc::clone(&managed_client), self.prepare_tools(tools))) + } + pub(crate) async fn reconnect_failed_startup(&self) { let Some(startup_reconnect) = self.startup_reconnect.as_ref() else { return; @@ -915,7 +959,7 @@ async fn start_server_task( let fetch_ticket = codex_apps_tools_cache_context .as_ref() .map(|cache_context| cache_context.begin_fetch(ConnectorRuntimeFetchSource::Startup)); - let tools = list_tools_for_client_uncached( + let client_tools = list_tools_for_client_uncached( &server_name, is_codex_apps_mcp_server, /*codex_apps_refresh_trigger*/ "initial", @@ -926,11 +970,13 @@ async fn start_server_task( .await .map_err(StartupOutcomeError::from)?; let server_info = mcp_server_info_from_implementation(initialize_result.server_info); - let 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, tools) - } - (None, None) => tools, + 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(); @@ -938,7 +984,7 @@ async fn start_server_task( tool_catalog_cache_context.as_ref(), tool_catalog_fetch_ticket, ) { - cache_context.publish_if_newest(fetch_ticket, &tools); + cache_context.publish_if_newest(fetch_ticket, &shared_tools); } if has_shared_tool_catalog { emit_duration( @@ -947,12 +993,10 @@ async fn start_server_task( &[("cache", "miss")], ); } - let tools = filter_tools(tools, &tool_filter); - let managed = ManagedClient { client: Arc::clone(&client), server_info, - tools, + tools: client_tools, tool_timeout: Some(tool_timeout), tool_filter, server_instructions: initialize_result.instructions,