send-aggregated output (#2364)

We want to send an aggregated output of stderr and stdout so we don't
have to aggregate it stderr+stdout as we lose order sometimes.

---------

Co-authored-by: Gabriel Peal <gpeal@users.noreply.github.com>

codex-rs/core/src/exec.rs

@@ -28,18 +28,17 @@ use crate::spawn::StdioPolicy;
 use crate::spawn::spawn_child_async;
 use serde_bytes::ByteBuf;
 
-// Maximum we send for each stream, which is either:
-// - 10KiB OR
-// - 256 lines
-const MAX_STREAM_OUTPUT: usize = 10 * 1024;
-const MAX_STREAM_OUTPUT_LINES: usize = 256;
-
 const DEFAULT_TIMEOUT_MS: u64 = 10_000;
 
 // Hardcode these since it does not seem worth including the libc crate just
 // for these.
 const SIGKILL_CODE: i32 = 9;
 const TIMEOUT_CODE: i32 = 64;
+const EXIT_CODE_SIGNAL_BASE: i32 = 128; // conventional shell: 128 + signal
+
+// I/O buffer sizing
+const READ_CHUNK_SIZE: usize = 8192; // bytes per read
+const AGGREGATE_BUFFER_INITIAL_CAPACITY: usize = 8 * 1024; // 8 KiB
 
 #[derive(Debug, Clone)]
 pub struct ExecParams {
@@ -153,6 +152,7 @@ pub async fn process_exec_tool_call(
                 exit_code,
                 stdout,
                 stderr,
+                aggregated_output: raw_output.aggregated_output.from_utf8_lossy(),
                 duration,
             })
         }
@@ -189,10 +189,11 @@ pub struct StreamOutput<T> {
     pub truncated_after_lines: Option<u32>,
 }
 #[derive(Debug)]
-pub struct RawExecToolCallOutput {
+struct RawExecToolCallOutput {
     pub exit_status: ExitStatus,
     pub stdout: StreamOutput<Vec<u8>>,
     pub stderr: StreamOutput<Vec<u8>>,
+    pub aggregated_output: StreamOutput<Vec<u8>>,
 }
 
 impl StreamOutput<String> {
@@ -213,11 +214,17 @@ impl StreamOutput<Vec<u8>> {
     }
 }
 
+#[inline]
+fn append_all(dst: &mut Vec<u8>, src: &[u8]) {
+    dst.extend_from_slice(src);
+}
+
 #[derive(Debug)]
 pub struct ExecToolCallOutput {
     pub exit_code: i32,
     pub stdout: StreamOutput<String>,
     pub stderr: StreamOutput<String>,
+    pub aggregated_output: StreamOutput<String>,
     pub duration: Duration,
 }
 
@@ -253,7 +260,7 @@ async fn exec(
 
 /// Consumes the output of a child process, truncating it so it is suitable for
 /// use as the output of a `shell` tool call. Also enforces specified timeout.
-pub(crate) async fn consume_truncated_output(
+async fn consume_truncated_output(
     mut child: Child,
     timeout: Duration,
     stdout_stream: Option<StdoutStream>,
@@ -273,19 +280,19 @@ pub(crate) async fn consume_truncated_output(
         ))
     })?;
 
+    let (agg_tx, agg_rx) = async_channel::unbounded::<Vec<u8>>();
+
     let stdout_handle = tokio::spawn(read_capped(
         BufReader::new(stdout_reader),
-        MAX_STREAM_OUTPUT,
-        MAX_STREAM_OUTPUT_LINES,
         stdout_stream.clone(),
         false,
+        Some(agg_tx.clone()),
     ));
     let stderr_handle = tokio::spawn(read_capped(
         BufReader::new(stderr_reader),
-        MAX_STREAM_OUTPUT,
-        MAX_STREAM_OUTPUT_LINES,
         stdout_stream.clone(),
         true,
+        Some(agg_tx.clone()),
     ));
 
     let exit_status = tokio::select! {
@@ -297,38 +304,48 @@ pub(crate) async fn consume_truncated_output(
                     // timeout
                     child.start_kill()?;
                     // Debatable whether `child.wait().await` should be called here.
-                    synthetic_exit_status(128 + TIMEOUT_CODE)
+                    synthetic_exit_status(EXIT_CODE_SIGNAL_BASE + TIMEOUT_CODE)
                 }
             }
         }
         _ = tokio::signal::ctrl_c() => {
             child.start_kill()?;
-            synthetic_exit_status(128 + SIGKILL_CODE)
+            synthetic_exit_status(EXIT_CODE_SIGNAL_BASE + SIGKILL_CODE)
         }
     };
 
     let stdout = stdout_handle.await??;
     let stderr = stderr_handle.await??;
 
+    drop(agg_tx);
+
+    let mut combined_buf = Vec::with_capacity(AGGREGATE_BUFFER_INITIAL_CAPACITY);
+    while let Ok(chunk) = agg_rx.recv().await {
+        append_all(&mut combined_buf, &chunk);
+    }
+    let aggregated_output = StreamOutput {
+        text: combined_buf,
+        truncated_after_lines: None,
+    };
+
     Ok(RawExecToolCallOutput {
         exit_status,
         stdout,
         stderr,
+        aggregated_output,
     })
 }
 
 async fn read_capped<R: AsyncRead + Unpin + Send + 'static>(
     mut reader: R,
-    max_output: usize,
-    max_lines: usize,
     stream: Option<StdoutStream>,
     is_stderr: bool,
+    aggregate_tx: Option<Sender<Vec<u8>>>,
 ) -> io::Result<StreamOutput<Vec<u8>>> {
-    let mut buf = Vec::with_capacity(max_output.min(8 * 1024));
-    let mut tmp = [0u8; 8192];
+    let mut buf = Vec::with_capacity(AGGREGATE_BUFFER_INITIAL_CAPACITY);
+    let mut tmp = [0u8; READ_CHUNK_SIZE];
 
-    let mut remaining_bytes = max_output;
-    let mut remaining_lines = max_lines;
+    // No caps: append all bytes
 
     loop {
         let n = reader.read(&mut tmp).await?;
@@ -355,33 +372,17 @@ async fn read_capped<R: AsyncRead + Unpin + Send + 'static>(
             let _ = stream.tx_event.send(event).await;
         }
 
-        // Copy into the buffer only while we still have byte and line budget.
-        if remaining_bytes > 0 && remaining_lines > 0 {
-            let mut copy_len = 0;
-            for &b in &tmp[..n] {
-                if remaining_bytes == 0 || remaining_lines == 0 {
-                    break;
-                }
-                copy_len += 1;
-                remaining_bytes -= 1;
-                if b == b'\n' {
-                    remaining_lines -= 1;
-                }
-            }
-            buf.extend_from_slice(&tmp[..copy_len]);
+        if let Some(tx) = &aggregate_tx {
+            let _ = tx.send(tmp[..n].to_vec()).await;
         }
-        // Continue reading to EOF to avoid back-pressure, but discard once caps are hit.
-    }
 
-    let truncated = remaining_lines == 0 || remaining_bytes == 0;
+        append_all(&mut buf, &tmp[..n]);
+        // Continue reading to EOF to avoid back-pressure
+    }
 
     Ok(StreamOutput {
         text: buf,
-        truncated_after_lines: if truncated {
-            Some((max_lines - remaining_lines) as u32)
-        } else {
-            None
-        },
+        truncated_after_lines: None,
     })
 }