From 04ae865e76e43e38bd529702ffcf5c9e448ee1f0 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Wed, 6 May 2026 12:22:50 -0700 Subject: [PATCH] support rewrites for legacy shell handlers --- codex-rs/core/src/tools/handlers/shell.rs | 45 ++++++++ .../core/src/tools/handlers/shell_tests.rs | 103 ++++++++++++++++++ 2 files changed, 148 insertions(+) diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 45365dfbfd..0bbf06f36b 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -89,6 +89,35 @@ 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, @@ -222,6 +251,14 @@ impl ToolHandler for ShellHandler { 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, @@ -300,6 +337,14 @@ impl ToolHandler for ContainerExecHandler { 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, diff --git a/codex-rs/core/src/tools/handlers/shell_tests.rs b/codex-rs/core/src/tools/handlers/shell_tests.rs index 8a32e5404b..4f44b953f4 100644 --- a/codex-rs/core/src/tools/handlers/shell_tests.rs +++ b/codex-rs/core/src/tools/handlers/shell_tests.rs @@ -2,6 +2,7 @@ use std::path::PathBuf; use std::sync::Arc; use codex_protocol::models::ShellCommandToolCallParams; +use codex_protocol::models::ShellToolCallParams; use core_test_support::PathBufExt; use core_test_support::test_path_buf; use pretty_assertions::assert_eq; @@ -16,8 +17,10 @@ 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; @@ -269,6 +272,106 @@ async fn shell_command_pre_tool_use_payload_uses_raw_command() { ); } +#[tokio::test] +async fn shell_handler_rewrites_hook_command_back_to_argv() { + let payload = ToolPayload::Function { + arguments: json!({ + "command": ["bash", "-lc", "printf old"], + "workdir": "subdir", + }) + .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 ToolPayload::Function { arguments } = invocation.payload else { + panic!("shell rewrite should preserve a function payload"); + }; + assert_eq!( + serde_json::from_str::(&arguments) + .expect("rewritten shell arguments should parse"), + ShellToolCallParams { + command: vec![ + "bash".to_string(), + "-lc".to_string(), + "printf new".to_string(), + ], + workdir: Some("subdir".to_string()), + timeout_ms: None, + sandbox_permissions: None, + additional_permissions: None, + prefix_rule: None, + justification: None, + } + ); +} + +#[tokio::test] +async fn container_exec_handler_rewrites_hook_command_back_to_argv() { + let payload = ToolPayload::Function { + arguments: json!({ + "command": ["bash", "-lc", "printf old"], + "workdir": "subdir", + }) + .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 ToolPayload::Function { arguments } = invocation.payload else { + panic!("container.exec rewrite should preserve a function payload"); + }; + assert_eq!( + serde_json::from_str::(&arguments) + .expect("rewritten container.exec arguments should parse"), + ShellToolCallParams { + command: vec![ + "bash".to_string(), + "-lc".to_string(), + "printf new".to_string(), + ], + workdir: Some("subdir".to_string()), + timeout_ms: None, + sandbox_permissions: None, + additional_permissions: None, + prefix_rule: None, + justification: None, + } + ); +} + #[tokio::test] async fn build_post_tool_use_payload_uses_tool_output_wire_value() { let payload = ToolPayload::Function {