Avoid full-string scans in token-budget truncation (#49712)

## Why

Token-budget truncation uses byte counts to estimate tokens, but previously scanned the entire string to count characters, including the discarded middle.

## What changed

Use UTF-8 boundary lookups to select the retained prefix and suffix for token-budget truncation. Skip character counting for a zero-token budget and defer it to paths that need character counts.

## Testing

Add regression coverage for multibyte characters crossing both retained boundaries, checking the truncated output and original token estimate.

GitOrigin-RevId: 37fee0c8d4ede8e74b491d3d641b0ea0e9003dd4
This commit is contained in:
jif
2026-09-30 19:14:56 +00:00
committed by copyberry
parent 596ae94fb0
commit 5e96aabd68
2 changed files with 23 additions and 3 deletions
+9 -3
View File
@@ -40,9 +40,8 @@ fn truncate_with_byte_estimate(s: &str, max_bytes: usize, use_tokens: bool) -> S
return String::new();
}
let total_chars = s.chars().count();
if max_bytes == 0 {
let total_chars = if use_tokens { 0 } else { s.chars().count() };
return format_truncation_marker(
use_tokens,
removed_units(use_tokens, s.len(), total_chars),
@@ -55,7 +54,14 @@ fn truncate_with_byte_estimate(s: &str, max_bytes: usize, use_tokens: bool) -> S
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 (removed_chars, left, right) = if use_tokens {
// Token markers depend only on byte counts; avoid scanning the discarded middle.
let prefix_end = s.floor_char_boundary(left_budget);
let suffix_start = s.ceil_char_boundary(total_bytes - right_budget);
(0, &s[..prefix_end], &s[suffix_start..])
} else {
split_string(s, left_budget, right_budget)
};
let marker = format_truncation_marker(
use_tokens,
removed_units(
@@ -109,6 +109,20 @@ fn truncate_middle_tokens_handles_utf8_content() {
assert_eq!(tokens, Some(16));
}
#[test]
fn truncate_middle_tokens_keeps_partial_characters_out_of_both_ends() {
for (input, expected) in [
("a😀b😀c", "a…2 tokens truncated…c"),
("é中😀é", "é…2 tokens truncated…é"),
("😀😀😀", "…2 tokens truncated…"),
] {
assert_eq!(
truncate_middle_with_token_budget(input, /*max_tokens*/ 1),
(expected.to_string(), Some(3)),
);
}
}
#[test]
fn truncate_middle_bytes_handles_utf8_content() {
let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with text\n";