zzos/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-28 08:43:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

Reuse MCP handlers across equivalent bindings (#47988)

Browse Source 
## Why

Refreshing the MCP catalog with unchanged tool metadata recreated handlers and forced the tool search index to rebuild.

## What changed

Cache handlers by canonical tool name and reuse them when tool metadata, agent plugin status, and the input schema size limit are unchanged. Replace handlers when those inputs change and evict tools removed from the catalog, preserving search index reuse across equivalent bindings.

## Testing

Extend the deferred tool world state test to verify that an unchanged catalog refresh captures a new binding without rebuilding the search index, and that removing all deferred tools stops advertising `tool_search`.

GitOrigin-RevId: 150866f635e9ee817d7ca44fb76e724a3f7608a5
This commit is contained in:
Henry Levy
2026-09-25 02:54:07 +00:00
committed by copyberry
parent 9ef08dcf2b
commit 0f0efab8d3
2 changed files with 119 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/mcp_tool_exposure.rs
+50 -39
View File
@@ -1,9 +1,8 @@
use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::hash_map::Entry;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::Weak;
use codex_connectors::AppToolPolicyEvaluator;
use codex_connectors::AppToolPolicyInput;
@@ -25,12 +24,14 @@ use crate::tools::registry::ToolRegistry;
#[derive(Default)]
pub(crate) struct McpHandlerCache {
cached: Mutex<Option<CachedMcpHandlers>>,
handlers: Mutex<HashMap<ToolName, CachedMcpHandler>>,
}
struct CachedMcpHandlers {
binding: Weak<McpBinding>,
handlers: HashMap<ToolName, Arc<McpHandler>>,
struct CachedMcpHandler {
tool_info: McpToolInfo,
agent_plugin: bool,
schema_max_bytes: Option<NonZeroUsize>,
handler: Arc<McpHandler>,
}
impl McpHandlerCache {
@@ -43,29 +44,17 @@ impl McpHandlerCache {
search_tool_enabled: bool,
registry: &mut ToolRegistry,
) -> HashSet<ToolName> {
let mut cached = self
.cached
let mut handlers = self
.handlers
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if !cached
.as_ref()
.and_then(|cached| cached.binding.upgrade())
.is_some_and(|cached_binding| Arc::ptr_eq(&cached_binding, binding))
{
*cached = None;
}
let cached = cached.get_or_insert_with(|| CachedMcpHandlers {
binding: Arc::downgrade(binding),
handlers: HashMap::new(),
});
append_mcp_tools(
binding.tools(),
config,
apps_enabled,
mcp_server_catalog,
search_tool_enabled,
&mut cached.handlers,
&mut handlers,
registry,
)
}
@@ -78,9 +67,15 @@ fn append_mcp_tools(
apps_enabled: bool,
mcp_server_catalog: &codex_mcp::ResolvedMcpCatalog,
search_tool_enabled: bool,
handlers: &mut HashMap<ToolName, Arc<McpHandler>>,
handlers: &mut HashMap<ToolName, CachedMcpHandler>,
registry: &mut ToolRegistry,
) -> HashSet<ToolName> {
let current_tools: HashSet<_> = all_mcp_tools
.iter()
.map(McpToolInfo::canonical_tool_name)
.collect();
handlers.retain(|name, _| current_tools.contains(name));
// Keep regular MCP tools first; Apps tools also require connector and policy checks.
let non_app_tools = filter_non_codex_apps_mcp_tools_only(all_mcp_tools);
let app_tools = apps_enabled
@@ -100,25 +95,41 @@ fn append_mcp_tools(
let agent_plugin = server.is_some_and(|server| server.source().is_agent_plugin());
let tool_input_schema_max_bytes =
server.and_then(|server| server.config().tool_input_schema_max_bytes);
let handler = match handlers.entry(tool_name.clone()) {
Entry::Occupied(entry) => Arc::clone(entry.get()),
Entry::Vacant(entry) => {
let handler = if agent_plugin {
McpHandler::new_agent_plugin(tool.clone())
} else if let Some(budget) = tool_input_schema_max_bytes {
McpHandler::new_with_schema_max_bytes(tool.clone(), budget.get())
} else {
McpHandler::new(tool.clone())
};
// Handlers contain immutable tool metadata, not a connection or authorization snapshot.
// Preserve their identity across equivalent bindings so the search index can also be reused.
let handler = if let Some(cached) = handlers.get(&tool_name).filter(|cached| {
cached.tool_info == *tool
&& cached.agent_plugin == agent_plugin
&& cached.schema_max_bytes == tool_input_schema_max_bytes
}) {
Arc::clone(&cached.handler)
} else {
handlers.remove(&tool_name);
let handler = if agent_plugin {
McpHandler::new_agent_plugin(tool.clone())
} else if let Some(budget) = tool_input_schema_max_bytes {
McpHandler::new_with_schema_max_bytes(tool.clone(), budget.get())
} else {
McpHandler::new(tool.clone())
};
match handler {
Ok(handler) => Arc::clone(entry.insert(Arc::new(handler))),
Err(err) => {
warn!("Skipping MCP tool `{tool_name}`: failed to build tool spec: {err}");
continue;
}
let handler = match handler {
Ok(handler) => Arc::new(handler),
Err(err) => {
warn!("Skipping MCP tool `{tool_name}`: failed to build tool spec: {err}");
continue;
}
}
};
handlers.insert(
tool_name.clone(),
CachedMcpHandler {
tool_info: tool.clone(),
agent_plugin,
schema_max_bytes: tool_input_schema_max_bytes,
handler: Arc::clone(&handler),
},
);
handler
};
let fits_agent_budget = if agent_plugin {
codex-rs/core/tests/suite/mcp_tool_exposure.rs
+69 -1
View File
@@ -59,10 +59,18 @@ use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tempfile::TempDir;
use tokio::sync::Semaphore;
use tracing::Subscriber;
use tracing::span::Attributes;
use tracing::span::Id;
use tracing_subscriber::Layer;
use tracing_subscriber::layer::Context as LayerContext;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use wiremock::Mock;
use wiremock::Request;
use wiremock::ResponseTemplate;
@@ -70,6 +78,27 @@ use wiremock::matchers::body_partial_json;
use wiremock::matchers::method;
use wiremock::matchers::path_regex;
#[derive(Clone, Default)]
struct McpCacheCounters {
binding_captures: Arc<AtomicUsize>,
search_index_builds: Arc<AtomicUsize>,
}
impl<S: Subscriber> Layer<S> for McpCacheCounters {
fn on_new_span(&self, attributes: &Attributes<'_>, _id: &Id, _context: LayerContext<'_, S>) {
let metadata = attributes.metadata();
match (metadata.target(), metadata.name()) {
("codex_mcp::connection_manager::tool_catalog", "capture_binding_with_metadata") => {
self.binding_captures.fetch_add(1, Ordering::SeqCst);
}
("codex_core::tools::handlers::tool_search", "new") => {
self.search_index_builds.fetch_add(1, Ordering::SeqCst);
}
_ => {}
}
}
}
struct McpResourceClientCapture {
client: Arc<Mutex<Option<McpResourceClient>>>,
}
@@ -1080,10 +1109,18 @@ async fn deferred_tool_world_state_is_disabled_by_default() -> Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
// Keep spawned tasks on the thread with the scoped tracing subscriber.
#[tokio::test(flavor = "current_thread")]
async fn deferred_tool_world_state_tracks_initial_unchanged_and_removed_namespaces() -> Result<()> {
skip_if_no_network!(Ok(()));
let counters = McpCacheCounters::default();
// Keep concurrent tests without a subscriber from caching these callsites as disabled.
let _interest_cache_guard =
tracing::Dispatch::new(tracing::subscriber::NoSubscriber::default());
let _tracing = tracing_subscriber::registry()
.with(counters.clone())
.set_default();
let server = responses::start_mock_server().await;
let apps_server = AppsTestServer::mount_searchable(&server).await?;
let response = mount_sse_sequence(&server, completed_response_sequence(/*count*/ 3)).await;
@@ -1094,7 +1131,29 @@ async fn deferred_tool_world_state_tracks_initial_unchanged_and_removed_namespac
test.submit_turn("inspect initially available deferred tools")
.await?;
let initial_captures = counters.binding_captures.load(Ordering::SeqCst);
let initial_index_builds = counters.search_index_builds.load(Ordering::SeqCst);
assert!(
initial_captures > 0,
"the initial turn must capture an MCP binding"
);
assert!(
initial_index_builds > 0,
"the initial turn must build the search index"
);
// Publish a new catalog revision with the same metadata from the ready client.
test.codex.refresh_codex_apps_tools().await?;
test.submit_turn("inspect unchanged deferred tools").await?;
assert!(
counters.binding_captures.load(Ordering::SeqCst) > initial_captures,
"the follow-up must capture a new binding after the refresh"
);
assert_eq!(
counters.search_index_builds.load(Ordering::SeqCst),
initial_index_builds,
"equivalent bindings must preserve MCP handlers and reuse the search index"
);
let mut refresh_config = test.config.clone();
let user_config_path = refresh_config.codex_home.join("config.toml");
@@ -1113,6 +1172,14 @@ enabled = false
let requests = response.requests();
assert_eq!(requests.len(), 3);
assert!(
requests[2].body_json()["tools"]
.as_array()
.expect("model request tools")
.iter()
.all(|tool| tool["type"] != "tool_search"),
"removing all deferred tools must stop advertising the cached search tool"
);
let tools_states = requests
.iter()
.map(tools_state_sections)
@@ -1133,6 +1200,7 @@ enabled = false
)
);
test.codex.shutdown_and_wait().await?;
Ok(())
}