diff --git a/codex-rs/core/src/unified_exec/head_tail_buffer.rs b/codex-rs/core/src/unified_exec/head_tail_buffer.rs index 27fa23472c..9bebfdcaf8 100644 --- a/codex-rs/core/src/unified_exec/head_tail_buffer.rs +++ b/codex-rs/core/src/unified_exec/head_tail_buffer.rs @@ -6,40 +6,17 @@ use std::collections::VecDeque; /// dropping the middle once it exceeds the configured maximum. The buffer is /// symmetric meaning 50% of the capacity is allocated to the head and 50% is /// allocated to the tail. -#[derive(Debug)] +#[derive(Debug, Default)] #[cfg_attr(test, derive(Eq, PartialEq))] -pub(crate) struct HeadTailBuffer { - max_bytes: usize, - head_budget: usize, - tail_budget: usize, +pub(crate) struct HeadTailBuffer { head: Vec, tail: VecDeque, omitted_bytes: usize, } -impl Default for HeadTailBuffer { - fn default() -> Self { - Self::new(UNIFIED_EXEC_OUTPUT_MAX_BYTES) - } -} - -impl HeadTailBuffer { - /// Create a new buffer that retains at most `max_bytes` of output. - /// - /// The retained output is split across a prefix ("head") and suffix ("tail") - /// budget, dropping bytes from the middle once the limit is exceeded. - pub(crate) fn new(max_bytes: usize) -> Self { - let head_budget = max_bytes / 2; - let tail_budget = max_bytes.saturating_sub(head_budget); - Self { - max_bytes, - head_budget, - tail_budget, - head: Vec::new(), - tail: VecDeque::new(), - omitted_bytes: 0, - } - } +impl HeadTailBuffer { + const HEAD_BUDGET: usize = MAX_BYTES / 2; + const TAIL_BUDGET: usize = MAX_BYTES.saturating_sub(Self::HEAD_BUDGET); // Used for tests. #[allow(dead_code)] @@ -69,13 +46,13 @@ impl HeadTailBuffer { if chunk.is_empty() { return; } - if self.max_bytes == 0 { + if MAX_BYTES == 0 { self.omitted_bytes = self.omitted_bytes.saturating_add(chunk.len()); return; } // Fill the head budget first, then keep a capped tail. - let remaining_head = self.head_budget.saturating_sub(self.head.len()); + let remaining_head = Self::HEAD_BUDGET.saturating_sub(self.head.len()); let head_len = remaining_head.min(chunk.len()); if head_len > 0 { self.head.extend_from_slice(&chunk[..head_len]); @@ -120,17 +97,14 @@ impl HeadTailBuffer { /// contents while preserving its configured capacity. pub(crate) fn drain(&mut self) -> Self { Self { - max_bytes: self.max_bytes, - head_budget: self.head_budget, - tail_budget: self.tail_budget, head: std::mem::take(&mut self.head), tail: std::mem::take(&mut self.tail), omitted_bytes: std::mem::take(&mut self.omitted_bytes), } } - /// Append retained output from another buffer and preserve any omissions it - /// already recorded. + /// Append a later buffer with the same budget. This preserves the summary + /// of the original concatenated output, including its omission count. pub(crate) fn push_buffer(&mut self, mut buffer: Self) { self.push_chunk(std::mem::take(&mut buffer.head)); self.push_chunk(buffer.tail.drain(..).collect()); @@ -141,15 +115,15 @@ impl HeadTailBuffer { if chunk.is_empty() { return; } - if self.tail_budget == 0 { + if Self::TAIL_BUDGET == 0 { self.omitted_bytes = self.omitted_bytes.saturating_add(chunk.len()); return; } - if chunk.len() >= self.tail_budget { + if chunk.len() >= Self::TAIL_BUDGET { // This single chunk is larger than the whole tail budget. Keep only the last // tail_budget bytes and drop everything else. - let start = chunk.len().saturating_sub(self.tail_budget); + let start = chunk.len().saturating_sub(Self::TAIL_BUDGET); let kept = &chunk[start..]; let dropped = chunk.len().saturating_sub(kept.len()); self.omitted_bytes = self @@ -166,7 +140,7 @@ impl HeadTailBuffer { } fn trim_tail_to_budget(&mut self) { - let excess = self.tail.len().saturating_sub(self.tail_budget); + let excess = self.tail.len().saturating_sub(Self::TAIL_BUDGET); if excess > 0 { drop(self.tail.drain(..excess)); self.omitted_bytes = self.omitted_bytes.saturating_add(excess); diff --git a/codex-rs/core/src/unified_exec/head_tail_buffer_tests.rs b/codex-rs/core/src/unified_exec/head_tail_buffer_tests.rs index 1fe97b26da..7110cddea7 100644 --- a/codex-rs/core/src/unified_exec/head_tail_buffer_tests.rs +++ b/codex-rs/core/src/unified_exec/head_tail_buffer_tests.rs @@ -4,7 +4,7 @@ use pretty_assertions::assert_eq; #[test] fn keeps_prefix_and_suffix_when_over_budget() { - let mut buf = HeadTailBuffer::new(/*max_bytes*/ 10); + let mut buf = HeadTailBuffer::<10>::default(); buf.push_chunk(b"0123456789".to_vec()); assert_eq!(buf.omitted_bytes(), 0); @@ -24,7 +24,7 @@ fn keeps_prefix_and_suffix_when_over_budget() { #[test] fn max_bytes_zero_drops_everything() { - let mut buf = HeadTailBuffer::new(/*max_bytes*/ 0); + let mut buf = HeadTailBuffer::<0>::default(); buf.push_chunk(b"abc".to_vec()); assert_eq!(buf.retained_bytes(), 0); @@ -34,7 +34,7 @@ fn max_bytes_zero_drops_everything() { #[test] fn head_budget_zero_keeps_only_last_byte_in_tail() { - let mut buf = HeadTailBuffer::new(/*max_bytes*/ 1); + let mut buf = HeadTailBuffer::<1>::default(); buf.push_chunk(b"abc".to_vec()); assert_eq!(buf.retained_bytes(), 1); @@ -44,12 +44,12 @@ fn head_budget_zero_keeps_only_last_byte_in_tail() { #[test] fn draining_resets_state_and_push_buffer_preserves_omissions() { - let mut buf = HeadTailBuffer::new(/*max_bytes*/ 10); + let mut buf = HeadTailBuffer::<10>::default(); buf.push_chunk(b"0123456789".to_vec()); buf.push_chunk(b"ab".to_vec()); let drained = buf.drain(); - let mut collected = HeadTailBuffer::new(/*max_bytes*/ 10); + let mut collected = HeadTailBuffer::<10>::default(); collected.push_buffer(drained); assert_eq!(buf.retained_bytes(), 0); @@ -62,7 +62,7 @@ fn draining_resets_state_and_push_buffer_preserves_omissions() { #[test] fn chunk_larger_than_tail_budget_keeps_only_tail_end() { - let mut buf = HeadTailBuffer::new(/*max_bytes*/ 10); + let mut buf = HeadTailBuffer::<10>::default(); buf.push_chunk(b"0123456789".to_vec()); // Tail budget is 5 bytes. This chunk should replace the tail and keep only its last 5 bytes. @@ -76,7 +76,7 @@ fn chunk_larger_than_tail_budget_keeps_only_tail_end() { #[test] fn fills_head_then_tail_across_multiple_chunks() { - let mut buf = HeadTailBuffer::new(/*max_bytes*/ 10); + let mut buf = HeadTailBuffer::<10>::default(); // Fill the 5-byte head budget across multiple chunks. buf.push_chunk(b"01".to_vec()); @@ -97,7 +97,7 @@ fn fills_head_then_tail_across_multiple_chunks() { #[test] fn empty_and_tiny_chunks_have_bounded_metadata() { - let mut buf = HeadTailBuffer::new(/*max_bytes*/ 10); + let mut buf = HeadTailBuffer::<10>::default(); for byte in b"0123456789ab" { buf.push_chunk(Vec::new()); diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index d72ebb8f54..682d789983 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -324,7 +324,7 @@ async fn write_stdin( #[test] fn push_chunk_preserves_prefix_and_suffix() { - let mut buffer = HeadTailBuffer::default(); + let mut buffer: HeadTailBuffer = Default::default(); buffer.push_chunk(vec![b'a'; UNIFIED_EXEC_OUTPUT_MAX_BYTES]); buffer.push_chunk(vec![b'b']); buffer.push_chunk(vec![b'c']); @@ -341,7 +341,7 @@ fn push_chunk_preserves_prefix_and_suffix() { #[test] fn head_tail_buffer_default_preserves_prefix_and_suffix() { - let mut buffer = HeadTailBuffer::default(); + let mut buffer: HeadTailBuffer = Default::default(); buffer.push_chunk(vec![b'a'; UNIFIED_EXEC_OUTPUT_MAX_BYTES]); buffer.push_chunk(b"bc".to_vec());