From 61724ae0832376f0d55f51751562a4dbd31570a3 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Thu, 30 Apr 2026 16:34:32 -0700 Subject: [PATCH] support hook input rewrites --- codex-rs/core/src/function_tool.rs | 2 + codex-rs/core/src/hook_runtime.rs | 41 +- codex-rs/core/src/mcp_tool_call.rs | 77 ++- codex-rs/core/src/mcp_tool_call_tests.rs | 74 +++ codex-rs/core/src/stream_events_utils.rs | 5 + codex-rs/core/src/tools/events.rs | 7 + .../core/src/tools/handlers/apply_patch.rs | 41 ++ codex-rs/core/src/tools/handlers/mcp.rs | 20 + codex-rs/core/src/tools/handlers/shell.rs | 85 ++++ .../core/src/tools/handlers/unified_exec.rs | 41 ++ codex-rs/core/src/tools/network_approval.rs | 14 +- codex-rs/core/src/tools/orchestrator.rs | 20 +- codex-rs/core/src/tools/registry.rs | 71 ++- .../tools/runtimes/shell/unix_escalation.rs | 16 +- codex-rs/core/src/tools/sandboxing.rs | 1 + codex-rs/core/src/unified_exec/errors.rs | 3 + .../core/src/unified_exec/process_manager.rs | 6 + codex-rs/core/tests/suite/hooks.rs | 438 ++++++++++++++++++ codex-rs/core/tests/suite/hooks_mcp.rs | 114 +++++ ...mission-request.command.output.schema.json | 2 +- codex-rs/hooks/src/engine/output_parser.rs | 70 ++- .../hooks/src/events/permission_request.rs | 47 +- codex-rs/hooks/src/events/pre_tool_use.rs | 49 ++ codex-rs/hooks/src/schema.rs | 4 +- 24 files changed, 1182 insertions(+), 66 deletions(-) diff --git a/codex-rs/core/src/function_tool.rs b/codex-rs/core/src/function_tool.rs index 240e04361c..5fd6639958 100644 --- a/codex-rs/core/src/function_tool.rs +++ b/codex-rs/core/src/function_tool.rs @@ -4,6 +4,8 @@ use thiserror::Error; pub enum FunctionCallError { #[error("{0}")] RespondToModel(String), + #[error("tool input rewritten by hook")] + UpdatedInput(serde_json::Value), #[error("LocalShellCall without call_id or id")] MissingLocalShellCallId, #[error("Fatal error: {0}")] diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 9a92854515..bed5252464 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -43,6 +43,11 @@ pub(crate) struct HookRuntimeOutcome { pub additional_contexts: Vec, } +pub(crate) enum PreToolUseHookResult { + Continue { updated_input: Option }, + Blocked(String), +} + pub(crate) enum PendingInputHookDisposition { Accepted(Box), Blocked { additional_contexts: Vec }, @@ -140,7 +145,7 @@ pub(crate) async fn run_pre_tool_use_hooks( tool_use_id: String, tool_name: &HookToolName, tool_input: &Value, -) -> Option { +) -> PreToolUseHookResult { let request = PreToolUseRequest { session_id: sess.conversation_id, turn_id: turn_context.sub_id.clone(), @@ -161,24 +166,32 @@ pub(crate) async fn run_pre_tool_use_hooks( hook_events, should_block, block_reason, + updated_input, } = hooks.run_pre_tool_use(request).await; emit_hook_completed_events(sess, turn_context, hook_events).await; if should_block { - block_reason.map(|reason| { - if (tool_name.name() == "Bash" || tool_name.name() == "apply_patch") - && let Some(command) = tool_input.get("command").and_then(Value::as_str) - { - format!("Command blocked by PreToolUse hook: {reason}. Command: {command}") - } else { - format!( - "Tool call blocked by PreToolUse hook: {reason}. Tool: {}", - tool_name.name() - ) - } - }) + block_reason.map_or( + PreToolUseHookResult::Continue { + updated_input: None, + }, + |reason| { + if (tool_name.name() == "Bash" || tool_name.name() == "apply_patch") + && let Some(command) = tool_input.get("command").and_then(Value::as_str) + { + PreToolUseHookResult::Blocked(format!( + "Command blocked by PreToolUse hook: {reason}. Command: {command}" + )) + } else { + PreToolUseHookResult::Blocked(format!( + "Tool call blocked by PreToolUse hook: {reason}. Tool: {}", + tool_name.name() + )) + } + }, + ) } else { - None + PreToolUseHookResult::Continue { updated_input } } } diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index ed8451e352..34f37a98aa 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -192,18 +192,45 @@ pub(crate) async fn handle_mcp_tool_call( }); notify_mcp_tool_call_event(sess.as_ref(), turn_context.as_ref(), tool_call_begin_event).await; - if let Some(decision) = maybe_request_mcp_tool_approval( - &sess, - turn_context, - &call_id, - &invocation, - &hook_tool_name, - metadata.as_ref(), - approval_mode, - ) - .await - { + let mut invocation = invocation; + let mut input_rewrites = 0; + loop { + let Some(decision) = maybe_request_mcp_tool_approval( + &sess, + turn_context, + &call_id, + &invocation, + &hook_tool_name, + metadata.as_ref(), + approval_mode, + ) + .await + else { + break; + }; + let tool_input = invocation + .arguments + .clone() + .unwrap_or_else(|| JsonValue::Object(serde_json::Map::new())); let result = match decision { + McpToolApprovalDecision::AcceptWithUpdatedInput(updated_input) => { + input_rewrites += 1; + if input_rewrites > 8 { + notify_mcp_tool_call_skip( + sess.as_ref(), + turn_context.as_ref(), + &call_id, + invocation, + mcp_app_resource_uri.clone(), + "hook input rewrite limit exceeded".to_string(), + /*already_started*/ true, + ) + .await + } else { + invocation.arguments = Some(updated_input); + continue; + } + } McpToolApprovalDecision::Accept | McpToolApprovalDecision::AcceptForSession | McpToolApprovalDecision::AcceptAndRemember => { @@ -270,8 +297,7 @@ pub(crate) async fn handle_mcp_tool_call( return HandledMcpToolCall { result: CallToolResult::from_result(result), - tool_input: arguments_value - .unwrap_or_else(|| JsonValue::Object(serde_json::Map::new())), + tool_input, }; } @@ -747,6 +773,7 @@ async fn maybe_track_codex_app_used( #[derive(Debug, Clone, PartialEq, Eq)] enum McpToolApprovalDecision { Accept, + AcceptWithUpdatedInput(JsonValue), AcceptForSession, AcceptAndRemember, Decline { message: Option }, @@ -1000,21 +1027,34 @@ async fn maybe_request_mcp_tool_approval( return Some(McpToolApprovalDecision::Accept); } + let permission_request_tool_input = invocation + .arguments + .clone() + .unwrap_or_else(|| serde_json::Value::Object(serde_json::Map::new())); match run_permission_request_hooks( sess, turn_context, call_id, PermissionRequestPayload { tool_name: HookToolName::new(hook_tool_name), - tool_input: invocation - .arguments - .clone() - .unwrap_or_else(|| serde_json::Value::Object(serde_json::Map::new())), + tool_input: permission_request_tool_input.clone(), }, ) .await { - Some(PermissionRequestDecision::Allow) => { + Some(PermissionRequestDecision::Allow { + updated_input: Some(updated_input), + }) => { + if updated_input == permission_request_tool_input { + return Some(McpToolApprovalDecision::Accept); + } + return Some(McpToolApprovalDecision::AcceptWithUpdatedInput( + updated_input, + )); + } + Some(PermissionRequestDecision::Allow { + updated_input: None, + }) => { return Some(McpToolApprovalDecision::Accept); } Some(PermissionRequestDecision::Deny { message }) => { @@ -1768,6 +1808,7 @@ async fn apply_mcp_tool_approval_decision( } } McpToolApprovalDecision::Accept + | McpToolApprovalDecision::AcceptWithUpdatedInput(_) | McpToolApprovalDecision::Decline { .. } | McpToolApprovalDecision::Cancel | McpToolApprovalDecision::BlockedBySafetyMonitor(_) => {} diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 6a14126681..a040883069 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -2069,6 +2069,80 @@ async fn permission_request_hook_allows_mcp_tool_call() { ); } +#[tokio::test] +async fn permission_request_hook_can_update_mcp_tool_input() { + let (mut session, turn_context) = make_session_and_context().await; + install_mcp_permission_request_hook( + &mut session, + &turn_context, + "mcp__memory__.*", + &serde_json::json!({ + "hookSpecificOutput": { + "hookEventName": "PermissionRequest", + "decision": { + "behavior": "allow", + "updatedInput": { + "entities": [{ + "name": "Grace", + "entityType": "person" + }] + } + } + } + }), + ); + let session = Arc::new(session); + let turn_context = Arc::new(turn_context); + let invocation = McpInvocation { + server: "memory".to_string(), + tool: "create_entities".to_string(), + arguments: Some(serde_json::json!({ + "entities": [{ + "name": "Ada", + "entityType": "person" + }] + })), + }; + let metadata = McpToolApprovalMetadata { + annotations: Some(annotations( + Some(false), + Some(true), + /*open_world*/ None, + )), + connector_id: None, + connector_name: None, + connector_description: None, + tool_title: Some("Create entities".to_string()), + tool_description: None, + mcp_app_resource_uri: None, + codex_apps_meta: None, + openai_file_input_params: None, + }; + + let decision = maybe_request_mcp_tool_approval( + &session, + &turn_context, + "call-mcp-hook", + &invocation, + "mcp__memory__create_entities", + Some(&metadata), + AppToolApproval::Auto, + ) + .await; + + assert_eq!( + decision, + Some(McpToolApprovalDecision::AcceptWithUpdatedInput( + serde_json::json!({ + "entities": [{ + "name": "Grace", + "entityType": "person" + }] + }) + )) + ); +} + #[tokio::test] async fn permission_request_hook_uses_hook_tool_name_without_metadata() { let (mut session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 5a31d18020..328bce0c14 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -336,6 +336,11 @@ pub(crate) async fn handle_output_item_done( output.needs_follow_up = true; } + Err(FunctionCallError::UpdatedInput(_)) => { + return Err(CodexErr::Fatal( + "updated tool input escaped dispatch".to_string(), + )); + } // A fatal error occurred; surface it back into history. Err(FunctionCallError::Fatal(message)) => { return Err(CodexErr::Fatal(message)); diff --git a/codex-rs/core/src/tools/events.rs b/codex-rs/core/src/tools/events.rs index 2b215a043d..fabc9a8141 100644 --- a/codex-rs/core/src/tools/events.rs +++ b/codex-rs/core/src/tools/events.rs @@ -353,6 +353,13 @@ impl ToolEmitter { let result = Err(FunctionCallError::RespondToModel(normalized)); (event, result) } + Err(ToolError::UpdatedInput(updated_input)) => { + let event = ToolEventStage::Failure(ToolEventFailure::Message( + "tool input rewritten by hook".to_string(), + )); + let result = Err(FunctionCallError::UpdatedInput(updated_input)); + (event, result) + } }; self.emit(ctx, event).await; result diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index d71eb7931a..71a40a4211 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -318,6 +318,47 @@ impl ToolHandler for ApplyPatchHandler { }) } + fn with_updated_hook_input( + &self, + 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(), + ) + })?; + 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::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/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 568e456158..7bda02f7a4 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -34,6 +34,26 @@ impl ToolHandler for McpHandler { }) } + 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, diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index b7512b7076..c0e3d93340 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -75,6 +75,17 @@ 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, @@ -212,6 +223,50 @@ impl ToolHandler for ShellHandler { }) } + 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(), + ) + })?; + 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}" + )) + })?, + } + } + 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, @@ -326,6 +381,36 @@ impl ToolHandler for ShellCommandHandler { }) } + 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(), + )); + }; + 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}" + )) + })?, + }; + Ok(invocation) + } + fn post_tool_use_payload( &self, invocation: &ToolInvocation, diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index 10c8deeb3f..1fd8cbf0a8 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -152,6 +152,44 @@ impl ToolHandler for UnifiedExecHandler { }) } + 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 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}" + )) + })?, + }; + Ok(invocation) + } + fn post_tool_use_payload( &self, invocation: &ToolInvocation, @@ -355,6 +393,9 @@ impl ToolHandler for UnifiedExecHandler { .await { Ok(response) => response, + Err(UnifiedExecError::UpdatedInput(updated_input)) => { + return Err(FunctionCallError::UpdatedInput(updated_input)); + } Err(UnifiedExecError::SandboxDenied { output, .. }) => { let output_text = output.aggregated_output.text; let original_token_count = approx_token_count(&output_text); diff --git a/codex-rs/core/src/tools/network_approval.rs b/codex-rs/core/src/tools/network_approval.rs index 14af2c9c5f..76d66f140b 100644 --- a/codex-rs/core/src/tools/network_approval.rs +++ b/codex-rs/core/src/tools/network_approval.rs @@ -473,7 +473,19 @@ impl NetworkApprovalService { .await { match permission_request_decision { - PermissionRequestDecision::Allow => { + PermissionRequestDecision::Allow { + updated_input: Some(_), + } => { + pending.set_decision(PendingApprovalDecision::Deny).await; + let mut pending_approvals = self.pending_host_approvals.lock().await; + pending_approvals.remove(&key); + return NetworkDecision::deny( + "updatedInput is not supported for network approval requests", + ); + } + PermissionRequestDecision::Allow { + updated_input: None, + } => { pending .set_decision(PendingApprovalDecision::AllowOnce) .await; diff --git a/codex-rs/core/src/tools/orchestrator.rs b/codex-rs/core/src/tools/orchestrator.rs index dcb42c36c6..05e0a530d9 100644 --- a/codex-rs/core/src/tools/orchestrator.rs +++ b/codex-rs/core/src/tools/orchestrator.rs @@ -397,6 +397,7 @@ impl ToolOrchestrator { if evaluate_permission_request_hooks && let Some(permission_request) = tool.permission_request_payload(req) { + let permission_request_tool_input = permission_request.tool_input.clone(); match run_permission_request_hooks( approval_ctx.session, approval_ctx.turn, @@ -405,7 +406,24 @@ impl ToolOrchestrator { ) .await { - Some(PermissionRequestDecision::Allow) => { + Some(PermissionRequestDecision::Allow { + updated_input: Some(updated_input), + }) => { + if updated_input != permission_request_tool_input { + return Err(ToolError::UpdatedInput(updated_input)); + } + let decision = ReviewDecision::Approved; + otel.tool_decision( + &tool_ctx.tool_name, + &tool_ctx.call_id, + &decision, + ToolDecisionSource::Config, + ); + return Ok(decision); + } + Some(PermissionRequestDecision::Allow { + updated_input: None, + }) => { let decision = ReviewDecision::Approved; otel.tool_decision( &tool_ctx.tool_name, diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index e1027c9fa9..e51d30f856 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -78,6 +78,20 @@ 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 @@ -166,6 +180,12 @@ trait AnyToolHandler: Send + Sync { 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, @@ -189,6 +209,14 @@ where 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) } @@ -197,9 +225,25 @@ where invocation: ToolInvocation, ) -> BoxFuture<'a, Result> { Box::pin(async move { + let mut invocation = invocation; + let mut rewrites = 0; + let output = loop { + match self.handle(invocation.clone()).await { + Err(FunctionCallError::UpdatedInput(updated_input)) => { + rewrites += 1; + if rewrites > 8 { + return Err(FunctionCallError::RespondToModel( + "hook input rewrite limit exceeded".to_string(), + )); + } + invocation = + ToolHandler::with_updated_hook_input(self, invocation, updated_input)?; + } + result => break result?, + } + }; let call_id = invocation.call_id.clone(); let payload = invocation.payload.clone(); - let output = self.handle(invocation.clone()).await?; let post_tool_use_payload = ToolHandler::post_tool_use_payload(self, &invocation, &output); Ok(AnyToolResult { @@ -264,7 +308,7 @@ impl ToolRegistry { )] pub(crate) async fn dispatch_any( &self, - invocation: ToolInvocation, + mut invocation: ToolInvocation, ) -> Result { let tool_name = invocation.tool_name.clone(); let display_name = tool_name.display(); @@ -354,8 +398,8 @@ impl ToolRegistry { return Err(err); } - if let Some(pre_tool_use_payload) = handler.pre_tool_use_payload(&invocation) - && let Some(message) = run_pre_tool_use_hooks( + if let Some(pre_tool_use_payload) = handler.pre_tool_use_payload(&invocation) { + match run_pre_tool_use_hooks( &invocation.session, &invocation.turn, invocation.call_id.clone(), @@ -363,10 +407,21 @@ impl ToolRegistry { &pre_tool_use_payload.tool_input, ) .await - { - let err = FunctionCallError::RespondToModel(message); - dispatch_trace.record_failed(&err); - return Err(err); + { + crate::hook_runtime::PreToolUseHookResult::Blocked(message) => { + let err = FunctionCallError::RespondToModel(message); + dispatch_trace.record_failed(&err); + return Err(err); + } + crate::hook_runtime::PreToolUseHookResult::Continue { + updated_input: Some(updated_input), + } => { + invocation = handler.with_updated_hook_input(invocation, updated_input)?; + } + crate::hook_runtime::PreToolUseHookResult::Continue { + updated_input: None, + } => {} + } } let is_mutating = handler.is_mutating(&invocation).await; diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs index 41aa552d89..074b4fdd16 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -420,7 +420,21 @@ impl CoreShellActionProvider { ) .await { - Some(PermissionRequestDecision::Allow) => { + Some(PermissionRequestDecision::Allow { + updated_input: Some(_), + }) => { + return PromptDecision { + decision: ReviewDecision::Denied, + guardian_review_id: None, + rejection_message: Some( + "updatedInput is not supported for intercepted exec approvals" + .to_string(), + ), + }; + } + Some(PermissionRequestDecision::Allow { + updated_input: None, + }) => { return PromptDecision { decision: ReviewDecision::Approved, guardian_review_id: None, diff --git a/codex-rs/core/src/tools/sandboxing.rs b/codex-rs/core/src/tools/sandboxing.rs index 122cd00fad..f6581368cf 100644 --- a/codex-rs/core/src/tools/sandboxing.rs +++ b/codex-rs/core/src/tools/sandboxing.rs @@ -350,6 +350,7 @@ pub(crate) struct ToolCtx { #[derive(Debug)] pub(crate) enum ToolError { Rejected(String), + UpdatedInput(serde_json::Value), Codex(CodexErr), } diff --git a/codex-rs/core/src/unified_exec/errors.rs b/codex-rs/core/src/unified_exec/errors.rs index 0576e7d7cf..501503e9df 100644 --- a/codex-rs/core/src/unified_exec/errors.rs +++ b/codex-rs/core/src/unified_exec/errors.rs @@ -1,4 +1,5 @@ use codex_protocol::exec_output::ExecToolCallOutput; +use serde_json::Value; use thiserror::Error; #[derive(Debug, Error)] @@ -23,6 +24,8 @@ pub(crate) enum UnifiedExecError { message: String, output: ExecToolCallOutput, }, + #[error("tool input rewritten by hook")] + UpdatedInput(Value), } impl UnifiedExecError { diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index c67abc48d6..f1d21b80ea 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -221,6 +221,9 @@ async fn finish_deferred_network_approval_for_session( fn network_approval_error_message(err: ToolError) -> String { match err { ToolError::Rejected(message) => message, + ToolError::UpdatedInput(_) => { + "updatedInput is not supported for deferred network approvals".to_string() + } ToolError::Codex(err) => err.to_string(), } } @@ -1065,6 +1068,9 @@ impl UnifiedExecProcessManager { }; UnifiedExecError::sandbox_denied(message, output) } + ToolError::UpdatedInput(updated_input) => { + UnifiedExecError::UpdatedInput(updated_input) + } other => UnifiedExecError::create_process(format!("{other:?}")), }) } diff --git a/codex-rs/core/tests/suite/hooks.rs b/codex-rs/core/tests/suite/hooks.rs index daf3053adc..d4aa02ee80 100644 --- a/codex-rs/core/tests/suite/hooks.rs +++ b/codex-rs/core/tests/suite/hooks.rs @@ -268,6 +268,54 @@ elif mode == "exit_2": Ok(()) } +fn write_updating_pre_tool_use_hook( + home: &Path, + matcher: &str, + updated_input: &Value, +) -> Result<()> { + let script_path = home.join("pre_tool_use_hook.py"); + let log_path = home.join("pre_tool_use_hook_log.jsonl"); + let updated_input_json = + serde_json::to_string(updated_input).context("serialize updated pre tool input")?; + let script = format!( + r#"import json +from pathlib import Path +import sys + +payload = json.load(sys.stdin) + +with Path(r"{log_path}").open("a", encoding="utf-8") as handle: + handle.write(json.dumps(payload) + "\n") + +print(json.dumps({{ + "hookSpecificOutput": {{ + "hookEventName": "PreToolUse", + "permissionDecision": "allow", + "updatedInput": {updated_input_json} + }} +}})) +"#, + log_path = log_path.display(), + updated_input_json = updated_input_json, + ); + let hooks = serde_json::json!({ + "hooks": { + "PreToolUse": [{ + "matcher": matcher, + "hooks": [{ + "type": "command", + "command": format!("python3 {}", script_path.display()), + "statusMessage": "rewriting pre tool input", + }] + }] + } + }); + + fs::write(&script_path, script).context("write updating pre tool use hook script")?; + fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?; + Ok(()) +} + fn write_pre_tool_use_hook_toml( home: &Path, script_name: &str, @@ -409,6 +457,56 @@ elif mode == "exit_2": Ok(()) } +fn write_updating_permission_request_hook( + home: &Path, + matcher: &str, + updated_input: &Value, +) -> Result<()> { + let script_path = home.join("permission_request_hook.py"); + let log_path = home.join("permission_request_hook_log.jsonl"); + let updated_input_json = + serde_json::to_string(updated_input).context("serialize updated permission input")?; + let script = format!( + r#"import json +from pathlib import Path +import sys + +payload = json.load(sys.stdin) + +with Path(r"{log_path}").open("a", encoding="utf-8") as handle: + handle.write(json.dumps(payload) + "\n") + +print(json.dumps({{ + "hookSpecificOutput": {{ + "hookEventName": "PermissionRequest", + "decision": {{ + "behavior": "allow", + "updatedInput": {updated_input_json} + }} + }} +}})) +"#, + log_path = log_path.display(), + updated_input_json = updated_input_json, + ); + let hooks = serde_json::json!({ + "hooks": { + "PermissionRequest": [{ + "matcher": matcher, + "hooks": [{ + "type": "command", + "command": format!("python3 {}", script_path.display()), + "statusMessage": "rewriting permission request input", + }] + }] + } + }); + + fs::write(&script_path, script).context("write updating permission hook script")?; + fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?; + Ok(()) +} + fn install_allow_permission_request_hook(home: &Path) -> Result<()> { write_permission_request_hook( home, @@ -1374,6 +1472,95 @@ async fn permission_request_hook_allows_shell_command_without_user_approval() -> Ok(()) } +#[tokio::test] +async fn permission_request_hook_rewrites_shell_command_before_execution() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let call_id = "permissionrequest-shell-command-rewrite"; + let original_marker = std::env::temp_dir().join("permissionrequest-shell-command-original"); + let rewritten_marker = std::env::temp_dir().join("permissionrequest-shell-command-rewritten"); + let original_command = format!("rm -f {}", original_marker.display()); + let rewritten_command = format!("printf rewritten > {}", rewritten_marker.display()); + let args = serde_json::json!({ "command": original_command }); + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + core_test_support::responses::ev_function_call( + call_id, + "shell_command", + &serde_json::to_string(&args)?, + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "permission hook rewrote it"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + let updated_input = serde_json::json!({ "command": rewritten_command.clone() }); + let mut builder = test_codex() + .with_pre_build_hook(move |home| { + if let Err(error) = + write_updating_permission_request_hook(home, "^Bash$", &updated_input) + { + panic!("failed to write updating permission hook fixture: {error}"); + } + }) + .with_config(|config| { + config + .features + .enable(Feature::CodexHooks) + .expect("test config should allow feature update"); + }); + let test = builder.build(&server).await?; + + if original_marker.exists() { + fs::remove_file(&original_marker).context("remove stale original permission marker")?; + } + if rewritten_marker.exists() { + fs::remove_file(&rewritten_marker).context("remove stale rewritten permission marker")?; + } + fs::write(&original_marker, "seed").context("create original permission marker")?; + + test.submit_turn_with_approval_and_permission_profile( + "run the rewritten shell command after hook approval", + AskForApproval::OnRequest, + PermissionProfile::Disabled, + ) + .await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + requests[1].function_call_output(call_id); + assert!( + original_marker.exists(), + "original command should not execute after permission rewrite" + ); + assert_eq!( + fs::read_to_string(&rewritten_marker).context("read rewritten permission marker")?, + "rewritten" + ); + + let hook_inputs = read_permission_request_hook_inputs(test.codex_home_path())?; + assert_eq!(hook_inputs.len(), 1); + assert_permission_request_hook_input( + &hook_inputs[0], + "Bash", + &original_command, + /*description*/ None, + ); + assert!(hook_inputs[0].get("tool_use_id").is_none()); + + Ok(()) +} + #[tokio::test] async fn permission_request_hook_allows_apply_patch_with_write_alias() -> Result<()> { skip_if_no_network!(Ok(())); @@ -1451,6 +1638,101 @@ async fn permission_request_hook_allows_apply_patch_with_write_alias() -> Result Ok(()) } +#[tokio::test] +async fn permission_request_hook_rewrites_apply_patch_before_execution() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let call_id = "permissionrequest-apply-patch-rewrite"; + let original_file = "permission_request_apply_patch_original.txt"; + let rewritten_file = "permission_request_apply_patch_rewritten.txt"; + let original_path = format!("../{original_file}"); + let rewritten_path = format!("../{rewritten_file}"); + let original_patch = format!( + r#"*** Begin Patch +*** Add File: {original_path} ++original +*** End Patch"# + ); + let rewritten_patch = format!( + r#"*** Begin Patch +*** Add File: {rewritten_path} ++rewritten +*** End Patch"# + ); + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_apply_patch_function_call(call_id, &original_patch), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "permission hook rewrote apply_patch"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + let updated_input = serde_json::json!({ "command": rewritten_patch.clone() }); + let mut builder = test_codex() + .with_pre_build_hook(move |home| { + if let Err(error) = + write_updating_permission_request_hook(home, "^apply_patch$", &updated_input) + { + panic!("failed to write updating permission hook fixture: {error}"); + } + }) + .with_config(|config| { + config.include_apply_patch_tool = true; + config + .features + .enable(Feature::CodexHooks) + .expect("test config should allow feature update"); + }); + let test = builder.build(&server).await?; + + test.submit_turn_with_approval_and_permission_profile( + "apply the rewritten patch after hook approval", + AskForApproval::OnRequest, + restrictive_workspace_write_profile(), + ) + .await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + requests[1].function_call_output(call_id); + assert!( + !test.workspace_path(&original_path).exists(), + "original patch should not create its target file" + ); + assert_eq!( + fs::read_to_string(test.workspace_path(&rewritten_path)) + .context("read rewritten permission apply_patch file")?, + "rewritten\n" + ); + + let hook_inputs = read_permission_request_hook_inputs(test.codex_home_path())?; + assert_eq!(hook_inputs.len(), 2); + assert_permission_request_hook_input( + &hook_inputs[0], + "apply_patch", + &original_patch, + /*description*/ None, + ); + assert_permission_request_hook_input( + &hook_inputs[1], + "apply_patch", + &rewritten_patch, + /*description*/ None, + ); + + Ok(()) +} + #[tokio::test] async fn permission_request_hook_sees_raw_exec_command_input() -> Result<()> { skip_if_no_network!(Ok(())); @@ -1831,6 +2113,85 @@ async fn pre_tool_use_blocks_shell_command_before_execution() -> Result<()> { Ok(()) } +#[tokio::test] +async fn pre_tool_use_rewrites_shell_command_before_execution() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let call_id = "pretooluse-shell-command-rewrite"; + let original_marker = std::env::temp_dir().join("pretooluse-shell-command-original-marker"); + let rewritten_marker = std::env::temp_dir().join("pretooluse-shell-command-rewritten-marker"); + let original_command = format!("printf original > {}", original_marker.display()); + let rewritten_command = format!("printf rewritten > {}", rewritten_marker.display()); + let args = serde_json::json!({ "command": original_command }); + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + core_test_support::responses::ev_function_call( + call_id, + "shell_command", + &serde_json::to_string(&args)?, + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "hook rewrote it"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + let updated_input = serde_json::json!({ "command": rewritten_command }); + let mut builder = test_codex() + .with_pre_build_hook(move |home| { + if let Err(error) = write_updating_pre_tool_use_hook(home, "^Bash$", &updated_input) { + panic!("failed to write updating pre tool use hook fixture: {error}"); + } + }) + .with_config(|config| { + config + .features + .enable(Feature::CodexHooks) + .expect("test config should allow feature update"); + }); + let test = builder.build(&server).await?; + + if original_marker.exists() { + fs::remove_file(&original_marker).context("remove stale original pre tool marker")?; + } + if rewritten_marker.exists() { + fs::remove_file(&rewritten_marker).context("remove stale rewritten pre tool marker")?; + } + + test.submit_turn_with_permission_profile( + "run the rewritten shell command", + PermissionProfile::Disabled, + ) + .await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + requests[1].function_call_output(call_id); + assert!( + !original_marker.exists(), + "original shell command should not execute after rewrite" + ); + assert_eq!( + fs::read_to_string(&rewritten_marker).context("read rewritten pre tool marker")?, + "rewritten" + ); + + let hook_inputs = read_pre_tool_use_hook_inputs(test.codex_home_path())?; + assert_eq!(hook_inputs.len(), 1); + assert_eq!(hook_inputs[0]["tool_input"]["command"], original_command); + + Ok(()) +} + #[tokio::test] async fn plugin_pre_tool_use_blocks_shell_command_before_execution() -> Result<()> { skip_if_no_network!(Ok(())); @@ -2417,6 +2778,83 @@ async fn pre_tool_use_blocks_apply_patch_before_execution() -> Result<()> { Ok(()) } +#[tokio::test] +async fn pre_tool_use_rewrites_apply_patch_before_execution() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let call_id = "pretooluse-apply-patch-rewrite"; + let original_file = "pre_tool_use_apply_patch_original.txt"; + let rewritten_file = "pre_tool_use_apply_patch_rewritten.txt"; + let original_patch = format!( + r#"*** Begin Patch +*** Add File: {original_file} ++original +*** End Patch"# + ); + let rewritten_patch = format!( + r#"*** Begin Patch +*** Add File: {rewritten_file} ++rewritten +*** End Patch"# + ); + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_apply_patch_function_call(call_id, &original_patch), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "apply_patch rewritten"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + let updated_input = serde_json::json!({ "command": rewritten_patch }); + let mut builder = test_codex() + .with_pre_build_hook(move |home| { + if let Err(error) = + write_updating_pre_tool_use_hook(home, "^apply_patch$", &updated_input) + { + panic!("failed to write updating pre tool use hook fixture: {error}"); + } + }) + .with_config(|config| { + config.include_apply_patch_tool = true; + config + .features + .enable(Feature::CodexHooks) + .expect("test config should allow feature update"); + }); + let test = builder.build(&server).await?; + + test.submit_turn("apply the rewritten patch").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + requests[1].function_call_output(call_id); + assert!( + !test.workspace_path(original_file).exists(), + "original patch should not create its target file" + ); + assert_eq!( + fs::read_to_string(test.workspace_path(rewritten_file)) + .context("read rewritten apply_patch file")?, + "rewritten\n" + ); + + let hook_inputs = read_pre_tool_use_hook_inputs(test.codex_home_path())?; + assert_eq!(hook_inputs.len(), 1); + assert_eq!(hook_inputs[0]["tool_input"]["command"], original_patch); + + Ok(()) +} + #[tokio::test] async fn pre_tool_use_blocks_apply_patch_with_write_alias() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/tests/suite/hooks_mcp.rs b/codex-rs/core/tests/suite/hooks_mcp.rs index 2157630e02..e24862b3c1 100644 --- a/codex-rs/core/tests/suite/hooks_mcp.rs +++ b/codex-rs/core/tests/suite/hooks_mcp.rs @@ -74,6 +74,50 @@ print(json.dumps({{ Ok(()) } +fn write_updating_pre_tool_use_hook(home: &Path, updated_message: &str) -> Result<()> { + let script_path = home.join("pre_tool_use_hook.py"); + let log_path = home.join("pre_tool_use_hook_log.jsonl"); + let updated_message_json = + serde_json::to_string(updated_message).context("serialize updated MCP message")?; + let script = format!( + r#"import json +from pathlib import Path +import sys + +payload = json.load(sys.stdin) + +with Path(r"{log_path}").open("a", encoding="utf-8") as handle: + handle.write(json.dumps(payload) + "\n") + +print(json.dumps({{ + "hookSpecificOutput": {{ + "hookEventName": "PreToolUse", + "permissionDecision": "allow", + "updatedInput": {{ "message": {updated_message_json} }} + }} +}})) +"#, + log_path = log_path.display(), + updated_message_json = updated_message_json, + ); + let hooks = serde_json::json!({ + "hooks": { + "PreToolUse": [{ + "matcher": RMCP_HOOK_MATCHER, + "hooks": [{ + "type": "command", + "command": format!("python3 {}", script_path.display()), + "statusMessage": "rewriting MCP pre tool input", + }] + }] + } + }); + + fs::write(&script_path, script).context("write updating pre tool use hook script")?; + fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?; + Ok(()) +} + fn write_post_tool_use_hook(home: &Path, additional_context: &str) -> Result<()> { let script_path = home.join("post_tool_use_hook.py"); let log_path = home.join("post_tool_use_hook_log.jsonl"); @@ -251,6 +295,76 @@ async fn pre_tool_use_blocks_mcp_tool_before_execution() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pre_tool_use_rewrites_mcp_tool_before_execution() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let call_id = "pretooluse-rmcp-echo-rewrite"; + let rewritten_message = "rewritten mcp hook input"; + let arguments = json!({ "message": RMCP_ECHO_MESSAGE }).to_string(); + let call_mock = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_function_call_with_namespace(call_id, RMCP_NAMESPACE, "echo", &arguments), + ev_completed("resp-1"), + ]), + ) + .await; + let final_mock = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "mcp pre hook rewrote it"), + ev_completed("resp-2"), + ]), + ) + .await; + + let rmcp_test_server_bin = stdio_server_bin()?; + let test = test_codex() + .with_pre_build_hook(move |home| { + if let Err(error) = write_updating_pre_tool_use_hook(home, rewritten_message) { + panic!("failed to write MCP updating pre tool use hook fixture: {error}"); + } + }) + .with_config(move |config| { + enable_hooks_and_rmcp_server(config, rmcp_test_server_bin, AppToolApproval::Approve); + }) + .build(&server) + .await?; + + test.submit_turn("call the rmcp echo tool with the MCP pre hook rewrite") + .await?; + + let final_request = final_mock.single_request(); + let output_item = final_request.function_call_output(call_id); + let output = output_item + .get("output") + .and_then(Value::as_str) + .expect("MCP tool output string"); + assert!( + output.contains(&format!("ECHOING: {rewritten_message}")), + "MCP tool should execute the rewritten input", + ); + assert!( + !output.contains(RMCP_ECHO_MESSAGE), + "MCP tool should not execute the original input", + ); + + let hook_inputs = read_hook_inputs(test.codex_home_path(), "pre_tool_use_hook_log.jsonl")?; + assert_eq!(hook_inputs.len(), 1); + assert_eq!( + hook_inputs[0]["tool_input"], + json!({ "message": RMCP_ECHO_MESSAGE }), + ); + + call_mock.single_request(); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn post_tool_use_records_mcp_tool_payload_and_context() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/hooks/schema/generated/permission-request.command.output.schema.json b/codex-rs/hooks/schema/generated/permission-request.command.output.schema.json index 21d45382b0..5b46925568 100644 --- a/codex-rs/hooks/schema/generated/permission-request.command.output.schema.json +++ b/codex-rs/hooks/schema/generated/permission-request.command.output.schema.json @@ -37,7 +37,7 @@ }, "updatedInput": { "default": null, - "description": "Reserved for a future input-rewrite capability.\n\nPermissionRequest hooks currently fail closed if this field is present." + "description": "Replaces the entire tool input object when `behavior` is `allow`." }, "updatedPermissions": { "default": null, diff --git a/codex-rs/hooks/src/engine/output_parser.rs b/codex-rs/hooks/src/engine/output_parser.rs index 0a3a994e19..5d8f14882f 100644 --- a/codex-rs/hooks/src/engine/output_parser.rs +++ b/codex-rs/hooks/src/engine/output_parser.rs @@ -16,13 +16,18 @@ pub(crate) struct SessionStartOutput { pub(crate) struct PreToolUseOutput { pub universal: UniversalOutput, pub block_reason: Option, + pub updated_input: Option, pub invalid_reason: Option, } #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) enum PermissionRequestDecision { - Allow, - Deny { message: String }, + Allow { + updated_input: Option, + }, + Deny { + message: String, + }, } #[derive(Debug, Clone)] @@ -123,10 +128,23 @@ pub(crate) fn parse_pre_tool_use(stdout: &str) -> Option { } else { None }; + let updated_input = if invalid_reason.is_none() { + hook_specific_output.and_then(|output| { + matches!( + output.permission_decision, + Some(PreToolUsePermissionDecisionWire::Allow) + ) + .then(|| output.updated_input.clone()) + .flatten() + }) + } else { + None + }; Some(PreToolUseOutput { universal, block_reason, + updated_input, invalid_reason, }) } @@ -298,7 +316,9 @@ fn unsupported_permission_request_hook_specific_output( decision: Option<&PermissionRequestDecisionWire>, ) -> Option { let decision = decision?; - if decision.updated_input.is_some() { + if decision.updated_input.is_some() + && !matches!(decision.behavior, PermissionRequestBehaviorWire::Allow) + { Some("PermissionRequest hook returned unsupported updatedInput".to_string()) } else if decision.updated_permissions.is_some() { Some("PermissionRequest hook returned unsupported updatedPermissions".to_string()) @@ -313,7 +333,9 @@ fn permission_request_decision( decision: &PermissionRequestDecisionWire, ) -> PermissionRequestDecision { match decision.behavior { - PermissionRequestBehaviorWire::Allow => PermissionRequestDecision::Allow, + PermissionRequestBehaviorWire::Allow => PermissionRequestDecision::Allow { + updated_input: decision.updated_input.clone(), + }, PermissionRequestBehaviorWire::Deny => PermissionRequestDecision::Deny { message: decision .message @@ -337,8 +359,13 @@ fn unsupported_post_tool_use_hook_specific_output( fn unsupported_pre_tool_use_hook_specific_output( output: &crate::schema::PreToolUseHookSpecificOutputWire, ) -> Option { - if output.updated_input.is_some() { - Some("PreToolUse hook returned unsupported updatedInput".to_string()) + if output.updated_input.is_some() + && !matches!( + output.permission_decision, + Some(PreToolUsePermissionDecisionWire::Allow) + ) + { + Some("PreToolUse hook returned updatedInput without permissionDecision:allow".to_string()) } else if output .additional_context .as_deref() @@ -348,9 +375,7 @@ fn unsupported_pre_tool_use_hook_specific_output( Some("PreToolUse hook returned unsupported additionalContext".to_string()) } else { match output.permission_decision { - Some(PreToolUsePermissionDecisionWire::Allow) => { - Some("PreToolUse hook returned unsupported permissionDecision:allow".to_string()) - } + Some(PreToolUsePermissionDecisionWire::Allow) => None, Some(PreToolUsePermissionDecisionWire::Ask) => { Some("PreToolUse hook returned unsupported permissionDecision:ask".to_string()) } @@ -424,7 +449,7 @@ mod tests { use super::parse_permission_request; #[test] - fn permission_request_rejects_reserved_updated_input_field() { + fn permission_request_accepts_allow_updated_input_field() { let parsed = parse_permission_request( &json!({ "continue": true, @@ -440,6 +465,31 @@ mod tests { ) .expect("permission request hook output should parse"); + assert_eq!( + parsed.decision, + Some(super::PermissionRequestDecision::Allow { + updated_input: Some(json!({})), + }) + ); + } + + #[test] + fn permission_request_rejects_deny_updated_input_field() { + let parsed = parse_permission_request( + &json!({ + "continue": true, + "hookSpecificOutput": { + "hookEventName": "PermissionRequest", + "decision": { + "behavior": "deny", + "updatedInput": {} + } + } + }) + .to_string(), + ) + .expect("permission request hook output should parse"); + assert_eq!( parsed.invalid_reason, Some("PermissionRequest hook returned unsupported updatedInput".to_string()) diff --git a/codex-rs/hooks/src/events/permission_request.rs b/codex-rs/hooks/src/events/permission_request.rs index 2cebd0e002..133461a26f 100644 --- a/codex-rs/hooks/src/events/permission_request.rs +++ b/codex-rs/hooks/src/events/permission_request.rs @@ -47,7 +47,7 @@ pub struct PermissionRequestRequest { #[derive(Debug, Clone, PartialEq, Eq)] pub enum PermissionRequestDecision { - Allow, + Allow { updated_input: Option }, Deny { message: String }, } @@ -154,8 +154,10 @@ fn resolve_permission_request_decision<'a>( let mut resolved_allow = None; for decision in decisions { match decision { - PermissionRequestDecision::Allow => { - resolved_allow = Some(PermissionRequestDecision::Allow); + PermissionRequestDecision::Allow { updated_input } => { + resolved_allow = Some(PermissionRequestDecision::Allow { + updated_input: updated_input.clone(), + }); } PermissionRequestDecision::Deny { message } => { return Some(PermissionRequestDecision::Deny { @@ -219,8 +221,8 @@ fn parse_completed( }); } else if let Some(parsed_decision) = parsed.decision { match parsed_decision { - output_parser::PermissionRequestDecision::Allow => { - decision = Some(PermissionRequestDecision::Allow); + output_parser::PermissionRequestDecision::Allow { updated_input } => { + decision = Some(PermissionRequestDecision::Allow { updated_input }); } output_parser::PermissionRequestDecision::Deny { message } => { status = HookRunStatus::Blocked; @@ -294,7 +296,9 @@ mod tests { #[test] fn permission_request_deny_overrides_earlier_allow() { let decisions = [ - PermissionRequestDecision::Allow, + PermissionRequestDecision::Allow { + updated_input: None, + }, PermissionRequestDecision::Deny { message: "repo deny".to_string(), }, @@ -311,13 +315,38 @@ mod tests { #[test] fn permission_request_returns_allow_when_no_handler_denies() { let decisions = [ - PermissionRequestDecision::Allow, - PermissionRequestDecision::Allow, + PermissionRequestDecision::Allow { + updated_input: None, + }, + PermissionRequestDecision::Allow { + updated_input: None, + }, ]; assert_eq!( resolve_permission_request_decision(decisions.iter()), - Some(PermissionRequestDecision::Allow) + Some(PermissionRequestDecision::Allow { + updated_input: None, + }) + ); + } + + #[test] + fn permission_request_keeps_highest_precedence_allow_update() { + let decisions = [ + PermissionRequestDecision::Allow { + updated_input: Some(serde_json::json!({ "command": "first" })), + }, + PermissionRequestDecision::Allow { + updated_input: Some(serde_json::json!({ "command": "second" })), + }, + ]; + + assert_eq!( + resolve_permission_request_decision(decisions.iter()), + Some(PermissionRequestDecision::Allow { + updated_input: Some(serde_json::json!({ "command": "second" })), + }) ); } diff --git a/codex-rs/hooks/src/events/pre_tool_use.rs b/codex-rs/hooks/src/events/pre_tool_use.rs index 6fe1555229..c83abf913c 100644 --- a/codex-rs/hooks/src/events/pre_tool_use.rs +++ b/codex-rs/hooks/src/events/pre_tool_use.rs @@ -37,12 +37,14 @@ pub struct PreToolUseOutcome { pub hook_events: Vec, pub should_block: bool, pub block_reason: Option, + pub updated_input: Option, } #[derive(Debug, Default, PartialEq, Eq)] struct PreToolUseHandlerData { should_block: bool, block_reason: Option, + updated_input: Option, } pub(crate) fn preview( @@ -78,6 +80,7 @@ pub(crate) async fn run( hook_events: Vec::new(), should_block: false, block_reason: None, + updated_input: None, }; } @@ -108,6 +111,14 @@ pub(crate) async fn run( let block_reason = results .iter() .find_map(|result| result.data.block_reason.clone()); + let updated_input = if should_block { + None + } else { + results + .iter() + .rev() + .find_map(|result| result.data.updated_input.clone()) + }; PreToolUseOutcome { hook_events: results @@ -118,6 +129,7 @@ pub(crate) async fn run( .collect(), should_block, block_reason, + updated_input, } } @@ -151,6 +163,7 @@ fn parse_completed( let mut status = HookRunStatus::Completed; let mut should_block = false; let mut block_reason = None; + let mut updated_input = None; match run_result.error.as_deref() { Some(error) => { @@ -185,6 +198,8 @@ fn parse_completed( kind: HookOutputEntryKind::Feedback, text: reason, }); + } else { + updated_input = parsed.updated_input; } } else if trimmed_stdout.starts_with('{') || trimmed_stdout.starts_with('[') { status = HookRunStatus::Failed; @@ -238,6 +253,7 @@ fn parse_completed( data: PreToolUseHandlerData { should_block, block_reason, + updated_input, }, } } @@ -247,6 +263,7 @@ fn serialization_failure_outcome(hook_events: Vec) -> PreToo hook_events, should_block: false, block_reason: None, + updated_input: None, } } @@ -298,6 +315,7 @@ mod tests { PreToolUseHandlerData { should_block: true, block_reason: Some("do not run that".to_string()), + updated_input: None, } ); assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); @@ -310,6 +328,30 @@ mod tests { ); } + #[test] + fn permission_decision_allow_can_update_input() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"allow","updatedInput":{"command":"echo rewritten"}}}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PreToolUseHandlerData { + should_block: false, + block_reason: None, + updated_input: Some(serde_json::json!({ "command": "echo rewritten" })), + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!(parsed.completed.run.entries, vec![]); + } + #[test] fn deprecated_block_decision_blocks_processing() { let parsed = parse_completed( @@ -327,6 +369,7 @@ mod tests { PreToolUseHandlerData { should_block: true, block_reason: Some("do not run that".to_string()), + updated_input: None, } ); assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); @@ -356,6 +399,7 @@ mod tests { PreToolUseHandlerData { should_block: false, block_reason: None, + updated_input: None, } ); assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); @@ -381,6 +425,7 @@ mod tests { PreToolUseHandlerData { should_block: false, block_reason: None, + updated_input: None, } ); assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); @@ -410,6 +455,7 @@ mod tests { PreToolUseHandlerData { should_block: false, block_reason: None, + updated_input: None, } ); assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); @@ -435,6 +481,7 @@ mod tests { PreToolUseHandlerData { should_block: false, block_reason: None, + updated_input: None, } ); assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); @@ -454,6 +501,7 @@ mod tests { PreToolUseHandlerData { should_block: false, block_reason: None, + updated_input: None, } ); assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); @@ -479,6 +527,7 @@ mod tests { PreToolUseHandlerData { should_block: true, block_reason: Some("blocked by policy".to_string()), + updated_input: None, } ); assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); diff --git a/codex-rs/hooks/src/schema.rs b/codex-rs/hooks/src/schema.rs index d08cce6ee2..7365529415 100644 --- a/codex-rs/hooks/src/schema.rs +++ b/codex-rs/hooks/src/schema.rs @@ -138,9 +138,7 @@ pub(crate) struct PermissionRequestHookSpecificOutputWire { #[serde(deny_unknown_fields)] pub(crate) struct PermissionRequestDecisionWire { pub behavior: PermissionRequestBehaviorWire, - /// Reserved for a future input-rewrite capability. - /// - /// PermissionRequest hooks currently fail closed if this field is present. + /// Replaces the entire tool input object when `behavior` is `allow`. #[serde(default)] pub updated_input: Option, /// Reserved for a future permission-rewrite capability.