openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-20 12:47:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Keep unified exec output collection bounded (#32150)

Browse Source 
## Why

Unified exec can drain process output multiple times while waiting for a command. Collecting those drains into an uncapped buffer allows large commands to exceed the output collection limit.

## What changed

- Accumulate drained output through the capped head/tail buffer so collection remains bounded while preserving the beginning and end of the stream.
- Include a `... N bytes omitted ...` marker when the collection cap drops output.
- Calculate the original token estimate from all observed bytes and preserve omission metadata through truncation and sandbox-denial responses.

## Testing

- Cover repeated drains, previously omitted output, omission markers, and large-output summaries.

GitOrigin-RevId: c6e2efeb875ed8f35914365eab7431f25fd3484c
This commit is contained in:
jif
2026-07-10 12:16:14 +00:00
committed by copyberry
parent 54c44b9ed4
commit 6138909d6e
13 changed files with 413 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
codex-rs/core/src/tools/context.rs
View File

@@ -7,6 +7,7 @@ use crate::tools::TELEMETRY_PREVIEW_MAX_BYTES;
use crate::tools::TELEMETRY_PREVIEW_MAX_LINES;
use crate::tools::TELEMETRY_PREVIEW_TRUNCATION_NOTICE;
use crate::turn_diff_tracker::TurnDiffTracker;
use crate::unified_exec::format_output_omission_marker;
use crate::unified_exec::resolve_max_tokens;
use codex_protocol::mcp::CallToolResult;
use codex_protocol::models::FunctionCallOutputBody;
@@ -17,10 +18,13 @@ use codex_protocol::models::function_call_output_content_items_to_text;
use codex_tools::LoadableToolSpec;
use codex_tools::ToolName;
use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::approx_token_count;
use codex_utils_output_truncation::formatted_truncate_text;
use codex_utils_output_truncation::truncate_text;
use codex_utils_string::take_bytes_at_char_boundary;
use serde::Serialize;
use serde_json::Value as JsonValue;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
@@ -319,6 +323,8 @@ pub struct ExecCommandToolOutput {
pub process_id: Option<i32>,
pub exit_code: Option<i32>,
pub original_token_count: Option<usize>,
/// Bytes omitted by the output collection cap before model-facing truncation.
pub output_omitted_bytes: Option<NonZeroUsize>,
pub hook_command: Option<String>,
}
@@ -406,7 +412,32 @@ impl ExecCommandToolOutput {
pub(crate) fn truncated_output(&self, max_tokens: usize) -> String {
let text = String::from_utf8_lossy(&self.raw_output).to_string();
formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens))
let policy = TruncationPolicy::Tokens(max_tokens);
let Some(omitted_bytes) = self.output_omitted_bytes else {
return formatted_truncate_text(&text, policy);
};
let marker = format_output_omission_marker(omitted_bytes.get());
if text.len() <= policy.byte_budget() {
return if text.contains(&marker) {
text
} else {
format!("{marker}\n{text}")
};
}
let original_token_count = self
.original_token_count
.unwrap_or_else(|| approx_token_count(&text));
let truncated = truncate_text(&text, policy);
let omission_notice = if truncated.contains(&marker) {
String::new()
} else {
format!("{marker}\n")
};
format!(
"Warning: truncated output (original token count: {original_token_count})\n{omission_notice}\n{truncated}"
)
}
fn response_text(&self) -> String {

40
codex-rs/core/src/tools/context_tests.rs
View File

@@ -434,6 +434,7 @@ fn exec_command_tool_output_formats_truncated_response() {
process_id: None,
exit_code: Some(0),
original_token_count: Some(10),
output_omitted_bytes: None,
hook_command: None,
}
.to_response_item("call-42", &payload);
@@ -461,3 +462,42 @@ fn exec_command_tool_output_formats_truncated_response() {
other => panic!("expected FunctionCallOutput, got {other:?}"),
}
}
#[test]
fn exec_command_tool_output_preserves_omission_metadata_when_truncated() {
let payload = ToolPayload::Function {
arguments: "{}".to_string(),
};
let marker = format_output_omission_marker(/*omitted_bytes*/ 123_456);
let raw_output = format!(
"HEAD-{}\n{marker}\nTAIL-{}",
"a".repeat(/*n*/ 100),
"z".repeat(/*n*/ 100)
)
.into_bytes();
let response = ExecCommandToolOutput {
event_call_id: "call-omitted".to_string(),
chunk_id: "abc123".to_string(),
wall_time: std::time::Duration::from_millis(/*millis*/ 1250),
raw_output,
truncation_policy: TruncationPolicy::Tokens(10_000),
max_output_tokens: Some(4),
process_id: None,
exit_code: Some(0),
original_token_count: Some(42_000),
output_omitted_bytes: NonZeroUsize::new(/*n*/ 123_456),
hook_command: None,
}
.to_response_item("call-omitted", &payload);
let ResponseInputItem::FunctionCallOutput { output, .. } = response else {
panic!("expected FunctionCallOutput");
};
let text = output
.body
.to_text()
.expect("exec output should serialize as text");
assert!(text.contains("Original token count: 42000"));
assert!(text.contains("Warning: truncated output (original token count: 42000)"));
assert_eq!(text.matches(&marker).count(), 1);
}

12
codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs
View File

@@ -335,6 +335,7 @@ impl ExecCommandHandler {
process_id: None,
exit_code: None,
original_token_count: None,
output_omitted_bytes: None,
hook_command: None,
}));
}
@@ -367,9 +368,15 @@ impl ExecCommandHandler {
.await
{
Ok(response) => Ok(boxed_tool_output(response)),
Err(UnifiedExecError::SandboxDenied { output, .. }) => {
Err(UnifiedExecError::SandboxDenied {
output,
original_token_count,
output_omitted_bytes,
..
}) => {
let output_text = output.aggregated_output.text;
let original_token_count = approx_token_count(&output_text);
let original_token_count =
original_token_count.unwrap_or_else(|| approx_token_count(&output_text));
Ok(boxed_tool_output(ExecCommandToolOutput {
event_call_id: context.call_id.clone(),
chunk_id: generate_chunk_id(),
@@ -382,6 +389,7 @@ impl ExecCommandHandler {
process_id: None,
exit_code: Some(output.exit_code),
original_token_count: Some(original_token_count),
output_omitted_bytes,
hook_command: Some(hook_command),
}))
}

6
codex-rs/core/src/tools/handlers/unified_exec_tests.rs
View File

@@ -301,6 +301,7 @@ async fn exec_command_post_tool_use_payload_uses_output_for_noninteractive_one_s
process_id: None,
exit_code: Some(0),
original_token_count: None,
output_omitted_bytes: None,
hook_command: Some("echo three".to_string()),
};
let invocation = invocation_for_payload("exec_command", "call-43", payload).await;
@@ -331,6 +332,7 @@ async fn exec_command_post_tool_use_payload_uses_output_for_interactive_completi
process_id: None,
exit_code: Some(0),
original_token_count: None,
output_omitted_bytes: None,
hook_command: Some("echo three".to_string()),
};
let invocation = invocation_for_payload("exec_command", "call-44", payload).await;
@@ -362,6 +364,7 @@ async fn exec_command_post_tool_use_payload_skips_running_sessions() {
process_id: Some(45),
exit_code: None,
original_token_count: None,
output_omitted_bytes: None,
hook_command: Some("echo three".to_string()),
};
let invocation = invocation_for_payload("exec_command", "call-45", payload).await;
@@ -388,6 +391,7 @@ async fn write_stdin_post_tool_use_payload_uses_original_exec_call_id_and_comman
process_id: None,
exit_code: Some(0),
original_token_count: None,
output_omitted_bytes: None,
hook_command: Some("sleep 1; echo finished".to_string()),
};
let invocation = invocation_for_payload("write_stdin", "write-stdin-call", payload).await;
@@ -419,6 +423,7 @@ async fn write_stdin_post_tool_use_payload_keeps_parallel_session_metadata_separ
process_id: None,
exit_code: Some(0),
original_token_count: None,
output_omitted_bytes: None,
hook_command: Some("sleep 2; echo alpha".to_string()),
};
let output_b = ExecCommandToolOutput {
@@ -431,6 +436,7 @@ async fn write_stdin_post_tool_use_payload_keeps_parallel_session_metadata_separ
process_id: None,
exit_code: Some(0),
original_token_count: None,
output_omitted_bytes: None,
hook_command: Some("sleep 1; echo beta".to_string()),
};
let invocation_b = invocation_for_payload("write_stdin", "write-call-b", payload.clone()).await;

2
codex-rs/core/src/unified_exec/async_watcher.rs
View File

@@ -326,7 +326,7 @@ async fn resolve_aggregated_output(
return fallback;
}
String::from_utf8_lossy(&guard.to_bytes()).to_string()
String::from_utf8_lossy(&guard.to_bytes_with_omission_marker()).to_string()
}
#[cfg(test)]

28
codex-rs/core/src/unified_exec/errors.rs
View File

@@ -1,5 +1,6 @@
use codex_protocol::exec_output::ExecToolCallOutput;
use codex_utils_path_uri::PathUri;
use std::num::NonZeroUsize;
use thiserror::Error;
#[derive(Debug, Error)]
@@ -23,6 +24,8 @@ pub(crate) enum UnifiedExecError {
SandboxDenied {
message: String,
output: ExecToolCallOutput,
original_token_count: Option<usize>,
output_omitted_bytes: Option<NonZeroUsize>,
},
#[error("{path} is not valid on {}", std::env::consts::OS)]
ForeignPath { path: PathUri },
@@ -38,6 +41,29 @@ impl UnifiedExecError {
}
pub(crate) fn sandbox_denied(message: String, output: ExecToolCallOutput) -> Self {
Self::SandboxDenied { message, output }
Self::SandboxDenied {
message,
output,
original_token_count: None,
output_omitted_bytes: None,
}
}
pub(crate) fn with_output_collection_metadata(
self,
original_token_count: usize,
output_omitted_bytes: Option<NonZeroUsize>,
) -> Self {
match self {
Self::SandboxDenied {
message, output, ..
} => Self::SandboxDenied {
message,
output,
original_token_count: Some(original_token_count),
output_omitted_bytes,
},
other => other,
}
}
}

149
codex-rs/core/src/unified_exec/head_tail_buffer.rs
View File

@@ -1,4 +1,5 @@
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"),
@@ -6,14 +7,13 @@ use std::collections::VecDeque;
/// symmetric meaning 50% of the capacity is allocated to the head and 50% is
/// allocated to the tail.
#[derive(Debug)]
#[cfg_attr(test, derive(Eq, PartialEq))]
pub(crate) struct HeadTailBuffer {
max_bytes: usize,
head_budget: usize,
tail_budget: usize,
head: VecDeque<Vec<u8>>,
tail: VecDeque<Vec<u8>>,
head_bytes: usize,
tail_bytes: usize,
head: Vec<u8>,
tail: VecDeque<u8>,
omitted_bytes: usize,
}
@@ -35,10 +35,8 @@ impl HeadTailBuffer {
max_bytes,
head_budget,
tail_budget,
head: VecDeque::new(),
head: Vec::new(),
tail: VecDeque::new(),
head_bytes: 0,
tail_bytes: 0,
omitted_bytes: 0,
}
}
@@ -47,7 +45,7 @@ impl HeadTailBuffer {
#[allow(dead_code)]
/// Total bytes currently retained by the buffer (head + tail).
pub(crate) fn retained_bytes(&self) -> usize {
self.head_bytes.saturating_add(self.tail_bytes)
self.head.len().saturating_add(self.tail.len())
}
// Used for tests.
@@ -57,37 +55,32 @@ impl HeadTailBuffer {
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: Vec<u8>) {
if chunk.is_empty() {
return;
}
if self.max_bytes == 0 {
self.omitted_bytes = self.omitted_bytes.saturating_add(chunk.len());
return;
}
// Fill the head budget first, then keep a capped tail.
if self.head_bytes < self.head_budget {
let remaining_head = self.head_budget.saturating_sub(self.head_bytes);
if chunk.len() <= remaining_head {
self.head_bytes = self.head_bytes.saturating_add(chunk.len());
self.head.push_back(chunk);
return;
}
// Split the chunk: part goes to head, remainder goes to tail.
let (head_part, tail_part) = chunk.split_at(remaining_head);
if !head_part.is_empty() {
self.head_bytes = self.head_bytes.saturating_add(head_part.len());
self.head.push_back(head_part.to_vec());
}
self.push_to_tail(tail_part.to_vec());
return;
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]);
}
self.push_to_tail(chunk);
self.push_to_tail(&chunk[head_len..]);
}
/// Snapshot the retained output as a list of chunks.
@@ -95,9 +88,13 @@ impl HeadTailBuffer {
/// The returned chunks are ordered as: head chunks first, then tail chunks.
/// Omitted bytes are not represented in the snapshot.
pub(crate) fn snapshot_chunks(&self) -> Vec<Vec<u8>> {
let mut out = Vec::new();
out.extend(self.head.iter().cloned());
out.extend(self.tail.iter().cloned());
let mut out = Vec::with_capacity(2);
if !self.head.is_empty() {
out.push(self.head.clone());
}
if !self.tail.is_empty() {
out.push(self.tail.iter().copied().collect());
}
out
}
@@ -107,29 +104,58 @@ impl HeadTailBuffer {
/// 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());
for chunk in self.head.iter() {
out.extend_from_slice(chunk);
}
for chunk in self.tail.iter() {
out.extend_from_slice(chunk);
}
out.extend_from_slice(&self.head);
out.extend(self.tail.iter().copied());
out
}
/// Drain all retained chunks from the buffer and reset its state.
///
/// The drained chunks are returned in head-then-tail order. Omitted bytes
/// are discarded along with the retained content.
pub(crate) fn drain_chunks(&mut self) -> Vec<Vec<u8>> {
let mut out: Vec<Vec<u8>> = self.head.drain(..).collect();
out.extend(self.tail.drain(..));
self.head_bytes = 0;
self.tail_bytes = 0;
self.omitted_bytes = 0;
/// 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
}
fn push_to_tail(&mut self, chunk: Vec<u8>) {
/// Drain the retained output and omission metadata, resetting this buffer's
/// 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.
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());
self.omitted_bytes = self.omitted_bytes.saturating_add(buffer.omitted_bytes);
}
fn push_to_tail(&mut self, chunk: &[u8]) {
if chunk.is_empty() {
return;
}
if self.tail_budget == 0 {
self.omitted_bytes = self.omitted_bytes.saturating_add(chunk.len());
return;
@@ -139,41 +165,26 @@ impl HeadTailBuffer {
// 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 kept = chunk[start..].to_vec();
let kept = &chunk[start..];
let dropped = chunk.len().saturating_sub(kept.len());
self.omitted_bytes = self
.omitted_bytes
.saturating_add(self.tail_bytes)
.saturating_add(self.tail.len())
.saturating_add(dropped);
self.tail.clear();
self.tail_bytes = kept.len();
self.tail.push_back(kept);
self.tail.extend(kept);
return;
}
self.tail_bytes = self.tail_bytes.saturating_add(chunk.len());
self.tail.push_back(chunk);
self.tail.extend(chunk);
self.trim_tail_to_budget();
}
fn trim_tail_to_budget(&mut self) {
let mut excess = self.tail_bytes.saturating_sub(self.tail_budget);
while excess > 0 {
match self.tail.front_mut() {
Some(front) if excess >= front.len() => {
excess -= front.len();
self.tail_bytes = self.tail_bytes.saturating_sub(front.len());
self.omitted_bytes = self.omitted_bytes.saturating_add(front.len());
self.tail.pop_front();
}
Some(front) => {
front.drain(..excess);
self.tail_bytes = self.tail_bytes.saturating_sub(excess);
self.omitted_bytes = self.omitted_bytes.saturating_add(excess);
break;
}
None => break,
}
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);
}
}
}

31
codex-rs/core/src/unified_exec/head_tail_buffer_tests.rs
View File

@@ -16,6 +16,10 @@ fn keeps_prefix_and_suffix_when_over_budget() {
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]
@@ -40,17 +44,21 @@ fn head_budget_zero_keeps_only_last_byte_in_tail() {
}
#[test]
fn draining_resets_state() {
fn draining_resets_state_and_push_buffer_preserves_omissions() {
let mut buf = HeadTailBuffer::new(/*max_bytes*/ 10);
buf.push_chunk(b"0123456789".to_vec());
buf.push_chunk(b"ab".to_vec());
let drained = buf.drain_chunks();
assert!(!drained.is_empty());
let drained = buf.drain();
let mut collected = HeadTailBuffer::new(/*max_bytes*/ 10);
collected.push_buffer(drained);
assert_eq!(buf.retained_bytes(), 0);
assert_eq!(buf.omitted_bytes(), 0);
assert_eq!(buf.to_bytes(), b"".to_vec());
assert_eq!(collected.to_bytes(), b"01234789ab".to_vec());
assert_eq!(collected.omitted_bytes(), 2);
assert_eq!(collected.total_bytes(), 12);
}
#[test]
@@ -87,3 +95,20 @@ fn fills_head_then_tail_across_multiple_chunks() {
assert_eq!(buf.to_bytes(), b"012346789a".to_vec());
assert_eq!(buf.omitted_bytes(), 1);
}
#[test]
fn empty_and_tiny_chunks_have_bounded_metadata() {
let mut buf = HeadTailBuffer::new(/*max_bytes*/ 10);
for byte in b"0123456789ab" {
buf.push_chunk(Vec::new());
buf.push_chunk(vec![*byte]);
}
assert_eq!(
buf.snapshot_chunks(),
vec![b"01234".to_vec(), b"789ab".to_vec()]
);
assert_eq!(buf.retained_bytes(), 10);
assert_eq!(buf.omitted_bytes(), 2);
}

4
codex-rs/core/src/unified_exec/mod.rs
View File

@@ -178,6 +178,10 @@ 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)

34
codex-rs/core/src/unified_exec/mod_tests.rs
View File

@@ -22,12 +22,13 @@ use codex_exec_server::WriteStatus;
use codex_sandboxing::SandboxType;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::approx_token_count;
use codex_utils_output_truncation::approx_tokens_from_byte_count;
use core_test_support::skip_if_no_remote_env;
use core_test_support::skip_if_sandbox;
use core_test_support::test_codex::test_env as remote_test_env;
use pretty_assertions::assert_eq;
use std::collections::HashMap;
use std::num::NonZeroUsize;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::Notify;
@@ -154,7 +155,7 @@ async fn exec_command_with_tty(
cancellation_token,
} = process.output_handles();
let deadline = started_at + Duration::from_millis(yield_time_ms);
let collected = UnifiedExecProcessManager::collect_output_until_deadline(
let collected_output = UnifiedExecProcessManager::collect_output_until_deadline(
&output_buffer,
&output_notify,
&output_closed,
@@ -165,7 +166,12 @@ async fn exec_command_with_tty(
)
.await;
let wall_time = Instant::now().saturating_duration_since(started_at);
let text = String::from_utf8_lossy(&collected).to_string();
let original_token_count = usize::try_from(approx_tokens_from_byte_count(
collected_output.total_bytes(),
))
.unwrap_or(usize::MAX);
let output_omitted_bytes = NonZeroUsize::new(collected_output.omitted_bytes());
let collected = collected_output.to_bytes_with_omission_marker();
let has_exited = process.has_exited();
let exit_code = process.exit_code();
let response_process_id = if process_started_alive && !has_exited {
@@ -196,7 +202,8 @@ async fn exec_command_with_tty(
max_output_tokens: None,
process_id: response_process_id,
exit_code,
original_token_count: Some(approx_token_count(&text)),
original_token_count: Some(original_token_count),
output_omitted_bytes,
hook_command: Some(cmd.to_string()),
})
}
@@ -327,17 +334,11 @@ fn push_chunk_preserves_prefix_and_suffix() {
assert_eq!(buffer.retained_bytes(), UNIFIED_EXEC_OUTPUT_MAX_BYTES);
let snapshot = buffer.snapshot_chunks();
let first = snapshot.first().expect("expected at least one chunk");
assert_eq!(first.first(), Some(&b'a'));
assert!(snapshot.iter().any(|chunk| chunk.as_slice() == b"b"));
assert_eq!(
snapshot
.last()
.expect("expected at least one chunk")
.as_slice(),
b"c"
);
let head_bytes = UNIFIED_EXEC_OUTPUT_MAX_BYTES / 2;
let tail_bytes = UNIFIED_EXEC_OUTPUT_MAX_BYTES - head_bytes;
let mut expected_tail = vec![b'a'; tail_bytes - 2];
expected_tail.extend_from_slice(b"bc");
assert_eq!(snapshot, vec![vec![b'a'; head_bytes], expected_tail]);
}
#[test]
@@ -876,7 +877,8 @@ async fn unified_exec_uses_remote_exec_server_when_configured() -> anyhow::Resul
/*pause_state*/ None,
Instant::now() + Duration::from_millis(2_500),
)
.await;
.await
.to_bytes_with_omission_marker();
assert!(String::from_utf8_lossy(&collected).contains("remote-unified-exec"));
Ok(())

60
codex-rs/core/src/unified_exec/process_manager.rs
View File

@@ -2,6 +2,7 @@ use rand::Rng;
use std::cmp::Reverse;
use std::collections::HashMap;
use std::collections::HashSet;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
@@ -63,7 +64,7 @@ use codex_protocol::error::SandboxErr;
use codex_protocol::protocol::ExecCommandSource;
use codex_sandboxing::SandboxCommand;
use codex_tools::ToolName;
use codex_utils_output_truncation::approx_token_count;
use codex_utils_output_truncation::approx_tokens_from_byte_count;
use codex_utils_path_uri::PathUri;
const UNIFIED_EXEC_ENV: [(&str, &str); 10] = [
@@ -487,7 +488,7 @@ impl UnifiedExecProcessManager {
cancellation_token,
} = process.output_handles();
let deadline = start + Duration::from_millis(yield_time_ms);
let collected = Self::collect_output_until_deadline(
let collected_output = Self::collect_output_until_deadline(
&output_buffer,
&output_notify,
&output_closed,
@@ -499,6 +500,12 @@ impl UnifiedExecProcessManager {
.await;
let wall_time = Instant::now().saturating_duration_since(start);
let original_token_count = usize::try_from(approx_tokens_from_byte_count(
collected_output.total_bytes(),
))
.unwrap_or(usize::MAX);
let output_omitted_bytes = NonZeroUsize::new(collected_output.omitted_bytes());
let collected = collected_output.to_bytes_with_omission_marker();
let text = String::from_utf8_lossy(&collected).to_string();
let chunk_id = generate_chunk_id();
if deferred_network_approval
@@ -565,7 +572,15 @@ impl UnifiedExecProcessManager {
{
return Err(fail_process_with_message(entry.process.as_ref(), message));
}
process.check_for_sandbox_denial_with_text(&text).await?;
process
.check_for_sandbox_denial_with_text(&text)
.await
.map_err(|err| {
err.with_output_collection_metadata(
original_token_count,
output_omitted_bytes,
)
})?;
(None, exit_code)
}
ProcessStatus::Unknown => {
@@ -612,11 +627,15 @@ impl UnifiedExecProcessManager {
.await;
self.release_process_id(request.process_id).await;
process.check_for_sandbox_denial_with_text(&text).await?;
process
.check_for_sandbox_denial_with_text(&text)
.await
.map_err(|err| {
err.with_output_collection_metadata(original_token_count, output_omitted_bytes)
})?;
(None, exit_code)
};
let original_token_count = approx_token_count(&text);
let response = ExecCommandToolOutput {
event_call_id: context.call_id.clone(),
chunk_id,
@@ -627,6 +646,7 @@ impl UnifiedExecProcessManager {
process_id: response_process_id,
exit_code,
original_token_count: Some(original_token_count),
output_omitted_bytes,
hook_command: Some(request.hook_command.clone()),
};
@@ -699,7 +719,7 @@ impl UnifiedExecProcessManager {
};
let start = Instant::now();
let deadline = start + Duration::from_millis(yield_time_ms);
let collected = Self::collect_output_until_deadline(
let collected_output = Self::collect_output_until_deadline(
&output_buffer,
&output_notify,
&output_closed,
@@ -711,8 +731,12 @@ impl UnifiedExecProcessManager {
.await;
let wall_time = Instant::now().saturating_duration_since(start);
let text = String::from_utf8_lossy(&collected).to_string();
let original_token_count = approx_token_count(&text);
let original_token_count = usize::try_from(approx_tokens_from_byte_count(
collected_output.total_bytes(),
))
.unwrap_or(usize::MAX);
let output_omitted_bytes = NonZeroUsize::new(collected_output.omitted_bytes());
let collected = collected_output.to_bytes_with_omission_marker();
let chunk_id = generate_chunk_id();
if network_approval
.as_ref()
@@ -782,6 +806,7 @@ impl UnifiedExecProcessManager {
process_id,
exit_code,
original_token_count: Some(original_token_count),
output_omitted_bytes,
hook_command: Some(hook_command),
};
@@ -1207,10 +1232,10 @@ impl UnifiedExecProcessManager {
cancellation_token: &CancellationToken,
mut pause_state: Option<watch::Receiver<bool>>,
mut deadline: Instant,
) -> Vec<u8> {
) -> HeadTailBuffer {
const POST_EXIT_CLOSE_WAIT_CAP: Duration = Duration::from_millis(50);
let mut collected: Vec<u8> = Vec::with_capacity(4096);
let mut collected = HeadTailBuffer::default();
let mut exit_signal_received = cancellation_token.is_cancelled();
let mut post_exit_deadline: Option<Instant> = None;
loop {
@@ -1220,17 +1245,20 @@ impl UnifiedExecProcessManager {
&mut post_exit_deadline,
)
.await;
let drained_chunks: Vec<Vec<u8>>;
let drained_output: HeadTailBuffer;
let has_drained_output: bool;
let mut wait_for_output = None;
{
let mut guard = output_buffer.lock().await;
drained_chunks = guard.drain_chunks();
if drained_chunks.is_empty() {
drained_output = guard.drain();
has_drained_output =
drained_output.retained_bytes() > 0 || drained_output.omitted_bytes() > 0;
if !has_drained_output {
wait_for_output = Some(output_notify.notified());
}
}
if drained_chunks.is_empty() {
if !has_drained_output {
exit_signal_received |= cancellation_token.is_cancelled();
if exit_signal_received && output_closed.load(std::sync::atomic::Ordering::Acquire)
{
@@ -1275,9 +1303,7 @@ impl UnifiedExecProcessManager {
continue;
}
for chunk in drained_chunks {
collected.extend_from_slice(&chunk);
}
collected.push_buffer(drained_output);
exit_signal_received |= cancellation_token.is_cancelled();
if Instant::now() >= deadline {

87
codex-rs/core/src/unified_exec/process_manager_tests.rs
View File

@@ -234,6 +234,93 @@ fn initial_exec_yield_time_has_no_platform_floor() {
);
}
#[tokio::test]
async fn output_collection_stays_bounded_across_repeated_drains() {
let output_buffer = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default()));
let output_notify = Arc::new(Notify::new());
let output_closed = Arc::new(AtomicBool::new(false));
let output_closed_notify = Arc::new(Notify::new());
let cancellation_token = CancellationToken::new();
let collect = UnifiedExecProcessManager::collect_output_until_deadline(
&output_buffer,
&output_notify,
&output_closed,
&output_closed_notify,
&cancellation_token,
/*pause_state*/ None,
Instant::now() + Duration::from_secs(5),
);
let produce = async {
for byte in [b'a', b'b', b'c'] {
output_buffer.lock().await.push_chunk(
vec![byte; crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES],
);
output_notify.notify_one();
tokio::time::timeout(Duration::from_secs(1), async {
loop {
if output_buffer.lock().await.retained_bytes() == 0 {
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("collector should drain each chunk");
}
output_closed.store(true, Ordering::Release);
cancellation_token.cancel();
output_closed_notify.notify_waiters();
output_notify.notify_waiters();
};
let (collected, ()) = tokio::join!(collect, produce);
let mut expected = HeadTailBuffer::default();
for byte in [b'a', b'b', b'c'] {
expected.push_chunk(vec![
byte;
crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES
]);
}
assert_eq!(collected, expected);
}
#[tokio::test]
async fn output_collection_preserves_omissions_from_drained_buffer() {
let mut buffered_output = HeadTailBuffer::default();
buffered_output.push_chunk(vec![
b'a';
crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES
]);
buffered_output.push_chunk(b"overflow".to_vec());
let mut expected = HeadTailBuffer::default();
expected.push_chunk(vec![
b'a';
crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES
]);
expected.push_chunk(b"overflow".to_vec());
let output_buffer = Arc::new(tokio::sync::Mutex::new(buffered_output));
let output_notify = Arc::new(Notify::new());
let output_closed = Arc::new(AtomicBool::new(true));
let output_closed_notify = Arc::new(Notify::new());
let cancellation_token = CancellationToken::new();
cancellation_token.cancel();
let collected = UnifiedExecProcessManager::collect_output_until_deadline(
&output_buffer,
&output_notify,
&output_closed,
&output_closed_notify,
&cancellation_token,
/*pause_state*/ None,
Instant::now() + Duration::from_secs(1),
)
.await;
assert_eq!(collected, expected);
}
#[tokio::test]
async fn network_denial_fallback_message_names_sandbox_network_proxy() {
let message = network_denial_message_for_session(/*session*/ None, /*deferred*/ None).await;