diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 9a92854515..5f27a7fc30 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -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); diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index e1027c9fa9..84d10fe77d 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -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>, specs: Vec, diff --git a/codex-rs/core/tests/suite/hooks.rs b/codex-rs/core/tests/suite/hooks.rs index 695e908ed8..20a9ce2654 100644 --- a/codex-rs/core/tests/suite/hooks.rs +++ b/codex-rs/core/tests/suite/hooks.rs @@ -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<()> { diff --git a/codex-rs/core/tests/suite/hooks_mcp.rs b/codex-rs/core/tests/suite/hooks_mcp.rs index 2157630e02..42fb20c493 100644 --- a/codex-rs/core/tests/suite/hooks_mcp.rs +++ b/codex-rs/core/tests/suite/hooks_mcp.rs @@ -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> { 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::(output)?, + json!({ + "content": [], + "structuredContent": { + "echo": "generic replacement" + }, + "isError": false, + }) + ); + + Ok(()) +} diff --git a/codex-rs/hooks/schema/generated/post-tool-use.command.output.schema.json b/codex-rs/hooks/schema/generated/post-tool-use.command.output.schema.json index 43a2a4828e..46bacc22bb 100644 --- a/codex-rs/hooks/schema/generated/post-tool-use.command.output.schema.json +++ b/codex-rs/hooks/schema/generated/post-tool-use.command.output.schema.json @@ -31,6 +31,9 @@ }, "updatedMCPToolOutput": { "default": null + }, + "updatedToolOutput": { + "default": null } }, "required": [ diff --git a/codex-rs/hooks/src/engine/output_parser.rs b/codex-rs/hooks/src/engine/output_parser.rs index 0a3a994e19..2c4f3c4623 100644 --- a/codex-rs/hooks/src/engine/output_parser.rs +++ b/codex-rs/hooks/src/engine/output_parser.rs @@ -40,6 +40,8 @@ pub(crate) struct PostToolUseOutput { pub invalid_block_reason: Option, pub additional_context: Option, pub invalid_reason: Option, + pub updated_tool_output: Option, + pub updated_mcp_tool_output: Option, } #[derive(Debug, Clone)] @@ -157,11 +159,7 @@ pub(crate) fn parse_permission_request(stdout: &str) -> Option Option { 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 { } 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 { 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 { - 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 { diff --git a/codex-rs/hooks/src/events/post_tool_use.rs b/codex-rs/hooks/src/events/post_tool_use.rs index 63045ef425..7961a4261d 100644 --- a/codex-rs/hooks/src/events/post_tool_use.rs +++ b/codex-rs/hooks/src/events/post_tool_use.rs @@ -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, pub additional_contexts: Vec, pub feedback_message: Option, + pub updated_tool_output: Option, } #[derive(Debug, Default, PartialEq, Eq)] @@ -48,6 +50,8 @@ struct PostToolUseHandlerData { stop_reason: Option, additional_contexts_for_model: Vec, feedback_messages_for_model: Vec, + updated_tool_output: Option, + updated_mcp_tool_output: Option, } 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) -> PostTo stop_reason: None, additional_contexts: Vec::new(), feedback_message: None, + updated_tool_output: None, + } +} + +fn select_updated_tool_output( + results: &mut [dispatcher::ParsedHandler], + is_mcp_tool: bool, + original_tool_response: &Value, +) -> Option { + 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::::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, } } } diff --git a/codex-rs/hooks/src/schema.rs b/codex-rs/hooks/src/schema.rs index d08cce6ee2..93da4857cb 100644 --- a/codex-rs/hooks/src/schema.rs +++ b/codex-rs/hooks/src/schema.rs @@ -173,6 +173,8 @@ pub(crate) struct PostToolUseHookSpecificOutputWire { #[serde(default)] pub additional_context: Option, #[serde(default)] + pub updated_tool_output: Option, + #[serde(default)] #[serde(rename = "updatedMCPToolOutput")] pub updated_mcp_tool_output: Option, }