diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 78960de43f..2d7dac04e1 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -46,8 +46,7 @@ use codex_protocol::openai_models::ToolMode; use codex_tools::ToolName; use codex_utils_audio::estimate_audio_token_count; use codex_utils_output_truncation::TruncationPolicy; -use codex_utils_output_truncation::formatted_truncate_text_content_items_with_policy; -use codex_utils_output_truncation::truncate_function_output_items_with_policy; +use codex_utils_output_truncation::formatted_truncate_function_output_items_with_policy; use delegate::CodeModeCellDelegate; use delegate::CodeModeDispatchBroker; @@ -321,16 +320,7 @@ fn truncate_code_mode_result( ) -> Vec { let max_output_tokens = resolve_max_tokens(max_output_tokens); let policy = TruncationPolicy::Tokens(max_output_tokens); - if items - .iter() - .all(|item| matches!(item, FunctionCallOutputContentItem::InputText { .. })) - { - let (truncated_items, _) = - formatted_truncate_text_content_items_with_policy(&items, policy); - return truncated_items; - } - - truncate_function_output_items_with_policy(&items, policy, estimate_audio_token_count) + formatted_truncate_function_output_items_with_policy(items, policy, estimate_audio_token_count) } // Submit synchronously so the recorder sees the call before the cell's dispatch gate closes. diff --git a/codex-rs/utils/output-truncation/src/lib.rs b/codex-rs/utils/output-truncation/src/lib.rs index 4301a6107d..7c49956547 100644 --- a/codex-rs/utils/output-truncation/src/lib.rs +++ b/codex-rs/utils/output-truncation/src/lib.rs @@ -6,11 +6,12 @@ use codex_protocol::mcp::CallToolResult; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; +use codex_utils_string::TextTruncation; 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::plan_token_truncation; use codex_utils_string::truncate_middle_chars; -use codex_utils_string::truncate_middle_with_token_budget; pub use codex_protocol::protocol::TruncationPolicy; @@ -25,18 +26,43 @@ pub fn formatted_truncate_text(content: &str, policy: TruncationPolicy) -> Strin return content.to_string(); } + format_truncated_text(content, policy, &mut []) +} + +fn format_truncated_text( + content: &str, + policy: TruncationPolicy, + boundaries: &mut [OutputBoundary], +) -> String { let original_token_count = approx_token_count(content); let total_lines = content.lines().count(); - let result = truncate_text(content, policy); - format!( - "Warning: truncated output (original token count: {original_token_count})\nTotal output lines: {total_lines}\n\n{result}" - ) + let header = format!( + "Warning: truncated output (original token count: {original_token_count})\nTotal output lines: {total_lines}\n\n" + ); + let (snippet, edit) = truncate_text_with_plan(content, policy); + if let Some(edit) = edit { + for boundary in boundaries { + boundary.text_offset = rendered_offset(&edit, boundary.text_offset, header.len()); + } + } + header + &snippet } pub fn truncate_text(content: &str, policy: TruncationPolicy) -> String { + truncate_text_with_plan(content, policy).0 +} + +// Only token-budget callers supply boundaries that need remapping. +fn truncate_text_with_plan( + content: &str, + policy: TruncationPolicy, +) -> (String, Option) { match policy { - TruncationPolicy::Bytes(bytes) => truncate_middle_chars(content, bytes), - TruncationPolicy::Tokens(tokens) => truncate_middle_with_token_budget(content, tokens).0, + TruncationPolicy::Bytes(bytes) => (truncate_middle_chars(content, bytes), None), + TruncationPolicy::Tokens(tokens) => match plan_token_truncation(content, tokens) { + Some(edit) => (edit.apply(content), Some(edit)), + None => (content.to_owned(), None), + }, } } @@ -104,67 +130,169 @@ pub fn truncate_function_output_payload( } } -pub fn formatted_truncate_text_content_items_with_policy( - items: &[FunctionCallOutputContentItem], +/// An uncharged boundary in the original logical output. Token-boundary +/// formatters update the item index and UTF-8 text offset to the rendered +/// position, including when content at that position was omitted. +/// +/// At entry, item indices are clamped to the end of the input, text offsets +/// round down to an in-bounds UTF-8 boundary, and offsets on non-text items +/// are zeroed. Positions preceding an earlier boundary move to that position, +/// preserving caller order. On return, boundaries refer to the rendered items; +/// an end boundary precedes any generated omission notices. +pub struct OutputBoundary { + pub item_index: usize, + pub text_offset: usize, +} + +/// Format text-only output as a whole observation, including the truncation +/// warning. Charge mixed text and audio against the budget in order while +/// preserving images. +pub fn formatted_truncate_function_output_items_with_policy<'a>( + items: impl Into>, policy: TruncationPolicy, -) -> (Vec, Option) { - let text_segments = items - .iter() - .filter_map(|item| match item { - FunctionCallOutputContentItem::InputText { text } => Some(text.as_str()), - FunctionCallOutputContentItem::InputImage { .. } - | FunctionCallOutputContentItem::InputAudio { .. } - | FunctionCallOutputContentItem::EncryptedContent { .. } => None, - }) - .collect::>(); + estimate_audio_token_count: impl Fn(&str) -> usize, +) -> Vec { + formatted_output(items.into(), policy, estimate_audio_token_count, &mut []) +} - if text_segments.is_empty() { - return (items.to_vec(), None); +/// As above, with zero-width boundaries for callers that have already dropped +/// output. Input positions are normalized as documented on [`OutputBoundary`]. +pub fn formatted_output_at_token_boundaries<'a>( + items: impl Into>, + max_tokens: usize, + estimate_audio_token_count: impl Fn(&str) -> usize, + boundaries: &mut [OutputBoundary], +) -> Vec { + let items = items.into(); + normalize_boundaries(&items, boundaries); + formatted_output( + items, + TruncationPolicy::Tokens(max_tokens), + estimate_audio_token_count, + boundaries, + ) +} + +fn formatted_output( + items: Cow<'_, [FunctionCallOutputContentItem]>, + policy: TruncationPolicy, + estimate_audio_token_count: impl Fn(&str) -> usize, + boundaries: &mut [OutputBoundary], +) -> Vec { + let mut combined_len = 0; + for item in items.iter() { + let FunctionCallOutputContentItem::InputText { text } = item else { + return truncate_output_items(items, policy, estimate_audio_token_count, boundaries); + }; + combined_len += usize::from(combined_len != 0) + text.len(); + } + if combined_len <= policy.byte_budget() { + return items.into_owned(); } - let mut combined = String::new(); - for text in &text_segments { - if !combined.is_empty() { - combined.push('\n'); + let mut combined = String::with_capacity(combined_len); + let mut next_boundary = 0; + for (index, item) in items.iter().enumerate() { + if let FunctionCallOutputContentItem::InputText { text } = item { + if !combined.is_empty() { + combined.push('\n'); + } + while let Some(boundary) = boundaries.get_mut(next_boundary) + && boundary.item_index == index + { + boundary.item_index = 0; + boundary.text_offset += combined.len(); + next_boundary += 1; + } + combined.push_str(text); } - combined.push_str(text); } - - if combined.len() <= policy.byte_budget() { - return (items.to_vec(), None); + // A boundary after the last item belongs at the end of the merged text. + for boundary in &mut boundaries[next_boundary..] { + boundary.item_index = 0; + boundary.text_offset = combined.len(); } - - let original_token_count = approx_token_count(&combined); - let mut out = vec![FunctionCallOutputContentItem::InputText { - text: formatted_truncate_text(&combined, policy), - }]; - out.extend(items.iter().filter_map(|item| match item { - FunctionCallOutputContentItem::InputImage { image, detail } => { - Some(FunctionCallOutputContentItem::InputImage { - image: image.clone(), - detail: *detail, - }) - } - FunctionCallOutputContentItem::InputAudio { audio_url } => { - Some(FunctionCallOutputContentItem::InputAudio { - audio_url: audio_url.clone(), - }) - } - FunctionCallOutputContentItem::EncryptedContent { encrypted_content } => { - Some(FunctionCallOutputContentItem::EncryptedContent { - encrypted_content: encrypted_content.clone(), - }) - } - FunctionCallOutputContentItem::InputText { .. } => None, - })); - - (out, Some(original_token_count)) + vec![FunctionCallOutputContentItem::InputText { + text: format_truncated_text(&combined, policy, boundaries), + }] } pub fn truncate_function_output_items_with_policy( items: &[FunctionCallOutputContentItem], policy: TruncationPolicy, estimate_audio_token_count: impl Fn(&str) -> usize, +) -> Vec { + truncate_output_items( + Cow::Borrowed(items), + policy, + estimate_audio_token_count, + &mut [], + ) +} + +/// Budget in input order even if a caller previously dropped every media item. +/// Input positions are normalized as documented on [`OutputBoundary`]. +pub fn mixed_output_at_token_boundaries<'a>( + items: impl Into>, + max_tokens: usize, + estimate_audio_token_count: impl Fn(&str) -> usize, + boundaries: &mut [OutputBoundary], +) -> Vec { + let items = items.into(); + normalize_boundaries(&items, boundaries); + truncate_output_items( + items, + TruncationPolicy::Tokens(max_tokens), + estimate_audio_token_count, + boundaries, + ) +} + +fn normalize_boundaries( + items: &[FunctionCallOutputContentItem], + boundaries: &mut [OutputBoundary], +) { + let mut previous = (0, 0); + for boundary in boundaries { + let item_index = boundary.item_index.min(items.len()); + let text_offset = match items.get(item_index) { + Some(FunctionCallOutputContentItem::InputText { text }) => { + text.floor_char_boundary(boundary.text_offset.min(text.len())) + } + _ => 0, + }; + previous = previous.max((item_index, text_offset)); + (boundary.item_index, boundary.text_offset) = previous; + } +} + +fn truncate_output_items( + items: Cow<'_, [FunctionCallOutputContentItem]>, + policy: TruncationPolicy, + estimate_audio_token_count: impl Fn(&str) -> usize, + boundaries: &mut [OutputBoundary], +) -> Vec { + match items { + Cow::Borrowed(items) => truncate_output_iter( + items.iter().map(Cow::Borrowed), + policy, + estimate_audio_token_count, + boundaries, + ), + Cow::Owned(items) => truncate_output_iter( + items.into_iter().map(Cow::Owned), + policy, + estimate_audio_token_count, + boundaries, + ), + } +} + +fn truncate_output_iter<'a>( + items: impl ExactSizeIterator>, + policy: TruncationPolicy, + estimate_audio_token_count: impl Fn(&str) -> usize, + boundaries: &mut [OutputBoundary], ) -> Vec { let mut out: Vec = Vec::with_capacity(items.len()); let mut remaining_budget = match policy { @@ -173,47 +301,44 @@ pub fn truncate_function_output_items_with_policy( }; let mut omitted_text_items = 0usize; let mut omitted_audio_items = 0usize; + let mut next_boundary = 0; - for item in items { - match item { + for (index, item) in items.enumerate() { + let item_out = out.len(); + let item_policy = match policy { + TruncationPolicy::Bytes(_) => TruncationPolicy::Bytes(remaining_budget), + TruncationPolicy::Tokens(_) => TruncationPolicy::Tokens(remaining_budget), + }; + let mut text_edit = None; + let keep = match item.as_ref() { + // Empty text contributes no model content but still consumes an API array slot. + FunctionCallOutputContentItem::InputText { text } if text.is_empty() => false, + FunctionCallOutputContentItem::InputText { .. } if remaining_budget == 0 => { + omitted_text_items += 1; + false + } FunctionCallOutputContentItem::InputText { text } => { - // Empty text contributes no model content but still consumes an API array slot. - if text.is_empty() { - continue; - } - 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); + true } 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); + let (snippet, edit) = truncate_text_with_plan(text, item_policy); + text_edit = edit; if snippet.is_empty() { omitted_text_items += 1; } else { out.push(FunctionCallOutputContentItem::InputText { text: snippet }); } remaining_budget = 0; + false } } - FunctionCallOutputContentItem::InputImage { image, detail } => { - out.push(FunctionCallOutputContentItem::InputImage { - image: image.clone(), - detail: *detail, - }); - } + FunctionCallOutputContentItem::InputImage { .. } + | FunctionCallOutputContentItem::EncryptedContent { .. } => true, FunctionCallOutputContentItem::InputAudio { audio_url } => { let token_cost = estimate_audio_token_count(audio_url); let cost = match policy { @@ -221,20 +346,35 @@ pub fn truncate_function_output_items_with_policy( TruncationPolicy::Tokens(_) => token_cost, }; if cost <= remaining_budget { - out.push(FunctionCallOutputContentItem::InputAudio { - audio_url: audio_url.clone(), - }); remaining_budget = remaining_budget.saturating_sub(cost); + true } else { omitted_audio_items += 1; + false } } - FunctionCallOutputContentItem::EncryptedContent { encrypted_content } => { - out.push(FunctionCallOutputContentItem::EncryptedContent { - encrypted_content: encrypted_content.clone(), - }); + }; + while let Some(boundary) = boundaries.get_mut(next_boundary) + && boundary.item_index == index + { + boundary.item_index = item_out; + if let Some(edit) = &text_edit { + boundary.text_offset = + rendered_offset(edit, boundary.text_offset, /*header_bytes*/ 0); + } else if !keep { + boundary.text_offset = 0; } + next_boundary += 1; } + if keep { + out.push(item.into_owned()); + } + } + // End boundaries precede generated omission counts: counts are not original + // content and must never move a collector notice away from its source. + for boundary in &mut boundaries[next_boundary..] { + boundary.item_index = out.len(); + boundary.text_offset = 0; } if omitted_text_items > 0 { @@ -247,10 +387,22 @@ pub fn truncate_function_output_items_with_policy( text: format!("[omitted {omitted_audio_items} audio items ...]"), }); } - out } +fn rendered_offset(edit: &TextTruncation, boundary: usize, header_bytes: usize) -> usize { + if boundary == 0 { + 0 + } else if boundary <= edit.removed.start { + header_bytes + boundary + } else { + header_bytes + + edit.removed.start + + edit.marker.len() + + boundary.saturating_sub(edit.removed.end) + } +} + pub fn approx_tokens_from_byte_count_i64(bytes: i64) -> i64 { if bytes <= 0 { return 0; diff --git a/codex-rs/utils/output-truncation/src/truncate_tests.rs b/codex-rs/utils/output-truncation/src/truncate_tests.rs index 813412634a..f3e1244b82 100644 --- a/codex-rs/utils/output-truncation/src/truncate_tests.rs +++ b/codex-rs/utils/output-truncation/src/truncate_tests.rs @@ -1,8 +1,11 @@ +use crate::OutputBoundary; use crate::TruncationPolicy; use crate::approx_token_count; use crate::approx_tokens_from_byte_count_i64; +use crate::formatted_output_at_token_boundaries; +use crate::formatted_truncate_function_output_items_with_policy; use crate::formatted_truncate_text; -use crate::formatted_truncate_text_content_items_with_policy; +use crate::mixed_output_at_token_boundaries; use crate::truncate_function_output_items_with_policy; use crate::truncate_mcp_tool_result; use crate::truncate_text; @@ -278,39 +281,28 @@ fn truncate_function_output_items_with_policy_discards_empty_text() { } #[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(), - }, - ]; +fn formatted_text_output_reuses_owned_items_under_limit() { + let items = vec![text("alpha"), text(""), text("beta")]; + let policy = TruncationPolicy::Bytes(32); - let (output, original_token_count) = - formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(32)); + let borrowed = formatted_truncate_function_output_items_with_policy(&items, policy, |_| 0); + assert_eq!(borrowed, items); - assert_eq!(output, items); - assert_eq!(original_token_count, None); + let allocation = items.as_ptr(); + let owned = formatted_truncate_function_output_items_with_policy(items, policy, |_| 0); + assert_eq!(owned, borrowed); + assert_eq!(owned.as_ptr(), allocation); } #[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(), - }, - ]; +fn formatted_text_output_preserves_empty_leading_text_behavior() { + let items = vec![text(""), text("abc")]; - let (output, original_token_count) = - formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(0)); + let output = formatted_truncate_function_output_items_with_policy( + items, + TruncationPolicy::Bytes(0), + |_| 0, + ); assert_eq!( output, @@ -318,101 +310,63 @@ fn formatted_truncate_text_content_items_with_policy_preserves_empty_leading_tex text: "Warning: truncated output (original token count: 1)\nTotal 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_media() { +fn formatted_mixed_output_budgets_text_and_audio_in_order() { + let audio = FunctionCallOutputContentItem::InputAudio { + audio_url: "audio:one".to_owned(), + }; + let file_image = FunctionCallOutputContentItem::InputImage { + image: ImageReference::File { + file_id: "file_123".to_owned(), + }, + detail: Some(DEFAULT_IMAGE_DETAIL), + }; let items = vec![ - FunctionCallOutputContentItem::InputText { - text: "abcd".to_string(), - }, - FunctionCallOutputContentItem::InputImage { - image: ImageReference::Inline { - image_url: "img:one".to_string(), - }, - detail: Some(DEFAULT_IMAGE_DETAIL), - }, - FunctionCallOutputContentItem::InputText { - text: "efgh".to_string(), - }, - FunctionCallOutputContentItem::InputAudio { - audio_url: "audio:one".to_string(), - }, - FunctionCallOutputContentItem::InputImage { - image: ImageReference::File { - file_id: "file_123".to_string(), - }, - detail: Some(DEFAULT_IMAGE_DETAIL), - }, - FunctionCallOutputContentItem::InputText { - text: "ijkl".to_string(), - }, - FunctionCallOutputContentItem::InputImage { - image: ImageReference::Inline { - image_url: "img:two".to_string(), - }, - detail: Some(DEFAULT_IMAGE_DETAIL), - }, + text("abcd"), + inline_image("img:one".to_owned()), + text("efgh"), + audio.clone(), + file_image.clone(), + text("ijkl"), + inline_image("img:two".to_owned()), ]; - let (output, original_token_count) = - formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(8)); + let output = formatted_truncate_function_output_items_with_policy( + items, + TruncationPolicy::Bytes(12), + |_| 1, + ); assert_eq!( output, vec![ - FunctionCallOutputContentItem::InputText { - text: "Warning: truncated output (original token count: 4)\nTotal output lines: 3\n\nabcd…6 chars truncated…ijkl".to_string(), - }, - FunctionCallOutputContentItem::InputImage { - image: ImageReference::Inline { image_url: "img:one".to_string() }, - detail: Some(DEFAULT_IMAGE_DETAIL), - }, - FunctionCallOutputContentItem::InputAudio { - audio_url: "audio:one".to_string(), - }, - FunctionCallOutputContentItem::InputImage { - image: ImageReference::File { - file_id: "file_123".to_string(), - }, - detail: Some(DEFAULT_IMAGE_DETAIL), - }, - FunctionCallOutputContentItem::InputImage { - image: ImageReference::Inline { image_url: "img:two".to_string() }, - detail: Some(DEFAULT_IMAGE_DETAIL), - }, + text("abcd"), + inline_image("img:one".to_owned()), + text("efgh"), + audio, + file_image, + inline_image("img:two".to_owned()), + text("[omitted 1 text items ...]"), ] ); - assert_eq!(original_token_count, Some(4)); } #[test] -fn formatted_truncate_text_content_items_with_policy_preserves_encrypted_content() { - let items = vec![ - FunctionCallOutputContentItem::InputText { - text: "abcdefgh".to_string(), - }, - FunctionCallOutputContentItem::EncryptedContent { - encrypted_content: "enc_opaque".to_string(), - }, - ]; +fn formatted_mixed_output_preserves_encrypted_content() { + let encrypted = FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: "enc_opaque".to_owned(), + }; + let items = vec![text("abcdefgh"), encrypted.clone()]; - let (output, original_token_count) = - formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(2)); - - assert_eq!( - output, - vec![ - FunctionCallOutputContentItem::InputText { - text: "Warning: truncated output (original token count: 2)\nTotal output lines: 1\n\na…6 chars truncated…h".to_string(), - }, - FunctionCallOutputContentItem::EncryptedContent { - encrypted_content: "enc_opaque".to_string(), - }, - ] + let output = formatted_truncate_function_output_items_with_policy( + items, + TruncationPolicy::Bytes(2), + |_| 0, ); - assert_eq!(original_token_count, Some(2)); + + assert_eq!(output, vec![text("a…6 chars truncated…h"), encrypted]); } #[test] @@ -475,18 +429,14 @@ fn truncate_function_output_items_with_policy_charges_audio_against_byte_budget( } #[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(), - }, - ]; +fn formatted_text_output_merges_all_text_for_token_budget() { + let items = vec![text("abcdefgh"), text("ijklmnop")]; - let (output, original_token_count) = - formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Tokens(2)); + let output = formatted_truncate_function_output_items_with_policy( + items, + TruncationPolicy::Tokens(2), + |_| 0, + ); assert_eq!( output, @@ -494,7 +444,6 @@ fn formatted_truncate_text_content_items_with_policy_merges_all_text_for_token_b text: "Warning: truncated output (original token count: 5)\nTotal output lines: 2\n\nabcd…3 tokens truncated…mnop".to_string(), }] ); - assert_eq!(original_token_count, Some(5)); } #[test] @@ -503,3 +452,113 @@ fn byte_count_conversion_clamps_non_positive_values() { assert_eq!(approx_tokens_from_byte_count_i64(/*bytes*/ 0), 0); assert_eq!(approx_tokens_from_byte_count_i64(/*bytes*/ 5), 2); } + +fn boundaries(positions: [(usize, usize); N]) -> [OutputBoundary; N] { + positions.map(|(item_index, text_offset)| OutputBoundary { + item_index, + text_offset, + }) +} + +fn positions(boundaries: [OutputBoundary; N]) -> [(usize, usize); N] { + boundaries.map(|b| (b.item_index, b.text_offset)) +} + +fn text(s: &str) -> FunctionCallOutputContentItem { + FunctionCallOutputContentItem::InputText { text: s.to_owned() } +} + +fn inline_image(image_url: String) -> FunctionCallOutputContentItem { + FunctionCallOutputContentItem::InputImage { + image: ImageReference::Inline { image_url }, + detail: Some(DEFAULT_IMAGE_DETAIL), + } +} + +#[test] +fn boundary_formatters_normalize_input_and_remap_truncated_text() { + let audio_cost: fn(&str) -> usize = |_| 1; + let items = [text("left-head"), text("right-tail")]; + let max_tokens = 2; + let mut boundary = boundaries([(0, 0), (0, 2), (1, 0), (1, 8), (2, 0)]); + let output = + formatted_output_at_token_boundaries(&items, max_tokens, audio_cost, &mut boundary); + let [FunctionCallOutputContentItem::InputText { text: rendered }] = output.as_slice() else { + panic!("expected one merged text item") + }; + let head = rendered.find("left").expect("expected retained prefix"); + let tail = rendered + .strip_suffix("tail") + .expect("expected retained suffix") + .len(); + assert_eq!( + positions(boundary), + [ + (0, 0), + (0, head + 2), + (0, tail), + (0, tail + 2), + (0, rendered.len()) + ] + ); + + for (middle, middle_end) in [(text("after"), 5), (inline_image("img:one".to_owned()), 0)] { + let items = [text("🙂é boundary"), middle, text("last")]; + for formatter in [ + formatted_output_at_token_boundaries, + mixed_output_at_token_boundaries, + ] { + for tokens in [0, 1, 2, 256] { + let mut b = boundaries([(0, 1), (1, usize::MAX), (0, 8), (usize::MAX, usize::MAX)]); + let output = formatter(&items, tokens, audio_cost, &mut b); + assert_eq!(output, formatter(&items, tokens, audio_cost, &mut [])); + let result = positions(b); + assert!(result.is_sorted()); + assert!( + result + .iter() + .all(|&(index, offset)| match output.get(index) { + Some(FunctionCallOutputContentItem::InputText { text }) => + text.is_char_boundary(offset), + _ => index <= output.len() && offset == 0, + }) + ); + if tokens == 256 { + assert_eq!(result, [(0, 0), (1, middle_end), (1, middle_end), (3, 0)]); + } + } + } + } +} + +#[test] +fn owned_mixed_output_moves_a_preserved_image_across_text_truncation() { + let image_url = "data:image/png;base64,".to_owned() + &"a".repeat(4096); + let url_ptr = image_url.as_ptr(); + let items = vec![ + text("long prefix"), + inline_image(image_url), + text("omitted"), + ]; + let mut borrowed = boundaries([(0, 5), (2, 3)]); + let mut owned = boundaries([(0, 5), (2, 3)]); + let max_tokens = 1; + let expected = formatted_output_at_token_boundaries(&items, max_tokens, |_| 0, &mut borrowed); + let actual = formatted_output_at_token_boundaries(items, max_tokens, |_| 0, &mut owned); + assert_eq!(actual, expected); + let projected = positions(owned); + assert_eq!(projected, positions(borrowed)); + let [(index, offset), omitted] = projected; + let FunctionCallOutputContentItem::InputText { text: snippet } = &actual[index] else { + panic!("expected boundary inside truncated text"); + }; + assert_eq!((index, &snippet[offset..], omitted), (0, "ix", (2, 0))); + let FunctionCallOutputContentItem::InputImage { + image: ImageReference::Inline { image_url }, + .. + } = &actual[1] + else { + panic!("expected preserved image") + }; + assert_eq!(image_url.as_ptr(), url_ptr); +} diff --git a/codex-rs/utils/string/src/truncate.rs b/codex-rs/utils/string/src/truncate.rs index acad740803..6ee66db0ff 100644 --- a/codex-rs/utils/string/src/truncate.rs +++ b/codex-rs/utils/string/src/truncate.rs @@ -12,6 +12,17 @@ pub struct TextTruncation { pub marker: String, } +impl TextTruncation { + /// Applies this edit to the string it was planned from. + pub fn apply(&self, original: &str) -> String { + assemble_truncated_output( + &original[..self.removed.start], + &original[self.removed.end..], + &self.marker, + ) + } +} + /// Plans the existing head/tail token truncation. The marker is outside the /// retained-text budget, just as it is in [`truncate_middle_with_token_budget`]. pub fn plan_token_truncation(s: &str, max_tokens: usize) -> Option { @@ -70,11 +81,7 @@ pub fn truncate_middle_with_token_budget(s: &str, max_tokens: usize) -> (String, let Some(truncation) = plan_token_truncation(s, max_tokens) else { return (s.to_string(), None); }; - let truncated = assemble_truncated_output( - &s[..truncation.removed.start], - &s[truncation.removed.end..], - &truncation.marker, - ); + let truncated = truncation.apply(s); let total_tokens = u64::try_from(approx_token_count(s)).unwrap_or(u64::MAX); if truncated == s {