Support PreToolUse updatedInput rewrites (#20527)

## Why `PreToolUse` already exposes `updatedInput` in its hook output schema, but Codex currently rejects it instead of applying the rewrite. That leaves hook authors unable to make the documented pre-execution adjustment to a tool call before it runs. ## What - Accept `updatedInput` from `PreToolUse` hooks when paired with `permissionDecision: "allow"`. - Apply the rewritten input before dispatch so the tool executes the updated payload, not the original one. - Preserve the stable hook-facing compatibility shapes that participating tool handlers expose: - Bash-like tools (`shell`, `container.exec`, `local_shell`, `shell_command`, `exec_command`) use `{ "command": ... }`. - `apply_patch` exposes its patch body through the same command-shaped hook contract. - MCP tools expose their JSON argument object directly. - Keep each participating tool handler responsible for translating hook-facing `updatedInput` back into its concrete invocation shape. ## Verification Direct Bash-like rewrite coverage: - `pre_tool_use_rewrites_shell_before_execution` - `pre_tool_use_rewrites_container_exec_before_execution` - `pre_tool_use_rewrites_local_shell_before_execution` - `pre_tool_use_rewrites_shell_command_before_execution` - `pre_tool_use_rewrites_exec_command_before_execution` These cases assert that each supported Bash-like surface runs only the rewritten command while the hook still observes the original `{ "command": ... }` input. `pre_tool_use_rewrites_apply_patch_before_execution` - Model emits one patch. - Hook swaps in a different patch. - Asserts only the rewritten file is created, and the hook saw the original patch. `pre_tool_use_rewrites_code_mode_nested_exec_command_before_execution` - Model runs one nested shell command from code mode. - Hook rewrites it. - Asserts only the rewritten command runs, and the hook saw the original nested input. `pre_tool_use_rewrites_mcp_tool_before_execution` - Model calls the RMCP echo tool. - Hook rewrites the MCP arguments. - Asserts the MCP server receives and returns the rewritten message, not the original one.
2026-07-01 00:31:56 +08:00 · 2026-05-11 22:27:24 -04:00
parent 17ed5ad0b0
commit d08906a944
22 changed files with 1021 additions and 47 deletions
@@ -1,6 +1,7 @@
 use std::path::Path;

-use futures::future::join_all;
+use futures::StreamExt;
+use futures::stream::FuturesUnordered;

 use codex_protocol::protocol::HookCompletedEvent;
 use codex_protocol::protocol::HookEventName;
@@ -20,6 +21,7 @@ use crate::events::common::matches_matcher;
 pub(crate) struct ParsedHandler<T> {
    pub completed: HookCompletedEvent,
    pub data: T,
+    pub completion_order: usize,
 }

 pub(crate) fn select_handlers(
@@ -90,18 +92,25 @@ pub(crate) async fn execute_handlers<T>(
    turn_id: Option<String>,
    parse: fn(&ConfiguredHandler, CommandRunResult, Option<String>) -> ParsedHandler<T>,
 ) -> Vec<ParsedHandler<T>> {
-    let results = join_all(
-        handlers
-            .iter()
-            .map(|handler| run_command(shell, handler, &input_json, cwd)),
-    )
-    .await;
+    let mut pending = FuturesUnordered::new();
+    for (configured_order, handler) in handlers.into_iter().enumerate() {
+        let input_json = input_json.clone();
+        let turn_id = turn_id.clone();
+        pending.push(async move {
+            let result = run_command(shell, &handler, &input_json, cwd).await;
+            (configured_order, parse(&handler, result, turn_id))
+        });
+    }

-    handlers
-        .into_iter()
-        .zip(results)
-        .map(|(handler, result)| parse(&handler, result, turn_id.clone()))
-        .collect()
+    let mut completed = Vec::new();
+    let mut completion_order = 0;
+    while let Some((configured_order, mut parsed)) = pending.next().await {
+        parsed.completion_order = completion_order;
+        completion_order += 1;
+        completed.push((configured_order, parsed));
+    }
+    completed.sort_by_key(|(configured_order, _)| *configured_order);
+    completed.into_iter().map(|(_, parsed)| parsed).collect()
 }

 pub(crate) fn completed_summary(
@@ -17,6 +17,7 @@ pub(crate) struct PreToolUseOutput {
    pub universal: UniversalOutput,
    pub block_reason: Option<String>,
    pub additional_context: Option<String>,
+    pub updated_input: Option<serde_json::Value>,
    pub invalid_reason: Option<String>,
 }

@@ -139,11 +140,24 @@ pub(crate) fn parse_pre_tool_use(stdout: &str) -> Option<PreToolUseOutput> {
    } 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,
        additional_context,
+        updated_input,
        invalid_reason,
    })
 }
@@ -377,12 +391,19 @@ fn unsupported_post_tool_use_hook_specific_output(
 fn unsupported_pre_tool_use_hook_specific_output(
    output: &crate::schema::PreToolUseHookSpecificOutputWire,
 ) -> Option<String> {
-    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 {
        match output.permission_decision {
            Some(PreToolUsePermissionDecisionWire::Allow) => {
-                Some("PreToolUse hook returned unsupported permissionDecision:allow".to_string())
+                output.updated_input.is_none().then(|| {
+                    "PreToolUse hook returned unsupported permissionDecision:allow".to_string()
+                })
            }
            Some(PreToolUsePermissionDecisionWire::Ask) => {
                Some("PreToolUse hook returned unsupported permissionDecision:ask".to_string())
@@ -299,6 +299,7 @@ fn parse_pre_completed(
            should_stop,
            stop_reason,
        },
+        completion_order: 0,
    }
 }

@@ -401,6 +402,7 @@ fn parse_completed(
            should_stop,
            stop_reason,
        },
+        completion_order: 0,
    }
 }

@@ -281,6 +281,7 @@ fn parse_completed(
    dispatcher::ParsedHandler {
        completed,
        data: PermissionRequestHandlerData { decision },
+        completion_order: 0,
    }
 }

@@ -298,6 +298,7 @@ fn parse_completed(
            additional_contexts_for_model,
            feedback_messages_for_model,
        },
+        completion_order: 0,
    }
 }

@@ -38,6 +38,7 @@ pub struct PreToolUseOutcome {
    pub should_block: bool,
    pub block_reason: Option<String>,
    pub additional_contexts: Vec<String>,
+    pub updated_input: Option<Value>,
 }

 #[derive(Debug, Default, PartialEq, Eq)]
@@ -45,6 +46,7 @@ struct PreToolUseHandlerData {
    should_block: bool,
    block_reason: Option<String>,
    additional_contexts_for_model: Vec<String>,
+    updated_input: Option<Value>,
 }

 pub(crate) fn preview(
@@ -81,6 +83,7 @@ pub(crate) async fn run(
            should_block: false,
            block_reason: None,
            additional_contexts: Vec::new(),
+            updated_input: None,
        };
    }

@@ -116,6 +119,11 @@ pub(crate) async fn run(
            .iter()
            .map(|result| result.data.additional_contexts_for_model.as_slice()),
    );
+    let updated_input = if should_block {
+        None
+    } else {
+        latest_updated_input(&results)
+    };

    PreToolUseOutcome {
        hook_events: results
@@ -127,9 +135,30 @@ pub(crate) async fn run(
        should_block,
        block_reason,
        additional_contexts,
+        updated_input,
    }
 }

+/// Chooses the rewrite from the hook that actually finished last.
+///
+/// Hook results stay in configured order for stable reporting, but the
+/// `PreToolUse` contract resolves competing rewrites by completion order.
+fn latest_updated_input(
+    results: &[dispatcher::ParsedHandler<PreToolUseHandlerData>],
+) -> Option<Value> {
+    results
+        .iter()
+        .filter_map(|result| {
+            result
+                .data
+                .updated_input
+                .clone()
+                .map(|updated_input| (result.completion_order, updated_input))
+        })
+        .max_by_key(|(completion_order, _)| *completion_order)
+        .map(|(_, updated_input)| updated_input)
+}
+
 /// Serializes command stdin for a selected `PreToolUse` hook.
 ///
 /// Handler selection may include internal matcher aliases, but hook stdin keeps
@@ -161,6 +190,7 @@ fn parse_completed(
    let mut should_block = false;
    let mut block_reason = None;
    let mut additional_contexts_for_model = Vec::new();
+    let mut updated_input = None;

    match run_result.error.as_deref() {
        Some(error) => {
@@ -204,6 +234,9 @@ fn parse_completed(
                                text: reason,
                            });
                        }
+                        if !should_block {
+                            updated_input = parsed.updated_input;
+                        }
                    }
                } else if output_parser::looks_like_json(&run_result.stdout) {
                    status = HookRunStatus::Failed;
@@ -258,7 +291,9 @@ fn parse_completed(
            should_block,
            block_reason,
            additional_contexts_for_model,
+            updated_input,
        },
+        completion_order: 0,
    }
 }

@@ -268,6 +303,7 @@ fn serialization_failure_outcome(hook_events: Vec<HookCompletedEvent>) -> PreToo
        should_block: false,
        block_reason: None,
        additional_contexts: Vec::new(),
+        updated_input: None,
    }
 }

@@ -284,6 +320,7 @@ mod tests {

    use super::PreToolUseHandlerData;
    use super::command_input_json;
+    use super::latest_updated_input;
    use super::parse_completed;
    use super::preview;
    use crate::engine::ConfiguredHandler;
@@ -320,6 +357,7 @@ mod tests {
                should_block: true,
                block_reason: Some("do not run that".to_string()),
                additional_contexts_for_model: Vec::new(),
+                updated_input: None,
            }
        );
        assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked);
@@ -332,6 +370,91 @@ 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,
+                additional_contexts_for_model: Vec::new(),
+                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 last_completed_updated_input_wins() {
+        let mut later_configured = parse_completed(
+            &handler(),
+            run_result(
+                Some(0),
+                r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"allow","updatedInput":{"command":"echo configured later"}}}"#,
+                "",
+            ),
+            Some("turn-1".to_string()),
+        );
+        later_configured.completion_order = 0;
+        let mut earlier_configured = parse_completed(
+            &handler(),
+            run_result(
+                Some(0),
+                r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"allow","updatedInput":{"command":"echo finished later"}}}"#,
+                "",
+            ),
+            Some("turn-1".to_string()),
+        );
+        earlier_configured.completion_order = 1;
+
+        assert_eq!(
+            latest_updated_input(&[later_configured, earlier_configured]),
+            Some(serde_json::json!({ "command": "echo finished later" }))
+        );
+    }
+
+    #[test]
+    fn permission_decision_allow_without_updated_input_fails_open() {
+        let parsed = parse_completed(
+            &handler(),
+            run_result(
+                Some(0),
+                r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"allow"}}"#,
+                "",
+            ),
+            Some("turn-1".to_string()),
+        );
+
+        assert_eq!(
+            parsed.data,
+            PreToolUseHandlerData {
+                should_block: false,
+                block_reason: None,
+                additional_contexts_for_model: Vec::new(),
+                updated_input: None,
+            }
+        );
+        assert_eq!(parsed.completed.run.status, HookRunStatus::Failed);
+        assert_eq!(
+            parsed.completed.run.entries,
+            vec![HookOutputEntry {
+                kind: HookOutputEntryKind::Error,
+                text: "PreToolUse hook returned unsupported permissionDecision:allow".to_string(),
+            }]
+        );
+    }
+
    #[test]
    fn deprecated_block_decision_blocks_processing() {
        let parsed = parse_completed(
@@ -350,6 +473,7 @@ mod tests {
                should_block: true,
                block_reason: Some("do not run that".to_string()),
                additional_contexts_for_model: Vec::new(),
+                updated_input: None,
            }
        );
        assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked);
@@ -380,6 +504,7 @@ mod tests {
                should_block: true,
                block_reason: Some("do not run that".to_string()),
                additional_contexts_for_model: vec!["remember this".to_string()],
+                updated_input: None,
            }
        );
        assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked);
@@ -416,6 +541,7 @@ mod tests {
                should_block: false,
                block_reason: None,
                additional_contexts_for_model: Vec::new(),
+                updated_input: None,
            }
        );
        assert_eq!(parsed.completed.run.status, HookRunStatus::Failed);
@@ -442,6 +568,7 @@ mod tests {
                should_block: false,
                block_reason: None,
                additional_contexts_for_model: Vec::new(),
+                updated_input: None,
            }
        );
        assert_eq!(parsed.completed.run.status, HookRunStatus::Failed);
@@ -472,6 +599,7 @@ mod tests {
                should_block: true,
                block_reason: Some("do not run that".to_string()),
                additional_contexts_for_model: vec!["nope".to_string()],
+                updated_input: None,
            }
        );
        assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked);
@@ -504,6 +632,7 @@ mod tests {
                should_block: false,
                block_reason: None,
                additional_contexts_for_model: Vec::new(),
+                updated_input: None,
            }
        );
        assert_eq!(parsed.completed.run.status, HookRunStatus::Completed);
@@ -524,6 +653,7 @@ mod tests {
                should_block: false,
                block_reason: None,
                additional_contexts_for_model: Vec::new(),
+                updated_input: None,
            }
        );
        assert_eq!(parsed.completed.run.status, HookRunStatus::Failed);
@@ -550,6 +680,7 @@ mod tests {
                should_block: true,
                block_reason: Some("blocked by policy".to_string()),
                additional_contexts_for_model: Vec::new(),
+                updated_input: None,
            }
        );
        assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked);
@@ -234,6 +234,7 @@ fn parse_completed(
            stop_reason,
            additional_contexts_for_model,
        },
+        completion_order: 0,
    }
 }

@@ -259,6 +259,7 @@ fn parse_completed(
            block_reason,
            continuation_fragments,
        },
+        completion_order: 0,
    }
 }

@@ -255,6 +255,7 @@ fn parse_completed(
            stop_reason,
            additional_contexts_for_model,
        },
+        completion_order: 0,
    }
 }