Avoid shifting bytes in streaming output buffers (#36194)

## Why

Removing each decoded prefix from a `Vec` shifts all remaining bytes. Repeating
that work is costly for streams containing many invalid UTF-8 bytes or many
framed messages in one relay record.

## What changed

- Buffer unified exec output in a `VecDeque` and use UTF-8 error offsets to
  consume the largest valid prefix while continuing to make progress on invalid
  input.
- Buffer Noise relay messages in `BytesMut` and advance past decoded frames
  without moving the remaining bytes.

## Testing

Added regression coverage for valid data surrounding invalid UTF-8, front
consumption without shifting, and reassembling many JSON-RPC messages from one
record.

GitOrigin-RevId: 661ccfffc68960ac6c9d37c738224da8ea83cbf0

codex-rs/core/src/unified_exec/async_watcher.rs

@@ -1,3 +1,4 @@
+use std::collections::VecDeque;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::sync::atomic::Ordering;
@@ -61,7 +62,7 @@ pub(crate) fn start_streaming_output(
     tokio::spawn(async move {
         use tokio::sync::broadcast::error::RecvError;
 
-        let mut pending = Vec::<u8>::new();
+        let mut pending = VecDeque::<u8>::new();
         let mut emitted_deltas: usize = 0;
 
         let mut grace_sleep: Option<Pin<Box<Sleep>>> = None;
@@ -220,7 +221,7 @@ pub(crate) fn spawn_exit_watcher(
 }
 
 async fn process_chunk(
-    pending: &mut Vec<u8>,
+    pending: &mut VecDeque<u8>,
     transcript: &Arc<Mutex<HeadTailBuffer>>,
     call_id: &str,
     session_ref: &Arc<Session>,
@@ -228,7 +229,7 @@ async fn process_chunk(
     emitted_deltas: &mut usize,
     chunk: Vec<u8>,
 ) {
-    pending.extend_from_slice(&chunk);
+    pending.extend(chunk);
     while let Some(prefix) = split_valid_utf8_prefix(pending) {
         {
             let mut guard = transcript.lock().await;
@@ -354,34 +355,24 @@ pub(crate) async fn emit_failed_exec_end_for_unified_exec(
         .await;
 }
 
-fn split_valid_utf8_prefix(buffer: &mut Vec<u8>) -> Option<Vec<u8>> {
+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 Vec<u8>, max_bytes: usize) -> Option<Vec<u8>> {
+fn split_valid_utf8_prefix_with_max(
+    buffer: &mut VecDeque<u8>,
+    max_bytes: usize,
+) -> Option<Vec<u8>> {
     if buffer.is_empty() {
         return None;
     }
 
     let max_len = buffer.len().min(max_bytes);
-    let mut split = max_len;
-    while split > 0 {
-        if std::str::from_utf8(&buffer[..split]).is_ok() {
-            let prefix = buffer[..split].to_vec();
-            buffer.drain(..split);
-            return Some(prefix);
-        }
-
-        if max_len - split > 4 {
-            break;
-        }
-        split -= 1;
-    }
-
-    // If no valid UTF-8 prefix was found, emit the first byte so the stream
-    // keeps making progress and the transcript reflects all bytes.
-    let byte = buffer.drain(..1).collect();
-    Some(byte)
+    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())
 }
 
 async fn resolve_aggregated_output(

codex-rs/core/src/unified_exec/async_watcher_tests.rs

@@ -1,3 +1,4 @@
+use std::collections::VecDeque;
 use std::sync::Arc;
 
 use super::TRAILING_OUTPUT_GRACE;
@@ -203,36 +204,76 @@ 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 = b"hello word!".to_vec();
+    let mut buf = VecDeque::from(b"hello word!".to_vec());
 
     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, b" word!".to_vec());
+    assert_eq!(buf, VecDeque::from(b" word!".to_vec()));
 
     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, b"!".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 = "ééé".as_bytes().to_vec();
+    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, "éé".as_bytes().to_vec());
+    assert_eq!(buf, VecDeque::from("éé".as_bytes().to_vec()));
 }
 
 #[test]
 fn split_valid_utf8_prefix_makes_progress_on_invalid_utf8() {
-    let mut buf = vec![0xff, b'a', b'b'];
+    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, b"ab".to_vec());
+    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));
+        }
+    }
+
+    assert!(buf.is_empty());
 }

codex-rs/exec-server/src/noise_relay/message_framing.rs

@@ -1,3 +1,5 @@
+use bytes::Buf;
+use bytes::BytesMut;
 use codex_exec_server_protocol::JSONRPCMessage;
 
 use crate::ExecServerError;
@@ -31,7 +33,7 @@ pub(crate) fn frame_jsonrpc_message(message: &JSONRPCMessage) -> Result<Vec<u8>,
 /// here so a bad authenticated peer cannot grow the reassembly buffer forever.
 #[derive(Default)]
 pub(crate) struct JsonRpcMessageDecoder {
-    buffered: Vec<u8>,
+    buffered: BytesMut,
 }
 
 impl JsonRpcMessageDecoder {
@@ -67,7 +69,7 @@ impl JsonRpcMessageDecoder {
             messages.push(serde_json::from_slice(
                 &self.buffered[LENGTH_PREFIX_BYTES..framed_len],
             )?);
-            self.buffered.drain(..framed_len);
+            self.buffered.advance(framed_len);
         }
 
         // Even before a message is complete, keep reassembly memory bounded.

codex-rs/exec-server/src/noise_relay/message_framing_tests.rs

@@ -50,3 +50,19 @@ fn rejects_oversized_plaintext_record() {
             if message == "Noise relay plaintext record exceeds maximum length"
     ));
 }
+
+#[test]
+fn reassembles_many_messages_from_one_record() {
+    let message = JSONRPCMessage::Notification(JSONRPCNotification {
+        method: "small/test".to_string(),
+        params: None,
+    });
+    let framed = frame_jsonrpc_message(&message).expect("frame message");
+    let message_count = NOISE_RECORD_PLAINTEXT_LEN / framed.len();
+    let record = framed.repeat(message_count);
+
+    let mut decoder = JsonRpcMessageDecoder::default();
+    let decoded = decoder.push(&record).expect("decode record");
+
+    assert_eq!(decoded, vec![message; message_count]);
+}