Move string truncation helpers into codex-utils-string (#15572)

- move the shared byte-based middle truncation logic from `core` into `codex-utils-string` - keep token-specific truncation in `codex-core` so rollout can reuse the shared helper in the next stacked PR --------- Co-authored-by: Codex <noreply@openai.com>
2026-07-01 00:31:56 +08:00 · 2026-03-24 15:45:40 -07:00
parent 0b619afc87
commit 062fa7a2bb
36 changed files with 551 additions and 487 deletions
@@ -0,0 +1,6 @@
+load("//:defs.bzl", "codex_rust_crate")
+
+codex_rust_crate(
+    name = "output-truncation",
+    crate_name = "codex_utils_output_truncation",
+)
@@ -0,0 +1,15 @@
+[package]
+edition.workspace = true
+license.workspace = true
+name = "codex-utils-output-truncation"
+version.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+codex-protocol = { workspace = true }
+codex-utils-string = { workspace = true }
+
+[dev-dependencies]
+pretty_assertions = { workspace = true }
@@ -0,0 +1,142 @@
+//! Helpers for truncating tool and exec output using [`TruncationPolicy`](codex_protocol::protocol::TruncationPolicy).
+
+use codex_protocol::models::FunctionCallOutputContentItem;
+pub use codex_utils_string::approx_bytes_for_tokens;
+pub use codex_utils_string::approx_token_count;
+pub use codex_utils_string::approx_tokens_from_byte_count;
+use codex_utils_string::truncate_middle_chars;
+use codex_utils_string::truncate_middle_with_token_budget;
+
+pub use codex_protocol::protocol::TruncationPolicy;
+
+pub fn formatted_truncate_text(content: &str, policy: TruncationPolicy) -> String {
+    if content.len() <= policy.byte_budget() {
+        return content.to_string();
+    }
+
+    let total_lines = content.lines().count();
+    let result = truncate_text(content, policy);
+    format!("Total output lines: {total_lines}\n\n{result}")
+}
+
+pub fn truncate_text(content: &str, policy: TruncationPolicy) -> String {
+    match policy {
+        TruncationPolicy::Bytes(bytes) => truncate_middle_chars(content, bytes),
+        TruncationPolicy::Tokens(tokens) => truncate_middle_with_token_budget(content, tokens).0,
+    }
+}
+
+pub fn formatted_truncate_text_content_items_with_policy(
+    items: &[FunctionCallOutputContentItem],
+    policy: TruncationPolicy,
+) -> (Vec<FunctionCallOutputContentItem>, Option<usize>) {
+    let text_segments = items
+        .iter()
+        .filter_map(|item| match item {
+            FunctionCallOutputContentItem::InputText { text } => Some(text.as_str()),
+            FunctionCallOutputContentItem::InputImage { .. } => None,
+        })
+        .collect::<Vec<_>>();
+
+    if text_segments.is_empty() {
+        return (items.to_vec(), None);
+    }
+
+    let mut combined = String::new();
+    for text in &text_segments {
+        if !combined.is_empty() {
+            combined.push('\n');
+        }
+        combined.push_str(text);
+    }
+
+    if combined.len() <= policy.byte_budget() {
+        return (items.to_vec(), None);
+    }
+
+    let mut out = vec![FunctionCallOutputContentItem::InputText {
+        text: formatted_truncate_text(&combined, policy),
+    }];
+    out.extend(items.iter().filter_map(|item| match item {
+        FunctionCallOutputContentItem::InputImage { image_url, detail } => {
+            Some(FunctionCallOutputContentItem::InputImage {
+                image_url: image_url.clone(),
+                detail: *detail,
+            })
+        }
+        FunctionCallOutputContentItem::InputText { .. } => None,
+    }));
+
+    (out, Some(approx_token_count(&combined)))
+}
+
+pub fn truncate_function_output_items_with_policy(
+    items: &[FunctionCallOutputContentItem],
+    policy: TruncationPolicy,
+) -> Vec<FunctionCallOutputContentItem> {
+    let mut out: Vec<FunctionCallOutputContentItem> = Vec::with_capacity(items.len());
+    let mut remaining_budget = match policy {
+        TruncationPolicy::Bytes(_) => policy.byte_budget(),
+        TruncationPolicy::Tokens(_) => policy.token_budget(),
+    };
+    let mut omitted_text_items = 0usize;
+
+    for item in items {
+        match item {
+            FunctionCallOutputContentItem::InputText { text } => {
+                if remaining_budget == 0 {
+                    omitted_text_items += 1;
+                    continue;
+                }
+
+                let cost = match policy {
+                    TruncationPolicy::Bytes(_) => text.len(),
+                    TruncationPolicy::Tokens(_) => approx_token_count(text),
+                };
+
+                if cost <= remaining_budget {
+                    out.push(FunctionCallOutputContentItem::InputText { text: text.clone() });
+                    remaining_budget = remaining_budget.saturating_sub(cost);
+                } else {
+                    let snippet_policy = match policy {
+                        TruncationPolicy::Bytes(_) => TruncationPolicy::Bytes(remaining_budget),
+                        TruncationPolicy::Tokens(_) => TruncationPolicy::Tokens(remaining_budget),
+                    };
+                    let snippet = truncate_text(text, snippet_policy);
+                    if snippet.is_empty() {
+                        omitted_text_items += 1;
+                    } else {
+                        out.push(FunctionCallOutputContentItem::InputText { text: snippet });
+                    }
+                    remaining_budget = 0;
+                }
+            }
+            FunctionCallOutputContentItem::InputImage { image_url, detail } => {
+                out.push(FunctionCallOutputContentItem::InputImage {
+                    image_url: image_url.clone(),
+                    detail: *detail,
+                });
+            }
+        }
+    }
+
+    if omitted_text_items > 0 {
+        out.push(FunctionCallOutputContentItem::InputText {
+            text: format!("[omitted {omitted_text_items} text items ...]"),
+        });
+    }
+
+    out
+}
+
+pub fn approx_tokens_from_byte_count_i64(bytes: i64) -> i64 {
+    if bytes <= 0 {
+        return 0;
+    }
+
+    let bytes = usize::try_from(bytes).unwrap_or(usize::MAX);
+    i64::try_from(approx_tokens_from_byte_count(bytes)).unwrap_or(i64::MAX)
+}
+
+#[cfg(test)]
+mod truncate_tests;
@@ -0,0 +1,281 @@
+use crate::TruncationPolicy;
+use crate::approx_token_count;
+use crate::approx_tokens_from_byte_count_i64;
+use crate::formatted_truncate_text;
+use crate::formatted_truncate_text_content_items_with_policy;
+use crate::truncate_function_output_items_with_policy;
+use crate::truncate_text;
+use codex_protocol::models::FunctionCallOutputContentItem;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn truncate_bytes_less_than_placeholder_returns_placeholder() {
+    let content = "example output";
+
+    assert_eq!(
+        "Total output lines: 1\n\n…13 chars truncated…t",
+        formatted_truncate_text(content, TruncationPolicy::Bytes(1)),
+    );
+}
+
+#[test]
+fn truncate_tokens_less_than_placeholder_returns_placeholder() {
+    let content = "example output";
+
+    assert_eq!(
+        "Total output lines: 1\n\nex…3 tokens truncated…ut",
+        formatted_truncate_text(content, TruncationPolicy::Tokens(1)),
+    );
+}
+
+#[test]
+fn truncate_tokens_under_limit_returns_original() {
+    let content = "example output";
+
+    assert_eq!(
+        content,
+        formatted_truncate_text(content, TruncationPolicy::Tokens(10)),
+    );
+}
+
+#[test]
+fn truncate_bytes_under_limit_returns_original() {
+    let content = "example output";
+
+    assert_eq!(
+        content,
+        formatted_truncate_text(content, TruncationPolicy::Bytes(20)),
+    );
+}
+
+#[test]
+fn truncate_tokens_over_limit_returns_truncated() {
+    let content = "this is an example of a long output that should be truncated";
+
+    assert_eq!(
+        "Total output lines: 1\n\nthis is an…10 tokens truncated… truncated",
+        formatted_truncate_text(content, TruncationPolicy::Tokens(5)),
+    );
+}
+
+#[test]
+fn truncate_bytes_over_limit_returns_truncated() {
+    let content = "this is an example of a long output that should be truncated";
+
+    assert_eq!(
+        "Total output lines: 1\n\nthis is an exam…30 chars truncated…ld be truncated",
+        formatted_truncate_text(content, TruncationPolicy::Bytes(30)),
+    );
+}
+
+#[test]
+fn truncate_bytes_reports_original_line_count_when_truncated() {
+    let content =
+        "this is an example of a long output that should be truncated\nalso some other line";
+
+    assert_eq!(
+        "Total output lines: 2\n\nthis is an exam…51 chars truncated…some other line",
+        formatted_truncate_text(content, TruncationPolicy::Bytes(30)),
+    );
+}
+
+#[test]
+fn truncate_tokens_reports_original_line_count_when_truncated() {
+    let content =
+        "this is an example of a long output that should be truncated\nalso some other line";
+
+    assert_eq!(
+        "Total output lines: 2\n\nthis is an example o…11 tokens truncated…also some other line",
+        formatted_truncate_text(content, TruncationPolicy::Tokens(10)),
+    );
+}
+
+#[test]
+fn truncate_middle_bytes_handles_utf8_content() {
+    let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with text\n";
+    let out = truncate_text(s, TruncationPolicy::Bytes(20));
+    assert_eq!(out, "😀😀…21 chars truncated…with text\n");
+}
+
+#[test]
+fn truncates_across_multiple_under_limit_texts_and_reports_omitted() {
+    let chunk = "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron pi rho sigma tau upsilon phi chi psi omega.\n";
+    let chunk_tokens = approx_token_count(chunk);
+    assert!(chunk_tokens > 0, "chunk must consume tokens");
+    let limit = chunk_tokens * 3;
+    let t1 = chunk.to_string();
+    let t2 = chunk.to_string();
+    let t3 = chunk.repeat(10);
+    let t4 = chunk.to_string();
+    let t5 = chunk.to_string();
+
+    let items = vec![
+        FunctionCallOutputContentItem::InputText { text: t1.clone() },
+        FunctionCallOutputContentItem::InputText { text: t2.clone() },
+        FunctionCallOutputContentItem::InputImage {
+            image_url: "img:mid".to_string(),
+            detail: None,
+        },
+        FunctionCallOutputContentItem::InputText { text: t3 },
+        FunctionCallOutputContentItem::InputText { text: t4 },
+        FunctionCallOutputContentItem::InputText { text: t5 },
+    ];
+
+    let output =
+        truncate_function_output_items_with_policy(&items, TruncationPolicy::Tokens(limit));
+
+    assert_eq!(output.len(), 5);
+
+    let first_text = match &output[0] {
+        FunctionCallOutputContentItem::InputText { text } => text,
+        other => panic!("unexpected first item: {other:?}"),
+    };
+    assert_eq!(first_text, &t1);
+
+    let second_text = match &output[1] {
+        FunctionCallOutputContentItem::InputText { text } => text,
+        other => panic!("unexpected second item: {other:?}"),
+    };
+    assert_eq!(second_text, &t2);
+
+    assert_eq!(
+        output[2],
+        FunctionCallOutputContentItem::InputImage {
+            image_url: "img:mid".to_string(),
+            detail: None,
+        }
+    );
+
+    let fourth_text = match &output[3] {
+        FunctionCallOutputContentItem::InputText { text } => text,
+        other => panic!("unexpected fourth item: {other:?}"),
+    };
+    assert!(
+        fourth_text.contains("tokens truncated"),
+        "expected marker in truncated snippet: {fourth_text}"
+    );
+
+    let summary_text = match &output[4] {
+        FunctionCallOutputContentItem::InputText { text } => text,
+        other => panic!("unexpected summary item: {other:?}"),
+    };
+    assert!(summary_text.contains("omitted 2 text items"));
+}
+
+#[test]
+fn formatted_truncate_text_content_items_with_policy_returns_original_under_limit() {
+    let items = vec![
+        FunctionCallOutputContentItem::InputText {
+            text: "alpha".to_string(),
+        },
+        FunctionCallOutputContentItem::InputText {
+            text: String::new(),
+        },
+        FunctionCallOutputContentItem::InputText {
+            text: "beta".to_string(),
+        },
+    ];
+
+    let (output, original_token_count) =
+        formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(32));
+
+    assert_eq!(output, items);
+    assert_eq!(original_token_count, None);
+}
+
+#[test]
+fn formatted_truncate_text_content_items_with_policy_preserves_empty_leading_text_behavior() {
+    let items = vec![
+        FunctionCallOutputContentItem::InputText {
+            text: String::new(),
+        },
+        FunctionCallOutputContentItem::InputText {
+            text: "abc".to_string(),
+        },
+    ];
+
+    let (output, original_token_count) =
+        formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(0));
+
+    assert_eq!(
+        output,
+        vec![FunctionCallOutputContentItem::InputText {
+            text: "Total output lines: 1\n\n…3 chars truncated…".to_string(),
+        }]
+    );
+    assert_eq!(original_token_count, Some(1));
+}
+
+#[test]
+fn formatted_truncate_text_content_items_with_policy_merges_text_and_appends_images() {
+    let items = vec![
+        FunctionCallOutputContentItem::InputText {
+            text: "abcd".to_string(),
+        },
+        FunctionCallOutputContentItem::InputImage {
+            image_url: "img:one".to_string(),
+            detail: None,
+        },
+        FunctionCallOutputContentItem::InputText {
+            text: "efgh".to_string(),
+        },
+        FunctionCallOutputContentItem::InputText {
+            text: "ijkl".to_string(),
+        },
+        FunctionCallOutputContentItem::InputImage {
+            image_url: "img:two".to_string(),
+            detail: None,
+        },
+    ];
+
+    let (output, original_token_count) =
+        formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(8));
+
+    assert_eq!(
+        output,
+        vec![
+            FunctionCallOutputContentItem::InputText {
+                text: "Total output lines: 3\n\nabcd…6 chars truncated…ijkl".to_string(),
+            },
+            FunctionCallOutputContentItem::InputImage {
+                image_url: "img:one".to_string(),
+                detail: None,
+            },
+            FunctionCallOutputContentItem::InputImage {
+                image_url: "img:two".to_string(),
+                detail: None,
+            },
+        ]
+    );
+    assert_eq!(original_token_count, Some(4));
+}
+
+#[test]
+fn formatted_truncate_text_content_items_with_policy_merges_all_text_for_token_budget() {
+    let items = vec![
+        FunctionCallOutputContentItem::InputText {
+            text: "abcdefgh".to_string(),
+        },
+        FunctionCallOutputContentItem::InputText {
+            text: "ijklmnop".to_string(),
+        },
+    ];
+
+    let (output, original_token_count) =
+        formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Tokens(2));
+
+    assert_eq!(
+        output,
+        vec![FunctionCallOutputContentItem::InputText {
+            text: "Total output lines: 2\n\nabcd…3 tokens truncated…mnop".to_string(),
+        }]
+    );
+    assert_eq!(original_token_count, Some(5));
+}
+
+#[test]
+fn byte_count_conversion_clamps_non_positive_values() {
+    assert_eq!(approx_tokens_from_byte_count_i64(/*bytes*/ -1), 0);
+    assert_eq!(approx_tokens_from_byte_count_i64(/*bytes*/ 0), 0);
+    assert_eq!(approx_tokens_from_byte_count_i64(/*bytes*/ 5), 2);
+}
@@ -1,3 +1,11 @@
+mod truncate;
+
+pub use truncate::approx_bytes_for_tokens;
+pub use truncate::approx_token_count;
+pub use truncate::approx_tokens_from_byte_count;
+pub use truncate::truncate_middle_chars;
+pub use truncate::truncate_middle_with_token_budget;
+
 // Truncate a &str to a byte budget at a char boundary (prefix)
 #[inline]
 pub fn take_bytes_at_char_boundary(s: &str, maxb: usize) -> &str {
@@ -112,6 +120,7 @@ fn parse_markdown_hash_location_point(point: &str) -> Option<(&str, Option<&str>
 }

 #[cfg(test)]
+#[allow(warnings, clippy::all)]
 mod tests {
    use super::find_uuids;
    use super::normalize_markdown_hash_location_suffix;
@@ -0,0 +1,156 @@
+//! Utilities for truncating large chunks of output while preserving a prefix
+//! and suffix on UTF-8 boundaries.
+
+const APPROX_BYTES_PER_TOKEN: usize = 4;
+
+/// Truncate a string to `max_bytes` using a character-count marker.
+pub fn truncate_middle_chars(s: &str, max_bytes: usize) -> String {
+    truncate_with_byte_estimate(s, max_bytes, /*use_tokens*/ false)
+}
+
+/// Truncate the middle of a UTF-8 string to at most `max_tokens` approximate
+/// tokens, preserving the beginning and the end. Returns the possibly
+/// truncated string and `Some(original_token_count)` if truncation occurred;
+/// otherwise returns the original string and `None`.
+pub fn truncate_middle_with_token_budget(s: &str, max_tokens: usize) -> (String, Option<u64>) {
+    if s.is_empty() {
+        return (String::new(), None);
+    }
+
+    if max_tokens > 0 && s.len() <= approx_bytes_for_tokens(max_tokens) {
+        return (s.to_string(), None);
+    }
+
+    let truncated = truncate_with_byte_estimate(
+        s,
+        approx_bytes_for_tokens(max_tokens),
+        /*use_tokens*/ true,
+    );
+    let total_tokens = u64::try_from(approx_token_count(s)).unwrap_or(u64::MAX);
+
+    if truncated == s {
+        (truncated, None)
+    } else {
+        (truncated, Some(total_tokens))
+    }
+}
+
+fn truncate_with_byte_estimate(s: &str, max_bytes: usize, use_tokens: bool) -> String {
+    if s.is_empty() {
+        return String::new();
+    }
+
+    let total_chars = s.chars().count();
+
+    if max_bytes == 0 {
+        return format_truncation_marker(
+            use_tokens,
+            removed_units(use_tokens, s.len(), total_chars),
+        );
+    }
+
+    if s.len() <= max_bytes {
+        return s.to_string();
+    }
+
+    let total_bytes = s.len();
+    let (left_budget, right_budget) = split_budget(max_bytes);
+    let (removed_chars, left, right) = split_string(s, left_budget, right_budget);
+    let marker = format_truncation_marker(
+        use_tokens,
+        removed_units(
+            use_tokens,
+            total_bytes.saturating_sub(max_bytes),
+            removed_chars,
+        ),
+    );
+
+    assemble_truncated_output(left, right, &marker)
+}
+
+pub fn approx_token_count(text: &str) -> usize {
+    let len = text.len();
+    len.saturating_add(APPROX_BYTES_PER_TOKEN.saturating_sub(1)) / APPROX_BYTES_PER_TOKEN
+}
+
+pub fn approx_bytes_for_tokens(tokens: usize) -> usize {
+    tokens.saturating_mul(APPROX_BYTES_PER_TOKEN)
+}
+
+pub fn approx_tokens_from_byte_count(bytes: usize) -> u64 {
+    let bytes_u64 = bytes as u64;
+    bytes_u64.saturating_add((APPROX_BYTES_PER_TOKEN as u64).saturating_sub(1))
+        / (APPROX_BYTES_PER_TOKEN as u64)
+}
+
+fn split_string(s: &str, beginning_bytes: usize, end_bytes: usize) -> (usize, &str, &str) {
+    if s.is_empty() {
+        return (0, "", "");
+    }
+
+    let len = s.len();
+    let tail_start_target = len.saturating_sub(end_bytes);
+    let mut prefix_end = 0usize;
+    let mut suffix_start = len;
+    let mut removed_chars = 0usize;
+    let mut suffix_started = false;
+
+    for (idx, ch) in s.char_indices() {
+        let char_end = idx + ch.len_utf8();
+        if char_end <= beginning_bytes {
+            prefix_end = char_end;
+            continue;
+        }
+
+        if idx >= tail_start_target {
+            if !suffix_started {
+                suffix_start = idx;
+                suffix_started = true;
+            }
+            continue;
+        }
+
+        removed_chars = removed_chars.saturating_add(1);
+    }
+
+    if suffix_start < prefix_end {
+        suffix_start = prefix_end;
+    }
+
+    let before = &s[..prefix_end];
+    let after = &s[suffix_start..];
+
+    (removed_chars, before, after)
+}
+
+fn split_budget(budget: usize) -> (usize, usize) {
+    let left = budget / 2;
+    (left, budget - left)
+}
+
+fn format_truncation_marker(use_tokens: bool, removed_count: u64) -> String {
+    if use_tokens {
+        format!("…{removed_count} tokens truncated…")
+    } else {
+        format!("…{removed_count} chars truncated…")
+    }
+}
+
+fn removed_units(use_tokens: bool, removed_bytes: usize, removed_chars: usize) -> u64 {
+    if use_tokens {
+        approx_tokens_from_byte_count(removed_bytes)
+    } else {
+        u64::try_from(removed_chars).unwrap_or(u64::MAX)
+    }
+}
+
+fn assemble_truncated_output(prefix: &str, suffix: &str, marker: &str) -> String {
+    let mut out = String::with_capacity(prefix.len() + marker.len() + suffix.len() + 1);
+    out.push_str(prefix);
+    out.push_str(marker);
+    out.push_str(suffix);
+    out
+}
+
+#[cfg(test)]
+mod tests;
@@ -0,0 +1,71 @@
+use super::split_string;
+use super::truncate_middle_chars;
+use super::truncate_middle_with_token_budget;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn split_string_works() {
+    assert_eq!(split_string("hello world", 5, 5), (1, "hello", "world"));
+    assert_eq!(split_string("abc", 0, 0), (3, "", ""));
+}
+
+#[test]
+fn split_string_handles_empty_string() {
+    assert_eq!(split_string("", 4, 4), (0, "", ""));
+}
+
+#[test]
+fn split_string_only_keeps_prefix_when_tail_budget_is_zero() {
+    assert_eq!(split_string("abcdef", 3, 0), (3, "abc", ""));
+}
+
+#[test]
+fn split_string_only_keeps_suffix_when_prefix_budget_is_zero() {
+    assert_eq!(split_string("abcdef", 0, 3), (3, "", "def"));
+}
+
+#[test]
+fn split_string_handles_overlapping_budgets_without_removal() {
+    assert_eq!(split_string("abcdef", 4, 4), (0, "abcd", "ef"));
+}
+
+#[test]
+fn split_string_respects_utf8_boundaries() {
+    assert_eq!(split_string("😀abc😀", 5, 5), (1, "😀a", "c😀"));
+
+    assert_eq!(split_string("😀😀😀😀😀", 1, 1), (5, "", ""));
+    assert_eq!(split_string("😀😀😀😀😀", 7, 7), (3, "😀", "😀"));
+    assert_eq!(split_string("😀😀😀😀😀", 8, 8), (1, "😀😀", "😀😀"));
+}
+
+#[test]
+fn truncate_with_token_budget_returns_original_when_under_limit() {
+    let s = "short output";
+    let limit = 100;
+    let (out, original) = truncate_middle_with_token_budget(s, limit);
+    assert_eq!(out, s);
+    assert_eq!(original, None);
+}
+
+#[test]
+fn truncate_with_token_budget_reports_truncation_at_zero_limit() {
+    let s = "abcdef";
+    let (out, original) = truncate_middle_with_token_budget(s, 0);
+    assert_eq!(out, "…2 tokens truncated…");
+    assert_eq!(original, Some(2));
+}
+
+#[test]
+fn truncate_middle_tokens_handles_utf8_content() {
+    let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with text\n";
+    let (out, tokens) = truncate_middle_with_token_budget(s, 8);
+    assert_eq!(out, "😀😀😀😀…8 tokens truncated… line with text\n");
+    assert_eq!(tokens, Some(16));
+}
+
+#[test]
+fn truncate_middle_bytes_handles_utf8_content() {
+    let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with text\n";
+    let out = truncate_middle_chars(s, 20);
+    assert_eq!(out, "😀😀…21 chars truncated…with text\n");
+}