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 }