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 e37fc62884..9ecc708371 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 eaa4150201..26ef021ad2 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 070b9270d5..4f828c8910 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 7797c8c61b..e08eb4b5fb 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 6ad0515e81..a384f2d8cc 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 5a709b287f..c8e2461e03 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 8000b2097c..29dc5199e3 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 cd94d88698..c4018fe936 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 34a4c9307c..f7575e1e52 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 3c0813329e..20691b8b5b 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 18b59c1b36..bb8ed10efd 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 5f700e3a98..f82f741346 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 0aea2b335e..71749e9cd1 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 b412f811bb..08b387376b 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 89d061d692..bd8172ac74 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 beac3a96f4..6459a899fe 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 1e9f5d9b07..9dab2d999d 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 ac212f7556..a6067e5b12 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 395e74f137..91f5990210 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 aa9d6f6581..d325edeeb0 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 c901df9ae0..7b575f8258 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 608d7f2f62..7d111f0418 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 527d5a29e7..8b0ee551e3 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 71909afb7e..584feec61f 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 1ab59cc3fe..3ad7b87145 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 5f246d6792..53c976c9d4 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 42f03e4e30..9868c77ca7 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 5d49ad861b..2f2e71677e 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 9e76352aa7..6073831e02 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 ae0a2c7694..7597cb6004 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 250f601ceb..54d74f13b7 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 091f0ffce7..feee6470a3 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 aaa3138588..579826ac04 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 e27c56576a..2df924b5f4 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 cc79615c14..29ee4b4ea3 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 1587e90ddb..7727fb4370 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 7e565bd312..730c3b7ef9 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 e27dd8fd31..b6c8eadadf 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 23f0207734..0f111922c6 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 921b52a86e..dd791d9135 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 92aeed0112..64dfdb8316 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 db530f3477..24f59a0f04 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 e40377e1d2..dd690b5631 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 cf7386fc1b..ca8739ccff 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 3badf1155a..cb0eff9ca9 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 f3d1375ac8..ed9665f228 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 41d08986e9..a9f0f6dc32 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 ccfa2520e4..a6d850468b 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 557a25fea8..5485eedc9f 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 e2bc8304da..1509c2344c 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 c706c609dc..06d8a1752c 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 8042a8d45e..f7cab7de6c 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 b11ee2fe7c..e8d1134a6d 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 4dc328523f..62237740ba 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; }