diff --git a/codex-rs/core/src/exec.rs b/codex-rs/core/src/exec.rs index 275f2fb85..ea23cb121 100644 --- a/codex-rs/core/src/exec.rs +++ b/codex-rs/core/src/exec.rs @@ -312,20 +312,7 @@ async fn exec_windows_sandbox( text: stderr_text, truncated_after_lines: None, }; - // Best-effort aggregate: stdout then stderr (capped). - let mut aggregated = Vec::with_capacity( - stdout - .text - .len() - .saturating_add(stderr.text.len()) - .min(EXEC_OUTPUT_MAX_BYTES), - ); - append_capped(&mut aggregated, &stdout.text, EXEC_OUTPUT_MAX_BYTES); - append_capped(&mut aggregated, &stderr.text, EXEC_OUTPUT_MAX_BYTES); - let aggregated_output = StreamOutput { - text: aggregated, - truncated_after_lines: None, - }; + let aggregated_output = aggregate_output(&stdout, &stderr); Ok(RawExecToolCallOutput { exit_status, @@ -519,6 +506,39 @@ fn append_capped(dst: &mut Vec, src: &[u8], max_bytes: usize) { dst.extend_from_slice(&src[..take]); } +fn aggregate_output( + stdout: &StreamOutput>, + stderr: &StreamOutput>, +) -> StreamOutput> { + let total_len = stdout.text.len().saturating_add(stderr.text.len()); + let max_bytes = EXEC_OUTPUT_MAX_BYTES; + let mut aggregated = Vec::with_capacity(total_len.min(max_bytes)); + + if total_len <= max_bytes { + aggregated.extend_from_slice(&stdout.text); + aggregated.extend_from_slice(&stderr.text); + return StreamOutput { + text: aggregated, + truncated_after_lines: None, + }; + } + + // Under contention, reserve 1/3 for stdout and 2/3 for stderr; rebalance unused stderr to stdout. + let want_stdout = stdout.text.len().min(max_bytes / 3); + let want_stderr = stderr.text.len(); + let stderr_take = want_stderr.min(max_bytes.saturating_sub(want_stdout)); + let remaining = max_bytes.saturating_sub(want_stdout + stderr_take); + let stdout_take = want_stdout + remaining.min(stdout.text.len().saturating_sub(want_stdout)); + + aggregated.extend_from_slice(&stdout.text[..stdout_take]); + aggregated.extend_from_slice(&stderr.text[..stderr_take]); + + StreamOutput { + text: aggregated, + truncated_after_lines: None, + } +} + #[derive(Clone, Debug)] pub struct ExecToolCallOutput { pub exit_code: i32, @@ -683,20 +703,7 @@ async fn consume_truncated_output( Duration::from_millis(IO_DRAIN_TIMEOUT_MS), ) .await?; - // Best-effort aggregate: stdout then stderr (capped). - let mut aggregated = Vec::with_capacity( - stdout - .text - .len() - .saturating_add(stderr.text.len()) - .min(EXEC_OUTPUT_MAX_BYTES), - ); - append_capped(&mut aggregated, &stdout.text, EXEC_OUTPUT_MAX_BYTES); - append_capped(&mut aggregated, &stderr.text, EXEC_OUTPUT_MAX_BYTES * 2); - let aggregated_output = StreamOutput { - text: aggregated, - truncated_after_lines: None, - }; + let aggregated_output = aggregate_output(&stdout, &stderr); Ok(RawExecToolCallOutput { exit_status, @@ -771,6 +778,7 @@ fn synthetic_exit_status(code: i32) -> ExitStatus { #[cfg(test)] mod tests { use super::*; + use pretty_assertions::assert_eq; use std::time::Duration; use tokio::io::AsyncWriteExt; @@ -846,6 +854,85 @@ mod tests { assert_eq!(out.text.len(), EXEC_OUTPUT_MAX_BYTES); } + #[test] + fn aggregate_output_prefers_stderr_on_contention() { + let stdout = StreamOutput { + text: vec![b'a'; EXEC_OUTPUT_MAX_BYTES], + truncated_after_lines: None, + }; + let stderr = StreamOutput { + text: vec![b'b'; EXEC_OUTPUT_MAX_BYTES], + truncated_after_lines: None, + }; + + let aggregated = aggregate_output(&stdout, &stderr); + let stdout_cap = EXEC_OUTPUT_MAX_BYTES / 3; + let stderr_cap = EXEC_OUTPUT_MAX_BYTES.saturating_sub(stdout_cap); + + assert_eq!(aggregated.text.len(), EXEC_OUTPUT_MAX_BYTES); + assert_eq!(aggregated.text[..stdout_cap], vec![b'a'; stdout_cap]); + assert_eq!(aggregated.text[stdout_cap..], vec![b'b'; stderr_cap]); + } + + #[test] + fn aggregate_output_fills_remaining_capacity_with_stderr() { + let stdout_len = EXEC_OUTPUT_MAX_BYTES / 10; + let stdout = StreamOutput { + text: vec![b'a'; stdout_len], + truncated_after_lines: None, + }; + let stderr = StreamOutput { + text: vec![b'b'; EXEC_OUTPUT_MAX_BYTES], + truncated_after_lines: None, + }; + + let aggregated = aggregate_output(&stdout, &stderr); + let stderr_cap = EXEC_OUTPUT_MAX_BYTES.saturating_sub(stdout_len); + + assert_eq!(aggregated.text.len(), EXEC_OUTPUT_MAX_BYTES); + assert_eq!(aggregated.text[..stdout_len], vec![b'a'; stdout_len]); + assert_eq!(aggregated.text[stdout_len..], vec![b'b'; stderr_cap]); + } + + #[test] + fn aggregate_output_rebalances_when_stderr_is_small() { + let stdout = StreamOutput { + text: vec![b'a'; EXEC_OUTPUT_MAX_BYTES], + truncated_after_lines: None, + }; + let stderr = StreamOutput { + text: vec![b'b'; 1], + truncated_after_lines: None, + }; + + let aggregated = aggregate_output(&stdout, &stderr); + let stdout_len = EXEC_OUTPUT_MAX_BYTES.saturating_sub(1); + + assert_eq!(aggregated.text.len(), EXEC_OUTPUT_MAX_BYTES); + assert_eq!(aggregated.text[..stdout_len], vec![b'a'; stdout_len]); + assert_eq!(aggregated.text[stdout_len..], vec![b'b'; 1]); + } + + #[test] + fn aggregate_output_keeps_stdout_then_stderr_when_under_cap() { + let stdout = StreamOutput { + text: vec![b'a'; 4], + truncated_after_lines: None, + }; + let stderr = StreamOutput { + text: vec![b'b'; 3], + truncated_after_lines: None, + }; + + let aggregated = aggregate_output(&stdout, &stderr); + let mut expected = Vec::new(); + expected.extend_from_slice(&stdout.text); + expected.extend_from_slice(&stderr.text); + + assert_eq!(aggregated.text, expected); + assert_eq!(aggregated.truncated_after_lines, None); + } + #[cfg(unix)] #[test] fn sandbox_detection_flags_sigsys_exit_code() {