From 775fb21d2af9b9936618fe22dd62e6f0cb3ba4a3 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 31 Jul 2026 17:18:49 +0000 Subject: [PATCH] Keep effective tool exposure in the registry (#36367) ## Why `ToolExecutor::exposure()` describes a tool's preferred exposure, while the host may apply step-specific policy when building a tool plan. ## What changed - Store each runtime together with its effective exposure in `ToolRegistry`. - Apply MCP, namespace, and collaboration-tool exposure overrides directly to registry entries instead of wrapping runtimes. - Use the registry-owned exposure when building model-visible specs, deferred tool search, code-mode executors, and parallel-call support. ## Testing - Cover that a hidden MCP tool remains routable but is not eligible for parallel tool calls. GitOrigin-RevId: 1bc2760add8d993f477954269fa13681bfcc2444 --- codex-rs/core/src/mcp_tool_exposure.rs | 22 ++- codex-rs/core/src/mcp_tool_exposure_test.rs | 40 ++++-- codex-rs/core/src/session/tests.rs | 25 ++-- .../core/src/stream_events_utils_tests.rs | 8 +- codex-rs/core/src/tools/registry.rs | 136 +++++------------- codex-rs/core/src/tools/registry_tests.rs | 28 ++-- codex-rs/core/src/tools/router.rs | 6 +- codex-rs/core/src/tools/router_tests.rs | 59 +++++--- codex-rs/core/src/tools/spec_plan.rs | 97 +++++++------ codex-rs/core/src/tools/spec_plan_tests.rs | 82 +++++------ codex-rs/tools/src/tool_executor.rs | 1 + 11 files changed, 235 insertions(+), 269 deletions(-) diff --git a/codex-rs/core/src/mcp_tool_exposure.rs b/codex-rs/core/src/mcp_tool_exposure.rs index 6c77604a31..218dde9509 100644 --- a/codex-rs/core/src/mcp_tool_exposure.rs +++ b/codex-rs/core/src/mcp_tool_exposure.rs @@ -13,16 +13,16 @@ use tracing::warn; use crate::config::Config; use crate::connectors; use crate::tools::handlers::McpHandler; -use crate::tools::registry::CoreToolRuntime; -use crate::tools::registry::override_tool_exposure; +use crate::tools::registry::ToolRegistry; #[instrument(level = "trace", skip_all)] -pub(crate) fn build_mcp_tool_runtimes<'a>( - all_mcp_tools: &'a [McpToolInfo], - connectors: Option<&'a [connectors::AppInfo]>, - config: &'a Config, +pub(crate) fn append_mcp_tools( + all_mcp_tools: &[McpToolInfo], + connectors: Option<&[connectors::AppInfo]>, + config: &Config, search_tool_enabled: bool, -) -> impl Iterator> + 'a { + registry: &mut ToolRegistry, +) { // 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 = connectors @@ -33,19 +33,17 @@ pub(crate) fn build_mcp_tool_runtimes<'a>( } else { ToolExposure::Direct }; - non_app_tools.chain(app_tools).filter_map(move |tool| { + for tool in non_app_tools.chain(app_tools) { let tool_name = tool.canonical_tool_name(); match McpHandler::new(tool.clone()) { Ok(handler) => { - let handler: Arc = Arc::new(handler); - Some(override_tool_exposure(handler, exposure)) + registry.register_external_with_exposure(Arc::new(handler), exposure); } Err(err) => { warn!("Skipping MCP tool `{tool_name}`: failed to build tool spec: {err}"); - None } } - }) + } } fn filter_non_codex_apps_mcp_tools_only( diff --git a/codex-rs/core/src/mcp_tool_exposure_test.rs b/codex-rs/core/src/mcp_tool_exposure_test.rs index 8a7a6b4230..844bf16632 100644 --- a/codex-rs/core/src/mcp_tool_exposure_test.rs +++ b/codex-rs/core/src/mcp_tool_exposure_test.rs @@ -12,6 +12,7 @@ use rmcp::model::Tool; use super::*; use crate::config::CONFIG_TOML_FILE; +use crate::config::Config; use crate::config::ConfigBuilder; use crate::config::test_config; use crate::connectors::AppInfo; @@ -91,11 +92,22 @@ fn expected_runtimes( } fn runtimes_by_name( - runtimes: impl IntoIterator>, + tools: &[ToolInfo], + connectors: Option<&[AppInfo]>, + config: &Config, + search_tool_enabled: bool, ) -> HashMap { - runtimes - .into_iter() - .map(|runtime| (runtime.tool_name(), runtime.exposure())) + let mut registry = ToolRegistry::default(); + append_mcp_tools( + tools, + connectors, + config, + search_tool_enabled, + &mut registry, + ); + registry + .entries() + .map(|tool| (tool.runtime.tool_name(), tool.exposure)) .collect() } @@ -114,12 +126,12 @@ async fn directly_exposes_effective_tool_sets_when_search_is_unavailable() { let config = test_config().await; let mcp_tools = numbered_mcp_tools(/*count*/ 2); - let runtimes = build_mcp_tool_runtimes( + let runtimes = runtimes_by_name( &mcp_tools, /*connectors*/ None, &config, /*search_tool_enabled*/ false, ); assert_eq!( - runtimes_by_name(runtimes), + runtimes, expected_runtimes(&mcp_tools, ToolExposure::Direct) ); } @@ -188,7 +200,7 @@ async fn excludes_tools_hidden_from_model_exposure() { ]; let connectors = vec![make_connector("calendar", "Calendar")]; - let runtimes = build_mcp_tool_runtimes( + let runtimes = runtimes_by_name( &mcp_tools, Some(connectors.as_slice()), &config, @@ -196,7 +208,7 @@ async fn excludes_tools_hidden_from_model_exposure() { ); assert_eq!( - runtimes_by_name(runtimes), + runtimes, expected_runtimes(&[visible_tool, visible_app_tool], ToolExposure::Direct) ); } @@ -239,7 +251,7 @@ enabled = true let connectors = vec![make_connector("calendar", "Calendar")]; let mcp_tools = [enabled_tool.clone(), disabled_tool]; - let runtimes = build_mcp_tool_runtimes( + let runtimes = runtimes_by_name( &mcp_tools, Some(connectors.as_slice()), &config, @@ -247,7 +259,7 @@ enabled = true ); assert_eq!( - runtimes_by_name(runtimes), + runtimes, expected_runtimes(&[enabled_tool], ToolExposure::Direct) ); } @@ -257,12 +269,12 @@ async fn defers_effective_tool_sets_when_search_is_available() { let config = test_config().await; let mcp_tools = numbered_mcp_tools(/*count*/ 2); - let runtimes = build_mcp_tool_runtimes( + let runtimes = runtimes_by_name( &mcp_tools, /*connectors*/ None, &config, /*search_tool_enabled*/ true, ); assert_eq!( - runtimes_by_name(runtimes), + runtimes, expected_runtimes(&mcp_tools, ToolExposure::Deferred) ); } @@ -290,7 +302,7 @@ async fn defers_apps_and_non_app_mcp_tools() { ]; let connectors = vec![make_connector("calendar", "Calendar")]; - let runtimes = build_mcp_tool_runtimes( + let runtimes = runtimes_by_name( &mcp_tools, Some(connectors.as_slice()), &config, @@ -298,7 +310,7 @@ async fn defers_apps_and_non_app_mcp_tools() { ); assert_eq!( - runtimes_by_name(runtimes), + runtimes, expected_runtimes(&mcp_tools, ToolExposure::Deferred) ); } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 83e037dd4c..22f07c0785 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -16,7 +16,6 @@ use crate::shell::default_user_shell; use crate::shell_snapshot::ShellSnapshot; use crate::test_support::models_manager_with_provider; use crate::tools::format_exec_output_str; -use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolRegistry; use codex_config::ConfigLayerStack; use codex_config::ConfigLayerStackOrdering; @@ -780,32 +779,32 @@ async fn preview_session_start_hooks( ) } -pub(crate) fn tool_runtimes_for_test_step( +pub(crate) fn tool_registry_for_test_step( step_context: &StepContext, -) -> (Vec>, Vec) { - let mut tool_runtimes = crate::tools::spec_plan::build_core_tool_runtimes( +) -> (ToolRegistry, Vec) { + let mut registry = crate::tools::spec_plan::build_core_tool_registry( step_context.turn.as_ref(), &step_context.environments, step_context.mcp.as_ref(), /*tool_suggest_candidates*/ None, /*wait_for_environment_tool_config*/ None, ); - let hosted_specs = crate::tools::spec_plan::append_source_tool_runtimes( + let hosted_specs = crate::tools::spec_plan::append_source_tools( step_context.turn.as_ref(), - &mut tool_runtimes, + &mut registry, Vec::new(), Vec::new(), &step_context.turn.dynamic_tools, ); - (tool_runtimes, hosted_specs) + (registry, hosted_specs) } fn test_tool_runtime(session: Arc, turn_context: Arc) -> ToolCallRuntime { let step_context = StepContext::for_test(Arc::clone(&turn_context)); - let (tool_runtimes, hosted_specs) = tool_runtimes_for_test_step(step_context.as_ref()); - let router = Arc::new(ToolRouter::from_tools( + let (registry, hosted_specs) = tool_registry_for_test_step(step_context.as_ref()); + let router = Arc::new(ToolRouter::from_registry( step_context.turn.as_ref(), - tool_runtimes, + registry, hosted_specs, &Default::default(), )); @@ -10881,10 +10880,10 @@ async fn abort_review_task_emits_exited_then_aborted_and_records_history() { async fn fatal_tool_error_stops_turn_and_reports_error() { let (session, turn_context, _rx) = make_session_and_context_with_rx().await; let step_context = StepContext::for_test(Arc::clone(&turn_context)); - let (tool_runtimes, hosted_specs) = tool_runtimes_for_test_step(step_context.as_ref()); - let router = ToolRouter::from_tools( + let (registry, hosted_specs) = tool_registry_for_test_step(step_context.as_ref()); + let router = ToolRouter::from_registry( step_context.turn.as_ref(), - tool_runtimes, + registry, hosted_specs, &Default::default(), ); diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index be31809f85..f4818d22bd 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/codex-rs/core/src/stream_events_utils_tests.rs @@ -8,7 +8,7 @@ use super::last_assistant_message_from_item; use super::response_item_may_include_external_context; use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; -use crate::session::tests::tool_runtimes_for_test_step; +use crate::session::tests::tool_registry_for_test_step; use crate::tools::ToolRouter; use crate::tools::parallel::ToolCallRuntime; use crate::turn_diff_tracker::TurnDiffTracker; @@ -278,10 +278,10 @@ async fn handle_output_item_done_returns_contributed_last_agent_message() { let session = Arc::new(session); let turn_context = Arc::new(turn_context); let step_context = StepContext::for_test(Arc::clone(&turn_context)); - let (tool_runtimes, hosted_specs) = tool_runtimes_for_test_step(step_context.as_ref()); - let router = Arc::new(ToolRouter::from_tools( + let (registry, hosted_specs) = tool_registry_for_test_step(step_context.as_ref()); + let router = Arc::new(ToolRouter::from_registry( step_context.turn.as_ref(), - tool_runtimes, + registry, hosted_specs, &Default::default(), )); diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index a3871064e5..3990fe049f 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -35,7 +35,6 @@ use codex_protocol::protocol::EventMsg; use codex_rollout::state_db; use codex_shell_command::parse_command::parse_shell_script; use codex_tools::ToolName; -use codex_tools::ToolSearchInfo; use codex_tools::ToolSpec; use futures::future::BoxFuture; use indexmap::IndexMap; @@ -243,94 +242,15 @@ pub(crate) struct PostToolUsePayload { pub(crate) tool_response: Value, } -pub(crate) fn override_tool_exposure( - handler: Arc, - exposure: ToolExposure, -) -> Arc { - if handler.exposure() == exposure { - return handler; - } - - Arc::new(ExposureOverride { handler, exposure }) -} - -struct ExposureOverride { - handler: Arc, - exposure: ToolExposure, -} - -impl ToolExecutor for ExposureOverride { - fn tool_name(&self) -> ToolName { - self.handler.tool_name() - } - - fn spec(&self) -> ToolSpec { - self.handler.spec() - } - - fn exposure(&self) -> ToolExposure { - self.exposure - } - - fn supports_parallel_tool_calls(&self) -> bool { - self.exposure != ToolExposure::Hidden && self.handler.supports_parallel_tool_calls() - } - - fn search_info(&self) -> Option { - self.handler.search_info() - } - - fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { - self.handler.handle(invocation) - } -} - -impl CoreToolRuntime for ExposureOverride { - fn wait_until_ready<'a>(&'a self, session: &'a Arc) -> Option> { - self.handler.wait_until_ready(session) - } - - fn matches_kind(&self, payload: &ToolPayload) -> bool { - self.handler.matches_kind(payload) - } - - fn waits_for_runtime_cancellation(&self) -> bool { - self.handler.waits_for_runtime_cancellation() - } - - fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { - self.handler.pre_tool_use_payload(invocation) - } - - fn post_tool_use_payload( - &self, - invocation: &ToolInvocation, - result: &dyn ToolOutput, - ) -> Option { - self.handler.post_tool_use_payload(invocation, result) - } - - fn with_updated_hook_input( - &self, - invocation: ToolInvocation, - updated_input: Value, - ) -> Result { - self.handler - .with_updated_hook_input(invocation, updated_input) - } - - fn telemetry_tags(&self, invocation: &ToolInvocation) -> ToolTelemetryTags { - self.handler.telemetry_tags(invocation) - } - - fn create_diff_consumer(&self) -> Option> { - self.handler.create_diff_consumer() - } +/// A tool runtime together with its effective exposure for the current step. +pub(crate) struct RegisteredTool { + pub(crate) runtime: Arc, + pub(crate) exposure: ToolExposure, } #[derive(Default)] pub struct ToolRegistry { - tools: IndexMap>, + tools: IndexMap, } impl ToolRegistry { @@ -358,13 +278,22 @@ impl ToolRegistry { where T: CoreToolRuntime + 'static, { - self.register_trusted(override_tool_exposure(Arc::new(handler), exposure)); + self.register_trusted_with_exposure(Arc::new(handler), exposure); } pub(crate) fn register_trusted(&mut self, runtime: Arc) { + let exposure = runtime.exposure(); + self.register_trusted_with_exposure(runtime, exposure); + } + + pub(crate) fn register_trusted_with_exposure( + &mut self, + runtime: Arc, + exposure: ToolExposure, + ) { match self.tools.entry(runtime.tool_name()) { Entry::Vacant(entry) => { - entry.insert(runtime); + entry.insert(RegisteredTool { runtime, exposure }); } Entry::Occupied(entry) => { let tool_name = entry.key(); @@ -380,10 +309,21 @@ impl ToolRegistry { return; } - self.tools.shift_insert(0, tool_name, runtime); + let exposure = runtime.exposure(); + self.tools + .shift_insert(0, tool_name, RegisteredTool { runtime, exposure }); } pub(crate) fn register_external(&mut self, runtime: Arc) -> bool { + let exposure = runtime.exposure(); + self.register_external_with_exposure(runtime, exposure) + } + + pub(crate) fn register_external_with_exposure( + &mut self, + runtime: Arc, + exposure: ToolExposure, + ) -> bool { let tool_name = runtime.tool_name(); if tool_name.namespace.is_none() && tool_name.name == "shell_command" { tracing::warn!(tool_name = %tool_name, "skipping external tool with reserved name"); @@ -392,7 +332,7 @@ impl ToolRegistry { match self.tools.entry(tool_name) { Entry::Vacant(entry) => { - entry.insert(runtime); + entry.insert(RegisteredTool { runtime, exposure }); true } Entry::Occupied(entry) => { @@ -406,21 +346,21 @@ impl ToolRegistry { } pub(crate) fn remove(&mut self, tool_name: &ToolName) -> Option> { - self.tools.shift_remove(tool_name) + self.tools.shift_remove(tool_name).map(|tool| tool.runtime) } - pub(crate) fn runtimes(&self) -> impl Iterator> { + pub(crate) fn entries(&self) -> impl Iterator { self.tools.values() } - pub(crate) fn runtimes_mut(&mut self) -> impl Iterator> { + pub(crate) fn entries_mut(&mut self) -> impl Iterator { self.tools.values_mut() } pub(crate) fn deferred_tool_namespaces(&self) -> BTreeMap { let mut namespaces = BTreeMap::::new(); for (name, tool) in &self.tools { - if tool.exposure() != ToolExposure::Deferred { + if tool.exposure != ToolExposure::Deferred { continue; } let Some(namespace) = &name.namespace else { @@ -430,7 +370,7 @@ impl ToolRegistry { if !existing_description.trim().is_empty() { continue; } - let description = match tool.spec() { + let description = match tool.runtime.spec() { ToolSpec::Namespace(namespace) => namespace.description, ToolSpec::Function(_) | ToolSpec::Freeform(_) @@ -458,7 +398,7 @@ impl ToolRegistry { } pub(crate) fn tool(&self, name: &ToolName) -> Option> { - self.tools.get(name).map(Arc::clone) + self.tools.get(name).map(|tool| Arc::clone(&tool.runtime)) } #[cfg(test)] @@ -470,7 +410,7 @@ impl ToolRegistry { #[cfg(test)] pub(crate) fn tool_exposure(&self, name: &ToolName) -> Option { - self.tools.get(name).map(|tool| tool.exposure()) + self.tools.get(name).map(|tool| tool.exposure) } pub(crate) fn create_diff_consumer( @@ -481,8 +421,8 @@ impl ToolRegistry { } pub(crate) fn supports_parallel_tool_calls(&self, name: &ToolName) -> Option { - let tool = self.tool(name)?; - Some(tool.supports_parallel_tool_calls()) + let tool = self.tools.get(name)?; + Some(tool.exposure != ToolExposure::Hidden && tool.runtime.supports_parallel_tool_calls()) } pub(crate) fn waits_for_runtime_cancellation(&self, name: &ToolName) -> Option { diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index 061c251fa6..4d8f1def8d 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -222,8 +222,8 @@ fn registry_preserves_external_winners_and_trusted_synthetic_order() { assert_eq!( registry - .runtimes() - .map(|runtime| runtime.tool_name()) + .entries() + .map(|tool| tool.runtime.tool_name()) .collect::>(), vec![synthetic_name, first_name.clone(), second_name], ); @@ -243,6 +243,10 @@ fn reserved_shell_command_rejects_external_runtimes_without_a_builtin() { let mut registry = ToolRegistry::default(); assert!(!registry.register_external(handler(shell_command_name.clone()))); + assert!(!registry.register_external_with_exposure( + handler(shell_command_name.clone()), + ToolExposure::Direct, + )); assert!(registry.tool(&shell_command_name).is_none()); let namespaced_handler = handler(namespaced_shell_command_name.clone()); @@ -255,7 +259,7 @@ fn reserved_shell_command_rejects_external_runtimes_without_a_builtin() { } #[tokio::test] -async fn readiness_selects_exact_tool_and_forwards_exposure_overrides() { +async fn readiness_selects_exact_tool_with_registry_owned_exposure() { let (session, _turn) = crate::session::tests::make_session_and_context().await; let session = Arc::new(session); let plain_name = codex_tools::ToolName::plain("echo"); @@ -275,16 +279,14 @@ async fn readiness_selects_exact_tool_and_forwards_exposure_overrides() { }, readiness_waits: Arc::clone(&plain_readiness_waits), }) as Arc; - let namespaced_handler = override_tool_exposure( - Arc::new(ReadinessTestHandler { - handler: TestHandler { - tool_name: namespaced_name.clone(), - }, - readiness_waits: Arc::clone(&namespaced_readiness_waits), - }), - ToolExposure::DirectModelOnly, - ); - let registry = ToolRegistry::from_tools([plain_handler, namespaced_handler]); + let namespaced_handler = Arc::new(ReadinessTestHandler { + handler: TestHandler { + tool_name: namespaced_name.clone(), + }, + readiness_waits: Arc::clone(&namespaced_readiness_waits), + }); + let mut registry = ToolRegistry::from_tools([plain_handler]); + registry.register_trusted_with_exposure(namespaced_handler, ToolExposure::DirectModelOnly); registry .tool(&plain_name) diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 746a298790..3acc12f7ea 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -84,15 +84,15 @@ pub(crate) struct ToolSuggestCandidates { impl ToolRouter { #[cfg(test)] - pub(crate) fn from_tools( + pub(crate) fn from_registry( turn_context: &TurnContext, - tool_runtimes: Vec>, + registry: ToolRegistry, hosted_specs: Vec, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> Self { finalize_tool_router( turn_context, - ToolRegistry::from_tools(tool_runtimes), + registry, hosted_specs, tool_search_handler_cache, ) diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index d40b5de1ff..c9654ec7a1 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -7,10 +7,10 @@ use crate::session::tests::make_session_and_context; use crate::tools::context::ToolPayload; use crate::tools::handlers::McpHandler; use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::RegisteredTool; use crate::tools::registry::ToolExposure; -use crate::tools::registry::override_tool_exposure; -use crate::tools::spec_plan::append_source_tool_runtimes; -use crate::tools::spec_plan::build_core_tool_runtimes; +use crate::tools::spec_plan::append_source_tools; +use crate::tools::spec_plan::build_core_tool_registry; use crate::tools::spec_plan::extension_tool_executors; use crate::turn_diff_tracker::TurnDiffTracker; use codex_extension_api::ExtensionData; @@ -128,27 +128,27 @@ fn extension_tool_test_registry() -> Arc> { fn test_tool_router( step_context: &StepContext, - mcp_tool_runtimes: Vec>, + mcp_tools: Vec, extension_tool_executors: impl IntoIterator>>, dynamic_tools: &[DynamicToolSpec], ) -> ToolRouter { - let mut tool_runtimes = build_core_tool_runtimes( + let mut registry = build_core_tool_registry( step_context.turn.as_ref(), &step_context.environments, step_context.mcp.as_ref(), /*tool_suggest_candidates*/ None, /*wait_for_environment_tool_config*/ None, ); - let hosted_specs = append_source_tool_runtimes( + let hosted_specs = append_source_tools( step_context.turn.as_ref(), - &mut tool_runtimes, - mcp_tool_runtimes, + &mut registry, + mcp_tools, extension_tool_executors, dynamic_tools, ); - ToolRouter::from_tools( + ToolRouter::from_registry( step_context.turn.as_ref(), - tool_runtimes, + registry, hosted_specs, &Default::default(), ) @@ -282,15 +282,24 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { "mcp__echo__", "query_with_delay", )), - override_tool_exposure( - mcp_runtime(mcp_tool_info( + RegisteredTool { + exposure: ToolExposure::DirectModelOnly, + ..mcp_runtime(mcp_tool_info( "hello_echo", /*supports_parallel_tool_calls*/ false, "mcp__hello_echo__", "query_with_delay", - )), - ToolExposure::DirectModelOnly, - ), + )) + }, + RegisteredTool { + exposure: ToolExposure::Hidden, + ..mcp_runtime(mcp_tool_info( + "hidden_echo", + /*supports_parallel_tool_calls*/ true, + "mcp__hidden_echo__", + "query_with_delay", + )) + }, ], Vec::new(), &turn.dynamic_tools, @@ -328,6 +337,17 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { Some(different_server_call.tool_name.clone()) ); + let hidden_call = ToolCall { + tool_name: ToolName::namespaced("mcp__hidden_echo__", "query_with_delay"), + call_id: "call-hidden".to_string(), + payload: ToolPayload::Function { + arguments: "{}".to_string(), + }, + encrypted_function_args: None, + }; + assert!(!router.tool_supports_parallel(&hidden_call)); + assert!(router.tool_runtime(&hidden_call).is_some()); + Ok(()) } @@ -435,8 +455,13 @@ fn mcp_tool_info( } } -fn mcp_runtime(tool_info: codex_mcp::ToolInfo) -> Arc { - Arc::new(McpHandler::new(tool_info).expect("MCP tool spec should build")) +fn mcp_runtime(tool_info: codex_mcp::ToolInfo) -> RegisteredTool { + let runtime = Arc::new(McpHandler::new(tool_info).expect("MCP tool spec should build")) + as Arc; + RegisteredTool { + exposure: runtime.exposure(), + runtime, + } } #[tokio::test] diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index c120bba8bf..157644d37d 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -2,7 +2,7 @@ use crate::agent::exceeds_thread_spawn_depth_limit; use crate::agent::next_thread_spawn_depth; use crate::connectors::AppInfo; use crate::environment_selection::TurnEnvironmentSnapshot; -use crate::mcp_tool_exposure::build_mcp_tool_runtimes; +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; @@ -56,9 +56,10 @@ use crate::tools::handlers::view_image_spec::ViewImageToolOptions; use crate::tools::hosted_spec::WebSearchToolOptions; use crate::tools::hosted_spec::create_web_search_tool; use crate::tools::registry::CoreToolRuntime; +#[cfg(test)] +use crate::tools::registry::RegisteredTool; use crate::tools::registry::ToolExposure; use crate::tools::registry::ToolRegistry; -use crate::tools::registry::override_tool_exposure; use crate::tools::router::ToolRouter; use codex_extension_api::ExtensionData; use codex_features::Feature; @@ -141,14 +142,13 @@ pub(crate) fn build_tool_router( { Vec::new() } else { - for runtime in build_mcp_tool_runtimes( + append_mcp_tools( mcp.tools(), connectors, &turn_context.config, search_tool_enabled(turn_context), - ) { - registry.register_external(runtime); - } + &mut registry, + ); let standalone_web_search_tool = append_extension_tool_executors( turn_context, extension_tool_executors(session, step_store), @@ -168,13 +168,13 @@ pub(crate) fn build_tool_router( #[cfg(test)] #[instrument(level = "trace", skip_all)] -pub(crate) fn build_core_tool_runtimes( +pub(crate) fn build_core_tool_registry( turn_context: &TurnContext, environments: &TurnEnvironmentSnapshot, mcp: &codex_mcp::McpBinding, tool_suggest_candidates: Option<&crate::tools::router::ToolSuggestCandidates>, wait_for_environment_tool_config: Option<&Arc>, -) -> Vec> { +) -> ToolRegistry { let default_agent_type_description = crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new()); let context = CoreToolPlanContext { @@ -188,15 +188,15 @@ pub(crate) fn build_core_tool_runtimes( }; let mut registry = ToolRegistry::default(); add_core_tool_sources(&context, &mut registry); - registry.runtimes().cloned().collect() + registry } #[cfg(test)] #[instrument(level = "trace", skip_all)] -pub(crate) fn append_source_tool_runtimes( +pub(crate) fn append_source_tools( turn_context: &TurnContext, - tool_runtimes: &mut Vec>, - mcp_tool_runtimes: Vec>, + registry: &mut ToolRegistry, + mcp_tools: Vec, extension_tool_executors: impl IntoIterator>>, dynamic_tools: &[DynamicToolSpec], ) -> Vec { @@ -204,14 +204,12 @@ pub(crate) fn append_source_tool_runtimes( return Vec::new(); } - let mut registry = ToolRegistry::from_tools(std::mem::take(tool_runtimes)); - for runtime in mcp_tool_runtimes { - registry.register_external(runtime); + for tool in mcp_tools { + registry.register_external_with_exposure(tool.runtime, tool.exposure); } let standalone_web_search_tool = - append_extension_tool_executors(turn_context, extension_tool_executors, &mut registry); - append_dynamic_tool_runtimes(dynamic_tools, &mut registry); - tool_runtimes.extend(registry.runtimes().cloned()); + append_extension_tool_executors(turn_context, extension_tool_executors, registry); + append_dynamic_tool_runtimes(dynamic_tools, registry); hosted_model_tool_specs(turn_context, standalone_web_search_tool.as_slice()) } @@ -252,10 +250,10 @@ pub(crate) fn finalize_tool_router( } let tool_search_name = ToolName::plain(TOOL_SEARCH_TOOL_NAME); if search_tool_enabled(turn_context) - && registry.runtimes().any(|runtime| { - runtime.tool_name() != tool_search_name - && runtime.exposure() == ToolExposure::Deferred - && runtime.search_info().is_some() + && registry.entries().any(|tool| { + tool.runtime.tool_name() != tool_search_name + && tool.exposure == ToolExposure::Deferred + && tool.runtime.search_info().is_some() }) { registry.remove(&tool_search_name); @@ -273,8 +271,9 @@ fn apply_direct_model_only_namespace_overrides( turn_context: &TurnContext, registry: &mut ToolRegistry, ) { - for runtime in registry.runtimes_mut() { - let configured = runtime + for tool in registry.entries_mut() { + let configured = tool + .runtime .tool_name() .namespace .as_ref() @@ -285,8 +284,8 @@ fn apply_direct_model_only_namespace_overrides( .direct_only_tool_namespaces .contains(namespace) }); - if configured && runtime.exposure().is_available_in_code_mode() { - *runtime = override_tool_exposure(Arc::clone(runtime), ToolExposure::DirectModelOnly); + if configured && tool.exposure.is_available_in_code_mode() { + tool.exposure = ToolExposure::DirectModelOnly; } } } @@ -299,12 +298,12 @@ fn build_model_visible_specs( hosted_specs: Vec, ) -> Vec { let mut specs = Vec::new(); - for runtime in registry.runtimes() { - let tool_name = runtime.tool_name(); - let exposure = runtime.exposure(); + 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 = runtime.spec(); + let spec = tool.runtime.spec(); specs.push(spec_for_model_request( turn_context, exposure, @@ -519,18 +518,18 @@ fn register_code_mode_executors( let mut exec_prompt_tool_specs = Vec::new(); let mut deferred_exec_prompt_tool_specs = Vec::new(); let deferred_tools_guidance_enabled = search_tool_enabled(turn_context); - for executor in registry.runtimes() { - let exposure = executor.exposure(); + for tool in registry.entries() { + let exposure = tool.exposure; if !exposure.is_available_in_code_mode() { continue; } - let tool_name = executor.tool_name(); + let tool_name = tool.runtime.tool_name(); if is_excluded_from_code_mode(turn_context, &tool_name) { continue; } - let spec = executor.spec(); + let spec = tool.runtime.spec(); let Some(code_mode_tool) = codex_tools::tool_spec_to_code_mode_tool_definition(&spec) else { continue; @@ -900,7 +899,7 @@ fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, registry: &mut Too agent_type_description(turn_context, context.default_agent_type_description); let hide_spawn_agent_metadata = turn_context.config.multi_agent_v2.hide_spawn_agent_metadata; - registry.register_trusted(override_tool_exposure( + registry.register_trusted_with_exposure( multi_agent_v2_handler( SpawnAgentHandlerV2::new(SpawnAgentToolOptions { available_models: turn_context.available_models.clone(), @@ -917,32 +916,32 @@ fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, registry: &mut Too tool_namespace, ), exposure, - )); - registry.register_trusted(override_tool_exposure( + ); + registry.register_trusted_with_exposure( multi_agent_v2_handler(SendMessageHandlerV2, tool_namespace), exposure, - )); - registry.register_trusted(override_tool_exposure( + ); + registry.register_trusted_with_exposure( multi_agent_v2_handler(FollowupTaskHandlerV2, tool_namespace), exposure, - )); + ); if turn_context.config.multi_agent_v2.wait_agent_enabled { - registry.register_trusted(override_tool_exposure( + registry.register_trusted_with_exposure( multi_agent_v2_handler( WaitAgentHandlerV2::new(context.wait_agent_timeouts), tool_namespace, ), exposure, - )); + ); } - registry.register_trusted(override_tool_exposure( + registry.register_trusted_with_exposure( multi_agent_v2_handler(InterruptAgentHandler, tool_namespace), exposure, - )); - registry.register_trusted(override_tool_exposure( + ); + registry.register_trusted_with_exposure( multi_agent_v2_handler(ListAgentsHandlerV2, tool_namespace), exposure, - )); + ); } else { let agent_type_description = agent_type_description(turn_context, context.default_agent_type_description); @@ -1012,9 +1011,9 @@ fn append_tool_search_executor( tool_search_handler_cache: &ToolSearchHandlerCache, ) { let search_infos = registry - .runtimes() - .filter(|executor| executor.exposure() == ToolExposure::Deferred) - .filter_map(|executor| executor.search_info()) + .entries() + .filter(|tool| tool.exposure == ToolExposure::Deferred) + .filter_map(|tool| tool.runtime.search_info()) .collect::>(); let source_listing = if turn_context .config diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 19a05b0365..441eea6ed2 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -41,18 +41,18 @@ use crate::tools::handlers::ToolSearchHandlerCache; use crate::tools::handlers::WaitForEnvironmentHandler; use crate::tools::handlers::multi_agents_spec::MULTI_AGENT_V1_NAMESPACE; use crate::tools::registry::CoreToolRuntime; -use crate::tools::registry::override_tool_exposure; +use crate::tools::registry::RegisteredTool; use crate::tools::router::ToolRouter; use crate::tools::router::ToolSuggestCandidates; use crate::tools::router::ToolSuggestPresentation; -use crate::tools::spec_plan::append_source_tool_runtimes; -use crate::tools::spec_plan::build_core_tool_runtimes; +use crate::tools::spec_plan::append_source_tools; +use crate::tools::spec_plan::build_core_tool_registry; const MULTI_AGENT_V2_NAMESPACE: &str = "collaboration"; #[derive(Default)] struct ToolPlanInputs { - tool_runtimes: Vec>, + tool_runtimes: Vec, tool_suggest_candidates: Option, extension_tool_executors: Vec>>, wait_for_environment_tool_config: Option>, @@ -190,23 +190,23 @@ async fn probe_with( configure_turn(&mut turn); let turn = Arc::new(turn); let step_context = StepContext::for_test(Arc::clone(&turn)); - let mut tool_runtimes = build_core_tool_runtimes( + let mut registry = build_core_tool_registry( step_context.turn.as_ref(), &step_context.environments, step_context.mcp.as_ref(), inputs.tool_suggest_candidates.as_ref(), inputs.wait_for_environment_tool_config.as_ref(), ); - let hosted_specs = append_source_tool_runtimes( + let hosted_specs = append_source_tools( step_context.turn.as_ref(), - &mut tool_runtimes, + &mut registry, inputs.tool_runtimes, inputs.extension_tool_executors, &inputs.dynamic_tools, ); - let router = ToolRouter::from_tools( + let router = ToolRouter::from_registry( step_context.turn.as_ref(), - tool_runtimes, + registry, hosted_specs, &Default::default(), ); @@ -396,11 +396,14 @@ fn mcp_runtime( namespace: &str, name: &str, exposure: ToolExposure, -) -> Arc { +) -> RegisteredTool { let handler: Arc = Arc::new( McpHandler::new(mcp_tool(server, namespace, name)).expect("MCP tool spec should build"), ); - override_tool_exposure(handler, exposure) + RegisteredTool { + runtime: handler, + exposure, + } } fn dynamic_tool(namespace: Option<&str>, name: &str, defer_loading: bool) -> DynamicToolSpec { @@ -914,9 +917,9 @@ async fn environment_tools_follow_the_step_context() { &turn.config, ))); - let plan = ToolPlanProbe::from_router(ToolRouter::from_tools( + let plan = ToolPlanProbe::from_router(ToolRouter::from_registry( turn.as_ref(), - build_core_tool_runtimes( + build_core_tool_registry( turn.as_ref(), &environments, mcp.as_ref(), @@ -1023,7 +1026,7 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() { &["lookup".to_string()] ); - let searchable_mcp = ToolPlanInputs { + let searchable_mcp = || ToolPlanInputs { tool_runtimes: vec![mcp_runtime( "searchable", "mcp__searchable", @@ -1037,10 +1040,7 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() { |turn| { turn.model_info.supports_search_tool = false; }, - ToolPlanInputs { - tool_runtimes: searchable_mcp.tool_runtimes.clone(), - ..ToolPlanInputs::default() - }, + searchable_mcp(), ) .await; missing_model_capability.assert_visible_lacks(&["tool_search"]); @@ -1062,10 +1062,7 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() { turn.model_info.supports_search_tool = true; use_bedrock_provider(turn); }, - ToolPlanInputs { - tool_runtimes: searchable_mcp.tool_runtimes.clone(), - ..ToolPlanInputs::default() - }, + searchable_mcp(), ) .await; bedrock_namespace_capability.assert_visible_contains(&["tool_search"]); @@ -1074,7 +1071,7 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() { |turn| { turn.model_info.supports_search_tool = true; }, - searchable_mcp, + searchable_mcp(), ) .await; enabled.assert_visible_contains(&["tool_search"]); @@ -1270,22 +1267,18 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { first_turn.model_info.supports_search_tool = true; let first_turn = Arc::new(first_turn); let first_step_context = StepContext::for_test(Arc::clone(&first_turn)); - let mut first_tool_runtimes = build_core_tool_runtimes( + let mut first_registry = build_core_tool_registry( first_step_context.turn.as_ref(), &first_step_context.environments, first_step_context.mcp.as_ref(), /*tool_suggest_candidates*/ None, /*wait_for_environment_tool_config*/ None, ); - first_tool_runtimes.push(mcp_runtime( - "first", - "mcp__first", - "lookup", - ToolExposure::Deferred, - )); - let first_router = ToolRouter::from_tools( + let first_tool = mcp_runtime("first", "mcp__first", "lookup", ToolExposure::Deferred); + first_registry.register_external_with_exposure(first_tool.runtime, first_tool.exposure); + let first_router = ToolRouter::from_registry( first_step_context.turn.as_ref(), - first_tool_runtimes, + first_registry, super::hosted_model_tool_specs(first_step_context.turn.as_ref(), &[]), &cache, ); @@ -1295,22 +1288,18 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { second_turn.model_info.supports_search_tool = true; let second_turn = Arc::new(second_turn); let second_step_context = StepContext::for_test(Arc::clone(&second_turn)); - let mut second_tool_runtimes = build_core_tool_runtimes( + let mut second_registry = build_core_tool_registry( second_step_context.turn.as_ref(), &second_step_context.environments, second_step_context.mcp.as_ref(), /*tool_suggest_candidates*/ None, /*wait_for_environment_tool_config*/ None, ); - second_tool_runtimes.push(mcp_runtime( - "second", - "mcp__second", - "lookup", - ToolExposure::Deferred, - )); - let second_router = ToolRouter::from_tools( + let second_tool = mcp_runtime("second", "mcp__second", "lookup", ToolExposure::Deferred); + second_registry.register_external_with_exposure(second_tool.runtime, second_tool.exposure); + let second_router = ToolRouter::from_registry( second_step_context.turn.as_ref(), - second_tool_runtimes, + second_registry, super::hosted_model_tool_specs(second_step_context.turn.as_ref(), &[]), &cache, ); @@ -1351,22 +1340,23 @@ async fn tool_search_cache_rebuilds_when_deferred_world_state_changes() { ); let turn = Arc::new(turn); let step_context = StepContext::for_test(Arc::clone(&turn)); - let mut tool_runtimes = build_core_tool_runtimes( + let mut registry = build_core_tool_registry( step_context.turn.as_ref(), &step_context.environments, step_context.mcp.as_ref(), /*tool_suggest_candidates*/ None, /*wait_for_environment_tool_config*/ None, ); - tool_runtimes.push(mcp_runtime( + let tool = mcp_runtime( "calendar", "mcp__calendar", "lookup", ToolExposure::Deferred, - )); - let router = ToolRouter::from_tools( + ); + registry.register_external_with_exposure(tool.runtime, tool.exposure); + let router = ToolRouter::from_registry( step_context.turn.as_ref(), - tool_runtimes, + registry, super::hosted_model_tool_specs(step_context.turn.as_ref(), &[]), &cache, ); diff --git a/codex-rs/tools/src/tool_executor.rs b/codex-rs/tools/src/tool_executor.rs index 508021ff91..c918abd603 100644 --- a/codex-rs/tools/src/tool_executor.rs +++ b/codex-rs/tools/src/tool_executor.rs @@ -60,6 +60,7 @@ pub trait ToolExecutor: Send + Sync { fn spec(&self) -> ToolSpec; + /// The preferred exposure before the host applies step-specific policy. fn exposure(&self) -> ToolExposure { ToolExposure::Direct }