[codex] Preserve raw code-mode exec output by default (#23564)

## Why
Code mode can use nested unified exec calls as data sources. When those
calls omit `max_output_tokens`, code mode should receive raw command
output so the script can parse or summarize it itself. When code mode
does provide `max_output_tokens`, that explicit nested budget should be
respected, including values above the default unified exec limit, rather
than being capped before code mode sees the result.

## What
- Preserve direct unified exec truncation behavior, while letting
code-mode exec/write_stdin keep `max_output_tokens` as `None` unless
explicitly supplied.
- Make code-mode tool results use raw output when no explicit limit is
present, and use the explicit nested limit directly when one is
specified.
- Refactor unified exec output formatting so `truncated_output` takes
the caller-selected token budget.
- Add e2e integration coverage for explicit nested exec limits, omitted
nested exec limits, outer exec limit propagation, omitted-limit outputs
that exceed both the default and a small truncation policy, explicit
nested limits above those caps, and high explicit limits that still
compact larger command output.
- Reuse the code-mode turn setup helper while directly asserting the
exact exec output item in each test.

## Testing
- `just fmt`
- `git diff --check`
- Not run locally per repo guidance; CI should validate the e2e
integration tests.
This commit is contained in:
Ahmed Ibrahim
2026-05-19 21:02:14 -07:00
committed by GitHub
Unverified
parent e43a2e297f
commit 5a4202ad90
10 changed files with 278 additions and 57 deletions
+14 -5
View File
@@ -311,6 +311,7 @@ pub struct ExecCommandToolOutput {
pub wall_time: Duration,
/// Raw bytes returned for this unified exec call before any truncation.
pub raw_output: Vec<u8>,
pub truncation_policy: TruncationPolicy,
pub max_output_tokens: Option<usize>,
pub process_id: Option<i32>,
pub exit_code: Option<i32>,
@@ -357,7 +358,9 @@ impl ToolOutput for ExecCommandToolOutput {
return None;
}
Some(JsonValue::String(self.truncated_output()))
Some(JsonValue::String(
self.truncated_output(self.model_output_max_tokens()),
))
}
fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue {
@@ -381,7 +384,10 @@ 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: match self.max_output_tokens {
Some(max_tokens) => self.truncated_output(max_tokens),
None => String::from_utf8_lossy(&self.raw_output).to_string(),
},
};
serde_json::to_value(result).unwrap_or_else(|err| {
@@ -391,9 +397,12 @@ impl ToolOutput for ExecCommandToolOutput {
}
impl ExecCommandToolOutput {
pub(crate) fn truncated_output(&self) -> String {
fn model_output_max_tokens(&self) -> usize {
resolve_max_tokens(self.max_output_tokens).min(self.truncation_policy.token_budget())
}
pub(crate) fn truncated_output(&self, max_tokens: usize) -> String {
let text = String::from_utf8_lossy(&self.raw_output).to_string();
let max_tokens = resolve_max_tokens(self.max_output_tokens);
formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens))
}
@@ -420,7 +429,7 @@ impl ExecCommandToolOutput {
}
sections.push("Output:".to_string());
sections.push(self.truncated_output());
sections.push(self.truncated_output(self.model_output_max_tokens()));
sections.join("\n")
}
+1
View File
@@ -429,6 +429,7 @@ fn exec_command_tool_output_formats_truncated_response() {
chunk_id: "abc123".to_string(),
wall_time: std::time::Duration::from_millis(1250),
raw_output: b"token one token two token three token four token five".to_vec(),
truncation_policy: TruncationPolicy::Tokens(10_000),
max_output_tokens: Some(4),
process_id: None,
exit_code: Some(0),
@@ -7,10 +7,8 @@ use crate::tools::context::ToolOutput;
use crate::tools::context::ToolPayload;
use crate::tools::hook_names::HookToolName;
use crate::tools::registry::PostToolUsePayload;
use crate::unified_exec::resolve_max_tokens;
use codex_protocol::models::AdditionalPermissionProfile;
use codex_tools::UnifiedExecShellMode;
use codex_utils_output_truncation::TruncationPolicy;
use serde::Deserialize;
use std::path::PathBuf;
use std::sync::Arc;
@@ -72,13 +70,6 @@ fn default_tty() -> bool {
false
}
fn effective_max_output_tokens(
max_output_tokens: Option<usize>,
truncation_policy: TruncationPolicy,
) -> usize {
resolve_max_tokens(max_output_tokens).min(truncation_policy.token_budget())
}
#[derive(Debug)]
pub(crate) struct ResolvedCommand {
pub(crate) command: Vec<String>,
@@ -36,7 +36,6 @@ 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::get_command;
use super::post_unified_exec_tool_use_payload;
@@ -162,8 +161,6 @@ impl ToolExecutor<ToolInvocation> for ExecCommandHandler {
prefix_rule,
..
} = args;
let max_output_tokens =
effective_max_output_tokens(max_output_tokens, turn.truncation_policy);
let exec_permission_approvals_enabled =
session.features().enabled(Feature::ExecPermissionApprovals);
@@ -241,7 +238,8 @@ 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),
truncation_policy: turn.truncation_policy,
max_output_tokens,
process_id: None,
exit_code: None,
original_token_count: None,
@@ -258,7 +256,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 +282,8 @@ 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),
truncation_policy: turn.truncation_policy,
max_output_tokens,
// Sandbox denial is terminal, so there is no live
// process for write_stdin to resume.
process_id: None,
@@ -14,7 +14,6 @@ use codex_tools::ToolSpec;
use serde::Deserialize;
use super::super::shell_spec::create_write_stdin_tool;
use super::effective_max_output_tokens;
use super::post_unified_exec_tool_use_payload;
#[derive(Debug, Deserialize)]
@@ -62,8 +61,6 @@ 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 response = session
.services
.unified_exec_manager
@@ -71,7 +68,8 @@ 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: args.max_output_tokens,
truncation_policy: turn.truncation_policy,
})
.await
.map_err(|err| {
@@ -4,6 +4,7 @@ use crate::shell::default_user_shell;
use codex_tools::UnifiedExecShellMode;
use codex_tools::ZshForkConfig;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_output_truncation::TruncationPolicy;
use pretty_assertions::assert_eq;
use std::sync::Arc;
@@ -17,6 +18,8 @@ use crate::tools::registry::CoreToolRuntime;
use crate::turn_diff_tracker::TurnDiffTracker;
use tokio::sync::Mutex;
const TEST_TRUNCATION_POLICY: TruncationPolicy = TruncationPolicy::Tokens(10_000);
async fn invocation_for_payload(
tool_name: &str,
call_id: &str,
@@ -258,6 +261,7 @@ async fn exec_command_post_tool_use_payload_uses_output_for_noninteractive_one_s
chunk_id: "chunk-1".to_string(),
wall_time: std::time::Duration::from_millis(498),
raw_output: b"three".to_vec(),
truncation_policy: TEST_TRUNCATION_POLICY,
max_output_tokens: None,
process_id: None,
exit_code: Some(0),
@@ -287,6 +291,7 @@ async fn exec_command_post_tool_use_payload_uses_output_for_interactive_completi
chunk_id: "chunk-1".to_string(),
wall_time: std::time::Duration::from_millis(498),
raw_output: b"three".to_vec(),
truncation_policy: TEST_TRUNCATION_POLICY,
max_output_tokens: None,
process_id: None,
exit_code: Some(0),
@@ -317,6 +322,7 @@ async fn exec_command_post_tool_use_payload_skips_running_sessions() {
chunk_id: "chunk-1".to_string(),
wall_time: std::time::Duration::from_millis(498),
raw_output: b"three".to_vec(),
truncation_policy: TEST_TRUNCATION_POLICY,
max_output_tokens: None,
process_id: Some(45),
exit_code: None,
@@ -342,6 +348,7 @@ async fn write_stdin_post_tool_use_payload_uses_original_exec_call_id_and_comman
chunk_id: "chunk-2".to_string(),
wall_time: std::time::Duration::from_millis(498),
raw_output: b"finished\n".to_vec(),
truncation_policy: TEST_TRUNCATION_POLICY,
max_output_tokens: None,
process_id: None,
exit_code: Some(0),
@@ -372,6 +379,7 @@ async fn write_stdin_post_tool_use_payload_keeps_parallel_session_metadata_separ
chunk_id: "chunk-a".to_string(),
wall_time: std::time::Duration::from_millis(498),
raw_output: b"alpha\n".to_vec(),
truncation_policy: TEST_TRUNCATION_POLICY,
max_output_tokens: None,
process_id: None,
exit_code: Some(0),
@@ -383,6 +391,7 @@ async fn write_stdin_post_tool_use_payload_keeps_parallel_session_metadata_separ
chunk_id: "chunk-b".to_string(),
wall_time: std::time::Duration::from_millis(498),
raw_output: b"beta\n".to_vec(),
truncation_policy: TEST_TRUNCATION_POLICY,
max_output_tokens: None,
process_id: None,
exit_code: Some(0),
+2
View File
@@ -31,6 +31,7 @@ use codex_exec_server::Environment;
use codex_network_proxy::NetworkProxy;
use codex_protocol::models::AdditionalPermissionProfile;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_output_truncation::TruncationPolicy;
use rand::Rng;
use rand::rng;
use tokio::sync::Mutex;
@@ -111,6 +112,7 @@ pub(crate) struct WriteStdinRequest<'a> {
pub input: &'a str,
pub yield_time_ms: u64,
pub max_output_tokens: Option<usize>,
pub truncation_policy: TruncationPolicy,
}
#[derive(Default)]
+31 -8
View File
@@ -10,6 +10,7 @@ use crate::tools::context::ExecCommandToolOutput;
use crate::unified_exec::WriteStdinRequest;
use crate::unified_exec::process::OutputHandles;
use codex_sandboxing::SandboxType;
use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::approx_token_count;
use core_test_support::get_remote_test_env;
use core_test_support::skip_if_sandbox;
@@ -162,6 +163,7 @@ async fn exec_command_with_tty(
chunk_id: generate_chunk_id(),
wall_time,
raw_output: collected,
truncation_policy: turn.truncation_policy,
max_output_tokens: None,
process_id: response_process_id,
exit_code,
@@ -195,6 +197,7 @@ async fn write_stdin(
input,
yield_time_ms,
max_output_tokens: None,
truncation_policy: TruncationPolicy::Tokens(10_000),
})
.await
}
@@ -260,7 +263,9 @@ async fn unified_exec_persists_across_requests() -> anyhow::Result<()> {
)
.await?;
assert!(
out_2.truncated_output().contains("codex"),
out_2
.truncated_output(DEFAULT_MAX_OUTPUT_TOKENS)
.contains("codex"),
"expected environment variable output"
);
@@ -301,7 +306,9 @@ async fn multi_unified_exec_sessions() -> anyhow::Result<()> {
"short command should not report a process id if it exits quickly"
);
assert!(
!out_2.truncated_output().contains("codex"),
!out_2
.truncated_output(DEFAULT_MAX_OUTPUT_TOKENS)
.contains("codex"),
"short command should run in a fresh shell"
);
@@ -313,7 +320,9 @@ async fn multi_unified_exec_sessions() -> anyhow::Result<()> {
)
.await?;
assert!(
out_3.truncated_output().contains("codex"),
out_3
.truncated_output(DEFAULT_MAX_OUTPUT_TOKENS)
.contains("codex"),
"session should preserve state"
);
@@ -350,7 +359,9 @@ async fn unified_exec_timeouts() -> anyhow::Result<()> {
)
.await?;
assert!(
!out_2.truncated_output().contains(TEST_VAR_VALUE),
!out_2
.truncated_output(DEFAULT_MAX_OUTPUT_TOKENS)
.contains(TEST_VAR_VALUE),
"timeout too short should yield incomplete output"
);
@@ -359,7 +370,9 @@ async fn unified_exec_timeouts() -> anyhow::Result<()> {
let out_3 = write_stdin(&session, process_id, "", /*yield_time_ms*/ 100).await?;
assert!(
out_3.truncated_output().contains(TEST_VAR_VALUE),
out_3
.truncated_output(DEFAULT_MAX_OUTPUT_TOKENS)
.contains(TEST_VAR_VALUE),
"subsequent poll should retrieve output"
);
@@ -394,7 +407,9 @@ async fn unified_exec_pause_blocks_yield_timeout() -> anyhow::Result<()> {
"pause should block the unified exec yield timeout"
);
assert!(
response.truncated_output().contains("unified-exec-done"),
response
.truncated_output(DEFAULT_MAX_OUTPUT_TOKENS)
.contains("unified-exec-done"),
"exec_command should wait for output after the pause lifts"
);
assert!(
@@ -420,7 +435,11 @@ async fn requests_with_large_timeout_are_capped() -> anyhow::Result<()> {
.await?;
assert!(result.process_id.is_some());
assert!(result.truncated_output().contains("codex"));
assert!(
result
.truncated_output(DEFAULT_MAX_OUTPUT_TOKENS)
.contains("codex")
);
Ok(())
}
@@ -442,7 +461,11 @@ async fn completed_commands_do_not_persist_sessions() -> anyhow::Result<()> {
result.process_id.is_some(),
"completed command should report a process id"
);
assert!(result.truncated_output().contains("codex"));
assert!(
result
.truncated_output(DEFAULT_MAX_OUTPUT_TOKENS)
.contains("codex")
);
assert!(
session
@@ -581,6 +581,7 @@ impl UnifiedExecProcessManager {
chunk_id,
wall_time,
raw_output: collected,
truncation_policy: context.turn.truncation_policy,
max_output_tokens: request.max_output_tokens,
process_id: response_process_id,
exit_code,
@@ -725,6 +726,7 @@ impl UnifiedExecProcessManager {
chunk_id,
wall_time,
raw_output: collected,
truncation_policy: request.truncation_policy,
max_output_tokens: request.max_output_tokens,
process_id,
exit_code,
+212 -25
View File
@@ -5,6 +5,7 @@ use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use codex_config::types::McpServerConfig;
use codex_config::types::McpServerTransportConfig;
use codex_core::config::Config;
use codex_features::Feature;
use codex_login::CodexAuth;
use codex_models_manager::bundled_models_response;
@@ -144,11 +145,21 @@ async fn run_code_mode_turn(
server: &MockServer,
prompt: &str,
code: &str,
) -> Result<(TestCodex, ResponseMock)> {
run_code_mode_turn_with_config(server, prompt, code, |_| {}).await
}
async fn run_code_mode_turn_with_config(
server: &MockServer,
prompt: &str,
code: &str,
configure: impl FnOnce(&mut Config) + Send + 'static,
) -> Result<(TestCodex, ResponseMock)> {
let mut builder = test_codex()
.with_model("test-gpt-5.1-codex")
.with_config(move |config| {
let _ = config.features.enable(Feature::CodeMode);
configure(config);
});
let test = builder.build(server).await?;
@@ -292,8 +303,7 @@ text(JSON.stringify(await tools.exec_command({ cmd: "printf code_mode_exec_marke
)
.await?;
let req = second_mock.single_request();
let items = custom_tool_output_items(&req, "call-1");
let items = custom_tool_output_items(&second_mock.single_request(), "call-1");
assert_eq!(items.len(), 2);
assert_regex_match(
concat!(
@@ -645,40 +655,217 @@ 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() -> 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
})));
"use exec_command from code mode",
r#"
const result = await tools.exec_command({
cmd: "printf '0123456789012345678901234567890123456789'",
max_output_tokens: 5
});
text(result.output);
"#,
)
.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"
assert_eq!(
text_item(
&custom_tool_output_items(&second_mock.single_request(), "call-1"),
/*index*/ 1
),
text_item(&items, /*index*/ 0),
"Total output lines: 1\n\n0123456789…5 tokens truncated…0123456789"
);
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_above_default_preserves_output() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (_test, second_mock) = run_code_mode_turn(
&server,
"use exec_command from code mode",
r#"// @exec: {"max_output_tokens": 20000}
const result = await tools.exec_command({
cmd: "python3 -c \"import sys; sys.stdout.write('x' * 50000)\"",
max_output_tokens: 20000
});
text(result.output);
"#,
)
.await?;
assert_eq!(
text_item(
&custom_tool_output_items(&second_mock.single_request(), "call-1"),
/*index*/ 1
),
"x".repeat(50_000)
);
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_above_default_truncates_larger_output() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (_test, second_mock) = run_code_mode_turn(
&server,
"use exec_command from code mode",
r#"// @exec: {"max_output_tokens": 25000}
const result = await tools.exec_command({
cmd: "python3 -c \"import sys; sys.stdout.write('A' * 90000)\"",
max_output_tokens: 20000
});
text(result.output);
"#,
)
.await?;
assert_eq!(
text_item(
&custom_tool_output_items(&second_mock.single_request(), "call-1"),
/*index*/ 1
),
format!(
"Total output lines: 1\n\n{}…2500 tokens truncated…{}",
"A".repeat(40_000),
"A".repeat(40_000)
)
);
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_above_truncation_policy_preserves_output() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (_test, second_mock) = run_code_mode_turn_with_config(
&server,
"use exec_command from code mode",
r#"// @exec: {"max_output_tokens": 20000}
const result = await tools.exec_command({
cmd: "python3 -c \"import sys; sys.stdout.write('x' * 50000)\"",
max_output_tokens: 20000
});
text(result.output);
"#,
|config| {
config.tool_output_token_limit = Some(50);
},
)
.await?;
assert_eq!(
text_item(
&custom_tool_output_items(&second_mock.single_request(), "call-1"),
/*index*/ 1
),
"x".repeat(50_000)
);
Ok(())
}
#[cfg_attr(windows, ignore = "no exec_command on Windows")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_exec_without_max_preserves_output_beyond_default() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (_test, second_mock) = run_code_mode_turn(
&server,
"use exec_command from code mode",
r#"// @exec: {"max_output_tokens": 20000}
const result = await tools.exec_command({
cmd: "python3 -c \"import sys; sys.stdout.write('x' * 50000)\""
});
text(result.output);
"#,
)
.await?;
assert_eq!(
text_item(
&custom_tool_output_items(&second_mock.single_request(), "call-1"),
/*index*/ 1
),
"x".repeat(50_000)
);
Ok(())
}
#[cfg_attr(windows, ignore = "no exec_command on Windows")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_exec_without_max_preserves_output_beyond_truncation_policy() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (_test, second_mock) = run_code_mode_turn_with_config(
&server,
"use exec_command from code mode",
r#"// @exec: {"max_output_tokens": 20000}
const result = await tools.exec_command({
cmd: "python3 -c \"import sys; sys.stdout.write('x' * 50000)\""
});
text(result.output);
"#,
|config| {
config.tool_output_token_limit = Some(50);
},
)
.await?;
assert_eq!(
text_item(
&custom_tool_output_items(&second_mock.single_request(), "call-1"),
/*index*/ 1
),
"x".repeat(50_000)
);
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() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (_test, second_mock) = run_code_mode_turn(
&server,
"use exec_command from code mode",
r#"// @exec: {"max_output_tokens": 5}
const result = await tools.exec_command({
cmd: "printf '0123456789012345678901234567890123456789'"
});
text(result.output);
"#,
)
.await?;
assert_eq!(
text_item(
&custom_tool_output_items(&second_mock.single_request(), "call-1"),
/*index*/ 1
),
"Total output lines: 1\n\n0123456789…5 tokens truncated…0123456789"
);
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));
Ok(())
}