From 6a94df6a311d9f325e0e10cfa634d8de0443b828 Mon Sep 17 00:00:00 2001 From: Michael Choi Date: Fri, 23 Oct 2026 14:48:07 +0000 Subject: [PATCH] fix(collect): read files over exec_with_output when a sandbox has no streaming exec Sandboxes that only implement exec_with_output (smol_vm, and the ARP provider in agentenvhub-sdk <= 0.34.18) failed in _collect_file with AttributeError. FDE-2526. Co-Authored-By: Claude Opus 5.5 --- .../task_step/task_steps/collect_artifacts.py | 58 +++++++++++-------- tst/unit/task_step/test_collect_integrity.py | 29 ++++++++++ 2 files changed, 62 insertions(+), 25 deletions(-) diff --git a/src/agent_env/task_step/task_steps/collect_artifacts.py b/src/agent_env/task_step/task_steps/collect_artifacts.py index 153f4e73..78d3570d 100644 --- a/src/agent_env/task_step/task_steps/collect_artifacts.py +++ b/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})") diff --git a/tst/unit/task_step/test_collect_integrity.py b/tst/unit/task_step/test_collect_integrity.py index e4266410..bf10c238 100644 --- a/tst/unit/task_step/test_collect_integrity.py +++ b/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