From c4c673164ed4d7fc32a3a19e139f9456e5ec4f6c Mon Sep 17 00:00:00 2001 From: jimmyfraiture Date: Sun, 7 Sep 2025 18:28:54 -0700 Subject: [PATCH] V9 --- .../core/src/exec_command/session_manager.rs | 153 +-------------- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/truncate.rs | 180 ++++++++++++++++++ codex-rs/core/src/unified_exec/mod.rs | 81 ++++---- codex-rs/core/src/unified_exec/path.rs | 77 +++++++- codex-rs/core/src/unified_exec/truncation.rs | 71 ------- 6 files changed, 291 insertions(+), 272 deletions(-) create mode 100644 codex-rs/core/src/truncate.rs delete mode 100644 codex-rs/core/src/unified_exec/truncation.rs diff --git a/codex-rs/core/src/exec_command/session_manager.rs b/codex-rs/core/src/exec_command/session_manager.rs index 9741b44b4e..c1253722db 100644 --- a/codex-rs/core/src/exec_command/session_manager.rs +++ b/codex-rs/core/src/exec_command/session_manager.rs @@ -19,6 +19,7 @@ use crate::exec_command::exec_command_params::ExecCommandParams; use crate::exec_command::exec_command_params::WriteStdinParams; use crate::exec_command::exec_command_session::ExecCommandSession; use crate::exec_command::session_id::SessionId; +use crate::truncate::truncate_middle; use codex_protocol::models::FunctionCallOutputPayload; #[derive(Debug, Default)] @@ -351,112 +352,6 @@ async fn create_exec_command_session( Ok((session, exit_rx)) } -/// 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)` (estimated at 4 bytes/token) -/// if truncation occurred; otherwise returns the original string and `None`. -fn truncate_middle(s: &str, max_bytes: usize) -> (String, Option) { - // No truncation needed - if s.len() <= max_bytes { - return (s.to_string(), None); - } - let est_tokens = (s.len() as u64).div_ceil(4); - if max_bytes == 0 { - // Cannot keep any content; still return a full marker (never truncated). - return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); - } - - // Helper to truncate a string to a given byte length on a char boundary. - fn truncate_on_boundary(input: &str, max_len: usize) -> &str { - if input.len() <= max_len { - return input; - } - let mut end = max_len; - while end > 0 && !input.is_char_boundary(end) { - end -= 1; - } - &input[..end] - } - - // Given a left/right budget, prefer newline boundaries; otherwise fall back - // to UTF-8 char boundaries. - fn pick_prefix_end(s: &str, left_budget: usize) -> usize { - if let Some(head) = s.get(..left_budget) - && let Some(i) = head.rfind('\n') - { - return i + 1; // keep the newline so suffix starts on a fresh line - } - truncate_on_boundary(s, left_budget).len() - } - - fn pick_suffix_start(s: &str, right_budget: usize) -> usize { - let start_tail = s.len().saturating_sub(right_budget); - if let Some(tail) = s.get(start_tail..) - && let Some(i) = tail.find('\n') - { - return start_tail + i + 1; // start after newline - } - // Fall back to a char boundary at or after start_tail. - let mut idx = start_tail.min(s.len()); - while idx < s.len() && !s.is_char_boundary(idx) { - idx += 1; - } - idx - } - - // Refine marker length and budgets until stable. Marker is never truncated. - let mut guess_tokens = est_tokens; // worst-case: everything truncated - for _ in 0..4 { - let marker = format!("…{guess_tokens} tokens truncated…"); - let marker_len = marker.len(); - let keep_budget = max_bytes.saturating_sub(marker_len); - if keep_budget == 0 { - // No room for any content within the cap; return a full, untruncated marker - // that reflects the entire truncated content. - return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); - } - - let left_budget = keep_budget / 2; - let right_budget = keep_budget - left_budget; - let prefix_end = pick_prefix_end(s, left_budget); - let mut suffix_start = pick_suffix_start(s, right_budget); - if suffix_start < prefix_end { - suffix_start = prefix_end; - } - let kept_content_bytes = prefix_end + (s.len() - suffix_start); - let truncated_content_bytes = s.len().saturating_sub(kept_content_bytes); - let new_tokens = (truncated_content_bytes as u64).div_ceil(4); - if new_tokens == guess_tokens { - let mut out = String::with_capacity(marker_len + kept_content_bytes + 1); - out.push_str(&s[..prefix_end]); - out.push_str(&marker); - // Place marker on its own line for symmetry when we keep line boundaries. - out.push('\n'); - out.push_str(&s[suffix_start..]); - return (out, Some(est_tokens)); - } - guess_tokens = new_tokens; - } - - // Fallback: use last guess to build output. - let marker = format!("…{guess_tokens} tokens truncated…"); - let marker_len = marker.len(); - let keep_budget = max_bytes.saturating_sub(marker_len); - if keep_budget == 0 { - return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); - } - let left_budget = keep_budget / 2; - let right_budget = keep_budget - left_budget; - let prefix_end = pick_prefix_end(s, left_budget); - let suffix_start = pick_suffix_start(s, right_budget); - let mut out = String::with_capacity(marker_len + prefix_end + (s.len() - suffix_start) + 1); - out.push_str(&s[..prefix_end]); - out.push_str(&marker); - out.push('\n'); - out.push_str(&s[suffix_start..]); - (out, Some(est_tokens)) -} - #[cfg(test)] mod tests { use super::*; @@ -620,50 +515,4 @@ Output: abc"#; assert_eq!(expected, text); } - - #[test] - fn truncate_middle_no_newlines_fallback() { - // A long string with no newlines that exceeds the cap. - let s = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"; - let max_bytes = 16; // force truncation - let (out, original) = truncate_middle(s, max_bytes); - // For very small caps, we return the full, untruncated marker, - // even if it exceeds the cap. - assert_eq!(out, "…16 tokens truncated…"); - // Original string length is 62 bytes => ceil(62/4) = 16 tokens. - assert_eq!(original, Some(16)); - } - - #[test] - fn truncate_middle_prefers_newline_boundaries() { - // Build a multi-line string of 20 numbered lines (each "NNN\n"). - let mut s = String::new(); - for i in 1..=20 { - s.push_str(&format!("{i:03}\n")); - } - // Total length: 20 lines * 4 bytes per line = 80 bytes. - assert_eq!(s.len(), 80); - - // Choose a cap that forces truncation while leaving room for - // a few lines on each side after accounting for the marker. - let max_bytes = 64; - // Expect exact output: first 4 lines, marker, last 4 lines, and correct token estimate (80/4 = 20). - assert_eq!( - truncate_middle(&s, max_bytes), - ( - r#"001 -002 -003 -004 -…12 tokens truncated… -017 -018 -019 -020 -"# - .to_string(), - Some(20) - ) - ); - } } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 0b9afa4923..d63f5c05e1 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -34,6 +34,7 @@ mod mcp_tool_call; mod message_history; mod model_provider_info; pub mod parse_command; +mod truncate; mod unified_exec; mod user_instructions; pub use model_provider_info::BUILT_IN_OSS_MODEL_PROVIDER_ID; diff --git a/codex-rs/core/src/truncate.rs b/codex-rs/core/src/truncate.rs new file mode 100644 index 0000000000..ab0158720b --- /dev/null +++ b/codex-rs/core/src/truncate.rs @@ -0,0 +1,180 @@ +//! Utilities for truncating large chunks of output while preserving a prefix +//! and suffix on UTF-8 boundaries. + +/// 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)` (estimated at 4 bytes/token) +/// if truncation occurred; otherwise returns the original string and `None`. +pub(crate) fn truncate_middle(s: &str, max_bytes: usize) -> (String, Option) { + if s.len() <= max_bytes { + return (s.to_string(), None); + } + + let est_tokens = (s.len() as u64).div_ceil(4); + if max_bytes == 0 { + return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); + } + + fn truncate_on_boundary(input: &str, max_len: usize) -> &str { + if input.len() <= max_len { + return input; + } + let mut end = max_len; + while end > 0 && !input.is_char_boundary(end) { + end -= 1; + } + &input[..end] + } + + fn pick_prefix_end(s: &str, left_budget: usize) -> usize { + if let Some(head) = s.get(..left_budget) + && let Some(i) = head.rfind('\n') + { + return i + 1; + } + truncate_on_boundary(s, left_budget).len() + } + + fn pick_suffix_start(s: &str, right_budget: usize) -> usize { + let start_tail = s.len().saturating_sub(right_budget); + if let Some(tail) = s.get(start_tail..) + && let Some(i) = tail.find('\n') + { + return start_tail + i + 1; + } + + let mut idx = start_tail.min(s.len()); + while idx < s.len() && !s.is_char_boundary(idx) { + idx += 1; + } + idx + } + + let mut guess_tokens = est_tokens; + for _ in 0..4 { + let marker = format!("…{guess_tokens} tokens truncated…"); + let marker_len = marker.len(); + let keep_budget = max_bytes.saturating_sub(marker_len); + if keep_budget == 0 { + return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); + } + + let left_budget = keep_budget / 2; + let right_budget = keep_budget - left_budget; + let prefix_end = pick_prefix_end(s, left_budget); + let mut suffix_start = pick_suffix_start(s, right_budget); + if suffix_start < prefix_end { + suffix_start = prefix_end; + } + + let kept_content_bytes = prefix_end + (s.len() - suffix_start); + let truncated_content_bytes = s.len().saturating_sub(kept_content_bytes); + let new_tokens = (truncated_content_bytes as u64).div_ceil(4); + + if new_tokens == guess_tokens { + let mut out = String::with_capacity(marker_len + kept_content_bytes + 1); + out.push_str(&s[..prefix_end]); + out.push_str(&marker); + out.push('\n'); + out.push_str(&s[suffix_start..]); + return (out, Some(est_tokens)); + } + + guess_tokens = new_tokens; + } + + let marker = format!("…{guess_tokens} tokens truncated…"); + let marker_len = marker.len(); + let keep_budget = max_bytes.saturating_sub(marker_len); + if keep_budget == 0 { + return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); + } + + let left_budget = keep_budget / 2; + let right_budget = keep_budget - left_budget; + let prefix_end = pick_prefix_end(s, left_budget); + let suffix_start = pick_suffix_start(s, right_budget); + + let mut out = String::with_capacity(marker_len + prefix_end + (s.len() - suffix_start) + 1); + out.push_str(&s[..prefix_end]); + out.push_str(&marker); + out.push('\n'); + out.push_str(&s[suffix_start..]); + (out, Some(est_tokens)) +} + +#[cfg(test)] +mod tests { + use super::truncate_middle; + + #[test] + fn truncate_middle_no_newlines_fallback() { + let s = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ*"; + let max_bytes = 32; + let (out, original) = truncate_middle(s, max_bytes); + assert!(out.starts_with("abc")); + assert!(out.contains("tokens truncated")); + assert!(out.ends_with("XYZ*")); + assert_eq!(original, Some((s.len() as u64).div_ceil(4))); + } + + #[test] + fn truncate_middle_prefers_newline_boundaries() { + let mut s = String::new(); + for i in 1..=20 { + s.push_str(&format!("{i:03}\n")); + } + assert_eq!(s.len(), 80); + + let max_bytes = 64; + let (out, tokens) = truncate_middle(&s, max_bytes); + assert!(out.starts_with("001\n002\n003\n004\n")); + assert!(out.contains("tokens truncated")); + assert!(out.ends_with("017\n018\n019\n020\n")); + assert_eq!(tokens, Some(20)); + } + + #[test] + fn truncate_middle_handles_utf8_content() { + let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with ascii text\n"; + let max_bytes = 32; + let (out, tokens) = truncate_middle(s, max_bytes); + + assert!(out.contains("tokens truncated")); + assert!(!out.contains('\u{fffd}')); + assert_eq!(tokens, Some((s.len() as u64).div_ceil(4))); + } + + #[test] + fn truncate_middle_prefers_newline_boundaries_2() { + // Build a multi-line string of 20 numbered lines (each "NNN\n"). + let mut s = String::new(); + for i in 1..=20 { + s.push_str(&format!("{i:03}\n")); + } + // Total length: 20 lines * 4 bytes per line = 80 bytes. + assert_eq!(s.len(), 80); + + // Choose a cap that forces truncation while leaving room for + // a few lines on each side after accounting for the marker. + let max_bytes = 64; + // Expect exact output: first 4 lines, marker, last 4 lines, and correct token estimate (80/4 = 20). + assert_eq!( + truncate_middle(&s, max_bytes), + ( + r#"001 +002 +003 +004 +…12 tokens truncated… +017 +018 +019 +020 +"# + .to_string(), + Some(20) + ) + ); + } +} diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 4b3ad68b9f..10ab10cc32 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -18,17 +18,16 @@ use tokio::time::Duration; use tokio::time::Instant; use crate::exec_command::ExecCommandSession; +use crate::truncate::truncate_middle; mod errors; mod path; -mod truncation; pub(crate) use errors::UnifiedExecError; use path::command_from_chunks; use path::join_input_chunks; use path::resolve_command_path; -use truncation::truncate_middle; const DEFAULT_TIMEOUT_MS: u64 = 250; const MAX_TIMEOUT_MS: u64 = 60_000; @@ -61,12 +60,40 @@ struct ManagedUnifiedExecSession { output_task: JoinHandle<()>, } -type OutputBuffer = Arc>>>; +#[derive(Debug, Default)] +struct OutputBufferState { + chunks: VecDeque>, + total_bytes: usize, +} + +impl OutputBufferState { + fn push_chunk(&mut self, chunk: Vec) { + self.total_bytes = self.total_bytes.saturating_add(chunk.len()); + self.chunks.push_back(chunk); + + while self.total_bytes > UNIFIED_EXEC_OUTPUT_MAX_BYTES { + match self.chunks.pop_front() { + Some(removed) => { + self.total_bytes = self.total_bytes.saturating_sub(removed.len()); + } + None => break, + } + } + } + + fn drain(&mut self) -> Vec> { + let drained: Vec> = self.chunks.drain(..).collect(); + self.total_bytes = 0; + drained + } +} + +type OutputBuffer = Arc>; type OutputHandles = (OutputBuffer, Arc); impl ManagedUnifiedExecSession { fn new(session: ExecCommandSession) -> Self { - let output_buffer = Arc::new(Mutex::new(VecDeque::new())); + let output_buffer = Arc::new(Mutex::new(OutputBufferState::default())); let output_notify = Arc::new(Notify::new()); let mut receiver = session.output_receiver(); let buffer_clone = Arc::clone(&output_buffer); @@ -74,7 +101,7 @@ impl ManagedUnifiedExecSession { let output_task = tokio::spawn(async move { while let Ok(chunk) = receiver.recv().await { let mut guard = buffer_clone.lock().await; - guard.push_back(chunk); + guard.push_chunk(chunk); drop(guard); notify_clone.notify_waiters(); } @@ -183,11 +210,7 @@ impl UnifiedExecSessionManager { loop { let drained_chunks = { let mut guard = output_buffer.lock().await; - let mut drained = Vec::new(); - while let Some(chunk) = guard.pop_front() { - drained.push(chunk); - } - drained + guard.drain() }; if drained_chunks.is_empty() { @@ -609,42 +632,4 @@ mod tests { Ok(()) } - - #[test] - fn truncate_middle_no_newlines_fallback() { - let s = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ*"; - let max_bytes = 32; - let (out, original) = truncate_middle(s, max_bytes); - assert_eq!(out, "abcdef[TRUNCATED CONTENT]UVWXYZ*"); - assert_eq!(original, Some((s.len() as u64).div_ceil(4))); - } - - #[test] - fn truncate_middle_prefers_newline_boundaries() { - let mut s = String::new(); - for i in 1..=20 { - s.push_str(&format!("{i:03}\n")); - } - assert_eq!(s.len(), 80); - - let max_bytes = 64; - let (out, tokens) = truncate_middle(&s, max_bytes); - assert!(out.starts_with("001\n002\n003\n004\n")); - assert!(out.contains("[TRUNCATED CONTENT]")); - assert!(out.ends_with("017\n018\n019\n020\n")); - assert_eq!(tokens, Some(20)); - } - - #[test] - // Ensure truncation is resilient against multi-bytes chars split. - fn truncate_middle_handles_utf8_content() { - let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with ascii text\n"; - let max_bytes = 32; - let (out, tokens) = truncate_middle(s, max_bytes); - - assert!(out.contains("[TRUNCATED CONTENT]")); - assert!(out.chars().any(|c| c == '😀')); - assert!(!out.contains('\u{fffd}')); - assert_eq!(tokens, Some((s.len() as u64).div_ceil(4))); - } } diff --git a/codex-rs/core/src/unified_exec/path.rs b/codex-rs/core/src/unified_exec/path.rs index 542aea7e05..b3b765760b 100644 --- a/codex-rs/core/src/unified_exec/path.rs +++ b/codex-rs/core/src/unified_exec/path.rs @@ -10,7 +10,7 @@ pub(crate) fn parse_command_line(line: &str) -> Result, UnifiedExecE return Err(UnifiedExecError::MissingCommandLine); } - match shlex::split(trimmed) { + match parse_command_line_impl(trimmed) { Some(parts) if !parts.is_empty() => Ok(parts), _ => Err(UnifiedExecError::InvalidCommandLine { command_line: trimmed.to_string(), @@ -18,6 +18,64 @@ pub(crate) fn parse_command_line(line: &str) -> Result, UnifiedExecE } } +#[cfg(not(windows))] +fn parse_command_line_impl(trimmed: &str) -> Option> { + shlex::split(trimmed) +} + +#[cfg(windows)] +fn parse_command_line_impl(trimmed: &str) -> Option> { + windows_split_command_line(trimmed) +} + +#[cfg(windows)] +fn windows_split_command_line(input: &str) -> Option> { + let mut args = Vec::new(); + let mut current = String::new(); + let mut in_quotes = false; + let mut backslashes = 0usize; + + for ch in input.chars() { + match ch { + '\\' => { + backslashes += 1; + } + '"' => { + current.extend(std::iter::repeat('\\').take(backslashes / 2)); + if backslashes % 2 == 0 { + in_quotes = !in_quotes; + } else { + current.push('"'); + } + backslashes = 0; + } + c if c.is_whitespace() && !in_quotes => { + current.extend(std::iter::repeat('\\').take(backslashes)); + backslashes = 0; + if !current.is_empty() { + args.push(std::mem::take(&mut current)); + } + } + other => { + current.extend(std::iter::repeat('\\').take(backslashes)); + backslashes = 0; + current.push(other); + } + } + } + + if in_quotes { + return None; + } + + current.extend(std::iter::repeat('\\').take(backslashes)); + if !current.is_empty() { + args.push(current); + } + + if args.is_empty() { None } else { Some(args) } +} + pub(crate) fn command_from_chunks(chunks: &[String]) -> Result, UnifiedExecError> { match chunks { [] => Err(UnifiedExecError::MissingCommandLine), @@ -117,3 +175,20 @@ fn is_executable(metadata: &std::fs::Metadata) -> bool { fn is_executable(metadata: &std::fs::Metadata) -> bool { metadata.is_file() } + +#[cfg(all(test, windows))] +mod tests { + use super::parse_command_line; + + #[test] + fn parses_windows_command_with_spaces() { + let out = parse_command_line(r#""C:\Program Files\Git\bin\bash.exe" -i"#).unwrap(); + assert_eq!( + out, + vec![ + "C:\\Program Files\\Git\\bin\\bash.exe".to_string(), + "-i".to_string(), + ] + ); + } +} diff --git a/codex-rs/core/src/unified_exec/truncation.rs b/codex-rs/core/src/unified_exec/truncation.rs deleted file mode 100644 index 540ac563be..0000000000 --- a/codex-rs/core/src/unified_exec/truncation.rs +++ /dev/null @@ -1,71 +0,0 @@ -const TRUNCATION_MARKER: &str = "[TRUNCATED CONTENT]"; - -pub(crate) fn truncate_middle(s: &str, max_bytes: usize) -> (String, Option) { - if s.len() <= max_bytes { - return (s.to_string(), None); - } - - let est_tokens = (s.len() as u64).div_ceil(4); - if max_bytes == 0 { - return (TRUNCATION_MARKER.to_string(), Some(est_tokens)); - } - - fn truncate_on_boundary(input: &str, max_len: usize) -> &str { - if input.len() <= max_len { - return input; - } - let mut end = max_len; - while end > 0 && !input.is_char_boundary(end) { - end -= 1; - } - &input[..end] - } - - fn pick_prefix_end(s: &str, left_budget: usize) -> usize { - if let Some(head) = s.get(..left_budget) - && let Some(i) = head.rfind('\n') - { - return i + 1; - } - truncate_on_boundary(s, left_budget).len() - } - - fn pick_suffix_start(s: &str, right_budget: usize) -> usize { - let start_tail = s.len().saturating_sub(right_budget); - if let Some(tail) = s.get(start_tail..) - && let Some(i) = tail.find('\n') - { - return start_tail + i + 1; - } - let mut idx = start_tail.min(s.len()); - while idx < s.len() && !s.is_char_boundary(idx) { - idx += 1; - } - idx - } - - let marker_len = TRUNCATION_MARKER.len(); - if marker_len >= max_bytes { - return (TRUNCATION_MARKER.to_string(), Some(est_tokens)); - } - - let keep_budget = max_bytes.saturating_sub(marker_len); - if keep_budget == 0 { - return (TRUNCATION_MARKER.to_string(), Some(est_tokens)); - } - - let left_budget = keep_budget / 2; - let right_budget = keep_budget - left_budget; - let prefix_end = pick_prefix_end(s, left_budget); - let mut suffix_start = pick_suffix_start(s, right_budget); - if suffix_start < prefix_end { - suffix_start = prefix_end; - } - - let mut out = String::with_capacity(marker_len + prefix_end + (s.len() - suffix_start)); - out.push_str(&s[..prefix_end]); - out.push_str(TRUNCATION_MARKER); - out.push_str(&s[suffix_start..]); - - (out, Some(est_tokens)) -}