From 066fb39cf1d050ed60eb88265bd9b7c2337a60dd Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 22 May 2026 18:58:38 +0100 Subject: [PATCH] try more things --- codex-rs/core/src/tools/parallel.rs | 3 +- codex-rs/core/src/tools/registry.rs | 169 +++------------- codex-rs/core/src/tools/registry_tests.rs | 6 +- codex-rs/core/src/tools/router.rs | 13 ++ codex-rs/core/src/tools/spec_plan.rs | 36 ++-- .../core/src/tools/spec_plan/planned_tools.rs | 189 ++++++++++++++++++ codex-rs/core/src/tools/spec_plan_tests.rs | 111 +++++++++- .../src/tools/tool_dispatch_trace_tests.rs | 13 +- 8 files changed, 360 insertions(+), 180 deletions(-) create mode 100644 codex-rs/core/src/tools/spec_plan/planned_tools.rs diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index b6ef07fff1..92711f63f7 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -328,8 +328,7 @@ mod tests { let handler = Arc::new(ImmediateHandler { tool_name: tool_name.clone(), }) as Arc; - let mut registry = ToolRegistry::default(); - registry.prepend([handler]); + let registry = ToolRegistry::from_tools([handler]); let router = Arc::new(ToolRouter::from_parts(registry, Vec::new())); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let runtime = ToolCallRuntime::new(router, session, turn_context, tracker); diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 319f8566c3..8f57671171 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -28,10 +28,7 @@ use crate::util::error_or_panic; use codex_extension_api::ToolCallOutcome; use codex_protocol::models::ResponseInputItem; use codex_protocol::protocol::EventMsg; -use codex_tools::ResponsesApiNamespace; -use codex_tools::ResponsesApiNamespaceTool; use codex_tools::ToolName; -use codex_tools::ToolSpec; use futures::future::BoxFuture; use serde_json::Value; use tracing::warn; @@ -166,160 +163,39 @@ pub(crate) struct PostToolUsePayload { #[derive(Default)] pub struct ToolRegistry { - ordered_tools: Vec>, - tools_by_name: HashMap>, -} - -pub(crate) struct RegisteredTool { - runtime: Arc, - tool_name: ToolName, - spec: ToolSpec, - exposure: ToolExposure, -} - -impl RegisteredTool { - fn new(runtime: T) -> Self - where - T: CoreToolRuntime + 'static, - { - Self::from_runtime(Arc::new(runtime)) - } - - fn from_runtime(runtime: Arc) -> Self { - Self { - tool_name: runtime.tool_name(), - spec: runtime.spec(), - exposure: runtime.exposure(), - runtime, - } - } - - fn with_exposure(mut self, exposure: ToolExposure) -> Self { - self.exposure = exposure; - self - } - - fn with_namespace(self, namespace: &str, description: &str) -> Self { - let Self { - runtime, - tool_name, - spec, - exposure, - } = self; - let spec = match spec { - ToolSpec::Function(tool) => ToolSpec::Namespace(ResponsesApiNamespace { - name: namespace.to_string(), - description: description.to_string(), - tools: vec![ResponsesApiNamespaceTool::Function(tool)], - }), - spec => spec, - }; - - Self { - runtime, - tool_name: ToolName::namespaced(namespace, tool_name.name), - spec, - exposure, - } - } - - pub(crate) fn tool_name(&self) -> ToolName { - self.tool_name.clone() - } - - pub(crate) fn spec(&self) -> ToolSpec { - self.spec.clone() - } - - pub(crate) fn exposure(&self) -> ToolExposure { - self.exposure - } - - pub(crate) fn search_info(&self) -> Option { - self.runtime.search_info() - } - - fn runtime(&self) -> Arc { - Arc::clone(&self.runtime) - } - - fn supports_parallel_tool_calls(&self) -> bool { - self.exposure != ToolExposure::Hidden && self.runtime.supports_parallel_tool_calls() - } + tools: HashMap>, } impl ToolRegistry { - pub(crate) fn add(&mut self, tool: T) - where - T: CoreToolRuntime + 'static, - { - self.add_registered(RegisteredTool::new(tool)); - } - - pub(crate) fn add_with_exposure(&mut self, tool: T, exposure: ToolExposure) - where - T: CoreToolRuntime + 'static, - { - self.add_registered(RegisteredTool::new(tool).with_exposure(exposure)); - } - - pub(crate) fn add_namespaced_with_exposure( - &mut self, - tool: T, - namespace: &str, - description: &str, - exposure: ToolExposure, - ) where - T: CoreToolRuntime + 'static, - { - self.add_registered( - RegisteredTool::new(tool) - .with_namespace(namespace, description) - .with_exposure(exposure), - ); - } - - fn add_registered(&mut self, tool: RegisteredTool) { - self.add_arc(Arc::new(tool)); - } - - pub(crate) fn prepend(&mut self, tools: impl IntoIterator>) { - let mut new_tools = Vec::new(); + pub(crate) fn from_tools(tools: impl IntoIterator>) -> Self { + let mut tools_by_name = HashMap::new(); for tool in tools { - let tool = Arc::new(RegisteredTool::from_runtime(tool)); - if self.insert(Arc::clone(&tool)) { - new_tools.push(tool); + let name = tool.tool_name(); + if tools_by_name.contains_key(&name) { + error_or_panic(format!("tool {name} already registered")); + continue; } + tools_by_name.insert(name, tool); } - self.ordered_tools.splice(0..0, new_tools); - } - - pub(crate) fn iter(&self) -> impl Iterator> { - self.ordered_tools.iter() - } - - fn add_arc(&mut self, tool: Arc) { - if self.insert(Arc::clone(&tool)) { - self.ordered_tools.push(tool); + Self { + tools: tools_by_name, } } - fn insert(&mut self, tool: Arc) -> bool { - let name = tool.tool_name(); - if self.tools_by_name.contains_key(&name) { - error_or_panic(format!("tool {name} already registered")); - return false; - } - self.tools_by_name.insert(name, tool); - true - } - fn tool(&self, name: &ToolName) -> Option> { - self.tools_by_name.get(name).map(|tool| tool.runtime()) + self.tools.get(name).map(Arc::clone) } - fn registered_tool(&self, name: &ToolName) -> Option> { - self.tools_by_name.get(name).map(Arc::clone) + #[cfg(test)] + pub(crate) fn tool_names_for_test(&self) -> Vec { + let mut names = self.tools.keys().cloned().collect::>(); + names.sort(); + names + } + + #[cfg(test)] + pub(crate) fn tool_exposure(&self, name: &ToolName) -> Option { + self.tools.get(name).map(|tool| tool.exposure()) } pub(crate) fn create_diff_consumer( @@ -330,10 +206,11 @@ impl ToolRegistry { } pub(crate) fn supports_parallel_tool_calls(&self, name: &ToolName) -> Option { - let tool = self.registered_tool(name)?; + let tool = self.tool(name)?; Some(tool.supports_parallel_tool_calls()) } + #[allow(dead_code)] pub(crate) async fn dispatch_any( &self, invocation: ToolInvocation, diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index 2e7076f82d..5f7e08e21b 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -145,8 +145,7 @@ fn handler_looks_up_namespaced_aliases_explicitly() { let namespaced_handler = Arc::new(TestHandler { tool_name: namespaced_name.clone(), }) as Arc; - let mut registry = ToolRegistry::default(); - registry.prepend([Arc::clone(&plain_handler), Arc::clone(&namespaced_handler)]); + let registry = ToolRegistry::from_tools([Arc::clone(&plain_handler), Arc::clone(&namespaced_handler)]); let plain = registry.tool(&plain_name); let namespaced = registry.tool(&namespaced_name); @@ -190,8 +189,7 @@ async fn dispatch_notifies_tool_lifecycle_contributors() -> anyhow::Result<()> { tool_name: failing_tool.clone(), result: LifecycleTestResult::Err, }) as Arc; - let mut registry = ToolRegistry::default(); - registry.prepend([ok_handler, failing_handler]); + let registry = ToolRegistry::from_tools([ok_handler, failing_handler]); let session = Arc::new(session); let turn = Arc::new(turn); diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index e01b9d13d0..cfd4776890 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -60,6 +60,19 @@ impl ToolRouter { self.model_visible_specs.clone() } + #[cfg(test)] + pub(crate) fn registered_tool_names_for_test(&self) -> Vec { + self.registry.tool_names_for_test() + } + + #[cfg(test)] + pub(crate) fn tool_exposure_for_test( + &self, + name: &ToolName, + ) -> Option { + self.registry.tool_exposure(name) + } + pub(crate) fn create_diff_consumer( &self, tool_name: &ToolName, diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 6de91ba03d..8ed4303a92 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -49,7 +49,6 @@ use crate::tools::hosted_spec::create_image_generation_tool; use crate::tools::hosted_spec::create_web_search_tool; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExposure; -use crate::tools::registry::ToolRegistry; use crate::tools::router::ToolRouter; use crate::tools::router::ToolRouterParams; use codex_features::Feature; @@ -80,6 +79,8 @@ use std::collections::HashSet; use std::sync::Arc; use tracing::warn; +use self::planned_tools::PlannedTools; + const MULTI_AGENT_V2_NAMESPACE_DESCRIPTION: &str = "Tools for spawning and managing sub-agents."; pub(crate) fn build_tool_router( @@ -93,7 +94,7 @@ pub(crate) fn build_tool_router( extension_tool_executors, dynamic_tools, } = params; - let mut tools = ToolRegistry::default(); + let mut tools = PlannedTools::default(); add_shell_tools(turn_context, &mut tools); add_core_utility_tools(turn_context, discoverable_tools.as_deref(), &mut tools); @@ -110,10 +111,10 @@ pub(crate) fn build_tool_router( prepend_code_mode_executors(turn_context, &mut tools); let model_visible_specs = model_visible_specs(turn_context, &tools); - ToolRouter::from_parts(tools, model_visible_specs) + ToolRouter::from_parts(tools.into_registry(), model_visible_specs) } -fn model_visible_specs(turn_context: &TurnContext, tools: &ToolRegistry) -> Vec { +fn model_visible_specs(turn_context: &TurnContext, tools: &PlannedTools) -> Vec { let mut specs = Vec::new(); for tool in tools.iter() { let tool_name = tool.tool_name(); @@ -295,7 +296,7 @@ fn is_hidden_by_code_mode_only( )) } -fn prepend_code_mode_executors(turn_context: &TurnContext, tools: &mut ToolRegistry) { +fn prepend_code_mode_executors(turn_context: &TurnContext, tools: &mut PlannedTools) { let deferred_tools_available = search_tool_enabled(turn_context) && tools .iter() @@ -307,7 +308,7 @@ fn prepend_code_mode_executors(turn_context: &TurnContext, tools: &mut ToolRegis fn build_code_mode_executors( turn_context: &TurnContext, - tools: &ToolRegistry, + tools: &PlannedTools, deferred_tools_available: bool, ) -> Vec> { if !code_mode_enabled(turn_context) { @@ -421,7 +422,7 @@ fn code_mode_namespace_descriptions( namespace_descriptions } -fn add_shell_tools(turn_context: &TurnContext, tools: &mut ToolRegistry) { +fn add_shell_tools(turn_context: &TurnContext, tools: &mut PlannedTools) { let features = turn_context.features.get(); let environment_mode = turn_context.tool_environment_mode(); if !environment_mode.has_environment() { @@ -448,10 +449,7 @@ fn add_shell_tools(turn_context: &TurnContext, tools: &mut ToolRegistry) { // Keep the legacy shell tool registered while unified exec is // model-visible. - tools.add_with_exposure( - ShellCommandHandler::new(shell_command_options), - ToolExposure::Hidden, - ); + tools.add_hidden(ShellCommandHandler::new(shell_command_options)); } ConfigShellToolType::Disabled => {} ConfigShellToolType::Default @@ -465,7 +463,7 @@ fn add_shell_tools(turn_context: &TurnContext, tools: &mut ToolRegistry) { fn add_core_utility_tools( turn_context: &TurnContext, discoverable_tools: Option<&[DiscoverableTool]>, - tools: &mut ToolRegistry, + tools: &mut PlannedTools, ) { let features = turn_context.features.get(); let environment_mode = turn_context.tool_environment_mode(); @@ -520,7 +518,7 @@ fn add_core_utility_tools( } } -fn add_collaboration_tools(turn_context: &TurnContext, tools: &mut ToolRegistry) { +fn add_collaboration_tools(turn_context: &TurnContext, tools: &mut PlannedTools) { if collab_tools_enabled(turn_context) { if multi_agent_v2_enabled(turn_context) { let exposure = if turn_context.config.multi_agent_v2.non_code_mode_only { @@ -603,7 +601,7 @@ fn add_collaboration_tools(turn_context: &TurnContext, tools: &mut ToolRegistry) } fn add_multi_agent_v2_tool( - tools: &mut ToolRegistry, + tools: &mut PlannedTools, tool: T, namespace: Option<&str>, exposure: ToolExposure, @@ -624,7 +622,7 @@ fn add_multi_agent_v2_tool( fn add_mcp_tools( mcp_tools: Option<&[ToolInfo]>, deferred_mcp_tools: Option<&[ToolInfo]>, - tools: &mut ToolRegistry, + tools: &mut PlannedTools, ) { if let Some(mcp_tools) = mcp_tools { tools.add(ListMcpResourcesHandler); @@ -655,7 +653,7 @@ fn add_mcp_tools( } } -fn add_dynamic_tools(dynamic_tools: &[DynamicToolSpec], tools: &mut ToolRegistry) { +fn add_dynamic_tools(dynamic_tools: &[DynamicToolSpec], tools: &mut PlannedTools) { for tool in dynamic_tools { let Some(handler) = DynamicToolHandler::new(tool) else { tracing::error!( @@ -669,7 +667,7 @@ fn add_dynamic_tools(dynamic_tools: &[DynamicToolSpec], tools: &mut ToolRegistry } } -fn append_tool_search_executor(turn_context: &TurnContext, tools: &mut ToolRegistry) { +fn append_tool_search_executor(turn_context: &TurnContext, tools: &mut PlannedTools) { if !(search_tool_enabled(turn_context) && namespace_tools_enabled(turn_context)) { return; } @@ -689,7 +687,7 @@ fn append_tool_search_executor(turn_context: &TurnContext, tools: &mut ToolRegis fn append_extension_tool_executors( turn_context: &TurnContext, executors: &[Arc>], - tools: &mut ToolRegistry, + tools: &mut PlannedTools, ) { // Extension ToolContributor implementations are resolved into executors // before planning. Core only adapts those executors into its runtime set. @@ -749,6 +747,8 @@ fn code_mode_namespace_name<'a>( .map(|namespace_description| namespace_description.name.as_str()) } +mod planned_tools; + #[cfg(test)] #[path = "spec_plan_tests.rs"] mod tests; diff --git a/codex-rs/core/src/tools/spec_plan/planned_tools.rs b/codex-rs/core/src/tools/spec_plan/planned_tools.rs new file mode 100644 index 0000000000..87cbc57049 --- /dev/null +++ b/codex-rs/core/src/tools/spec_plan/planned_tools.rs @@ -0,0 +1,189 @@ +use std::sync::Arc; + +use codex_tools::ResponsesApiNamespace; +use codex_tools::ResponsesApiNamespaceTool; +use codex_tools::ToolExecutor; +use codex_tools::ToolName; +use codex_tools::ToolSpec; +use futures::future::BoxFuture; +use serde_json::Value; + +use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; +use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::PostToolUsePayload; +use crate::tools::registry::PreToolUsePayload; +use crate::tools::registry::ToolArgumentDiffConsumer; +use crate::tools::registry::ToolExposure; +use crate::tools::registry::ToolRegistry; +use crate::tools::registry::ToolTelemetryTags; +use crate::tools::tool_search_entry::ToolSearchInfo; + +type PlannedRuntime = Arc; + +#[derive(Default)] +pub(super) struct PlannedTools { + runtimes: Vec, +} + +impl PlannedTools { + pub(super) fn add(&mut self, runtime: T) + where + T: CoreToolRuntime + 'static, + { + self.runtimes.push(Arc::new(runtime)); + } + + pub(super) fn add_with_exposure(&mut self, runtime: T, exposure: ToolExposure) + where + T: CoreToolRuntime + 'static, + { + let runtime: PlannedRuntime = Arc::new(runtime); + if runtime.exposure() == exposure { + self.runtimes.push(runtime); + } else { + self.runtimes.push(Arc::new(OverriddenTool { + runtime, + exposure, + namespace: None, + })); + } + } + + pub(super) fn add_hidden(&mut self, runtime: T) + where + T: CoreToolRuntime + 'static, + { + self.add_with_exposure(runtime, ToolExposure::Hidden); + } + + pub(super) fn add_namespaced_with_exposure( + &mut self, + runtime: T, + namespace: &str, + description: &str, + exposure: ToolExposure, + ) where + T: CoreToolRuntime + 'static, + { + self.runtimes.push(Arc::new(OverriddenTool { + runtime: Arc::new(runtime), + exposure, + namespace: Some(NamespaceOverride { + name: namespace.to_string(), + description: description.to_string(), + }), + })); + } + + pub(super) fn prepend(&mut self, runtimes: impl IntoIterator) { + self.runtimes.splice(0..0, runtimes); + } + + pub(super) fn iter(&self) -> impl Iterator { + self.runtimes.iter() + } + + pub(super) fn into_registry(self) -> ToolRegistry { + ToolRegistry::from_tools(self.runtimes) + } +} + +struct NamespaceOverride { + name: String, + description: String, +} + +/// A runtime whose name, spec, and exposure are projected for one turn. +/// +/// The wrapper remains a runtime so the registry never stores identity or +/// presentation data separately from the executable tool. +struct OverriddenTool { + runtime: PlannedRuntime, + exposure: ToolExposure, + namespace: Option, +} + +#[async_trait::async_trait] +impl ToolExecutor for OverriddenTool { + fn tool_name(&self) -> ToolName { + match &self.namespace { + Some(namespace) => { + ToolName::namespaced(namespace.name.clone(), self.runtime.tool_name().name) + } + None => self.runtime.tool_name(), + } + } + + fn spec(&self) -> ToolSpec { + match (&self.namespace, self.runtime.spec()) { + (Some(namespace), ToolSpec::Function(tool)) => { + ToolSpec::Namespace(ResponsesApiNamespace { + name: namespace.name.clone(), + description: namespace.description.clone(), + tools: vec![ResponsesApiNamespaceTool::Function(tool)], + }) + } + (_, spec) => spec, + } + } + + fn exposure(&self) -> ToolExposure { + self.exposure + } + + fn supports_parallel_tool_calls(&self) -> bool { + self.exposure != ToolExposure::Hidden && self.runtime.supports_parallel_tool_calls() + } + + async fn handle( + &self, + invocation: ToolInvocation, + ) -> Result, FunctionCallError> { + self.runtime.handle(invocation).await + } +} + +impl CoreToolRuntime for OverriddenTool { + fn search_info(&self) -> Option { + self.runtime.search_info() + } + + fn matches_kind(&self, payload: &ToolPayload) -> bool { + self.runtime.matches_kind(payload) + } + + fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { + self.runtime.pre_tool_use_payload(invocation) + } + + fn post_tool_use_payload( + &self, + invocation: &ToolInvocation, + result: &dyn ToolOutput, + ) -> Option { + self.runtime.post_tool_use_payload(invocation, result) + } + + fn with_updated_hook_input( + &self, + invocation: ToolInvocation, + updated_input: Value, + ) -> Result { + self.runtime + .with_updated_hook_input(invocation, updated_input) + } + + fn telemetry_tags<'a>( + &'a self, + invocation: &'a ToolInvocation, + ) -> BoxFuture<'a, ToolTelemetryTags> { + self.runtime.telemetry_tags(invocation) + } + + fn create_diff_consumer(&self) -> Option> { + self.runtime.create_diff_consumer() + } +} diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 24f1d6da8a..5186e60b87 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -20,6 +20,8 @@ use codex_tools::DiscoverablePluginInfo; use codex_tools::DiscoverableTool; use codex_tools::ResponsesApiNamespaceTool; use codex_tools::ResponsesApiTool; +use codex_tools::ToolExposure; +use codex_tools::ToolName; use codex_tools::ToolSpec; use pretty_assertions::assert_eq; use serde_json::json; @@ -42,6 +44,8 @@ struct ToolPlanProbe { visible_specs: Vec, visible_names: Vec, namespace_functions: BTreeMap>, + registered_names: Vec, + exposures: BTreeMap, } impl ToolPlanProbe { @@ -71,10 +75,25 @@ impl ToolPlanProbe { | ToolSpec::Freeform(_) => None, }) .collect::>(); + let registered_tool_names = router.registered_tool_names_for_test(); + let registered_names = registered_tool_names + .iter() + .map(ToString::to_string) + .collect::>(); + let exposures = registered_tool_names + .iter() + .filter_map(|name| { + router + .tool_exposure_for_test(name) + .map(|exposure| (name.to_string(), exposure)) + }) + .collect::>(); Self { visible_specs, visible_names, namespace_functions, + registered_names, + exposures, } } @@ -98,6 +117,31 @@ impl ToolPlanProbe { } } + fn assert_registered_contains(&self, expected: &[&str]) { + for name in expected { + assert!( + self.registered_names + .iter() + .any(|registered| registered == name), + "expected registered tool `{name}` in {:?}", + self.registered_names + ); + } + } + + fn assert_registered_lacks(&self, expected_absent: &[&str]) { + for name in expected_absent { + assert!( + !self + .registered_names + .iter() + .any(|registered| registered == name), + "expected registered tool `{name}` to be absent from {:?}", + self.registered_names + ); + } + } + fn namespace_function_names(&self, namespace: &str) -> &[String] { self.namespace_functions .get(namespace) @@ -110,6 +154,13 @@ impl ToolPlanProbe { .find(|spec| spec.name() == name) .unwrap_or_else(|| panic!("expected visible spec `{name}` in {:?}", self.visible_names)) } + + fn exposure(&self, name: &str) -> ToolExposure { + *self + .exposures + .get(name) + .unwrap_or_else(|| panic!("expected registered tool `{name}`")) + } } async fn probe_with( @@ -287,7 +338,7 @@ fn apply_patch_accepts_environment_id(spec: &ToolSpec) -> bool { } #[tokio::test] -async fn shell_family_exposes_unified_exec_instead_of_legacy_shell() { +async fn shell_family_registers_visible_unified_exec_and_hidden_legacy_shell() { let plan = probe(|turn| { set_features(turn, &[Feature::ShellTool, Feature::UnifiedExec]); set_feature(turn, Feature::ShellZshFork, /*enabled*/ false); @@ -297,6 +348,8 @@ async fn shell_family_exposes_unified_exec_instead_of_legacy_shell() { plan.assert_visible_contains(&["exec_command", "write_stdin"]); plan.assert_visible_lacks(&["shell_command"]); + plan.assert_registered_contains(&["exec_command", "write_stdin", "shell_command"]); + assert_eq!(plan.exposure("shell_command"), ToolExposure::Hidden); } #[tokio::test] @@ -313,6 +366,12 @@ async fn environment_count_controls_environment_backed_tools() { "apply_patch", "view_image", ]); + no_environment.assert_registered_lacks(&[ + "shell_command", + "exec_command", + "apply_patch", + "view_image", + ]); let multiple_environments = probe(|turn| { duplicate_primary_environment(turn); set_feature(turn, Feature::ShellTool, /*enabled*/ true); @@ -451,10 +510,11 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() { ) .await; enabled.assert_visible_contains(&["tool_search"]); + enabled.assert_registered_contains(&["tool_search", "mcp__searchable__lookup"]); } #[tokio::test] -async fn invalid_mcp_tools_are_not_exposed() { +async fn invalid_mcp_tools_are_not_registered() { let plan = probe_with( |_| {}, ToolPlanInputs { @@ -469,6 +529,7 @@ async fn invalid_mcp_tools_are_not_exposed() { .await; plan.assert_visible_lacks(&["mcp__invalid__"]); + plan.assert_registered_lacks(&["mcp__invalid__lookup"]); } #[tokio::test] @@ -700,6 +761,10 @@ async fn multi_agent_feature_selects_one_agent_tool_family() { }) .await; direct_model_only.assert_visible_contains(&["spawn_agent", "send_message", "wait_agent"]); + assert_eq!( + direct_model_only.exposure("spawn_agent"), + ToolExposure::DirectModelOnly + ); } #[tokio::test] @@ -719,6 +784,27 @@ async fn v1_multi_agent_tools_defer_when_tool_search_available() { "wait_agent", "close_agent", ]); + for tool_name in [ + "spawn_agent", + "send_input", + "resume_agent", + "wait_agent", + "close_agent", + ] { + let namespaced_tool_name = ToolName::namespaced(MULTI_AGENT_V1_NAMESPACE, tool_name); + let namespaced_tool_name = namespaced_tool_name.to_string(); + assert!( + plan.registered_names.contains(&namespaced_tool_name), + "expected namespaced runtime for {tool_name}" + ); + assert!( + !plan + .registered_names + .contains(&ToolName::plain(tool_name).to_string()), + "expected no plain runtime for deferred {tool_name}" + ); + assert_eq!(plan.exposure(&namespaced_tool_name), ToolExposure::Deferred); + } let ToolSpec::ToolSearch { description, .. } = plan.visible_spec("tool_search") else { panic!("expected visible tool_search spec"); }; @@ -745,6 +831,18 @@ async fn multi_agent_v2_can_use_configured_tool_namespace() { "list_agents", ] { namespaced.assert_visible_lacks(&[tool_name]); + assert!( + namespaced + .registered_names + .contains(&ToolName::namespaced("agents", tool_name).to_string()), + "expected namespaced runtime for {tool_name}" + ); + assert!( + !namespaced + .registered_names + .contains(&ToolName::plain(tool_name).to_string()), + "expected no plain runtime for {tool_name}" + ); assert!( namespaced .namespace_function_names("agents") @@ -768,6 +866,15 @@ async fn multi_agent_v2_namespace_is_ignored_without_provider_namespace_support( plan.assert_visible_contains(&["spawn_agent", "send_message", "list_agents"]); plan.assert_visible_lacks(&["agents"]); + assert!( + plan.registered_names + .contains(&ToolName::plain("spawn_agent").to_string()) + ); + assert!( + !plan + .registered_names + .contains(&ToolName::namespaced("agents", "spawn_agent").to_string()) + ); } #[tokio::test] diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index c918de5df3..163a2b3730 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -72,10 +72,9 @@ async fn dispatch_lifecycle_trace_records_direct_and_code_mode_requesters() -> a "await tools.test_tool({})", ); - let mut registry = ToolRegistry::default(); - registry.add(TestHandler { + let registry = ToolRegistry::from_tools([Arc::new(TestHandler { tool_name: codex_tools::ToolName::plain("test_tool"), - }); + }) as Arc]); let session = Arc::new(session); let turn = Arc::new(turn); @@ -184,10 +183,9 @@ async fn dispatch_lifecycle_trace_records_incompatible_payload_failures() -> any let (mut session, turn) = make_session_and_context().await; attach_test_trace(&mut session, &turn, temp.path())?; - let mut registry = ToolRegistry::default(); - registry.add(TestHandler { + let registry = ToolRegistry::from_tools([Arc::new(TestHandler { tool_name: codex_tools::ToolName::plain("test_tool"), - }); + }) as Arc]); let session = Arc::new(session); let turn = Arc::new(turn); @@ -219,8 +217,7 @@ async fn missing_code_mode_wait_traces_only_the_wait_tool_call() -> anyhow::Resu let (mut session, turn) = make_session_and_context().await; attach_test_trace(&mut session, &turn, temp.path())?; - let mut registry = ToolRegistry::default(); - registry.add(CodeModeWaitHandler); + let registry = ToolRegistry::from_tools([Arc::new(CodeModeWaitHandler) as Arc]); let session = Arc::new(session); let turn = Arc::new(turn);