fix(sandbox): Keep artifact checksums in sync

Buffer local artifact contents while hashing so recorded checksums match the exact bytes written into the sandbox. Cover both LocalFile and LocalDir with focused mutation tests.

Author: matthewflint

src/agents/sandbox/entries/artifacts.py

@@ -6,24 +6,23 @@
 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
+from typing import TYPE_CHECKING, Literal, cast
 
 from pydantic import Field, field_serializer, field_validator
 
 from ..errors import (
     GitCloneError,
     GitCopyError,
     GitMissingInImageError,
-    LocalChecksumError,
     LocalDirReadError,
     LocalFileReadError,
 )
 from ..materialization import MaterializedFile, gather_in_order
 from ..types import ExecResult, User
-from ..util.checksums import sha256_file
 from .base import BaseEntry
 
 if TYPE_CHECKING:
@@ -32,16 +31,27 @@
 _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 _sha256_handle(handle: io.BufferedReader) -> str:
+def _buffer_handle_with_checksum(handle: io.BufferedReader) -> tuple[io.IOBase, str]:
     digest = hashlib.sha256()
-    while True:
-        chunk = handle.read(1024 * 1024)
-        if not chunk:
-            break
-        digest.update(chunk)
-    return digest.hexdigest()
+    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 Dir(BaseEntry):
@@ -111,13 +121,16 @@ async def apply(
     ) -> list[MaterializedFile]:
         src = (base_dir / self.src).resolve()
         try:
-            checksum = sha256_file(src)
+            with src.open("rb") as f:
+                buffered, checksum = _buffer_handle_with_checksum(f)
         except OSError as e:
-            raise LocalChecksumError(src=src, cause=e) from e
+            raise LocalFileReadError(src=src, cause=e) from e
         await session.mkdir(Path(dest).parent, parents=True)
         try:
-            with src.open("rb") as f:
-                await session.write(dest, f)
+            try:
+                await session.write(dest, buffered)
+            finally:
+                buffered.close()
         except OSError as e:
             raise LocalFileReadError(src=src, cause=e) from e
         await self._apply_metadata(session, dest)
@@ -349,10 +362,12 @@ async def _copy_local_dir_file(
             )
             with os.fdopen(fd, "rb") as f:
                 fd = None
-                checksum = _sha256_handle(f)
-                f.seek(0)
+                buffered, checksum = _buffer_handle_with_checksum(f)
                 await session.mkdir(child_dest.parent, parents=True, user=user)
-                await session.write(child_dest, f, user=user)
+                try:
+                    await session.write(child_dest, buffered, user=user)
+                finally:
+                    buffered.close()
         except OSError as e:
             raise LocalFileReadError(src=src, cause=e) from e
         finally:

tests/sandbox/test_entries.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import hashlib
 import io
 import os
 from collections.abc import Awaitable, Callable, Sequence
@@ -98,6 +99,19 @@ async def _exec_internal(
         return ExecResult(stdout=b"", stderr=b"", exit_code=0)
 
 
+class _MutatingWriteSession(_RecordingSession):
+    def __init__(self, mutate_before_read: Callable[[], None]) -> None:
+        super().__init__()
+        self._mutate_before_read = mutate_before_read
+        self._mutated = False
+
+    async def write(self, path: Path, data: io.IOBase, *, user: object = None) -> None:
+        if not self._mutated:
+            self._mutate_before_read()
+            self._mutated = True
+        await super().write(path, data, user=user)
+
+
 @pytest.mark.asyncio
 async def test_base_sandbox_session_uses_current_working_directory_for_local_file_sources(
     monkeypatch: pytest.MonkeyPatch,
@@ -118,6 +132,25 @@ async def test_base_sandbox_session_uses_current_working_directory_for_local_fil
     assert session.writes[Path("/workspace/copied.txt")] == b"hello"
 
 
+@pytest.mark.asyncio
+async def test_local_file_checksum_matches_written_bytes_when_source_changes(
+    tmp_path: Path,
+) -> None:
+    source = tmp_path / "source.txt"
+    source.write_bytes(b"original")
+    session = _MutatingWriteSession(lambda: source.write_bytes(b"mutated"))
+
+    result = await LocalFile(src=Path("source.txt")).apply(
+        session,
+        Path("/workspace/copied.txt"),
+        tmp_path,
+    )
+
+    written = session.writes[Path("/workspace/copied.txt")]
+    assert written == b"original"
+    assert result[0].sha256 == hashlib.sha256(written).hexdigest()
+
+
 @pytest.mark.asyncio
 async def test_local_dir_copy_falls_back_when_safe_dir_fd_open_unavailable(
     monkeypatch: pytest.MonkeyPatch,
@@ -145,6 +178,30 @@ async def test_local_dir_copy_falls_back_when_safe_dir_fd_open_unavailable(
     assert session.writes[Path("/workspace/copied/safe.txt")] == b"safe"
 
 
+@pytest.mark.asyncio
+async def test_local_dir_checksum_matches_written_bytes_when_source_changes(
+    tmp_path: Path,
+) -> None:
+    src_root = tmp_path / "src"
+    src_root.mkdir()
+    src_file = src_root / "safe.txt"
+    src_file.write_bytes(b"original")
+    session = _MutatingWriteSession(lambda: src_file.write_bytes(b"mutated"))
+    local_dir = LocalDir(src=Path("src"))
+
+    result = await local_dir._copy_local_dir_file(
+        base_dir=tmp_path,
+        session=session,
+        src_root=src_root,
+        src=src_file,
+        dest_root=Path("/workspace/copied"),
+    )
+
+    written = session.writes[Path("/workspace/copied/safe.txt")]
+    assert written == b"original"
+    assert result.sha256 == hashlib.sha256(written).hexdigest()
+
+
 @pytest.mark.asyncio
 async def test_local_dir_copy_revalidates_swapped_paths_during_open(
     monkeypatch: pytest.MonkeyPatch,