openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-21 17:48:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Keep refreshed MCP tool catalogs with their clients (#43031)

Browse Source 
## Why

Refreshed Apps catalogs need to survive connection reuse. Catalog revisions scoped to the connection set can lose refreshed tools during reconciliation and invalidate prepared calls for unrelated servers.

## What changed

- Store tools, revisions, and refresh coordination on each MCP client, and key cached bindings by client identity and revision.
- Preserve existing binding snapshots while rejecting stale calls before preparation. Wait for active calls to finish before publishing a refreshed catalog.
- Add `CodexThread::refresh_codex_apps_tools` to reconcile the runtime and return the refreshed client's raw Apps inventory alongside tool names eligible under that runtime's generic MCP policy. Keep hidden and disabled tools in the inventory for callers to apply app-specific policy.

## Testing

Add regression coverage for refreshed catalogs surviving client reuse, unrelated calls remaining valid, binding cache invalidation, and refresh snapshots retaining the client's own inventory when another client wins the shared-cache race. Update stale-call and preparation tests to use client-owned catalogs.

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