openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-24 18:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

Extract the head-tail buffer into a utility crate (#56535)

Browse Source 
## What changed

Move `HeadTailBuffer` and `format_output_omission_marker` into the new
`codex-utils-head-tail-buffer` crate, exposing the buffer API with an explicit
const-generic byte budget. Update core callers to use the crate and add Cargo
and Bazel definitions.

Move the existing buffer tests alongside the implementation, preserving the
prefix/suffix retention and omission-marker behavior.

GitOrigin-RevId: ff488d0a0a62c39c0e11c11ebb539166cc4c6e74
This commit is contained in:
Sean Huang
2026-10-23 17:35:25 +00:00
committed by copyberry
parent 0efc474605
commit 7c568b331b
11 changed files with 51 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tools/context.rs
+1 -1
View File
@@ -4,7 +4,6 @@ use crate::session::step_context::StepContext;
use crate::session::turn_context::TurnContext;
use crate::session::turn_context::TurnEnvironment;
use crate::turn_diff_tracker::TurnDiffTracker;
use crate::unified_exec::format_output_omission_marker;
use crate::unified_exec::resolve_max_tokens;
use codex_protocol::ResponseItemId;
use codex_protocol::mcp::CallToolResult;
@@ -17,6 +16,7 @@ use codex_protocol::models::function_call_output_content_items_to_text;
use codex_tools::LoadableToolSpec;
use codex_tools::ToolName;
use codex_utils_audio::estimate_audio_token_count;
use codex_utils_head_tail_buffer::format_output_omission_marker;
use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::approx_token_count;
use codex_utils_output_truncation::formatted_truncate_text;
codex-rs/core/src/unified_exec/head_tail_buffer.rs
-169
View File
@@ -1,169 +0,0 @@
use crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES;
use crate::unified_exec::format_output_omission_marker;
use std::collections::VecDeque;
/// A capped buffer that preserves a stable prefix ("head") and suffix ("tail"),
/// 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, Default)]
#[cfg_attr(test, derive(Eq, PartialEq))]
pub(crate) struct HeadTailBuffer<const MAX_BYTES: usize = UNIFIED_EXEC_OUTPUT_MAX_BYTES> {
head: Vec<u8>,
tail: VecDeque<u8>,
omitted_bytes: usize,
}
impl<const MAX_BYTES: usize> HeadTailBuffer<MAX_BYTES> {
const HEAD_BUDGET: usize = MAX_BYTES / 2;
const TAIL_BUDGET: usize = MAX_BYTES.saturating_sub(Self::HEAD_BUDGET);
// Used for tests.
#[allow(dead_code)]
/// Total bytes currently retained by the buffer (head + tail).
pub(crate) fn retained_bytes(&self) -> usize {
self.head.len().saturating_add(self.tail.len())
}
// Used for tests.
#[allow(dead_code)]
/// Total bytes that were dropped from the middle due to the size cap.
pub(crate) fn omitted_bytes(&self) -> usize {
self.omitted_bytes
}
/// Total bytes observed by the buffer, including bytes omitted by the cap.
pub(crate) fn total_bytes(&self) -> usize {
self.retained_bytes().saturating_add(self.omitted_bytes)
}
/// Append a chunk of bytes to the buffer.
///
/// Bytes are first added to the head until the head budget is full; any
/// remaining bytes are added to the tail, with older tail bytes being
/// dropped to preserve the tail budget.
pub(crate) fn push_chunk(&mut self, chunk: &[u8]) {
let chunk = self.fill_head(chunk);
self.push_tail(chunk);
}
/// Fill the stable prefix and return the bytes that did not fit.
fn fill_head<'a>(&mut self, chunk: &'a [u8]) -> &'a [u8] {
let Self {
head,
tail: _,
omitted_bytes: _,
} = self;
let remaining_head = Self::HEAD_BUDGET.saturating_sub(head.len());
// A shorter chunk fits entirely in the head.
let (chunk_head, chunk_tail) = chunk
.split_at_checked(remaining_head)
.unwrap_or((chunk, &[]));
head.extend_from_slice(chunk_head);
chunk_tail
}
/// Append bytes known not to belong in the head, keeping the newest tail bytes.
fn push_tail(&mut self, chunk: &[u8]) {
let Self {
head: _,
tail,
omitted_bytes,
} = self;
let remaining_tail = Self::TAIL_BUDGET.saturating_sub(tail.len());
let excess_tail = chunk.len().saturating_sub(remaining_tail);
*omitted_bytes = omitted_bytes.saturating_add(excess_tail);
// Discard old tail bytes first, then skip any excess incoming bytes.
let chunk = match excess_tail.checked_sub(tail.len()) {
None => {
tail.drain(..excess_tail);
chunk
}
Some(skip) => {
tail.clear();
&chunk[skip..]
}
};
tail.extend(chunk);
}
/// Return the retained output as a single byte vector.
///
/// The output is formed by concatenating head chunks, then tail chunks.
/// Omitted bytes are not represented in the returned value.
pub(crate) fn to_bytes(&self) -> Vec<u8> {
let mut out = Vec::with_capacity(self.retained_bytes());
out.extend_from_slice(&self.head);
out.extend(self.tail.iter().copied());
out
}
/// Return the retained output with an explicit marker between the head and
/// tail when bytes were omitted.
pub(crate) fn to_bytes_with_omission_marker(&self) -> Vec<u8> {
if self.omitted_bytes == 0 {
return self.to_bytes();
}
let marker = format_output_omission_marker(self.omitted_bytes);
let marker_delimiter_bytes = 2;
let mut out = Vec::with_capacity(
self.retained_bytes()
.saturating_add(marker.len())
.saturating_add(marker_delimiter_bytes),
);
out.extend_from_slice(&self.head);
out.push(b'\n');
out.extend_from_slice(marker.as_bytes());
out.push(b'\n');
out.extend(self.tail.iter().copied());
out
}
/// 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, buffer: Self) {
let Self {
head,
tail,
omitted_bytes,
} = buffer;
self.omitted_bytes = self.omitted_bytes.saturating_add(omitted_bytes);
// Preserve an existing prefix; otherwise reuse the source head.
let overflow = if self.head.is_empty() {
self.head = head;
&[]
} else {
self.fill_head(&head)
};
// A full source tail displaces both the old tail and the unused source head.
if tail.len() == Self::TAIL_BUDGET {
self.omitted_bytes = self
.omitted_bytes
.saturating_add(self.tail.len())
.saturating_add(overflow.len());
self.tail = tail;
} else {
self.push_tail(overflow);
// An empty destination can take a partial source tail without copying it.
if self.tail.is_empty() {
self.tail = tail;
} else {
// A nonempty source tail means its head, and now ours, is full.
let (first, second) = tail.as_slices();
self.push_tail(first);
self.push_tail(second);
}
}
}
}
#[cfg(test)]
#[path = "head_tail_buffer_tests.rs"]
mod tests;
codex-rs/core/src/unified_exec/head_tail_buffer_tests.rs
-91
View File
@@ -1,91 +0,0 @@
use super::HeadTailBuffer;
use pretty_assertions::assert_eq;
#[test]
fn keeps_prefix_and_suffix_when_over_budget() {
let mut buf = HeadTailBuffer::<10>::default();
buf.push_chunk(b"0123456789");
assert_eq!(buf.omitted_bytes(), 0);
// Exceeds max by 2; we should keep head+tail and omit the middle.
buf.push_chunk(b"ab");
assert!(buf.omitted_bytes() > 0);
let rendered = String::from_utf8_lossy(&buf.to_bytes()).to_string();
assert!(rendered.starts_with("01234"));
assert!(rendered.ends_with("89ab"));
assert_eq!(
String::from_utf8_lossy(&buf.to_bytes_with_omission_marker()),
"01234\n... 2 bytes omitted ...\n789ab"
);
}
#[test]
fn max_bytes_zero_drops_everything() {
let mut buf = HeadTailBuffer::<0>::default();
buf.push_chunk(b"abc");
assert_eq!(buf.retained_bytes(), 0);
assert_eq!(buf.omitted_bytes(), 3);
assert_eq!(buf.to_bytes(), b"");
}
#[test]
fn head_budget_zero_keeps_only_last_byte_in_tail() {
let mut buf = HeadTailBuffer::<1>::default();
buf.push_chunk(b"abc");
assert_eq!(buf.retained_bytes(), 1);
assert_eq!(buf.omitted_bytes(), 2);
assert_eq!(buf.to_bytes(), b"c");
}
#[test]
fn chunk_larger_than_tail_budget_keeps_only_tail_end() {
let mut buf = HeadTailBuffer::<10>::default();
buf.push_chunk(b"0123456789");
// Tail budget is 5 bytes. This chunk should replace the tail and keep only its last 5 bytes.
buf.push_chunk(b"ABCDEFGHIJK");
let out = String::from_utf8_lossy(&buf.to_bytes()).to_string();
assert!(out.starts_with("01234"));
assert!(out.ends_with("GHIJK"));
assert!(buf.omitted_bytes() > 0);
}
#[test]
fn fills_head_then_tail_across_multiple_chunks() {
let mut buf = HeadTailBuffer::<10>::default();
// Fill the 5-byte head budget across multiple chunks.
buf.push_chunk(b"01");
buf.push_chunk(b"234");
assert_eq!(buf.to_bytes(), b"01234");
// Then fill the 5-byte tail budget.
buf.push_chunk(b"567");
buf.push_chunk(b"89");
assert_eq!(buf.to_bytes(), b"0123456789");
assert_eq!(buf.omitted_bytes(), 0);
// One more byte causes the tail to drop its oldest byte.
buf.push_chunk(b"a");
assert_eq!(buf.to_bytes(), b"012346789a");
assert_eq!(buf.omitted_bytes(), 1);
}
#[test]
fn empty_and_tiny_chunks_have_bounded_metadata() {
let mut buf = HeadTailBuffer::<10>::default();
for byte in b"0123456789ab" {
buf.push_chunk(&[]);
buf.push_chunk(&[*byte]);
}
assert_eq!(buf.retained_bytes(), 10);
assert_eq!(buf.omitted_bytes(), 2);
}
codex-rs/core/src/unified_exec/mod.rs
-5
View File
@@ -48,7 +48,6 @@ use codex_core_plugins::PluginMetricsSidecar;
mod async_watcher;
mod errors;
mod head_tail_buffer;
mod oneshot;
mod process;
mod process_manager;
@@ -232,10 +231,6 @@ pub(crate) fn resolve_max_tokens(max_tokens: Option<usize>) -> usize {
max_tokens.unwrap_or(DEFAULT_MAX_OUTPUT_TOKENS)
}
pub(crate) fn format_output_omission_marker(omitted_bytes: usize) -> String {
format!("... {omitted_bytes} bytes omitted ...")
}
pub(crate) fn generate_chunk_id() -> String {
let mut rng = rng();
(0..6)
codex-rs/core/src/unified_exec/process.rs
+1 -1
View File
@@ -24,6 +24,7 @@ use codex_protocol::protocol::TruncationPolicy;
use codex_sandboxing::SandboxType;
use codex_sandboxing::is_likely_sandbox_denied;
use codex_sandboxing::record_filesystem_sandbox_violation;
use codex_utils_head_tail_buffer::HeadTailBuffer;
use codex_utils_output_truncation::formatted_truncate_text;
use codex_utils_pty::ExecCommandSession;
use codex_utils_pty::ProcessSignal as PtyProcessSignal;
@@ -32,7 +33,6 @@ use codex_utils_pty::SpawnedPty;
use super::UNIFIED_EXEC_OUTPUT_MAX_BYTES;
use super::UNIFIED_EXEC_OUTPUT_MAX_TOKENS;
use super::UnifiedExecError;
use super::head_tail_buffer::HeadTailBuffer;
use super::process_state::ProcessState;
use crate::shell_snapshot::ShellSnapshotFile;
codex-rs/core/src/unified_exec/process_manager.rs
+1 -1
View File
@@ -64,7 +64,6 @@ use crate::unified_exec::async_watcher::spawn_exit_watcher;
use crate::unified_exec::async_watcher::start_streaming_output;
use crate::unified_exec::clamp_yield_time;
use crate::unified_exec::generate_chunk_id;
use crate::unified_exec::head_tail_buffer::HeadTailBuffer;
use crate::unified_exec::process::OutputBuffers;
use crate::unified_exec::process::OutputHandles;
use crate::unified_exec::process::UnifiedExecProcess;
@@ -89,6 +88,7 @@ use codex_protocol::shell_environment::is_non_inheritable_env_var;
use codex_sandboxing::SandboxCommand;
use codex_shell_command::is_dangerous_command::DangerousCommandPlatform;
use codex_tools::ToolName;
use codex_utils_head_tail_buffer::HeadTailBuffer;
use codex_utils_output_truncation::approx_tokens_from_byte_count;
use codex_utils_path_uri::PathUri;