Move shell to use truncate_text (#6842)

Move shell to use the configurable `truncate_text`

---------

Co-authored-by: pakrym-oai <[email protected]>

Author: aibrahim-oai

codex-rs/core/src/client_common.rs

@@ -165,11 +165,9 @@ fn build_structured_output(parsed: &ExecOutputJson) -> String {
     ));
 
     let mut output = parsed.output.clone();
-    if let Some(total_lines) = extract_total_output_lines(&parsed.output) {
+    if let Some((stripped, total_lines)) = strip_total_output_header(&parsed.output) {
         sections.push(format!("Total output lines: {total_lines}"));
-        if let Some(stripped) = strip_total_output_header(&output) {
-            output = stripped.to_string();
-        }
+        output = stripped.to_string();
     }
 
     sections.push("Output:".to_string());
@@ -178,19 +176,12 @@ fn build_structured_output(parsed: &ExecOutputJson) -> String {
     sections.join("\n")
 }
 
-fn extract_total_output_lines(output: &str) -> Option<u32> {
-    let marker_start = output.find("[... omitted ")?;
-    let marker = &output[marker_start..];
-    let (_, after_of) = marker.split_once(" of ")?;
-    let (total_segment, _) = after_of.split_once(' ')?;
-    total_segment.parse::<u32>().ok()
-}
-
-fn strip_total_output_header(output: &str) -> Option<&str> {
+fn strip_total_output_header(output: &str) -> Option<(&str, u32)> {
     let after_prefix = output.strip_prefix("Total output lines: ")?;
-    let (_, remainder) = after_prefix.split_once('\n')?;
+    let (total_segment, remainder) = after_prefix.split_once('\n')?;
+    let total_lines = total_segment.parse::<u32>().ok()?;
     let remainder = remainder.strip_prefix('\n').unwrap_or(remainder);
-    Some(remainder)
+    Some((remainder, total_lines))
 }
 
 #[derive(Debug)]

codex-rs/core/src/codex.rs

@@ -2479,8 +2479,9 @@ mod tests {
             duration: StdDuration::from_secs(1),
             timed_out: true,
         };
+        let (_, turn_context) = make_session_and_context();
 
-        let out = format_exec_output_str(&exec);
+        let out = format_exec_output_str(&exec, turn_context.truncation_policy);
 
         assert_eq!(
             out,

codex-rs/core/src/context_manager/history.rs

@@ -145,13 +145,17 @@ impl ContextManager {
     }
 
     fn process_item(&self, item: &ResponseItem, policy: TruncationPolicy) -> ResponseItem {
+        let policy_with_serialization_budget = policy.mul(1.2);
         match item {
             ResponseItem::FunctionCallOutput { call_id, output } => {
-                let truncated = truncate_text(output.content.as_str(), policy);
-                let truncated_items = output
-                    .content_items
-                    .as_ref()
-                    .map(|items| truncate_function_output_items_with_policy(items, policy));
+                let truncated =
+                    truncate_text(output.content.as_str(), policy_with_serialization_budget);
+                let truncated_items = output.content_items.as_ref().map(|items| {
+                    truncate_function_output_items_with_policy(
+                        items,
+                        policy_with_serialization_budget,
+                    )
+                });
                 ResponseItem::FunctionCallOutput {
                     call_id: call_id.clone(),
                     output: FunctionCallOutputPayload {
@@ -162,7 +166,7 @@ impl ContextManager {
                 }
             }
             ResponseItem::CustomToolCallOutput { call_id, output } => {
-                let truncated = truncate_text(output, policy);
+                let truncated = truncate_text(output, policy_with_serialization_budget);
                 ResponseItem::CustomToolCallOutput {
                     call_id: call_id.clone(),
                     output: truncated,

codex-rs/core/src/context_manager/history_tests.rs

@@ -12,8 +12,8 @@ use codex_protocol::models::ReasoningItemReasoningSummary;
 use pretty_assertions::assert_eq;
 use regex_lite::Regex;
 
-const EXEC_FORMAT_MAX_LINES: usize = 256;
 const EXEC_FORMAT_MAX_BYTES: usize = 10_000;
+const EXEC_FORMAT_MAX_TOKENS: usize = 2_500;
 
 fn assistant_msg(text: &str) -> ResponseItem {
     ResponseItem::Message {
@@ -56,6 +56,10 @@ fn reasoning_msg(text: &str) -> ResponseItem {
     }
 }
 
+fn truncate_exec_output(content: &str) -> String {
+    truncate::truncate_text(content, TruncationPolicy::Tokens(EXEC_FORMAT_MAX_TOKENS))
+}
+
 #[test]
 fn filters_non_api_messages() {
     let mut h = ContextManager::default();
@@ -228,7 +232,7 @@ fn normalization_retains_local_shell_outputs() {
         ResponseItem::FunctionCallOutput {
             call_id: "shell-1".to_string(),
             output: FunctionCallOutputPayload {
-                content: "ok".to_string(),
+                content: "Total output lines: 1\n\nok".to_string(),
                 ..Default::default()
             },
         },
@@ -330,8 +334,8 @@ fn record_items_respects_custom_token_limit() {
     assert!(stored.content.contains("tokens truncated"));
 }
 
-fn assert_truncated_message_matches(message: &str, line: &str, total_lines: usize) {
-    let pattern = truncated_message_pattern(line, total_lines);
+fn assert_truncated_message_matches(message: &str, line: &str, expected_removed: usize) {
+    let pattern = truncated_message_pattern(line);
     let regex = Regex::new(&pattern).unwrap_or_else(|err| {
         panic!("failed to compile regex {pattern}: {err}");
     });
@@ -347,51 +351,37 @@ fn assert_truncated_message_matches(message: &str, line: &str, total_lines: usiz
         "body exceeds byte limit: {} bytes",
         body.len()
     );
+    let removed: usize = captures
+        .name("removed")
+        .expect("missing removed capture")
+        .as_str()
+        .parse()
+        .unwrap_or_else(|err| panic!("invalid removed tokens: {err}"));
+    assert_eq!(removed, expected_removed, "mismatched removed token count");
 }
 
-fn truncated_message_pattern(line: &str, total_lines: usize) -> String {
-    let head_lines = EXEC_FORMAT_MAX_LINES / 2;
-    let tail_lines = EXEC_FORMAT_MAX_LINES - head_lines;
-    let head_take = head_lines.min(total_lines);
-    let tail_take = tail_lines.min(total_lines.saturating_sub(head_take));
-    let omitted = total_lines.saturating_sub(head_take + tail_take);
+fn truncated_message_pattern(line: &str) -> String {
     let escaped_line = regex_lite::escape(line);
-    if omitted == 0 {
-        return format!(
-            r"(?s)^Total output lines: {total_lines}\n\n(?P<body>{escaped_line}.*\n\[\.{{3}} removed \d+ bytes to fit {EXEC_FORMAT_MAX_BYTES} byte limit \.{{3}}]\n\n.*)$",
-        );
-    }
-    format!(
-        r"(?s)^Total output lines: {total_lines}\n\n(?P<body>{escaped_line}.*\n\[\.{{3}} omitted {omitted} of {total_lines} lines \.{{3}}]\n\n.*)$",
-    )
+    format!(r"(?s)^(?P<body>{escaped_line}.*?)(?:\r?)?…(?P<removed>\d+) tokens truncated…(?:.*)?$")
 }
 
 #[test]
 fn format_exec_output_truncates_large_error() {
     let line = "very long execution error line that should trigger truncation\n";
     let large_error = line.repeat(2_500); // way beyond both byte and line limits
 
-    let truncated = truncate::truncate_with_line_bytes_budget(&large_error, EXEC_FORMAT_MAX_BYTES);
+    let truncated = truncate_exec_output(&large_error);
 
-    let total_lines = large_error.lines().count();
-    assert_truncated_message_matches(&truncated, line, total_lines);
+    assert_truncated_message_matches(&truncated, line, 36250);
     assert_ne!(truncated, large_error);
 }
 
 #[test]
 fn format_exec_output_marks_byte_truncation_without_omitted_lines() {
-    let long_line = "a".repeat(EXEC_FORMAT_MAX_BYTES + 50);
-    let truncated = truncate::truncate_with_line_bytes_budget(&long_line, EXEC_FORMAT_MAX_BYTES);
-
+    let long_line = "a".repeat(EXEC_FORMAT_MAX_BYTES + 10000);
+    let truncated = truncate_exec_output(&long_line);
     assert_ne!(truncated, long_line);
-    let removed_bytes = long_line.len().saturating_sub(EXEC_FORMAT_MAX_BYTES);
-    let marker_line = format!(
-        "[... removed {removed_bytes} bytes to fit {EXEC_FORMAT_MAX_BYTES} byte limit ...]"
-    );
-    assert!(
-        truncated.contains(&marker_line),
-        "missing byte truncation marker: {truncated}"
-    );
+    assert_truncated_message_matches(&truncated, "a", 2500);
     assert!(
         !truncated.contains("omitted"),
         "line omission marker should not appear when no lines were dropped: {truncated}"
@@ -401,34 +391,25 @@ fn format_exec_output_marks_byte_truncation_without_omitted_lines() {
 #[test]
 fn format_exec_output_returns_original_when_within_limits() {
     let content = "example output\n".repeat(10);
-
-    assert_eq!(
-        truncate::truncate_with_line_bytes_budget(&content, EXEC_FORMAT_MAX_BYTES),
-        content
-    );
+    assert_eq!(truncate_exec_output(&content), content);
 }
 
 #[test]
 fn format_exec_output_reports_omitted_lines_and_keeps_head_and_tail() {
-    let total_lines = EXEC_FORMAT_MAX_LINES + 100;
+    let total_lines = 2_000;
+    let filler = "x".repeat(64);
     let content: String = (0..total_lines)
-        .map(|idx| format!("line-{idx}\n"))
+        .map(|idx| format!("line-{idx}-{filler}\n"))
         .collect();
 
-    let truncated = truncate::truncate_with_line_bytes_budget(&content, EXEC_FORMAT_MAX_BYTES);
-    let omitted = total_lines - EXEC_FORMAT_MAX_LINES;
-    let expected_marker = format!("[... omitted {omitted} of {total_lines} lines ...]");
-
-    assert!(
-        truncated.contains(&expected_marker),
-        "missing omitted marker: {truncated}"
-    );
+    let truncated = truncate_exec_output(&content);
+    assert_truncated_message_matches(&truncated, "line-0-", 34_723);
     assert!(
-        truncated.contains("line-0\n"),
+        truncated.contains("line-0-"),
         "expected head line to remain: {truncated}"
     );
 
-    let last_line = format!("line-{}\n", total_lines - 1);
+    let last_line = format!("line-{}-", total_lines - 1);
     assert!(
         truncated.contains(&last_line),
         "expected tail line to remain: {truncated}"
@@ -437,22 +418,15 @@ fn format_exec_output_reports_omitted_lines_and_keeps_head_and_tail() {
 
 #[test]
 fn format_exec_output_prefers_line_marker_when_both_limits_exceeded() {
-    let total_lines = EXEC_FORMAT_MAX_LINES + 42;
+    let total_lines = 300;
     let long_line = "x".repeat(256);
     let content: String = (0..total_lines)
         .map(|idx| format!("line-{idx}-{long_line}\n"))
         .collect();
 
-    let truncated = truncate::truncate_with_line_bytes_budget(&content, EXEC_FORMAT_MAX_BYTES);
+    let truncated = truncate_exec_output(&content);
 
-    assert!(
-        truncated.contains("[... omitted 42 of 298 lines ...]"),
-        "expected omitted marker when line count exceeds limit: {truncated}"
-    );
-    assert!(
-        !truncated.contains("byte limit"),
-        "line omission marker should take precedence over byte marker: {truncated}"
-    );
+    assert_truncated_message_matches(&truncated, "line-0-", 17_423);
 }
 
 //TODO(aibrahim): run CI in release mode.

codex-rs/core/src/context_manager/mod.rs

@@ -1,5 +1,4 @@
 mod history;
 mod normalize;
 
-pub(crate) use crate::truncate::truncate_with_line_bytes_budget;
 pub(crate) use history::ContextManager;

codex-rs/core/src/tasks/user_shell.rs

@@ -122,7 +122,11 @@ impl SessionTask for UserShellCommandTask {
                     duration: Duration::ZERO,
                     timed_out: false,
                 };
-                let output_items = [user_shell_command_record_item(&raw_command, &exec_output)];
+                let output_items = [user_shell_command_record_item(
+                    &raw_command,
+                    &exec_output,
+                    &turn_context,
+                )];
                 session
                     .record_conversation_items(turn_context.as_ref(), &output_items)
                     .await;
@@ -164,12 +168,19 @@ impl SessionTask for UserShellCommandTask {
                             aggregated_output: output.aggregated_output.text.clone(),
                             exit_code: output.exit_code,
                             duration: output.duration,
-                            formatted_output: format_exec_output_str(&output),
+                            formatted_output: format_exec_output_str(
+                                &output,
+                                turn_context.truncation_policy,
+                            ),
                         }),
                     )
                     .await;
 
-                let output_items = [user_shell_command_record_item(&raw_command, &output)];
+                let output_items = [user_shell_command_record_item(
+                    &raw_command,
+                    &output,
+                    &turn_context,
+                )];
                 session
                     .record_conversation_items(turn_context.as_ref(), &output_items)
                     .await;
@@ -201,11 +212,18 @@ impl SessionTask for UserShellCommandTask {
                             aggregated_output: exec_output.aggregated_output.text.clone(),
                             exit_code: exec_output.exit_code,
                             duration: exec_output.duration,
-                            formatted_output: format_exec_output_str(&exec_output),
+                            formatted_output: format_exec_output_str(
+                                &exec_output,
+                                turn_context.truncation_policy,
+                            ),
                         }),
                     )
                     .await;
-                let output_items = [user_shell_command_record_item(&raw_command, &exec_output)];
+                let output_items = [user_shell_command_record_item(
+                    &raw_command,
+                    &exec_output,
+                    &turn_context,
+                )];
                 session
                     .record_conversation_items(turn_context.as_ref(), &output_items)
                     .await;

codex-rs/core/src/tools/events.rs

@@ -242,12 +242,16 @@ impl ToolEmitter {
         self.emit(ctx, ToolEventStage::Begin).await;
     }
 
-    fn format_exec_output_for_model(&self, output: &ExecToolCallOutput) -> String {
+    fn format_exec_output_for_model(
+        &self,
+        output: &ExecToolCallOutput,
+        ctx: ToolEventCtx<'_>,
+    ) -> String {
         match self {
             Self::Shell { freeform: true, .. } => {
-                super::format_exec_output_for_model_freeform(output)
+                super::format_exec_output_for_model_freeform(output, ctx.turn.truncation_policy)
             }
-            _ => super::format_exec_output_for_model_structured(output),
+            _ => super::format_exec_output_for_model_structured(output, ctx.turn.truncation_policy),
         }
     }
 
@@ -258,7 +262,7 @@ impl ToolEmitter {
     ) -> Result<String, FunctionCallError> {
         let (event, result) = match out {
             Ok(output) => {
-                let content = self.format_exec_output_for_model(&output);
+                let content = self.format_exec_output_for_model(&output, ctx);
                 let exit_code = output.exit_code;
                 let event = ToolEventStage::Success(output);
                 let result = if exit_code == 0 {
@@ -270,7 +274,7 @@ impl ToolEmitter {
             }
             Err(ToolError::Codex(CodexErr::Sandbox(SandboxErr::Timeout { output })))
             | Err(ToolError::Codex(CodexErr::Sandbox(SandboxErr::Denied { output }))) => {
-                let response = self.format_exec_output_for_model(&output);
+                let response = self.format_exec_output_for_model(&output, ctx);
                 let event = ToolEventStage::Failure(ToolEventFailure::Output(*output));
                 let result = Err(FunctionCallError::RespondToModel(response));
                 (event, result)
@@ -359,7 +363,7 @@ async fn emit_exec_stage(
                 aggregated_output: output.aggregated_output.text.clone(),
                 exit_code: output.exit_code,
                 duration: output.duration,
-                formatted_output: format_exec_output_str(&output),
+                formatted_output: format_exec_output_str(&output, ctx.turn.truncation_policy),
             };
             emit_exec_end(ctx, exec_input, exec_result).await;
         }