Bound unified exec output delta frames (#39937)

## What changed

- Cap each `ExecCommandOutputDelta` frame at 8 KiB and stop building frames once the per-call delta quota is exhausted.
- Preserve complete UTF-8 scalars across producer chunks while passing malformed or incomplete trailing bytes through without losing them.
- Append all process output to the transcript independently of delta emission.

## Testing

- Cover split multibyte characters, malformed and trailing bytes, frame and event limits, transcript preservation, and interactive PTY output.

GitOrigin-RevId: ef4100e353df01569a8fbae650c70f439920a2fd
This commit is contained in:
Tamir Duberstein
2026-08-21 14:04:33 +00:00
committed by copyberry
parent 7f9832d0d0
commit 748d8ac834
3 changed files with 340 additions and 162 deletions
+130 -70
View File
@@ -1,4 +1,3 @@
use std::collections::VecDeque;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::Ordering;
@@ -41,6 +40,19 @@ pub(crate) const TRAILING_OUTPUT_GRACE: Duration = Duration::from_millis(100);
/// process arbitrarily large delta payloads.
const UNIFIED_EXEC_OUTPUT_DELTA_MAX_BYTES: usize = 8192;
struct Emitter {
remaining_deltas: usize,
session: Arc<Session>,
turn: Arc<TurnContext>,
call_id: String,
}
struct Buffer<const MAX_BYTES: usize = UNIFIED_EXEC_OUTPUT_DELTA_MAX_BYTES> {
pending: Vec<u8>,
transcript: Arc<Mutex<HeadTailBuffer>>,
emitter: Emitter,
}
/// Spawn a background task that continuously reads from the PTY, appends to the
/// shared transcript, and emits ExecCommandOutputDelta events on UTF‑8
/// boundaries.
@@ -58,15 +70,21 @@ pub(crate) fn start_streaming_output(
..
} = process.output_handles().clone();
let session_ref = Arc::clone(&context.session);
let turn_ref = Arc::clone(&context.step_context.turn);
let call_id = context.call_id.clone();
let emitter = Emitter {
remaining_deltas: MAX_EXEC_OUTPUT_DELTAS_PER_CALL,
session: Arc::clone(&context.session),
turn: Arc::clone(&context.step_context.turn),
call_id: context.call_id.clone(),
};
tokio::spawn(async move {
use tokio::sync::broadcast::error::RecvError;
let mut pending = VecDeque::<u8>::new();
let mut emitted_deltas: usize = 0;
let mut output: Buffer = Buffer {
pending: Vec::new(),
transcript,
emitter,
};
let mut grace_sleep: Option<Pin<Box<Sleep>>> = None;
let output_closed_notified = output_closed_notify.notified();
@@ -112,15 +130,7 @@ pub(crate) fn start_streaming_output(
}
};
process_chunk(
&mut pending,
&transcript,
&call_id,
&session_ref,
&turn_ref,
&mut emitted_deltas,
chunk,
).await;
output.push(chunk).await;
}
}
}
@@ -140,18 +150,11 @@ pub(crate) fn start_streaming_output(
) => break,
};
process_chunk(
&mut pending,
&transcript,
&call_id,
&session_ref,
&turn_ref,
&mut emitted_deltas,
chunk,
)
.await;
output.push(chunk).await;
}
}
output.finish().await;
output_drained.notify_one();
});
}
@@ -236,35 +239,92 @@ pub(crate) fn spawn_exit_watcher(
});
}
async fn process_chunk(
pending: &mut VecDeque<u8>,
transcript: &Arc<Mutex<HeadTailBuffer>>,
call_id: &str,
session_ref: &Arc<Session>,
turn_ref: &Arc<TurnContext>,
emitted_deltas: &mut usize,
chunk: Vec<u8>,
) {
pending.extend(chunk);
while let Some(prefix) = split_valid_utf8_prefix(pending) {
{
let mut guard = transcript.lock().await;
guard.push_chunk(&prefix);
}
if *emitted_deltas >= MAX_EXEC_OUTPUT_DELTAS_PER_CALL {
continue;
}
let event = ExecCommandOutputDeltaEvent {
call_id: call_id.to_string(),
stream: ExecOutputStream::Stdout,
chunk: prefix,
impl<const MAX_BYTES: usize> Buffer<MAX_BYTES> {
async fn push(&mut self, mut bytes: Vec<u8>) {
const {
assert!(
MAX_BYTES >= char::MAX.len_utf8(),
"a frame must fit one UTF-8 scalar"
)
};
session_ref
.send_event(turn_ref.as_ref(), EventMsg::ExecCommandOutputDelta(event))
.await;
*emitted_deltas += 1;
let Self {
pending,
transcript,
emitter,
} = self;
transcript.lock().await.push_chunk(&bytes);
// Reuse a producer chunk when it fits, retaining only an incomplete
// UTF-8 suffix for the next push.
if pending.is_empty() && bytes.len() <= MAX_BYTES {
emitter
.emit(|| {
let complete = utf8_boundary(&bytes);
pending.extend(bytes.drain(complete..));
bytes
})
.await;
return;
}
let mut bytes = bytes.as_slice();
let mut next_chunk = || {
let space = MAX_BYTES.saturating_sub(pending.len());
let (prefix, rest) = bytes.split_at_checked(space).unwrap_or((bytes, &[]));
let mut chunk = Vec::with_capacity(pending.len().saturating_add(prefix.len()));
chunk.append(pending); // Empties pending.
chunk.extend_from_slice(prefix);
bytes = rest;
let complete = utf8_boundary(&chunk);
// Only the incomplete suffix passes through pending.
pending.extend(chunk.drain(complete..));
chunk
};
while emitter.emit(&mut next_chunk).await {}
}
async fn finish(self) {
let Self {
pending,
transcript: _,
mut emitter,
} = self;
debug_assert!(
pending.len() < char::MAX.len_utf8(),
"only an incomplete UTF-8 scalar can remain"
);
emitter.emit(|| pending).await;
}
}
impl Emitter {
/// Build a frame only while quota remains. Returns whether a nonempty frame was sent.
async fn emit(&mut self, make_chunk: impl FnOnce() -> Vec<u8>) -> bool {
let Self {
remaining_deltas,
session,
turn,
call_id,
} = self;
let Some(remaining) = remaining_deltas.checked_sub(1) else {
return false;
};
let chunk = make_chunk();
let emit = !chunk.is_empty();
if emit {
let event = ExecCommandOutputDeltaEvent {
call_id: call_id.clone(),
stream: ExecOutputStream::Stdout,
chunk,
};
session
.send_event(turn.as_ref(), EventMsg::ExecCommandOutputDelta(event))
.await;
*remaining_deltas = remaining;
}
emit
}
}
@@ -371,24 +431,24 @@ pub(crate) async fn emit_failed_exec_end_for_unified_exec(
.await;
}
fn split_valid_utf8_prefix(buffer: &mut VecDeque<u8>) -> Option<Vec<u8>> {
split_valid_utf8_prefix_with_max(buffer, UNIFIED_EXEC_OUTPUT_DELTA_MAX_BYTES)
}
fn split_valid_utf8_prefix_with_max(
buffer: &mut VecDeque<u8>,
max_bytes: usize,
) -> Option<Vec<u8>> {
if buffer.is_empty() {
return None;
/// Keep only a potentially incomplete UTF-8 suffix; malformed bytes remain raw.
/// A UTF-8 scalar spans at most four bytes, so its incomplete tail fits in three.
fn utf8_boundary(bytes: &[u8]) -> usize {
let mut boundary = bytes.len().saturating_sub(char::MAX.len_utf8() - 1);
while boundary < bytes.len() {
match std::str::from_utf8(&bytes[boundary..]) {
Ok(s) => return boundary + s.len(),
Err(error) => {
boundary += error.valid_up_to();
if let Some(invalid_len) = error.error_len() {
boundary += invalid_len;
} else {
return boundary;
}
}
}
}
let max_len = buffer.len().min(max_bytes);
let split = match std::str::from_utf8(&buffer.make_contiguous()[..max_len]) {
Ok(_) => max_len,
Err(error) => error.valid_up_to().max(1),
};
Some(buffer.drain(..split).collect())
bytes.len()
}
async fn resolve_aggregated_output(
@@ -1,10 +1,11 @@
use std::collections::VecDeque;
use std::sync::Arc;
use super::Buffer;
use super::Emitter;
use super::TRAILING_OUTPUT_GRACE;
use super::spawn_exit_watcher;
use super::split_valid_utf8_prefix_with_max;
use super::start_streaming_output;
use super::utf8_boundary;
use crate::session::tests::make_session_and_context_with_rx;
use crate::unified_exec::UnifiedExecContext;
use crate::unified_exec::head_tail_buffer::HeadTailBuffer;
@@ -14,6 +15,8 @@ use codex_protocol::items::CommandExecutionStatus;
use codex_protocol::items::TurnItem;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::ExecCommandOutputDeltaEvent;
use codex_protocol::protocol::ExecOutputStream;
use codex_sandboxing::SandboxType;
use pretty_assertions::assert_eq;
@@ -67,6 +70,49 @@ async fn streaming_output_harness() -> anyhow::Result<StreamingOutputHarness> {
})
}
#[tokio::test]
async fn streaming_output_preserves_multibyte_characters_across_chunks() -> anyhow::Result<()> {
let StreamingOutputHarness {
process,
stdout_tx,
exit_tx,
transcript,
rx_event,
..
} = streaming_output_harness().await?;
let output_drained = process.output_drained_notify();
let drained = output_drained.notified();
tokio::pin!(drained);
stdout_tx.send(vec![0xc3]).expect("send UTF-8 lead byte");
stdout_tx
.send(vec![0xa9])
.expect("send UTF-8 continuation byte");
drop(stdout_tx);
exit_tx.send(0).expect("send exit");
(&mut drained).await;
let event = rx_event.recv().await.expect("receive output delta");
let EventMsg::ExecCommandOutputDelta(delta) = event.msg else {
panic!("expected ExecCommandOutputDelta");
};
assert_eq!(
delta,
ExecCommandOutputDeltaEvent {
call_id: "streaming-output-test".to_string(),
stream: ExecOutputStream::Stdout,
chunk: "é".as_bytes().to_vec(),
}
);
assert_eq!(
transcript.lock().await.to_bytes_with_omission_marker(),
"é".as_bytes()
);
assert!(rx_event.try_recv().is_err());
Ok(())
}
#[tokio::test]
async fn streaming_output_finishes_on_close_without_waiting_for_grace() -> anyhow::Result<()> {
let StreamingOutputHarness {
@@ -86,7 +132,7 @@ async fn streaming_output_finishes_on_close_without_waiting_for_grace() -> anyho
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(50)).await;
stdout_tx
.send(b"LATE-OUTPUT-MARKER".to_vec())
.send(b"LATE-OUTPUT-MARKER\xc3".to_vec())
.expect("send late output");
});
@@ -100,7 +146,7 @@ async fn streaming_output_finishes_on_close_without_waiting_for_grace() -> anyho
);
assert_eq!(
transcript.lock().await.to_bytes_with_omission_marker(),
b"LATE-OUTPUT-MARKER"
b"LATE-OUTPUT-MARKER\xc3"
);
Ok(())
@@ -110,8 +156,10 @@ async fn streaming_output_finishes_on_close_without_waiting_for_grace() -> anyho
async fn streaming_output_keeps_grace_as_fallback_without_close() -> anyhow::Result<()> {
let StreamingOutputHarness {
process,
stdout_tx: _stdout_tx,
stdout_tx,
exit_tx,
transcript,
rx_event,
..
} = streaming_output_harness().await?;
let output_drained = process.output_drained_notify();
@@ -120,6 +168,7 @@ async fn streaming_output_keeps_grace_as_fallback_without_close() -> anyhow::Res
tokio::time::pause();
let exited_at = Instant::now();
stdout_tx.send(vec![0xc3]).expect("send UTF-8 lead byte");
exit_tx.send(0).expect("send exit");
(&mut drained).await;
let elapsed = Instant::now().saturating_duration_since(exited_at);
@@ -130,6 +179,23 @@ async fn streaming_output_keeps_grace_as_fallback_without_close() -> anyhow::Res
&& elapsed <= TRAILING_OUTPUT_GRACE + Duration::from_millis(10),
"missing output close should use the grace fallback: {elapsed:?}"
);
assert_eq!(
transcript.lock().await.to_bytes_with_omission_marker(),
vec![0xc3]
);
let event = rx_event.try_recv().expect("receive final output delta");
let EventMsg::ExecCommandOutputDelta(delta) = event.msg else {
panic!("expected ExecCommandOutputDelta");
};
assert_eq!(
delta,
ExecCommandOutputDeltaEvent {
call_id: "streaming-output-test".to_string(),
stream: ExecOutputStream::Stdout,
chunk: vec![0xc3],
}
);
assert!(rx_event.try_recv().is_err());
Ok(())
}
@@ -208,77 +274,70 @@ async fn exit_watcher_waits_for_late_network_denial_before_classifying_end() ->
}
#[test]
fn split_valid_utf8_prefix_respects_max_bytes_for_ascii() {
let mut buf = VecDeque::from(b"hello word!".to_vec());
fn utf8_boundary_preserves_complete_characters() {
assert_eq!(utf8_boundary(b"hello"), 5);
let first =
split_valid_utf8_prefix_with_max(&mut buf, /*max_bytes*/ 5).expect("expected prefix");
assert_eq!(first, b"hello".to_vec());
assert_eq!(buf, VecDeque::from(b" word!".to_vec()));
let bytes = "ééé".as_bytes();
assert_eq!(utf8_boundary(&bytes[..3]), 2);
let second =
split_valid_utf8_prefix_with_max(&mut buf, /*max_bytes*/ 5).expect("expected prefix");
assert_eq!(second, b" word".to_vec());
assert_eq!(buf, VecDeque::from(b"!".to_vec()));
}
#[test]
fn split_valid_utf8_prefix_avoids_splitting_utf8_codepoints() {
// "é" is 2 bytes in UTF-8. With a max of 3 bytes, we should only emit 1 char (2 bytes).
let mut buf = VecDeque::from("ééé".as_bytes().to_vec());
let first =
split_valid_utf8_prefix_with_max(&mut buf, /*max_bytes*/ 3).expect("expected prefix");
assert_eq!(std::str::from_utf8(&first).unwrap(), "é");
assert_eq!(buf, VecDeque::from("éé".as_bytes().to_vec()));
}
#[test]
fn split_valid_utf8_prefix_makes_progress_on_invalid_utf8() {
let mut buf = VecDeque::from(vec![0xff, b'a', b'b']);
let first =
split_valid_utf8_prefix_with_max(&mut buf, /*max_bytes*/ 2).expect("expected prefix");
assert_eq!(first, vec![0xff]);
assert_eq!(buf, VecDeque::from(b"ab".to_vec()));
}
#[test]
fn split_valid_utf8_prefix_consumes_all_valid_bytes_before_invalid_utf8() {
let mut bytes = vec![b'a'; 4096];
bytes.push(0xff);
bytes.extend(vec![b'b'; 4096]);
let mut buf = VecDeque::from(bytes);
let first =
split_valid_utf8_prefix_with_max(&mut buf, /*max_bytes*/ 8192).expect("expected prefix");
assert_eq!(first, vec![b'a'; 4096]);
let second =
split_valid_utf8_prefix_with_max(&mut buf, /*max_bytes*/ 8192).expect("expected prefix");
assert_eq!(second, vec![0xff]);
let third =
split_valid_utf8_prefix_with_max(&mut buf, /*max_bytes*/ 8192).expect("expected prefix");
assert_eq!(third, vec![b'b'; 4096]);
assert!(buf.is_empty());
}
#[test]
fn split_invalid_utf8_advances_without_shifting_remaining_bytes() {
let mut buf = VecDeque::from(vec![0xff; 1024]);
let initial = buf.as_slices().0.as_ptr();
for offset in 0..1024 {
assert_eq!(
split_valid_utf8_prefix_with_max(&mut buf, /*max_bytes*/ 128),
Some(vec![0xff])
);
if let Some(first) = buf.as_slices().0.first() {
assert_eq!(first, &0xff);
assert_eq!(buf.as_slices().0.as_ptr(), initial.wrapping_add(offset + 1));
}
let bytes = "😀".as_bytes();
assert_eq!(utf8_boundary(bytes), bytes.len());
for len in 1..bytes.len() {
assert_eq!(utf8_boundary(&bytes[..len]), 0);
}
assert!(buf.is_empty());
assert_eq!(utf8_boundary(&[0xf0, 0x9f, 0x98, 0x80, 0xc3]), 4);
}
#[test]
fn utf8_boundary_batches_malformed_output() {
assert_eq!(utf8_boundary(&[0xff, b'a', b'b']), 3);
assert_eq!(utf8_boundary(&[0xff, 0xc3]), 1);
assert_eq!(utf8_boundary(&[0xff, 0xc3, 0xa9]), 3);
assert_eq!(utf8_boundary(&[0xe0, 0x80]), 2);
assert_eq!(utf8_boundary(b"a\xffbbb"), 5);
}
#[tokio::test]
async fn streaming_output_bounds_invalid_bytes_and_keeps_the_full_transcript() {
let (session, turn, rx_event) = make_session_and_context_with_rx().await;
let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default()));
let mut output = Buffer::<8> {
pending: Vec::new(),
transcript: Arc::clone(&transcript),
emitter: Emitter {
remaining_deltas: 2,
session,
turn,
call_id: "bounded-output-test".to_string(),
},
};
// The first frame splits 😀; the last allowed frame leaves é incomplete.
let bytes = b"\xff\xff\xff\xff\xff\xff\xf0\x9f\x98\x80\xff\xff\xff\xc3\xa9";
output.push(bytes.to_vec()).await;
output.push(vec![0xfe, 0xfe]).await;
output.finish().await;
let mut chunks = Vec::new();
while let Ok(event) = rx_event.try_recv() {
let EventMsg::ExecCommandOutputDelta(delta) = event.msg else {
panic!("expected ExecCommandOutputDelta");
};
chunks.push(delta.chunk);
}
assert_eq!(
chunks,
vec![
b"\xff\xff\xff\xff\xff\xff".to_vec(),
b"\xf0\x9f\x98\x80\xff\xff\xff".to_vec(),
]
);
let mut expected_transcript = bytes.to_vec();
expected_transcript.extend([0xfe, 0xfe]);
assert_eq!(
transcript.lock().await.to_bytes_with_omission_marker(),
expected_transcript
);
}
+76 -17
View File
@@ -917,8 +917,21 @@ async fn unified_exec_emits_output_delta_for_exec_command() -> Result<()> {
let call_id = "uexec-delta-1";
let args = json!({
"cmd": "printf 'HELLO-UEXEC'",
"yield_time_ms": 1000,
"cmd": "stty -echo; read start; printf 'HELLO-UEXEC\\303'; read finish; printf '\\251\\303'",
"yield_time_ms": 250,
"tty": true,
});
let start_call_id = "uexec-delta-start";
let start_args = json!({
"chars": "start\n",
"session_id": 1000,
"yield_time_ms": 250,
});
let finish_call_id = "uexec-delta-finish";
let finish_args = json!({
"chars": "finish\n",
"session_id": 1000,
"yield_time_ms": 250,
});
let responses = vec![
@@ -929,30 +942,76 @@ async fn unified_exec_emits_output_delta_for_exec_command() -> Result<()> {
]),
sse(vec![
ev_response_created("resp-2"),
ev_assistant_message("msg-1", "finished"),
ev_function_call(
start_call_id,
"write_stdin",
&serde_json::to_string(&start_args)?,
),
ev_completed("resp-2"),
]),
sse(vec![
ev_response_created("resp-3"),
ev_function_call(
finish_call_id,
"write_stdin",
&serde_json::to_string(&finish_args)?,
),
ev_completed("resp-3"),
]),
sse(vec![
ev_response_created("resp-4"),
ev_assistant_message("msg-1", "finished"),
ev_completed("resp-4"),
]),
];
mount_sse_sequence(&server, responses).await;
submit_unified_exec_turn(&test, "emit delta", PermissionProfile::Disabled).await?;
let event = wait_for_event_match(&test.codex, |msg| match msg {
EventMsg::ExecCommandEnd(ev) if ev.call_id == call_id => Some(ev.clone()),
_ => None,
})
.await;
let mut deltas = Vec::new();
let mut end_event = None;
let mut turn_completed = false;
while !turn_completed || end_event.is_none() {
match wait_for_event(&test.codex, |_| true).await {
EventMsg::ExecCommandOutputDelta(event) => {
if event.call_id != call_id {
continue;
}
deltas.push(event.chunk);
}
EventMsg::ExecCommandEnd(event) => {
if event.call_id != call_id {
continue;
}
end_event = Some(event);
}
EventMsg::TurnComplete(_) => turn_completed = true,
_ => {}
}
}
let text = event.stdout;
assert!(
text.contains("HELLO-UEXEC"),
"delta chunk missing expected text: {text:?}",
assert_eq!(
(
deltas.concat(),
deltas
.iter()
.map(|bytes| String::from_utf8_lossy(bytes))
.collect::<String>(),
),
(
b"HELLO-UEXEC\xc3\xa9\xc3".to_vec(),
"HELLO-UEXECé�".to_string(),
)
);
let end_event = end_event.expect("expected command completion");
assert_eq!(
(
end_event.exit_code,
end_event.stdout,
end_event.aggregated_output
),
(0, "HELLO-UEXECé�".to_string(), "HELLO-UEXECé�".to_string())
);
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
Ok(())
}