fix(sandbox): Stream artifact checksum writes

Avoid spooling full artifact contents into a temporary file before write. Hash the same stream as it is consumed by session.write and update the checksum tests to assert that recorded hashes match the bytes that were actually materialized.

Author: matthewflint

src/agents/sandbox/entries/artifacts.py

@@ -6,11 +6,10 @@
 import os
 import re
 import stat
-import tempfile
 import uuid
 from collections.abc import Awaitable, Callable, Mapping
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal, cast
+from typing import TYPE_CHECKING, Literal
 
 from pydantic import Field, field_serializer, field_validator
 
@@ -31,27 +30,52 @@
 _COMMIT_REF_RE = re.compile(r"[0-9a-fA-F]{7,40}")
 _OPEN_SUPPORTS_DIR_FD = os.open in os.supports_dir_fd
 _HAS_O_DIRECTORY = hasattr(os, "O_DIRECTORY")
-_LOCAL_COPY_BUFFER_MAX_MEMORY_BYTES = 8 * 1024 * 1024
-
-
-def _buffer_handle_with_checksum(handle: io.BufferedReader) -> tuple[io.IOBase, str]:
-    digest = hashlib.sha256()
-    buffer = tempfile.SpooledTemporaryFile(
-        max_size=_LOCAL_COPY_BUFFER_MAX_MEMORY_BYTES,
-        mode="w+b",
-    )
-    try:
-        while True:
-            chunk = handle.read(1024 * 1024)
-            if not chunk:
-                break
-            digest.update(chunk)
-            buffer.write(chunk)
-        buffer.seek(0)
-        return cast(io.IOBase, buffer), digest.hexdigest()
-    except BaseException:
-        buffer.close()
-        raise
+
+
+class _HashingReader(io.IOBase):
+    def __init__(self, stream: io.BufferedReader) -> None:
+        self._stream = stream
+        self._digest = hashlib.sha256()
+        self._started = False
+        self._finished = False
+
+    def readable(self) -> bool:
+        return True
+
+    def read(self, size: int = -1) -> bytes:
+        chunk = self._stream.read(size)
+        if chunk is None:
+            self._finished = True
+            return b""
+        if isinstance(chunk, bytearray):
+            chunk = bytes(chunk)
+        self._started = True
+        if not chunk:
+            self._finished = True
+            return b""
+        self._digest.update(chunk)
+        if size < 0 or len(chunk) < size:
+            self._finished = True
+        return chunk
+
+    def readinto(self, b: bytearray) -> int:
+        data = self.read(len(b))
+        n = len(data)
+        b[:n] = data
+        return n
+
+    def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
+        if self._started:
+            raise io.UnsupportedOperation("cannot seek after reads begin")
+        return int(self._stream.seek(offset, whence))
+
+    def tell(self) -> int:
+        return int(self._stream.tell())
+
+    def hexdigest(self) -> str:
+        if not self._finished:
+            raise RuntimeError("checksum is not available until the stream is fully consumed")
+        return self._digest.hexdigest()
 
 
 class Dir(BaseEntry):
@@ -120,21 +144,15 @@ async def apply(
         base_dir: Path,
     ) -> list[MaterializedFile]:
         src = (base_dir / self.src).resolve()
-        try:
-            with src.open("rb") as f:
-                buffered, checksum = _buffer_handle_with_checksum(f)
-        except OSError as e:
-            raise LocalFileReadError(src=src, cause=e) from e
         await session.mkdir(Path(dest).parent, parents=True)
         try:
-            try:
-                await session.write(dest, buffered)
-            finally:
-                buffered.close()
+            with src.open("rb") as f:
+                hashing_reader = _HashingReader(f)
+                await session.write(dest, hashing_reader)
         except OSError as e:
             raise LocalFileReadError(src=src, cause=e) from e
         await self._apply_metadata(session, dest)
-        return [MaterializedFile(path=dest, sha256=checksum)]
+        return [MaterializedFile(path=dest, sha256=hashing_reader.hexdigest())]
 
 
 class LocalDir(BaseEntry):
@@ -362,18 +380,15 @@ async def _copy_local_dir_file(
             )
             with os.fdopen(fd, "rb") as f:
                 fd = None
-                buffered, checksum = _buffer_handle_with_checksum(f)
+                hashing_reader = _HashingReader(f)
                 await session.mkdir(child_dest.parent, parents=True, user=user)
-                try:
-                    await session.write(child_dest, buffered, user=user)
-                finally:
-                    buffered.close()
+                await session.write(child_dest, hashing_reader, user=user)
         except OSError as e:
             raise LocalFileReadError(src=src, cause=e) from e
         finally:
             if fd is not None:
                 os.close(fd)
-        return MaterializedFile(path=child_dest, sha256=checksum)
+        return MaterializedFile(path=child_dest, sha256=hashing_reader.hexdigest())
 
     def _open_local_dir_file_for_copy(
         self, *, base_dir: Path, src_root: Path, rel_child: Path

tests/sandbox/test_entries.py

@@ -112,6 +112,27 @@ async def write(self, path: Path, data: io.IOBase, *, user: object = None) -> No
         await super().write(path, data, user=user)
 
 
+class _ChunkedMutatingWriteSession(_RecordingSession):
+    def __init__(self, mutate_after_first_chunk: Callable[[], None]) -> None:
+        super().__init__()
+        self._mutate_after_first_chunk = mutate_after_first_chunk
+        self._mutated = False
+
+    async def write(self, path: Path, data: io.IOBase, *, user: object = None) -> None:
+        _ = user
+        chunks: list[bytes] = []
+        first = data.read(4)
+        if isinstance(first, bytes):
+            chunks.append(first)
+        if not self._mutated:
+            self._mutate_after_first_chunk()
+            self._mutated = True
+        rest = data.read()
+        if isinstance(rest, bytes):
+            chunks.append(rest)
+        self.writes[path] = b"".join(chunks)
+
+
 @pytest.mark.asyncio
 async def test_base_sandbox_session_uses_current_working_directory_for_local_file_sources(
     monkeypatch: pytest.MonkeyPatch,
@@ -138,7 +159,7 @@ async def test_local_file_checksum_matches_written_bytes_when_source_changes(
 ) -> None:
     source = tmp_path / "source.txt"
     source.write_bytes(b"original")
-    session = _MutatingWriteSession(lambda: source.write_bytes(b"mutated"))
+    session = _ChunkedMutatingWriteSession(lambda: source.write_bytes(b"mutated"))
 
     result = await LocalFile(src=Path("source.txt")).apply(
         session,
@@ -147,7 +168,6 @@ async def test_local_file_checksum_matches_written_bytes_when_source_changes(
     )
 
     written = session.writes[Path("/workspace/copied.txt")]
-    assert written == b"original"
     assert result[0].sha256 == hashlib.sha256(written).hexdigest()
 
 
@@ -186,7 +206,7 @@ async def test_local_dir_checksum_matches_written_bytes_when_source_changes(
     src_root.mkdir()
     src_file = src_root / "safe.txt"
     src_file.write_bytes(b"original")
-    session = _MutatingWriteSession(lambda: src_file.write_bytes(b"mutated"))
+    session = _ChunkedMutatingWriteSession(lambda: src_file.write_bytes(b"mutated"))
     local_dir = LocalDir(src=Path("src"))
 
     result = await local_dir._copy_local_dir_file(
@@ -198,7 +218,6 @@ async def test_local_dir_checksum_matches_written_bytes_when_source_changes(
     )
 
     written = session.writes[Path("/workspace/copied/safe.txt")]
-    assert written == b"original"
     assert result.sha256 == hashlib.sha256(written).hexdigest()