From 1cd3e86056d062af01a1157fcd8ddb17f84d19a7 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Wed, 6 May 2026 14:16:38 -0700 Subject: [PATCH] Centralize pre-hook tool compatibility --- .../core/src/tools/handlers/apply_patch.rs | 44 +-- .../src/tools/handlers/apply_patch_tests.rs | 8 +- codex-rs/core/src/tools/handlers/mcp.rs | 58 +--- codex-rs/core/src/tools/handlers/mod.rs | 47 +-- codex-rs/core/src/tools/handlers/shell.rs | 142 +-------- .../core/src/tools/handlers/shell_tests.rs | 75 ++--- .../core/src/tools/handlers/unified_exec.rs | 41 +-- .../src/tools/handlers/unified_exec_tests.rs | 8 +- codex-rs/core/src/tools/hook_compat.rs | 276 ++++++++++++++++++ codex-rs/core/src/tools/mod.rs | 1 + codex-rs/core/src/tools/registry.rs | 43 +-- 11 files changed, 330 insertions(+), 413 deletions(-) create mode 100644 codex-rs/core/src/tools/hook_compat.rs diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 1dc06780ae..cccc517ba3 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -22,12 +22,9 @@ use crate::tools::events::ToolEmitter; use crate::tools::events::ToolEventCtx; use crate::tools::handlers::apply_granted_turn_permissions; use crate::tools::handlers::parse_arguments; -use crate::tools::handlers::rewrite_function_string_argument; -use crate::tools::handlers::updated_hook_command; use crate::tools::hook_names::HookToolName; use crate::tools::orchestrator::ToolOrchestrator; use crate::tools::registry::PostToolUsePayload; -use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; @@ -243,12 +240,8 @@ fn write_permissions_for_paths( normalize_additional_permissions(permissions).ok() } -/// Extracts the raw patch text used as the command-shaped hook input for apply_patch. -/// -/// The apply_patch tool can arrive as the older JSON/function shape or as a -/// freeform custom tool call. Both represent the same file edit operation, so -/// hooks see the raw patch body in `tool_input.command` either way. -fn apply_patch_payload_command(payload: &ToolPayload) -> Option { +/// Extracts the raw patch text from either supported apply_patch payload form. +pub(crate) fn apply_patch_payload_command(payload: &ToolPayload) -> Option { match payload { ToolPayload::Function { arguments } => parse_arguments::(arguments) .ok() @@ -318,39 +311,6 @@ impl ToolHandler for ApplyPatchHandler { Some(Box::::default()) } - fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { - apply_patch_payload_command(&invocation.payload).map(|command| PreToolUsePayload { - tool_name: HookToolName::apply_patch(), - tool_input: serde_json::json!({ "command": command }), - }) - } - - // Hooks expose apply_patch through the stable `{ "command": ... }` shape, - // while the underlying tool stores the patch as either the function - // argument `input` or freeform custom-tool input. - fn with_updated_hook_input( - &self, - mut invocation: ToolInvocation, - updated_input: serde_json::Value, - ) -> Result { - let patch = updated_hook_command(&updated_input)?; - invocation.payload = match invocation.payload { - ToolPayload::Function { arguments } => ToolPayload::Function { - arguments: rewrite_function_string_argument( - &arguments, - "apply_patch", - "input", - patch, - )?, - }, - ToolPayload::Custom { .. } => ToolPayload::Custom { - input: patch.to_string(), - }, - payload => payload, - }; - Ok(invocation) - } - fn post_tool_use_payload( &self, invocation: &ToolInvocation, diff --git a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs index 04472e4623..4ee6bd2450 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs @@ -49,10 +49,8 @@ async fn pre_tool_use_payload_uses_json_patch_input() { arguments: json!({ "input": patch }).to_string(), }; let invocation = invocation_for_payload(payload).await; - let handler = ApplyPatchHandler; - assert_eq!( - handler.pre_tool_use_payload(&invocation), + crate::tools::hook_compat::pre_tool_use_payload(&invocation), Some(PreToolUsePayload { tool_name: HookToolName::apply_patch(), tool_input: json!({ "command": patch }), @@ -67,10 +65,8 @@ async fn pre_tool_use_payload_uses_freeform_patch_input() { input: patch.to_string(), }; let invocation = invocation_for_payload(payload).await; - let handler = ApplyPatchHandler; - assert_eq!( - handler.pre_tool_use_payload(&invocation), + crate::tools::hook_compat::pre_tool_use_payload(&invocation), Some(PreToolUsePayload { tool_name: HookToolName::apply_patch(), tool_input: json!({ "command": patch }), diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 7397f5a45b..1b9d6a613b 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -10,11 +10,9 @@ use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; use crate::tools::hook_names::HookToolName; use crate::tools::registry::PostToolUsePayload; -use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; use codex_tools::ToolName; -use serde_json::Value; pub struct McpHandler { tool_name: ToolName, @@ -37,40 +35,6 @@ impl ToolHandler for McpHandler { ToolKind::Mcp } - fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { - let ToolPayload::Mcp { raw_arguments, .. } = &invocation.payload else { - return None; - }; - - Some(PreToolUsePayload { - tool_name: HookToolName::new(self.tool_name.display()), - tool_input: mcp_hook_tool_input(raw_arguments), - }) - } - - // MCP hooks expose the full arguments object, so rewrites replace the - // serialized raw argument payload wholesale rather than patching one - // handler-owned field. - fn with_updated_hook_input( - &self, - mut invocation: ToolInvocation, - updated_input: Value, - ) -> Result { - invocation.payload = match invocation.payload { - ToolPayload::Mcp { server, tool, .. } => ToolPayload::Mcp { - server, - tool, - raw_arguments: serde_json::to_string(&updated_input).map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to serialize rewritten MCP arguments: {err}" - )) - })?, - }, - payload => payload, - }; - Ok(invocation) - } - fn post_tool_use_payload( &self, invocation: &ToolInvocation, @@ -137,14 +101,6 @@ impl ToolHandler for McpHandler { } } -fn mcp_hook_tool_input(raw_arguments: &str) -> Value { - if raw_arguments.trim().is_empty() { - return Value::Object(serde_json::Map::new()); - } - - serde_json::from_str(raw_arguments).unwrap_or_else(|_| Value::String(raw_arguments.to_string())) -} - #[cfg(test)] mod tests { use super::*; @@ -170,13 +126,8 @@ mod tests { .to_string(), }; let (session, turn) = make_session_and_context().await; - let handler = McpHandler::new(codex_tools::ToolName::namespaced( - "mcp__memory__", - "create_entities", - )); - assert_eq!( - handler.pre_tool_use_payload(&ToolInvocation { + crate::tools::hook_compat::pre_tool_use_payload(&ToolInvocation { session: session.into(), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), @@ -186,7 +137,7 @@ mod tests { source: ToolCallSource::Direct, payload, }), - Some(PreToolUsePayload { + Some(crate::tools::registry::PreToolUsePayload { tool_name: HookToolName::new("mcp__memory__create_entities"), tool_input: json!({ "entities": [{ @@ -262,6 +213,9 @@ mod tests { #[test] fn mcp_hook_tool_input_defaults_empty_args_to_object() { - assert_eq!(mcp_hook_tool_input(" "), json!({})); + assert_eq!( + crate::tools::hook_compat::mcp_hook_tool_input(" "), + json!({}) + ); } } diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index 8a6cb164cf..60ea73303b 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -24,7 +24,6 @@ use codex_sandboxing::policy_transforms::normalize_additional_permissions; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_absolute_path::AbsolutePathBufGuard; use serde::Deserialize; -use serde_json::Map; use serde_json::Value; use std::path::Path; @@ -54,6 +53,9 @@ pub use shell::ContainerExecHandler; pub use shell::LocalShellHandler; pub use shell::ShellCommandHandler; pub use shell::ShellHandler; +pub(crate) use shell::local_shell_payload_command; +pub(crate) use shell::shell_command_payload_command; +pub(crate) use shell::shell_function_payload_command; pub use test_sync::TestSyncHandler; pub use tool_search::ToolSearchHandler; pub use unavailable_tool::UnavailableToolHandler; @@ -62,7 +64,7 @@ pub use unified_exec::ExecCommandHandler; pub use unified_exec::WriteStdinHandler; pub use view_image::ViewImageHandler; -fn parse_arguments(arguments: &str) -> Result +pub(crate) fn parse_arguments(arguments: &str) -> Result where T: for<'de> Deserialize<'de>, { @@ -82,47 +84,6 @@ where parse_arguments(arguments) } -fn updated_hook_command(updated_input: &Value) -> Result<&str, FunctionCallError> { - updated_input - .get("command") - .and_then(Value::as_str) - .ok_or_else(|| { - FunctionCallError::RespondToModel( - "hook returned updatedInput without string field `command`".to_string(), - ) - }) -} - -fn rewrite_function_arguments( - arguments: &str, - tool_name: &str, - rewrite: impl FnOnce(&mut Map), -) -> Result { - let mut arguments: Value = parse_arguments(arguments)?; - let Value::Object(arguments) = &mut arguments else { - return Err(FunctionCallError::RespondToModel(format!( - "{tool_name} arguments must be an object" - ))); - }; - rewrite(arguments); - serde_json::to_string(&arguments).map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to serialize rewritten {tool_name} arguments: {err}" - )) - }) -} - -fn rewrite_function_string_argument( - arguments: &str, - tool_name: &str, - field_name: &str, - value: &str, -) -> Result { - rewrite_function_arguments(arguments, tool_name, |arguments| { - arguments.insert(field_name.to_string(), Value::String(value.to_string())); - }) -} - fn resolve_workdir_base_path( arguments: &str, default_cwd: &AbsolutePathBuf, diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 0bbf06f36b..26c4e9f820 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -25,13 +25,9 @@ use crate::tools::handlers::normalize_and_validate_additional_permissions; use crate::tools::handlers::parse_arguments; use crate::tools::handlers::parse_arguments_with_base_path; use crate::tools::handlers::resolve_workdir_base_path; -use crate::tools::handlers::rewrite_function_arguments; -use crate::tools::handlers::rewrite_function_string_argument; -use crate::tools::handlers::updated_hook_command; use crate::tools::hook_names::HookToolName; use crate::tools::orchestrator::ToolOrchestrator; use crate::tools::registry::PostToolUsePayload; -use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; use crate::tools::runtimes::shell::ShellRequest; @@ -59,7 +55,7 @@ pub struct ShellCommandHandler { backend: ShellCommandBackend, } -fn shell_function_payload_command(payload: &ToolPayload) -> Option { +pub(crate) fn shell_function_payload_command(payload: &ToolPayload) -> Option { let ToolPayload::Function { arguments } = payload else { return None; }; @@ -69,7 +65,7 @@ fn shell_function_payload_command(payload: &ToolPayload) -> Option { .map(|params| codex_shell_command::parse_command::shlex_join(¶ms.command)) } -fn local_shell_payload_command(payload: &ToolPayload) -> Option { +pub(crate) fn local_shell_payload_command(payload: &ToolPayload) -> Option { let ToolPayload::LocalShell { params } = payload else { return None; }; @@ -79,7 +75,7 @@ fn local_shell_payload_command(payload: &ToolPayload) -> Option { )) } -fn shell_command_payload_command(payload: &ToolPayload) -> Option { +pub(crate) fn shell_command_payload_command(payload: &ToolPayload) -> Option { let ToolPayload::Function { arguments } = payload else { return None; }; @@ -89,35 +85,6 @@ fn shell_command_payload_command(payload: &ToolPayload) -> Option { .map(|params| params.command) } -// Hooks expose legacy function shell tools as joined command strings, while -// their function payload stores argv. Split on the way back in to invert the -// hook-facing representation. -fn rewrite_shell_function_updated_hook_input( - mut invocation: ToolInvocation, - updated_input: JsonValue, - tool_name: &str, -) -> Result { - let ToolPayload::Function { arguments } = invocation.payload else { - return Err(FunctionCallError::RespondToModel(format!( - "hook input rewrite received unsupported {tool_name} payload" - ))); - }; - let command = shlex::split(updated_hook_command(&updated_input)?).ok_or_else(|| { - FunctionCallError::RespondToModel( - "hook returned shell input with an invalid command string".to_string(), - ) - })?; - invocation.payload = ToolPayload::Function { - arguments: rewrite_function_arguments(&arguments, tool_name, |arguments| { - arguments.insert( - "command".to_string(), - JsonValue::Array(command.into_iter().map(JsonValue::String).collect()), - ); - })?, - }; - Ok(invocation) -} - struct RunExecLikeArgs { tool_name: String, exec_params: ExecParams, @@ -247,18 +214,6 @@ impl ToolHandler for ShellHandler { .unwrap_or(true) } - fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { - shell_function_pre_tool_use_payload(invocation) - } - - fn with_updated_hook_input( - &self, - invocation: ToolInvocation, - updated_input: JsonValue, - ) -> Result { - rewrite_shell_function_updated_hook_input(invocation, updated_input, "shell") - } - fn post_tool_use_payload( &self, invocation: &ToolInvocation, @@ -333,18 +288,6 @@ impl ToolHandler for ContainerExecHandler { .unwrap_or(true) } - fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { - shell_function_pre_tool_use_payload(invocation) - } - - fn with_updated_hook_input( - &self, - invocation: ToolInvocation, - updated_input: JsonValue, - ) -> Result { - rewrite_shell_function_updated_hook_input(invocation, updated_input, "container.exec") - } - fn post_tool_use_payload( &self, invocation: &ToolInvocation, @@ -417,50 +360,6 @@ impl ToolHandler for LocalShellHandler { !is_known_safe_command(¶ms.command) } - fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { - local_shell_payload_command(&invocation.payload).map(|command| PreToolUsePayload { - tool_name: HookToolName::bash(), - tool_input: serde_json::json!({ "command": command }), - }) - } - - // Hooks see a joined shell command string, but local_shell stores argv. - // Split on the way back in to invert the hook-facing representation. - fn with_updated_hook_input( - &self, - mut invocation: ToolInvocation, - updated_input: JsonValue, - ) -> Result { - let command = updated_hook_command(&updated_input)?; - invocation.payload = match invocation.payload { - ToolPayload::Function { arguments } => { - let command = shlex::split(command).ok_or_else(|| { - FunctionCallError::RespondToModel( - "hook returned shell input with an invalid command string".to_string(), - ) - })?; - ToolPayload::Function { - arguments: rewrite_function_arguments(&arguments, "shell", |arguments| { - arguments.insert( - "command".to_string(), - JsonValue::Array(command.into_iter().map(JsonValue::String).collect()), - ); - })?, - } - } - ToolPayload::LocalShell { mut params } => { - params.command = shlex::split(command).ok_or_else(|| { - FunctionCallError::RespondToModel( - "hook returned shell input with an invalid command string".to_string(), - ) - })?; - ToolPayload::LocalShell { params } - } - payload => payload, - }; - Ok(invocation) - } - fn post_tool_use_payload( &self, invocation: &ToolInvocation, @@ -512,13 +411,6 @@ impl ToolHandler for LocalShellHandler { } } -fn shell_function_pre_tool_use_payload(invocation: &ToolInvocation) -> Option { - shell_function_payload_command(&invocation.payload).map(|command| PreToolUsePayload { - tool_name: HookToolName::bash(), - tool_input: serde_json::json!({ "command": command }), - }) -} - fn shell_function_post_tool_use_payload( invocation: &ToolInvocation, result: &FunctionToolOutput, @@ -569,34 +461,6 @@ impl ToolHandler for ShellCommandHandler { .unwrap_or(true) } - fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { - shell_command_payload_command(&invocation.payload).map(|command| PreToolUsePayload { - tool_name: HookToolName::bash(), - tool_input: serde_json::json!({ "command": command }), - }) - } - - fn with_updated_hook_input( - &self, - mut invocation: ToolInvocation, - updated_input: JsonValue, - ) -> Result { - let ToolPayload::Function { arguments } = invocation.payload else { - return Err(FunctionCallError::RespondToModel( - "hook input rewrite received unsupported shell_command payload".to_string(), - )); - }; - invocation.payload = ToolPayload::Function { - arguments: rewrite_function_string_argument( - &arguments, - "shell_command", - "command", - updated_hook_command(&updated_input)?, - )?, - }; - Ok(invocation) - } - fn post_tool_use_payload( &self, invocation: &ToolInvocation, diff --git a/codex-rs/core/src/tools/handlers/shell_tests.rs b/codex-rs/core/src/tools/handlers/shell_tests.rs index 4f44b953f4..449f6c94bd 100644 --- a/codex-rs/core/src/tools/handlers/shell_tests.rs +++ b/codex-rs/core/src/tools/handlers/shell_tests.rs @@ -17,10 +17,7 @@ use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolCallSource; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; -use crate::tools::handlers::ContainerExecHandler; -use crate::tools::handlers::LocalShellHandler; use crate::tools::handlers::ShellCommandHandler; -use crate::tools::handlers::ShellHandler; use crate::tools::hook_names::HookToolName; use crate::tools::registry::ToolHandler; use crate::turn_diff_tracker::TurnDiffTracker; @@ -224,10 +221,8 @@ async fn local_shell_pre_tool_use_payload_uses_joined_command() { }, }; let (session, turn) = make_session_and_context().await; - let handler = LocalShellHandler; - assert_eq!( - handler.pre_tool_use_payload(&ToolInvocation { + crate::tools::hook_compat::pre_tool_use_payload(&ToolInvocation { session: session.into(), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), @@ -250,12 +245,8 @@ async fn shell_command_pre_tool_use_payload_uses_raw_command() { arguments: json!({ "command": "printf shell command" }).to_string(), }; let (session, turn) = make_session_and_context().await; - let handler = ShellCommandHandler { - backend: super::ShellCommandBackend::Classic, - }; - assert_eq!( - handler.pre_tool_use_payload(&ToolInvocation { + crate::tools::hook_compat::pre_tool_use_payload(&ToolInvocation { session: session.into(), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), @@ -282,23 +273,20 @@ async fn shell_handler_rewrites_hook_command_back_to_argv() { .to_string(), }; let (session, turn) = make_session_and_context().await; - let handler = ShellHandler; - - let invocation = handler - .with_updated_hook_input( - ToolInvocation { - session: session.into(), - turn: turn.into(), - cancellation_token: tokio_util::sync::CancellationToken::new(), - tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), - call_id: "call-43".to_string(), - tool_name: codex_tools::ToolName::plain("shell"), - source: ToolCallSource::Direct, - payload, - }, - json!({ "command": "bash -lc 'printf new'" }), - ) - .expect("shell rewrite should succeed"); + let invocation = crate::tools::hook_compat::apply_updated_input( + ToolInvocation { + session: session.into(), + turn: turn.into(), + cancellation_token: tokio_util::sync::CancellationToken::new(), + tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), + call_id: "call-43".to_string(), + tool_name: codex_tools::ToolName::plain("shell"), + source: ToolCallSource::Direct, + payload, + }, + json!({ "command": "bash -lc 'printf new'" }), + ) + .expect("shell rewrite should succeed"); let ToolPayload::Function { arguments } = invocation.payload else { panic!("shell rewrite should preserve a function payload"); @@ -332,23 +320,20 @@ async fn container_exec_handler_rewrites_hook_command_back_to_argv() { .to_string(), }; let (session, turn) = make_session_and_context().await; - let handler = ContainerExecHandler; - - let invocation = handler - .with_updated_hook_input( - ToolInvocation { - session: session.into(), - turn: turn.into(), - cancellation_token: tokio_util::sync::CancellationToken::new(), - tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), - call_id: "call-44".to_string(), - tool_name: codex_tools::ToolName::plain("container.exec"), - source: ToolCallSource::Direct, - payload, - }, - json!({ "command": "bash -lc 'printf new'" }), - ) - .expect("container.exec rewrite should succeed"); + let invocation = crate::tools::hook_compat::apply_updated_input( + ToolInvocation { + session: session.into(), + turn: turn.into(), + cancellation_token: tokio_util::sync::CancellationToken::new(), + tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), + call_id: "call-44".to_string(), + tool_name: codex_tools::ToolName::plain("container.exec"), + source: ToolCallSource::Direct, + payload, + }, + json!({ "command": "bash -lc 'printf new'" }), + ) + .expect("container.exec rewrite should succeed"); let ToolPayload::Function { arguments } = invocation.payload else { panic!("container.exec rewrite should preserve a function payload"); diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index 1f56841425..224c0aaacd 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -14,11 +14,8 @@ use crate::tools::handlers::normalize_and_validate_additional_permissions; use crate::tools::handlers::parse_arguments; use crate::tools::handlers::parse_arguments_with_base_path; use crate::tools::handlers::resolve_tool_environment; -use crate::tools::handlers::rewrite_function_string_argument; -use crate::tools::handlers::updated_hook_command; use crate::tools::hook_names::HookToolName; use crate::tools::registry::PostToolUsePayload; -use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; use crate::unified_exec::ExecCommandRequest; @@ -48,7 +45,7 @@ pub struct WriteStdinHandler; #[derive(Debug, Deserialize)] pub(crate) struct ExecCommandArgs { - cmd: String, + pub(crate) cmd: String, #[serde(default)] pub(crate) workdir: Option, #[serde(default)] @@ -151,42 +148,6 @@ impl ToolHandler for ExecCommandHandler { !is_known_safe_command(&command) } - fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { - let ToolPayload::Function { arguments } = &invocation.payload else { - return None; - }; - - parse_arguments::(arguments) - .ok() - .map(|args| PreToolUsePayload { - tool_name: HookToolName::bash(), - tool_input: serde_json::json!({ "command": args.cmd }), - }) - } - - // Hooks normalize Bash-like tools to `{ "command": ... }`, while the - // exec_command wire schema still names the field `cmd`. - fn with_updated_hook_input( - &self, - mut invocation: ToolInvocation, - updated_input: serde_json::Value, - ) -> Result { - let ToolPayload::Function { arguments } = invocation.payload else { - return Err(FunctionCallError::RespondToModel( - "hook input rewrite received unsupported exec_command payload".to_string(), - )); - }; - invocation.payload = ToolPayload::Function { - arguments: rewrite_function_string_argument( - &arguments, - "exec_command", - "cmd", - updated_hook_command(&updated_input)?, - )?, - }; - Ok(invocation) - } - fn post_tool_use_payload( &self, invocation: &ToolInvocation, diff --git a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs index 8818b2e344..451c5fa3b3 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs @@ -184,10 +184,8 @@ async fn exec_command_pre_tool_use_payload_uses_raw_command() { arguments: serde_json::json!({ "cmd": "printf exec command" }).to_string(), }; let (session, turn) = make_session_and_context().await; - let handler = ExecCommandHandler; - assert_eq!( - handler.pre_tool_use_payload(&ToolInvocation { + crate::tools::hook_compat::pre_tool_use_payload(&ToolInvocation { session: session.into(), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), @@ -210,10 +208,8 @@ async fn exec_command_pre_tool_use_payload_skips_write_stdin() { arguments: serde_json::json!({ "chars": "echo hi" }).to_string(), }; let (session, turn) = make_session_and_context().await; - let handler = WriteStdinHandler; - assert_eq!( - handler.pre_tool_use_payload(&ToolInvocation { + crate::tools::hook_compat::pre_tool_use_payload(&ToolInvocation { session: session.into(), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), diff --git a/codex-rs/core/src/tools/hook_compat.rs b/codex-rs/core/src/tools/hook_compat.rs new file mode 100644 index 0000000000..5821598e66 --- /dev/null +++ b/codex-rs/core/src/tools/hook_compat.rs @@ -0,0 +1,276 @@ +use serde_json::Map; +use serde_json::Value; + +use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolPayload; +use crate::tools::handlers::apply_patch::apply_patch_payload_command; +use crate::tools::handlers::local_shell_payload_command; +use crate::tools::handlers::parse_arguments; +use crate::tools::handlers::shell_command_payload_command; +use crate::tools::handlers::shell_function_payload_command; +use crate::tools::handlers::unified_exec::ExecCommandArgs; +use crate::tools::hook_names::HookToolName; +use crate::tools::registry::PreToolUsePayload; + +/// Projects native tool payloads into the stable input shape exposed to hooks. +/// +/// The hook protocol intentionally hides a few native tool-schema differences: +/// +/// - Bash-like tools are exposed as `Bash` with `{ "command": }`, +/// even though their native payloads store commands as argv, `command`, or +/// `cmd` depending on the concrete tool. +/// - `apply_patch` exposes its raw patch body through `{ "command": }` +/// for compatibility with existing hook consumers. +/// - MCP tools already use arbitrary JSON argument objects, so hooks see those +/// arguments directly rather than through a compatibility shape. +pub(crate) fn pre_tool_use_payload(invocation: &ToolInvocation) -> Option { + match invocation.tool_name.name.as_str() { + "shell" | "container.exec" => { + shell_function_payload_command(&invocation.payload).map(bash_payload) + } + "local_shell" => local_shell_payload_command(&invocation.payload).map(bash_payload), + "shell_command" => shell_command_payload_command(&invocation.payload).map(bash_payload), + "exec_command" => exec_command_payload_command(&invocation.payload).map(bash_payload), + "apply_patch" => { + apply_patch_payload_command(&invocation.payload).map(|command| PreToolUsePayload { + tool_name: HookToolName::apply_patch(), + tool_input: serde_json::json!({ "command": command }), + }) + } + _ => mcp_payload(invocation), + } +} + +/// Rebuilds native tool payloads from hook-facing `updatedInput`. +/// +/// This is the inverse of [`pre_tool_use_payload`]: Bash-like and +/// `apply_patch` updates come back through the compatibility `{ "command": ... }` +/// shape and must be written into each tool's native schema, while MCP updates +/// replace the raw argument object wholesale because the hook-facing and native +/// representations are the same. +pub(crate) fn apply_updated_input( + invocation: ToolInvocation, + updated_input: Value, +) -> Result { + match invocation.tool_name.name.as_str() { + "shell" => rewrite_shell_function_updated_input(invocation, updated_input, "shell"), + "container.exec" => { + rewrite_shell_function_updated_input(invocation, updated_input, "container.exec") + } + "local_shell" => rewrite_local_shell_updated_input(invocation, updated_input), + "shell_command" => rewrite_shell_command_updated_input(invocation, updated_input), + "exec_command" => rewrite_exec_command_updated_input(invocation, updated_input), + "apply_patch" => rewrite_apply_patch_updated_input(invocation, updated_input), + _ => rewrite_mcp_updated_input(invocation, updated_input), + } +} + +fn bash_payload(command: String) -> PreToolUsePayload { + PreToolUsePayload { + tool_name: HookToolName::bash(), + tool_input: serde_json::json!({ "command": command }), + } +} + +fn exec_command_payload_command(payload: &ToolPayload) -> Option { + let ToolPayload::Function { arguments } = payload else { + return None; + }; + + parse_arguments::(arguments) + .ok() + .map(|args| args.cmd) +} + +fn mcp_payload(invocation: &ToolInvocation) -> Option { + let ToolPayload::Mcp { raw_arguments, .. } = &invocation.payload else { + return None; + }; + + Some(PreToolUsePayload { + tool_name: HookToolName::new(invocation.tool_name.display()), + tool_input: mcp_hook_tool_input(raw_arguments), + }) +} + +fn updated_hook_command(updated_input: &Value) -> Result<&str, FunctionCallError> { + updated_input + .get("command") + .and_then(Value::as_str) + .ok_or_else(|| { + FunctionCallError::RespondToModel( + "hook returned updatedInput without string field `command`".to_string(), + ) + }) +} + +fn rewrite_function_arguments( + arguments: &str, + tool_name: &str, + rewrite: impl FnOnce(&mut Map), +) -> Result { + let mut arguments: Value = parse_arguments(arguments)?; + let Value::Object(arguments) = &mut arguments else { + return Err(FunctionCallError::RespondToModel(format!( + "{tool_name} arguments must be an object" + ))); + }; + rewrite(arguments); + serde_json::to_string(&arguments).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to serialize rewritten {tool_name} arguments: {err}" + )) + }) +} + +fn rewrite_function_string_argument( + arguments: &str, + tool_name: &str, + field_name: &str, + value: &str, +) -> Result { + rewrite_function_arguments(arguments, tool_name, |arguments| { + arguments.insert(field_name.to_string(), Value::String(value.to_string())); + }) +} + +/// Rehydrates legacy function-style shell tools from hook-facing command text. +fn rewrite_shell_function_updated_input( + mut invocation: ToolInvocation, + updated_input: Value, + tool_name: &str, +) -> Result { + let ToolPayload::Function { arguments } = invocation.payload else { + return Err(FunctionCallError::RespondToModel(format!( + "hook input rewrite received unsupported {tool_name} payload" + ))); + }; + let command = shlex::split(updated_hook_command(&updated_input)?).ok_or_else(|| { + FunctionCallError::RespondToModel( + "hook returned shell input with an invalid command string".to_string(), + ) + })?; + invocation.payload = ToolPayload::Function { + arguments: rewrite_function_arguments(&arguments, tool_name, |arguments| { + arguments.insert( + "command".to_string(), + Value::Array(command.into_iter().map(Value::String).collect()), + ); + })?, + }; + Ok(invocation) +} + +/// Rehydrates `local_shell` argv from hook-facing command text. +fn rewrite_local_shell_updated_input( + mut invocation: ToolInvocation, + updated_input: Value, +) -> Result { + let command = updated_hook_command(&updated_input)?; + invocation.payload = match invocation.payload { + ToolPayload::LocalShell { mut params } => { + params.command = shlex::split(command).ok_or_else(|| { + FunctionCallError::RespondToModel( + "hook returned shell input with an invalid command string".to_string(), + ) + })?; + ToolPayload::LocalShell { params } + } + payload => payload, + }; + Ok(invocation) +} + +/// Stores hook-facing command text back into the native `shell_command.command`. +fn rewrite_shell_command_updated_input( + mut invocation: ToolInvocation, + updated_input: Value, +) -> Result { + let ToolPayload::Function { arguments } = invocation.payload else { + return Err(FunctionCallError::RespondToModel( + "hook input rewrite received unsupported shell_command payload".to_string(), + )); + }; + invocation.payload = ToolPayload::Function { + arguments: rewrite_function_string_argument( + &arguments, + "shell_command", + "command", + updated_hook_command(&updated_input)?, + )?, + }; + Ok(invocation) +} + +/// Stores hook-facing command text back into the native `exec_command.cmd`. +fn rewrite_exec_command_updated_input( + mut invocation: ToolInvocation, + updated_input: Value, +) -> Result { + let ToolPayload::Function { arguments } = invocation.payload else { + return Err(FunctionCallError::RespondToModel( + "hook input rewrite received unsupported exec_command payload".to_string(), + )); + }; + invocation.payload = ToolPayload::Function { + arguments: rewrite_function_string_argument( + &arguments, + "exec_command", + "cmd", + updated_hook_command(&updated_input)?, + )?, + }; + Ok(invocation) +} + +/// Stores hook-facing patch text back into the native apply_patch payload form. +fn rewrite_apply_patch_updated_input( + mut invocation: ToolInvocation, + updated_input: Value, +) -> Result { + let patch = updated_hook_command(&updated_input)?; + invocation.payload = match invocation.payload { + ToolPayload::Function { arguments } => ToolPayload::Function { + arguments: rewrite_function_string_argument(&arguments, "apply_patch", "input", patch)?, + }, + ToolPayload::Custom { .. } => ToolPayload::Custom { + input: patch.to_string(), + }, + payload => payload, + }; + Ok(invocation) +} + +/// Replaces MCP raw arguments directly because MCP hooks expose that JSON object as-is. +fn rewrite_mcp_updated_input( + mut invocation: ToolInvocation, + updated_input: Value, +) -> Result { + invocation.payload = match invocation.payload { + ToolPayload::Mcp { server, tool, .. } => ToolPayload::Mcp { + server, + tool, + raw_arguments: serde_json::to_string(&updated_input).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to serialize rewritten MCP arguments: {err}" + )) + })?, + }, + payload => { + return Err(FunctionCallError::RespondToModel(format!( + "tool {} does not support hook input rewriting for payload {payload:?}", + invocation.tool_name.display() + ))); + } + }; + Ok(invocation) +} + +pub(crate) fn mcp_hook_tool_input(raw_arguments: &str) -> Value { + if raw_arguments.trim().is_empty() { + return Value::Object(Map::new()); + } + + serde_json::from_str(raw_arguments).unwrap_or_else(|_| Value::String(raw_arguments.to_string())) +} diff --git a/codex-rs/core/src/tools/mod.rs b/codex-rs/core/src/tools/mod.rs index 659a7d3e54..eeea95687e 100644 --- a/codex-rs/core/src/tools/mod.rs +++ b/codex-rs/core/src/tools/mod.rs @@ -2,6 +2,7 @@ pub(crate) mod code_mode; pub(crate) mod context; pub(crate) mod events; pub(crate) mod handlers; +pub(crate) mod hook_compat; pub(crate) mod hook_names; pub(crate) mod network_approval; pub(crate) mod orchestrator; diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 4547fe8ebf..c3aa087313 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -16,6 +16,7 @@ use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; +use crate::tools::hook_compat; use crate::tools::hook_names::HookToolName; use crate::tools::tool_dispatch_trace::ToolDispatchTrace; use codex_hooks::HookEvent; @@ -69,10 +70,6 @@ pub trait ToolHandler: Send + Sync { async { false } } - fn pre_tool_use_payload(&self, _invocation: &ToolInvocation) -> Option { - None - } - fn post_tool_use_payload( &self, _invocation: &ToolInvocation, @@ -81,20 +78,6 @@ pub trait ToolHandler: Send + Sync { None } - /// Rebuilds a tool invocation from hook-facing `tool_input`. - /// - /// Tools that opt into input-rewriting hooks should invert the same stable - /// hook contract they expose from `pre_tool_use_payload`. - fn with_updated_hook_input( - &self, - _invocation: ToolInvocation, - _updated_input: Value, - ) -> Result { - Err(FunctionCallError::RespondToModel( - "tool does not support hook input rewriting".to_string(), - )) - } - /// Creates an optional consumer for streamed tool argument diffs. fn create_diff_consumer(&self) -> Option> { None @@ -181,14 +164,6 @@ trait AnyToolHandler: Send + Sync { fn is_mutating<'a>(&'a self, invocation: &'a ToolInvocation) -> BoxFuture<'a, bool>; - fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option; - - fn with_updated_hook_input( - &self, - invocation: ToolInvocation, - updated_input: Value, - ) -> Result; - fn create_diff_consumer(&self) -> Option>; fn handle_any<'a>( &'a self, @@ -208,18 +183,6 @@ where Box::pin(ToolHandler::is_mutating(self, invocation)) } - fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { - ToolHandler::pre_tool_use_payload(self, invocation) - } - - fn with_updated_hook_input( - &self, - invocation: ToolInvocation, - updated_input: Value, - ) -> Result { - ToolHandler::with_updated_hook_input(self, invocation, updated_input) - } - fn create_diff_consumer(&self) -> Option> { ToolHandler::create_diff_consumer(self) } @@ -378,7 +341,7 @@ impl ToolRegistry { return Err(err); } - if let Some(pre_tool_use_payload) = handler.pre_tool_use_payload(&invocation) { + if let Some(pre_tool_use_payload) = hook_compat::pre_tool_use_payload(&invocation) { match run_pre_tool_use_hooks( &invocation.session, &invocation.turn, @@ -396,7 +359,7 @@ impl ToolRegistry { crate::hook_runtime::PreToolUseHookResult::Continue { updated_input: Some(updated_input), } => { - invocation = handler.with_updated_hook_input(invocation, updated_input)?; + invocation = hook_compat::apply_updated_input(invocation, updated_input)?; } crate::hook_runtime::PreToolUseHookResult::Continue { updated_input: None,