Let code-mode exec use explicit output budgets

This commit is contained in:
Ahmed Ibrahim
2026-05-19 20:09:36 -07:00
parent 8e69d9ca12
commit 1b36ab6622
3 changed files with 88 additions and 17 deletions
+10 -6
View File
@@ -358,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 {
@@ -383,7 +385,7 @@ impl ToolOutput for ExecCommandToolOutput {
session_id: self.process_id,
original_token_count: self.original_token_count,
output: match self.max_output_tokens {
Some(_) => self.truncated_output(),
Some(max_tokens) => self.truncated_output(max_tokens),
None => String::from_utf8_lossy(&self.raw_output).to_string(),
},
};
@@ -395,10 +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).min(self.truncation_policy.token_budget());
formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens))
}
@@ -425,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")
}
+28 -8
View File
@@ -263,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"
);
@@ -304,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"
);
@@ -316,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"
);
@@ -353,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"
);
@@ -362,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"
);
@@ -397,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!(
@@ -423,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(())
}
@@ -445,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
+50 -3
View File
@@ -47,7 +47,7 @@ use wiremock::MockServer;
const CODE_MODE_EXEC_LONG_OUTPUT_COMMAND: &str =
"printf '0123456789012345678901234567890123456789'";
const CODE_MODE_EXEC_RAW_OUTPUT_LEN: usize = 50_000;
const CODE_MODE_EXEC_OUTER_MAX_OUTPUT_TOKENS: usize = 20_000;
const CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS: usize = 20_000;
const CODE_MODE_EXEC_SMALL_TOOL_OUTPUT_LIMIT: usize = 50;
fn custom_tool_output_items(req: &ResponsesRequest, call_id: &str) -> Vec<Value> {
@@ -714,6 +714,53 @@ text(result.output);
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 command = code_mode_exec_repeated_output_command(CODE_MODE_EXEC_RAW_OUTPUT_LEN);
let code = format!(
r#"// @exec: {{"max_output_tokens": {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}}}
const result = await tools.exec_command({{
cmd: {command:?},
max_output_tokens: {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}
}});
text(result.output);
"#
);
let output = code_mode_exec_output_text(&code).await?;
assert_eq!(output.len(), CODE_MODE_EXEC_RAW_OUTPUT_LEN);
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 command = code_mode_exec_repeated_output_command(CODE_MODE_EXEC_RAW_OUTPUT_LEN);
let code = format!(
r#"// @exec: {{"max_output_tokens": {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}}}
const result = await tools.exec_command({{
cmd: {command:?},
max_output_tokens: {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}
}});
text(result.output);
"#
);
let output = code_mode_exec_output_text_with_config(&code, |config| {
config.tool_output_token_limit = Some(CODE_MODE_EXEC_SMALL_TOOL_OUTPUT_LIMIT);
})
.await?;
assert_eq!(output.len(), CODE_MODE_EXEC_RAW_OUTPUT_LEN);
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<()> {
@@ -721,7 +768,7 @@ async fn code_mode_exec_without_max_preserves_output_beyond_default() -> Result<
let command = code_mode_exec_repeated_output_command(CODE_MODE_EXEC_RAW_OUTPUT_LEN);
let code = format!(
r#"// @exec: {{"max_output_tokens": {CODE_MODE_EXEC_OUTER_MAX_OUTPUT_TOKENS}}}
r#"// @exec: {{"max_output_tokens": {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}}}
const result = await tools.exec_command({{
cmd: {command:?}
}});
@@ -742,7 +789,7 @@ async fn code_mode_exec_without_max_preserves_output_beyond_truncation_policy()
let command = code_mode_exec_repeated_output_command(CODE_MODE_EXEC_RAW_OUTPUT_LEN);
let code = format!(
r#"// @exec: {{"max_output_tokens": {CODE_MODE_EXEC_OUTER_MAX_OUTPUT_TOKENS}}}
r#"// @exec: {{"max_output_tokens": {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}}}
const result = await tools.exec_command({{
cmd: {command:?}
}});