Preserve raw code-mode exec output by default

This commit is contained in:
Ahmed Ibrahim
2026-05-19 15:10:56 -07:00
parent c2141c7ce0
commit 5c07818724
8 changed files with 135 additions and 39 deletions
+8 -1
View File
@@ -381,7 +381,7 @@ impl ToolOutput for ExecCommandToolOutput {
exit_code: self.exit_code,
session_id: self.process_id,
original_token_count: self.original_token_count,
output: self.truncated_output(),
output: self.code_mode_output(),
};
serde_json::to_value(result).unwrap_or_else(|err| {
@@ -391,6 +391,13 @@ impl ToolOutput for ExecCommandToolOutput {
}
impl ExecCommandToolOutput {
fn code_mode_output(&self) -> String {
match self.max_output_tokens {
Some(_) => self.truncated_output(),
None => String::from_utf8_lossy(&self.raw_output).to_string(),
}
}
pub(crate) fn truncated_output(&self) -> String {
let text = String::from_utf8_lossy(&self.raw_output).to_string();
let max_tokens = resolve_max_tokens(self.max_output_tokens);
@@ -2,6 +2,7 @@ use crate::sandboxing::SandboxPermissions;
use crate::shell::Shell;
use crate::shell::ShellType;
use crate::shell::get_shell_by_model_provided_path;
use crate::tools::context::ToolCallSource;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolOutput;
use crate::tools::context::ToolPayload;
@@ -79,6 +80,20 @@ fn effective_max_output_tokens(
resolve_max_tokens(max_output_tokens).min(truncation_policy.token_budget())
}
fn effective_max_output_tokens_for_source(
source: &ToolCallSource,
max_output_tokens: Option<usize>,
truncation_policy: TruncationPolicy,
) -> Option<usize> {
match source {
ToolCallSource::Direct => Some(effective_max_output_tokens(
max_output_tokens,
truncation_policy,
)),
ToolCallSource::CodeMode { .. } => max_output_tokens,
}
}
#[derive(Debug)]
pub(crate) struct ResolvedCommand {
pub(crate) command: Vec<String>,
@@ -36,7 +36,7 @@ use super::super::shell_spec::CommandToolOptions;
use super::super::shell_spec::create_exec_command_tool_with_environment_id;
use super::ExecCommandArgs;
use super::ExecCommandEnvironmentArgs;
use super::effective_max_output_tokens;
use super::effective_max_output_tokens_for_source;
use super::get_command;
use super::post_unified_exec_tool_use_payload;
@@ -98,6 +98,7 @@ impl ToolExecutor<ToolInvocation> for ExecCommandHandler {
turn,
tracker,
call_id,
source,
payload,
..
} = invocation;
@@ -162,8 +163,11 @@ impl ToolExecutor<ToolInvocation> for ExecCommandHandler {
prefix_rule,
..
} = args;
let max_output_tokens =
effective_max_output_tokens(max_output_tokens, turn.truncation_policy);
let max_output_tokens = effective_max_output_tokens_for_source(
&source,
max_output_tokens,
turn.truncation_policy,
);
let exec_permission_approvals_enabled =
session.features().enabled(Feature::ExecPermissionApprovals);
@@ -241,7 +245,7 @@ impl ToolExecutor<ToolInvocation> for ExecCommandHandler {
chunk_id: String::new(),
wall_time: std::time::Duration::ZERO,
raw_output: output.into_text().into_bytes(),
max_output_tokens: Some(max_output_tokens),
max_output_tokens,
process_id: None,
exit_code: None,
original_token_count: None,
@@ -258,7 +262,7 @@ impl ToolExecutor<ToolInvocation> for ExecCommandHandler {
hook_command: hook_command.clone(),
process_id,
yield_time_ms,
max_output_tokens: Some(max_output_tokens),
max_output_tokens,
cwd,
sandbox_cwd: turn_environment.cwd.clone(),
environment,
@@ -284,7 +288,7 @@ impl ToolExecutor<ToolInvocation> for ExecCommandHandler {
chunk_id: generate_chunk_id(),
wall_time: output.duration,
raw_output: output_text.into_bytes(),
max_output_tokens: Some(max_output_tokens),
max_output_tokens,
// Sandbox denial is terminal, so there is no live
// process for write_stdin to resume.
process_id: None,
@@ -14,7 +14,7 @@ use codex_tools::ToolSpec;
use serde::Deserialize;
use super::super::shell_spec::create_write_stdin_tool;
use super::effective_max_output_tokens;
use super::effective_max_output_tokens_for_source;
use super::post_unified_exec_tool_use_payload;
#[derive(Debug, Deserialize)]
@@ -48,6 +48,7 @@ impl ToolExecutor<ToolInvocation> for WriteStdinHandler {
let ToolInvocation {
session,
turn,
source,
payload,
..
} = invocation;
@@ -62,8 +63,11 @@ impl ToolExecutor<ToolInvocation> for WriteStdinHandler {
};
let args: WriteStdinArgs = parse_arguments(&arguments)?;
let max_output_tokens =
effective_max_output_tokens(args.max_output_tokens, turn.truncation_policy);
let max_output_tokens = effective_max_output_tokens_for_source(
&source,
args.max_output_tokens,
turn.truncation_policy,
);
let response = session
.services
.unified_exec_manager
@@ -71,7 +75,7 @@ impl ToolExecutor<ToolInvocation> for WriteStdinHandler {
process_id: args.session_id,
input: &args.chars,
yield_time_ms: args.yield_time_ms,
max_output_tokens: Some(max_output_tokens),
max_output_tokens,
})
.await
.map_err(|err| {
+75 -28
View File
@@ -43,6 +43,9 @@ use std::time::Duration;
use std::time::Instant;
use wiremock::MockServer;
const CODE_MODE_EXEC_LONG_OUTPUT_COMMAND: &str =
"printf '0123456789012345678901234567890123456789'";
fn custom_tool_output_items(req: &ResponsesRequest, call_id: &str) -> Vec<Value> {
match req.custom_tool_call_output(call_id).get("output") {
Some(Value::Array(items)) => items.clone(),
@@ -175,6 +178,20 @@ async fn run_code_mode_turn(
Ok((test, second_mock))
}
async fn assert_code_mode_exec_output_snapshot(snapshot_name: &str, code: &str) -> Result<()> {
let server = responses::start_mock_server().await;
let (_test, second_mock) =
run_code_mode_turn(&server, "use exec_command from code mode", code).await?;
let req = second_mock.single_request();
let items = custom_tool_output_items(&req, "call-1");
assert_eq!(items.len(), 2);
let output = text_item(&items, /*index*/ 1);
insta::assert_snapshot!(snapshot_name, output);
Ok(())
}
async fn run_code_mode_turn_with_rmcp(
server: &MockServer,
prompt: &str,
@@ -645,40 +662,70 @@ text(JSON.stringify(results));
#[cfg_attr(windows, ignore = "no exec_command on Windows")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_can_truncate_final_result_with_configured_budget() -> Result<()> {
async fn code_mode_exec_command_explicit_max_output_tokens_truncates_snapshot() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (_test, second_mock) = run_code_mode_turn(
&server,
"use exec to truncate the final result",
r#"// @exec: {"max_output_tokens": 6}
text(JSON.stringify(await tools.exec_command({
cmd: "printf 'token one token two token three token four token five token six token seven'",
max_output_tokens: 100
})));
"#,
let code = format!(
r#"
const result = await tools.exec_command({{
cmd: {CODE_MODE_EXEC_LONG_OUTPUT_COMMAND:?},
max_output_tokens: 5
}});
text(result.output);
"#
);
assert_code_mode_exec_output_snapshot(
"code_mode_exec_command_explicit_max_output_tokens_truncates",
&code,
)
.await?;
let req = second_mock.single_request();
let items = custom_tool_output_items(&req, "call-1");
assert_eq!(items.len(), 2);
assert_regex_match(
concat!(
r"(?s)\A",
r"Script completed\nWall time \d+\.\d seconds\nOutput:\n\z"
),
text_item(&items, /*index*/ 0),
Ok(())
}
#[cfg_attr(windows, ignore = "no exec_command on Windows")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_exec_command_without_max_output_tokens_uses_raw_output_snapshot() -> Result<()> {
skip_if_no_network!(Ok(()));
let code = format!(
r#"
const result = await tools.exec_command({{
cmd: {CODE_MODE_EXEC_LONG_OUTPUT_COMMAND:?}
}});
text(result.output);
"#
);
let expected_pattern = r#"(?sx)
\A
Total\ output\ lines:\ 1\n
\n
.*…\d+\ tokens\ truncated….*
\z
"#;
assert_regex_match(expected_pattern, text_item(&items, /*index*/ 1));
assert_code_mode_exec_output_snapshot(
"code_mode_exec_command_without_max_output_tokens_uses_raw_output",
&code,
)
.await?;
Ok(())
}
#[cfg_attr(windows, ignore = "no exec_command on Windows")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_exec_explicit_max_output_tokens_truncates_snapshot() -> Result<()> {
skip_if_no_network!(Ok(()));
let code = format!(
r#"// @exec: {{"max_output_tokens": 5}}
const result = await tools.exec_command({{
cmd: {CODE_MODE_EXEC_LONG_OUTPUT_COMMAND:?}
}});
text(result.output);
"#
);
assert_code_mode_exec_output_snapshot(
"code_mode_exec_explicit_max_output_tokens_truncates",
&code,
)
.await?;
Ok(())
}
@@ -0,0 +1,7 @@
---
source: core/tests/suite/code_mode.rs
expression: output
---
Total output lines: 1
0123456789…5 tokens truncated…0123456789
@@ -0,0 +1,5 @@
---
source: core/tests/suite/code_mode.rs
expression: output
---
0123456789012345678901234567890123456789
@@ -0,0 +1,7 @@
---
source: core/tests/suite/code_mode.rs
expression: output
---
Total output lines: 1
0123456789…5 tokens truncated…0123456789