diff --git a/codex-rs/core/src/unified_exec/async_watcher.rs b/codex-rs/core/src/unified_exec/async_watcher.rs index 62791dda7f..d20a85843b 100644 --- a/codex-rs/core/src/unified_exec/async_watcher.rs +++ b/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::::new(); + let mut pending = VecDeque::::new(); let mut emitted_deltas: usize = 0; let mut grace_sleep: Option>> = None; @@ -220,7 +221,7 @@ pub(crate) fn spawn_exit_watcher( } async fn process_chunk( - pending: &mut Vec, + pending: &mut VecDeque, transcript: &Arc>, call_id: &str, session_ref: &Arc, @@ -228,7 +229,7 @@ async fn process_chunk( emitted_deltas: &mut usize, chunk: Vec, ) { - 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) -> Option> { +fn split_valid_utf8_prefix(buffer: &mut VecDeque) -> Option> { split_valid_utf8_prefix_with_max(buffer, UNIFIED_EXEC_OUTPUT_DELTA_MAX_BYTES) } -fn split_valid_utf8_prefix_with_max(buffer: &mut Vec, max_bytes: usize) -> Option> { +fn split_valid_utf8_prefix_with_max( + buffer: &mut VecDeque, + max_bytes: usize, +) -> Option> { 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( diff --git a/codex-rs/core/src/unified_exec/async_watcher_tests.rs b/codex-rs/core/src/unified_exec/async_watcher_tests.rs index fe731c40bd..66fd8dba11 100644 --- a/codex-rs/core/src/unified_exec/async_watcher_tests.rs +++ b/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()); } diff --git a/codex-rs/exec-server/src/noise_relay/message_framing.rs b/codex-rs/exec-server/src/noise_relay/message_framing.rs index 402b68691a..b15673fe32 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing.rs +++ b/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, /// here so a bad authenticated peer cannot grow the reassembly buffer forever. #[derive(Default)] pub(crate) struct JsonRpcMessageDecoder { - buffered: Vec, + 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. diff --git a/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs b/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs index 3b9ac7939e..563e9e3bad 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs +++ b/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]); +}