fix(sandbox): stream checksums after safe local file opens

Author: matthewflint

src/agents/sandbox/entries/artifacts.py

@@ -17,9 +17,10 @@
     GitCloneError,
     GitCopyError,
     GitMissingInImageError,
-    LocalChecksumError,
+    LocalArtifactError,
     LocalDirReadError,
     LocalFileReadError,
+    WorkspaceArchiveWriteError,
 )
 from ..materialization import MaterializedFile, gather_in_order
 from ..types import ExecResult, User
@@ -33,14 +34,91 @@
 _HAS_O_DIRECTORY = hasattr(os, "O_DIRECTORY")
 
 
-def _sha256_handle(handle: io.BufferedReader) -> str:
-    digest = hashlib.sha256()
-    while True:
-        chunk = handle.read(1024 * 1024)
+class _HashingReader(io.IOBase):
+    def __init__(
+        self,
+        stream: io.BufferedReader,
+        *,
+        read_error_factory: Callable[[OSError], BaseException] | None = None,
+    ) -> None:
+        self._stream = stream
+        self._digest = hashlib.sha256()
+        self._started = False
+        self._finished = False
+        self._read_error_factory = read_error_factory
+
+    def readable(self) -> bool:
+        return True
+
+    def read(self, size: int = -1) -> bytes:
+        try:
+            chunk = self._stream.read(size)
+        except OSError as e:
+            if self._read_error_factory is not None:
+                raise self._read_error_factory(e) from e
+            raise
+        if chunk is None:
+            self._finished = True
+            return b""
+        if isinstance(chunk, bytearray):
+            chunk = bytes(chunk)
+        self._started = True
         if not chunk:
-            break
-        digest.update(chunk)
-    return digest.hexdigest()
+            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")
+        try:
+            return int(self._stream.seek(offset, whence))
+        except OSError as e:
+            if self._read_error_factory is not None:
+                raise self._read_error_factory(e) from e
+            raise
+
+    def tell(self) -> int:
+        try:
+            return int(self._stream.tell())
+        except OSError as e:
+            if self._read_error_factory is not None:
+                raise self._read_error_factory(e) from e
+            raise
+
+    def hexdigest(self) -> str:
+        if not self._finished:
+            raise RuntimeError("checksum is not available until the stream is fully consumed")
+        return self._digest.hexdigest()
+
+
+def _find_nested_local_artifact_error(exc: BaseException) -> LocalArtifactError | None:
+    seen: set[int] = set()
+    current: BaseException | None = exc
+    while current is not None and id(current) not in seen:
+        if isinstance(current, LocalArtifactError):
+            return current
+        seen.add(id(current))
+        next_exc = getattr(current, "cause", None)
+        if not isinstance(next_exc, BaseException):
+            next_exc = current.__cause__
+        current = next_exc if isinstance(next_exc, BaseException) else None
+    return None
+
+
+def _reraise_nested_local_artifact_error(exc: BaseException) -> None:
+    nested_local_artifact_error = _find_nested_local_artifact_error(exc)
+    if nested_local_artifact_error is not None:
+        raise nested_local_artifact_error
 
 
 class Dir(BaseEntry):
@@ -122,13 +200,15 @@ async def apply(
             )
             with os.fdopen(fd, "rb") as f:
                 fd = None
-                try:
-                    checksum = _sha256_handle(f)
-                    f.seek(0)
-                except OSError as e:
-                    raise LocalChecksumError(src=src, cause=e) from e
+                hashing_reader = _HashingReader(
+                    f,
+                    read_error_factory=lambda e: LocalFileReadError(src=src, cause=e),
+                )
                 await session.mkdir(Path(dest).parent, parents=True)
-                await session.write(dest, f)
+                await session.write(dest, hashing_reader)
+        except WorkspaceArchiveWriteError as e:
+            _reraise_nested_local_artifact_error(e)
+            raise
         except LocalDirReadError as e:
             context = dict(e.context)
             context.pop("src", None)
@@ -139,7 +219,7 @@ async def apply(
             if fd is not None:
                 os.close(fd)
         await self._apply_metadata(session, dest)
-        return [MaterializedFile(path=dest, sha256=checksum)]
+        return [MaterializedFile(path=dest, sha256=hashing_reader.hexdigest())]
 
 
 class LocalDir(BaseEntry):
@@ -367,16 +447,21 @@ async def _copy_local_dir_file(
             )
             with os.fdopen(fd, "rb") as f:
                 fd = None
-                checksum = _sha256_handle(f)
-                f.seek(0)
+                hashing_reader = _HashingReader(
+                    f,
+                    read_error_factory=lambda e: LocalFileReadError(src=src, cause=e),
+                )
                 await session.mkdir(child_dest.parent, parents=True, user=user)
-                await session.write(child_dest, f, user=user)
+                await session.write(child_dest, hashing_reader, user=user)
+        except WorkspaceArchiveWriteError as e:
+            _reraise_nested_local_artifact_error(e)
+            raise
         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

@@ -16,10 +16,12 @@
     InvalidManifestPathError,
     LocalDirReadError,
     LocalFileReadError,
+    WorkspaceArchiveWriteError,
 )
 from agents.sandbox.manifest import Manifest
 from agents.sandbox.materialization import MaterializedFile
 from agents.sandbox.session.base_sandbox_session import BaseSandboxSession
+from agents.sandbox.session.workspace_payloads import coerce_write_payload
 from agents.sandbox.snapshot import NoopSnapshot
 from agents.sandbox.types import ExecResult, User
 from tests.utils.factories import TestSessionState
@@ -154,6 +156,77 @@ def test_resolve_workspace_path_rejects_absolute_symlink_escape_for_host_root(
     assert exc_info.value.context == {"rel": escaped.as_posix(), "reason": "absolute"}
 
 
+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)
+
+
+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)
+
+
+class _PayloadWrappingWriteSession(_RecordingSession):
+    async def write(self, path: Path, data: io.IOBase, *, user: object = None) -> None:
+        _ = user
+        payload = coerce_write_payload(path=path, data=data)
+        chunks: list[bytes] = []
+        try:
+            while True:
+                chunk = payload.stream.read(4)
+                if not chunk:
+                    break
+                chunks.append(chunk)
+        except Exception as e:
+            raise WorkspaceArchiveWriteError(path=path, cause=e) from e
+        self.writes[path] = b"".join(chunks)
+
+
+class _FailAfterChunkStream(io.BytesIO):
+    def __init__(self, data: bytes, *, owned_fd: int | None = None) -> None:
+        super().__init__(data)
+        self._owned_fd = owned_fd
+        self._read_count = 0
+
+    def read(self, size: int | None = -1) -> bytes:
+        if self._read_count > 0:
+            raise OSError("source read failed")
+        self._read_count += 1
+        return super().read(-1 if size is None else size)
+
+    def close(self) -> None:
+        try:
+            super().close()
+        finally:
+            if self._owned_fd is not None:
+                os.close(self._owned_fd)
+                self._owned_fd = None
+
+
 def _symlink_or_skip(path: Path, target: Path, *, target_is_directory: bool = False) -> None:
     try:
         path.symlink_to(target, target_is_directory=target_is_directory)
@@ -184,6 +257,28 @@ 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")
+
+    def mutate_source() -> None:
+        source.write_bytes(b"mutated")
+
+    session = _ChunkedMutatingWriteSession(mutate_source)
+
+    result = await LocalFile(src=Path("source.txt")).apply(
+        session,
+        Path("/workspace/copied.txt"),
+        tmp_path,
+    )
+
+    written = session.writes[Path("/workspace/copied.txt")]
+    assert result[0].sha256 == hashlib.sha256(written).hexdigest()
+
+
 @pytest.mark.asyncio
 async def test_local_file_rejects_symlinked_source_ancestors(tmp_path: Path) -> None:
     target_dir = tmp_path / "secret-dir"
@@ -271,6 +366,93 @@ 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")
+
+    def mutate_source() -> None:
+        src_file.write_bytes(b"mutated")
+
+    session = _ChunkedMutatingWriteSession(mutate_source)
+    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 result.sha256 == hashlib.sha256(written).hexdigest()
+
+
+@pytest.mark.asyncio
+async def test_local_file_preserves_local_read_error_when_write_wraps_stream_failures(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    source = (tmp_path / "source.txt").resolve()
+    source.write_bytes(b"original")
+    session = _PayloadWrappingWriteSession()
+
+    def failing_fdopen(
+        fd: int,
+        *args: object,
+        **kwargs: object,
+    ) -> io.IOBase:
+        _ = args, kwargs
+        return _FailAfterChunkStream(b"original", owned_fd=fd)
+
+    monkeypatch.setattr("agents.sandbox.entries.artifacts.os.fdopen", failing_fdopen)
+
+    with pytest.raises(LocalFileReadError) as excinfo:
+        await LocalFile(src=Path("source.txt")).apply(
+            session,
+            Path("/workspace/copied.txt"),
+            tmp_path,
+        )
+
+    assert excinfo.value.context["src"] == str(source)
+    assert isinstance(excinfo.value.cause, OSError)
+
+
+@pytest.mark.asyncio
+async def test_local_dir_copy_preserves_local_read_error_when_write_wraps_stream_failures(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    src_root = tmp_path / "src"
+    src_root.mkdir()
+    src_file = (src_root / "safe.txt").resolve()
+    src_file.write_bytes(b"original")
+    session = _PayloadWrappingWriteSession()
+    local_dir = LocalDir(src=Path("src"))
+
+    def failing_fdopen(fd: int, *args: object, **kwargs: object) -> io.IOBase:
+        return _FailAfterChunkStream(b"original", owned_fd=fd)
+
+    monkeypatch.setattr("agents.sandbox.entries.artifacts.os.fdopen", failing_fdopen)
+
+    with pytest.raises(LocalFileReadError) as excinfo:
+        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"),
+        )
+
+    assert excinfo.value.context["src"] == str(src_file)
+    assert isinstance(excinfo.value.cause, OSError)
+
+
 @pytest.mark.asyncio
 async def test_local_dir_copy_revalidates_swapped_paths_during_open(
     monkeypatch: pytest.MonkeyPatch,