chore: apply truncation policy to unified_exec (#19247)

we were not respecting turn's `truncation_policy` to clamp output tokens
for `unified_exec` and `write_stdin`.

this meant truncation was only being applied by `ContextManager` before
the output was stored in-memory (so it _was_ being truncated from
model-visible context), but the full output was persisted to rollout on
disk.

now we respect that `truncation_policy` and `ContextManager`-level
truncation remains a backup.

### Tests
added tests, tested locally.
This commit is contained in:
sayan-oai
2026-04-24 00:17:39 -07:00
committed by GitHub
parent ac8c9fc49c
commit e083b6c757
2 changed files with 189 additions and 4 deletions
@@ -25,6 +25,7 @@ use crate::unified_exec::UnifiedExecError;
use crate::unified_exec::UnifiedExecProcessManager;
use crate::unified_exec::WriteStdinRequest;
use crate::unified_exec::generate_chunk_id;
use crate::unified_exec::resolve_max_tokens;
use codex_features::Feature;
use codex_otel::SessionTelemetry;
use codex_otel::TOOL_CALL_UNIFIED_EXEC_METRIC;
@@ -33,6 +34,7 @@ use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::TerminalInteractionEvent;
use codex_shell_command::is_safe_command::is_known_safe_command;
use codex_tools::UnifiedExecShellMode;
use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::approx_token_count;
use serde::Deserialize;
use std::path::PathBuf;
@@ -89,6 +91,13 @@ 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())
}
impl ToolHandler for UnifiedExecHandler {
type Output = ExecCommandToolOutput;
@@ -231,6 +240,8 @@ impl ToolHandler for UnifiedExecHandler {
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);
@@ -311,7 +322,7 @@ impl ToolHandler for UnifiedExecHandler {
chunk_id: String::new(),
wall_time: std::time::Duration::ZERO,
raw_output: output.into_text().into_bytes(),
max_output_tokens: None,
max_output_tokens: Some(max_output_tokens),
process_id: None,
exit_code: None,
original_token_count: None,
@@ -327,7 +338,7 @@ impl ToolHandler for UnifiedExecHandler {
hook_command: hook_command.clone(),
process_id,
yield_time_ms,
max_output_tokens,
max_output_tokens: Some(max_output_tokens),
workdir,
network: context.turn.network.clone(),
tty,
@@ -352,7 +363,7 @@ impl ToolHandler for UnifiedExecHandler {
chunk_id: generate_chunk_id(),
wall_time: output.duration,
raw_output: output_text.into_bytes(),
max_output_tokens,
max_output_tokens: Some(max_output_tokens),
// Sandbox denial is terminal, so there is no live
// process for write_stdin to resume.
process_id: None,
@@ -370,12 +381,14 @@ impl ToolHandler for UnifiedExecHandler {
}
"write_stdin" => {
let args: WriteStdinArgs = parse_arguments(&arguments)?;
let max_output_tokens =
effective_max_output_tokens(args.max_output_tokens, turn.truncation_policy);
let response = manager
.write_stdin(WriteStdinRequest {
process_id: args.session_id,
input: &args.chars,
yield_time_ms: args.yield_time_ms,
max_output_tokens: args.max_output_tokens,
max_output_tokens: Some(max_output_tokens),
})
.await
.map_err(|err| {