diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index bed5252464..e078cbebf1 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -48,6 +48,8 @@ pub(crate) enum PreToolUseHookResult { Blocked(String), } +pub(crate) const MAX_HOOK_INPUT_REWRITES: usize = 8; + pub(crate) enum PendingInputHookDisposition { Accepted(Box), Blocked { additional_contexts: Vec }, diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 34f37a98aa..3925e6a31c 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -25,6 +25,7 @@ use crate::guardian::guardian_timeout_message; use crate::guardian::new_guardian_review_id; use crate::guardian::review_approval_request; use crate::guardian::routes_approval_to_guardian; +use crate::hook_runtime::MAX_HOOK_INPUT_REWRITES; use crate::hook_runtime::run_permission_request_hooks; use crate::mcp_openai_file::rewrite_mcp_tool_arguments_for_openai_files; use crate::mcp_tool_approval_templates::RenderedMcpToolApprovalParam; @@ -215,7 +216,7 @@ pub(crate) async fn handle_mcp_tool_call( let result = match decision { McpToolApprovalDecision::AcceptWithUpdatedInput(updated_input) => { input_rewrites += 1; - if input_rewrites > 8 { + if input_rewrites > MAX_HOOK_INPUT_REWRITES { notify_mcp_tool_call_skip( sess.as_ref(), turn_context.as_ref(), diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 71a40a4211..581dd7747e 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -22,6 +22,8 @@ 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; @@ -323,34 +325,16 @@ impl ToolHandler for ApplyPatchHandler { mut invocation: ToolInvocation, updated_input: serde_json::Value, ) -> Result { - let patch = updated_input - .get("command") - .and_then(serde_json::Value::as_str) - .ok_or_else(|| { - FunctionCallError::RespondToModel( - "hook returned updatedInput without string field `command`".to_string(), - ) - })?; + let patch = updated_hook_command(&updated_input)?; invocation.payload = match invocation.payload { - ToolPayload::Function { arguments } => { - let mut arguments: serde_json::Value = parse_arguments(&arguments)?; - let serde_json::Value::Object(arguments) = &mut arguments else { - return Err(FunctionCallError::RespondToModel( - "apply_patch arguments must be an object".to_string(), - )); - }; - arguments.insert( - "input".to_string(), - serde_json::Value::String(patch.to_string()), - ); - ToolPayload::Function { - arguments: serde_json::to_string(&arguments).map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to serialize rewritten apply_patch arguments: {err}" - )) - })?, - } - } + ToolPayload::Function { arguments } => ToolPayload::Function { + arguments: rewrite_function_string_argument( + &arguments, + "apply_patch", + "input", + patch, + )?, + }, ToolPayload::Custom { .. } => ToolPayload::Custom { input: patch.to_string(), }, diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index 0ddd1e5062..bd53311556 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -25,6 +25,7 @@ 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; @@ -74,6 +75,48 @@ 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<(), FunctionCallError>, +) -> 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())); + Ok(()) + }) +} + 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 c0e3d93340..18ab33b2bd 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -25,6 +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; @@ -75,17 +78,6 @@ fn shell_command_payload_command(payload: &ToolPayload) -> Option { .map(|params| params.command) } -fn updated_hook_command(updated_input: &JsonValue) -> Result<&str, FunctionCallError> { - updated_input - .get("command") - .and_then(JsonValue::as_str) - .ok_or_else(|| { - FunctionCallError::RespondToModel( - "hook returned updatedInput without string field `command`".to_string(), - ) - }) -} - struct RunExecLikeArgs { tool_name: String, exec_params: ExecParams, @@ -236,21 +228,13 @@ impl ToolHandler for ShellHandler { "hook returned shell input with an invalid command string".to_string(), ) })?; - let mut arguments: JsonValue = parse_arguments(&arguments)?; - let JsonValue::Object(arguments) = &mut arguments else { - return Err(FunctionCallError::RespondToModel( - "shell arguments must be an object".to_string(), - )); - }; - arguments.insert( - "command".to_string(), - JsonValue::Array(command.into_iter().map(JsonValue::String).collect()), - ); ToolPayload::Function { - arguments: serde_json::to_string(&arguments).map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to serialize rewritten shell arguments: {err}" - )) + arguments: rewrite_function_arguments(&arguments, "shell", |arguments| { + arguments.insert( + "command".to_string(), + JsonValue::Array(command.into_iter().map(JsonValue::String).collect()), + ); + Ok(()) })?, } } @@ -391,22 +375,13 @@ impl ToolHandler for ShellCommandHandler { "hook input rewrite received unsupported shell_command payload".to_string(), )); }; - let mut arguments: JsonValue = parse_arguments(&arguments)?; - let JsonValue::Object(arguments) = &mut arguments else { - return Err(FunctionCallError::RespondToModel( - "shell_command arguments must be an object".to_string(), - )); - }; - arguments.insert( - "command".to_string(), - JsonValue::String(updated_hook_command(&updated_input)?.to_string()), - ); invocation.payload = ToolPayload::Function { - arguments: serde_json::to_string(&arguments).map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to serialize rewritten shell_command arguments: {err}" - )) - })?, + arguments: rewrite_function_string_argument( + &arguments, + "shell_command", + "command", + updated_hook_command(&updated_input)?, + )?, }; Ok(invocation) } diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index 1fd8cbf0a8..9c2853aab0 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -14,6 +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_workdir_base_path; +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; @@ -162,30 +164,13 @@ impl ToolHandler for UnifiedExecHandler { "hook input rewrite received unsupported unified_exec payload".to_string(), )); }; - let command = updated_input - .get("command") - .and_then(serde_json::Value::as_str) - .ok_or_else(|| { - FunctionCallError::RespondToModel( - "hook returned updatedInput without string field `command`".to_string(), - ) - })?; - let mut arguments: serde_json::Value = parse_arguments(&arguments)?; - let serde_json::Value::Object(arguments) = &mut arguments else { - return Err(FunctionCallError::RespondToModel( - "exec_command arguments must be an object".to_string(), - )); - }; - arguments.insert( - "cmd".to_string(), - serde_json::Value::String(command.to_string()), - ); invocation.payload = ToolPayload::Function { - arguments: serde_json::to_string(&arguments).map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to serialize rewritten exec_command arguments: {err}" - )) - })?, + arguments: rewrite_function_string_argument( + &arguments, + "exec_command", + "cmd", + updated_hook_command(&updated_input)?, + )?, }; Ok(invocation) } diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index e51d30f856..21ce0c4f73 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -5,6 +5,7 @@ use std::time::Instant; use crate::function_tool::FunctionCallError; use crate::goals::GoalRuntimeEvent; +use crate::hook_runtime::MAX_HOOK_INPUT_REWRITES; use crate::hook_runtime::record_additional_contexts; use crate::hook_runtime::run_post_tool_use_hooks; use crate::hook_runtime::run_pre_tool_use_hooks; @@ -231,7 +232,7 @@ where match self.handle(invocation.clone()).await { Err(FunctionCallError::UpdatedInput(updated_input)) => { rewrites += 1; - if rewrites > 8 { + if rewrites > MAX_HOOK_INPUT_REWRITES { return Err(FunctionCallError::RespondToModel( "hook input rewrite limit exceeded".to_string(), ));