Defer exec output truncation defaults to formatting

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<u8>,
+    pub truncation_policy: TruncationPolicy,
     pub max_output_tokens: Option<usize>,
     pub process_id: Option<i32>,
     pub exit_code: Option<i32>,
@@ -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))
     }
 

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),

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<usize>,
-    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<usize>,
-    truncation_policy: TruncationPolicy,
-) -> Option<usize> {
-    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<String>,

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<ToolInvocation> for ExecCommandHandler {
             turn,
             tracker,
             call_id,
-            source,
             payload,
             ..
         } = invocation;
@@ -163,11 +161,6 @@ impl ToolExecutor<ToolInvocation> 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<ToolInvocation> 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<ToolInvocation> 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<ToolInvocation> 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.

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<ToolInvocation> for WriteStdinHandler {
         let ToolInvocation {
             session,
             turn,
-            source,
             payload,
             ..
         } = invocation;
@@ -63,11 +61,6 @@ impl ToolExecutor<ToolInvocation> 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<ToolInvocation> 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| {

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),

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<usize>,
+    pub truncation_policy: TruncationPolicy,
     pub cwd: AbsolutePathBuf,
     pub sandbox_cwd: AbsolutePathBuf,
     pub environment: Arc<Environment>,
@@ -111,6 +113,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)]

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
 }

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,

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)]