agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Support PreToolUse updatedInput rewrites (#20527)

Browse Source 
## 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.
This commit is contained in:
Abhinav
2026-05-11 22:27:24 -04:00
committed by GitHub
parent 17ed5ad0b0
commit d08906a944
22 changed files with 1021 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/tests/suite/hooks.rs
+391
View File
@@ -26,6 +26,7 @@ use core_test_support::managed_network_requirements_loader;
use core_test_support::responses::ev_apply_patch_custom_tool_call;
use core_test_support::responses::ev_assistant_message;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_custom_tool_call;
use core_test_support::responses::ev_function_call;
use core_test_support::responses::ev_message_item_added;
use core_test_support::responses::ev_output_text_delta;
@@ -307,6 +308,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,
@@ -2081,6 +2130,274 @@ async fn blocked_pre_tool_use_records_additional_context_for_shell_command() ->
!marker.exists(),
"blocked command should not create marker file"
);
Ok(())
}
#[derive(Clone, Copy)]
enum BashRewriteSurface {
ContainerExec,
ExecCommand,
LocalShell,
Shell,
ShellCommand,
}
impl BashRewriteSurface {
fn slug(self) -> &'static str {
match self {
BashRewriteSurface::ContainerExec => "container-exec",
BashRewriteSurface::ExecCommand => "exec-command",
BashRewriteSurface::LocalShell => "local-shell",
BashRewriteSurface::Shell => "shell",
BashRewriteSurface::ShellCommand => "shell-command",
}
}
fn tool_call(self, call_id: &str, command: &[String], command_text: &str) -> Result<Value> {
match self {
BashRewriteSurface::ContainerExec => Ok(ev_function_call(
call_id,
"container.exec",
&serde_json::to_string(&serde_json::json!({ "command": command }))?,
)),
BashRewriteSurface::ExecCommand => Ok(ev_function_call(
call_id,
"exec_command",
&serde_json::to_string(&serde_json::json!({ "cmd": command_text }))?,
)),
BashRewriteSurface::LocalShell => {
Ok(core_test_support::responses::ev_local_shell_call(
call_id,
"completed",
command.iter().map(String::as_str).collect(),
))
}
BashRewriteSurface::Shell => Ok(ev_function_call(
call_id,
"shell",
&serde_json::to_string(&serde_json::json!({ "command": command }))?,
)),
BashRewriteSurface::ShellCommand => Ok(ev_function_call(
call_id,
"shell_command",
&serde_json::to_string(&serde_json::json!({ "command": command_text }))?,
)),
}
}
fn original_command(self, marker: &Path) -> (Vec<String>, String) {
let command_text = format!("printf original > {}", marker.display());
match self {
BashRewriteSurface::ContainerExec
| BashRewriteSurface::LocalShell
| BashRewriteSurface::Shell => {
let command = vec!["/bin/sh".to_string(), "-c".to_string(), command_text];
let command_text = codex_shell_command::parse_command::shlex_join(&command);
(command, command_text)
}
BashRewriteSurface::ExecCommand | BashRewriteSurface::ShellCommand => {
(Vec::new(), command_text)
}
}
}
fn rewritten_command(self, marker: &Path) -> String {
let command_text = format!("printf rewritten > {}", marker.display());
match self {
BashRewriteSurface::ContainerExec
| BashRewriteSurface::LocalShell
| BashRewriteSurface::Shell => codex_shell_command::parse_command::shlex_join(&[
"/bin/sh".to_string(),
"-c".to_string(),
command_text,
]),
BashRewriteSurface::ExecCommand | BashRewriteSurface::ShellCommand => command_text,
}
}
fn configure(self, config: &mut Config) {
trust_discovered_hooks(config);
if matches!(self, BashRewriteSurface::ExecCommand) {
config.use_experimental_unified_exec_tool = true;
if let Err(error) = config.features.enable(Feature::UnifiedExec) {
panic!("test config should allow feature update: {error}");
}
}
}
}
async fn assert_pre_tool_use_rewrites_bash_surface(surface: BashRewriteSurface) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let slug = surface.slug();
let call_id = format!("pretooluse-{slug}-rewrite");
let original_marker = std::env::temp_dir().join(format!("pretooluse-{slug}-original-marker"));
let rewritten_marker = std::env::temp_dir().join(format!("pretooluse-{slug}-rewritten-marker"));
let (tool_command, original_command) = surface.original_command(&original_marker);
let rewritten_command = surface.rewritten_command(&rewritten_marker);
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-1"),
surface.tool_call(&call_id, &tool_command, &original_command)?,
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(move |config| surface.configure(config));
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(
&format!("run the rewritten {slug} 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 {slug} 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 pre_tool_use_rewrites_shell_before_execution() -> Result<()> {
assert_pre_tool_use_rewrites_bash_surface(BashRewriteSurface::Shell).await
}
#[tokio::test]
async fn pre_tool_use_rewrites_container_exec_before_execution() -> Result<()> {
assert_pre_tool_use_rewrites_bash_surface(BashRewriteSurface::ContainerExec).await
}
#[tokio::test]
async fn pre_tool_use_rewrites_local_shell_before_execution() -> Result<()> {
assert_pre_tool_use_rewrites_bash_surface(BashRewriteSurface::LocalShell).await
}
#[tokio::test]
async fn pre_tool_use_rewrites_shell_command_before_execution() -> Result<()> {
assert_pre_tool_use_rewrites_bash_surface(BashRewriteSurface::ShellCommand).await
}
#[tokio::test]
async fn pre_tool_use_rewrites_exec_command_before_execution() -> Result<()> {
assert_pre_tool_use_rewrites_bash_surface(BashRewriteSurface::ExecCommand).await
}
#[tokio::test]
async fn pre_tool_use_rewrites_code_mode_nested_exec_command_before_execution() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let call_id = "pretooluse-code-mode-rewrite";
let original_marker = std::env::temp_dir().join("pretooluse-code-mode-original-marker");
let rewritten_marker = std::env::temp_dir().join("pretooluse-code-mode-rewritten-marker");
let original_command = format!("printf original > {}", original_marker.display());
let rewritten_command = format!("printf rewritten > {}", rewritten_marker.display());
let original_command_json =
serde_json::to_string(&original_command).context("serialize original command")?;
let code = format!(
r#"
const output = await tools.exec_command({{ cmd: {original_command_json} }});
text(output.output);
"#
);
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-1"),
ev_custom_tool_call(call_id, "exec", &code),
ev_completed("resp-1"),
]),
sse(vec![
ev_response_created("resp-2"),
ev_assistant_message("msg-1", "hook rewrote the nested command"),
ev_completed("resp-2"),
]),
],
)
.await;
let updated_input = serde_json::json!({ "command": rewritten_command });
let mut builder = test_codex()
.with_model("test-gpt-5.1-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| {
let _ = config.features.enable(Feature::CodeMode);
trust_discovered_hooks(config);
});
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 from code mode",
PermissionProfile::Disabled,
)
.await?;
let requests = responses.requests();
assert_eq!(requests.len(), 2);
requests[1].custom_tool_call_output(call_id);
assert!(
!original_marker.exists(),
"original nested shell command should not execute after rewrite"
);
assert_eq!(
fs::read_to_string(&rewritten_marker)
.context("read rewritten code mode 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(())
}
@@ -2675,6 +2992,80 @@ 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_custom_tool_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;
trust_discovered_hooks(config);
});
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].custom_tool_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(()));
codex-rs/core/tests/suite/hooks_mcp.rs
+114
View File
@@ -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");
@@ -249,6 +293,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(()));