From 3c7f013f9735e67796c70d95f75f436b7f97e3ec Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 2 Apr 2026 16:03:52 -0700 Subject: [PATCH] core: cut codex-core compile time 63% with native async ToolHandler (#16630) ## Why `ToolHandler` was still paying a large compile-time tax from `#[async_trait]` on every concrete handler impl, even though the only object-safe boundary the registry actually stores is the internal `AnyToolHandler` adapter. This PR removes that macro-generated async wrapper layer from concrete `ToolHandler` impls while keeping the existing object-safe shim in `AnyToolHandler`. In practice, that gets essentially the same compile-time win as the larger type-erasure refactor in #16627, but with a much smaller diff and without changing the public shape of `ToolHandler`. That tradeoff matters here because this is a broad `codex-core` hotspot and reviewers should be able to judge the compile-time impact from hard numbers, not vibes. ## Headline result On a clean `codex-core` package rebuild (`cargo clean -p codex-core` before each command), rustc `total` dropped from **187.15s to 68.98s** versus the shared `0bd31dc382bd` baseline: **-63.1%**. The biggest hot passes dropped by roughly **71-72%**: | Metric | Baseline `0bd31dc382bd` | This PR `41f7ac0adeac` | Delta | |---|---:|---:|---:| | `total` | 187.15s | 68.98s | **-63.1%** | | `generate_crate_metadata` | 84.53s | 24.49s | **-71.0%** | | `MIR_borrow_checking` | 84.13s | 24.58s | **-70.8%** | | `monomorphization_collector_graph_walk` | 79.74s | 22.19s | **-72.2%** | | `evaluate_obligation` self-time | 180.62s | 46.91s | **-74.0%** | Important caveat: `-Z time-passes` timings are nested, so `generate_crate_metadata` and `monomorphization_collector_graph_walk` are mostly overlapping, not additive. ## Why this PR over #16627 #16627 already proved that the `ToolHandler` stack was the right hotspot, but it got there by making `ToolHandler` object-safe and changing every handler to return `BoxFuture>` directly. This PR keeps the lower-churn shape: - `ToolHandler` remains generic over `type Output`. - Concrete handlers use native RPITIT futures with explicit `Send` bounds. - `AnyToolHandler` remains the only object-safe adapter and still does the boxing at the registry boundary, as before. - The implementation diff is only **33 files, +28/-77**. The measurements are at least comparable, and in this run this PR is slightly faster than #16627 on the pass-level total: | Metric | #16627 | This PR | Delta | |---|---:|---:|---:| | `total` | 79.90s | 68.98s | **-13.7%** | | `generate_crate_metadata` | 25.88s | 24.49s | **-5.4%** | | `monomorphization_collector_graph_walk` | 23.54s | 22.19s | **-5.7%** | | `evaluate_obligation` self-time | 43.29s | 46.91s | +8.4% | ## Profile data ### Crate-level timings `cargo +nightly build -p codex-core --lib -Z unstable-options --timings=json` after `cargo clean -p codex-core`. Baseline data below is reused from the shared parent `0bd31dc382bd` profile because this PR and #16627 are both one commit on top of that same parent. | Crate | Baseline `duration` | This PR `duration` | Delta | Baseline `rmeta_time` | This PR `rmeta_time` | Delta | |---|---:|---:|---:|---:|---:|---:| | `codex_core` | 187.380776583s | 69.171113833s | **-63.1%** | 174.474507208s | 55.873015583s | **-68.0%** | | `starlark` | 17.90s | 16.773824125s | -6.3% | n/a | 8.8999965s | n/a | ### Pass-level timings `cargo +nightly rustc -p codex-core --lib -- -Z time-passes -Z time-passes-format=json` after `cargo clean -p codex-core`. | Pass | Baseline | This PR | Delta | |---|---:|---:|---:| | `total` | 187.150662083s | 68.978770375s | **-63.1%** | | `generate_crate_metadata` | 84.531864625s | 24.487462958s | **-71.0%** | | `MIR_borrow_checking` | 84.131389375s | 24.575553875s | **-70.8%** | | `monomorphization_collector_graph_walk` | 79.737515042s | 22.190207417s | **-72.2%** | | `codegen_crate` | 12.362532292s | 12.695237625s | +2.7% | | `type_check_crate` | 4.4765405s | 5.442019542s | +21.6% | | `coherence_checking` | 3.311121208s | 4.239935292s | +28.0% | | process `real` / `user` / `sys` | 187.70s / 201.87s / 4.99s | 69.52s / 85.90s / 2.92s | n/a | ### Self-profile query summary `cargo +nightly rustc -p codex-core --lib -- -Z self-profile=... -Z self-profile-events=default,query-keys,args,llvm,artifact-sizes` after `cargo clean -p codex-core`, summarized with `measureme summarize -p 0.5`. | Query / phase | Baseline self time | This PR self time | Delta | Baseline total time | This PR total time | Baseline item count | This PR item count | Baseline cache hits | This PR cache hits | |---|---:|---:|---:|---:|---:|---:|---:|---:|---:| | `evaluate_obligation` | 180.62s | 46.91s | **-74.0%** | 182.08s | 48.37s | 572,234 | 388,659 | 1,130,998 | 1,058,553 | | `mir_borrowck` | 1.42s | 1.49s | +4.9% | 93.77s | 29.59s | n/a | 6,184 | n/a | 15,298 | | `typeck` | 1.84s | 1.87s | +1.6% | 2.38s | 2.44s | n/a | 9,367 | n/a | 79,247 | | `LLVM_module_codegen_emit_obj` | n/a | 17.12s | n/a | 17.01s | 17.12s | n/a | 256 | n/a | 0 | | `LLVM_passes` | n/a | 13.07s | n/a | 12.95s | 13.07s | n/a | 1 | n/a | 0 | | `codegen_module` | n/a | 12.33s | n/a | 12.22s | 13.64s | n/a | 256 | n/a | 0 | | `items_of_instance` | n/a | 676.00ms | n/a | n/a | 24.96s | n/a | 99,990 | n/a | 0 | | `type_op_prove_predicate` | n/a | 660.79ms | n/a | n/a | 24.78s | n/a | 78,762 | n/a | 235,877 | | Summary | Baseline | This PR | |---|---:|---:| | `evaluate_obligation` % of total CPU | 70.821% | 38.880% | | self-profile total CPU time | 255.042999997s | 120.661175956s | | process `real` / `user` / `sys` | 220.96s / 235.02s / 7.09s | 86.35s / 103.66s / 3.54s | ### Artifact sizes From the same `measureme summarize` output: | Artifact | Baseline | This PR | Delta | |---|---:|---:|---:| | `crate_metadata` | 26,534,471 bytes | 26,545,248 bytes | +10,777 | | `dep_graph` | 253,181,425 bytes | 239,240,806 bytes | -13,940,619 | | `linked_artifact` | 565,366,624 bytes | 562,673,176 bytes | -2,693,448 | | `object_file` | 513,127,264 bytes | 510,464,096 bytes | -2,663,168 | | `query_cache` | 137,440,945 bytes | 136,982,566 bytes | -458,379 | | `cgu_instructions` | 3,586,307 bytes | 3,575,121 bytes | -11,186 | | `codegen_unit_size_estimate` | 2,084,846 bytes | 2,078,773 bytes | -6,073 | | `work_product_index` | 19,565 bytes | 19,565 bytes | 0 | ### Baseline hotspots before this change These are the top normalized obligation buckets from the shared baseline profile: | Obligation bucket | Samples | Duration | |---|---:|---:| | `outlives:tasks::review::ReviewTask` | 1,067 | 6.33s | | `outlives:tools::handlers::unified_exec::UnifiedExecHandler` | 896 | 5.63s | | `trait:T as tools::registry::ToolHandler` | 876 | 5.45s | | `outlives:tools::handlers::shell::ShellHandler` | 888 | 5.37s | | `outlives:tools::handlers::shell::ShellCommandHandler` | 870 | 5.29s | | `outlives:tools::runtimes::shell::unix_escalation::CoreShellActionProvider` | 637 | 3.73s | | `outlives:tools::handlers::mcp::McpHandler` | 695 | 3.61s | | `outlives:tasks::regular::RegularTask` | 726 | 3.57s | Top `items_of_instance` entries before this change were mostly concrete async handler/task impls: | Instance | Duration | |---|---:| | `tasks::regular::{impl#2}::run` | 3.79s | | `tools::handlers::mcp::{impl#0}::handle` | 3.27s | | `tools::runtimes::shell::unix_escalation::{impl#2}::determine_action` | 3.09s | | `tools::handlers::agent_jobs::{impl#11}::handle` | 3.07s | | `tools::handlers::multi_agents::spawn::{impl#1}::handle` | 2.84s | | `tasks::review::{impl#4}::run` | 2.82s | | `tools::handlers::multi_agents_v2::spawn::{impl#2}::handle` | 2.80s | | `tools::handlers::multi_agents::resume_agent::{impl#1}::handle` | 2.73s | | `tools::handlers::unified_exec::{impl#2}::handle` | 2.54s | | `tasks::compact::{impl#4}::run` | 2.45s | ## What changed Relevant pre-change registry shape: [`codex-rs/core/src/tools/registry.rs`](https://github.com/openai/codex/blob/0bd31dc382bd1c33dc2bb6b97069c76aa10ba14b/codex-rs/core/src/tools/registry.rs#L38-L219) Current registry shape in this PR: [`codex-rs/core/src/tools/registry.rs`](https://github.com/openai/codex/blob/41f7ac0adeac81d667541853d6546267d6083613/codex-rs/core/src/tools/registry.rs#L38-L203) - `ToolHandler::{is_mutating, handle}` now return native `impl Future + Send` futures instead of using `#[async_trait]`. - `AnyToolHandler` remains the object-safe adapter and boxes those futures at the registry boundary with explicit lifetimes. - Concrete handlers and the registry test handler drop `#[async_trait]` but otherwise keep their async method bodies intact. - Representative examples: [`codex-rs/core/src/tools/handlers/shell.rs`](https://github.com/openai/codex/blob/41f7ac0adeac81d667541853d6546267d6083613/codex-rs/core/src/tools/handlers/shell.rs#L223-L379), [`codex-rs/core/src/tools/handlers/unified_exec.rs`](https://github.com/openai/codex/blob/41f7ac0adeac81d667541853d6546267d6083613/codex-rs/core/src/tools/handlers/unified_exec.rs), [`codex-rs/core/src/tools/registry_tests.rs`](https://github.com/openai/codex/blob/41f7ac0adeac81d667541853d6546267d6083613/codex-rs/core/src/tools/registry_tests.rs) ## Tradeoff This is intentionally less invasive than #16627: it does **not** move result boxing into every concrete handler and does **not** change `ToolHandler` into an object-safe trait. Instead, it keeps the existing registry-level type-erasure boundary and only removes the macro-generated async wrapper layer from concrete impls. So the runtime boxing story stays basically the same as before, while the compile-time savings are still large. ## Verification Existing verification for this branch still applies: - Ran `cargo test -p codex-core`; this change compiled and the suite reached the known unrelated `config::tests::*guardian*` failures, with no local diff under `codex-rs/core/src/config/`. Profiling commands used for the tables above: - `cargo clean -p codex-core` - `cargo +nightly build -p codex-core --lib -Z unstable-options --timings=json` - `cargo +nightly rustc -p codex-core --lib -- -Z time-passes -Z time-passes-format=json` - `cargo +nightly rustc -p codex-core --lib -- -Z self-profile=... -Z self-profile-events=default,query-keys,args,llvm,artifact-sizes` - `measureme summarize -p 0.5` --- .../src/tools/code_mode/execute_handler.rs | 3 -- .../core/src/tools/code_mode/wait_handler.rs | 2 - .../core/src/tools/handlers/agent_jobs.rs | 2 - .../core/src/tools/handlers/apply_patch.rs | 2 - codex-rs/core/src/tools/handlers/dynamic.rs | 2 - codex-rs/core/src/tools/handlers/js_repl.rs | 3 -- codex-rs/core/src/tools/handlers/list_dir.rs | 2 - codex-rs/core/src/tools/handlers/mcp.rs | 2 - .../core/src/tools/handlers/mcp_resource.rs | 2 - .../core/src/tools/handlers/multi_agents.rs | 1 - .../handlers/multi_agents/close_agent.rs | 1 - .../handlers/multi_agents/resume_agent.rs | 1 - .../tools/handlers/multi_agents/send_input.rs | 1 - .../src/tools/handlers/multi_agents/spawn.rs | 1 - .../src/tools/handlers/multi_agents/wait.rs | 1 - .../src/tools/handlers/multi_agents_v2.rs | 1 - .../handlers/multi_agents_v2/close_agent.rs | 1 - .../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 | 1 - .../tools/handlers/multi_agents_v2/wait.rs | 1 - codex-rs/core/src/tools/handlers/plan.rs | 2 - .../src/tools/handlers/request_permissions.rs | 2 - .../src/tools/handlers/request_user_input.rs | 2 - codex-rs/core/src/tools/handlers/shell.rs | 3 -- codex-rs/core/src/tools/handlers/test_sync.rs | 2 - .../core/src/tools/handlers/tool_search.rs | 2 - .../core/src/tools/handlers/tool_suggest.rs | 2 - .../core/src/tools/handlers/unified_exec.rs | 2 - .../core/src/tools/handlers/view_image.rs | 2 - codex-rs/core/src/tools/registry.rs | 51 ++++++++++--------- codex-rs/core/src/tools/registry_tests.rs | 2 - 33 files changed, 28 insertions(+), 77 deletions(-) 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 3f77216c1..c46c3e045 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -1,5 +1,3 @@ -use async_trait::async_trait; - use crate::function_tool::FunctionCallError; use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; @@ -53,7 +51,6 @@ impl CodeModeExecuteHandler { } } -#[async_trait] impl ToolHandler 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 f319985a8..8a9d24c1f 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -1,4 +1,3 @@ -use async_trait::async_trait; use serde::Deserialize; use crate::function_tool::FunctionCallError; @@ -39,7 +38,6 @@ where }) } -#[async_trait] impl ToolHandler for CodeModeWaitHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index c607730b8..044e2072f 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -13,7 +13,6 @@ use crate::tools::handlers::multi_agents::build_agent_spawn_config; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; -use async_trait::async_trait; use codex_protocol::ThreadId; use codex_protocol::protocol::AgentStatus; use codex_protocol::protocol::SessionSource; @@ -178,7 +177,6 @@ impl JobProgressEmitter { } } -#[async_trait] impl ToolHandler for BatchJobHandler { 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 801cff56f..00b1173aa 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -21,7 +21,6 @@ use crate::tools::registry::ToolKind; use crate::tools::runtimes::apply_patch::ApplyPatchRequest; use crate::tools::runtimes::apply_patch::ApplyPatchRuntime; use crate::tools::sandboxing::ToolCtx; -use async_trait::async_trait; use codex_apply_patch::ApplyPatchAction; use codex_apply_patch::ApplyPatchFileChange; use codex_protocol::models::FileSystemPermissions; @@ -122,7 +121,6 @@ async fn effective_patch_permissions( ) } -#[async_trait] impl ToolHandler 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 23285bc90..5f73f096e 100644 --- a/codex-rs/core/src/tools/handlers/dynamic.rs +++ b/codex-rs/core/src/tools/handlers/dynamic.rs @@ -7,7 +7,6 @@ use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; -use async_trait::async_trait; use codex_protocol::dynamic_tools::DynamicToolCallRequest; use codex_protocol::dynamic_tools::DynamicToolResponse; use codex_protocol::models::FunctionCallOutputContentItem; @@ -20,7 +19,6 @@ use tracing::warn; pub struct DynamicToolHandler; -#[async_trait] impl ToolHandler for DynamicToolHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/js_repl.rs b/codex-rs/core/src/tools/handlers/js_repl.rs index c5906abaf..cf8d0e862 100644 --- a/codex-rs/core/src/tools/handlers/js_repl.rs +++ b/codex-rs/core/src/tools/handlers/js_repl.rs @@ -1,4 +1,3 @@ -use async_trait::async_trait; use serde_json::Value as JsonValue; use std::sync::Arc; use std::time::Duration; @@ -92,7 +91,6 @@ async fn emit_js_repl_exec_end( }; emitter.emit(ctx, stage).await; } -#[async_trait] impl ToolHandler for JsReplHandler { type Output = FunctionToolOutput; @@ -182,7 +180,6 @@ impl ToolHandler for JsReplHandler { } } -#[async_trait] impl ToolHandler for JsReplResetHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/list_dir.rs b/codex-rs/core/src/tools/handlers/list_dir.rs index fd461e82e..be4e356c5 100644 --- a/codex-rs/core/src/tools/handlers/list_dir.rs +++ b/codex-rs/core/src/tools/handlers/list_dir.rs @@ -4,7 +4,6 @@ use std::fs::FileType; use std::path::Path; use std::path::PathBuf; -use async_trait::async_trait; use codex_utils_string::take_bytes_at_char_boundary; use serde::Deserialize; use tokio::fs; @@ -45,7 +44,6 @@ struct ListDirArgs { depth: usize, } -#[async_trait] impl ToolHandler for ListDirHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 18e0df25c..6eaccaf81 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -1,4 +1,3 @@ -use async_trait::async_trait; use std::sync::Arc; use crate::function_tool::FunctionCallError; @@ -10,7 +9,6 @@ use crate::tools::registry::ToolKind; use codex_protocol::mcp::CallToolResult; pub struct McpHandler; -#[async_trait] impl ToolHandler for McpHandler { type Output = CallToolResult; diff --git a/codex-rs/core/src/tools/handlers/mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource.rs index 158aee612..d0b7ddc5a 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource.rs @@ -3,7 +3,6 @@ use std::sync::Arc; use std::time::Duration; use std::time::Instant; -use async_trait::async_trait; use codex_protocol::mcp::CallToolResult; use codex_protocol::models::function_call_output_content_items_to_text; use rmcp::model::ListResourceTemplatesResult; @@ -178,7 +177,6 @@ struct ReadResourcePayload { result: ReadResourceResult, } -#[async_trait] impl ToolHandler for McpResourceHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/multi_agents.rs b/codex-rs/core/src/tools/handlers/multi_agents.rs index 166dbd287..df66543da 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents.rs @@ -17,7 +17,6 @@ pub(crate) use crate::tools::handlers::multi_agents_common::*; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; -use async_trait::async_trait; use codex_protocol::ThreadId; use codex_protocol::models::ResponseInputItem; use codex_protocol::openai_models::ReasoningEffort; 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 efc8ec378..779a85e19 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 @@ -2,7 +2,6 @@ use super::*; pub(crate) struct Handler; -#[async_trait] impl ToolHandler for Handler { type Output = CloseAgentResult; 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 09526182f..9991a9067 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 @@ -4,7 +4,6 @@ use std::sync::Arc; pub(crate) struct Handler; -#[async_trait] impl ToolHandler for Handler { type Output = ResumeAgentResult; 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 3c8527712..4ae3240cb 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 @@ -3,7 +3,6 @@ use crate::agent::control::render_input_preview; pub(crate) struct Handler; -#[async_trait] impl ToolHandler 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 308ec49d8..8e4bfb5b5 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -10,7 +10,6 @@ use crate::agent::next_thread_spawn_depth; pub(crate) struct Handler; -#[async_trait] impl ToolHandler for Handler { type Output = SpawnAgentResult; 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 2fe33f1ed..8102b9732 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/wait.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/wait.rs @@ -14,7 +14,6 @@ use tokio::time::timeout_at; pub(crate) struct Handler; -#[async_trait] impl ToolHandler for Handler { type Output = WaitAgentResult; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2.rs index 44191ed85..51b80e4a7 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2.rs @@ -11,7 +11,6 @@ use crate::tools::handlers::multi_agents_common::*; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; -use async_trait::async_trait; use codex_protocol::AgentPath; use codex_protocol::models::ResponseInputItem; use codex_protocol::openai_models::ReasoningEffort; 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 2296ae5dd..8074f7fe0 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 @@ -2,7 +2,6 @@ use super::*; pub(crate) struct Handler; -#[async_trait] impl ToolHandler for Handler { type Output = CloseAgentResult; 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 d0f87f342..312e17a31 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 @@ -6,7 +6,6 @@ use super::*; pub(crate) struct Handler; -#[async_trait] impl ToolHandler for Handler { type Output = MessageToolResult; 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 b3b3f7160..579c44199 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 @@ -3,7 +3,6 @@ use crate::agent::control::ListedAgent; pub(crate) struct Handler; -#[async_trait] impl ToolHandler 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 a16aebd4c..a490d1100 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 @@ -6,7 +6,6 @@ use super::*; pub(crate) struct Handler; -#[async_trait] impl ToolHandler for Handler { type Output = MessageToolResult; 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 3dc8d7202..9b3b81851 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 @@ -11,7 +11,6 @@ use codex_protocol::protocol::Op; pub(crate) struct Handler; -#[async_trait] impl ToolHandler for Handler { type Output = SpawnAgentResult; 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 bf3ad8e93..777fe1c2e 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 @@ -6,7 +6,6 @@ use tokio::time::timeout_at; pub(crate) struct Handler; -#[async_trait] impl ToolHandler 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 ce0b98b1c..fb1c708ad 100644 --- a/codex-rs/core/src/tools/handlers/plan.rs +++ b/codex-rs/core/src/tools/handlers/plan.rs @@ -6,7 +6,6 @@ use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; -use async_trait::async_trait; use codex_protocol::config_types::ModeKind; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; @@ -44,7 +43,6 @@ impl ToolOutput for PlanToolOutput { } } -#[async_trait] impl ToolHandler 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 e0755deac..2590b9d7d 100644 --- a/codex-rs/core/src/tools/handlers/request_permissions.rs +++ b/codex-rs/core/src/tools/handlers/request_permissions.rs @@ -1,4 +1,3 @@ -use async_trait::async_trait; use codex_protocol::request_permissions::RequestPermissionsArgs; use codex_sandboxing::policy_transforms::normalize_additional_permissions; @@ -12,7 +11,6 @@ use crate::tools::registry::ToolKind; pub struct RequestPermissionsHandler; -#[async_trait] impl ToolHandler for RequestPermissionsHandler { 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 f12defd0e..dcc9445a9 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input.rs @@ -5,7 +5,6 @@ use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; -use async_trait::async_trait; use codex_protocol::request_user_input::RequestUserInputArgs; use codex_tools::REQUEST_USER_INPUT_TOOL_NAME; use codex_tools::normalize_request_user_input_args; @@ -15,7 +14,6 @@ pub struct RequestUserInputHandler { pub default_mode_request_user_input: bool, } -#[async_trait] impl ToolHandler for RequestUserInputHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 34cdd5b97..6f662ff49 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -1,4 +1,3 @@ -use async_trait::async_trait; use codex_protocol::ThreadId; use codex_protocol::models::ShellCommandToolCallParams; use codex_protocol::models::ShellToolCallParams; @@ -178,7 +177,6 @@ impl From for ShellCommandHandler { } } -#[async_trait] impl ToolHandler for ShellHandler { type Output = FunctionToolOutput; @@ -279,7 +277,6 @@ impl ToolHandler for ShellHandler { } } -#[async_trait] impl ToolHandler 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 2d6e35148..ad2647243 100644 --- a/codex-rs/core/src/tools/handlers/test_sync.rs +++ b/codex-rs/core/src/tools/handlers/test_sync.rs @@ -4,7 +4,6 @@ use std::sync::Arc; use std::sync::OnceLock; use std::time::Duration; -use async_trait::async_trait; use serde::Deserialize; use tokio::sync::Barrier; use tokio::time::sleep; @@ -54,7 +53,6 @@ fn barrier_map() -> &'static tokio::sync::Mutex> { BARRIERS.get_or_init(|| tokio::sync::Mutex::new(HashMap::new())) } -#[async_trait] impl ToolHandler 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 aa9d8e92f..dd203e93b 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -4,7 +4,6 @@ use crate::tools::context::ToolPayload; use crate::tools::context::ToolSearchOutput; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; -use async_trait::async_trait; use bm25::Document; use bm25::Language; use bm25::SearchEngineBuilder; @@ -25,7 +24,6 @@ impl ToolSearchHandler { } } -#[async_trait] impl ToolHandler for ToolSearchHandler { type Output = ToolSearchOutput; diff --git a/codex-rs/core/src/tools/handlers/tool_suggest.rs b/codex-rs/core/src/tools/handlers/tool_suggest.rs index 45c77a350..380aebd51 100644 --- a/codex-rs/core/src/tools/handlers/tool_suggest.rs +++ b/codex-rs/core/src/tools/handlers/tool_suggest.rs @@ -1,6 +1,5 @@ use std::collections::HashSet; -use async_trait::async_trait; use codex_app_server_protocol::AppInfo; use codex_mcp::mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_rmcp_client::ElicitationAction; @@ -28,7 +27,6 @@ use crate::tools::registry::ToolKind; pub struct ToolSuggestHandler; -#[async_trait] impl ToolHandler for ToolSuggestHandler { type Output = FunctionToolOutput; diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index 98d7d717f..dfe8c4275 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -22,7 +22,6 @@ use crate::unified_exec::ExecCommandRequest; use crate::unified_exec::UnifiedExecContext; use crate::unified_exec::UnifiedExecProcessManager; use crate::unified_exec::WriteStdinRequest; -use async_trait::async_trait; use codex_features::Feature; use codex_otel::SessionTelemetry; use codex_otel::metrics::names::TOOL_CALL_UNIFIED_EXEC_METRIC; @@ -86,7 +85,6 @@ fn default_tty() -> bool { false } -#[async_trait] impl ToolHandler for UnifiedExecHandler { 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 e87638bba..ec34fafb3 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -1,4 +1,3 @@ -use async_trait::async_trait; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; @@ -37,7 +36,6 @@ enum ViewImageDetail { Original, } -#[async_trait] impl ToolHandler for ViewImageHandler { type Output = ViewImageOutput; diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index f1ebb38c0..2723d8a9a 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -13,7 +13,6 @@ use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; -use async_trait::async_trait; use codex_hooks::HookEvent; use codex_hooks::HookEventAfterToolUse; use codex_hooks::HookPayload; @@ -26,6 +25,7 @@ use codex_protocol::protocol::SandboxPolicy; use codex_tools::ConfiguredToolSpec; use codex_tools::ToolSpec; use codex_utils_readiness::Readiness; +use futures::future::BoxFuture; use serde_json::Value; use tracing::warn; @@ -35,7 +35,6 @@ pub enum ToolKind { Mcp, } -#[async_trait] pub trait ToolHandler: Send + Sync { type Output: ToolOutput + 'static; @@ -54,8 +53,11 @@ pub trait ToolHandler: Send + Sync { /// user (through file system, OS operations, ...). /// This function must remains defensive and return `true` if a doubt exist on the /// exact effect of a ToolInvocation. - async fn is_mutating(&self, _invocation: &ToolInvocation) -> bool { - false + fn is_mutating( + &self, + _invocation: &ToolInvocation, + ) -> impl std::future::Future + Send { + async { false } } fn pre_tool_use_payload(&self, _invocation: &ToolInvocation) -> Option { @@ -73,7 +75,10 @@ pub trait ToolHandler: Send + Sync { /// Perform the actual [ToolInvocation] and returns a [ToolOutput] containing /// the final output to return to the model. - async fn handle(&self, invocation: ToolInvocation) -> Result; + fn handle( + &self, + invocation: ToolInvocation, + ) -> impl std::future::Future> + Send; } pub(crate) struct AnyToolResult { @@ -112,11 +117,10 @@ pub(crate) struct PostToolUsePayload { pub(crate) tool_response: Value, } -#[async_trait] trait AnyToolHandler: Send + Sync { fn matches_kind(&self, payload: &ToolPayload) -> bool; - async fn is_mutating(&self, invocation: &ToolInvocation) -> bool; + fn is_mutating<'a>(&'a self, invocation: &'a ToolInvocation) -> BoxFuture<'a, bool>; fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option; @@ -127,13 +131,12 @@ trait AnyToolHandler: Send + Sync { result: &dyn ToolOutput, ) -> Option; - async fn handle_any( - &self, + fn handle_any<'a>( + &'a self, invocation: ToolInvocation, - ) -> Result; + ) -> BoxFuture<'a, Result>; } -#[async_trait] impl AnyToolHandler for T where T: ToolHandler, @@ -142,8 +145,8 @@ where ToolHandler::matches_kind(self, payload) } - async fn is_mutating(&self, invocation: &ToolInvocation) -> bool { - ToolHandler::is_mutating(self, invocation).await + fn is_mutating<'a>(&'a self, invocation: &'a ToolInvocation) -> BoxFuture<'a, bool> { + Box::pin(ToolHandler::is_mutating(self, invocation)) } fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { @@ -159,17 +162,19 @@ where ToolHandler::post_tool_use_payload(self, call_id, payload, result) } - async fn handle_any( - &self, + fn handle_any<'a>( + &'a self, invocation: ToolInvocation, - ) -> Result { - let call_id = invocation.call_id.clone(); - let payload = invocation.payload.clone(); - let output = self.handle(invocation).await?; - Ok(AnyToolResult { - call_id, - payload, - result: Box::new(output), + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let call_id = invocation.call_id.clone(); + let payload = invocation.payload.clone(); + let output = self.handle(invocation).await?; + Ok(AnyToolResult { + call_id, + payload, + result: Box::new(output), + }) }) } } diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index 687364872..287996d7d 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -1,10 +1,8 @@ use super::*; -use async_trait::async_trait; use pretty_assertions::assert_eq; struct TestHandler; -#[async_trait] impl ToolHandler for TestHandler { type Output = crate::tools::context::FunctionToolOutput;