core: use native async ToolHandler trait

2026-10-26 18:38:47 +00:00 · 2026-04-02 15:26:42 -07:00
parent 0bd31dc382
commit 41f7ac0ade
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;