diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 80ec343159..46df387891 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -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, + pub truncation_policy: TruncationPolicy, pub max_output_tokens: Option, pub process_id: Option, pub exit_code: Option, @@ -381,7 +382,10 @@ impl ToolOutput for ExecCommandToolOutput { exit_code: self.exit_code, session_id: self.process_id, original_token_count: self.original_token_count, - output: self.code_mode_output(), + output: match self.max_output_tokens { + Some(_) => self.truncated_output(), + None => String::from_utf8_lossy(&self.raw_output).to_string(), + }, }; serde_json::to_value(result).unwrap_or_else(|err| { @@ -391,16 +395,10 @@ 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); + let max_tokens = + resolve_max_tokens(self.max_output_tokens).min(self.truncation_policy.token_budget()); formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens)) } diff --git a/codex-rs/core/src/tools/context_tests.rs b/codex-rs/core/src/tools/context_tests.rs index fe40341e95..6ac5b6978f 100644 --- a/codex-rs/core/src/tools/context_tests.rs +++ b/codex-rs/core/src/tools/context_tests.rs @@ -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), diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index d44dc70c6f..170c1bd72f 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -2,16 +2,13 @@ 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; 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; @@ -73,27 +70,6 @@ fn default_tty() -> bool { false } -fn effective_max_output_tokens( - max_output_tokens: Option, - truncation_policy: TruncationPolicy, -) -> usize { - resolve_max_tokens(max_output_tokens).min(truncation_policy.token_budget()) -} - -fn effective_max_output_tokens_for_source( - source: &ToolCallSource, - max_output_tokens: Option, - truncation_policy: TruncationPolicy, -) -> Option { - 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, diff --git a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs index feab190996..b5813ebfe5 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs @@ -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_for_source; use super::get_command; use super::post_unified_exec_tool_use_payload; @@ -98,7 +97,6 @@ impl ToolExecutor for ExecCommandHandler { turn, tracker, call_id, - source, payload, .. } = invocation; @@ -163,11 +161,6 @@ impl ToolExecutor for ExecCommandHandler { prefix_rule, .. } = args; - 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); @@ -245,6 +238,7 @@ impl ToolExecutor for ExecCommandHandler { chunk_id: String::new(), wall_time: std::time::Duration::ZERO, raw_output: output.into_text().into_bytes(), + truncation_policy: turn.truncation_policy, max_output_tokens, process_id: None, exit_code: None, @@ -263,6 +257,7 @@ impl ToolExecutor for ExecCommandHandler { process_id, yield_time_ms, max_output_tokens, + truncation_policy: turn.truncation_policy, cwd, sandbox_cwd: turn_environment.cwd.clone(), environment, @@ -288,6 +283,7 @@ impl ToolExecutor for ExecCommandHandler { chunk_id: generate_chunk_id(), wall_time: output.duration, raw_output: output_text.into_bytes(), + truncation_policy: turn.truncation_policy, max_output_tokens, // Sandbox denial is terminal, so there is no live // process for write_stdin to resume. diff --git a/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs b/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs index b7712265a1..eb4323b947 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs @@ -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_for_source; use super::post_unified_exec_tool_use_payload; #[derive(Debug, Deserialize)] @@ -48,7 +47,6 @@ impl ToolExecutor for WriteStdinHandler { let ToolInvocation { session, turn, - source, payload, .. } = invocation; @@ -63,11 +61,6 @@ impl ToolExecutor for WriteStdinHandler { }; let args: WriteStdinArgs = parse_arguments(&arguments)?; - 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 @@ -75,7 +68,8 @@ impl ToolExecutor for WriteStdinHandler { process_id: args.session_id, input: &args.chars, yield_time_ms: args.yield_time_ms, - max_output_tokens, + max_output_tokens: args.max_output_tokens, + truncation_policy: turn.truncation_policy, }) .await .map_err(|err| { diff --git a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs index 50a8e44214..0eb9ebbedb 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs @@ -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), diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 25abaf2835..6c409a60f6 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -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; @@ -93,6 +94,7 @@ pub(crate) struct ExecCommandRequest { pub process_id: i32, pub yield_time_ms: u64, pub max_output_tokens: Option, + pub truncation_policy: TruncationPolicy, pub cwd: AbsolutePathBuf, pub sandbox_cwd: AbsolutePathBuf, pub environment: Arc, @@ -111,6 +113,7 @@ pub(crate) struct WriteStdinRequest<'a> { pub input: &'a str, pub yield_time_ms: u64, pub max_output_tokens: Option, + pub truncation_policy: TruncationPolicy, } #[derive(Default)] diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index 6926f9cbf1..e524d09722 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -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 } diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 8c0cc89680..42d64adc21 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -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, diff --git a/codex-rs/core/src/unified_exec/process_manager_tests.rs b/codex-rs/core/src/unified_exec/process_manager_tests.rs index cd83ce7a9c..8debfea154 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -170,6 +170,7 @@ async fn failed_initial_end_for_unstored_process_uses_fallback_output() { process_id: 123, yield_time_ms: 1000, max_output_tokens: None, + truncation_policy: turn.truncation_policy, #[allow(deprecated)] cwd: turn.cwd.clone(), #[allow(deprecated)]