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 5cd47544..a13ddfab 100644 --- a/src/agent_env/task_step/task_steps/collect_artifacts.py +++ b/src/agent_env/task_step/task_steps/collect_artifacts.py @@ -67,6 +67,7 @@ from __future__ import annotations import asyncio +import base64 import functools import hashlib import logging @@ -101,6 +102,11 @@ _CUA_SINGLE_SHOT_MAX = 16 * 1024 * 1024 # 16 MB _CUA_CHUNK_BYTES = 8 * 1024 * 1024 # 8 MB per ranged read +# A sandbox exec stream can buffer a bounded amount of output (about 16 MiB on one remote VM provider) +# and drop the oldest events when the reader lags, with exit 0. A file larger than this is read in ranges of this +# many raw bytes (~5.4 MB of base64) so no single exec nears that bound. +_COLLECT_CHUNK_BYTES = 4 * 1024 * 1024 + # artifact_paths entries with a URL scheme are not VM files — they're gold/source # URLs that downstream consumers (evaluator, artifact previewer) fetch directly. # Collect skips them rather than joining them to base_path (which would mangle @@ -458,6 +464,41 @@ async def _get_file_size(self, sandbox, container: Optional[str], source_path: s raise RuntimeError(f"could not size {source_path} (stat exit {exit_code}): {stderr[:200]}") return int(stdout.strip()) + @staticmethod + async def _exec_base64_into(sandbox, args: tuple, fh, source_path: str) -> int: + """Run a command that prints base64, decode its stdout into ``fh`` and return the bytes written. + + stdout is decoded incrementally (bounds memory) and stderr is drained in parallel so the shared + output pump can't backpressure-stall. Raw binary over a websocket exec transport is silently + dropped, which is why the sandbox encodes it as ASCII.""" + process = await sandbox.exec(*args) + stderr_task = asyncio.create_task(process.stderr.read()) + stdout = process.stdout + written = 0 + 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: + written += fh.write(base64.b64decode(pending[:n])) + del pending[:n] + if pending: + written += 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") + written += 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})") + return written + async def _collect_file( self, sandbox, container: Optional[str], source_path: str, object_name: str, store, artifact_id: str, version: int, content_type: str, @@ -465,12 +506,11 @@ async def _collect_file( """Read a file out of the sandbox and upload it to S3 — one path, any size. 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.""" - import base64 - bash_cmd = f"base64 < {shlex.quote(source_path)}" - args = _exec_args(sandbox, container, ("bash", "-c", bash_cmd)) + silently dropped; ASCII survives), decode into a temp file (bounds memory), then + multipart-upload. A file over ``_COLLECT_CHUNK_BYTES`` is read in ranges, one exec each, + so no exec stream outgrows the transport's output buffer; smaller files take one exec.""" + quoted = shlex.quote(source_path) + src_size = await self._get_file_size(sandbox, container, source_path) local_path = None handed_off = False @@ -478,42 +518,33 @@ 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)) + if src_size <= _COLLECT_CHUNK_BYTES: + args = _exec_args(sandbox, container, ("bash", "-c", f"base64 < {quoted}")) + await self._exec_base64_into(sandbox, args, fh, source_path) 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})") + written = 0 + while written < src_size: + want = min(_COLLECT_CHUNK_BYTES, src_size - written) + script = f"tail -c +{written + 1} {quoted} | head -c {want} | base64" + args = _exec_args(sandbox, container, ("bash", "-c", script)) + got = await self._exec_base64_into(sandbox, args, fh, source_path) + if got != want: + raise RuntimeError( + f"short read for {source_path}: got {got} of {want} bytes at offset {written} " + f"(exec-channel truncation) — refusing to upload a corrupt artifact" + ) + written += got # Integrity check: base64-over-exec can silently truncate large streams on a # websocket exec transport under load. Fail loudly on a size mismatch rather than upload a # truncated artifact (retryable failure > undetectable short file). - src_size = await self._get_file_size(sandbox, container, source_path) got_size = os.path.getsize(local_path) - if src_size >= 0 and got_size != src_size: + final_size = await self._get_file_size(sandbox, container, source_path) + if final_size >= 0 and got_size != final_size: raise RuntimeError( f"collect integrity check failed for {source_path}: decoded {got_size} bytes but " - f"source is {src_size} (exec-channel truncation) — refusing to upload a corrupt artifact" + f"source is {final_size} (exec-channel truncation) — refusing to upload a corrupt artifact" ) path = local_path @@ -582,8 +613,6 @@ async def _controller_get_file(self, deployed_env, path: str) -> str: lenient skip). Small files come back in one ``cua_get_file`` call; large ones stream in byte ranges straight to disk, so neither the transfer nor this process ever holds the whole file in memory.""" - import base64 - size = await self._controller_file_size(deployed_env, path) if size < 0: raise FileNotFoundError(f"file not found on CUA VM: {path}") @@ -608,8 +637,6 @@ async def _controller_stream_to_file(self, deployed_env, path: str, size: int, l than scans. Each range is decoded and written straight to the file so memory stays bounded to one chunk. base64 may wrap, so whitespace is stripped before decoding. A short read fails loudly rather than leaving a truncated artifact.""" - import base64 - logger.info( f"Streaming {path} ({size} bytes) off the CUA VM in {_CUA_CHUNK_BYTES}-byte chunks" ) diff --git a/tst/unit/task_step/test_collect_chunked.py b/tst/unit/task_step/test_collect_chunked.py new file mode 100644 index 00000000..4e7d8f4e --- /dev/null +++ b/tst/unit/task_step/test_collect_chunked.py @@ -0,0 +1,145 @@ +"""`_collect_file` reads a large file in bounded chunks so no single exec streams more than a +transport's output buffer holds (a sandbox exec channel can drop the oldest output silently).""" +import base64 +import hashlib +import os +import re + +import pytest + +from agent_env.task_step.task_steps import collect_artifacts +from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep + +CHUNK = 1000 +RANGED = re.compile(r"tail -c \+(\d+) \S+ \| head -c (\d+) \| base64") + + +class _Stdout: + def __init__(self, data: bytes): + self._data = data + + def __aiter__(self): + self._it = iter([self._data[i:i + 4096] for i in range(0, len(self._data), 4096)]) + return self + + async def __anext__(self): + try: + return next(self._it) + except StopIteration: + raise StopAsyncIteration + + +class _Stderr: + async def read(self): + return b"" + + +class _Proc: + def __init__(self, data: bytes): + self.stdout = _Stdout(data) + self.stderr = _Stderr() + + async def wait(self): + return 0 + + +class _Sandbox: + """Serves `base64 < file` whole and `tail | head | base64` ranges, recording each command.""" + mode = "container" + + def __init__(self, content: bytes, serve_short: int = 0, append_after_first_read: bytes = b""): + self.content = content + self.serve_short = serve_short + self.append_after_first_read = append_after_first_read + self.commands: list[str] = [] + + async def exec(self, *args): + cmd = args[-1] + self.commands.append(cmd) + match = RANGED.search(cmd) + if match: + start, length = int(match.group(1)), int(match.group(2)) + data = self.content[start - 1:start - 1 + length] + if self.serve_short: + data = data[:-self.serve_short] + else: + data = self.content + if self.append_after_first_read: + self.content += self.append_after_first_read + self.append_after_first_read = b"" + return _Proc(base64.encodebytes(data)) + + async def exec_with_output(self, *args): + return (0, f"{len(self.content)}\n", "") + + +class _Store: + def __init__(self): + self.uploaded = None + + def put_object_file(self, **kw): + with open(kw["file_path"], "rb") as fh: + self.uploaded = fh.read() + return "s3://bucket/ok" + + +@pytest.fixture(autouse=True) +def _small_chunks(monkeypatch): + monkeypatch.setattr(collect_artifacts, "_COLLECT_CHUNK_BYTES", CHUNK) + + +async def _collect(sandbox): + store = _Store() + step = CollectArtifactsTaskStep(id="collect", version=None) + await step._collect_file(sandbox, None, "/app/artifact/a.bin", "a.bin", store, "uid", 1, + "application/octet-stream") + return store.uploaded + + +def _ranged(sandbox): + return [c for c in sandbox.commands if RANGED.search(c)] + + +@pytest.mark.asyncio +async def test_multi_chunk_file_is_reassembled_exactly(): + content = os.urandom(CHUNK * 3 + 123) + sandbox = _Sandbox(content) + got = await _collect(sandbox) + assert hashlib.sha256(got).digest() == hashlib.sha256(content).digest() + assert len(_ranged(sandbox)) == 4 + + +@pytest.mark.asyncio +async def test_exact_multiple_of_the_chunk_size_reads_no_extra_chunk(): + content = os.urandom(CHUNK * 2) + sandbox = _Sandbox(content) + assert await _collect(sandbox) == content + assert len(_ranged(sandbox)) == 2 + + +@pytest.mark.asyncio +async def test_small_file_takes_a_single_unranged_exec(): + content = os.urandom(CHUNK) + sandbox = _Sandbox(content) + assert await _collect(sandbox) == content + assert len(sandbox.commands) == 1 and not _ranged(sandbox) + + +@pytest.mark.asyncio +async def test_empty_file_collects_as_zero_bytes(): + sandbox = _Sandbox(b"") + assert await _collect(sandbox) == b"" + + +@pytest.mark.asyncio +async def test_short_chunk_read_raises_and_uploads_nothing(): + sandbox = _Sandbox(os.urandom(CHUNK * 3), serve_short=1) + with pytest.raises(RuntimeError, match="short read"): + await _collect(sandbox) + + +@pytest.mark.asyncio +async def test_a_file_that_grows_during_collection_fails_the_integrity_check(): + sandbox = _Sandbox(os.urandom(CHUNK * 3), append_after_first_read=b"late bytes") + with pytest.raises(RuntimeError, match="integrity check failed"): + await _collect(sandbox)