clone/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-05-24 04:54:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

support PostToolUse output rewrites

Browse Source
This commit is contained in:
Abhinav Vedmala
2026-05-01 16:24:07 -07:00
parent d55479488e
commit fa92ad8b0e
8 changed files with 487 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
codex-rs/core/src/hook_runtime.rs
View File

@@ -231,6 +231,7 @@ pub(crate) async fn run_post_tool_use_hooks(
tool_input: Value,
tool_response: Value,
) -> PostToolUseOutcome {
let is_mcp_tool = tool_name.starts_with("mcp__");
let request = PostToolUseRequest {
session_id: sess.conversation_id,
turn_id: turn_context.sub_id.clone(),
@@ -243,6 +244,7 @@ pub(crate) async fn run_post_tool_use_hooks(
tool_use_id,
tool_input,
tool_response,
is_mcp_tool,
};
let hooks = sess.hooks();
let preview_runs = hooks.preview_post_tool_use(&request);

15
codex-rs/core/src/tools/registry.rs
View File

@@ -478,6 +478,14 @@ impl ToolRegistry {
/*success*/ None,
));
}
} else if let Some(updated_tool_output) = &outcome.updated_tool_output {
let mut guard = response_cell.lock().await;
if let Some(result) = guard.as_mut() {
result.result = Box::new(FunctionToolOutput::from_text(
post_tool_use_output_text(updated_tool_output),
Some(true),
));
}
}
}
@@ -514,6 +522,13 @@ impl ToolRegistry {
}
}
fn post_tool_use_output_text(output: &Value) -> String {
match output {
Value::String(text) => text.clone(),
_ => output.to_string(),
}
}
pub struct ToolRegistryBuilder {
handlers: HashMap<ToolName, Arc<dyn AnyToolHandler>>,
specs: Vec<ConfiguredToolSpec>,

151
codex-rs/core/tests/suite/hooks.rs
View File

@@ -459,6 +459,20 @@ elif mode == "continue_false":
"continue": False,
"stopReason": reason
}}))
elif mode == "updated_tool_output":
print(json.dumps({{
"hookSpecificOutput": {{
"hookEventName": "PostToolUse",
"updatedToolOutput": reason
}}
}}))
elif mode == "invalid_updated_tool_output":
print(json.dumps({{
"hookSpecificOutput": {{
"hookEventName": "PostToolUse",
"updatedToolOutput": {{"redacted": reason}}
}}
}}))
elif mode == "exit_2":
sys.stderr.write(reason + "\n")
raise SystemExit(2)
@@ -2729,6 +2743,143 @@ async fn post_tool_use_block_decision_replaces_shell_command_output_with_reason(
Ok(())
}
#[tokio::test]
async fn post_tool_use_updated_tool_output_replaces_shell_command_output() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let call_id = "posttooluse-shell-command-updated-output";
let command = "printf original-output".to_string();
let args = serde_json::json!({ "command": 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", "post hook output rewrite observed"),
ev_completed("resp-2"),
]),
],
)
.await;
let rewritten_output = "[redacted]";
let mut builder = test_codex()
.with_pre_build_hook(|home| {
if let Err(error) = write_post_tool_use_hook(
home,
Some("^Bash$"),
"updated_tool_output",
rewritten_output,
) {
panic!("failed to write post tool use hook test fixture: {error}");
}
})
.with_config(|config| {
config
.features
.enable(Feature::CodexHooks)
.expect("test config should allow feature update");
});
let test = builder.build(&server).await?;
test.submit_turn("run the shell command with output rewrite")
.await?;
let requests = responses.requests();
assert_eq!(requests.len(), 2);
let output_item = requests[1].function_call_output(call_id);
let output = output_item
.get("output")
.and_then(Value::as_str)
.expect("shell command output string");
assert_eq!(output, rewritten_output);
let hook_inputs = read_post_tool_use_hook_inputs(test.codex_home_path())?;
assert_eq!(hook_inputs.len(), 1);
assert_eq!(
hook_inputs[0]["tool_response"],
Value::String("original-output".to_string())
);
Ok(())
}
#[tokio::test]
async fn post_tool_use_ignores_updated_tool_output_with_wrong_builtin_kind() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let call_id = "posttooluse-shell-command-invalid-updated-output";
let command = "printf original-output".to_string();
let args = serde_json::json!({ "command": 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", "post hook invalid rewrite ignored"),
ev_completed("resp-2"),
]),
],
)
.await;
let mut builder = test_codex()
.with_pre_build_hook(|home| {
if let Err(error) = write_post_tool_use_hook(
home,
Some("^Bash$"),
"invalid_updated_tool_output",
"[redacted]",
) {
panic!("failed to write post tool use hook test fixture: {error}");
}
})
.with_config(|config| {
config
.features
.enable(Feature::CodexHooks)
.expect("test config should allow feature update");
});
let test = builder.build(&server).await?;
test.submit_turn("run the shell command with invalid output rewrite")
.await?;
let requests = responses.requests();
assert_eq!(requests.len(), 2);
let output_item = requests[1].function_call_output(call_id);
let output = output_item
.get("output")
.and_then(Value::as_str)
.expect("shell command output string");
assert!(
output.contains("original-output"),
"invalid replacement should preserve the original output, got {output:?}",
);
Ok(())
}
#[tokio::test]
async fn post_tool_use_continue_false_replaces_shell_command_output_with_stop_reason() -> Result<()>
{

103
codex-rs/core/tests/suite/hooks_mcp.rs
View File

@@ -117,6 +117,48 @@ print(json.dumps({{
Ok(())
}
fn write_post_tool_use_output_hook(home: &Path) -> Result<()> {
let script_path = home.join("post_tool_use_output_hook.py");
let script = r#"import json
print(json.dumps({
"hookSpecificOutput": {
"hookEventName": "PostToolUse",
"updatedToolOutput": {
"content": [],
"structuredContent": {
"echo": "generic replacement"
},
"isError": False
},
"updatedMCPToolOutput": {
"content": [],
"structuredContent": {
"echo": "mcp replacement"
},
"isError": False
}
}
}))
"#;
let hooks = serde_json::json!({
"hooks": {
"PostToolUse": [{
"matcher": RMCP_HOOK_MATCHER,
"hooks": [{
"type": "command",
"command": format!("python3 {}", script_path.display()),
"statusMessage": "rewriting MCP post tool use output",
}]
}]
}
});
fs::write(&script_path, script).context("write post tool use output hook script")?;
fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?;
Ok(())
}
fn read_hook_inputs(home: &Path, log_name: &str) -> Result<Vec<Value>> {
fs::read_to_string(home.join(log_name))
.with_context(|| format!("read {log_name}"))?
@@ -348,3 +390,64 @@ async fn post_tool_use_records_mcp_tool_payload_and_context() -> Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn post_tool_use_updated_tool_output_wins_for_mcp_tool() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let call_id = "posttooluse-rmcp-echo-output-rewrite";
let arguments = json!({ "message": RMCP_ECHO_MESSAGE }).to_string();
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-1"),
ev_function_call_with_namespace(call_id, RMCP_NAMESPACE, "echo", &arguments),
ev_completed("resp-1"),
]),
sse(vec![
ev_response_created("resp-2"),
ev_assistant_message("msg-1", "mcp post hook rewrite observed"),
ev_completed("resp-2"),
]),
],
)
.await;
let rmcp_test_server_bin = stdio_server_bin()?;
let test = test_codex()
.with_pre_build_hook(|home| {
if let Err(error) = write_post_tool_use_output_hook(home) {
panic!("failed to write MCP post tool use output 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 post hook rewrite")
.await?;
let requests = responses.requests();
assert_eq!(requests.len(), 2);
let output_item = requests[1].function_call_output(call_id);
let output = output_item
.get("output")
.and_then(Value::as_str)
.expect("MCP tool output string");
assert_eq!(
serde_json::from_str::<Value>(output)?,
json!({
"content": [],
"structuredContent": {
"echo": "generic replacement"
},
"isError": false,
})
);
Ok(())
}

3
codex-rs/hooks/schema/generated/post-tool-use.command.output.schema.json generated
View File

@@ -31,6 +31,9 @@
},
"updatedMCPToolOutput": {
"default": null
},
"updatedToolOutput": {
"default": null
}
},
"required": [

31
codex-rs/hooks/src/engine/output_parser.rs
View File

@@ -40,6 +40,8 @@ pub(crate) struct PostToolUseOutput {
pub invalid_block_reason: Option<String>,
pub additional_context: Option<String>,
pub invalid_reason: Option<String>,
pub updated_tool_output: Option<serde_json::Value>,
pub updated_mcp_tool_output: Option<serde_json::Value>,
}
#[derive(Debug, Clone)]
@@ -157,11 +159,7 @@ pub(crate) fn parse_permission_request(stdout: &str) -> Option<PermissionRequest
pub(crate) fn parse_post_tool_use(stdout: &str) -> Option<PostToolUseOutput> {
let wire: PostToolUseCommandOutputWire = parse_json(stdout)?;
let universal = UniversalOutput::from(wire.universal);
let invalid_reason = unsupported_post_tool_use_universal(&universal).or_else(|| {
wire.hook_specific_output
.as_ref()
.and_then(unsupported_post_tool_use_hook_specific_output)
});
let invalid_reason = unsupported_post_tool_use_universal(&universal);
let should_block = matches!(wire.decision, Some(BlockDecisionWire::Block));
let invalid_block_reason = if should_block
&& match wire.reason.as_deref() {
@@ -174,9 +172,16 @@ pub(crate) fn parse_post_tool_use(stdout: &str) -> Option<PostToolUseOutput> {
} else {
None
};
let additional_context = wire
let (additional_context, updated_tool_output, updated_mcp_tool_output) = wire
.hook_specific_output
.and_then(|output| output.additional_context);
.map(|output| {
(
output.additional_context,
output.updated_tool_output,
output.updated_mcp_tool_output,
)
})
.unwrap_or_default();
Some(PostToolUseOutput {
universal,
@@ -185,6 +190,8 @@ pub(crate) fn parse_post_tool_use(stdout: &str) -> Option<PostToolUseOutput> {
invalid_block_reason,
additional_context,
invalid_reason,
updated_tool_output,
updated_mcp_tool_output,
})
}
@@ -324,16 +331,6 @@ fn permission_request_decision(
}
}
fn unsupported_post_tool_use_hook_specific_output(
output: &crate::schema::PostToolUseHookSpecificOutputWire,
) -> Option<String> {
if output.updated_mcp_tool_output.is_some() {
Some("PostToolUse hook returned unsupported updatedMCPToolOutput".to_string())
} else {
None
}
}
fn unsupported_pre_tool_use_hook_specific_output(
output: &crate::schema::PreToolUseHookSpecificOutputWire,
) -> Option<String> {

210
codex-rs/hooks/src/events/post_tool_use.rs
View File

@@ -31,6 +31,7 @@ pub struct PostToolUseRequest {
pub tool_use_id: String,
pub tool_input: Value,
pub tool_response: Value,
pub is_mcp_tool: bool,
}
#[derive(Debug)]
@@ -40,6 +41,7 @@ pub struct PostToolUseOutcome {
pub stop_reason: Option<String>,
pub additional_contexts: Vec<String>,
pub feedback_message: Option<String>,
pub updated_tool_output: Option<Value>,
}
#[derive(Debug, Default, PartialEq, Eq)]
@@ -48,6 +50,8 @@ struct PostToolUseHandlerData {
stop_reason: Option<String>,
additional_contexts_for_model: Vec<String>,
feedback_messages_for_model: Vec<String>,
updated_tool_output: Option<Value>,
updated_mcp_tool_output: Option<Value>,
}
pub(crate) fn preview(
@@ -85,6 +89,7 @@ pub(crate) async fn run(
stop_reason: None,
additional_contexts: Vec::new(),
feedback_message: None,
updated_tool_output: None,
};
}
@@ -111,6 +116,9 @@ pub(crate) async fn run(
)
.await;
let mut results = results;
let updated_tool_output =
select_updated_tool_output(&mut results, request.is_mcp_tool, &request.tool_response);
let additional_contexts = common::flatten_additional_contexts(
results
.iter()
@@ -138,6 +146,7 @@ pub(crate) async fn run(
stop_reason,
additional_contexts,
feedback_message,
updated_tool_output,
}
}
@@ -174,6 +183,8 @@ fn parse_completed(
let mut stop_reason = None;
let mut additional_contexts_for_model = Vec::new();
let mut feedback_messages_for_model = Vec::new();
let mut updated_tool_output = None;
let mut updated_mcp_tool_output = None;
match run_result.error.as_deref() {
Some(error) => {
@@ -223,17 +234,17 @@ fn parse_completed(
.and_then(common::trimmed_non_empty)
.unwrap_or(stop_text);
feedback_messages_for_model.push(model_feedback);
} else if let Some(invalid_reason) = parsed.invalid_reason {
} else if let Some(invalid_reason) = &parsed.invalid_reason {
status = HookRunStatus::Failed;
entries.push(HookOutputEntry {
kind: HookOutputEntryKind::Error,
text: invalid_reason,
text: invalid_reason.clone(),
});
} else if let Some(invalid_block_reason) = parsed.invalid_block_reason {
} else if let Some(invalid_block_reason) = &parsed.invalid_block_reason {
status = HookRunStatus::Failed;
entries.push(HookOutputEntry {
kind: HookOutputEntryKind::Error,
text: invalid_block_reason,
text: invalid_block_reason.clone(),
});
} else if parsed.should_block {
status = HookRunStatus::Blocked;
@@ -245,6 +256,13 @@ fn parse_completed(
feedback_messages_for_model.push(reason);
}
}
let can_rewrite_output = parsed.universal.continue_processing
&& parsed.invalid_reason.is_none()
&& parsed.invalid_block_reason.is_none();
if can_rewrite_output {
updated_tool_output = parsed.updated_tool_output;
updated_mcp_tool_output = parsed.updated_mcp_tool_output;
}
} else if trimmed_stdout.starts_with('{') || trimmed_stdout.starts_with('[') {
status = HookRunStatus::Failed;
entries.push(HookOutputEntry {
@@ -297,6 +315,8 @@ fn parse_completed(
stop_reason,
additional_contexts_for_model,
feedback_messages_for_model,
updated_tool_output,
updated_mcp_tool_output,
},
}
}
@@ -308,6 +328,64 @@ fn serialization_failure_outcome(hook_events: Vec<HookCompletedEvent>) -> PostTo
stop_reason: None,
additional_contexts: Vec::new(),
feedback_message: None,
updated_tool_output: None,
}
}
fn select_updated_tool_output(
results: &mut [dispatcher::ParsedHandler<PostToolUseHandlerData>],
is_mcp_tool: bool,
original_tool_response: &Value,
) -> Option<Value> {
let mut selected = None;
for result in results {
let candidate = if let Some(updated_tool_output) = result.data.updated_tool_output.clone() {
Some(updated_tool_output)
} else if is_mcp_tool {
result.data.updated_mcp_tool_output.clone()
} else if result.data.updated_mcp_tool_output.is_some() {
result.completed.run.entries.push(HookOutputEntry {
kind: HookOutputEntryKind::Warning,
text: "ignored updatedMCPToolOutput for non-MCP tool".to_string(),
});
None
} else {
None
};
let Some(candidate) = candidate else {
continue;
};
if is_mcp_tool || json_kind_matches(original_tool_response, &candidate) {
selected = Some(candidate);
} else {
result.completed.run.entries.push(HookOutputEntry {
kind: HookOutputEntryKind::Warning,
text: format!(
"ignored updatedToolOutput: expected {} to match tool_response shape",
json_kind_name(original_tool_response)
),
});
}
}
selected
}
fn json_kind_matches(left: &Value, right: &Value) -> bool {
json_kind_name(left) == json_kind_name(right)
}
fn json_kind_name(value: &Value) -> &'static str {
match value {
Value::Null => "null",
Value::Bool(_) => "boolean",
Value::Number(_) => "number",
Value::String(_) => "string",
Value::Array(_) => "array",
Value::Object(_) => "object",
}
}
@@ -362,6 +440,8 @@ mod tests {
stop_reason: None,
additional_contexts_for_model: Vec::new(),
feedback_messages_for_model: vec!["bash output looked sketchy".to_string()],
updated_tool_output: None,
updated_mcp_tool_output: None,
}
);
assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked);
@@ -386,6 +466,8 @@ mod tests {
stop_reason: None,
additional_contexts_for_model: vec!["Remember the bash cleanup note.".to_string()],
feedback_messages_for_model: Vec::new(),
updated_tool_output: None,
updated_mcp_tool_output: None,
}
);
assert_eq!(
@@ -398,7 +480,7 @@ mod tests {
}
#[test]
fn unsupported_updated_mcp_tool_output_fails_open() {
fn updated_mcp_tool_output_is_recorded() {
let parsed = parse_completed(
&handler(),
run_result(
@@ -416,16 +498,12 @@ mod tests {
stop_reason: None,
additional_contexts_for_model: Vec::new(),
feedback_messages_for_model: Vec::new(),
updated_tool_output: None,
updated_mcp_tool_output: Some(json!({"ok": true})),
}
);
assert_eq!(parsed.completed.run.status, HookRunStatus::Failed);
assert_eq!(
parsed.completed.run.entries,
vec![HookOutputEntry {
kind: HookOutputEntryKind::Error,
text: "PostToolUse hook returned unsupported updatedMCPToolOutput".to_string(),
}]
);
assert_eq!(parsed.completed.run.status, HookRunStatus::Completed);
assert_eq!(parsed.completed.run.entries, Vec::<HookOutputEntry>::new());
}
#[test]
@@ -443,6 +521,8 @@ mod tests {
stop_reason: None,
additional_contexts_for_model: Vec::new(),
feedback_messages_for_model: vec!["post hook says pause".to_string()],
updated_tool_output: None,
updated_mcp_tool_output: None,
}
);
assert_eq!(parsed.completed.run.status, HookRunStatus::Completed);
@@ -467,6 +547,8 @@ mod tests {
stop_reason: Some("halt after bash output".to_string()),
additional_contexts_for_model: Vec::new(),
feedback_messages_for_model: vec!["post-tool hook says stop".to_string()],
updated_tool_output: None,
updated_mcp_tool_output: None,
}
);
assert_eq!(parsed.completed.run.status, HookRunStatus::Stopped);
@@ -494,6 +576,8 @@ mod tests {
stop_reason: None,
additional_contexts_for_model: Vec::new(),
feedback_messages_for_model: Vec::new(),
updated_tool_output: None,
updated_mcp_tool_output: None,
}
);
assert_eq!(parsed.completed.run.status, HookRunStatus::Completed);
@@ -540,6 +624,105 @@ mod tests {
assert_eq!(completed[0].run.id, runs[0].id);
}
#[test]
fn select_updated_tool_output_prefers_last_valid_replacement() {
let mut results = vec![
parse_completed(
&handler(),
run_result(
Some(0),
r#"{"hookSpecificOutput":{"hookEventName":"PostToolUse","updatedToolOutput":"first"}}"#,
"",
),
Some("turn-1".to_string()),
),
parse_completed(
&handler(),
run_result(
Some(0),
r#"{"hookSpecificOutput":{"hookEventName":"PostToolUse","updatedToolOutput":"second"}}"#,
"",
),
Some("turn-1".to_string()),
),
];
assert_eq!(
super::select_updated_tool_output(
&mut results,
/*is_mcp_tool*/ false,
&json!("old")
),
Some(json!("second"))
);
}
#[test]
fn select_updated_tool_output_ignores_wrong_builtin_json_kind() {
let mut results = vec![parse_completed(
&handler(),
run_result(
Some(0),
r#"{"hookSpecificOutput":{"hookEventName":"PostToolUse","updatedToolOutput":{"ok":true}}}"#,
"",
),
Some("turn-1".to_string()),
)];
assert_eq!(
super::select_updated_tool_output(
&mut results,
/*is_mcp_tool*/ false,
&json!("old")
),
None
);
assert_eq!(
results[0].completed.run.entries,
vec![HookOutputEntry {
kind: HookOutputEntryKind::Warning,
text: "ignored updatedToolOutput: expected string to match tool_response shape"
.to_string(),
}]
);
}
#[test]
fn select_updated_tool_output_prefers_generic_mcp_replacement() {
let mut results = vec![parse_completed(
&handler(),
run_result(
Some(0),
r#"{"hookSpecificOutput":{"hookEventName":"PostToolUse","updatedToolOutput":{"source":"generic"},"updatedMCPToolOutput":{"source":"mcp"}}}"#,
"",
),
Some("turn-1".to_string()),
)];
assert_eq!(
super::select_updated_tool_output(&mut results, /*is_mcp_tool*/ true, &json!({})),
Some(json!({"source": "generic"}))
);
}
#[test]
fn select_updated_tool_output_accepts_mcp_fallback_replacement() {
let mut results = vec![parse_completed(
&handler(),
run_result(
Some(0),
r#"{"hookSpecificOutput":{"hookEventName":"PostToolUse","updatedMCPToolOutput":{"source":"mcp"}}}"#,
"",
),
Some("turn-1".to_string()),
)];
assert_eq!(
super::select_updated_tool_output(&mut results, /*is_mcp_tool*/ true, &json!({})),
Some(json!({"source": "mcp"}))
);
}
fn handler() -> ConfiguredHandler {
ConfiguredHandler {
event_name: HookEventName::PostToolUse,
@@ -579,6 +762,7 @@ mod tests {
tool_use_id: tool_use_id.to_string(),
tool_input: json!({ "command": "echo hello" }),
tool_response: json!({"ok": true}),
is_mcp_tool: false,
}
}
}

2
codex-rs/hooks/src/schema.rs
View File

@@ -173,6 +173,8 @@ pub(crate) struct PostToolUseHookSpecificOutputWire {
#[serde(default)]
pub additional_context: Option<String>,
#[serde(default)]
pub updated_tool_output: Option<Value>,
#[serde(default)]
#[serde(rename = "updatedMCPToolOutput")]
pub updated_mcp_tool_output: Option<Value>,
}