From 7562b347c012500b45af484b3fc1970855ab0a11 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Fri, 22 May 2026 11:33:00 -0700 Subject: [PATCH] Preserve code-mode results for post-tool rewrites --- codex-rs/core/src/tools/context.rs | 41 ++++------------------- codex-rs/core/src/tools/registry.rs | 4 +-- codex-rs/core/src/tools/registry_tests.rs | 18 ++++++---- codex-rs/core/tests/suite/code_mode.rs | 7 ++-- codex-rs/core/tests/suite/hooks.rs | 14 ++------ 5 files changed, 25 insertions(+), 59 deletions(-) diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 379e9a7502..e885b861eb 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -96,16 +96,6 @@ pub trait ToolOutput: Send { fn code_mode_result(&self, payload: &ToolPayload) -> JsonValue { response_input_to_code_mode_result(self.to_response_item("", payload)) } - - /// Returns the caller-visible value that code mode should receive after a - /// `PostToolUse` hook accepts `updatedToolOutput`. - fn rewritten_code_mode_result( - &self, - _payload: &ToolPayload, - updated_tool_output: &JsonValue, - ) -> JsonValue { - updated_tool_output.clone() - } } impl ToolOutput for CallToolResult { @@ -299,14 +289,14 @@ impl ToolOutput for FunctionToolOutput { } } -/// Preserves the original typed tool output while rewriting what model-authored -/// callers see. -pub(crate) struct CallerVisibleRewriteOutput { +/// Preserves the original typed tool output while rewriting the model-visible +/// response item. +pub(crate) struct ModelVisibleRewriteOutput { original_tool_output: Box, updated_tool_output: JsonValue, } -impl CallerVisibleRewriteOutput { +impl ModelVisibleRewriteOutput { pub(crate) fn new( original_tool_output: Box, updated_tool_output: JsonValue, @@ -318,7 +308,7 @@ impl CallerVisibleRewriteOutput { } } -impl ToolOutput for CallerVisibleRewriteOutput { +impl ToolOutput for ModelVisibleRewriteOutput { fn log_preview(&self) -> String { self.original_tool_output.log_preview() } @@ -339,8 +329,7 @@ impl ToolOutput for CallerVisibleRewriteOutput { } fn code_mode_result(&self, payload: &ToolPayload) -> JsonValue { - self.original_tool_output - .rewritten_code_mode_result(payload, &self.updated_tool_output) + self.original_tool_output.code_mode_result(payload) } } @@ -379,9 +368,6 @@ impl ToolOutput for ApplyPatchToolOutput { } fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue { - // Native code-mode callers only need the side effect here. If a - // `PostToolUse` hook installs a caller-visible rewrite, the wrapper - // returns that replacement instead of this empty native result. JsonValue::Object(serde_json::Map::new()) } } @@ -489,21 +475,6 @@ impl ToolOutput for ExecCommandToolOutput { JsonValue::String(format!("failed to serialize exec result: {err}")) }) } - - // Code mode consumes exec results as a structured envelope, so preserve - // metadata such as exit code and timing while replacing only the - // caller-visible command output. - fn rewritten_code_mode_result( - &self, - payload: &ToolPayload, - updated_tool_output: &JsonValue, - ) -> JsonValue { - let mut result = self.code_mode_result(payload); - if let JsonValue::Object(fields) = &mut result { - fields.insert("output".to_string(), updated_tool_output.clone()); - } - result - } } impl ExecCommandToolOutput { diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index f0fddc8be0..62d94d3eb3 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -12,8 +12,8 @@ use crate::memory_usage::emit_metric_for_tool_read; use crate::sandbox_tags::permission_profile_policy_tag; use crate::sandbox_tags::permission_profile_sandbox_tag; use crate::session::turn_context::TurnContext; -use crate::tools::context::CallerVisibleRewriteOutput; use crate::tools::context::FunctionToolOutput; +use crate::tools::context::ModelVisibleRewriteOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; @@ -519,7 +519,7 @@ impl ToolRegistry { } else if let Some(updated_tool_output) = &outcome.updated_tool_output { let mut guard = response_cell.lock().await; if let Some(mut result) = guard.take() { - result.result = Box::new(CallerVisibleRewriteOutput::new( + result.result = Box::new(ModelVisibleRewriteOutput::new( result.result, updated_tool_output.clone(), )); diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index 74cf852313..1d9808748e 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -1,6 +1,6 @@ use super::*; -use crate::tools::context::CallerVisibleRewriteOutput; use crate::tools::context::McpToolOutput; +use crate::tools::context::ModelVisibleRewriteOutput; use crate::tools::handlers::GetGoalHandler; use crate::tools::handlers::goal_spec::GET_GOAL_TOOL_NAME; use crate::tools::handlers::goal_spec::create_get_goal_tool; @@ -68,8 +68,8 @@ fn handler_looks_up_namespaced_aliases_explicitly() { } #[test] -fn caller_visible_rewrite_reaches_direct_and_code_mode_callers() { - let result = mcp_result_with_caller_visible_rewrite(); +fn model_visible_rewrite_preserves_code_mode_result() { + let result = mcp_result_with_model_visible_rewrite(); match result.into_response() { ResponseInputItem::FunctionCallOutput { call_id, output } => { @@ -83,20 +83,24 @@ fn caller_visible_rewrite_reaches_direct_and_code_mode_callers() { } assert_eq!( - mcp_result_with_caller_visible_rewrite().code_mode_result(), + mcp_result_with_model_visible_rewrite().code_mode_result(), json!({ - "echo": "rewritten", + "content": [], + "structuredContent": { + "echo": "original", + }, + "isError": false, }) ); } -fn mcp_result_with_caller_visible_rewrite() -> AnyToolResult { +fn mcp_result_with_model_visible_rewrite() -> AnyToolResult { AnyToolResult { call_id: "mcp-call-1".to_string(), payload: ToolPayload::Function { arguments: "{}".to_string(), }, - result: Box::new(CallerVisibleRewriteOutput::new( + result: Box::new(ModelVisibleRewriteOutput::new( Box::new(McpToolOutput { result: CallToolResult { content: Vec::new(), diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 2146416bc0..03d489ffdd 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -18,6 +18,7 @@ use codex_protocol::protocol::Op; use codex_protocol::user_input::UserInput; use core_test_support::apps_test_server::AppsTestServer; use core_test_support::assert_regex_match; +use core_test_support::hooks::trust_discovered_hooks; use core_test_support::responses; use core_test_support::responses::ResponseMock; use core_test_support::responses::ResponsesRequest; @@ -352,7 +353,7 @@ text(JSON.stringify(await tools.exec_command({ cmd: "printf code_mode_exec_marke #[cfg_attr(windows, ignore = "no exec_command on Windows")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn code_mode_post_tool_use_updated_tool_output_rewrites_exec_command_output() -> Result<()> { +async fn code_mode_post_tool_use_updated_tool_output_preserves_exec_command_result() -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -364,8 +365,8 @@ async fn code_mode_post_tool_use_updated_tool_output_rewrites_exec_command_outpu .expect("write post tool use hook"); }) .with_config(|config| { + trust_discovered_hooks(config); let _ = config.features.enable(Feature::CodeMode); - let _ = config.features.enable(Feature::CodexHooks); }); let test = builder.build(&server).await?; @@ -406,7 +407,7 @@ text(JSON.stringify(await tools.exec_command({ cmd: "printf original-output" })) let parsed: Value = serde_json::from_str(text_item(&items, /*index*/ 1))?; assert_eq!( parsed.get("output").and_then(Value::as_str), - Some(rewritten_output), + Some("original-output"), ); assert_eq!(parsed.get("exit_code").and_then(Value::as_i64), Some(0)); assert!(parsed.get("wall_time_seconds").is_some()); diff --git a/codex-rs/core/tests/suite/hooks.rs b/codex-rs/core/tests/suite/hooks.rs index 8bcd6e5dda..74c25cbbe2 100644 --- a/codex-rs/core/tests/suite/hooks.rs +++ b/codex-rs/core/tests/suite/hooks.rs @@ -3415,12 +3415,7 @@ async fn post_tool_use_updated_tool_output_replaces_shell_command_output() -> Re 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"); - }); + .with_config(trust_discovered_hooks); let test = builder.build(&server).await?; test.submit_turn("run the shell command with output rewrite") @@ -3485,12 +3480,7 @@ async fn post_tool_use_ignores_updated_tool_output_with_wrong_builtin_kind() -> 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"); - }); + .with_config(trust_discovered_hooks); let test = builder.build(&server).await?; test.submit_turn("run the shell command with invalid output rewrite")