diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 8f8200caad..189b3aa7a5 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -1,8 +1,8 @@ use crate::codex::TurnContext; use crate::context_manager::normalize; -use crate::context_manager::truncate; -use crate::context_manager::truncate::format_output_for_model_body; -use crate::context_manager::truncate::globally_truncate_function_output_items; +use crate::truncate; +use crate::truncate::format_output_for_model_body; +use crate::truncate::globally_truncate_function_output_items; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::TokenUsage; diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index a1aeacd322..c81749c2c1 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -1,6 +1,6 @@ use super::*; use crate::context_manager::MODEL_FORMAT_MAX_LINES; -use crate::context_manager::truncate; +use crate::truncate; use codex_git::GhostCommit; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputContentItem; diff --git a/codex-rs/core/src/context_manager/mod.rs b/codex-rs/core/src/context_manager/mod.rs index d945cb95b0..ab0d2e8168 100644 --- a/codex-rs/core/src/context_manager/mod.rs +++ b/codex-rs/core/src/context_manager/mod.rs @@ -1,8 +1,7 @@ mod history; mod normalize; -mod truncate; +pub(crate) use crate::truncate::MODEL_FORMAT_MAX_BYTES; +pub(crate) use crate::truncate::MODEL_FORMAT_MAX_LINES; +pub(crate) use crate::truncate::format_output_for_model_body; pub(crate) use history::ContextManager; -pub(crate) use truncate::MODEL_FORMAT_MAX_BYTES; -pub(crate) use truncate::MODEL_FORMAT_MAX_LINES; -pub(crate) use truncate::format_output_for_model_body; diff --git a/codex-rs/core/src/context_manager/truncate.rs b/codex-rs/core/src/context_manager/truncate.rs deleted file mode 100644 index 88027ad15d..0000000000 --- a/codex-rs/core/src/context_manager/truncate.rs +++ /dev/null @@ -1,148 +0,0 @@ -use codex_protocol::models::FunctionCallOutputContentItem; -use codex_utils_string::take_bytes_at_char_boundary; -use codex_utils_string::take_last_bytes_at_char_boundary; - -use crate::util::error_or_panic; - -// Model-formatting limits: clients get full streams; only content sent to the model is truncated. -pub const MODEL_FORMAT_MAX_BYTES: usize = 10 * 1024; // 10 KiB -pub const MODEL_FORMAT_MAX_LINES: usize = 256; // lines - -pub(crate) fn globally_truncate_function_output_items( - items: &[FunctionCallOutputContentItem], -) -> Vec { - let mut out: Vec = Vec::with_capacity(items.len()); - let mut remaining = MODEL_FORMAT_MAX_BYTES; - let mut omitted_text_items = 0usize; - - for it in items { - match it { - FunctionCallOutputContentItem::InputText { text } => { - if remaining == 0 { - omitted_text_items += 1; - continue; - } - - let len = text.len(); - if len <= remaining { - out.push(FunctionCallOutputContentItem::InputText { text: text.clone() }); - remaining -= len; - } else { - let slice = take_bytes_at_char_boundary(text, remaining); - if !slice.is_empty() { - out.push(FunctionCallOutputContentItem::InputText { - text: slice.to_string(), - }); - } - remaining = 0; - } - } - // todo(aibrahim): handle input images; resize - FunctionCallOutputContentItem::InputImage { image_url } => { - out.push(FunctionCallOutputContentItem::InputImage { - image_url: image_url.clone(), - }); - } - } - } - - if omitted_text_items > 0 { - out.push(FunctionCallOutputContentItem::InputText { - text: format!("[omitted {omitted_text_items} text items ...]"), - }); - } - - out -} - -pub(crate) fn format_output_for_model_body( - content: &str, - limit_bytes: usize, - limit_lines: usize, -) -> String { - // Head+tail truncation for the model: show the beginning and end with an elision. - // Clients still receive full streams; only this formatted summary is capped. - let total_lines = content.lines().count(); - if content.len() <= limit_bytes && total_lines <= limit_lines { - return content.to_string(); - } - let output = truncate_formatted_exec_output(content, total_lines, limit_bytes, limit_lines); - format!("Total output lines: {total_lines}\n\n{output}") -} - -fn truncate_formatted_exec_output( - content: &str, - total_lines: usize, - limit_bytes: usize, - limit_lines: usize, -) -> String { - debug_panic_on_double_truncation(content); - let head_lines: usize = limit_lines / 2; - let tail_lines: usize = limit_lines - head_lines; // 128 - let head_bytes: usize = limit_bytes / 2; - let segments: Vec<&str> = content.split_inclusive('\n').collect(); - let head_take = head_lines.min(segments.len()); - let tail_take = tail_lines.min(segments.len().saturating_sub(head_take)); - let omitted = segments.len().saturating_sub(head_take + tail_take); - - let head_slice_end: usize = segments - .iter() - .take(head_take) - .map(|segment| segment.len()) - .sum(); - let tail_slice_start: usize = if tail_take == 0 { - content.len() - } else { - content.len() - - segments - .iter() - .rev() - .take(tail_take) - .map(|segment| segment.len()) - .sum::() - }; - let head_slice = &content[..head_slice_end]; - let tail_slice = &content[tail_slice_start..]; - let truncated_by_bytes = content.len() > limit_bytes; - // this is a bit wrong. We are counting metadata lines and not just shell output lines. - let marker = if omitted > 0 { - Some(format!( - "\n[... omitted {omitted} of {total_lines} lines ...]\n\n" - )) - } else if truncated_by_bytes { - Some(format!( - "\n[... output truncated to fit {limit_bytes} bytes ...]\n\n" - )) - } else { - None - }; - - let marker_len = marker.as_ref().map_or(0, String::len); - let base_head_budget = head_bytes.min(limit_bytes); - let head_budget = base_head_budget.min(limit_bytes.saturating_sub(marker_len)); - let head_part = take_bytes_at_char_boundary(head_slice, head_budget); - let mut result = String::with_capacity(limit_bytes.min(content.len())); - - result.push_str(head_part); - if let Some(marker_text) = marker.as_ref() { - result.push_str(marker_text); - } - - let remaining = limit_bytes.saturating_sub(result.len()); - if remaining == 0 { - return result; - } - - let tail_part = take_last_bytes_at_char_boundary(tail_slice, remaining); - result.push_str(tail_part); - - result -} - -fn debug_panic_on_double_truncation(content: &str) { - if content.contains("Total output lines:") && content.contains("omitted") { - error_or_panic(format!( - "FunctionCallOutput content was already truncated before ContextManager::record_items; this would cause double truncation {content}" - )); - } -} diff --git a/codex-rs/core/src/truncate.rs b/codex-rs/core/src/truncate.rs index 3f0be8fcaf..c82e088560 100644 --- a/codex-rs/core/src/truncate.rs +++ b/codex-rs/core/src/truncate.rs @@ -1,8 +1,195 @@ //! Utilities for truncating large chunks of output while preserving a prefix -//! and suffix on UTF-8 boundaries. +//! and suffix on UTF-8 boundaries, and helpers for line/token‑based truncation +//! used across the core crate. +use crate::util::error_or_panic; +use codex_protocol::models::FunctionCallOutputContentItem; +use codex_utils_string::take_bytes_at_char_boundary; +use codex_utils_string::take_last_bytes_at_char_boundary; use codex_utils_tokenizer::Tokenizer; +/// Model-formatting limits: clients get full streams; only content sent to the model is truncated. +pub const MODEL_FORMAT_MAX_BYTES: usize = 10 * 1024; // 10 KiB +pub const MODEL_FORMAT_MAX_LINES: usize = 256; // lines + +/// Globally truncate function output items to fit within `MODEL_FORMAT_MAX_BYTES` +/// by preserving as many text/image items as possible and appending a summary +/// for any omitted text items. +pub(crate) fn globally_truncate_function_output_items( + items: &[FunctionCallOutputContentItem], +) -> Vec { + let mut out: Vec = Vec::with_capacity(items.len()); + let mut remaining = MODEL_FORMAT_MAX_BYTES; + let mut omitted_text_items = 0usize; + + for it in items { + match it { + FunctionCallOutputContentItem::InputText { text } => { + if remaining == 0 { + omitted_text_items += 1; + continue; + } + + let len = text.len(); + if len <= remaining { + out.push(FunctionCallOutputContentItem::InputText { text: text.clone() }); + remaining -= len; + } else { + let slice = take_bytes_at_char_boundary(text, remaining); + if !slice.is_empty() { + out.push(FunctionCallOutputContentItem::InputText { + text: slice.to_string(), + }); + } + remaining = 0; + } + } + // todo(aibrahim): handle input images; resize + FunctionCallOutputContentItem::InputImage { image_url } => { + out.push(FunctionCallOutputContentItem::InputImage { + image_url: image_url.clone(), + }); + } + } + } + + if omitted_text_items > 0 { + out.push(FunctionCallOutputContentItem::InputText { + text: format!("[omitted {omitted_text_items} text items ...]"), + }); + } + + out +} + +/// Format a block of exec/tool output for model consumption, truncating by +/// lines and bytes while preserving head and tail segments. +pub(crate) fn format_output_for_model_body( + content: &str, + limit_bytes: usize, + limit_lines: usize, +) -> String { + // Head+tail truncation for the model: show the beginning and end with an elision. + // Clients still receive full streams; only this formatted summary is capped. + let total_lines = content.lines().count(); + if content.len() <= limit_bytes && total_lines <= limit_lines { + return content.to_string(); + } + let output = truncate_formatted_exec_output(content, total_lines, limit_bytes, limit_lines); + format!("Total output lines: {total_lines}\n\n{output}") +} + +fn truncate_formatted_exec_output( + content: &str, + total_lines: usize, + limit_bytes: usize, + limit_lines: usize, +) -> String { + debug_panic_on_double_truncation(content); + let head_lines: usize = limit_lines / 2; + let tail_lines: usize = limit_lines - head_lines; // 128 + let head_bytes: usize = limit_bytes / 2; + let segments: Vec<&str> = content.split_inclusive('\n').collect(); + let head_take = head_lines.min(segments.len()); + let tail_take = tail_lines.min(segments.len().saturating_sub(head_take)); + let omitted = segments.len().saturating_sub(head_take + tail_take); + + let head_slice_end: usize = segments + .iter() + .take(head_take) + .map(|segment| segment.len()) + .sum(); + let tail_slice_start: usize = if tail_take == 0 { + content.len() + } else { + content.len() + - segments + .iter() + .rev() + .take(tail_take) + .map(|segment| segment.len()) + .sum::() + }; + let head_slice = &content[..head_slice_end]; + let tail_slice = &content[tail_slice_start..]; + let truncated_by_bytes = content.len() > limit_bytes; + // this is a bit wrong. We are counting metadata lines and not just shell output lines. + let marker = if omitted > 0 { + Some(format!( + "\n[... omitted {omitted} of {total_lines} lines ...]\n\n" + )) + } else if truncated_by_bytes { + Some(format!( + "\n[... output truncated to fit {limit_bytes} bytes ...]\n\n" + )) + } else { + None + }; + + let marker_len = marker.as_ref().map_or(0, String::len); + let base_head_budget = head_bytes.min(limit_bytes); + let head_budget = base_head_budget.min(limit_bytes.saturating_sub(marker_len)); + let head_part = take_bytes_at_char_boundary(head_slice, head_budget); + let mut result = String::with_capacity(limit_bytes.min(content.len())); + + result.push_str(head_part); + if let Some(marker_text) = marker.as_ref() { + result.push_str(marker_text); + } + + let remaining = limit_bytes.saturating_sub(result.len()); + if remaining == 0 { + return result; + } + + let tail_part = take_last_bytes_at_char_boundary(tail_slice, remaining); + result.push_str(tail_part); + + result +} + +fn debug_panic_on_double_truncation(content: &str) { + if content.contains("Total output lines:") && content.contains("omitted") { + error_or_panic(format!( + "FunctionCallOutput content was already truncated before ContextManager::record_items; this would cause double truncation {content}" + )); + } +} + +/// Truncate an output string to a maximum number of “tokens”, where tokens are +/// approximated as individual `char`s. Preserves a prefix and suffix with an +/// elision marker describing how many tokens were omitted. +pub(crate) fn truncate_output_to_tokens( + output: &str, + max_tokens: usize, +) -> (String, Option) { + if max_tokens == 0 { + let total_tokens = output.chars().count(); + let message = format!("…{total_tokens} tokens truncated…"); + return (message, Some(total_tokens)); + } + + let tokens: Vec = output.chars().collect(); + let total_tokens = tokens.len(); + if total_tokens <= max_tokens { + return (output.to_string(), None); + } + + let half = max_tokens / 2; + if half == 0 { + let truncated = total_tokens.saturating_sub(max_tokens); + let message = format!("…{truncated} tokens truncated…"); + return (message, Some(total_tokens)); + } + + let truncated = total_tokens.saturating_sub(half * 2); + let mut truncated_output = String::new(); + truncated_output.extend(&tokens[..half]); + truncated_output.push_str(&format!("…{truncated} tokens truncated…")); + truncated_output.extend(&tokens[total_tokens - half..]); + (truncated_output, Some(total_tokens)) +} + /// Truncate the middle of a UTF-8 string to at most `max_bytes` bytes, /// preserving the beginning and the end. Returns the possibly truncated /// string and `Some(original_token_count)` (counted with the local tokenizer; @@ -136,8 +323,33 @@ pub(crate) fn truncate_middle(s: &str, max_bytes: usize) -> (String, Option #[cfg(test)] mod tests { + use super::MODEL_FORMAT_MAX_BYTES; + use super::MODEL_FORMAT_MAX_LINES; + use super::format_output_for_model_body; + use super::globally_truncate_function_output_items; use super::truncate_middle; + use super::truncate_output_to_tokens; + use codex_protocol::models::FunctionCallOutputContentItem; use codex_utils_tokenizer::Tokenizer; + use pretty_assertions::assert_eq; + use regex_lite::Regex; + + fn truncated_message_pattern(line: &str, total_lines: usize) -> String { + let head_lines = MODEL_FORMAT_MAX_LINES / 2; + let tail_lines = MODEL_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); + let escaped_line = regex_lite::escape(line); + if omitted == 0 { + return format!( + r"(?s)^Total output lines: {total_lines}\n\n(?P{escaped_line}.*\n\[\.{{3}} output truncated to fit {MODEL_FORMAT_MAX_BYTES} bytes \.{{3}}]\n\n.*)$", + ); + } + format!( + r"(?s)^Total output lines: {total_lines}\n\n(?P{escaped_line}.*\n\[\.{{3}} omitted {omitted} of {total_lines} lines \.{{3}}]\n\n.*)$", + ) + } #[test] fn truncate_middle_no_newlines_fallback() { @@ -197,4 +409,208 @@ mod tests { assert!(out.ends_with("017\n018\n019\n020\n")); assert_eq!(total, Some(tok.count(&s) as u64)); } + + #[test] + fn truncate_output_to_tokens_returns_original_when_under_limit() { + let s = "short output"; + let (truncated, original) = truncate_output_to_tokens(s, 100); + assert_eq!(truncated, s); + assert_eq!(original, None); + } + + #[test] + fn truncate_output_to_tokens_reports_truncation_at_zero_limit() { + let s = "abcdef"; + let (truncated, original) = truncate_output_to_tokens(s, 0); + assert!(truncated.contains("tokens truncated")); + assert_eq!(original, Some(s.chars().count())); + } + + #[test] + fn truncate_output_to_tokens_preserves_prefix_and_suffix() { + let s = "abcdefghijklmnopqrstuvwxyz"; + let max_tokens = 10; + let (truncated, original) = truncate_output_to_tokens(s, max_tokens); + assert!(truncated.starts_with("abcde")); + assert!(truncated.ends_with("vwxyz")); + assert_eq!(original, Some(s.chars().count())); + } + + #[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 = format_output_for_model_body( + &large_error, + MODEL_FORMAT_MAX_BYTES, + MODEL_FORMAT_MAX_LINES, + ); + + let total_lines = large_error.lines().count(); + let pattern = truncated_message_pattern(line, total_lines); + let regex = Regex::new(&pattern).unwrap_or_else(|err| { + panic!("failed to compile regex {pattern}: {err}"); + }); + let captures = regex + .captures(&truncated) + .unwrap_or_else(|| panic!("message failed to match pattern {pattern}: {truncated}")); + let body = captures + .name("body") + .expect("missing body capture") + .as_str(); + assert!( + body.len() <= MODEL_FORMAT_MAX_BYTES, + "body exceeds byte limit: {} bytes", + body.len() + ); + assert_ne!(truncated, large_error); + } + + #[test] + fn format_exec_output_marks_byte_truncation_without_omitted_lines() { + let long_line = "a".repeat(MODEL_FORMAT_MAX_BYTES + 50); + let truncated = format_output_for_model_body( + &long_line, + MODEL_FORMAT_MAX_BYTES, + MODEL_FORMAT_MAX_LINES, + ); + + assert_ne!(truncated, long_line); + let marker_line = + format!("[... output truncated to fit {MODEL_FORMAT_MAX_BYTES} bytes ...]"); + assert!( + truncated.contains(&marker_line), + "missing byte truncation marker: {truncated}" + ); + assert!( + !truncated.contains("omitted"), + "line omission marker should not appear when no lines were dropped: {truncated}" + ); + } + + #[test] + fn format_exec_output_returns_original_when_within_limits() { + let content = "example output\n".repeat(10); + + assert_eq!( + format_output_for_model_body(&content, MODEL_FORMAT_MAX_BYTES, MODEL_FORMAT_MAX_LINES), + content + ); + } + + #[test] + fn format_exec_output_reports_omitted_lines_and_keeps_head_and_tail() { + let total_lines = MODEL_FORMAT_MAX_LINES + 100; + let content: String = (0..total_lines) + .map(|idx| format!("line-{idx}\n")) + .collect(); + + let truncated = + format_output_for_model_body(&content, MODEL_FORMAT_MAX_BYTES, MODEL_FORMAT_MAX_LINES); + + let omitted = total_lines - MODEL_FORMAT_MAX_LINES; + let expected_marker = format!("[... omitted {omitted} of {total_lines} lines ...]"); + + assert!( + truncated.contains(&expected_marker), + "missing omitted marker: {truncated}" + ); + assert!( + truncated.contains("line-0\n"), + "expected head line to remain: {truncated}" + ); + + let last_line = format!("line-{}\n", total_lines - 1); + assert!( + truncated.contains(&last_line), + "expected tail line to remain: {truncated}" + ); + } + + #[test] + fn format_exec_output_prefers_line_marker_when_both_limits_exceeded() { + let total_lines = MODEL_FORMAT_MAX_LINES + 42; + let long_line = "x".repeat(256); + let content: String = (0..total_lines) + .map(|idx| format!("line-{idx}-{long_line}\n")) + .collect(); + + let truncated = + format_output_for_model_body(&content, MODEL_FORMAT_MAX_BYTES, MODEL_FORMAT_MAX_LINES); + + assert!( + truncated.contains("[... omitted 42 of 298 lines ...]"), + "expected omitted marker when line count exceeds limit: {truncated}" + ); + assert!( + !truncated.contains("output truncated to fit"), + "line omission marker should take precedence over byte marker: {truncated}" + ); + } + + #[test] + fn truncates_across_multiple_under_limit_texts_and_reports_omitted() { + // Arrange: several text items, none exceeding per-item limit, but total exceeds budget. + let budget = MODEL_FORMAT_MAX_BYTES; + let t1_len = (budget / 2).saturating_sub(10); + let t2_len = (budget / 2).saturating_sub(10); + let remaining_after_t1_t2 = budget.saturating_sub(t1_len + t2_len); + let t3_len = 50; // gets truncated to remaining_after_t1_t2 + let t4_len = 5; // omitted + let t5_len = 7; // omitted + + let t1 = "a".repeat(t1_len); + let t2 = "b".repeat(t2_len); + let t3 = "c".repeat(t3_len); + let t4 = "d".repeat(t4_len); + let t5 = "e".repeat(t5_len); + + let items = vec![ + FunctionCallOutputContentItem::InputText { text: t1 }, + FunctionCallOutputContentItem::InputText { text: t2 }, + FunctionCallOutputContentItem::InputImage { + image_url: "img:mid".to_string(), + }, + FunctionCallOutputContentItem::InputText { text: t3 }, + FunctionCallOutputContentItem::InputText { text: t4 }, + FunctionCallOutputContentItem::InputText { text: t5 }, + ]; + + let output = globally_truncate_function_output_items(&items); + + // Expect: t1 (full), t2 (full), image, t3 (truncated), summary mentioning 2 omitted. + 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.len(), t1_len); + + let second_text = match &output[1] { + FunctionCallOutputContentItem::InputText { text } => text, + other => panic!("unexpected second item: {other:?}"), + }; + assert_eq!(second_text.len(), t2_len); + + assert_eq!( + output[2], + FunctionCallOutputContentItem::InputImage { + image_url: "img:mid".to_string() + } + ); + + let fourth_text = match &output[3] { + FunctionCallOutputContentItem::InputText { text } => text, + other => panic!("unexpected fourth item: {other:?}"), + }; + assert_eq!(fourth_text.len(), remaining_after_t1_t2); + + let summary_text = match &output[4] { + FunctionCallOutputContentItem::InputText { text } => text, + other => panic!("unexpected summary item: {other:?}"), + }; + assert!(summary_text.contains("omitted 2 text items")); + } } diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 8674564f35..f77b744497 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -124,37 +124,6 @@ pub(crate) fn generate_chunk_id() -> String { .collect() } -pub(crate) fn truncate_output_to_tokens( - output: &str, - max_tokens: usize, -) -> (String, Option) { - if max_tokens == 0 { - let total_tokens = output.chars().count(); - let message = format!("…{total_tokens} tokens truncated…"); - return (message, Some(total_tokens)); - } - - let tokens: Vec = output.chars().collect(); - let total_tokens = tokens.len(); - if total_tokens <= max_tokens { - return (output.to_string(), None); - } - - let half = max_tokens / 2; - if half == 0 { - let truncated = total_tokens.saturating_sub(max_tokens); - let message = format!("…{truncated} tokens truncated…"); - return (message, Some(total_tokens)); - } - - let truncated = total_tokens.saturating_sub(half * 2); - let mut truncated_output = String::new(); - truncated_output.extend(&tokens[..half]); - truncated_output.push_str(&format!("…{truncated} tokens truncated…")); - truncated_output.extend(&tokens[total_tokens - half..]); - (truncated_output, Some(total_tokens)) -} - #[cfg(test)] #[cfg(unix)] mod tests { diff --git a/codex-rs/core/src/unified_exec/session_manager.rs b/codex-rs/core/src/unified_exec/session_manager.rs index e7daee88ed..fee46df8b8 100644 --- a/codex-rs/core/src/unified_exec/session_manager.rs +++ b/codex-rs/core/src/unified_exec/session_manager.rs @@ -36,7 +36,7 @@ use super::generate_chunk_id; use super::resolve_max_tokens; use super::session::OutputBuffer; use super::session::UnifiedExecSession; -use super::truncate_output_to_tokens; +use crate::truncate::truncate_output_to_tokens; impl UnifiedExecSessionManager { pub(crate) async fn exec_command(