github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-21 05:22:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bind MCP calls to captured catalog revisions (#34588)

Browse Source 
## 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
This commit is contained in:
jif
2026-07-21 18:13:41 +00:00
committed by copyberry
parent 51200321eb
commit 65f8bf6853
11 changed files with 1495 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-mcp/src/binding.rs
+294
View File
@@ -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<McpConnectionManager>,
clients: Arc<McpBindingClients>,
config: Arc<McpConfig>,
plugins_available: bool,
tools: Vec<ToolInfo>,
calls: HashMap<(String, String), PreparedMcpCall>,
}
impl McpBinding {
/// Creates an empty binding for tests and callers without a materialized runtime.
pub fn empty(config: Arc<McpConfig>) -> 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<McpConnectionManager>,
clients: Arc<McpBindingClients>,
config: Arc<McpConfig>,
plugins_available: bool,
tools: Vec<ToolInfo>,
calls: HashMap<(String, String), PreparedMcpCall>,
) -> Self {
Self {
connections,
clients,
config,
plugins_available,
tools,
calls,
}
}
pub fn config(&self) -> &Arc<McpConfig> {
&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<PreparedMcpCall> {
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<PaginatedRequestParams>,
) -> Result<ListResourcesResult> {
self.clients.list_resources(server, params).await
}
pub async fn list_all_resources(
&self,
include_server: impl Fn(&str) -> bool,
) -> HashMap<String, Vec<Resource>> {
self.clients.list_all_resources(include_server).await
}
pub async fn list_resource_templates(
&self,
server: &str,
params: Option<PaginatedRequestParams>,
) -> Result<ListResourceTemplatesResult> {
self.clients.list_resource_templates(server, params).await
}
pub async fn list_all_resource_templates(
&self,
include_server: impl Fn(&str) -> bool,
) -> HashMap<String, Vec<ResourceTemplate>> {
self.clients
.list_all_resource_templates(include_server)
.await
}
pub async fn read_resource(
&self,
server: &str,
params: ReadResourceRequestParams,
) -> Result<ReadResourceResult> {
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<McpConnectionManager>,
client: Arc<ManagedClient>,
catalog_revision: u64,
catalog_revision_source: Arc<RwLock<u64>>,
tool_info: ToolInfo,
server_name: String,
server_metadata: McpServerMetadata,
plugin_id: Option<String>,
selected_plugin_server: bool,
}
impl PreparedMcpCall {
#[expect(
clippy::too_many_arguments,
reason = "the exact call authority stays together"
)]
pub(crate) fn new(
connections: Arc<McpConnectionManager>,
client: Arc<ManagedClient>,
catalog_revision: u64,
catalog_revision_source: Arc<RwLock<u64>>,
tool_info: ToolInfo,
server_metadata: McpServerMetadata,
plugin_id: Option<String>,
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<bool> {
Ok(self.client.server_supports_sandbox_state_meta_capability)
}
pub async fn call(
&self,
arguments: Option<JsonValue>,
meta: Option<JsonValue>,
) -> Result<CallToolResult> {
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<F, Fut>(&self, prepare: F) -> Result<CallToolResult>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<(Option<JsonValue>, Option<JsonValue>)>>,
{
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("<content>".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;
codex-rs/codex-mcp/src/binding_clients.rs
+91 -1
View File
@@ -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<String, Arc<ManagedClient>>,
identity: McpBindingClientIdentity,
}
#[derive(Clone)]
pub(crate) struct McpBindingClientIdentity(Vec<(String, Weak<RmcpClient>)>);
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<String, Arc<ManagedClient>>) -> Self {
Self { clients }
let mut identity = clients
.iter()
.map(|(server, client)| (server.clone(), Arc::downgrade(&client.client)))
.collect::<Vec<_>>();
identity.sort_by(|left, right| left.0.cmp(&right.0));
Self {
clients,
identity: McpBindingClientIdentity(identity),
}
}
pub(crate) fn client(&self, server: &str) -> Option<Arc<ManagedClient>> {
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<PaginatedRequestParams>,
) -> Result<ListResourcesResult> {
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<PaginatedRequestParams>,
) -> Result<ListResourceTemplatesResult> {
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<ReadResourceResult> {
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(
codex-rs/codex-mcp/src/binding_tests.rs
+351
View File
@@ -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<DuplexStream>> {
async {
let (client_stream, _server_stream) = tokio::io::duplex(1);
Ok(client_stream)
}
.boxed()
}
}
struct TestStep {
step: Arc<McpBinding>,
client: Arc<RmcpClient>,
tool_catalog_revision: Arc<tokio::sync::RwLock<u64>>,
}
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());
}
codex-rs/codex-mcp/src/connection_manager.rs
+19
View File
@@ -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<String, AsyncManagedClient>,
server_metadata: HashMap<String, McpServerMetadata>,
required_servers: Vec<String>,
tool_catalog_revision: Arc<RwLock<u64>>,
codex_apps_tools_override: RwLock<Option<Vec<ToolInfo>>>,
codex_apps_refresh_lock: Mutex<()>,
tool_plugin_provenance: Arc<ToolPluginProvenance>,
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()
}
codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs
+118 -10
View File
@@ -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<Self>,
config: Arc<crate::McpConfig>,
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<Self>,
tool_info: &ToolInfo,
client: Arc<ManagedClient>,
tool_catalog_revision: u64,
) -> Option<PreparedMcpCall> {
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<Vec<ToolInfo>> {
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(),
codex-rs/codex-mcp/src/connection_manager_tests.rs
+487 -3
View File
@@ -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<ToolInfo>, tools:
);
}
async fn capture_binding(manager: &Arc<McpConnectionManager>) -> 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<Arc<Notify>>,
release_list: Option<Arc<Notify>>,
}
impl ServerHandler for RefreshTestTransportFactory {
fn get_info(&self) -> ServerInfo {
ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
}
async fn list_tools(
&self,
_request: Option<PaginatedRequestParams>,
_context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
) -> Result<ListToolsResult, McpError> {
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<DuplexStream>> {
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<ToolInfo>) -> ManagedClient {
ManagedClient {
client: Arc::new(
@@ -153,6 +224,83 @@ async fn create_ready_async_managed_client(tools: Vec<ToolInfo>) -> AsyncManaged
}
}
async fn create_test_manager_with_ready_apps_client(
cache_context: ConnectorRuntimeContext<ToolInfo>,
tool_name: &str,
list_started: Option<Arc<Notify>>,
release_list: Option<Arc<Notify>>,
) -> anyhow::Result<Arc<McpConnectionManager>> {
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::<Result<ManagedClient, StartupOutcomeError>>(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<ToolInfo>,
reconnect_factory: Arc<dyn Fn() -> 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<_>>(),
vec!["shared_cached_tool"]
);
let step = capture_binding(&manager).await;
assert_eq!(
step.tools()
.iter()
.map(|tool| tool.callable_name.as_str())
.collect::<Vec<_>>(),
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::<ToolInfo>::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<_>>(),
vec!["b_only"]
);
assert_eq!(
tools_a
.iter()
.map(|tool| tool.callable_name.as_str())
.collect::<Vec<_>>(),
vec!["b_only"]
);
assert_eq!(
cache_context_a
.current_tools()
.expect("shared cache tools")
.iter()
.map(|tool| tool.callable_name.as_str())
.collect::<Vec<_>>(),
vec!["b_only"]
);
assert_eq!(
capture_binding(&manager_a)
.await
.tools()
.iter()
.map(|tool| tool.callable_name.as_str())
.collect::<Vec<_>>(),
vec!["a_only"]
);
assert_eq!(
capture_binding(&manager_b)
.await
.tools()
.iter()
.map(|tool| tool.callable_name.as_str())
.collect::<Vec<_>>(),
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<_>>(),
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<_>>(),
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<_>>(),
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<_>>(),
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);
}
codex-rs/codex-mcp/src/lib.rs
+3
View File
@@ -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;
codex-rs/codex-mcp/src/mcp/mod.rs
+1 -1
View File
@@ -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;
codex-rs/codex-mcp/src/mcp/mod_tests.rs
+1 -1
View File
@@ -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,
codex-rs/codex-mcp/src/resource_client.rs
+75 -24
View File
@@ -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<ResourceContent>,
}
/// 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<McpRuntime>,
source: McpResourceSource,
}
/// Opaque identity for the connection set currently used by an MCP resource client.
#[derive(Clone)]
pub struct McpResourceClientCacheKey(Weak<McpConnectionManager>);
pub struct McpResourceClientCacheKey(McpResourceClientCacheKeyInner);
#[derive(Clone)]
enum McpResourceClientCacheKeyInner {
Latest(Weak<McpConnectionManager>),
Exact(McpBindingClientIdentity),
}
#[derive(Clone)]
enum McpResourceSource {
Latest(Arc<McpRuntime>),
Exact(Arc<McpBindingClients>),
}
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<McpRuntime>) -> 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<McpBindingClients>) -> 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<McpResourcePage> {
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<McpResourceReadResult> {
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()
codex-rs/codex-mcp/src/rmcp_client.rs
+55 -11
View File
@@ -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<ManagedClientFuture> {
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<Vec<ToolInfo>>,
) -> Option<(Arc<ManagedClient>, Vec<ToolInfo>)> {
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,