From e1831db7c31ec3ec3a88f29e95cc336858038bad Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 6 Aug 2026 13:31:14 +0000 Subject: [PATCH] Reuse MCP handlers across sampling steps (#37273) ## Why MCP tool handlers and their Code Mode definitions are immutable for the lifetime of a stable MCP binding, so rebuilding them for every sampling step repeats schema construction unnecessarily. ## What changed - Cache MCP handlers per session and reuse them while the same binding remains active, invalidating the cache when the binding changes. - Lazily cache schema-free Code Mode definitions on MCP handlers while continuing to rebuild definitions for step-local tools whose specifications may change. - Reapply current app enablement, tool policy, exposure, and plugin schema budgets whenever cached handlers are registered for a step. ## Testing - Cover handler and Code Mode definition reuse, binding invalidation, changing app policy, stable continuation schemas, and step-local custom tool updates. GitOrigin-RevId: 69c5011fb97c08580ac476e7fb2286b1e92e0359 --- codex-rs/core/src/mcp_tool_exposure.rs | 127 +++++++++++++----- codex-rs/core/src/mcp_tool_exposure_test.rs | 87 ++++++++++++ codex-rs/core/src/session/mod.rs | 2 +- codex-rs/core/src/session/session.rs | 1 + codex-rs/core/src/session/tests.rs | 31 +++++ codex-rs/core/src/session/turn.rs | 2 +- codex-rs/core/src/state/service.rs | 3 + .../src/tools/code_mode/execute_handler.rs | 29 +++- codex-rs/core/src/tools/handlers/mcp.rs | 61 ++++++++- codex-rs/core/src/tools/registry.rs | 25 +++- codex-rs/core/src/tools/spec_plan.rs | 92 +++++++++---- codex-rs/core/tests/suite/code_mode.rs | 50 +++++-- codex-rs/core/tests/suite/mcp_tool_cache.rs | 14 +- 13 files changed, 439 insertions(+), 85 deletions(-) diff --git a/codex-rs/core/src/mcp_tool_exposure.rs b/codex-rs/core/src/mcp_tool_exposure.rs index 84d14ebcd3..43b490ecc2 100644 --- a/codex-rs/core/src/mcp_tool_exposure.rs +++ b/codex-rs/core/src/mcp_tool_exposure.rs @@ -1,9 +1,14 @@ +use std::collections::HashMap; use std::collections::HashSet; +use std::collections::hash_map::Entry; use std::sync::Arc; +use std::sync::Mutex; +use std::sync::Weak; use codex_connectors::AppToolPolicyEvaluator; use codex_connectors::AppToolPolicyInput; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_mcp::McpBinding; use codex_mcp::ToolInfo as McpToolInfo; use codex_mcp::tool_is_model_visible; use codex_tools::ToolExposure; @@ -18,13 +23,62 @@ use crate::config::Config; use crate::tools::handlers::McpHandler; use crate::tools::registry::ToolRegistry; +#[derive(Default)] +pub(crate) struct McpHandlerCache { + cached: Mutex>, +} + +struct CachedMcpHandlers { + binding: Weak, + handlers: HashMap>, +} + +impl McpHandlerCache { + pub(crate) fn append_mcp_tools( + &self, + binding: &Arc, + config: &Config, + apps_enabled: bool, + mcp_server_catalog: &codex_mcp::ResolvedMcpCatalog, + search_tool_enabled: bool, + registry: &mut ToolRegistry, + ) -> HashSet { + let mut cached = self + .cached + .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, + registry, + ) + } +} + #[instrument(level = "trace", skip_all)] -pub(crate) fn append_mcp_tools( +fn append_mcp_tools( all_mcp_tools: &[McpToolInfo], config: &Config, apps_enabled: bool, mcp_server_catalog: &codex_mcp::ResolvedMcpCatalog, search_tool_enabled: bool, + handlers: &mut HashMap>, registry: &mut ToolRegistry, ) -> HashSet { // Keep regular MCP tools first; Apps tools also require connector and policy checks. @@ -45,43 +99,48 @@ pub(crate) fn append_mcp_tools( let agent_plugin = mcp_server_catalog .server(&tool.server_name) .is_some_and(|server| server.source().is_agent_plugin()); - let handler = if agent_plugin { - McpHandler::new_agent_plugin(tool.clone()) - } else { - McpHandler::new(tool.clone()) - }; - match handler { - Ok(handler) => { - let fits_agent_budget = if agent_plugin { - handler.model_spec_bytes().is_ok_and(|bytes| { - if bytes > MAX_AGENT_PLUGIN_MCP_SPEC_BYTES { - return false; - } - let next = agent_plugin_bytes.saturating_add(bytes); - if next <= MAX_AGENT_PLUGIN_MCP_TOTAL_BYTES { - agent_plugin_bytes = next; - true - } else { - false - } - }) + 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 { - true + McpHandler::new(tool.clone()) }; - let tool_exposure = if fits_agent_budget { - exposure - } else { - ToolExposure::Hidden - }; - if registry.register_external_with_exposure(Arc::new(handler), tool_exposure) - && fits_agent_budget - { - registered_tools.insert(tool_name); + + 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; + } } } - Err(err) => { - warn!("Skipping MCP tool `{tool_name}`: failed to build tool spec: {err}"); - } + }; + + let fits_agent_budget = if agent_plugin { + handler.model_spec_bytes().is_ok_and(|bytes| { + if bytes > MAX_AGENT_PLUGIN_MCP_SPEC_BYTES { + return false; + } + let next = agent_plugin_bytes.saturating_add(bytes); + if next <= MAX_AGENT_PLUGIN_MCP_TOTAL_BYTES { + agent_plugin_bytes = next; + true + } else { + false + } + }) + } else { + true + }; + let tool_exposure = if fits_agent_budget { + exposure + } else { + ToolExposure::Hidden + }; + if registry.register_external_with_exposure(handler, tool_exposure) && fits_agent_budget { + registered_tools.insert(tool_name); } } registered_tools diff --git a/codex-rs/core/src/mcp_tool_exposure_test.rs b/codex-rs/core/src/mcp_tool_exposure_test.rs index 9f709f8dcb..c94be8ce55 100644 --- a/codex-rs/core/src/mcp_tool_exposure_test.rs +++ b/codex-rs/core/src/mcp_tool_exposure_test.rs @@ -95,6 +95,7 @@ fn runtimes_by_name_with_catalog( mcp_server_catalog: &ResolvedMcpCatalog, search_tool_enabled: bool, ) -> HashMap { + let mut handlers = HashMap::new(); let mut registry = ToolRegistry::default(); append_mcp_tools( tools, @@ -102,6 +103,7 @@ fn runtimes_by_name_with_catalog( apps_enabled, mcp_server_catalog, search_tool_enabled, + &mut handlers, &mut registry, ); registry @@ -220,6 +222,89 @@ async fn directly_exposes_effective_tool_sets_when_search_is_unavailable() { ); } +#[tokio::test] +async fn cached_app_handlers_still_obey_current_apps_enablement_and_tool_policy() { + let config = test_config().await; + let codex_home = tempdir().expect("create restrictive config directory"); + std::fs::write( + codex_home.path().join(CONFIG_TOML_FILE), + "[apps.calendar]\ndefault_tools_enabled = false\n", + ) + .expect("write restrictive app policy"); + let restricted_config = ConfigBuilder::without_managed_config_for_tests() + .codex_home(codex_home.path().to_path_buf()) + .build() + .await + .expect("build restrictive app policy"); + let tools = [make_mcp_tool( + CODEX_APPS_MCP_SERVER_NAME, + "events/create", + "mcp__codex_apps__calendar", + "create", + Some("calendar"), + Some("Calendar"), + )]; + let mut handlers = HashMap::new(); + let catalog = ResolvedMcpCatalog::default(); + let mut allowed_registry = ToolRegistry::default(); + let allowed = append_mcp_tools( + &tools, + &config, + /*apps_enabled*/ true, + &catalog, + /*search_tool_enabled*/ false, + &mut handlers, + &mut allowed_registry, + ); + let cached_handler = &allowed_registry + .entries() + .next() + .expect("allowed app tool should be registered") + .runtime; + + let mut disabled_registry = ToolRegistry::default(); + let disabled = append_mcp_tools( + &tools, + &config, + /*apps_enabled*/ false, + &catalog, + /*search_tool_enabled*/ false, + &mut handlers, + &mut disabled_registry, + ); + let mut restricted_registry = ToolRegistry::default(); + let restricted = append_mcp_tools( + &tools, + &restricted_config, + /*apps_enabled*/ true, + &catalog, + /*search_tool_enabled*/ false, + &mut handlers, + &mut restricted_registry, + ); + let mut restored_registry = ToolRegistry::default(); + let restored = append_mcp_tools( + &tools, + &config, + /*apps_enabled*/ true, + &catalog, + /*search_tool_enabled*/ true, + &mut handlers, + &mut restored_registry, + ); + let restored_handler = restored_registry + .entries() + .next() + .expect("restored app tool should be registered"); + + assert_eq!(allowed, HashSet::from([tools[0].canonical_tool_name()])); + assert!(disabled.is_empty()); + assert!(restricted.is_empty()); + assert_eq!(restored, allowed); + assert!(Arc::ptr_eq(cached_handler, &restored_handler.runtime)); + assert_eq!(restored_handler.exposure, ToolExposure::Deferred); +} + #[tokio::test] async fn excludes_tools_hidden_from_model_exposure() { let config = test_config().await; @@ -339,6 +424,7 @@ async fn app_tool_registration_uses_trusted_catalog_metadata_and_preserves_sourc synthetic_app_tool.clone(), regular_tool.clone(), ]; + let mut handlers = HashMap::new(); let mut registry = ToolRegistry::default(); append_mcp_tools( @@ -347,6 +433,7 @@ async fn app_tool_registration_uses_trusted_catalog_metadata_and_preserves_sourc /*apps_enabled*/ true, &ResolvedMcpCatalog::default(), /*search_tool_enabled*/ false, + &mut handlers, &mut registry, ); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index f69ff50476..4db2ba4dac 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3143,7 +3143,7 @@ impl Session { self.as_ref(), turn_context.as_ref(), &environments, - mcp.as_ref(), + &mcp, &extension_data, prepared_recommendations, ) diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 522055b7b2..ed99e8d01f 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1115,6 +1115,7 @@ impl Session { // Start with an empty connection set. The initialized set is // published after SessionConfigured so MCP events follow it. mcp_runtime, + mcp_handler_cache: Default::default(), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 3a0ae6850a..8048bf525e 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5750,6 +5750,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { .then(|| Arc::new(crate::state::ExecutedToolCallRecorder::default())); let services = SessionServices { mcp_runtime, + mcp_handler_cache: Default::default(), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), @@ -7968,6 +7969,7 @@ where .then(|| Arc::new(crate::state::ExecutedToolCallRecorder::default())); let services = SessionServices { mcp_runtime, + mcp_handler_cache: Default::default(), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), @@ -8183,6 +8185,18 @@ async fn refresh_mcp_servers_uses_latest_state_for_existing_turns() { .capture_step_context(next_turn, &CancellationToken::new()) .await .expect("a fresh cancellation token cannot be cancelled"); + assert!( + !Arc::ptr_eq(&old_step.mcp, &new_step.mcp), + "publishing a new MCP runtime must invalidate cached bindings" + ); + let refreshed_old_step = session + .capture_step_context(Arc::clone(&turn_context), &CancellationToken::new()) + .await + .expect("capture an existing turn after its MCP runtime is republished"); + assert!( + Arc::ptr_eq(&new_step.mcp, &refreshed_old_step.mcp), + "existing turns should reuse the newly published immutable MCP binding" + ); let rematerialized_old = session .mcp_runtime_for_step( &turn_context, @@ -8220,12 +8234,17 @@ async fn refresh_mcp_servers_uses_latest_state_for_existing_turns() { async fn refreshed_mcp_binding_captures_current_approval_authority() { let (session, old_turn) = make_session_and_context().await; let session = Arc::new(session); + let old_turn = Arc::new(old_turn); let previous_policy = old_turn.approval_policy(); assert_ne!(previous_policy, AskForApproval::Never); assert_eq!( old_turn.config.permissions.approval_policy.value(), previous_policy ); + let old_step = session + .capture_step_context(Arc::clone(&old_turn), &CancellationToken::new()) + .await + .expect("capture initial sampling step"); session .update_settings(SessionSettingsUpdate { @@ -8244,6 +8263,18 @@ async fn refreshed_mcp_binding_captures_current_approval_authority() { .current_binding() .await .expect("refreshed runtime should be available"); + assert!( + !Arc::ptr_eq(&old_step.mcp, &binding), + "changed approval authority must invalidate the cached MCP binding" + ); + let refreshed_step = session + .capture_step_context(Arc::clone(&old_turn), &CancellationToken::new()) + .await + .expect("capture existing turn after its approval authority changes"); + assert!( + Arc::ptr_eq(&binding, &refreshed_step.mcp), + "existing turns must use the MCP binding with current approval authority" + ); let config = binding.config(); assert_eq!( ( diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index eb18e3b24f..5cb61ef7b7 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1475,7 +1475,7 @@ pub(crate) async fn built_tools( sess: &Session, turn_context: &TurnContext, environments: &TurnEnvironmentSnapshot, - mcp: &codex_mcp::McpBinding, + mcp: &Arc, step_store: &ExtensionData, prepared_recommendations: PreparedToolRecommendations, ) -> CodexResult> { diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index dd17a076c1..1e3cb00bca 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -14,6 +14,7 @@ use crate::environment_selection::ThreadEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejectionCircuitBreaker; use crate::mcp::McpManager; +use crate::mcp_tool_exposure::McpHandlerCache; use crate::tools::ExecutedToolCallRecorder; use crate::tools::code_mode::CodeModeService; use crate::tools::handlers::ToolSearchHandlerCache; @@ -45,6 +46,8 @@ use tokio::sync::Mutex; pub(crate) struct SessionServices { /// The single owner of live MCP connections for this thread. pub(crate) mcp_runtime: Arc, + /// Immutable MCP handlers scoped to this thread's current binding. + pub(crate) mcp_handler_cache: McpHandlerCache, pub(crate) unified_exec_manager: UnifiedExecProcessManager, pub(crate) elicitations: ElicitationService, #[cfg_attr(not(unix), allow(dead_code))] diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index e6ea6ff461..87d260be4e 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -7,6 +7,7 @@ use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; use codex_tools::ToolName; use codex_tools::ToolSpec; +use std::sync::Arc; use super::ExecContext; use super::PUBLIC_TOOL_NAME; @@ -14,13 +15,15 @@ use super::handle_runtime_response; use super::is_exec_tool_name; use super::telemetry::CodeModeToolCallGuard; +type CodeModeNestedTool = (Arc, Option>); + pub struct CodeModeExecuteHandler { spec: ToolSpec, - nested_tool_specs: Vec, + nested_tool_specs: Vec, } impl CodeModeExecuteHandler { - pub(crate) fn new(spec: ToolSpec, nested_tool_specs: Vec) -> Self { + pub(crate) fn new(spec: ToolSpec, nested_tool_specs: Vec) -> Self { Self { spec, nested_tool_specs, @@ -38,8 +41,26 @@ impl CodeModeExecuteHandler { let args = codex_code_mode::parse_exec_source(&code).map_err(FunctionCallError::RespondToModel)?; let exec = ExecContext { session, turn }; - let enabled_tools = - codex_tools::collect_code_mode_tool_definitions(&self.nested_tool_specs); + let mut enabled_tools = Vec::with_capacity(self.nested_tool_specs.len()); + for (spec, cached_runtime) in &self.nested_tool_specs { + if let Some(cached_definitions) = cached_runtime + .as_ref() + .and_then(|runtime| runtime.cached_code_mode_definitions()) + { + enabled_tools.extend_from_slice(cached_definitions); + continue; + } + + let definitions = + codex_tools::collect_code_mode_tool_definitions(std::iter::once(spec.as_ref())); + enabled_tools.extend(definitions.into_iter().map(|mut definition| { + definition.input_schema = None; + definition.output_schema = None; + definition + })); + } + enabled_tools.sort_by(|left, right| left.name.cmp(&right.name)); + enabled_tools.dedup_by(|left, right| left.name == right.name); let started_at = std::time::Instant::now(); let started_cell = exec .session diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 7a67372737..b6c9324022 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::sync::OnceLock; use std::time::Instant; use crate::function_tool::FunctionCallError; @@ -37,7 +38,8 @@ const MAX_MCP_NAMESPACE_DESCRIPTION_BYTES: usize = 512 * 1024; pub struct McpHandler { tool_info: ToolInfo, - spec: ToolSpec, + spec: Arc, + code_mode_tool_definitions: OnceLock>, } impl McpHandler { @@ -66,8 +68,12 @@ impl McpHandler { .to_string() }); } - let spec = create_tool_spec(&tool_info, agent_plugin)?; - Ok(Self { tool_info, spec }) + let spec = Arc::new(create_tool_spec(&tool_info, agent_plugin)?); + Ok(Self { + tool_info, + spec, + code_mode_tool_definitions: OnceLock::new(), + }) } pub(crate) fn model_spec_bytes(&self) -> Result { @@ -104,7 +110,7 @@ impl ToolExecutor for McpHandler { } fn spec(&self) -> ToolSpec { - self.spec.clone() + self.spec.as_ref().clone() } fn supports_parallel_tool_calls(&self) -> bool { @@ -197,6 +203,27 @@ impl McpHandler { } impl CoreToolRuntime for McpHandler { + fn immutable_spec(&self) -> Option<&Arc> { + Some(&self.spec) + } + + fn cached_code_mode_definitions(&self) -> Option<&[codex_code_mode::ToolDefinition]> { + Some( + self.code_mode_tool_definitions + .get_or_init(|| { + let mut definitions = codex_tools::collect_code_mode_tool_definitions( + std::iter::once(self.spec.as_ref()), + ); + for definition in &mut definitions { + definition.input_schema = None; + definition.output_schema = None; + } + definitions + }) + .as_slice(), + ) + } + fn wait_until_ready<'a>(&'a self, session: &'a Arc) -> Option> { Some(Box::pin(async move { session @@ -536,6 +563,32 @@ mod tests { ); } + #[test] + fn mcp_code_mode_definitions_are_cached_lazily() { + let handler = McpHandler::new(tool_info("filesystem", "mcp__filesystem", "read_file")) + .expect("MCP tool spec should build"); + + assert!(handler.code_mode_tool_definitions.get().is_none()); + assert!(Arc::ptr_eq( + handler + .immutable_spec() + .expect("MCP spec should be immutable"), + &handler.spec, + )); + + let first = handler + .cached_code_mode_definitions() + .expect("MCP definitions should be cached"); + assert_eq!(first.len(), 1); + assert!(first[0].input_schema.is_none()); + assert!(first[0].output_schema.is_none()); + + let second = handler + .cached_code_mode_definitions() + .expect("MCP definitions should be cached"); + assert!(std::ptr::eq(first, second)); + } + #[test] fn mcp_read_only_hint_supports_parallel_calls_without_server_opt_in() { let mut read_only_info = tool_info("foo", "mcp__foo__", "read"); diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index fd41dcb11c..b603512dd6 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -51,6 +51,16 @@ pub use codex_tools::ToolExposure; /// Implementers provide the shared `ToolExecutor` behavior plus optional /// core-owned metadata for hooks, telemetry, tool search, and argument diffs. pub(crate) trait CoreToolRuntime: ToolExecutor { + /// Returns a shared spec when both the spec and search metadata are immutable. + fn immutable_spec(&self) -> Option<&Arc> { + None + } + + /// Returns lazily cached Code Mode definitions owned by this runtime. + fn cached_code_mode_definitions(&self) -> Option<&[codex_code_mode::ToolDefinition]> { + None + } + /// Returns a readiness wait for this exact tool before taking the execution gate. fn wait_until_ready<'a>(&'a self, _session: &'a Arc) -> Option> { None @@ -385,15 +395,22 @@ impl ToolRegistry { if !existing_description.trim().is_empty() { continue; } - let description = match tool.runtime.spec() { - ToolSpec::Namespace(namespace) => namespace.description, + let owned_spec; + let spec = if let Some(spec) = tool.runtime.immutable_spec() { + spec.as_ref() + } else { + owned_spec = tool.runtime.spec(); + &owned_spec + }; + let description = match spec { + ToolSpec::Namespace(namespace) => namespace.description.as_str(), ToolSpec::Function(_) | ToolSpec::Freeform(_) | ToolSpec::ToolSearch { .. } - | ToolSpec::WebSearch { .. } => String::new(), + | ToolSpec::WebSearch { .. } => "", }; if !description.trim().is_empty() { - *existing_description = description; + *existing_description = description.to_string(); } } namespaces diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 0e94988ec1..1b11729175 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -2,7 +2,6 @@ use crate::agent::exceeds_thread_spawn_depth_limit; use crate::agent::next_thread_spawn_depth; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::image_preparation::unified_image_budget_enabled; -use crate::mcp_tool_exposure::append_mcp_tools; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::tools::code_mode::default_exec_yield_time_override_ms; @@ -120,7 +119,7 @@ pub(crate) fn build_tool_router( session: &Session, turn_context: &TurnContext, environments: &TurnEnvironmentSnapshot, - mcp: &codex_mcp::McpBinding, + mcp: &Arc, apps_enabled: bool, step_store: &ExtensionData, tool_suggest_candidates: Option<&crate::tools::router::ToolSuggestCandidates>, @@ -147,8 +146,8 @@ pub(crate) fn build_tool_router( { Vec::new() } else { - let registered_mcp_tools = append_mcp_tools( - mcp.tools(), + let registered_mcp_tools = session.services.mcp_handler_cache.append_mcp_tools( + mcp, &turn_context.config, apps_enabled, &mcp.config().mcp_server_catalog, @@ -350,7 +349,14 @@ pub(crate) fn finalize_tool_router( let conflicting_tool_names = registry .entries() .filter_map(|tool| { - let ToolSpec::Namespace(namespace) = tool.runtime.spec() else { + let owned_spec; + let spec = if let Some(spec) = tool.runtime.immutable_spec() { + spec.as_ref() + } else { + owned_spec = tool.runtime.spec(); + &owned_spec + }; + let ToolSpec::Namespace(namespace) = spec else { return None; }; (namespace.name == tool_search_name.name).then(|| tool.runtime.tool_name()) @@ -375,15 +381,22 @@ pub(crate) fn finalize_tool_router( let mut namespace_descriptions = BTreeMap::new(); for tool in registry.entries() { - let ToolSpec::Namespace(namespace) = tool.runtime.spec() else { + let owned_spec; + let spec = if let Some(spec) = tool.runtime.immutable_spec() { + spec.as_ref() + } else { + owned_spec = tool.runtime.spec(); + &owned_spec + }; + let ToolSpec::Namespace(namespace) = spec else { continue; }; if namespace.description.trim().is_empty() { continue; } - match namespace_descriptions.entry(namespace.name) { + match namespace_descriptions.entry(namespace.name.clone()) { Entry::Vacant(entry) => { - entry.insert(namespace.description); + entry.insert(namespace.description.clone()); } Entry::Occupied(entry) if entry.get() != &namespace.description => { return Err(CodexErrorDetails::ToolCollision(entry.key().clone()).into()); @@ -402,6 +415,15 @@ fn apply_direct_model_only_namespace_overrides( turn_context: &TurnContext, registry: &mut ToolRegistry, ) { + if turn_context + .config + .code_mode + .direct_only_tool_namespaces + .is_empty() + { + return; + } + for tool in registry.entries_mut() { let configured = tool .runtime @@ -431,19 +453,24 @@ fn build_model_visible_specs( ) -> Vec { let mut specs = Vec::new(); for tool in registry.entries() { - let tool_name = tool.runtime.tool_name(); let exposure = tool.exposure; - if exposure.is_direct() && !is_hidden_by_code_mode_only(turn_context, &tool_name, exposure) - { - let spec = tool.runtime.spec(); - specs.push(spec_for_model_request( - turn_context, - exposure, - &tool_name, - code_mode_tool_names, - spec, - )); + if !exposure.is_direct() { + continue; } + + let tool_name = tool.runtime.tool_name(); + if is_hidden_by_code_mode_only(turn_context, &tool_name, exposure) { + continue; + } + + let spec = tool.runtime.spec(); + specs.push(spec_for_model_request( + turn_context, + exposure, + &tool_name, + code_mode_tool_names, + spec, + )); } specs.extend(hosted_specs); @@ -653,6 +680,7 @@ fn register_code_mode_executors( let mut code_mode_nested_tool_specs = Vec::new(); let mut exec_prompt_tool_specs = Vec::new(); let mut deferred_exec_prompt_tool_specs = Vec::new(); + let mut included_deferred_mcp_output_schema = false; let deferred_tools_guidance_enabled = search_tool_enabled(turn_context); for tool in registry.entries() { let exposure = tool.exposure; @@ -665,9 +693,13 @@ fn register_code_mode_executors( continue; } - let spec = tool.runtime.spec(); + let immutable_spec = tool.runtime.immutable_spec(); + let cached_runtime = immutable_spec.map(|_| Arc::clone(&tool.runtime)); + let spec = immutable_spec + .map(Arc::clone) + .unwrap_or_else(|| Arc::new(tool.runtime.spec())); // Derive the name without serializing and augmenting every tool schema. - let code_mode_name = match &spec { + let code_mode_name = match spec.as_ref() { ToolSpec::Function(_) | ToolSpec::Freeform(_) => { codex_tools::code_mode_name_for_tool_name(&tool_name) } @@ -698,13 +730,18 @@ fn register_code_mode_executors( } if exposure == ToolExposure::Deferred { - if deferred_tools_guidance_enabled { - deferred_exec_prompt_tool_specs.push(spec.clone()); + if deferred_tools_guidance_enabled + && (cached_runtime.is_none() || !included_deferred_mcp_output_schema) + { + if cached_runtime.is_some() { + included_deferred_mcp_output_schema = true; + } + deferred_exec_prompt_tool_specs.push(Arc::clone(&spec)); } } else { - exec_prompt_tool_specs.push(spec.clone()); + exec_prompt_tool_specs.push(spec.as_ref().clone()); } - code_mode_nested_tool_specs.push(spec); + code_mode_nested_tool_specs.push((spec, cached_runtime)); } if turn_context.model_info.use_responses_lite { @@ -716,10 +753,11 @@ fn register_code_mode_executors( let namespace_descriptions = code_mode_namespace_descriptions(&exec_prompt_tool_specs); let mut enabled_tools = collect_code_mode_exec_prompt_tool_definitions(exec_prompt_tool_specs.iter()); + let deferred_tools = collect_code_mode_exec_prompt_tool_definitions( + deferred_exec_prompt_tool_specs.iter().map(Arc::as_ref), + ); enabled_tools .sort_by(|left, right| compare_code_mode_tools(left, right, &namespace_descriptions)); - let deferred_tools = - collect_code_mode_exec_prompt_tool_definitions(deferred_exec_prompt_tool_specs.iter()); let default_exec_yield_time_ms = default_exec_yield_time_override_ms(&turn_context.config.features) .unwrap_or(codex_code_mode::DEFAULT_EXEC_YIELD_TIME_MS); diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 9e3ade1f61..d8b6fabbb3 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -80,6 +80,8 @@ use std::fs; use std::io::Cursor; use std::path::Path; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; use test_case::test_case; @@ -4175,7 +4177,10 @@ text(JSON.stringify({ Ok(()) } -struct NamespacedCustomTool; +struct NamespacedCustomTool { + generation: usize, + generations: Arc, +} impl ToolContributor for NamespacedCustomTool { fn tools( @@ -4183,7 +4188,19 @@ impl ToolContributor for NamespacedCustomTool { _session_store: &ExtensionData, _thread_store: &ExtensionData, ) -> Vec>> { - vec![Arc::new(Self)] + Vec::new() + } + + fn tools_for_step( + &self, + _session_store: &ExtensionData, + _thread_store: &ExtensionData, + _step_store: &ExtensionData, + ) -> Vec>> { + vec![Arc::new(Self { + generation: self.generations.fetch_add(1, Ordering::Relaxed) + 1, + generations: Arc::clone(&self.generations), + })] } } @@ -4198,7 +4215,7 @@ impl ToolExecutor for NamespacedCustomTool { description: "Editing tools.".to_string(), tools: vec![ResponsesApiNamespaceTool::Custom(FreeformTool { name: "apply_patch".to_string(), - description: "Apply a raw editor patch.".to_string(), + description: format!("Apply a raw editor patch (step {}).", self.generation), defer_loading: None, format: FreeformToolFormat { r#type: "grammar".to_string(), @@ -4220,6 +4237,7 @@ impl ToolExecutor for NamespacedCustomTool { "namespace": call.tool_name.namespace, "name": call.tool_name.name, "input": input, + "generation": self.generation, }))) as Box) }) } @@ -4231,7 +4249,10 @@ async fn code_mode_exposes_and_dispatches_namespaced_custom_tools() -> Result<() let server = responses::start_mock_server().await; let mut extensions = ExtensionRegistryBuilder::::new(); - extensions.tool_contributor(Arc::new(NamespacedCustomTool)); + extensions.tool_contributor(Arc::new(NamespacedCustomTool { + generation: 0, + generations: Arc::new(AtomicUsize::new(0)), + })); let mut builder = test_codex() .with_model("test-gpt-5.1-codex") .with_extensions(Arc::new(extensions.build())) @@ -4283,8 +4304,11 @@ text(JSON.stringify({ let declaration = "declare const tools: { editor__apply_patch(input: string): Promise; };"; - let description = - format!("Apply a raw editor patch.\n\nexec tool declaration:\n```ts\n{declaration}\n```"); + let description = |generation| { + format!( + "Apply a raw editor patch (step {generation}).\n\nexec tool declaration:\n```ts\n{declaration}\n```" + ) + }; let first_body = requests[0].body_json(); let namespaced_custom_tool = namespace_child_tool(&first_body, "editor", "apply_patch") .expect("namespaced custom tool should be included in the model request"); @@ -4293,7 +4317,7 @@ text(JSON.stringify({ &serde_json::json!({ "type": "custom", "name": "apply_patch", - "description": description, + "description": description(1), "format": { "type": "grammar", "syntax": "lark", @@ -4301,6 +4325,14 @@ text(JSON.stringify({ }, }) ); + let second_body = requests[1].body_json(); + let second_description = description(2); + assert_eq!( + namespace_child_tool(&second_body, "editor", "apply_patch") + .and_then(|tool| tool.get("description")) + .and_then(Value::as_str), + Some(second_description.as_str()) + ); let (direct_output, direct_success) = custom_tool_output_body_and_success(&requests[1], "call-direct"); @@ -4314,6 +4346,7 @@ text(JSON.stringify({ "namespace": "editor", "name": "apply_patch", "input": "direct patch", + "generation": 1, }) ); @@ -4327,11 +4360,12 @@ text(JSON.stringify({ exec_output, serde_json::json!({ "name": "editor__apply_patch", - "description": format!("Editing tools.\n\n{description}"), + "description": format!("Editing tools.\n\n{second_description}"), "result": { "namespace": "editor", "name": "apply_patch", "input": "nested patch", + "generation": 2, }, }) ); diff --git a/codex-rs/core/tests/suite/mcp_tool_cache.rs b/codex-rs/core/tests/suite/mcp_tool_cache.rs index 7d257d5714..8dd4863d9e 100644 --- a/codex-rs/core/tests/suite/mcp_tool_cache.rs +++ b/codex-rs/core/tests/suite/mcp_tool_cache.rs @@ -222,8 +222,13 @@ async fn mcp_calls_stay_bound_to_each_thread() -> anyhow::Result<()> { wait_for_event(thread, |event| matches!(event, EventMsg::TurnComplete(_))).await; let request = call_response.single_request(); assert!(request.tool_by_name(NAMESPACE, "echo").is_some()); - let output = completion_response - .single_request() + let completion_request = completion_response.single_request(); + assert_eq!( + request.body_json()["tools"], + completion_request.body_json()["tools"], + "MCP tool schemas must remain unchanged across a same-turn continuation" + ); + let output = completion_request .function_call_output_text(call_id) .expect("MCP result should be returned to the model"); assert!(output.contains(&process)); @@ -502,6 +507,11 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow .await?; let expected_error = format!("MCP tool `{SERVER_NAME}/cwd` is not available to the model"); assert_eq!(cached_turn.await??, second_process); + assert_definition( + &cached_done_response, + &format!("Use the tools from {second_process}."), + &format!("Echo from {second_process}."), + ); let output = cached_done_response .single_request() .function_call_output_text(app_only_call_id)