fix(collect): read files over exec_with_output when a sandbox has no streaming exec

src/agent_env/task_step/task_steps/collect_artifacts.py

@@ -467,7 +467,8 @@ async def _collect_file(
         base64-encode in the sandbox (raw binary over a websocket exec transport is
         silently dropped; ASCII survives), stream stdout + decode incrementally to
         a temp file (bounds memory), then multipart-upload. stderr is drained in
-        parallel so the shared output pump can't backpressure-stall."""
+        parallel so the shared output pump can't backpressure-stall. A sandbox with
+        only buffered ``exec_with_output`` reads the encoding in one piece instead."""
         import base64
         bash_cmd = f"base64 < {shlex.quote(source_path)}"
         args = _exec_args(sandbox, container, ("bash", "-c", bash_cmd))
@@ -478,30 +479,37 @@ async def _collect_file(
             with tempfile.NamedTemporaryFile(delete=False, suffix=f"-{object_name.replace('/', '_')}") as tmp:
                 local_path = tmp.name
 
-            process = await sandbox.exec(*args)
-            stderr_task = asyncio.create_task(process.stderr.read())
-            stdout = process.stdout
-            with open(local_path, "wb") as fh:
-                if hasattr(stdout, "__aiter__"):
-                    # Async-iterable stdout (websocket exec): stream, decoding 4-char-aligned base64.
-                    pending = bytearray()
-                    async for chunk in stdout:
-                        pending += chunk.translate(None, b"\r\n")  # strip 76-col wrap
-                        n = (len(pending) // 4) * 4
-                        if n:
-                            fh.write(base64.b64decode(pending[:n]))
-                            del pending[:n]
-                    if pending:
-                        fh.write(base64.b64decode(pending))
-                else:
-                    # Modal stdout has read() but no __aiter__: read all, decode once.
-                    data = (await stdout.read()).translate(None, b"\r\n")
-                    fh.write(base64.b64decode(data))
-            exit_code = await process.wait()
-            try:
-                await stderr_task
-            except Exception:
-                pass
+            if not hasattr(sandbox, "exec"):
+                # A sandbox with only buffered exec returns the whole encoding at once.
+                exit_code, encoded, _ = await sandbox.exec_with_output(*args)
+                if exit_code == 0:
+                    with open(local_path, "wb") as fh:
+                        fh.write(base64.b64decode(encoded.translate({ord("\r"): None, ord("\n"): None})))
+            else:
+                process = await sandbox.exec(*args)
+                stderr_task = asyncio.create_task(process.stderr.read())
+                stdout = process.stdout
+                with open(local_path, "wb") as fh:
+                    if hasattr(stdout, "__aiter__"):
+                        # Async-iterable stdout (websocket exec): stream, decoding 4-char-aligned base64.
+                        pending = bytearray()
+                        async for chunk in stdout:
+                            pending += chunk.translate(None, b"\r\n")  # strip 76-col wrap
+                            n = (len(pending) // 4) * 4
+                            if n:
+                                fh.write(base64.b64decode(pending[:n]))
+                                del pending[:n]
+                        if pending:
+                            fh.write(base64.b64decode(pending))
+                    else:
+                        # Modal stdout has read() but no __aiter__: read all, decode once.
+                        data = (await stdout.read()).translate(None, b"\r\n")
+                        fh.write(base64.b64decode(data))
+                exit_code = await process.wait()
+                try:
+                    await stderr_task
+                except Exception:
+                    pass
             if exit_code != 0:
                 raise RuntimeError(f"base64 read failed for {source_path} (exit {exit_code})")
 

tst/unit/task_step/test_collect_integrity.py

@@ -125,3 +125,32 @@ def put_object_file(self, **kw):
     assert uploaded == [src]
     assert "Uploading collected /app/artifact/a.bin finished after its caller was cancelled" in caplog.text
     assert list(tmp_path.iterdir()) == []
+
+
+class _BufferedOnlySandbox:
+    """A sandbox with exec_with_output and no streaming exec."""
+    mode = "container"
+    def __init__(self, src: bytes):
+        self._src = src
+    async def exec_with_output(self, *args):
+        if "base64" in args[-1]:
+            return (0, base64.encodebytes(self._src).decode(), "")  # 76-col wrapped, like base64(1)
+        return (0, f"{len(self._src)}\n", "")
+
+
+@pytest.mark.asyncio
+async def test_a_sandbox_without_streaming_exec_collects_over_buffered_exec():
+    src = bytes(range(256)) * 40
+    uploaded = {}
+
+    class _Store:
+        def put_object_file(self, **kw):
+            with open(kw["file_path"], "rb") as fh:
+                uploaded["bytes"] = fh.read()
+            return "s3://bucket/ok"
+
+    step = CollectArtifactsTaskStep(id="collect", version=None)
+    out = await step._collect_file(_BufferedOnlySandbox(src), None, "/app/artifact/out.bin", "out.bin",
+                                   _Store(), "uid", 1, "application/octet-stream")
+    assert out == "s3://bucket/ok"
+    assert uploaded["bytes"] == src