From 6d65686313d484db0bb1212cf1a8e1915282024d Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 14 May 2026 11:23:57 +0200 Subject: [PATCH] feat: make ToolExecutor an async trait (#22560) ## Why `codex_tools::ToolExecutor` keeps a tool spec attached to its runtime handler, but extension tools still carried a parallel `ExtensionToolFuture` / `ExtensionToolExecutor` shape. That made extension-owned tools look different from host tools even though routing, registration, and execution need the same abstraction. This PR makes the shared executor contract directly async and lets extension tools implement it too, so host tools and extension tools can move through the same registration path. ## What changed - Changed `ToolExecutor::handle` to an `async fn` using `async-trait`, and updated built-in tool handlers to implement the async trait directly. - Replaced the bespoke `ExtensionToolFuture` contract with a marker `ExtensionToolExecutor` over `ToolExecutor`, re-exporting `ToolExecutor` from `codex-extension-api`. - Updated the memories extension tools to implement the shared executor trait. - Split tool-router construction into collected executors plus hosted model specs, keeping hosted tools like web search and image generation separate from executable handlers. - Updated spec/router tests and extension-tool stubs for the new executor shape. ## Verification - Not run locally. --- codex-rs/Cargo.lock | 2 + .../src/tools/code_mode/execute_handler.rs | 1 + .../core/src/tools/code_mode/wait_handler.rs | 1 + .../agent_jobs/report_agent_job_result.rs | 1 + .../agent_jobs/spawn_agents_on_csv.rs | 1 + .../core/src/tools/handlers/apply_patch.rs | 1 + codex-rs/core/src/tools/handlers/dynamic.rs | 1 + .../src/tools/handlers/extension_tools.rs | 12 +- .../src/tools/handlers/goal/create_goal.rs | 1 + .../core/src/tools/handlers/goal/get_goal.rs | 1 + .../src/tools/handlers/goal/update_goal.rs | 1 + codex-rs/core/src/tools/handlers/mcp.rs | 1 + .../list_mcp_resource_templates.rs | 1 + .../mcp_resource/list_mcp_resources.rs | 1 + .../mcp_resource/read_mcp_resource.rs | 1 + .../handlers/multi_agents/close_agent.rs | 8 +- .../handlers/multi_agents/resume_agent.rs | 8 +- .../tools/handlers/multi_agents/send_input.rs | 1 + .../src/tools/handlers/multi_agents/spawn.rs | 8 +- .../src/tools/handlers/multi_agents/wait.rs | 1 + .../handlers/multi_agents_v2/close_agent.rs | 8 +- .../handlers/multi_agents_v2/followup_task.rs | 1 + .../handlers/multi_agents_v2/list_agents.rs | 1 + .../handlers/multi_agents_v2/send_message.rs | 1 + .../tools/handlers/multi_agents_v2/spawn.rs | 8 +- .../tools/handlers/multi_agents_v2/wait.rs | 1 + codex-rs/core/src/tools/handlers/plan.rs | 1 + .../src/tools/handlers/request_permissions.rs | 1 + .../tools/handlers/request_plugin_install.rs | 1 + .../src/tools/handlers/request_user_input.rs | 1 + .../src/tools/handlers/shell/shell_command.rs | 1 + codex-rs/core/src/tools/handlers/test_sync.rs | 1 + .../core/src/tools/handlers/tool_search.rs | 1 + .../handlers/unified_exec/exec_command.rs | 1 + .../handlers/unified_exec/write_stdin.rs | 1 + .../core/src/tools/handlers/view_image.rs | 1 + codex-rs/core/src/tools/registry.rs | 7 - codex-rs/core/src/tools/registry_tests.rs | 5 +- codex-rs/core/src/tools/router.rs | 15 +- codex-rs/core/src/tools/router_tests.rs | 28 +-- codex-rs/core/src/tools/spec.rs | 22 ++- codex-rs/core/src/tools/spec_plan.rs | 162 ++++++++++-------- codex-rs/core/src/tools/spec_plan_tests.rs | 50 +++--- codex-rs/core/src/tools/spec_tests.rs | 26 ++- .../src/tools/tool_dispatch_trace_tests.rs | 1 + .../ext/extension-api/src/contributors.rs | 1 - .../extension-api/src/contributors/tools.rs | 31 +--- codex-rs/ext/extension-api/src/lib.rs | 2 +- codex-rs/ext/memories/Cargo.toml | 1 + codex-rs/ext/memories/src/tools/list.rs | 43 ++--- codex-rs/ext/memories/src/tools/read.rs | 37 ++-- codex-rs/ext/memories/src/tools/search.rs | 27 +-- codex-rs/tools/Cargo.toml | 1 + codex-rs/tools/src/tool_executor.rs | 8 +- 54 files changed, 313 insertions(+), 237 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index e37fc6288..9ecc70837 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3137,6 +3137,7 @@ dependencies = [ name = "codex-memories-extension" version = "0.0.0" dependencies = [ + "async-trait", "codex-core", "codex-extension-api", "codex-features", @@ -3729,6 +3730,7 @@ dependencies = [ name = "codex-tools" version = "0.0.0" dependencies = [ + "async-trait", "codex-app-server-protocol", "codex-code-mode", "codex-features", diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index eaa415020..26ef021ad 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -87,6 +87,7 @@ impl CodeModeExecuteHandler { } } +#[async_trait::async_trait] impl ToolExecutor for CodeModeExecuteHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index 070b9270d..4f828c891 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -41,6 +41,7 @@ where }) } +#[async_trait::async_trait] impl ToolExecutor for CodeModeWaitHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs b/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs index 7797c8c61..e08eb4b5f 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs @@ -12,6 +12,7 @@ use super::*; pub struct ReportAgentJobResultHandler; +#[async_trait::async_trait] impl ToolExecutor for ReportAgentJobResultHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs b/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs index 6ad0515e8..a384f2d8c 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs @@ -13,6 +13,7 @@ use super::*; pub struct SpawnAgentsOnCsvHandler; +#[async_trait::async_trait] impl ToolExecutor for SpawnAgentsOnCsvHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 5a709b287..c8e2461e0 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -297,6 +297,7 @@ async fn effective_patch_permissions( ) } +#[async_trait::async_trait] impl ToolExecutor for ApplyPatchHandler { type Output = ApplyPatchToolOutput; diff --git a/codex-rs/core/src/tools/handlers/dynamic.rs b/codex-rs/core/src/tools/handlers/dynamic.rs index 8000b2097..29dc5199e 100644 --- a/codex-rs/core/src/tools/handlers/dynamic.rs +++ b/codex-rs/core/src/tools/handlers/dynamic.rs @@ -60,6 +60,7 @@ impl DynamicToolHandler { } } +#[async_trait::async_trait] impl ToolExecutor for DynamicToolHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index cd94d8869..c4018fe93 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -35,6 +35,7 @@ impl ExtensionToolHandler { } } +#[async_trait::async_trait] impl ToolExecutor for ExtensionToolHandler { type Output = ExtensionToolOutput; @@ -115,7 +116,10 @@ mod tests { struct StubExtensionExecutor; - impl codex_extension_api::ExtensionToolExecutor for StubExtensionExecutor { + #[async_trait::async_trait] + impl codex_extension_api::ToolExecutor for StubExtensionExecutor { + type Output = codex_tools::JsonToolOutput; + fn tool_name(&self) -> codex_tools::ToolName { codex_tools::ToolName::plain("extension_echo") } @@ -141,11 +145,11 @@ mod tests { )) } - fn handle( + async fn handle( &self, _call: codex_tools::ToolCall, - ) -> codex_extension_api::ExtensionToolFuture<'_> { - Box::pin(async { Ok(codex_tools::JsonToolOutput::new(json!({ "ok": true }))) }) + ) -> Result { + Ok(codex_tools::JsonToolOutput::new(json!({ "ok": true }))) } } diff --git a/codex-rs/core/src/tools/handlers/goal/create_goal.rs b/codex-rs/core/src/tools/handlers/goal/create_goal.rs index 34a4c9307..f7575e1e5 100644 --- a/codex-rs/core/src/tools/handlers/goal/create_goal.rs +++ b/codex-rs/core/src/tools/handlers/goal/create_goal.rs @@ -18,6 +18,7 @@ use super::goal_response; pub struct CreateGoalHandler; +#[async_trait::async_trait] impl ToolExecutor for CreateGoalHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/goal/get_goal.rs b/codex-rs/core/src/tools/handlers/goal/get_goal.rs index 3c0813329..20691b8b5 100644 --- a/codex-rs/core/src/tools/handlers/goal/get_goal.rs +++ b/codex-rs/core/src/tools/handlers/goal/get_goal.rs @@ -15,6 +15,7 @@ use super::goal_response; pub struct GetGoalHandler; +#[async_trait::async_trait] impl ToolExecutor for GetGoalHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/goal/update_goal.rs b/codex-rs/core/src/tools/handlers/goal/update_goal.rs index 18b59c1b3..bb8ed10ef 100644 --- a/codex-rs/core/src/tools/handlers/goal/update_goal.rs +++ b/codex-rs/core/src/tools/handlers/goal/update_goal.rs @@ -20,6 +20,7 @@ use super::goal_response; pub struct UpdateGoalHandler; +#[async_trait::async_trait] impl ToolExecutor for UpdateGoalHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 5f700e3a9..f82f74134 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -45,6 +45,7 @@ impl McpHandler { } } +#[async_trait::async_trait] impl ToolExecutor for McpHandler { type Output = McpToolOutput; diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs index 0aea2b335..71749e9cd 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs @@ -26,6 +26,7 @@ use super::serialize_function_output; pub struct ListMcpResourceTemplatesHandler; +#[async_trait::async_trait] impl ToolExecutor for ListMcpResourceTemplatesHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs index b412f811b..08b387376 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs @@ -26,6 +26,7 @@ use super::serialize_function_output; pub struct ListMcpResourcesHandler; +#[async_trait::async_trait] impl ToolExecutor for ListMcpResourcesHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs index 89d061d69..bd8172ac7 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs @@ -26,6 +26,7 @@ use super::serialize_function_output; pub struct ReadMcpResourceHandler; +#[async_trait::async_trait] impl ToolExecutor for ReadMcpResourceHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs index beac3a96f..6459a899f 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs @@ -5,6 +5,7 @@ use codex_tools::ToolSpec; pub(crate) struct Handler; +#[async_trait::async_trait] impl ToolExecutor for Handler { type Output = CloseAgentResult; @@ -16,11 +17,8 @@ impl ToolExecutor for Handler { Some(create_close_agent_tool_v1()) } - fn handle( - &self, - invocation: ToolInvocation, - ) -> impl std::future::Future> + Send { - Box::pin(handle_close_agent(invocation)) + async fn handle(&self, invocation: ToolInvocation) -> Result { + handle_close_agent(invocation).await } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs index 1e9f5d9b0..9dab2d999 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs @@ -7,6 +7,7 @@ use std::sync::Arc; pub(crate) struct Handler; +#[async_trait::async_trait] impl ToolExecutor for Handler { type Output = ResumeAgentResult; @@ -18,11 +19,8 @@ impl ToolExecutor for Handler { Some(create_resume_agent_tool()) } - fn handle( - &self, - invocation: ToolInvocation, - ) -> impl std::future::Future> + Send { - Box::pin(handle_resume_agent(invocation)) + async fn handle(&self, invocation: ToolInvocation) -> Result { + handle_resume_agent(invocation).await } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs b/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs index ac212f755..a6067e5b1 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs @@ -6,6 +6,7 @@ use codex_tools::ToolSpec; pub(crate) struct Handler; +#[async_trait::async_trait] impl ToolExecutor for Handler { type Output = SendInputResult; diff --git a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs index 395e74f13..91f599021 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -22,6 +22,7 @@ impl Handler { } } +#[async_trait::async_trait] impl ToolExecutor for Handler { type Output = SpawnAgentResult; @@ -33,11 +34,8 @@ impl ToolExecutor for Handler { Some(create_spawn_agent_tool_v1(self.options.clone())) } - fn handle( - &self, - invocation: ToolInvocation, - ) -> impl std::future::Future> + Send { - Box::pin(handle_spawn_agent(invocation)) + async fn handle(&self, invocation: ToolInvocation) -> Result { + handle_spawn_agent(invocation).await } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents/wait.rs b/codex-rs/core/src/tools/handlers/multi_agents/wait.rs index aa9d6f658..d325edeeb 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/wait.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/wait.rs @@ -27,6 +27,7 @@ impl Handler { } } +#[async_trait::async_trait] impl ToolExecutor for Handler { type Output = WaitAgentResult; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs index c901df9ae..7b575f825 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs @@ -5,6 +5,7 @@ use codex_tools::ToolSpec; pub(crate) struct Handler; +#[async_trait::async_trait] impl ToolExecutor for Handler { type Output = CloseAgentResult; @@ -16,11 +17,8 @@ impl ToolExecutor for Handler { Some(create_close_agent_tool_v2()) } - fn handle( - &self, - invocation: ToolInvocation, - ) -> impl std::future::Future> + Send { - Box::pin(handle_close_agent(invocation)) + async fn handle(&self, invocation: ToolInvocation) -> Result { + handle_close_agent(invocation).await } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/followup_task.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/followup_task.rs index 608d7f2f6..7d111f041 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/followup_task.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/followup_task.rs @@ -8,6 +8,7 @@ use codex_tools::ToolSpec; pub(crate) struct Handler; +#[async_trait::async_trait] impl ToolExecutor for Handler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/list_agents.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/list_agents.rs index 527d5a29e..8b0ee551e 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/list_agents.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/list_agents.rs @@ -5,6 +5,7 @@ use codex_tools::ToolSpec; pub(crate) struct Handler; +#[async_trait::async_trait] impl ToolExecutor for Handler { type Output = ListAgentsResult; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/send_message.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/send_message.rs index 71909afb7..584feec61 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/send_message.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/send_message.rs @@ -8,6 +8,7 @@ use codex_tools::ToolSpec; pub(crate) struct Handler; +#[async_trait::async_trait] impl ToolExecutor for Handler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs index 1ab59cc3f..3ad7b8714 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs @@ -24,6 +24,7 @@ impl Handler { } } +#[async_trait::async_trait] impl ToolExecutor for Handler { type Output = SpawnAgentResult; @@ -35,11 +36,8 @@ impl ToolExecutor for Handler { Some(create_spawn_agent_tool_v2(self.options.clone())) } - fn handle( - &self, - invocation: ToolInvocation, - ) -> impl std::future::Future> + Send { - Box::pin(handle_spawn_agent(invocation)) + async fn handle(&self, invocation: ToolInvocation) -> Result { + handle_spawn_agent(invocation).await } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs index 5f246d679..53c976c9d 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs @@ -19,6 +19,7 @@ impl Handler { } } +#[async_trait::async_trait] impl ToolExecutor for Handler { type Output = WaitAgentResult; diff --git a/codex-rs/core/src/tools/handlers/plan.rs b/codex-rs/core/src/tools/handlers/plan.rs index 42f03e4e3..9868c77ca 100644 --- a/codex-rs/core/src/tools/handlers/plan.rs +++ b/codex-rs/core/src/tools/handlers/plan.rs @@ -44,6 +44,7 @@ impl ToolOutput for PlanToolOutput { } } +#[async_trait::async_trait] impl ToolExecutor for PlanHandler { type Output = PlanToolOutput; diff --git a/codex-rs/core/src/tools/handlers/request_permissions.rs b/codex-rs/core/src/tools/handlers/request_permissions.rs index 5d49ad861..2f2e71677 100644 --- a/codex-rs/core/src/tools/handlers/request_permissions.rs +++ b/codex-rs/core/src/tools/handlers/request_permissions.rs @@ -15,6 +15,7 @@ use codex_tools::ToolSpec; pub struct RequestPermissionsHandler; +#[async_trait::async_trait] impl ToolExecutor for RequestPermissionsHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install.rs b/codex-rs/core/src/tools/handlers/request_plugin_install.rs index 9e76352aa..6073831e0 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install.rs @@ -50,6 +50,7 @@ impl RequestPluginInstallHandler { } } +#[async_trait::async_trait] impl ToolExecutor for RequestPluginInstallHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/request_user_input.rs b/codex-rs/core/src/tools/handlers/request_user_input.rs index ae0a2c769..7597cb600 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input.rs @@ -19,6 +19,7 @@ pub struct RequestUserInputHandler { pub available_modes: Vec, } +#[async_trait::async_trait] impl ToolExecutor for RequestUserInputHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/shell/shell_command.rs b/codex-rs/core/src/tools/handlers/shell/shell_command.rs index 250f601ce..54d74f13b 100644 --- a/codex-rs/core/src/tools/handlers/shell/shell_command.rs +++ b/codex-rs/core/src/tools/handlers/shell/shell_command.rs @@ -127,6 +127,7 @@ impl From for ShellCommandHandler { } } +#[async_trait::async_trait] impl ToolExecutor for ShellCommandHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/test_sync.rs b/codex-rs/core/src/tools/handlers/test_sync.rs index 091f0ffce..feee6470a 100644 --- a/codex-rs/core/src/tools/handlers/test_sync.rs +++ b/codex-rs/core/src/tools/handlers/test_sync.rs @@ -56,6 +56,7 @@ fn barrier_map() -> &'static tokio::sync::Mutex> { BARRIERS.get_or_init(|| tokio::sync::Mutex::new(HashMap::new())) } +#[async_trait::async_trait] impl ToolExecutor for TestSyncHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index aaa313858..579826ac0 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -52,6 +52,7 @@ impl ToolSearchHandler { } } +#[async_trait::async_trait] impl ToolExecutor for ToolSearchHandler { type Output = ToolSearchOutput; diff --git a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs index e27c56576..2df924b5f 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs @@ -68,6 +68,7 @@ impl ExecCommandHandler { } } +#[async_trait::async_trait] impl ToolExecutor for ExecCommandHandler { type Output = ExecCommandToolOutput; diff --git a/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs b/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs index cc79615c1..29ee4b4ea 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs @@ -31,6 +31,7 @@ struct WriteStdinArgs { pub struct WriteStdinHandler; +#[async_trait::async_trait] impl ToolExecutor for WriteStdinHandler { type Output = ExecCommandToolOutput; diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index 1587e90dd..7727fb437 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -62,6 +62,7 @@ enum ViewImageDetail { Original, } +#[async_trait::async_trait] impl ToolExecutor for ViewImageHandler { type Output = ViewImageOutput; diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 7e565bd31..730c3b7ef 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -637,13 +637,6 @@ impl ToolRegistryBuilder { self.specs.push(spec); } - pub fn register_handler(&mut self, handler: Arc) - where - H: ToolHandler + 'static, - { - self.register_tool(handler); - } - pub(crate) fn register_tool(&mut self, handler: Arc) { self.register_tool_internal(handler, /*include_spec*/ true); } diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index e27dd8fd3..b6c8eadad 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -8,6 +8,7 @@ struct TestHandler { tool_name: codex_tools::ToolName, } +#[async_trait::async_trait] impl ToolExecutor for TestHandler { type Output = crate::tools::context::FunctionToolOutput; @@ -65,9 +66,9 @@ fn handler_looks_up_namespaced_aliases_explicitly() { } #[test] -fn register_handler_adds_handler_and_spec() { +fn register_tool_adds_executor_and_spec() { let mut builder = ToolRegistryBuilder::new(); - builder.register_handler(Arc::new(GetGoalHandler)); + builder.register_tool(Arc::new(GetGoalHandler)); let (specs, registry) = builder.build(); diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 23f020773..0f111922c 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -5,10 +5,12 @@ use crate::tools::context::SharedTurnDiffTracker; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::registry::AnyToolResult; +use crate::tools::registry::RegisteredTool; use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::registry::ToolExposure; use crate::tools::registry::ToolRegistry; -use crate::tools::spec::build_specs_with_discoverable_tools; +use crate::tools::spec::collect_tool_router_parts; +use crate::tools::spec_plan::build_tool_registry_builder_from_executors; use codex_extension_api::ExtensionToolExecutor; use codex_mcp::ToolInfo; use codex_protocol::dynamic_tools::DynamicToolSpec; @@ -53,7 +55,7 @@ impl ToolRouter { extension_tool_executors, dynamic_tools, } = params; - let builder = build_specs_with_discoverable_tools( + let parts = collect_tool_router_parts( config, mcp_tools, deferred_mcp_tools, @@ -61,6 +63,15 @@ impl ToolRouter { &extension_tool_executors, dynamic_tools, ); + Self::from_executors(config, parts.executors, parts.hosted_specs) + } + + pub(crate) fn from_executors( + config: &ToolsConfig, + executors: Vec>, + hosted_specs: Vec, + ) -> Self { + let builder = build_tool_registry_builder_from_executors(config, executors, hosted_specs); let (specs, registry) = builder.build(); let model_visible_specs = specs .into_iter() diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index 921b52a86..dd791d913 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -8,9 +8,9 @@ use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionRegistry; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ExtensionToolExecutor; -use codex_extension_api::ExtensionToolOutput; use codex_extension_api::ResponsesApiTool; use codex_extension_api::ToolCall as ExtensionToolCall; +use codex_extension_api::ToolExecutor; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::ResponseInputItem; @@ -44,7 +44,10 @@ impl codex_extension_api::ToolContributor for ExtensionEchoContributor { struct ExtensionEchoExecutor; -impl ExtensionToolExecutor for ExtensionEchoExecutor { +#[async_trait::async_trait] +impl ToolExecutor for ExtensionEchoExecutor { + type Output = codex_tools::JsonToolOutput; + fn tool_name(&self) -> ToolName { ToolName::namespaced("extension/", "echo") } @@ -72,16 +75,17 @@ impl ExtensionToolExecutor for ExtensionEchoExecutor { })) } - fn handle(&self, call: ExtensionToolCall) -> codex_extension_api::ExtensionToolFuture<'_> { - Box::pin(async move { - let arguments: serde_json::Value = serde_json::from_str(call.function_arguments()?) - .expect("test arguments should parse"); - Ok(ExtensionToolOutput::new(json!({ - "arguments": arguments, - "callId": call.call_id.clone(), - "ok": true, - }))) - }) + async fn handle( + &self, + call: ExtensionToolCall, + ) -> Result { + let arguments: serde_json::Value = + serde_json::from_str(call.function_arguments()?).expect("test arguments should parse"); + Ok(codex_tools::JsonToolOutput::new(json!({ + "arguments": arguments, + "callId": call.call_id, + "ok": true, + }))) } } diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 92aeed011..64dfdb831 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -7,8 +7,9 @@ use crate::tools::handlers::multi_agents_common::DEFAULT_WAIT_TIMEOUT_MS; use crate::tools::handlers::multi_agents_common::MAX_WAIT_TIMEOUT_MS; use crate::tools::handlers::multi_agents_common::MIN_WAIT_TIMEOUT_MS; use crate::tools::handlers::multi_agents_spec::WaitAgentTimeoutOptions; -use crate::tools::registry::ToolRegistryBuilder; -use crate::tools::spec_plan::build_tool_registry_builder; +use crate::tools::registry::RegisteredTool; +use crate::tools::spec_plan::collect_tool_executors; +use crate::tools::spec_plan::hosted_model_tool_specs; use crate::tools::spec_plan_types::ToolRegistryBuildParams; use codex_extension_api::ExtensionToolExecutor; use codex_mcp::ToolInfo; @@ -28,14 +29,19 @@ pub(crate) fn tool_user_shell_type(user_shell: &Shell) -> ToolUserShellType { } } -pub(crate) fn build_specs_with_discoverable_tools( +pub(crate) struct ToolRouterParts { + pub(crate) executors: Vec>, + pub(crate) hosted_specs: Vec, +} + +pub(crate) fn collect_tool_router_parts( config: &ToolsConfig, mcp_tools: Option>, deferred_mcp_tools: Option>, discoverable_tools: Option>, extension_tool_executors: &[Arc], dynamic_tools: &[DynamicToolSpec], -) -> ToolRegistryBuilder { +) -> ToolRouterParts { let default_agent_type_description = crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new()); let (min_wait_timeout_ms, max_wait_timeout_ms, default_wait_timeout_ms) = @@ -61,7 +67,7 @@ pub(crate) fn build_specs_with_discoverable_tools( DEFAULT_WAIT_TIMEOUT_MS, ) }; - build_tool_registry_builder( + let executors = collect_tool_executors( config, ToolRegistryBuildParams { mcp_tools: mcp_tools.as_deref(), @@ -76,7 +82,11 @@ pub(crate) fn build_specs_with_discoverable_tools( max_timeout_ms: max_wait_timeout_ms, }, }, - ) + ); + ToolRouterParts { + executors, + hosted_specs: hosted_model_tool_specs(config), + } } #[cfg(test)] diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index db530f347..24f59a0f0 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -62,51 +62,46 @@ use std::collections::HashSet; use std::sync::Arc; use tracing::warn; -pub fn build_tool_registry_builder( +pub(crate) fn build_tool_registry_builder_from_executors( config: &ToolsConfig, - params: ToolRegistryBuildParams<'_>, + executors: Vec>, + hosted_specs: Vec, ) -> ToolRegistryBuilder { let mut builder = ToolRegistryBuilder::new(); - let handlers = collect_handler_tools(config, params); - let deferred_tools_available = handlers + let deferred_tools_available = executors .iter() - .any(|handler| handler.exposure() == ToolExposure::Deferred); + .any(|executor| executor.exposure() == ToolExposure::Deferred); - for handler in build_code_mode_handlers( + for executor in build_code_mode_executors( config, - &handlers, + &executors, config.search_tool && deferred_tools_available, ) { - builder.register_tool(handler); + builder.register_tool(executor); } let mut non_deferred_specs = Vec::new(); let mut deferred_search_infos = Vec::new(); - for handler in &handlers { - match handler.exposure() { + for executor in &executors { + match executor.exposure() { ToolExposure::Direct | ToolExposure::DirectModelOnly => { - if let Some(spec) = handler.spec() { - non_deferred_specs.push((spec, handler.exposure())); + if let Some(spec) = executor.spec() { + non_deferred_specs.push((spec, executor.exposure())); } } ToolExposure::Deferred => { - if let Some(search_info) = handler.search_info() { + if let Some(search_info) = executor.search_info() { deferred_search_infos.push(search_info); } } } } - if let Some(web_search_tool) = create_web_search_tool(WebSearchToolOptions { - web_search_mode: config.web_search_mode, - web_search_config: config.web_search_config.as_ref(), - web_search_tool_type: config.web_search_tool_type, - }) { - non_deferred_specs.push((web_search_tool, ToolExposure::Direct)); - } - if config.image_gen_tool { - non_deferred_specs.push((create_image_generation_tool("png"), ToolExposure::Direct)); - } + non_deferred_specs.extend( + hosted_specs + .into_iter() + .map(|spec| (spec, ToolExposure::Direct)), + ); let non_deferred_specs = non_deferred_specs .into_iter() @@ -126,34 +121,49 @@ pub fn build_tool_registry_builder( builder.push_spec(spec); } - for handler in handlers { - builder.register_tool_without_spec(handler); + for executor in executors { + builder.register_tool_without_spec(executor); } if config.search_tool && config.namespace_tools && !deferred_search_infos.is_empty() { - builder.register_handler(Arc::new(ToolSearchHandler::new(deferred_search_infos))); + builder.register_tool(Arc::new(ToolSearchHandler::new(deferred_search_infos))); } builder } -fn build_code_mode_handlers( +pub(crate) fn hosted_model_tool_specs(config: &ToolsConfig) -> Vec { + let mut specs = Vec::new(); + if let Some(web_search_tool) = create_web_search_tool(WebSearchToolOptions { + web_search_mode: config.web_search_mode, + web_search_config: config.web_search_config.as_ref(), + web_search_tool_type: config.web_search_tool_type, + }) { + specs.push(web_search_tool); + } + if config.image_gen_tool { + specs.push(create_image_generation_tool("png")); + } + specs +} + +fn build_code_mode_executors( config: &ToolsConfig, - handlers: &[Arc], + executors: &[Arc], deferred_tools_available: bool, ) -> Vec> { if !config.code_mode_enabled { return vec![]; } - let code_mode_nested_tool_specs = handlers + let code_mode_nested_tool_specs = executors .iter() - .filter_map(|handler| { - if handler.exposure() == ToolExposure::DirectModelOnly { + .filter_map(|executor| { + if executor.exposure() == ToolExposure::DirectModelOnly { return None; } - handler.spec() + executor.spec() }) .collect::>(); let namespace_descriptions = code_mode_namespace_descriptions(&code_mode_nested_tool_specs); @@ -244,32 +254,32 @@ fn code_mode_namespace_descriptions( namespace_descriptions } -fn collect_handler_tools( +pub(crate) fn collect_tool_executors( config: &ToolsConfig, params: ToolRegistryBuildParams<'_>, ) -> Vec> { let exec_permission_approvals_enabled = config.exec_permission_approvals_enabled; - let mut handlers = Vec::>::new(); + let mut executors = Vec::>::new(); if config.environment_mode.has_environment() { let include_environment_id = matches!(config.environment_mode, ToolEnvironmentMode::Multiple); match &config.shell_type { ConfigShellToolType::UnifiedExec => { - handlers.push(Arc::new(ExecCommandHandler::new( + executors.push(Arc::new(ExecCommandHandler::new( ExecCommandHandlerOptions { allow_login_shell: config.allow_login_shell, exec_permission_approvals_enabled, include_environment_id, }, ))); - handlers.push(Arc::new(WriteStdinHandler)); + executors.push(Arc::new(WriteStdinHandler)); } ConfigShellToolType::Disabled => {} ConfigShellToolType::Default | ConfigShellToolType::Local | ConfigShellToolType::ShellCommand => { - handlers.push(Arc::new(ShellCommandHandler::new( + executors.push(Arc::new(ShellCommandHandler::new( ShellCommandHandlerOptions { backend_config: config.shell_command_backend, allow_login_shell: config.allow_login_shell, @@ -285,7 +295,7 @@ fn collect_handler_tools( { match &config.shell_type { ConfigShellToolType::UnifiedExec => { - handlers.push(Arc::new(ShellCommandHandler::from( + executors.push(Arc::new(ShellCommandHandler::from( config.shell_command_backend, ))); } @@ -297,31 +307,31 @@ fn collect_handler_tools( } if params.mcp_tools.is_some() { - handlers.push(Arc::new(ListMcpResourcesHandler)); - handlers.push(Arc::new(ListMcpResourceTemplatesHandler)); - handlers.push(Arc::new(ReadMcpResourceHandler)); + executors.push(Arc::new(ListMcpResourcesHandler)); + executors.push(Arc::new(ListMcpResourceTemplatesHandler)); + executors.push(Arc::new(ReadMcpResourceHandler)); } - handlers.push(Arc::new(PlanHandler)); + executors.push(Arc::new(PlanHandler)); if config.goal_tools { - handlers.push(Arc::new(GetGoalHandler)); - handlers.push(Arc::new(CreateGoalHandler)); - handlers.push(Arc::new(UpdateGoalHandler)); + executors.push(Arc::new(GetGoalHandler)); + executors.push(Arc::new(CreateGoalHandler)); + executors.push(Arc::new(UpdateGoalHandler)); } - handlers.push(Arc::new(RequestUserInputHandler { + executors.push(Arc::new(RequestUserInputHandler { available_modes: config.request_user_input_available_modes.clone(), })); if config.request_permissions_tool_enabled { - handlers.push(Arc::new(RequestPermissionsHandler)); + executors.push(Arc::new(RequestPermissionsHandler)); } if config.tool_suggest && let Some(discoverable_tools) = params.discoverable_tools.filter(|tools| !tools.is_empty()) { - handlers.push(Arc::new(RequestPluginInstallHandler::new( + executors.push(Arc::new(RequestPluginInstallHandler::new( discoverable_tools, ))); } @@ -329,7 +339,7 @@ fn collect_handler_tools( if config.environment_mode.has_environment() && config.apply_patch_tool_type.is_some() { let include_environment_id = matches!(config.environment_mode, ToolEnvironmentMode::Multiple); - handlers.push(Arc::new(ApplyPatchHandler::new(include_environment_id))); + executors.push(Arc::new(ApplyPatchHandler::new(include_environment_id))); } if config @@ -337,13 +347,13 @@ fn collect_handler_tools( .iter() .any(|tool| tool == "test_sync_tool") { - handlers.push(Arc::new(TestSyncHandler)); + executors.push(Arc::new(TestSyncHandler)); } if config.environment_mode.has_environment() { let include_environment_id = matches!(config.environment_mode, ToolEnvironmentMode::Multiple); - handlers.push(Arc::new(ViewImageHandler::new(ViewImageToolOptions { + executors.push(Arc::new(ViewImageHandler::new(ViewImageToolOptions { can_request_original_image_detail: config.can_request_original_image_detail, include_environment_id, }))); @@ -358,7 +368,7 @@ fn collect_handler_tools( }; let agent_type_description = agent_type_description(config, params.default_agent_type_description); - handlers.push(multi_agent_v2_handler( + executors.push(multi_agent_v2_handler( SpawnAgentHandlerV2::new(SpawnAgentToolOptions { available_models: config.available_models.clone(), agent_type_description, @@ -369,18 +379,18 @@ fn collect_handler_tools( }), exposure, )); - handlers.push(multi_agent_v2_handler(SendMessageHandlerV2, exposure)); - handlers.push(multi_agent_v2_handler(FollowupTaskHandlerV2, exposure)); - handlers.push(multi_agent_v2_handler( + executors.push(multi_agent_v2_handler(SendMessageHandlerV2, exposure)); + executors.push(multi_agent_v2_handler(FollowupTaskHandlerV2, exposure)); + executors.push(multi_agent_v2_handler( WaitAgentHandlerV2::new(params.wait_agent_timeouts), exposure, )); - handlers.push(multi_agent_v2_handler(CloseAgentHandlerV2, exposure)); - handlers.push(multi_agent_v2_handler(ListAgentsHandlerV2, exposure)); + executors.push(multi_agent_v2_handler(CloseAgentHandlerV2, exposure)); + executors.push(multi_agent_v2_handler(ListAgentsHandlerV2, exposure)); } else { let agent_type_description = agent_type_description(config, params.default_agent_type_description); - handlers.push(Arc::new(SpawnAgentHandler::new(SpawnAgentToolOptions { + executors.push(Arc::new(SpawnAgentHandler::new(SpawnAgentToolOptions { available_models: config.available_models.clone(), agent_type_description, hide_agent_type_model_reasoning: config.hide_spawn_agent_metadata, @@ -388,29 +398,29 @@ fn collect_handler_tools( usage_hint_text: config.spawn_agent_usage_hint_text.clone(), max_concurrent_threads_per_session: config.max_concurrent_threads_per_session, }))); - handlers.push(Arc::new(SendInputHandler)); - handlers.push(Arc::new(ResumeAgentHandler)); - handlers.push(Arc::new(WaitAgentHandler::new(params.wait_agent_timeouts))); - handlers.push(Arc::new(CloseAgentHandler)); + executors.push(Arc::new(SendInputHandler)); + executors.push(Arc::new(ResumeAgentHandler)); + executors.push(Arc::new(WaitAgentHandler::new(params.wait_agent_timeouts))); + executors.push(Arc::new(CloseAgentHandler)); } } if config.agent_jobs_tools { - handlers.push(Arc::new(SpawnAgentsOnCsvHandler)); + executors.push(Arc::new(SpawnAgentsOnCsvHandler)); if config.agent_jobs_worker_tools { - handlers.push(Arc::new(ReportAgentJobResultHandler)); + executors.push(Arc::new(ReportAgentJobResultHandler)); } } if let Some(mcp_tools) = params.mcp_tools { for tool in mcp_tools { - handlers.push(Arc::new(McpHandler::new(tool.clone()))); + executors.push(Arc::new(McpHandler::new(tool.clone()))); } } if let Some(deferred_mcp_tools) = params.deferred_mcp_tools { for tool in deferred_mcp_tools { - handlers.push(Arc::new(McpHandler::with_exposure( + executors.push(Arc::new(McpHandler::with_exposure( tool.clone(), ToolExposure::Deferred, ))); @@ -426,26 +436,26 @@ fn collect_handler_tools( continue; }; - handlers.push(handler); + executors.push(handler); } - append_extension_tool_handlers(config, params.extension_tool_executors, &mut handlers); + append_extension_tool_executors(config, params.extension_tool_executors, &mut executors); - handlers + executors } -fn append_extension_tool_handlers( +fn append_extension_tool_executors( config: &ToolsConfig, executors: &[Arc], - handlers: &mut Vec>, + registered_executors: &mut Vec>, ) { if executors.is_empty() { return; } - let mut reserved_tool_names = handlers + let mut reserved_tool_names = registered_executors .iter() - .map(|handler| handler.tool_name()) + .map(|executor| executor.tool_name()) .collect::>(); if config.code_mode_enabled { reserved_tool_names.insert(ToolName::plain(codex_code_mode::PUBLIC_TOOL_NAME)); @@ -453,9 +463,9 @@ fn append_extension_tool_handlers( } if config.search_tool && config.namespace_tools - && handlers + && registered_executors .iter() - .any(|handler| handler.exposure() == ToolExposure::Deferred) + .any(|executor| executor.exposure() == ToolExposure::Deferred) { reserved_tool_names.insert(ToolName::plain(TOOL_SEARCH_TOOL_NAME)); } @@ -466,7 +476,7 @@ fn append_extension_tool_handlers( warn!("Skipping extension tool `{tool_name}`: handler already registered"); continue; } - handlers.push(Arc::new(ExtensionToolHandler::new(executor))); + registered_executors.push(Arc::new(ExtensionToolHandler::new(executor))); } } diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index e40377e1d..dd690b563 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -28,6 +28,7 @@ use crate::tools::registry::ToolRegistry; use codex_app_server_protocol::AppInfo; use codex_extension_api::ExtensionToolExecutor; use codex_extension_api::ToolCall as ExtensionToolCall; +use codex_extension_api::ToolExecutor; use codex_features::Feature; use codex_features::Features; use codex_mcp::ToolInfo; @@ -79,7 +80,10 @@ fn extension_tool_executor(name: &str, description: &str) -> Arc for SpecOnlyExtensionExecutor { + type Output = codex_tools::JsonToolOutput; + fn tool_name(&self) -> ToolName { ToolName::plain(self.name.as_str()) } @@ -102,8 +106,11 @@ fn extension_tool_executor(name: &str, description: &str) -> Arc codex_extension_api::ExtensionToolFuture<'_> { - Box::pin(async { panic!("spec planning should not execute extension tools") }) + async fn handle( + &self, + _call: ExtensionToolCall, + ) -> Result { + panic!("spec planning should not execute extension tools") } } @@ -131,7 +138,7 @@ fn extension_tools_do_not_replace_builtin_tools() { "update_plan", "Extension attempt to replace a built-in tool.", )]; - let (tools, _) = build_specs_with_discoverable_tools( + let (tools, _) = build_specs_with_inputs_for_test( &tools_config, /*mcp_tools*/ None, /*deferred_mcp_tools*/ None, @@ -1882,7 +1889,7 @@ fn request_plugin_install_is_not_registered_without_feature_flag() { permission_profile: &PermissionProfile::Disabled, windows_sandbox_level: WindowsSandboxLevel::Disabled, }); - let (tools, _) = build_specs_with_discoverable_tools( + let (tools, _) = build_specs_with_inputs_for_test( &tools_config, /*mcp_tools*/ None, /*deferred_mcp_tools*/ None, @@ -1923,7 +1930,7 @@ fn request_plugin_install_can_be_registered_without_search_tool() { permission_profile: &PermissionProfile::Disabled, windows_sandbox_level: WindowsSandboxLevel::Disabled, }); - let (tools, _) = build_specs_with_discoverable_tools( + let (tools, _) = build_specs_with_inputs_for_test( &tools_config, /*mcp_tools*/ None, /*deferred_mcp_tools*/ None, @@ -1992,7 +1999,7 @@ fn request_plugin_install_description_lists_discoverable_tools() { })), ]; - let (tools, registry) = build_specs_with_discoverable_tools( + let (tools, registry) = build_specs_with_inputs_for_test( &tools_config, /*mcp_tools*/ None, /*deferred_mcp_tools*/ None, @@ -2293,7 +2300,7 @@ fn code_mode_only_exec_description_includes_extension_tool_details() { "extension_echo", "Echoes arguments through an extension tool.", )]; - let (tools, _) = build_specs_with_discoverable_tools( + let (tools, _) = build_specs_with_inputs_for_test( &tools_config, /*mcp_tools*/ None, /*deferred_mcp_tools*/ None, @@ -2391,7 +2398,7 @@ fn build_specs( deferred_mcp_tools: Option>, dynamic_tools: &[DynamicToolSpec], ) -> (Vec, ToolRegistry) { - build_specs_with_discoverable_tools( + build_specs_with_inputs_for_test( config, mcp_tools, deferred_mcp_tools, @@ -2401,7 +2408,7 @@ fn build_specs( ) } -fn build_specs_with_discoverable_tools( +fn build_specs_with_inputs_for_test( config: &ToolsConfig, mcp_tools: Option>, deferred_mcp_tools: Option>, @@ -2415,17 +2422,20 @@ fn build_specs_with_discoverable_tools( .map(|(name, tool)| tool_info_from_parts(name, tool.clone())) .collect::>() }); - let builder = build_tool_registry_builder( + let params = ToolRegistryBuildParams { + mcp_tools: mcp_tool_inputs.as_deref(), + deferred_mcp_tools: deferred_mcp_tools.as_deref(), + discoverable_tools: discoverable_tools.as_deref(), + extension_tool_executors, + dynamic_tools, + default_agent_type_description: DEFAULT_AGENT_TYPE_DESCRIPTION, + wait_agent_timeouts: wait_agent_timeout_options(), + }; + let executors = collect_tool_executors(config, params); + let builder = build_tool_registry_builder_from_executors( config, - ToolRegistryBuildParams { - mcp_tools: mcp_tool_inputs.as_deref(), - deferred_mcp_tools: deferred_mcp_tools.as_deref(), - discoverable_tools: discoverable_tools.as_deref(), - extension_tool_executors, - dynamic_tools, - default_agent_type_description: DEFAULT_AGENT_TYPE_DESCRIPTION, - wait_agent_timeouts: wait_agent_timeout_options(), - }, + executors, + hosted_model_tool_specs(config), ); builder.build() } diff --git a/codex-rs/core/src/tools/spec_tests.rs b/codex-rs/core/src/tools/spec_tests.rs index cf7386fc1..ca8739ccf 100644 --- a/codex-rs/core/src/tools/spec_tests.rs +++ b/codex-rs/core/src/tools/spec_tests.rs @@ -3,7 +3,9 @@ use crate::shell::Shell; use crate::shell::ShellType; use crate::test_support::construct_model_info_offline; use crate::tools::ToolRouter; +use crate::tools::registry::ToolRegistryBuilder; use crate::tools::router::ToolRouterParams; +use crate::tools::spec_plan::build_tool_registry_builder_from_executors; use codex_app_server_protocol::AppInfo; use codex_features::Feature; use codex_features::Features; @@ -36,6 +38,7 @@ use core_test_support::assert_regex_match; use pretty_assertions::assert_eq; use std::collections::BTreeMap; use std::path::PathBuf; +use std::sync::Arc; use super::*; @@ -270,7 +273,7 @@ fn build_specs( deferred_mcp_tools: Option>, dynamic_tools: &[DynamicToolSpec], ) -> ToolRegistryBuilder { - build_specs_with_discoverable_tools( + build_specs_with_inputs_for_test( config, mcp_tools, deferred_mcp_tools, @@ -280,6 +283,25 @@ fn build_specs( ) } +fn build_specs_with_inputs_for_test( + config: &ToolsConfig, + mcp_tools: Option>, + deferred_mcp_tools: Option>, + discoverable_tools: Option>, + extension_tool_executors: &[Arc], + dynamic_tools: &[DynamicToolSpec], +) -> ToolRegistryBuilder { + let parts = collect_tool_router_parts( + config, + mcp_tools, + deferred_mcp_tools, + discoverable_tools, + extension_tool_executors, + dynamic_tools, + ); + build_tool_registry_builder_from_executors(config, parts.executors, parts.hosted_specs) +} + #[tokio::test] async fn get_memory_requires_feature_flag() { let config = test_config().await; @@ -803,7 +825,7 @@ async fn request_plugin_install_requires_apps_and_plugins_features() { permission_profile: &PermissionProfile::Disabled, windows_sandbox_level: WindowsSandboxLevel::Disabled, }); - let (tools, _) = build_specs_with_discoverable_tools( + let (tools, _) = build_specs_with_inputs_for_test( &tools_config, /*mcp_tools*/ None, /*deferred_mcp_tools*/ None, 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 3badf1155..cb0eff9ca 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -30,6 +30,7 @@ struct TestHandler { tool_name: codex_tools::ToolName, } +#[async_trait::async_trait] impl ToolExecutor for TestHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index f3d1375ac..ed9665f22 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -18,7 +18,6 @@ pub use thread_lifecycle::ThreadResumeInput; pub use thread_lifecycle::ThreadStartInput; pub use thread_lifecycle::ThreadStopInput; pub use tools::ExtensionToolExecutor; -pub use tools::ExtensionToolFuture; pub use tools::ExtensionToolOutput; pub use turn_lifecycle::TurnAbortInput; pub use turn_lifecycle::TurnStartInput; diff --git a/codex-rs/ext/extension-api/src/contributors/tools.rs b/codex-rs/ext/extension-api/src/contributors/tools.rs index 41d08986e..a9f0f6dc3 100644 --- a/codex-rs/ext/extension-api/src/contributors/tools.rs +++ b/codex-rs/ext/extension-api/src/contributors/tools.rs @@ -1,32 +1,15 @@ -use std::future::Future; -use std::pin::Pin; - -use codex_tools::FunctionCallError; use codex_tools::JsonToolOutput; use codex_tools::ToolCall; -use codex_tools::ToolName; -use codex_tools::ToolSpec; +use codex_tools::ToolExecutor; /// Model-facing output returned by extension-owned tools. pub type ExtensionToolOutput = JsonToolOutput; -/// Future returned by extension-owned tool execution. -pub type ExtensionToolFuture<'a> = - Pin> + Send + 'a>>; - -/// Object-safe runtime contract for extension-owned model-visible tools. +/// Thin alias for extension-owned executable tools. /// -/// Implementations keep an extension tool's model-visible spec attached to the -/// executable runtime that handles calls for that tool. -pub trait ExtensionToolExecutor: Send + Sync { - /// The concrete tool name handled by this extension runtime. - fn tool_name(&self) -> ToolName; +/// Extensions implement the shared `ToolExecutor` contract directly; +/// the marker keeps contributor signatures readable while preserving one +/// executable-tool abstraction across host and extension tools. +pub trait ExtensionToolExecutor: ToolExecutor {} - /// The model-visible spec for this extension tool. - fn spec(&self) -> Option { - None - } - - /// Execute one extension tool invocation. - fn handle(&self, call: ToolCall) -> ExtensionToolFuture<'_>; -} +impl ExtensionToolExecutor for T where T: ToolExecutor {} diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index ccfa2520e..a6d850468 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -9,6 +9,7 @@ pub use codex_tools::FunctionCallError; pub use codex_tools::JsonToolOutput; pub use codex_tools::ResponsesApiTool; pub use codex_tools::ToolCall; +pub use codex_tools::ToolExecutor; pub use codex_tools::ToolName; pub use codex_tools::ToolPayload; pub use codex_tools::ToolSpec; @@ -18,7 +19,6 @@ pub use contributors::ApprovalReviewFuture; pub use contributors::ConfigContributor; pub use contributors::ContextContributor; pub use contributors::ExtensionToolExecutor; -pub use contributors::ExtensionToolFuture; pub use contributors::ExtensionToolOutput; pub use contributors::PromptFragment; pub use contributors::PromptSlot; diff --git a/codex-rs/ext/memories/Cargo.toml b/codex-rs/ext/memories/Cargo.toml index 557a25fea..5485eedc9 100644 --- a/codex-rs/ext/memories/Cargo.toml +++ b/codex-rs/ext/memories/Cargo.toml @@ -13,6 +13,7 @@ doctest = false workspace = true [dependencies] +async-trait = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-features = { workspace = true } diff --git a/codex-rs/ext/memories/src/tools/list.rs b/codex-rs/ext/memories/src/tools/list.rs index e2bc8304d..1509c2344 100644 --- a/codex-rs/ext/memories/src/tools/list.rs +++ b/codex-rs/ext/memories/src/tools/list.rs @@ -1,7 +1,6 @@ -use codex_extension_api::ExtensionToolExecutor; -use codex_extension_api::ExtensionToolFuture; use codex_extension_api::JsonToolOutput; use codex_extension_api::ToolCall; +use codex_extension_api::ToolExecutor; use codex_extension_api::ToolName; use codex_extension_api::ToolSpec; use schemars::JsonSchema; @@ -35,10 +34,13 @@ pub(super) struct ListTool { pub(super) backend: B, } -impl ExtensionToolExecutor for ListTool +#[async_trait::async_trait] +impl ToolExecutor for ListTool where B: MemoriesBackend, { + type Output = JsonToolOutput; + fn tool_name(&self) -> ToolName { memory_tool_name(LIST_TOOL_NAME) } @@ -50,23 +52,24 @@ where )) } - fn handle(&self, call: ToolCall) -> ExtensionToolFuture<'_> { + async fn handle( + &self, + call: ToolCall, + ) -> Result { let backend = self.backend.clone(); - Box::pin(async move { - let args: ListArgs = parse_args(&call)?; - let response = backend - .list(ListMemoriesRequest { - path: args.path, - cursor: args.cursor, - max_results: clamp_max_results( - args.max_results, - DEFAULT_LIST_MAX_RESULTS, - MAX_LIST_RESULTS, - ), - }) - .await - .map_err(backend_error_to_function_call)?; - Ok(JsonToolOutput::new(json!(response))) - }) + let args: ListArgs = parse_args(&call)?; + let response = backend + .list(ListMemoriesRequest { + path: args.path, + cursor: args.cursor, + max_results: clamp_max_results( + args.max_results, + DEFAULT_LIST_MAX_RESULTS, + MAX_LIST_RESULTS, + ), + }) + .await + .map_err(backend_error_to_function_call)?; + Ok(JsonToolOutput::new(json!(response))) } } diff --git a/codex-rs/ext/memories/src/tools/read.rs b/codex-rs/ext/memories/src/tools/read.rs index c706c609d..06d8a1752 100644 --- a/codex-rs/ext/memories/src/tools/read.rs +++ b/codex-rs/ext/memories/src/tools/read.rs @@ -1,7 +1,6 @@ -use codex_extension_api::ExtensionToolExecutor; -use codex_extension_api::ExtensionToolFuture; use codex_extension_api::JsonToolOutput; use codex_extension_api::ToolCall; +use codex_extension_api::ToolExecutor; use codex_extension_api::ToolName; use codex_extension_api::ToolSpec; use schemars::JsonSchema; @@ -34,10 +33,13 @@ pub(super) struct ReadTool { pub(super) backend: B, } -impl ExtensionToolExecutor for ReadTool +#[async_trait::async_trait] +impl ToolExecutor for ReadTool where B: MemoriesBackend, { + type Output = JsonToolOutput; + fn tool_name(&self) -> ToolName { memory_tool_name(READ_TOOL_NAME) } @@ -49,20 +51,21 @@ where )) } - fn handle(&self, call: ToolCall) -> ExtensionToolFuture<'_> { + async fn handle( + &self, + call: ToolCall, + ) -> Result { let backend = self.backend.clone(); - Box::pin(async move { - let args: ReadArgs = parse_args(&call)?; - let response = backend - .read(ReadMemoryRequest { - path: args.path, - line_offset: args.line_offset.unwrap_or(1), - max_lines: args.max_lines, - max_tokens: DEFAULT_READ_MAX_TOKENS, - }) - .await - .map_err(backend_error_to_function_call)?; - Ok(JsonToolOutput::new(json!(response))) - }) + let args: ReadArgs = parse_args(&call)?; + let response = backend + .read(ReadMemoryRequest { + path: args.path, + line_offset: args.line_offset.unwrap_or(1), + max_lines: args.max_lines, + max_tokens: DEFAULT_READ_MAX_TOKENS, + }) + .await + .map_err(backend_error_to_function_call)?; + Ok(JsonToolOutput::new(json!(response))) } } diff --git a/codex-rs/ext/memories/src/tools/search.rs b/codex-rs/ext/memories/src/tools/search.rs index 8042a8d45..f7cab7de6 100644 --- a/codex-rs/ext/memories/src/tools/search.rs +++ b/codex-rs/ext/memories/src/tools/search.rs @@ -1,7 +1,6 @@ -use codex_extension_api::ExtensionToolExecutor; -use codex_extension_api::ExtensionToolFuture; use codex_extension_api::JsonToolOutput; use codex_extension_api::ToolCall; +use codex_extension_api::ToolExecutor; use codex_extension_api::ToolName; use codex_extension_api::ToolSpec; use schemars::JsonSchema; @@ -43,10 +42,13 @@ pub(super) struct SearchTool { pub(super) backend: B, } -impl ExtensionToolExecutor for SearchTool +#[async_trait::async_trait] +impl ToolExecutor for SearchTool where B: MemoriesBackend, { + type Output = JsonToolOutput; + fn tool_name(&self) -> ToolName { memory_tool_name(SEARCH_TOOL_NAME) } @@ -58,16 +60,17 @@ where )) } - fn handle(&self, call: ToolCall) -> ExtensionToolFuture<'_> { + async fn handle( + &self, + call: ToolCall, + ) -> Result { let backend = self.backend.clone(); - Box::pin(async move { - let args: SearchArgs = parse_args(&call)?; - let response = backend - .search(args.into_request()) - .await - .map_err(backend_error_to_function_call)?; - Ok(JsonToolOutput::new(json!(response))) - }) + let args: SearchArgs = parse_args(&call)?; + let response = backend + .search(args.into_request()) + .await + .map_err(backend_error_to_function_call)?; + Ok(JsonToolOutput::new(json!(response))) } } diff --git a/codex-rs/tools/Cargo.toml b/codex-rs/tools/Cargo.toml index b11ee2fe7..e8d1134a6 100644 --- a/codex-rs/tools/Cargo.toml +++ b/codex-rs/tools/Cargo.toml @@ -8,6 +8,7 @@ version.workspace = true workspace = true [dependencies] +async-trait = { workspace = true } codex-app-server-protocol = { workspace = true } codex-code-mode = { workspace = true } codex-features = { workspace = true } diff --git a/codex-rs/tools/src/tool_executor.rs b/codex-rs/tools/src/tool_executor.rs index 4dc328523..62237740b 100644 --- a/codex-rs/tools/src/tool_executor.rs +++ b/codex-rs/tools/src/tool_executor.rs @@ -1,5 +1,3 @@ -use std::future::Future; - use crate::FunctionCallError; use crate::ToolName; use crate::ToolOutput; @@ -36,6 +34,7 @@ impl ToolExposure { /// Implementations keep the model-visible spec tied to the executable runtime. /// Host crates can layer routing, hooks, telemetry, or other orchestration on /// top without reopening the spec/runtime split. +#[async_trait::async_trait] pub trait ToolExecutor: Send + Sync { type Output: ToolOutput + 'static; @@ -54,8 +53,5 @@ pub trait ToolExecutor: Send + Sync { false } - fn handle( - &self, - invocation: Invocation, - ) -> impl Future> + Send; + async fn handle(&self, invocation: Invocation) -> Result; }