mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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:
@@ -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| {
|
||||
|
||||
Reference in New Issue
Block a user