From fdda59c00b95a7656d29f1fcfdee88e7a7884f66 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 13 May 2026 18:16:51 +0200 Subject: [PATCH] Introduce tool exposure for deferred registration (#22489) ## Why Deferred tools were tracked with separate side-channel filtering after tool specs had already been assembled. That made the registry responsible for executing tools while the router/spec planner separately decided whether those same tools should be exposed to the model up front. This PR makes exposure part of the tool handler contract so direct versus deferred availability travels with the executable tool registration. Next step will be to simplify registration ## What Changed - Adds `ToolExposure` to `codex-tools` and exposes it through `ToolExecutor`, defaulting tools to `Direct`. - Teaches dynamic tools and MCP handlers to mark deferred tools as `Deferred` at construction time. - Renames the registry object-safe wrapper from `AnyToolHandler` to `RegisteredTool` and uses `ToolExposure` when deciding whether to include a handler's spec in the initial model-visible tool list. - Refactors tool spec planning to derive direct specs and deferred search entries from registered handlers, removing the router's special-case deferred dynamic tool filtering. ## Verification - Not run. --- codex-rs/core/src/tools/handlers/dynamic.rs | 11 ++++ codex-rs/core/src/tools/handlers/mcp.rs | 15 +++++- codex-rs/core/src/tools/registry.rs | 54 +++++++++++--------- codex-rs/core/src/tools/registry_tests.rs | 4 +- codex-rs/core/src/tools/router.rs | 52 ++----------------- codex-rs/core/src/tools/spec_plan.rs | 56 ++++++++++----------- codex-rs/tools/src/lib.rs | 1 + codex-rs/tools/src/tool_executor.rs | 12 +++++ 8 files changed, 100 insertions(+), 105 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/dynamic.rs b/codex-rs/core/src/tools/handlers/dynamic.rs index 3f3e72e2f3..8000b2097c 100644 --- a/codex-rs/core/src/tools/handlers/dynamic.rs +++ b/codex-rs/core/src/tools/handlers/dynamic.rs @@ -6,6 +6,7 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolExecutor; +use crate::tools::registry::ToolExposure; use crate::tools::registry::ToolHandler; use crate::tools::tool_search_entry::ToolSearchInfo; use crate::turn_timing::now_unix_timestamp_ms; @@ -30,6 +31,7 @@ use tracing::warn; pub struct DynamicToolHandler { tool_name: ToolName, spec: Option, + exposure: ToolExposure, search_text: String, } @@ -48,6 +50,11 @@ impl DynamicToolHandler { Some(Self { tool_name, spec: Some(spec), + exposure: if tool.defer_loading { + ToolExposure::Deferred + } else { + ToolExposure::Direct + }, search_text: build_dynamic_search_text(tool), }) } @@ -64,6 +71,10 @@ impl ToolExecutor for DynamicToolHandler { self.spec.clone() } + fn exposure(&self) -> ToolExposure { + self.exposure + } + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 2cd02e8a0a..5f700e3a98 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -13,6 +13,7 @@ use crate::tools::hook_names::HookToolName; use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolExecutor; +use crate::tools::registry::ToolExposure; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolTelemetryTags; use crate::tools::tool_search_entry::ToolSearchInfo; @@ -28,11 +29,19 @@ use serde_json::Value; pub struct McpHandler { tool_info: ToolInfo, + exposure: ToolExposure, } impl McpHandler { pub fn new(tool_info: ToolInfo) -> Self { - Self { tool_info } + Self::with_exposure(tool_info, ToolExposure::Direct) + } + + pub fn with_exposure(tool_info: ToolInfo, exposure: ToolExposure) -> Self { + Self { + tool_info, + exposure, + } } } @@ -71,6 +80,10 @@ impl ToolExecutor for McpHandler { })) } + fn exposure(&self) -> ToolExposure { + self.exposure + } + fn supports_parallel_tool_calls(&self) -> bool { self.tool_info.supports_parallel_tool_calls } diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index b8480600f7..ad54b5fb77 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -35,6 +35,7 @@ use tracing::warn; pub(crate) type ToolTelemetryTags = Vec<(&'static str, String)>; pub use codex_tools::ToolExecutor; +pub use codex_tools::ToolExposure; pub trait ToolHandler: ToolExecutor { fn search_info(&self) -> Option { @@ -155,11 +156,17 @@ pub(crate) struct PostToolUsePayload { pub(crate) tool_response: Value, } -pub(crate) trait AnyToolHandler: Send + Sync { +/// Object-safe registry entry for heterogeneous tool handlers. +/// +/// Concrete handlers keep their typed `ToolExecutor::Output`; the registry +/// boxes that output only after typed hooks have run. +pub(crate) trait RegisteredTool: Send + Sync { fn tool_name(&self) -> ToolName; fn spec(&self) -> Option; + fn exposure(&self) -> ToolExposure; + fn search_info(&self) -> Option; fn supports_parallel_tool_calls(&self) -> bool; @@ -186,7 +193,7 @@ pub(crate) trait AnyToolHandler: Send + Sync { ) -> BoxFuture<'a, Result>; } -impl AnyToolHandler for T +impl RegisteredTool for T where T: ToolHandler, { @@ -198,6 +205,10 @@ where ToolExecutor::spec(self) } + fn exposure(&self) -> ToolExposure { + ToolExecutor::exposure(self) + } + fn search_info(&self) -> Option { ToolHandler::search_info(self) } @@ -253,11 +264,11 @@ where } pub struct ToolRegistry { - handlers: HashMap>, + handlers: HashMap>, } impl ToolRegistry { - fn new(handlers: HashMap>) -> Self { + fn new(handlers: HashMap>) -> Self { Self { handlers } } @@ -272,10 +283,10 @@ impl ToolRegistry { T: ToolHandler + 'static, { let name = handler.tool_name(); - Self::new(HashMap::from([(name, handler as Arc)])) + Self::new(HashMap::from([(name, handler as Arc)])) } - fn handler(&self, name: &ToolName) -> Option> { + fn handler(&self, name: &ToolName) -> Option> { self.handlers.get(name).map(Arc::clone) } @@ -534,7 +545,7 @@ impl ToolRegistry { } pub struct ToolRegistryBuilder { - handlers: HashMap>, + handlers: HashMap>, specs: Vec, } @@ -554,29 +565,28 @@ impl ToolRegistryBuilder { where H: ToolHandler + 'static, { - self.register_any_handler(handler); + self.register_tool(handler); } - pub(crate) fn register_any_handler(&mut self, handler: Arc) { - self.register_any_handler_internal(handler, /*include_spec*/ true); + pub(crate) fn register_tool(&mut self, handler: Arc) { + self.register_tool_internal(handler, /*include_spec*/ true); } - pub(crate) fn register_any_handler_without_spec(&mut self, handler: Arc) { - self.register_any_handler_internal(handler, /*include_spec*/ false); + pub(crate) fn register_tool_without_spec(&mut self, handler: Arc) { + self.register_tool_internal(handler, /*include_spec*/ false); } - fn register_any_handler_internal( - &mut self, - handler: Arc, - include_spec: bool, - ) { + fn register_tool_internal(&mut self, handler: Arc, include_spec: bool) { let name = handler.tool_name(); if self.handlers.contains_key(&name) { error_or_panic(format!("handler for tool {name} already registered")); return; } - if include_spec && let Some(spec) = handler.spec() { + if include_spec + && handler.exposure() == ToolExposure::Direct + && let Some(spec) = handler.spec() + { self.push_spec(spec); } @@ -590,12 +600,8 @@ impl ToolRegistryBuilder { return; } - if let Some(spec) = executor.spec() { - self.push_spec(spec); - } - - let handler: Arc = Arc::new(ExtensionToolHandler::new(executor)); - self.handlers.insert(tool_name, handler); + let handler: Arc = Arc::new(ExtensionToolHandler::new(executor)); + self.register_tool_internal(handler, /*include_spec*/ true); } pub fn build(self) -> (Vec, ToolRegistry) { diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index 1cc68be44d..e27dd8fd31 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -33,10 +33,10 @@ fn handler_looks_up_namespaced_aliases_explicitly() { let namespaced_name = codex_tools::ToolName::namespaced(namespace, tool_name); let plain_handler = Arc::new(TestHandler { tool_name: plain_name.clone(), - }) as Arc; + }) as Arc; let namespaced_handler = Arc::new(TestHandler { tool_name: namespaced_name.clone(), - }) as Arc; + }) as Arc; let registry = ToolRegistry::new(HashMap::from([ (plain_name.clone(), Arc::clone(&plain_handler)), (namespaced_name.clone(), Arc::clone(&namespaced_handler)), diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index abb0e18976..d386d6dae3 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -17,11 +17,9 @@ use codex_protocol::models::ResponseItem; use codex_protocol::models::SearchToolCallParams; use codex_protocol::models::ShellToolCallParams; use codex_tools::DiscoverableTool; -use codex_tools::ResponsesApiNamespaceTool; use codex_tools::ToolName; use codex_tools::ToolSpec; use codex_tools::ToolsConfig; -use std::collections::HashSet; use std::sync::Arc; use tokio_util::sync::CancellationToken; use tracing::instrument; @@ -66,21 +64,11 @@ impl ToolRouter { dynamic_tools, ); let (specs, registry) = builder.build(); - let deferred_dynamic_tools = dynamic_tools - .iter() - .filter(|tool| tool.defer_loading) - .map(|tool| ToolName::new(tool.namespace.clone(), tool.name.clone())) - .collect::>(); let model_visible_specs = specs - .iter() - .filter_map(|spec| { - if config.code_mode_only_enabled - && codex_code_mode::is_code_mode_nested_tool(spec.name()) - { - return None; - } - - filter_deferred_dynamic_tool_spec(spec.clone(), &deferred_dynamic_tools) + .into_iter() + .filter(|spec| { + !config.code_mode_only_enabled + || !codex_code_mode::is_code_mode_nested_tool(spec.name()) }) .collect(); @@ -232,38 +220,6 @@ pub(crate) fn extension_tool_executors(session: &Session) -> Vec, -) -> Option { - if deferred_dynamic_tools.is_empty() { - return Some(spec); - } - - match spec { - ToolSpec::Function(tool) => { - if deferred_dynamic_tools.contains(&ToolName::plain(tool.name.as_str())) { - None - } else { - Some(ToolSpec::Function(tool)) - } - } - ToolSpec::Namespace(mut namespace) => { - let namespace_name = namespace.name.clone(); - namespace.tools.retain(|tool| match tool { - ResponsesApiNamespaceTool::Function(tool) => !deferred_dynamic_tools.contains( - &ToolName::namespaced(namespace_name.as_str(), tool.name.as_str()), - ), - }); - if namespace.tools.is_empty() { - None - } else { - Some(ToolSpec::Namespace(namespace)) - } - } - spec => Some(spec), - } -} #[cfg(test)] #[path = "router_tests.rs"] mod tests; diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 86ff84f8e8..0ac670921b 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -44,7 +44,8 @@ use crate::tools::handlers::view_image_spec::ViewImageToolOptions; use crate::tools::hosted_spec::WebSearchToolOptions; use crate::tools::hosted_spec::create_image_generation_tool; use crate::tools::hosted_spec::create_web_search_tool; -use crate::tools::registry::AnyToolHandler; +use crate::tools::registry::RegisteredTool; +use crate::tools::registry::ToolExposure; use crate::tools::registry::ToolRegistryBuilder; use crate::tools::spec_plan_types::ToolRegistryBuildParams; use crate::tools::spec_plan_types::agent_type_description; @@ -52,13 +53,11 @@ use codex_extension_api::ExtensionToolExecutor; use codex_protocol::openai_models::ConfigShellToolType; use codex_tools::ResponsesApiNamespaceTool; use codex_tools::ToolEnvironmentMode; -use codex_tools::ToolName; use codex_tools::ToolSpec; use codex_tools::ToolsConfig; use codex_tools::collect_code_mode_exec_prompt_tool_definitions; use codex_tools::default_namespace_description; use std::collections::BTreeMap; -use std::collections::HashSet; use std::sync::Arc; pub fn build_tool_registry_builder( @@ -66,40 +65,34 @@ pub fn build_tool_registry_builder( params: ToolRegistryBuildParams<'_>, ) -> ToolRegistryBuilder { let mut builder = ToolRegistryBuilder::new(); - let all_deferred_tools = params - .deferred_mcp_tools - .into_iter() - .flatten() - .map(codex_mcp::ToolInfo::canonical_tool_name) - .chain( - params - .dynamic_tools - .iter() - .filter(|tool| tool.defer_loading) - .map(|tool| ToolName::new(tool.namespace.clone(), tool.name.clone())), - ) - .collect::>(); let handlers = collect_handler_tools(config, params); + let deferred_tools_available = handlers + .iter() + .any(|handler| handler.exposure() == ToolExposure::Deferred); for handler in build_code_mode_handlers( config, &handlers, params.extension_tool_executors, - config.search_tool && !all_deferred_tools.is_empty(), + config.search_tool && deferred_tools_available, ) { - builder.register_any_handler(handler); + builder.register_tool(handler); } let mut non_deferred_specs = Vec::new(); let mut deferred_search_infos = Vec::new(); for handler in &handlers { - let tool_name = handler.tool_name(); - if all_deferred_tools.contains(&tool_name) { - if let Some(search_info) = handler.search_info() { - deferred_search_infos.push(search_info); + match handler.exposure() { + ToolExposure::Direct => { + if let Some(spec) = handler.spec() { + non_deferred_specs.push(spec); + } + } + ToolExposure::Deferred => { + if let Some(search_info) = handler.search_info() { + deferred_search_infos.push(search_info); + } } - } else if let Some(spec) = handler.spec() { - non_deferred_specs.push(spec); } } @@ -127,7 +120,7 @@ pub fn build_tool_registry_builder( } for handler in handlers { - builder.register_any_handler_without_spec(handler); + builder.register_tool_without_spec(handler); } if config.search_tool && config.namespace_tools && !deferred_search_infos.is_empty() { @@ -143,10 +136,10 @@ pub fn build_tool_registry_builder( fn build_code_mode_handlers( config: &ToolsConfig, - handlers: &[Arc], + handlers: &[Arc], extension_tool_executors: &[Arc], deferred_tools_available: bool, -) -> Vec> { +) -> Vec> { if !config.code_mode_enabled { return vec![]; } @@ -251,9 +244,9 @@ fn code_mode_namespace_descriptions( fn collect_handler_tools( config: &ToolsConfig, params: ToolRegistryBuildParams<'_>, -) -> Vec> { +) -> Vec> { let exec_permission_approvals_enabled = config.exec_permission_approvals_enabled; - let mut handlers = Vec::>::new(); + let mut handlers = Vec::>::new(); if config.environment_mode.has_environment() { let include_environment_id = @@ -430,7 +423,10 @@ fn collect_handler_tools( if let Some(deferred_mcp_tools) = params.deferred_mcp_tools { for tool in deferred_mcp_tools { - handlers.push(Arc::new(McpHandler::new(tool.clone()))); + handlers.push(Arc::new(McpHandler::with_exposure( + tool.clone(), + ToolExposure::Deferred, + ))); } } diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index 86e4750b32..2a87ba35d7 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -79,6 +79,7 @@ pub use tool_discovery::ToolSearchSourceInfo; pub use tool_discovery::collect_request_plugin_install_entries; pub use tool_discovery::filter_request_plugin_install_discoverable_tools_for_client; pub use tool_executor::ToolExecutor; +pub use tool_executor::ToolExposure; pub use tool_output::JsonToolOutput; pub use tool_output::ToolOutput; pub use tool_payload::ToolPayload; diff --git a/codex-rs/tools/src/tool_executor.rs b/codex-rs/tools/src/tool_executor.rs index 8c38e9fc8c..1d7169ca1d 100644 --- a/codex-rs/tools/src/tool_executor.rs +++ b/codex-rs/tools/src/tool_executor.rs @@ -5,6 +5,14 @@ use crate::ToolName; use crate::ToolOutput; use crate::ToolSpec; +/// Controls whether a tool is exposed in the initial model-visible tool list +/// or registered for later discovery. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ToolExposure { + Direct, + Deferred, +} + /// Shared runtime contract for model-visible tools. /// /// Implementations keep the model-visible spec tied to the executable runtime. @@ -20,6 +28,10 @@ pub trait ToolExecutor: Send + Sync { None } + fn exposure(&self) -> ToolExposure { + ToolExposure::Direct + } + fn supports_parallel_tool_calls(&self) -> bool { false }