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<Output = T>`. 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<Result<AnyToolResult, _>>` 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`
2026-07-01 00:31:56 +08:00 · 2026-04-02 16:03:52 -07:00
parent 0bd31dc382
commit 3c7f013f97
33 changed files with 28 additions and 77 deletions
@@ -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;

@@ -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;

@@ -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;

@@ -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;

@@ -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;

@@ -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;

@@ -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;

@@ -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;

@@ -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;

@@ -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;
@@ -2,7 +2,6 @@ use super::*;

 pub(crate) struct Handler;

-#[async_trait]
 impl ToolHandler for Handler {
    type Output = CloseAgentResult;

@@ -4,7 +4,6 @@ use std::sync::Arc;

 pub(crate) struct Handler;

-#[async_trait]
 impl ToolHandler for Handler {
    type Output = ResumeAgentResult;

@@ -3,7 +3,6 @@ use crate::agent::control::render_input_preview;

 pub(crate) struct Handler;

-#[async_trait]
 impl ToolHandler for Handler {
    type Output = SendInputResult;

@@ -10,7 +10,6 @@ use crate::agent::next_thread_spawn_depth;

 pub(crate) struct Handler;

-#[async_trait]
 impl ToolHandler for Handler {
    type Output = SpawnAgentResult;

@@ -14,7 +14,6 @@ use tokio::time::timeout_at;

 pub(crate) struct Handler;

-#[async_trait]
 impl ToolHandler for Handler {
    type Output = WaitAgentResult;

@@ -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;
@@ -2,7 +2,6 @@ use super::*;

 pub(crate) struct Handler;

-#[async_trait]
 impl ToolHandler for Handler {
    type Output = CloseAgentResult;

@@ -6,7 +6,6 @@ use super::*;

 pub(crate) struct Handler;

-#[async_trait]
 impl ToolHandler for Handler {
    type Output = MessageToolResult;

@@ -3,7 +3,6 @@ use crate::agent::control::ListedAgent;

 pub(crate) struct Handler;

-#[async_trait]
 impl ToolHandler for Handler {
    type Output = ListAgentsResult;

@@ -6,7 +6,6 @@ use super::*;

 pub(crate) struct Handler;

-#[async_trait]
 impl ToolHandler for Handler {
    type Output = MessageToolResult;

@@ -11,7 +11,6 @@ use codex_protocol::protocol::Op;

 pub(crate) struct Handler;

-#[async_trait]
 impl ToolHandler for Handler {
    type Output = SpawnAgentResult;

@@ -6,7 +6,6 @@ use tokio::time::timeout_at;

 pub(crate) struct Handler;

-#[async_trait]
 impl ToolHandler for Handler {
    type Output = WaitAgentResult;

@@ -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;

@@ -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;

@@ -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;

@@ -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<ShellCommandBackendConfig> 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;

@@ -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<HashMap<String, BarrierState>> {
    BARRIERS.get_or_init(|| tokio::sync::Mutex::new(HashMap::new()))
 }

-#[async_trait]
 impl ToolHandler for TestSyncHandler {
    type Output = FunctionToolOutput;

@@ -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;

@@ -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;

@@ -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;

@@ -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;

@@ -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<Output = bool> + Send {
+        async { false }
    }

    fn pre_tool_use_payload(&self, _invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
@@ -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<Self::Output, FunctionCallError>;
+    fn handle(
+        &self,
+        invocation: ToolInvocation,
+    ) -> impl std::future::Future<Output = Result<Self::Output, FunctionCallError>> + 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<PreToolUsePayload>;

@@ -127,13 +131,12 @@ trait AnyToolHandler: Send + Sync {
        result: &dyn ToolOutput,
    ) -> Option<PostToolUsePayload>;

-    async fn handle_any(
-        &self,
+    fn handle_any<'a>(
+        &'a self,
        invocation: ToolInvocation,
-    ) -> Result<AnyToolResult, FunctionCallError>;
+    ) -> BoxFuture<'a, Result<AnyToolResult, FunctionCallError>>;
 }

-#[async_trait]
 impl<T> 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<PreToolUsePayload> {
@@ -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<AnyToolResult, FunctionCallError> {
-        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<AnyToolResult, FunctionCallError>> {
+        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),
+            })
        })
    }
 }
@@ -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;