core: add a full-buffer exec capture policy (#15254)

codex-rs/core/src/exec_tests.rs

@@ -1,6 +1,7 @@
 use super::*;
 use codex_protocol::config_types::WindowsSandboxLevel;
 use pretty_assertions::assert_eq;
+use std::collections::HashMap;
 use std::time::Duration;
 use tokio::io::AsyncWriteExt;
 
@@ -91,14 +92,16 @@ fn sandbox_detection_ignores_network_policy_text_with_zero_exit_code() {
 }
 
 #[tokio::test]
-async fn read_capped_limits_retained_bytes() {
+async fn read_output_limits_retained_bytes_for_shell_capture() {
     let (mut writer, reader) = tokio::io::duplex(1024);
     let bytes = vec![b'a'; EXEC_OUTPUT_MAX_BYTES.saturating_add(128 * 1024)];
     tokio::spawn(async move {
         writer.write_all(&bytes).await.expect("write");
     });
 
-    let out = read_capped(reader, None, false).await.expect("read");
+    let out = read_output(reader, None, false, Some(EXEC_OUTPUT_MAX_BYTES))
+        .await
+        .expect("read");
     assert_eq!(out.text.len(), EXEC_OUTPUT_MAX_BYTES);
 }
 
@@ -113,7 +116,7 @@ fn aggregate_output_prefers_stderr_on_contention() {
         truncated_after_lines: None,
     };
 
-    let aggregated = aggregate_output(&stdout, &stderr);
+    let aggregated = aggregate_output(&stdout, &stderr, Some(EXEC_OUTPUT_MAX_BYTES));
     let stdout_cap = EXEC_OUTPUT_MAX_BYTES / 3;
     let stderr_cap = EXEC_OUTPUT_MAX_BYTES.saturating_sub(stdout_cap);
 
@@ -134,7 +137,7 @@ fn aggregate_output_fills_remaining_capacity_with_stderr() {
         truncated_after_lines: None,
     };
 
-    let aggregated = aggregate_output(&stdout, &stderr);
+    let aggregated = aggregate_output(&stdout, &stderr, Some(EXEC_OUTPUT_MAX_BYTES));
     let stderr_cap = EXEC_OUTPUT_MAX_BYTES.saturating_sub(stdout_len);
 
     assert_eq!(aggregated.text.len(), EXEC_OUTPUT_MAX_BYTES);
@@ -153,7 +156,7 @@ fn aggregate_output_rebalances_when_stderr_is_small() {
         truncated_after_lines: None,
     };
 
-    let aggregated = aggregate_output(&stdout, &stderr);
+    let aggregated = aggregate_output(&stdout, &stderr, Some(EXEC_OUTPUT_MAX_BYTES));
     let stdout_len = EXEC_OUTPUT_MAX_BYTES.saturating_sub(1);
 
     assert_eq!(aggregated.text.len(), EXEC_OUTPUT_MAX_BYTES);
@@ -172,7 +175,7 @@ fn aggregate_output_keeps_stdout_then_stderr_when_under_cap() {
         truncated_after_lines: None,
     };
 
-    let aggregated = aggregate_output(&stdout, &stderr);
+    let aggregated = aggregate_output(&stdout, &stderr, Some(EXEC_OUTPUT_MAX_BYTES));
     let mut expected = Vec::new();
     expected.extend_from_slice(&stdout.text);
     expected.extend_from_slice(&stderr.text);
@@ -181,6 +184,192 @@ fn aggregate_output_keeps_stdout_then_stderr_when_under_cap() {
     assert_eq!(aggregated.truncated_after_lines, None);
 }
 
+#[tokio::test]
+async fn read_output_retains_all_bytes_for_full_buffer_capture() {
+    let (mut writer, reader) = tokio::io::duplex(1024);
+    let bytes = vec![b'a'; EXEC_OUTPUT_MAX_BYTES.saturating_add(128 * 1024)];
+    let expected_len = bytes.len();
+    // The duplex pipe is smaller than `bytes`, so the writer must run concurrently
+    // with `read_output()` or `write_all()` will block once the buffer fills up.
+    tokio::spawn(async move {
+        writer.write_all(&bytes).await.expect("write");
+    });
+
+    let out = read_output(reader, None, false, None).await.expect("read");
+    assert_eq!(out.text.len(), expected_len);
+}
+
+#[test]
+fn aggregate_output_keeps_all_bytes_when_uncapped() {
+    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, None);
+
+    assert_eq!(aggregated.text.len(), EXEC_OUTPUT_MAX_BYTES * 2);
+    assert_eq!(
+        aggregated.text[..EXEC_OUTPUT_MAX_BYTES],
+        vec![b'a'; EXEC_OUTPUT_MAX_BYTES]
+    );
+    assert_eq!(
+        aggregated.text[EXEC_OUTPUT_MAX_BYTES..],
+        vec![b'b'; EXEC_OUTPUT_MAX_BYTES]
+    );
+}
+
+#[test]
+fn full_buffer_capture_policy_disables_caps_and_exec_expiration() {
+    assert_eq!(ExecCapturePolicy::FullBuffer.retained_bytes_cap(), None);
+    assert_eq!(
+        ExecCapturePolicy::FullBuffer.io_drain_timeout(),
+        Duration::from_millis(IO_DRAIN_TIMEOUT_MS)
+    );
+    assert!(!ExecCapturePolicy::FullBuffer.uses_expiration());
+}
+
+#[tokio::test]
+async fn exec_full_buffer_capture_ignores_expiration() -> Result<()> {
+    #[cfg(windows)]
+    let command = vec![
+        "powershell.exe".to_string(),
+        "-NonInteractive".to_string(),
+        "-NoLogo".to_string(),
+        "-Command".to_string(),
+        "Start-Sleep -Milliseconds 50; [Console]::Out.Write('hello')".to_string(),
+    ];
+    #[cfg(not(windows))]
+    let command = vec![
+        "/bin/sh".to_string(),
+        "-c".to_string(),
+        "sleep 0.05; printf hello".to_string(),
+    ];
+
+    let env: HashMap<String, String> = std::env::vars().collect();
+    let output = exec(
+        ExecParams {
+            command,
+            cwd: std::env::current_dir()?,
+            expiration: 1.into(),
+            capture_policy: ExecCapturePolicy::FullBuffer,
+            env,
+            network: None,
+            sandbox_permissions: SandboxPermissions::UseDefault,
+            windows_sandbox_level: WindowsSandboxLevel::Disabled,
+            windows_sandbox_private_desktop: false,
+            justification: None,
+            arg0: None,
+        },
+        SandboxType::None,
+        &SandboxPolicy::DangerFullAccess,
+        &FileSystemSandboxPolicy::unrestricted(),
+        NetworkSandboxPolicy::Enabled,
+        /*stdout_stream*/ None,
+        /*after_spawn*/ None,
+    )
+    .await?;
+
+    assert_eq!(output.stdout.from_utf8_lossy().text.trim(), "hello");
+    assert!(!output.timed_out);
+
+    Ok(())
+}
+
+#[cfg(unix)]
+#[tokio::test]
+async fn exec_full_buffer_capture_keeps_io_drain_timeout_when_descendant_holds_pipe_open()
+-> Result<()> {
+    let output = tokio::time::timeout(
+        Duration::from_millis(IO_DRAIN_TIMEOUT_MS * 3),
+        exec(
+            ExecParams {
+                command: vec![
+                    "/bin/sh".to_string(),
+                    "-c".to_string(),
+                    "printf hello; sleep 30 &".to_string(),
+                ],
+                cwd: std::env::current_dir()?,
+                expiration: 1.into(),
+                capture_policy: ExecCapturePolicy::FullBuffer,
+                env: std::env::vars().collect(),
+                network: None,
+                sandbox_permissions: SandboxPermissions::UseDefault,
+                windows_sandbox_level: WindowsSandboxLevel::Disabled,
+                windows_sandbox_private_desktop: false,
+                justification: None,
+                arg0: None,
+            },
+            SandboxType::None,
+            &SandboxPolicy::DangerFullAccess,
+            &FileSystemSandboxPolicy::unrestricted(),
+            NetworkSandboxPolicy::Enabled,
+            /*stdout_stream*/ None,
+            /*after_spawn*/ None,
+        ),
+    )
+    .await
+    .expect("full-buffer exec should return once the I/O drain guard fires")?;
+
+    assert!(!output.timed_out);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn process_exec_tool_call_preserves_full_buffer_capture_policy() -> Result<()> {
+    let byte_count = EXEC_OUTPUT_MAX_BYTES.saturating_add(128 * 1024);
+    #[cfg(windows)]
+    let command = vec![
+        "powershell.exe".to_string(),
+        "-NonInteractive".to_string(),
+        "-NoLogo".to_string(),
+        "-Command".to_string(),
+        format!("Start-Sleep -Milliseconds 50; [Console]::Out.Write('a' * {byte_count})"),
+    ];
+    #[cfg(not(windows))]
+    let command = vec![
+        "/bin/sh".to_string(),
+        "-c".to_string(),
+        format!("sleep 0.05; head -c {byte_count} /dev/zero | tr '\\0' 'a'"),
+    ];
+
+    let cwd = std::env::current_dir()?;
+    let sandbox_policy = SandboxPolicy::DangerFullAccess;
+    let output = process_exec_tool_call(
+        ExecParams {
+            command,
+            cwd: cwd.clone(),
+            expiration: 1.into(),
+            capture_policy: ExecCapturePolicy::FullBuffer,
+            env: std::env::vars().collect(),
+            network: None,
+            sandbox_permissions: SandboxPermissions::UseDefault,
+            windows_sandbox_level: WindowsSandboxLevel::Disabled,
+            windows_sandbox_private_desktop: false,
+            justification: None,
+            arg0: None,
+        },
+        &sandbox_policy,
+        &FileSystemSandboxPolicy::from(&sandbox_policy),
+        NetworkSandboxPolicy::Enabled,
+        cwd.as_path(),
+        &None,
+        false,
+        None,
+    )
+    .await?;
+
+    assert!(!output.timed_out);
+    assert_eq!(output.stdout.text.len(), byte_count);
+
+    Ok(())
+}
+
 #[test]
 fn windows_restricted_token_skips_external_sandbox_policies() {
     let policy = SandboxPolicy::ExternalSandbox {
@@ -396,6 +585,7 @@ async fn kill_child_process_group_kills_grandchildren_on_timeout() -> Result<()>
         command,
         cwd: std::env::current_dir()?,
         expiration: 500.into(),
+        capture_policy: ExecCapturePolicy::ShellTool,
         env,
         network: None,
         sandbox_permissions: SandboxPermissions::UseDefault,
@@ -453,6 +643,7 @@ async fn process_exec_tool_call_respects_cancellation_token() -> Result<()> {
         command,
         cwd: cwd.clone(),
         expiration: ExecExpiration::Cancellation(cancel_token),
+        capture_policy: ExecCapturePolicy::ShellTool,
         env,
         network: None,
         sandbox_permissions: SandboxPermissions::UseDefault,