From c24fea4f074ab37c52a53e850662d85336fd9b4e Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Fri, 22 May 2026 09:18:06 -0700 Subject: [PATCH] Make function tool hooks default --- .../core/src/tools/code_mode/wait_handler.rs | 15 ++++- .../handlers/unified_exec/write_stdin.rs | 5 +- codex-rs/core/src/tools/registry.rs | 58 +++++-------------- codex-rs/core/src/tools/registry_tests.rs | 15 +++++ 4 files changed, 46 insertions(+), 47 deletions(-) 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 7f3a923c64..1630249716 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -2,9 +2,12 @@ use serde::Deserialize; use crate::function_tool::FunctionCallError; use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; use crate::tools::context::boxed_tool_output; use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::PostToolUsePayload; +use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolExecutor; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -111,7 +114,15 @@ impl ToolExecutor for CodeModeWaitHandler { } impl CoreToolRuntime for CodeModeWaitHandler { - fn supports_default_function_tool_hooks(&self) -> bool { - false + fn pre_tool_use_payload(&self, _invocation: &ToolInvocation) -> Option { + None + } + + fn post_tool_use_payload( + &self, + _invocation: &ToolInvocation, + _result: &dyn ToolOutput, + ) -> Option { + None } } 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 2bd8f2b268..0491ace02d 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 @@ -5,6 +5,7 @@ use crate::tools::context::boxed_tool_output; use crate::tools::handlers::parse_arguments; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::PostToolUsePayload; +use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolExecutor; use crate::unified_exec::WriteStdinRequest; use codex_protocol::protocol::EventMsg; @@ -101,8 +102,8 @@ impl CoreToolRuntime for WriteStdinHandler { matches!(payload, ToolPayload::Function { .. }) } - fn supports_default_function_tool_hooks(&self) -> bool { - false + fn pre_tool_use_payload(&self, _invocation: &ToolInvocation) -> Option { + None } fn post_tool_use_payload( diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index b6143878c7..8443162a4f 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -68,10 +68,6 @@ pub(crate) trait CoreToolRuntime: ToolExecutor { invocation: &ToolInvocation, result: &dyn ToolOutput, ) -> Option { - if !self.supports_default_function_tool_hooks() { - return None; - } - let ToolPayload::Function { arguments } = &invocation.payload else { return None; }; @@ -102,10 +98,6 @@ pub(crate) trait CoreToolRuntime: ToolExecutor { } fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { - if !self.supports_default_function_tool_hooks() { - return None; - } - let ToolPayload::Function { arguments } = &invocation.payload else { return None; }; @@ -125,38 +117,22 @@ pub(crate) trait CoreToolRuntime: ToolExecutor { invocation: ToolInvocation, updated_input: Value, ) -> Result { - if self.supports_default_function_tool_hooks() { - let ToolPayload::Function { .. } = &invocation.payload else { - return Err(FunctionCallError::RespondToModel( - "hook input rewrite received unsupported function tool payload".to_string(), - )); - }; + let ToolPayload::Function { .. } = &invocation.payload else { + return Err(FunctionCallError::RespondToModel( + "hook input rewrite received unsupported function tool payload".to_string(), + )); + }; - let arguments = serde_json::to_string(&updated_input).map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to serialize rewritten {} arguments: {err}", - flat_tool_name(&invocation.tool_name) - )) - })?; - return Ok(ToolInvocation { - payload: ToolPayload::Function { arguments }, - ..invocation - }); - } - - Err(FunctionCallError::RespondToModel( - "tool does not support hook input rewriting".to_string(), - )) - } - - /// Returns whether this tool uses the generic function-tool hook contract. - /// - /// Most local function tools expose their JSON arguments directly to hooks. - /// Tools with compatibility-specific hook contracts can override the hook - /// payload methods instead, while function tools that should not run hooks - /// can opt out here. - fn supports_default_function_tool_hooks(&self) -> bool { - true + let arguments = serde_json::to_string(&updated_input).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to serialize rewritten {} arguments: {err}", + flat_tool_name(&invocation.tool_name) + )) + })?; + Ok(ToolInvocation { + payload: ToolPayload::Function { arguments }, + ..invocation + }) } /// Creates an optional consumer for streamed tool argument diffs. @@ -310,10 +286,6 @@ impl CoreToolRuntime for ExposureOverride { .with_updated_hook_input(invocation, updated_input) } - fn supports_default_function_tool_hooks(&self) -> bool { - self.handler.supports_default_function_tool_hooks() - } - fn telemetry_tags<'a>( &'a self, invocation: &'a ToolInvocation, diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index dbcfe480d4..a3099d284a 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -257,6 +257,21 @@ async fn code_mode_wait_does_not_expose_default_hook_payloads() { assert_eq!(wait.post_tool_use_payload(&wait_invocation, &output), None); } +#[tokio::test] +async fn write_stdin_does_not_expose_default_pre_tool_use_payload() { + let (session, turn) = crate::session::tests::make_session_and_context().await; + + let write_stdin = crate::tools::handlers::WriteStdinHandler; + let invocation = test_invocation( + Arc::new(session), + Arc::new(turn), + "write-stdin-call", + write_stdin.tool_name(), + ); + + assert_eq!(write_stdin.pre_tool_use_payload(&invocation), None); +} + #[test] fn model_visible_override_keeps_code_mode_result_typed() { let result = AnyToolResult {