diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index 0ba8ad0d5b..546a03dfb6 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -44,6 +44,7 @@ from ....sandbox.manifest import Manifest from ....sandbox.session import SandboxSession, SandboxSessionState from ....sandbox.session.base_sandbox_session import BaseSandboxSession +from ....sandbox.session.bounded_read import collect_bounded from ....sandbox.session.dependencies import Dependencies from ....sandbox.session.manager import Instrumentation from ....sandbox.session.pty_output import collect_pty_output @@ -506,6 +507,25 @@ async def mkdir( cause=e, ) from e + async def _read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + workspace_path = await self._validate_path_access(path) + filesystem = self._sandbox.fs + client = filesystem.get_client() + path_arg = filesystem.format_path(sandbox_path_str(workspace_path)) + async with client.stream( + "GET", + f"{filesystem.url}/filesystem/{path_arg}", + headers={"Accept": "application/octet-stream"}, + ) as response: + if response.status_code == 404: + raise WorkspaceReadNotFoundError(path=path) + if response.status_code != 200: + raise WorkspaceArchiveReadError( + path=path, + retryable=True if response.status_code in TRANSIENT_HTTP_STATUS_CODES else None, + ) + return await collect_bounded(response.aiter_bytes(chunk_size=65536), max_bytes) + async def read(self, path: Path | str, *, user: str | User | None = None) -> io.IOBase: error_path = posix_path_as_path(coerce_posix_path(path)) if user is not None: diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index 2445cd7a8d..94a3d9e093 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -53,6 +53,7 @@ from ....sandbox.manifest import Manifest from ....sandbox.session import SandboxSession, SandboxSessionState from ....sandbox.session.base_sandbox_session import BaseSandboxSession +from ....sandbox.session.bounded_read import collect_bounded from ....sandbox.session.dependencies import Dependencies from ....sandbox.session.manager import Instrumentation from ....sandbox.session.mount_lifecycle import ( @@ -1265,6 +1266,43 @@ async def pty_terminate_all(self) -> None: for entry in entries: await self._terminate_pty_entry(entry) + async def _read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + workspace_path = await self._validate_path_access(path) + url_path = quote(sandbox_path_str(workspace_path).lstrip("/"), safe="/") + async with self._session().get( + self._url(f"file/{url_path}"), timeout=self._request_timeout() + ) as response: + if response.status == 404: + raise WorkspaceReadNotFoundError(path=path) + if response.status != 200: + raise WorkspaceArchiveReadError( + path=path, + retryable=False + if response.status == 403 + else _cloudflare_retryability_for_status(response.status), + ) + # Existing Workers return either bytes or an SSE-encoded file. Bound + # the wire representation too, before the existing decoder allocates. + try: + prefix = await response.content.readexactly(7) + except asyncio.IncompleteReadError as error: + prefix = error.partial + if prefix != b"data: {": + if len(prefix) >= max_bytes: + return prefix[:max_bytes] + return prefix + await collect_bounded( + response.content.iter_chunked(65536), max_bytes - len(prefix) + ) + wire_limit = 8 * max_bytes + 65536 + body = prefix + await collect_bounded( + response.content.iter_chunked(65536), wire_limit + 1 - len(prefix) + ) + if len(body) > wire_limit: + raise WorkspaceArchiveReadError( + path=path, context={"reason": "bounded_read_wire_limit"} + ) + return self._decode_streamed_payload(body)[:max_bytes] + async def read(self, path: Path | str, *, user: str | User | None = None) -> io.IOBase: if user is not None: await self._check_read_with_exec(path, user=user) diff --git a/src/agents/extensions/sandbox/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index e7c9cb8918..e6c22b9af3 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -34,6 +34,7 @@ ExecTransportError, ExposedPortUnavailableError, InvalidManifestPathError as InvalidManifestPathError, + SandboxError, WorkspaceArchiveReadError, WorkspaceArchiveWriteError, WorkspaceReadNotFoundError, @@ -43,6 +44,7 @@ from ....sandbox.manifest import Manifest from ....sandbox.session import SandboxSession, SandboxSessionState from ....sandbox.session.base_sandbox_session import BaseSandboxSession +from ....sandbox.session.bounded_read import collect_bounded from ....sandbox.session.dependencies import Dependencies from ....sandbox.session.manager import Instrumentation from ....sandbox.session.pty_output import collect_pty_output @@ -940,6 +942,31 @@ async def _terminate_pty_entry(self, entry: _DaytonaPtySessionEntry) -> None: except asyncio.TimeoutError: pass + async def _read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + workspace_path = await self._validate_path_access(path) + # The high-level download buffers the response. The generated toolbox + # client exposes the same endpoint without preloading its body. + try: + response = await self._sandbox.fs._api_client.download_file_without_preload_content( + path=sandbox_path_str(workspace_path), + _request_timeout=float(self.state.timeouts.file_download_s), + ) + try: + if response.status == 404: + raise WorkspaceReadNotFoundError(path=path) + if response.status != 200: + raise WorkspaceArchiveReadError( + path=path, retryable=_DAYTONA_HTTP_STATUS_RETRYABLE.get(response.status) + ) + return await collect_bounded(response.content.iter_chunked(65536), max_bytes) + finally: + response.close() + except SandboxError: + raise + except Exception as error: + retryable, _ = _daytona_provider_retryability(error) + raise WorkspaceArchiveReadError(path=path, retryable=retryable) from None + async def read(self, path: Path | str, *, user: str | User | None = None) -> io.IOBase: error_path = posix_path_as_path(coerce_posix_path(path)) if user is not None: diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index 22eb301a62..150a658a3d 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -51,6 +51,7 @@ from ....sandbox.manifest import Manifest from ....sandbox.session import SandboxSession, SandboxSessionState from ....sandbox.session.base_sandbox_session import BaseSandboxSession +from ....sandbox.session.bounded_read import collect_bounded from ....sandbox.session.dependencies import Dependencies from ....sandbox.session.manager import Instrumentation from ....sandbox.session.pty_output import collect_pty_output @@ -1114,6 +1115,20 @@ async def pty_terminate_all(self) -> None: for entry in entries: await self._terminate_pty_entry(entry) + async def _read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + workspace_path = await self._validate_path_access(path) + try: + stream = await _sandbox_read_file( + self._sandbox, sandbox_path_str(workspace_path), format="stream" + ) + async with cast(Any, stream) as chunks: + return await collect_bounded(chunks, max_bytes) + except _e2b_not_found_error_types(): + raise WorkspaceReadNotFoundError(path=path) from None + except Exception as error: + retryable, _ = _e2b_provider_retryability(error) + raise WorkspaceArchiveReadError(path=path, retryable=retryable) from None + async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase: if user is not None: await self._check_read_with_exec(path, user=user) diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index f0fe85e589..b0c8880011 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -1248,6 +1248,39 @@ async def _terminate_pty_entry(self, entry: _ModalPtyProcessEntry) -> None: return_exceptions=True, ) + async def _read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + try: + workspace_path = await self._validate_path_access(path) + await self._ensure_sandbox() + assert self._sandbox is not None + # Each read starts a remote operation. Use Modal's 100 MiB per-read + # ceiling while respecting the caller's remaining byte budget. + stream = await self._sandbox.open.aio(sandbox_path_str(workspace_path), "rb") + completed = False + try: + result = bytearray() + while len(result) < max_bytes: + chunk = await stream.read.aio(min(100 * 1024 * 1024, max_bytes - len(result))) + if not chunk: + break + result.extend(chunk) + payload = bytes(result) + completed = True + return payload + finally: + try: + await asyncio.wait_for(stream.close.aio(), timeout=5.0) + except Exception: + # Preserve an active read failure or cancellation. A close failure + # still fails a read that would otherwise have completed. + if completed: + raise + except (FileNotFoundError, SandboxError): + raise + except Exception as error: + retryable, _ = _modal_provider_retryability(error) + raise WorkspaceArchiveReadError(path=path, retryable=retryable) from None + async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase: if user is not None: await self._check_read_with_exec(path, user=user) diff --git a/src/agents/extensions/sandbox/runloop/sandbox.py b/src/agents/extensions/sandbox/runloop/sandbox.py index 2662a55db4..2d3477218c 100644 --- a/src/agents/extensions/sandbox/runloop/sandbox.py +++ b/src/agents/extensions/sandbox/runloop/sandbox.py @@ -48,6 +48,7 @@ from ....sandbox.manifest import Manifest from ....sandbox.session import SandboxSession, SandboxSessionState from ....sandbox.session.base_sandbox_session import BaseSandboxSession +from ....sandbox.session.bounded_read import collect_bounded from ....sandbox.session.dependencies import Dependencies from ....sandbox.session.manager import Instrumentation from ....sandbox.session.runtime_helpers import RESOLVE_WORKSPACE_PATH_HELPER, RuntimeHelperScript @@ -955,6 +956,24 @@ async def _resolve_exposed_port(self, port: int) -> ExposedPortEndpoint: cause=e, ) from e + async def _read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + normalized_path = await self._validate_path_access(path) + try: + async with self._sdk.api.devboxes.with_streaming_response.download_file( + self.devbox_id, + path=sandbox_path_str(normalized_path), + timeout=self.state.timeouts.file_download_s, + ) as response: + return await collect_bounded(response.iter_bytes(chunk_size=65536), max_bytes) + except Exception as error: + if _is_runloop_not_found(error): + raise WorkspaceReadNotFoundError(path=path) from None + if _is_runloop_provider_error(error): + raise WorkspaceArchiveReadError( + path=path, retryable=_runloop_provider_retryability(error) + ) from None + raise + async def read(self, path: Path | str, *, user: str | User | None = None) -> io.IOBase: """Read a file via Runloop's binary file API.""" error_path = posix_path_as_path(coerce_posix_path(path)) diff --git a/src/agents/extensions/sandbox/vercel/sandbox.py b/src/agents/extensions/sandbox/vercel/sandbox.py index cf2d91afe6..8a91d65ab8 100644 --- a/src/agents/extensions/sandbox/vercel/sandbox.py +++ b/src/agents/extensions/sandbox/vercel/sandbox.py @@ -56,6 +56,7 @@ from ....sandbox.materialization import MaterializationResult from ....sandbox.session import SandboxSession, SandboxSessionState, manifest_ops from ....sandbox.session.base_sandbox_session import BaseSandboxSession +from ....sandbox.session.bounded_read import collect_bounded from ....sandbox.session.dependencies import Dependencies from ....sandbox.session.manager import Instrumentation from ....sandbox.session.mount_lifecycle import ( @@ -1239,6 +1240,35 @@ async def _resolve_exposed_port(self, port: int) -> ExposedPortEndpoint: tls=tls, ) + @redact_mount_error_data + async def _read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + async with self._s3_mount_operation(): + normalized_path = await self._validate_path_access(path) + sandbox = await self._ensure_sandbox() + try: + chunks = await sandbox.iter_file( + sandbox_path_str(normalized_path), chunk_size=65536 + ) + completed = False + try: + payload = await collect_bounded(chunks, max_bytes) + completed = True + return payload + finally: + try: + await chunks.aclose() + except Exception: + # Preserve the primary read failure or cancellation. + # A close failure is primary only after a successful read. + if completed: + raise + except vercel_sandbox.SandboxNotFoundError: + raise WorkspaceReadNotFoundError(path=path) from None + except Exception as error: + raise WorkspaceArchiveReadError( + path=path, retryable=_vercel_provider_retryability(error) + ) from None + @redact_mount_error_data async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase: async with self._s3_mount_operation(): diff --git a/src/agents/sandbox/sandboxes/_unix_local_file_ops.py b/src/agents/sandbox/sandboxes/_unix_local_file_ops.py index 90d228b4ac..348fe7c537 100644 --- a/src/agents/sandbox/sandboxes/_unix_local_file_ops.py +++ b/src/agents/sandbox/sandboxes/_unix_local_file_ops.py @@ -60,6 +60,19 @@ def read(self, path: Path) -> io.IOBase: os.close(fd) raise + def read_bounded(self, path: Path, max_bytes: int) -> bytes: + with self.parent(path) as (parent_fd, name): + fd = os.open(name, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=parent_fd) + try: + if not stat.S_ISREG(os.fstat(fd).st_mode): + raise OSError("Bounded reads require a regular file") + stream = os.fdopen(fd, "rb") + except BaseException: + os.close(fd) + raise + with stream: + return stream.read(max_bytes) + def write(self, path: Path, stream: io.IOBase) -> None: with self.parent(path, for_write=True, create_parents=True) as (parent_fd, name): fd = os.open( diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index fc7c430bd0..68bc8346f6 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -52,6 +52,7 @@ MountConfigError, WorkspaceArchiveReadError, WorkspaceArchiveWriteError, + WorkspaceReadNotFoundError, ) from ..manifest import Manifest from ..session import SandboxSession, SandboxSessionState @@ -836,6 +837,27 @@ async def _prepare_user_pty_pid_path(self, *, path: Path, user: str | None) -> N error_path=path, ) + async def _read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + workspace_path = await self._validate_path_access(path) + # Docker writes already require POSIX sh and head -c. Restrict this + # internal read to image-owned utilities, independent of manifest PATH. + result = await self.exec( + "/bin/sh", + "-c", + "PATH=/usr/bin:/bin; export PATH; " + '[ -e "$1" ] || exit 44; [ -f "$1" ] || exit 45; head -c "$2" < "$1"', + "sh", + sandbox_path_str(workspace_path), + str(max_bytes), + shell=False, + timeout=30.0, + ) + if result.exit_code == 44: + raise WorkspaceReadNotFoundError(path=path) + if not result.ok(): + raise WorkspaceArchiveReadError(path=path) + return result.stdout + async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase: workspace_path = await self._validate_path_access(path) diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 2a808a082f..c0ca8bfaa2 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -1030,6 +1030,9 @@ async def rm( except OSError as e: raise WorkspaceArchiveWriteError(path=normalized, cause=e) from e + async def _read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + return self._files.read_bounded(self.normalize_path(path), max_bytes) + async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase: if user is not None: await self._check_read_with_exec(path, user=user) diff --git a/src/agents/sandbox/session/base_sandbox_session.py b/src/agents/sandbox/session/base_sandbox_session.py index d377bea9ef..ecb12c3291 100644 --- a/src/agents/sandbox/session/base_sandbox_session.py +++ b/src/agents/sandbox/session/base_sandbox_session.py @@ -25,6 +25,7 @@ InvalidManifestPathError, MountConfigError, PtySessionNotFoundError, + SandboxError, WorkspaceArchiveReadError, WorkspaceArchiveWriteError, WorkspaceReadNotFoundError, @@ -929,6 +930,55 @@ async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase :raises: FileNotFoundError: If the path does not exist. """ + async def read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + """Read a prefix of at most ``max_bytes`` bytes. + + Uses the session identity and the same path scope as ``read(path)``. + A result of exactly ``max_bytes`` bytes does not indicate EOF. Backends + provided by the SDK bound acquisition from the source. For compatibility, + custom backends inherit a fallback that consumes and closes the stream + returned by ``read()``. That fallback cannot bound any data acquired inside + ``read()`` before it returns. Override ``_read_bounded`` to provide that + guarantee. Encoded transports may reject responses exceeding their bounded + framing budget. Failures use payload-free ``WorkspaceArchiveReadError`` + diagnostics. + """ + if max_bytes <= 0: + raise ValueError("max_bytes must be positive") + # Raise outside the handler: provider exceptions may retain response bodies. + reason = "bounded_read_failed" + retryable: bool | None = None + missing = False + try: + return await self._read_bounded(path, max_bytes=max_bytes) + except (FileNotFoundError, WorkspaceReadNotFoundError): + missing = True + except SandboxError as error: + retryable = error.retryable + if error.context.get("reason") == "bounded_read_wire_limit": + reason = "bounded_read_wire_limit" + except Exception: + # Discard payload-bearing provider diagnostics at this boundary. + reason = "bounded_read_failed" + if missing: + raise WorkspaceReadNotFoundError(path=path) + raise WorkspaceArchiveReadError(path=path, context={"reason": reason}, retryable=retryable) + + async def _read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + stream = await self.read(path) + try: + result = bytearray() + while len(result) < max_bytes: + remaining = max_bytes - len(result) + payload = stream.read(remaining) + if not payload: + break + chunk = payload.encode("utf-8") if isinstance(payload, str) else bytes(payload) + result.extend(chunk[:remaining]) + return bytes(result) + finally: + stream.close() + @abc.abstractmethod async def write( self, diff --git a/src/agents/sandbox/session/bounded_read.py b/src/agents/sandbox/session/bounded_read.py new file mode 100644 index 0000000000..c45af0db85 --- /dev/null +++ b/src/agents/sandbox/session/bounded_read.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from collections.abc import AsyncIterable + + +async def collect_bounded(chunks: AsyncIterable[bytes], max_bytes: int) -> bytes: + """Collect a prefix; the caller owns and closes the underlying stream.""" + result = bytearray() + async for chunk in chunks: + result.extend(chunk[: max_bytes - len(result)]) + if len(result) == max_bytes: + break + return bytes(result) diff --git a/src/agents/sandbox/session/sandbox_session.py b/src/agents/sandbox/session/sandbox_session.py index 923f025857..bea749e41c 100644 --- a/src/agents/sandbox/session/sandbox_session.py +++ b/src/agents/sandbox/session/sandbox_session.py @@ -667,6 +667,13 @@ async def _read( async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase: return await self._read(path, user=user) + async def read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + return await self._annotate( + op="read", + start_data={**_read_start_data(self, path), "max_bytes": max_bytes}, + run=lambda: self._inner.read_bounded(path, max_bytes=max_bytes), + ) + @instrumented_op("write", data=_write_start_data) async def write( self, diff --git a/src/agents/sandbox/session/sinks.py b/src/agents/sandbox/session/sinks.py index 0781d8407d..ca2a2c994e 100644 --- a/src/agents/sandbox/session/sinks.py +++ b/src/agents/sandbox/session/sinks.py @@ -7,10 +7,10 @@ from collections.abc import Callable from pathlib import Path from types import ModuleType -from typing import Literal, Protocol, runtime_checkable +from typing import Literal, NoReturn, Protocol, runtime_checkable from urllib.request import Request, urlopen -from ..errors import WorkspaceReadNotFoundError +from ..errors import WorkspaceArchiveReadError, WorkspaceReadNotFoundError from .base_sandbox_session import BaseSandboxSession from .events import EventPayloadPolicy, SandboxSessionEvent from .utils import event_to_json_line @@ -144,9 +144,13 @@ class WorkspaceJsonlSink(EventSink): """ Append events to a JSONL file inside the session workspace (under manifest.root). - This sink still runs in the client process, but writes into the session via - `SandboxSession.write()`, so it works across sandboxes (Docker/Modal) - without requiring host-mounted volumes. + This sink reads and replaces the outbox through the session's file APIs. + With a finite `max_bytes`, SDK backends limit acquisition from the source. + Legacy custom backends use their existing `read()` API; data acquired inside that method + is outside this sink's size limit. Backends own read deadlines and cleanup; + this sink does not impose an additional timeout. Workspace logs are + workload-modifiable and are not authoritative audit records. Use + `JsonlOutboxSink` or `HttpProxySink` for storage outside the workspace. """ def __init__( @@ -158,6 +162,7 @@ def __init__( on_error: OnErrorPolicy = "log", payload_policy: EventPayloadPolicy | None = None, flush_every: int = 1, + max_bytes: int | None = None, ) -> None: """ Args: @@ -168,7 +173,20 @@ def __init__( Example: Path("logs/events-{session_id}.jsonl") + max_bytes: Optional maximum replacement-file and pending-buffer size in bytes. + The default `None` preserves unlimited delivery and whole-file reads. + Set a positive budget to bound acquisition of workload-modifiable logs. + Exceeding either budget clears pending events + and permanently stops this sink, reporting one error through + `on_error`. Subsequent events are ignored, including after rebind. + Create a new sink for a new outbox or with a larger finite budget, + or use a host/HTTP sink for longer-lived logs. No file is truncated + or rotated when the budget is exceeded. """ + if max_bytes is not None and max_bytes <= 0: + raise ValueError("max_bytes must be positive") + self.max_bytes = max_bytes + self._disabled = False self.workspace_relpath = workspace_relpath self.ephemeral = ephemeral self.mode = mode @@ -204,7 +222,10 @@ def bind(self, session: BaseSandboxSession) -> None: self._session.register_persist_workspace_skip_path(relpath) def _buffer_event(self, event: SandboxSessionEvent) -> bool: - self._buf.extend(event_to_json_line(event).encode("utf-8")) + line = event_to_json_line(event).encode("utf-8") + if self.max_bytes is not None and len(self._buf) + len(line) > self.max_bytes: + self._stop_at_limit() + self._buf.extend(line) self._seen += 1 if self._seen % self._flush_every == 0: @@ -238,26 +259,38 @@ async def _flush_buffer(self) -> None: relpath = self._resolved_workspace_relpath or self.workspace_relpath existing = await self._read_existing_outbox(relpath) pending = bytes(self._buf) + if self.max_bytes is not None and len(existing) + len(pending) > self.max_bytes: + self._stop_at_limit() await self._session.write(relpath, io.BytesIO(existing + pending)) self._buf.clear() + def _stop_at_limit(self) -> NoReturn: + self._disabled = True + self._buf.clear() + raise RuntimeError( + f"WorkspaceJsonlSink exceeded max_bytes={self.max_bytes}; delivery stopped. " + "Use a new outbox and sink, or JsonlOutboxSink/HttpProxySink." + ) + async def _read_existing_outbox(self, relpath: Path) -> bytes: if self._session is None: return b"" try: - existing = await self._session.read(relpath) + if self.max_bytes is None: + existing = await self._session.read(relpath) + try: + payload = existing.read() + finally: + existing.close() + return payload.encode("utf-8") if isinstance(payload, str) else bytes(payload) + return await self._session.read_bounded(relpath, max_bytes=self.max_bytes + 1) except (FileNotFoundError, WorkspaceReadNotFoundError): return b"" - - try: - payload = existing.read() - finally: - existing.close() - - if isinstance(payload, str): - return payload.encode("utf-8") - return bytes(payload) + except WorkspaceArchiveReadError as error: + if error.context.get("reason") != "bounded_read_wire_limit": + raise + self._stop_at_limit() async def handle(self, event: SandboxSessionEvent) -> None: # If unbound (e.g., audit event emission used without a SandboxSession wrapper), @@ -266,6 +299,8 @@ async def handle(self, event: SandboxSessionEvent) -> None: return async with self._lock: + if self._disabled: + return if not self._buffer_event(event): return diff --git a/tests/conftest.py b/tests/conftest.py index 8e83cafa70..1a7ea46f86 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -80,6 +80,7 @@ def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: "sandbox/test_session_sinks.py", "sandbox/test_snapshot.py", "sandbox/test_unix_local.py", + "sandbox/test_workspace_jsonl_sink.py", ] ) diff --git a/tests/extensions/sandbox/test_modal.py b/tests/extensions/sandbox/test_modal.py index a84a79d75e..2be8e9334d 100644 --- a/tests/extensions/sandbox/test_modal.py +++ b/tests/extensions/sandbox/test_modal.py @@ -27,7 +27,7 @@ R2Mount, S3Mount, ) -from agents.sandbox.entries.mounts.base import InContainerMountAdapter +from agents.sandbox.entries.mounts.base import InContainerMountAdapter, MountStrategyBase from agents.sandbox.errors import ( InvalidManifestPathError, MountConfigError, @@ -445,11 +445,34 @@ class _FakeModalNotFoundError(_FakeModalError): monkeypatch.setitem(sys.modules, "modal.exception", fake_modal_exception) monkeypatch.setitem(sys.modules, "modal.config", fake_modal_config) monkeypatch.setitem(sys.modules, "modal.container_process", fake_container_process) - sys.modules.pop("agents.extensions.sandbox.modal.sandbox", None) - sys.modules.pop("agents.extensions.sandbox.modal.mounts", None) - sys.modules.pop("agents.extensions.sandbox.modal", None) + # Re-importing an adapter also replaces its registered model classes. + # Record undo before clearing these entries for fresh class registration; + # unrelated registrations must survive teardown. + monkeypatch.setitem( + SandboxSessionState._subclass_registry, + "modal", + SandboxSessionState._subclass_registry.get("modal", SandboxSessionState), + ) + del SandboxSessionState._subclass_registry["modal"] + monkeypatch.setitem( + MountStrategyBase._subclass_registry, + "modal_cloud_bucket", + MountStrategyBase._subclass_registry.get("modal_cloud_bucket", MountStrategyBase), + ) + del MountStrategyBase._subclass_registry["modal_cloud_bucket"] + module_names = ( + "agents.extensions.sandbox.modal.sandbox", + "agents.extensions.sandbox.modal.mounts", + "agents.extensions.sandbox.modal", + ) + for name in module_names: + monkeypatch.delitem(sys.modules, name, raising=False) module: Any = importlib.import_module("agents.extensions.sandbox.modal.sandbox") + # Track the fresh imports too, so teardown removes them before restoring + # any original modules alongside the real provider SDK. + for name in module_names: + monkeypatch.setitem(sys.modules, name, sys.modules.pop(name)) return module, create_calls, registry_tags diff --git a/tests/extensions/sandbox/test_runloop.py b/tests/extensions/sandbox/test_runloop.py index 80711bf3b2..5f6efd85e8 100644 --- a/tests/extensions/sandbox/test_runloop.py +++ b/tests/extensions/sandbox/test_runloop.py @@ -36,6 +36,7 @@ from agents.sandbox.session.base_sandbox_session import BaseSandboxSession from agents.sandbox.session.dependencies import Dependencies from agents.sandbox.session.sandbox_client import BaseSandboxClientOptions +from agents.sandbox.session.sandbox_session_state import SandboxSessionState from agents.sandbox.snapshot import NoopSnapshot, SnapshotBase from agents.sandbox.types import ExposedPortEndpoint from tests.utils.factories import make_run_state @@ -1277,9 +1278,28 @@ def _load_runloop_module(monkeypatch: pytest.MonkeyPatch) -> Any: "runloop_api_client.types.shared.launch_parameters", fake_launch_parameters_module, ) - sys.modules.pop("agents.extensions.sandbox.runloop.sandbox", None) - sys.modules.pop("agents.extensions.sandbox.runloop", None) - return importlib.import_module("agents.extensions.sandbox.runloop.sandbox") + # Re-importing an adapter also replaces its registered model classes. + # Record undo before clearing these entries for fresh class registration; + # unrelated registrations must survive teardown. + monkeypatch.setitem( + SandboxSessionState._subclass_registry, + "runloop", + SandboxSessionState._subclass_registry.get("runloop", SandboxSessionState), + ) + del SandboxSessionState._subclass_registry["runloop"] + module_names = ( + "agents.extensions.sandbox.runloop.sandbox", + "agents.extensions.sandbox.runloop", + ) + for name in module_names: + monkeypatch.delitem(sys.modules, name, raising=False) + + module: Any = importlib.import_module("agents.extensions.sandbox.runloop.sandbox") + # Track the fresh imports too, so teardown removes them before restoring + # any original modules alongside the real provider SDK. + for name in module_names: + monkeypatch.setitem(sys.modules, name, sys.modules.pop(name)) + return module def _build_tar_bytes(files: dict[str, bytes]) -> bytes: diff --git a/tests/extensions/sandbox/test_vercel.py b/tests/extensions/sandbox/test_vercel.py index 118d64423f..ca3c63df10 100644 --- a/tests/extensions/sandbox/test_vercel.py +++ b/tests/extensions/sandbox/test_vercel.py @@ -27,7 +27,7 @@ RcloneMountPattern, S3Mount, ) -from agents.sandbox.entries.mounts.base import InContainerMountAdapter +from agents.sandbox.entries.mounts.base import InContainerMountAdapter, MountStrategyBase from agents.sandbox.errors import ( ConfigurationError, ErrorCode, @@ -45,6 +45,7 @@ from agents.sandbox.session.base_sandbox_session import BaseSandboxSession from agents.sandbox.session.dependencies import Dependencies from agents.sandbox.session.manager import Instrumentation +from agents.sandbox.session.sandbox_session_state import SandboxSessionState from agents.sandbox.session.sinks import CallbackSink from agents.sandbox.snapshot import NoopSnapshot, SnapshotBase from agents.sandbox.types import User @@ -545,11 +546,35 @@ def _load_vercel_module(monkeypatch: pytest.MonkeyPatch) -> Any: monkeypatch.setitem(sys.modules, "vercel", fake_vercel) monkeypatch.setitem(sys.modules, "vercel.sandbox", fake_vercel_sandbox) - sys.modules.pop("agents.extensions.sandbox.vercel.mounts", None) - sys.modules.pop("agents.extensions.sandbox.vercel.sandbox", None) - sys.modules.pop("agents.extensions.sandbox.vercel", None) - - return importlib.import_module("agents.extensions.sandbox.vercel.sandbox") + # Re-importing an adapter also replaces its registered model classes. + # Record undo before clearing these entries for fresh class registration; + # unrelated registrations must survive teardown. + monkeypatch.setitem( + SandboxSessionState._subclass_registry, + "vercel", + SandboxSessionState._subclass_registry.get("vercel", SandboxSessionState), + ) + del SandboxSessionState._subclass_registry["vercel"] + monkeypatch.setitem( + MountStrategyBase._subclass_registry, + "vercel_cloud_bucket", + MountStrategyBase._subclass_registry.get("vercel_cloud_bucket", MountStrategyBase), + ) + del MountStrategyBase._subclass_registry["vercel_cloud_bucket"] + module_names = ( + "agents.extensions.sandbox.vercel.sandbox", + "agents.extensions.sandbox.vercel.mounts", + "agents.extensions.sandbox.vercel", + ) + for name in module_names: + monkeypatch.delitem(sys.modules, name, raising=False) + + module: Any = importlib.import_module("agents.extensions.sandbox.vercel.sandbox") + # Track the fresh imports too, so teardown removes them before restoring + # any original modules alongside the real provider SDK. + for name in module_names: + monkeypatch.setitem(sys.modules, name, sys.modules.pop(name)) + return module async def _noop_sleep(*_args: object, **_kwargs: object) -> None: diff --git a/tests/sandbox/_filesystem_test_session.py b/tests/sandbox/_filesystem_test_session.py index 83f1d4d400..54f34d4035 100644 --- a/tests/sandbox/_filesystem_test_session.py +++ b/tests/sandbox/_filesystem_test_session.py @@ -18,12 +18,13 @@ from agents.sandbox.manifest import Manifest from agents.sandbox.sandboxes.unix_local import ( UnixLocalSandboxClient, + UnixLocalSandboxSession, UnixLocalSandboxSessionState, ) from agents.sandbox.session import SandboxSession from agents.sandbox.session.base_sandbox_session import BaseSandboxSession from agents.sandbox.session.sandbox_session_state import SandboxSessionState -from agents.sandbox.snapshot import NoopSnapshot, SnapshotBase, SnapshotSpec +from agents.sandbox.snapshot import LocalSnapshot, NoopSnapshot, SnapshotBase, SnapshotSpec from agents.sandbox.types import ExecResult, Permissions, User @@ -75,6 +76,10 @@ async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase except OSError as error: raise WorkspaceArchiveReadError(path=path, cause=error) from error + async def _read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + with self.normalize_path(path).open("rb") as stream: + return stream.read(max_bytes) + async def write( self, path: Path, @@ -210,3 +215,42 @@ async def resume(self, state: SandboxSessionState) -> SandboxSession: FilesystemTestSandboxSession(state=unix_state), instrumentation=self._instrumentation, ) + + +def _build_unix_local_session( + tmp_path: Path, + *, + manifest: Manifest | None = None, + exposed_ports: tuple[int, ...] = (), +) -> UnixLocalSandboxSession: + workspace = tmp_path / "workspace" + snapshot = LocalSnapshot(id=str(uuid.uuid4()), base_path=tmp_path) + session_manifest = ( + manifest.model_copy(update={"root": str(workspace)}, deep=True) + if manifest is not None + else Manifest(root=str(workspace)) + ) + state = UnixLocalSandboxSessionState( + manifest=session_manifest, + snapshot=snapshot, + exposed_ports=exposed_ports, + ) + return UnixLocalSandboxSession.from_state(state) + + +def _build_filesystem_test_session( + tmp_path: Path, + *, + manifest: Manifest | None = None, +) -> FilesystemTestSandboxSession: + workspace = tmp_path / "workspace" + session_manifest = ( + manifest.model_copy(update={"root": str(workspace)}, deep=True) + if manifest is not None + else Manifest(root=str(workspace)) + ) + state = UnixLocalSandboxSessionState( + manifest=session_manifest, + snapshot=LocalSnapshot(id=str(uuid.uuid4()), base_path=tmp_path), + ) + return FilesystemTestSandboxSession(state=state) diff --git a/tests/sandbox/test_bounded_read.py b/tests/sandbox/test_bounded_read.py new file mode 100644 index 0000000000..9fbdba60cb --- /dev/null +++ b/tests/sandbox/test_bounded_read.py @@ -0,0 +1,714 @@ +from __future__ import annotations + +import asyncio +import base64 +import sys +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from agents.sandbox.errors import WorkspaceArchiveReadError, WorkspaceReadNotFoundError +from agents.sandbox.session import ( + CallbackSink, + Instrumentation, + SandboxSession, + SandboxSessionEvent, + SandboxSessionFinishEvent, +) +from agents.sandbox.types import ExecResult + + +class _Content: + def __init__(self, data: bytes, *, fail: bool = False) -> None: + self.data = data + self.offset = 0 + self.closed = False + self.fail = fail + + async def chunks(self, chunk_size: int = 3) -> AsyncIterator[bytes]: + while self.offset < len(self.data): + if self.fail: + raise ValueError("synthetic-private-response") + end = min(self.offset + min(chunk_size, 3), len(self.data)) + chunk = self.data[self.offset : end] + self.offset = end + yield chunk + + def __aiter__(self) -> AsyncIterator[bytes]: + return self.chunks() + + def iter_bytes(self, chunk_size: int) -> AsyncIterator[bytes]: + return self.chunks(chunk_size) + + def aiter_bytes(self, chunk_size: int) -> AsyncIterator[bytes]: + return self.chunks(chunk_size) + + def iter_chunked(self, chunk_size: int) -> AsyncIterator[bytes]: + return self.chunks(chunk_size) + + async def readexactly(self, size: int) -> bytes: + if self.fail: + raise ValueError("synthetic-private-response") + data = self.data[self.offset : self.offset + size] + self.offset += len(data) + if len(data) < size: + raise asyncio.IncompleteReadError(data, size) + return data + + async def __aenter__(self) -> _Content: + return self + + async def __aexit__(self, *args: object) -> None: + self.closed = True + + def close(self) -> None: + self.closed = True + + +PROVIDERS = ["e2b", "runloop", "vercel", "daytona", "blaxel", "cloudflare"] +CLASSES = { + "e2b": "E2BSandboxSession", + "runloop": "RunloopSandboxSession", + "vercel": "VercelSandboxSession", + "daytona": "DaytonaSandboxSession", + "blaxel": "BlaxelSandboxSession", + "cloudflare": "CloudflareSandboxSession", +} + + +def _session(provider: str, content: _Content, *, status: int = 200) -> Any: + module = pytest.importorskip(f"agents.extensions.sandbox.{provider}.sandbox") + cls = getattr(module, CLASSES[provider]) + session = object.__new__(cls) + session._validate_path_access = AsyncMock(return_value=Path("/workspace/out.jsonl")) + # The same response object supports the documented provider streaming shapes. + content.content = content # type: ignore[attr-defined] + content.status = status # type: ignore[attr-defined] + content.status_code = status # type: ignore[attr-defined] + if provider == "e2b": + session._sandbox = SimpleNamespace( + files=SimpleNamespace(read=AsyncMock(return_value=content)) + ) + elif provider == "runloop": + session.state = SimpleNamespace( + devbox_id="test", timeouts=SimpleNamespace(file_download_s=30) + ) + session._sdk = SimpleNamespace( + api=SimpleNamespace( + devboxes=SimpleNamespace( + with_streaming_response=SimpleNamespace( + download_file=MagicMock(return_value=content) + ) + ) + ) + ) + elif provider == "vercel": + + @asynccontextmanager + async def mount() -> AsyncIterator[None]: + session.mount_active = True + try: + yield + finally: + session.mount_active = False + + async def chunks() -> AsyncIterator[bytes]: + try: + async for chunk in content: + assert session.mount_active + yield chunk + finally: + assert session.mount_active + content.closed = True + + session._s3_mount_operation = mount + session._ensure_sandbox = AsyncMock( + return_value=SimpleNamespace(iter_file=AsyncMock(return_value=chunks())) + ) + elif provider == "daytona": + session.state = SimpleNamespace(timeouts=SimpleNamespace(file_download_s=30)) + session._sandbox = SimpleNamespace( + fs=SimpleNamespace( + _api_client=SimpleNamespace( + download_file_without_preload_content=AsyncMock(return_value=content) + ) + ) + ) + elif provider == "blaxel": + session._sandbox = SimpleNamespace( + fs=SimpleNamespace( + url="https://example.invalid", + format_path=lambda p: p, + get_client=lambda: SimpleNamespace(stream=MagicMock(return_value=content)), + ) + ) + else: + session._session = lambda: SimpleNamespace(get=MagicMock(return_value=content)) + session._url = lambda p: "https://example.invalid/" + p + session._request_timeout = lambda: None + return session + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", PROVIDERS) +@pytest.mark.parametrize("limit", [1, 5, 100]) +async def test_bounded_provider_read_closes_at_prefix(provider: str, limit: int) -> None: + content = _Content(b"\x00\xffabcdefghijk") + session = _session(provider, content) + assert await session.read_bounded(Path("out.jsonl"), max_bytes=limit) == content.data[:limit] + assert content.closed + assert content.offset <= max(9, limit + 2) + if provider == "vercel": + assert not session.mount_active + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", PROVIDERS) +async def test_bounded_provider_read_closes_and_discards_failed_response(provider: str) -> None: + content = _Content(b"private-response", fail=True) + session = _session(provider, content) + with pytest.raises(WorkspaceArchiveReadError) as caught: + await session.read_bounded(Path("out.jsonl"), max_bytes=5) + assert content.closed + assert caught.value.__context__ is None + assert "synthetic-private-response" not in repr(vars(caught.value)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ["daytona", "blaxel", "cloudflare"]) +async def test_bounded_http_read_missing_file_closes_without_body(provider: str) -> None: + content = _Content(b"private-error-body") + session = _session(provider, content, status=404) + with pytest.raises(WorkspaceReadNotFoundError): + await session.read_bounded(Path("out.jsonl"), max_bytes=5) + assert content.closed + assert content.offset == 0 + + +@pytest.mark.asyncio +async def test_cloudflare_bounded_read_decodes_existing_sse_format() -> None: + expected = b"\x00\xffbinary\n" + wire = ( + b'data: {"type":"metadata","isBinary":true}\n\n' + b'data: {"type":"chunk","data":"' + base64.b64encode(expected) + b'"}\n\n' + b'data: {"type":"complete"}\n\n' + ) + content = _Content(wire) + session = _session("cloudflare", content) + assert await session.read_bounded(Path("out.jsonl"), max_bytes=4) == expected[:4] + assert content.closed + + +@pytest.mark.asyncio +async def test_cloudflare_bounded_read_limits_encoded_response() -> None: + content = _Content(b'data: {"type":"chunk","data":"' + b"x" * 100000) + session = _session("cloudflare", content) + with pytest.raises(WorkspaceArchiveReadError) as caught: + await session.read_bounded(Path("out.jsonl"), max_bytes=4) + assert caught.value.context["reason"] == "bounded_read_wire_limit" + assert caught.value.__context__ is None + assert content.offset < 66000 + assert content.closed + + +@pytest.mark.asyncio +async def test_docker_bounded_read_uses_trusted_utilities_without_encoding() -> None: + pytest.importorskip("docker") + from agents.sandbox.sandboxes.docker import DockerSandboxSession + + session = object.__new__(DockerSandboxSession) + session._validate_path_access = AsyncMock(return_value=Path("/workspace/out.jsonl")) + session.exec = AsyncMock( + return_value=ExecResult(stdout=b"\x00\xffabc", stderr=b"", exit_code=0) + ) + assert await session.read_bounded(Path("out.jsonl"), max_bytes=5) == b"\x00\xffabc" + args = session.exec.call_args.args + assert args[0] == "/bin/sh" + assert "PATH=/usr/bin:/bin; export PATH" in args[2] + assert "base64" not in args[2] + assert args[-2:] == ("/workspace/out.jsonl", "5") + assert session.exec.call_args.kwargs == {"shell": False, "timeout": 30.0} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "close_error", [None, RuntimeError("Close failed"), asyncio.TimeoutError()] +) +async def test_modal_bounded_read_closes_provider_descriptor(close_error: Exception | None) -> None: + pytest.importorskip("modal") + from agents.extensions.sandbox.modal.sandbox import ModalSandboxSession + + session = object.__new__(ModalSandboxSession) + session._validate_path_access = AsyncMock(return_value=Path("/workspace/out.jsonl")) + session._ensure_sandbox = AsyncMock() + stream = SimpleNamespace( + read=SimpleNamespace(aio=AsyncMock(return_value=b"abc")), + close=SimpleNamespace(aio=AsyncMock(side_effect=close_error)), + ) + session._sandbox = SimpleNamespace(open=SimpleNamespace(aio=AsyncMock(return_value=stream))) + if close_error is None: + assert await session.read_bounded(Path("out.jsonl"), max_bytes=3) == b"abc" + else: + with pytest.raises(WorkspaceArchiveReadError) as caught: + await session.read_bounded(Path("out.jsonl"), max_bytes=3) + assert caught.value.context["reason"] == "bounded_read_failed" + assert caught.value.__context__ is None + stream.read.aio.assert_awaited_once_with(3) + stream.close.aio.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "limit,expected_requests", + [ + (4 * 1024 * 1024, [4 * 1024 * 1024]), + (8 * 1024 * 1024 + 1, [8 * 1024 * 1024 + 1, 4 * 1024 * 1024 + 1]), + (101 * 1024 * 1024, [100 * 1024 * 1024, 97 * 1024 * 1024]), + ], +) +async def test_modal_bounded_read_uses_large_bounded_requests( + limit: int, expected_requests: list[int] +) -> None: + pytest.importorskip("modal") + from agents.extensions.sandbox.modal.sandbox import ModalSandboxSession + + session = object.__new__(ModalSandboxSession) + session._validate_path_access = AsyncMock(return_value=Path("/workspace/out.jsonl")) + session._ensure_sandbox = AsyncMock() + history = b"x" * (4 * 1024 * 1024) + remaining = history + + async def read(size: int) -> bytes: + nonlocal remaining + chunk, remaining = remaining[:size], remaining[size:] + return chunk + + stream = SimpleNamespace( + read=SimpleNamespace(aio=AsyncMock(side_effect=read)), + close=SimpleNamespace(aio=AsyncMock()), + ) + session._sandbox = SimpleNamespace(open=SimpleNamespace(aio=AsyncMock(return_value=stream))) + assert await session.read_bounded(Path("out.jsonl"), max_bytes=limit) == history + requests = [call.args[0] for call in stream.read.aio.await_args_list] + assert requests == expected_requests + assert all(size <= 100 * 1024 * 1024 for size in requests) + stream.close.aio.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["acquisition", "body"]) +async def test_daytona_bounded_read_timeout_preserves_retryability(phase: str) -> None: + pytest.importorskip("daytona") + from agents.extensions.sandbox.daytona.sandbox import DaytonaSandboxSessionState + from agents.sandbox.manifest import Manifest + from agents.sandbox.snapshot import NoopSnapshot + + class TimeoutContent(_Content): + async def chunks(self, chunk_size: int = 3) -> AsyncIterator[bytes]: + yield b"x" + raise asyncio.TimeoutError("synthetic-private-response") + + content = TimeoutContent(b"fixture") + inner = _session("daytona", content) + inner.state = DaytonaSandboxSessionState( + sandbox_id="test", manifest=Manifest(), snapshot=NoopSnapshot(id="test") + ) + if phase == "acquisition": + inner._sandbox.fs._api_client.download_file_without_preload_content.side_effect = ( + asyncio.TimeoutError("synthetic-private-response") + ) + events: list[SandboxSessionEvent] = [] + wrapper = SandboxSession( + inner, + instrumentation=Instrumentation( + sinks=[CallbackSink(lambda event, _: events.append(event), mode="sync")] + ), + ) + with pytest.raises(WorkspaceArchiveReadError) as caught: + await wrapper.read_bounded(Path("out.jsonl"), max_bytes=5) + assert caught.value.retryable is True + assert caught.value.cause is caught.value.__cause__ is caught.value.__context__ is None + assert "synthetic-private-response" not in str(caught.value) + assert content.closed is (phase == "body") + finish = next(event for event in events if event.op == "read" and event.phase == "finish") + assert isinstance(finish, SandboxSessionFinishEvent) + assert finish.error_retryable is True + assert "synthetic-private-response" not in finish.model_dump_json() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "provider,retryable", + [ + (provider, retryable) + for provider in ["cloudflare", "vercel", "runloop", "modal", "daytona", "e2b"] + for retryable in [True, False] + ] + + [("blaxel", True), ("blaxel", None)], +) +async def test_bounded_provider_read_preserves_retry_policy( + provider: str, retryable: bool | None +) -> None: + import httpx + + from agents.sandbox.manifest import Manifest + from agents.sandbox.session.sandbox_session_state import SandboxSessionState + from agents.sandbox.snapshot import NoopSnapshot + + content = _Content(b"synthetic-private-response") + close = AsyncMock() + inner: Any + if provider == "modal": + modal = pytest.importorskip("modal") + from agents.extensions.sandbox.modal.sandbox import ModalSandboxSession + + inner = object.__new__(ModalSandboxSession) + inner._validate_path_access = AsyncMock(return_value=Path("/workspace/out.jsonl")) + inner._ensure_sandbox = AsyncMock() + error_cls = ( + modal.exception.InternalError if retryable else modal.exception.PermissionDeniedError + ) + stream = SimpleNamespace( + read=SimpleNamespace( + aio=AsyncMock(side_effect=error_cls("synthetic-private-response")) + ), + close=SimpleNamespace(aio=close), + ) + inner._sandbox = SimpleNamespace(open=SimpleNamespace(aio=AsyncMock(return_value=stream))) + else: + inner = _session(provider, content, status=503 if retryable else 403) + inner.state = SandboxSessionState( + type="test", manifest=Manifest(), snapshot=NoopSnapshot(id="test") + ) + if provider == "e2b": + from e2b import exceptions as sdk_errors + + error_cls = ( + sdk_errors.RateLimitException if retryable else sdk_errors.AuthenticationException + ) + inner._sandbox.files.read.side_effect = error_cls("synthetic-private-response") + elif provider == "daytona": + from agents.extensions.sandbox.daytona.sandbox import DaytonaSandboxSessionState + + inner.state = DaytonaSandboxSessionState( + sandbox_id="test", manifest=Manifest(), snapshot=NoopSnapshot(id="test") + ) + elif provider == "vercel": + from vercel import sandbox as sdk + + error_cls = sdk.SandboxRateLimitError if retryable else sdk.SandboxPermissionError + inner._ensure_sandbox.return_value.iter_file.side_effect = error_cls( + httpx.Response(429 if retryable else 403), "synthetic-private-response" + ) + elif provider == "runloop": + import runloop_api_client + + from agents.extensions.sandbox.runloop.sandbox import RunloopSandboxSessionState + + inner.state = RunloopSandboxSessionState( + devbox_id="test", manifest=Manifest(), snapshot=NoopSnapshot(id="test") + ) + error_cls = ( + runloop_api_client.RateLimitError + if retryable + else runloop_api_client.PermissionDeniedError + ) + inner._sdk.api.devboxes.with_streaming_response.download_file.side_effect = error_cls( + "synthetic-private-response", + response=httpx.Response( + 429 if retryable else 403, request=httpx.Request("GET", "https://example.invalid") + ), + body={"detail": "synthetic-private-response"}, + ) + events: list[SandboxSessionEvent] = [] + wrapper = SandboxSession( + inner, + instrumentation=Instrumentation( + sinks=[CallbackSink(lambda event, _: events.append(event), mode="sync")] + ), + ) + with pytest.raises(WorkspaceArchiveReadError) as caught: + await wrapper.read_bounded(Path("out.jsonl"), max_bytes=5) + assert caught.value.retryable is retryable + assert caught.value.cause is caught.value.__cause__ is caught.value.__context__ is None + finish = next(event for event in events if event.op == "read" and event.phase == "finish") + assert isinstance(finish, SandboxSessionFinishEvent) + assert finish.error_retryable is retryable + assert "synthetic-private-response" not in finish.model_dump_json() + if provider in {"cloudflare", "daytona", "blaxel"}: + assert content.closed + assert content.offset == 0 + elif provider == "modal": + close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_read_bounded_wrapper_forwards_limit() -> None: + content = _Content(b"abcdefgh") + inner = _session("daytona", content) + # The public wrapper's instrumentation needs a real session state. + from agents.sandbox.manifest import Manifest + from agents.sandbox.session.sandbox_session_state import SandboxSessionState + from agents.sandbox.snapshot import NoopSnapshot + + inner.state = SandboxSessionState( + type="test", manifest=Manifest(), snapshot=NoopSnapshot(id="test") + ) + inner.read_bounded = AsyncMock(return_value=b"abc") + wrapper = SandboxSession(inner) + assert await wrapper.read_bounded(Path("out.jsonl"), max_bytes=3) == b"abc" + inner.read_bounded.assert_awaited_once_with(Path("out.jsonl"), max_bytes=3) + + +@pytest.mark.asyncio +@pytest.mark.skipif(sys.platform == "win32", reason="UnixLocal is not available on Windows") +async def test_native_bounded_read_needs_no_process(tmp_path: Path) -> None: + from agents.sandbox.manifest import Manifest + from agents.sandbox.sandboxes.unix_local import ( + UnixLocalSandboxSession, + UnixLocalSandboxSessionState, + ) + from agents.sandbox.snapshot import NoopSnapshot + + (tmp_path / "out.jsonl").write_bytes(b"\x00\xffbinary data") + session = UnixLocalSandboxSession( + state=UnixLocalSandboxSessionState( + manifest=Manifest(root=str(tmp_path)), snapshot=NoopSnapshot(id="test") + ) + ) + session.exec = AsyncMock(side_effect=AssertionError("No process required")) + assert await session.read_bounded(Path("out.jsonl"), max_bytes=4) == b"\x00\xffbi" + with pytest.raises(WorkspaceReadNotFoundError): + await session.read_bounded(Path("missing"), max_bytes=4) + with pytest.raises(ValueError): + await session.read_bounded(Path("out.jsonl"), max_bytes=0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", PROVIDERS) +async def test_bounded_provider_read_cancellation_closes_response(provider: str) -> None: + entered = asyncio.Event() + + class PausedContent(_Content): + async def chunks(self, chunk_size: int = 3) -> AsyncIterator[bytes]: + entered.set() + await asyncio.Event().wait() + yield b"unreachable" + + async def readexactly(self, size: int) -> bytes: + entered.set() + await asyncio.Event().wait() + return b"unreachable" + + content = PausedContent(b"fixture") + session = _session(provider, content) + task = asyncio.create_task(session.read_bounded(Path("out.jsonl"), max_bytes=5)) + await asyncio.wait_for(entered.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert content.closed + if provider == "vercel": + assert not session.mount_active + + +@pytest.mark.asyncio +async def test_runloop_bounded_read_uses_streaming_sdk_response() -> None: + import httpx + + pytest.importorskip("runloop_api_client") + from runloop_api_client import AsyncRunloop + + consumed = 0 + closed = False + + class Body(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + nonlocal consumed + for _ in range(10): + consumed += 1 + yield b"x" * 65536 + + async def aclose(self) -> None: + nonlocal closed + closed = True + + def respond(request: httpx.Request) -> httpx.Response: + assert request.url.path.endswith("/download_file") + return httpx.Response(200, stream=Body()) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http: + async with AsyncRunloop(bearer_token="synthetic-test-key", http_client=http) as api: + session = _session("runloop", _Content(b"unused")) + session._sdk = SimpleNamespace(api=api) + assert await session.read_bounded(Path("out.jsonl"), max_bytes=5) == b"xxxxx" + assert consumed == 1 + assert closed + + +@pytest.mark.asyncio +async def test_e2b_bounded_read_closes_real_sdk_stream() -> None: + import httpx + + pytest.importorskip("e2b") + from e2b.sandbox.filesystem.filesystem import AsyncFileStreamReader + + consumed = 0 + closed = False + + class Body(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + nonlocal consumed + for _ in range(10): + consumed += 1 + yield b"x" * 65536 + + async def aclose(self) -> None: + nonlocal closed + closed = True + + stream = AsyncFileStreamReader(httpx.Response(200, stream=Body())) + session = _session("e2b", _Content(b"unused")) + session._sandbox.files.read = AsyncMock(return_value=stream) + assert await session.read_bounded(Path("out.jsonl"), max_bytes=5) == b"xxxxx" + assert consumed == 1 + assert closed + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "close_error", [None, RuntimeError("Close failed"), asyncio.TimeoutError()] +) +async def test_modal_bounded_read_cancellation_closes_descriptor( + close_error: Exception | None, +) -> None: + pytest.importorskip("modal") + from agents.extensions.sandbox.modal.sandbox import ModalSandboxSession + + entered = asyncio.Event() + + async def read(size: int) -> bytes: + entered.set() + await asyncio.Event().wait() + return b"unreachable" + + session: Any = object.__new__(ModalSandboxSession) + session._validate_path_access = AsyncMock(return_value=Path("/workspace/out.jsonl")) + session._ensure_sandbox = AsyncMock() + close = AsyncMock(side_effect=close_error) + stream = SimpleNamespace(read=SimpleNamespace(aio=read), close=SimpleNamespace(aio=close)) + session._sandbox = SimpleNamespace(open=SimpleNamespace(aio=AsyncMock(return_value=stream))) + task = asyncio.create_task(session.read_bounded(Path("out.jsonl"), max_bytes=5)) + await asyncio.wait_for(entered.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + close.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status,retryable", [(400, False), (429, True), (500, True), (418, None)]) +async def test_daytona_bounded_read_classifies_http_status( + status: int, retryable: bool | None +) -> None: + content = _Content(b"synthetic-private-response") + session = _session("daytona", content, status=status) + with pytest.raises(WorkspaceArchiveReadError) as caught: + await session.read_bounded(Path("out.jsonl"), max_bytes=5) + assert caught.value.retryable is retryable + assert content.closed + assert content.offset == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["success", "failure", "cancellation"]) +async def test_vercel_bounded_read_preserves_primary_failure_during_close(outcome: str) -> None: + import httpx + + pytest.importorskip("vercel.sandbox") + from vercel import sandbox as sdk + + session = _session("vercel", _Content(b"fixture")) + entered = asyncio.Event() + closed = False + + class Stream: + def __aiter__(self) -> Stream: + return self + + async def __anext__(self) -> bytes: + entered.set() + if outcome == "failure": + raise sdk.SandboxPermissionError(httpx.Response(403), "synthetic-private-read") + if outcome == "cancellation": + await asyncio.Event().wait() + return b"fixture" + + async def aclose(self) -> None: + nonlocal closed + assert session.mount_active + closed = True + raise sdk.SandboxRateLimitError(httpx.Response(429), "synthetic-private-close") + + session._ensure_sandbox.return_value.iter_file.return_value = Stream() + task = asyncio.create_task(session.read_bounded(Path("out.jsonl"), max_bytes=5)) + await asyncio.wait_for(entered.wait(), timeout=1) + if outcome == "cancellation": + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + with pytest.raises(WorkspaceArchiveReadError) as caught: + await task + assert caught.value.retryable is (outcome == "success") + assert caught.value.__cause__ is caught.value.__context__ is None + assert "synthetic-private" not in repr(vars(caught.value)) + assert closed + assert not session.mount_active + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error_name,retryable", + [("RateLimitException", True), ("TimeoutException", True), ("AuthenticationException", False)], +) +async def test_e2b_bounded_stream_failure_preserves_retryability( + error_name: str, retryable: bool +) -> None: + pytest.importorskip("e2b") + from e2b import exceptions as sdk_errors + + error_cls = getattr(sdk_errors, error_name) + + class FailedContent(_Content): + async def chunks(self, chunk_size: int = 3) -> AsyncIterator[bytes]: + yield b"x" + raise error_cls("synthetic-private-response") + + content = FailedContent(b"fixture") + session = _session("e2b", content) + with pytest.raises(WorkspaceArchiveReadError) as caught: + await session.read_bounded(Path("out.jsonl"), max_bytes=5) + assert caught.value.retryable is retryable + assert caught.value.__cause__ is caught.value.__context__ is None + assert "synthetic-private-response" not in repr(vars(caught.value)) + assert content.closed + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status,retryable", [(429, None), (500, True), (502, True), (504, True)]) +async def test_blaxel_bounded_read_preserves_existing_status_policy( + status: int, retryable: bool | None +) -> None: + content = _Content(b"synthetic-private-response") + session = _session("blaxel", content, status=status) + with pytest.raises(WorkspaceArchiveReadError) as caught: + await session.read_bounded(Path("out.jsonl"), max_bytes=5) + assert caught.value.retryable is retryable + assert content.closed + assert content.offset == 0 diff --git a/tests/sandbox/test_session_sinks.py b/tests/sandbox/test_session_sinks.py index 2fa520e002..2975e67bff 100644 --- a/tests/sandbox/test_session_sinks.py +++ b/tests/sandbox/test_session_sinks.py @@ -3,7 +3,6 @@ import asyncio import io import json -import tarfile import uuid from pathlib import Path from unittest.mock import MagicMock, patch @@ -11,13 +10,11 @@ import pytest from inline_snapshot import snapshot -from agents.sandbox.entries import Dir, File -from agents.sandbox.errors import WorkspaceReadNotFoundError -from agents.sandbox.manifest import Manifest -from agents.sandbox.sandboxes.unix_local import ( - UnixLocalSandboxSession, - UnixLocalSandboxSessionState, +from agents.sandbox.entries import File +from agents.sandbox.errors import ( + WorkspaceReadNotFoundError, ) +from agents.sandbox.manifest import Manifest from agents.sandbox.session import ( CallbackSink, ChainedSink, @@ -29,56 +26,18 @@ SandboxSessionEvent, SandboxSessionFinishEvent, SandboxSessionStartEvent, - WorkspaceJsonlSink, ) from agents.sandbox.session.base_sandbox_session import BaseSandboxSession from agents.sandbox.session.sandbox_session import _read_with_expected_span_errors -from agents.sandbox.snapshot import LocalSnapshot from agents.sandbox.types import ExecResult from agents.tracing import custom_span, trace -from tests.sandbox._filesystem_test_session import FilesystemTestSandboxSession +from tests.sandbox._filesystem_test_session import ( + _build_filesystem_test_session, + _build_unix_local_session, +) from tests.testing_processor import fetch_normalized_spans, fetch_ordered_spans -def _build_unix_local_session( - tmp_path: Path, - *, - manifest: Manifest | None = None, - exposed_ports: tuple[int, ...] = (), -) -> UnixLocalSandboxSession: - workspace = tmp_path / "workspace" - snapshot = LocalSnapshot(id=str(uuid.uuid4()), base_path=tmp_path) - session_manifest = ( - manifest.model_copy(update={"root": str(workspace)}, deep=True) - if manifest is not None - else Manifest(root=str(workspace)) - ) - state = UnixLocalSandboxSessionState( - manifest=session_manifest, - snapshot=snapshot, - exposed_ports=exposed_ports, - ) - return UnixLocalSandboxSession.from_state(state) - - -def _build_filesystem_test_session( - tmp_path: Path, - *, - manifest: Manifest | None = None, -) -> FilesystemTestSandboxSession: - workspace = tmp_path / "workspace" - session_manifest = ( - manifest.model_copy(update={"root": str(workspace)}, deep=True) - if manifest is not None - else Manifest(root=str(workspace)) - ) - state = UnixLocalSandboxSessionState( - manifest=session_manifest, - snapshot=LocalSnapshot(id=str(uuid.uuid4()), base_path=tmp_path), - ) - return FilesystemTestSandboxSession(state=state) - - @pytest.mark.asyncio async def test_filesystem_test_session_rejects_process_backed_operations(tmp_path: Path) -> None: session = _build_filesystem_test_session(tmp_path) @@ -228,215 +187,6 @@ def _callback(_event: SandboxSessionEvent, _session: BaseSandboxSession) -> None assert seen == [1, 2] -@pytest.mark.asyncio -@pytest.mark.requires_native_macos_sandbox -async def test_workspace_jsonl_sink_writes_into_workspace_and_persists(tmp_path: Path) -> None: - inner = _build_unix_local_session(tmp_path) - instrumentation = Instrumentation( - sinks=[WorkspaceJsonlSink(mode="sync", on_error="raise", ephemeral=False)] - ) - wrapped = SandboxSession(inner, instrumentation=instrumentation) - - async with wrapped as session: - await session.exec("echo hi") - - outbox_stream = await inner.read(Path(f"logs/events-{inner.state.session_id}.jsonl")) - lines = outbox_stream.read().decode("utf-8").splitlines() - assert any(json.loads(line)["op"] == "exec" for line in lines) - - snapshot_path = tmp_path / f"{inner.state.snapshot.id}.tar" - with tarfile.open(snapshot_path, mode="r:*") as tar: - names = [member.name for member in tar.getmembers()] - assert any(f"logs/events-{inner.state.session_id}.jsonl" in name for name in names) - - -@pytest.mark.asyncio -@pytest.mark.requires_native_macos_sandbox -async def test_workspace_jsonl_sink_supports_session_id_template(tmp_path: Path) -> None: - inner = _build_unix_local_session(tmp_path) - relpath = Path("logs/events-{session_id}.jsonl") - instrumentation = Instrumentation( - sinks=[ - WorkspaceJsonlSink( - mode="sync", - on_error="raise", - ephemeral=False, - workspace_relpath=relpath, - ) - ] - ) - wrapped = SandboxSession(inner, instrumentation=instrumentation) - - async with wrapped as session: - await session.exec("echo hi") - - expected_path = Path(f"logs/events-{inner.state.session_id}.jsonl") - outbox_stream = await inner.read(expected_path) - lines = outbox_stream.read().decode("utf-8").splitlines() - assert any(json.loads(line)["op"] == "exec" for line in lines) - - -@pytest.mark.asyncio -async def test_workspace_jsonl_sink_preserves_preexisting_outbox_contents(tmp_path: Path) -> None: - inner = _build_filesystem_test_session(tmp_path) - relpath = Path(f"logs/events-{inner.state.session_id}.jsonl") - old_line = b'{"old":true}\n' - - async with inner: - await inner.write(relpath, io.BytesIO(old_line)) - sink = WorkspaceJsonlSink(mode="sync", on_error="raise", ephemeral=False) - sink.bind(inner) - - start = SandboxSessionStartEvent( - session_id=inner.state.session_id, - seq=1, - op="write", - span_id=str(uuid.uuid4()), - ) - finish = SandboxSessionFinishEvent( - session_id=inner.state.session_id, - seq=2, - op="write", - span_id=start.span_id, - ok=True, - duration_ms=0.0, - ) - - await sink.handle(start) - await sink.handle(finish) - - outbox_stream = await inner.read(relpath) - lines = outbox_stream.read().decode("utf-8").splitlines() - - assert len(lines) == 3 - assert json.loads(lines[0]) == {"old": True} - assert json.loads(lines[1])["seq"] == 1 - assert json.loads(lines[2])["seq"] == 2 - - -@pytest.mark.asyncio -async def test_workspace_jsonl_sink_does_not_duplicate_lines_across_flushes( - tmp_path: Path, -) -> None: - inner = _build_filesystem_test_session(tmp_path) - relpath = Path(f"logs/events-{inner.state.session_id}.jsonl") - - async with inner: - sink = WorkspaceJsonlSink(mode="sync", on_error="raise", ephemeral=False, flush_every=1) - sink.bind(inner) - - for seq in (1, 2, 3): - await sink.handle( - SandboxSessionStartEvent( - session_id=inner.state.session_id, - seq=seq, - op="write", - span_id=str(uuid.uuid4()), - ) - ) - - outbox_stream = await inner.read(relpath) - lines = outbox_stream.read().decode("utf-8").splitlines() - - assert [json.loads(line)["seq"] for line in lines] == [1, 2, 3] - - -@pytest.mark.asyncio -async def test_workspace_jsonl_sink_clears_flushed_buffer(tmp_path: Path) -> None: - inner = _build_filesystem_test_session(tmp_path) - relpath = Path(f"logs/events-{inner.state.session_id}.jsonl") - - async with inner: - sink = WorkspaceJsonlSink(mode="sync", on_error="raise", ephemeral=False, flush_every=1) - sink.bind(inner) - - for seq in (1, 2): - await sink.handle( - SandboxSessionStartEvent( - session_id=inner.state.session_id, - seq=seq, - op="write", - span_id=str(uuid.uuid4()), - ) - ) - assert sink._buf == bytearray() - - outbox_stream = await inner.read(relpath) - lines = outbox_stream.read().decode("utf-8").splitlines() - - assert [json.loads(line)["seq"] for line in lines] == [1, 2] - - -@pytest.mark.asyncio -@pytest.mark.requires_native_macos_sandbox -async def test_workspace_jsonl_sink_ephemeral_excludes_runtime_outbox_with_existing_parent( - tmp_path: Path, -) -> None: - inner = _build_unix_local_session( - tmp_path, - manifest=Manifest( - entries={ - "logs": Dir( - children={ - "keep.txt": File(content=b"keep"), - } - ) - } - ), - ) - instrumentation = Instrumentation( - sinks=[WorkspaceJsonlSink(mode="sync", on_error="raise", ephemeral=True)] - ) - wrapped = SandboxSession(inner, instrumentation=instrumentation) - - async with wrapped as session: - await session.exec("echo hi") - relpath = Path(f"logs/events-{inner.state.session_id}.jsonl") - outbox_stream = await inner.read(relpath) - assert outbox_stream.read() - - logs_entry = inner.state.manifest.entries["logs"] - assert isinstance(logs_entry, Dir) - assert {str(child) for child in logs_entry.children.keys()} == {"keep.txt"} - - snapshot_path = tmp_path / f"{inner.state.snapshot.id}.tar" - with tarfile.open(snapshot_path, mode="r:*") as tar: - names = [member.name for member in tar.getmembers()] - assert any(name.endswith("logs/keep.txt") for name in names) - assert not any(f"logs/events-{inner.state.session_id}.jsonl" in name for name in names) - - -@pytest.mark.asyncio -@pytest.mark.requires_native_macos_sandbox -async def test_workspace_jsonl_sink_flushes_on_stop_when_flush_every_gt_one( - tmp_path: Path, -) -> None: - inner = _build_unix_local_session(tmp_path) - instrumentation = Instrumentation( - sinks=[ - WorkspaceJsonlSink( - mode="sync", - on_error="raise", - ephemeral=False, - flush_every=10, - ) - ] - ) - wrapped = SandboxSession(inner, instrumentation=instrumentation) - - async with wrapped as session: - await session.exec("echo hi") - - outbox_stream = await inner.read(Path(f"logs/events-{inner.state.session_id}.jsonl")) - lines = outbox_stream.read().decode("utf-8").splitlines() - assert lines - - snapshot_path = tmp_path / f"{inner.state.snapshot.id}.tar" - with tarfile.open(snapshot_path, mode="r:*") as tar: - names = [member.name for member in tar.getmembers()] - assert any(f"logs/events-{inner.state.session_id}.jsonl" in name for name in names) - - @pytest.mark.asyncio @pytest.mark.requires_native_macos_sandbox async def test_callback_sink_receives_bound_inner_session(tmp_path: Path) -> None: diff --git a/tests/sandbox/test_workspace_jsonl_sink.py b/tests/sandbox/test_workspace_jsonl_sink.py new file mode 100644 index 0000000000..d083fa51db --- /dev/null +++ b/tests/sandbox/test_workspace_jsonl_sink.py @@ -0,0 +1,748 @@ +from __future__ import annotations + +import asyncio +import io +import json +import tarfile +import traceback +import uuid +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest + +from agents import _debug +from agents.sandbox.entries import Dir, File, InContainerMountStrategy, RcloneMountPattern, S3Mount +from agents.sandbox.errors import ( + ExecTimeoutError, + ExecTransportError, + OpName, + WorkspaceArchiveReadError, +) +from agents.sandbox.manifest import Manifest +from agents.sandbox.sandboxes.unix_local import ( + UnixLocalSandboxSessionState, +) +from agents.sandbox.session import ( + CallbackSink, + Instrumentation, + SandboxSession, + SandboxSessionEvent, + SandboxSessionFinishEvent, + SandboxSessionStartEvent, + WorkspaceJsonlSink, +) +from agents.sandbox.session.base_sandbox_session import BaseSandboxSession +from agents.sandbox.session.sinks import OnErrorPolicy +from agents.sandbox.session.utils import event_to_json_line +from tests.sandbox._filesystem_test_session import ( + FilesystemTestSandboxSession, + _build_filesystem_test_session, + _build_unix_local_session, +) + + +class _BoundedReadSession(FilesystemTestSandboxSession): + """Record bounded API usage without exposing a whole-file read API.""" + + def __init__(self, state: UnixLocalSandboxSessionState) -> None: + super().__init__(state) + self.requests: list[tuple[Path, int]] = [] + + async def _read_bounded(self, path: Path, *, max_bytes: int) -> bytes: + self.requests.append((path, max_bytes)) + return await super()._read_bounded(path, max_bytes=max_bytes) + + +def _build_bounded_read_session(tmp_path: Path) -> _BoundedReadSession: + return _BoundedReadSession(_build_filesystem_test_session(tmp_path).state) + + +def _outbox_event(inner: BaseSandboxSession, *, op: OpName = "write") -> SandboxSessionStartEvent: + return SandboxSessionStartEvent( + session_id=inner.state.session_id, seq=1, op=op, span_id="test-span" + ) + + +class _LegacyReadSession(FilesystemTestSandboxSession): + # Model a custom backend that implements only the released read/write APIs. + _read_bounded = BaseSandboxSession._read_bounded + + +class _ShortReadStream(io.BytesIO): + def read(self, size: int = -1) -> bytes: + return super().read(min(size, 3)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("history_kind", ["missing", "binary", "text", "short"]) +async def test_workspace_jsonl_sink_legacy_backend_delivers( + tmp_path: Path, history_kind: str +) -> None: + inner = _LegacyReadSession(_build_filesystem_test_session(tmp_path).state) + sink = WorkspaceJsonlSink( + max_bytes=8 * 1024 * 1024, + workspace_relpath=Path("out.jsonl"), + mode="sync", + on_error="raise", + ) + instrumentation = Instrumentation(sinks=[sink]) + SandboxSession(inner, instrumentation=instrumentation) + old = "日本語\n".encode() if history_kind == "text" else b"\x00\xff\n" + stream: io.IOBase + if history_kind == "text": + stream = io.StringIO(old.decode()) + elif history_kind == "short": + old += b"short reads must preserve all history\n" + stream = _ShortReadStream(old) + else: + stream = io.BytesIO(old) + async with inner: + if history_kind == "missing": + old = b"" + await instrumentation.emit(_outbox_event(inner)) + stream.close() + else: + await inner.write(Path("out.jsonl"), io.BytesIO(old)) + with patch.object(inner, "read", return_value=stream): + await instrumentation.emit(_outbox_event(inner)) + content = (Path(inner.state.manifest.root) / "out.jsonl").read_bytes() + assert stream.closed + assert content.startswith(old) + assert json.loads(content[len(old) :])["seq"] == 1 + assert not sink._buf + + +@pytest.mark.asyncio +@pytest.mark.parametrize("text_stream", [False, True]) +async def test_workspace_jsonl_sink_legacy_backend_stops_at_limit( + tmp_path: Path, text_stream: bool +) -> None: + inner = _LegacyReadSession(_build_filesystem_test_session(tmp_path).state) + sink = WorkspaceJsonlSink(workspace_relpath=Path("out.jsonl"), max_bytes=1024) + sink.bind(inner) + old = "日本語\n" * 1024 + stream = io.StringIO(old) if text_stream else io.BytesIO(old.encode()) + async with inner: + await inner.write(Path("out.jsonl"), io.BytesIO(old.encode())) + with patch.object(inner, "read", return_value=stream): + with pytest.raises(RuntimeError, match="delivery stopped"): + await sink.handle(_outbox_event(inner)) + assert stream.closed + assert (Path(inner.state.manifest.root) / "out.jsonl").read_bytes() == old.encode() + assert not sink._buf + + +@pytest.mark.asyncio +async def test_workspace_jsonl_sink_legacy_backend_read_failure_closes_stream( + tmp_path: Path, +) -> None: + inner = _LegacyReadSession(_build_filesystem_test_session(tmp_path).state) + sink = WorkspaceJsonlSink( + max_bytes=8 * 1024 * 1024, + workspace_relpath=Path("out.jsonl"), + mode="sync", + on_error="raise", + ) + instrumentation = Instrumentation(sinks=[sink]) + SandboxSession(inner, instrumentation=instrumentation) + stream = io.BytesIO(b"original\n") + async with inner: + await inner.write(Path("out.jsonl"), io.BytesIO(b"original\n")) + with ( + patch.object(inner, "read", return_value=stream), + patch.object(stream, "read", side_effect=OSError("synthetic-private-payload")), + ): + with pytest.raises(RuntimeError, match="sandbox event sink failed") as caught: + await instrumentation.emit(_outbox_event(inner)) + assert stream.closed + assert (Path(inner.state.manifest.root) / "out.jsonl").read_bytes() == b"original\n" + assert sink._buf + error = caught.value.__context__ + assert isinstance(error, WorkspaceArchiveReadError) + assert error.__context__ is None + assert "synthetic-private-payload" not in str(error) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("slack", [-1, 0, 1]) +async def test_workspace_jsonl_sink_replacement_budget(tmp_path: Path, slack: int) -> None: + inner = _build_bounded_read_session(tmp_path) + event = _outbox_event(inner) + old = b'{"old":true}\n' + budget = len(old) + len(event_to_json_line(event).encode()) + slack + sink = WorkspaceJsonlSink(workspace_relpath=Path("out.jsonl"), max_bytes=budget) + sink.bind(inner) + async with inner: + await inner.write(Path("out.jsonl"), io.BytesIO(old)) + with patch.object(inner, "read", side_effect=AssertionError("Unbounded download")): + if slack < 0: + with pytest.raises(RuntimeError, match="delivery stopped"): + await sink.handle(event) + else: + await sink.handle(event) + content = (Path(inner.state.manifest.root) / "out.jsonl").read_bytes() + assert content.startswith(old) + if slack < 0: + assert content == old + else: + assert json.loads(content[len(old) :])["seq"] == 1 + assert not sink._buf + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", ["raise", "log", "ignore"]) +async def test_workspace_jsonl_sink_exhaustion_stops_buffering( + tmp_path: Path, policy: OnErrorPolicy, caplog: pytest.LogCaptureFixture +) -> None: + inner = _build_bounded_read_session(tmp_path) + sink = WorkspaceJsonlSink(workspace_relpath=Path("out.jsonl"), max_bytes=1024) + sink.mode = "sync" + sink.on_error = policy + instrumentation = Instrumentation(sinks=[sink]) + SandboxSession(inner, instrumentation=instrumentation) + old = b"synthetic-history\n" * 100 + async with inner: + await inner.write(Path("out.jsonl"), io.BytesIO(old)) + if policy == "raise": + with pytest.raises(RuntimeError, match="sandbox event sink failed"): + await instrumentation.emit(_outbox_event(inner)) + else: + await instrumentation.emit(_outbox_event(inner)) + for _ in range(20): + await instrumentation.emit(_outbox_event(inner)) + sink.bind(inner) + await instrumentation.emit(_outbox_event(inner)) + assert len(inner.requests) == 1 + assert not sink._buf + assert (Path(inner.state.manifest.root) / "out.jsonl").read_bytes() == old + assert len(caplog.records) == (1 if policy == "log" else 0) + assert "synthetic-history" not in caplog.text + + +@pytest.mark.asyncio +async def test_workspace_jsonl_sink_pending_budget_before_flush(tmp_path: Path) -> None: + inner = _build_bounded_read_session(tmp_path) + event = _outbox_event(inner) + sink = WorkspaceJsonlSink(max_bytes=len(event_to_json_line(event).encode()), flush_every=100) + sink.bind(inner) + async with inner: + await sink.handle(event) + with pytest.raises(RuntimeError, match="max_bytes"): + await sink.handle(event) + await sink.handle(event) + assert not inner.requests + assert not sink._buf + + +@pytest.mark.asyncio +async def test_workspace_jsonl_sink_retries_delivery_and_flushes_lifecycle(tmp_path: Path) -> None: + inner = _build_bounded_read_session(tmp_path) + sink = WorkspaceJsonlSink(workspace_relpath=Path("out.jsonl"), flush_every=100, ephemeral=True) + sink.bind(inner) + async with inner: + await sink.handle(_outbox_event(inner)) + with patch.object(inner, "write", side_effect=OSError("temporary failure")): + with pytest.raises(OSError): + await sink.handle(_outbox_event(inner, op="persist_workspace")) + assert sink._buf + await sink.handle(_outbox_event(inner, op="stop")) + content = (Path(inner.state.manifest.root) / "out.jsonl").read_text() + assert [json.loads(line)["op"] for line in content.splitlines()] == [ + "write", + "persist_workspace", + "stop", + ] + assert not sink._buf + assert inner._persist_workspace_skip_relpaths() == {Path("out.jsonl")} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("root_mount", [False, True]) +async def test_workspace_jsonl_sink_preserves_mounted_write_and_rebind( + tmp_path: Path, root_mount: bool +) -> None: + inner = _build_bounded_read_session(tmp_path) + inner.state.manifest.entries["storage"] = S3Mount( + bucket="test-bucket", + mount_path=inner.state.manifest.root if root_mount else "logs", + mount_strategy=InContainerMountStrategy(pattern=RcloneMountPattern()), + ) + original_exclusions = inner._persist_workspace_skip_relpaths() + sink = WorkspaceJsonlSink(workspace_relpath=Path("logs/out.jsonl")) + sink.bind(inner) + assert inner._persist_workspace_skip_relpaths() == original_exclusions + async with inner: + await sink.handle(_outbox_event(inner)) + sink = WorkspaceJsonlSink(workspace_relpath=Path("logs/out.jsonl")) + sink.bind(inner) + await sink.handle(_outbox_event(inner)) + content = (Path(inner.state.manifest.root) / "logs/out.jsonl").read_text() + assert len(content.splitlines()) == 2 + assert all(json.loads(line)["seq"] == 1 for line in content.splitlines()) + + +@pytest.mark.asyncio +async def test_workspace_jsonl_sink_bounds_read_before_transfer(tmp_path: Path) -> None: + inner = _build_bounded_read_session(tmp_path) + sink = WorkspaceJsonlSink(workspace_relpath=Path("out.jsonl"), max_bytes=1024) + sink.bind(inner) + async with inner: + old = b"ordinary fixture\n" * 1024 + await inner.write(Path("out.jsonl"), io.BytesIO(old)) + with patch.object(inner, "read", side_effect=AssertionError("Unbounded download")): + assert len(await sink._read_existing_outbox(Path("out.jsonl"))) == 1025 + with pytest.raises(RuntimeError, match="delivery stopped"): + await sink.handle(_outbox_event(inner)) + assert (Path(inner.state.manifest.root) / "out.jsonl").read_bytes() == old + + +@pytest.mark.asyncio +async def test_workspace_jsonl_sink_preserves_bytes_and_reports_failure( + tmp_path: Path, +) -> None: + inner = _build_bounded_read_session(tmp_path) + sink = WorkspaceJsonlSink(workspace_relpath=Path("out.jsonl")) + sink.bind(inner) + async with inner: + await sink.handle(_outbox_event(inner)) + old = b"\x00\xff\n" + "日本語\n".encode() + b"x" * 20000 + b"\n" + await inner.write(Path("out.jsonl"), io.BytesIO(old)) + await sink.handle(_outbox_event(inner)) + content = (Path(inner.state.manifest.root) / "out.jsonl").read_bytes() + assert content.startswith(old) + assert json.loads(content[len(old) :])["seq"] == 1 + # Read failures must not cause a replacement write. + await inner.mkdir(Path("directory")) + sink = WorkspaceJsonlSink(workspace_relpath=Path("directory")) + sink.bind(inner) + with pytest.raises(WorkspaceArchiveReadError): + await sink.handle(_outbox_event(inner)) + + +@pytest.mark.asyncio +async def test_workspace_jsonl_sink_failed_read_without_writing(tmp_path: Path) -> None: + inner = _build_bounded_read_session(tmp_path) + sink = WorkspaceJsonlSink( + max_bytes=8 * 1024 * 1024, + ) + sink.bind(inner) + async with inner: + with ( + patch.object(inner, "_read_bounded", side_effect=ValueError("synthetic-private-value")), + patch.object(inner, "write", new_callable=AsyncMock) as write, + ): + with pytest.raises(WorkspaceArchiveReadError) as caught: + await sink.handle(_outbox_event(inner)) + write.assert_not_called() + assert "synthetic-private-value" not in str(caught.value) + assert caught.value.__context__ is None + assert sink._buf + + +@pytest.mark.asyncio +@pytest.mark.parametrize("timeout_error", [False, True]) +@pytest.mark.parametrize("policy", ["raise", "log"]) +@pytest.mark.parametrize("redact", [False, True]) +async def test_workspace_jsonl_sink_errors_have_no_pending_payload( + tmp_path: Path, + timeout_error: bool, + policy: OnErrorPolicy, + redact: bool, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + inner = _build_bounded_read_session(tmp_path) + sink = WorkspaceJsonlSink(max_bytes=8 * 1024 * 1024, mode="sync", on_error=policy) + instrumentation = Instrumentation(sinks=[sink]) + SandboxSession(inner, instrumentation=instrumentation) + monkeypatch.setattr(_debug, "DONT_LOG_TOOL_DATA", redact) + event = _outbox_event(inner).model_copy(update={"data": {"secret": "synthetic-private-value"}}) + + async def fail(path: Path, *, max_bytes: int) -> bytes: + command = ("provider-file-read", str(path)) + if timeout_error: + raise ExecTimeoutError(command=command, timeout_s=30.0) + raise ExecTransportError(command=command) + + error: BaseException | None = None + async with inner: + with patch.object(inner, "_read_bounded", side_effect=fail): + if policy == "raise": + with pytest.raises(RuntimeError, match="sandbox event sink failed") as caught: + await instrumentation.emit(event) + error = caught.value + else: + await instrumentation.emit(event) + for record in caplog.records: + assert "synthetic-private-value" not in repr(vars(record)) + if redact: + assert record.exc_info is None + elif record.exc_info: + error = record.exc_info[1] + while error is not None: + assert "synthetic-private-value" not in repr(vars(error)) + assert "synthetic-private-value" not in "".join(traceback.format_exception(error)) + error = error.__context__ + assert sink._buf + + +def test_workspace_jsonl_sink_requires_positive_budget() -> None: + with pytest.raises(ValueError, match="max_bytes must be positive"): + WorkspaceJsonlSink(max_bytes=0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("explicit_none", [False, True]) +async def test_workspace_jsonl_sink_default_keeps_delivering_after_eight_mib( + tmp_path: Path, explicit_none: bool +) -> None: + inner = _build_filesystem_test_session(tmp_path) + relpath = Path("out.jsonl") + sink = ( + WorkspaceJsonlSink(workspace_relpath=relpath, mode="sync", on_error="raise", max_bytes=None) + if explicit_none + else WorkspaceJsonlSink(workspace_relpath=relpath, mode="sync", on_error="raise") + ) + instrumentation = Instrumentation(sinks=[sink]) + SandboxSession(inner, instrumentation=instrumentation) + old = b'{"old":"' + b"x" * (8 * 1024 * 1024) + b'"}\n' + async with inner: + await inner.write(relpath, io.BytesIO(old)) + await instrumentation.emit(_outbox_event(inner)) + await instrumentation.emit(_outbox_event(inner, op="stop")) + content = inner.normalize_path(relpath).read_bytes() + assert content.startswith(old) + assert [json.loads(line)["op"] for line in content[len(old) :].splitlines()] == [ + "write", + "stop", + ] + assert not sink._buf + + +@pytest.mark.asyncio +@pytest.mark.requires_native_macos_sandbox +async def test_workspace_jsonl_sink_writes_into_workspace_and_persists(tmp_path: Path) -> None: + inner = _build_unix_local_session(tmp_path) + instrumentation = Instrumentation( + sinks=[WorkspaceJsonlSink(mode="sync", on_error="raise", ephemeral=False)] + ) + wrapped = SandboxSession(inner, instrumentation=instrumentation) + + async with wrapped as session: + await session.exec("echo hi") + + outbox_stream = await inner.read(Path(f"logs/events-{inner.state.session_id}.jsonl")) + lines = outbox_stream.read().decode("utf-8").splitlines() + assert any(json.loads(line)["op"] == "exec" for line in lines) + + snapshot_path = tmp_path / f"{inner.state.snapshot.id}.tar" + with tarfile.open(snapshot_path, mode="r:*") as tar: + names = [member.name for member in tar.getmembers()] + assert any(f"logs/events-{inner.state.session_id}.jsonl" in name for name in names) + + +@pytest.mark.asyncio +@pytest.mark.requires_native_macos_sandbox +async def test_workspace_jsonl_sink_supports_session_id_template(tmp_path: Path) -> None: + inner = _build_unix_local_session(tmp_path) + relpath = Path("logs/events-{session_id}.jsonl") + instrumentation = Instrumentation( + sinks=[ + WorkspaceJsonlSink( + mode="sync", + on_error="raise", + ephemeral=False, + workspace_relpath=relpath, + ) + ] + ) + wrapped = SandboxSession(inner, instrumentation=instrumentation) + + async with wrapped as session: + await session.exec("echo hi") + + expected_path = Path(f"logs/events-{inner.state.session_id}.jsonl") + outbox_stream = await inner.read(expected_path) + lines = outbox_stream.read().decode("utf-8").splitlines() + assert any(json.loads(line)["op"] == "exec" for line in lines) + + +@pytest.mark.asyncio +async def test_workspace_jsonl_sink_preserves_preexisting_outbox_contents(tmp_path: Path) -> None: + inner = _build_bounded_read_session(tmp_path) + relpath = Path(f"logs/events-{inner.state.session_id}.jsonl") + old_line = b'{"old":true}\n' + + async with inner: + await inner.write(relpath, io.BytesIO(old_line)) + sink = WorkspaceJsonlSink(mode="sync", on_error="raise", ephemeral=False) + sink.bind(inner) + + start = SandboxSessionStartEvent( + session_id=inner.state.session_id, + seq=1, + op="write", + span_id=str(uuid.uuid4()), + ) + finish = SandboxSessionFinishEvent( + session_id=inner.state.session_id, + seq=2, + op="write", + span_id=start.span_id, + ok=True, + duration_ms=0.0, + ) + + await sink.handle(start) + await sink.handle(finish) + + outbox_stream = await inner.read(relpath) + lines = outbox_stream.read().decode("utf-8").splitlines() + + assert len(lines) == 3 + assert json.loads(lines[0]) == {"old": True} + assert json.loads(lines[1])["seq"] == 1 + assert json.loads(lines[2])["seq"] == 2 + + +@pytest.mark.asyncio +async def test_workspace_jsonl_sink_does_not_duplicate_lines_across_flushes( + tmp_path: Path, +) -> None: + inner = _build_bounded_read_session(tmp_path) + relpath = Path(f"logs/events-{inner.state.session_id}.jsonl") + + async with inner: + sink = WorkspaceJsonlSink(mode="sync", on_error="raise", ephemeral=False, flush_every=1) + sink.bind(inner) + + for seq in (1, 2, 3): + await sink.handle( + SandboxSessionStartEvent( + session_id=inner.state.session_id, + seq=seq, + op="write", + span_id=str(uuid.uuid4()), + ) + ) + + outbox_stream = await inner.read(relpath) + lines = outbox_stream.read().decode("utf-8").splitlines() + + assert [json.loads(line)["seq"] for line in lines] == [1, 2, 3] + + +@pytest.mark.asyncio +async def test_workspace_jsonl_sink_clears_flushed_buffer(tmp_path: Path) -> None: + inner = _build_bounded_read_session(tmp_path) + relpath = Path(f"logs/events-{inner.state.session_id}.jsonl") + + async with inner: + sink = WorkspaceJsonlSink(mode="sync", on_error="raise", ephemeral=False, flush_every=1) + sink.bind(inner) + + for seq in (1, 2): + await sink.handle( + SandboxSessionStartEvent( + session_id=inner.state.session_id, + seq=seq, + op="write", + span_id=str(uuid.uuid4()), + ) + ) + assert sink._buf == bytearray() + + outbox_stream = await inner.read(relpath) + lines = outbox_stream.read().decode("utf-8").splitlines() + + assert [json.loads(line)["seq"] for line in lines] == [1, 2] + + +@pytest.mark.asyncio +@pytest.mark.requires_native_macos_sandbox +async def test_workspace_jsonl_sink_ephemeral_excludes_runtime_outbox_with_existing_parent( + tmp_path: Path, +) -> None: + inner = _build_unix_local_session( + tmp_path, + manifest=Manifest( + entries={ + "logs": Dir( + children={ + "keep.txt": File(content=b"keep"), + } + ) + } + ), + ) + instrumentation = Instrumentation( + sinks=[WorkspaceJsonlSink(mode="sync", on_error="raise", ephemeral=True)] + ) + wrapped = SandboxSession(inner, instrumentation=instrumentation) + + async with wrapped as session: + await session.exec("echo hi") + relpath = Path(f"logs/events-{inner.state.session_id}.jsonl") + outbox_stream = await inner.read(relpath) + assert outbox_stream.read() + + logs_entry = inner.state.manifest.entries["logs"] + assert isinstance(logs_entry, Dir) + assert {str(child) for child in logs_entry.children.keys()} == {"keep.txt"} + + snapshot_path = tmp_path / f"{inner.state.snapshot.id}.tar" + with tarfile.open(snapshot_path, mode="r:*") as tar: + names = [member.name for member in tar.getmembers()] + assert any(name.endswith("logs/keep.txt") for name in names) + assert not any(f"logs/events-{inner.state.session_id}.jsonl" in name for name in names) + + +@pytest.mark.asyncio +@pytest.mark.requires_native_macos_sandbox +async def test_workspace_jsonl_sink_flushes_on_stop_when_flush_every_gt_one( + tmp_path: Path, +) -> None: + inner = _build_unix_local_session(tmp_path) + instrumentation = Instrumentation( + sinks=[ + WorkspaceJsonlSink( + mode="sync", + on_error="raise", + ephemeral=False, + flush_every=10, + ) + ] + ) + wrapped = SandboxSession(inner, instrumentation=instrumentation) + + async with wrapped as session: + await session.exec("echo hi") + + outbox_stream = await inner.read(Path(f"logs/events-{inner.state.session_id}.jsonl")) + lines = outbox_stream.read().decode("utf-8").splitlines() + assert lines + + snapshot_path = tmp_path / f"{inner.state.snapshot.id}.tar" + with tarfile.open(snapshot_path, mode="r:*") as tar: + names = [member.name for member in tar.getmembers()] + assert any(f"logs/events-{inner.state.session_id}.jsonl" in name for name in names) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("retryable", "reason", "expected_reason"), + [ + (True, "provider_failure", "bounded_read_failed"), + (False, "bounded_read_wire_limit", "bounded_read_wire_limit"), + (None, "provider_failure", "bounded_read_failed"), + ], +) +async def test_bounded_read_preserves_retryability_without_provider_diagnostics( + tmp_path: Path, retryable: bool | None, reason: str, expected_reason: str +) -> None: + events: list[SandboxSessionEvent] = [] + instrumentation = Instrumentation( + sinks=[CallbackSink(lambda event, _: events.append(event), mode="sync")] + ) + inner = _build_bounded_read_session(tmp_path) + failure = WorkspaceArchiveReadError( + path=Path("out.jsonl"), + context={"reason": reason, "response": "synthetic-private-payload"}, + cause=OSError("synthetic-private-payload"), + retryable=retryable, + ) + async with SandboxSession(inner, instrumentation=instrumentation) as session: + with patch.object(inner, "_read_bounded", side_effect=failure): + with pytest.raises(WorkspaceArchiveReadError) as caught: + await session.read_bounded(Path("out.jsonl"), max_bytes=100) + + error = caught.value + assert error is not failure + assert error.retryable is retryable + assert error.context == {"path": "out.jsonl", "reason": expected_reason} + assert error.cause is error.__cause__ is error.__context__ is None + finish = next(event for event in events if event.op == "read" and event.phase == "finish") + assert isinstance(finish, SandboxSessionFinishEvent) + assert finish.error_retryable is retryable + assert "synthetic-private-payload" not in finish.model_dump_json() + + +@pytest.mark.asyncio +async def test_bounded_read_preserves_classified_transport_retryability(tmp_path: Path) -> None: + inner = _build_bounded_read_session(tmp_path) + failure = ExecTransportError( + command=("read-helper", "synthetic-private-payload"), retryable=True + ) + with patch.object(inner, "_read_bounded", side_effect=failure): + with pytest.raises(WorkspaceArchiveReadError) as caught: + await inner.read_bounded(Path("out.jsonl"), max_bytes=100) + assert caught.value.retryable is True + assert caught.value.cause is caught.value.__cause__ is caught.value.__context__ is None + assert "synthetic-private-payload" not in str(caught.value) + + +@pytest.mark.asyncio +async def test_workspace_jsonl_sink_wire_budget_stops_delivery(tmp_path: Path) -> None: + inner = _build_bounded_read_session(tmp_path) + sink = WorkspaceJsonlSink( + max_bytes=8 * 1024 * 1024, + ) + sink.bind(inner) + async with inner: + with ( + patch.object( + inner, + "_read_bounded", + side_effect=WorkspaceArchiveReadError( + path=Path("out.jsonl"), context={"reason": "bounded_read_wire_limit"} + ), + ) as read, + patch.object(inner, "write", new_callable=AsyncMock) as write, + ): + with pytest.raises(RuntimeError, match="delivery stopped"): + await sink.handle(_outbox_event(inner)) + await sink.handle(_outbox_event(inner)) + read.assert_awaited_once() + write.assert_not_called() + assert not sink._buf + + +@pytest.mark.asyncio +async def test_workspace_sink_waits_for_backend_read_cleanup( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + inner = _build_bounded_read_session(tmp_path) + sink = WorkspaceJsonlSink( + max_bytes=8 * 1024 * 1024, workspace_relpath=Path("out.jsonl"), on_error="raise" + ) + sink.bind(inner) + inner.running = AsyncMock(return_value=True) # type: ignore[method-assign] + loop = asyncio.get_running_loop() + real_time = loop.time + offset = 0.0 + monkeypatch.setattr(loop, "time", lambda: real_time() + offset) + cleaned = False + + async def read(path: Path, *, max_bytes: int) -> bytes: + nonlocal offset, cleaned + # Advance beyond the former sink deadline without waiting in real time. + # The backend still owns its deadline and must finish cleanup first. + offset = 31.0 + for _ in range(4): + await asyncio.sleep(0) + cleaned = True + raise WorkspaceArchiveReadError(path=path, retryable=True) + + inner._read_bounded = read # type: ignore[method-assign] + inner.write = AsyncMock() # type: ignore[method-assign] + with pytest.raises(WorkspaceArchiveReadError): + await sink.handle(_outbox_event(inner)) + assert cleaned + inner.write.assert_not_awaited() + # A failed read retains the event for delivery once the backend recovers. + inner._read_bounded = AsyncMock(return_value=b"") # type: ignore[method-assign] + await sink.handle(_outbox_event(inner, op="stop")) + inner.write.assert_awaited_once() + written = inner.write.call_args.args[1].getvalue().splitlines() + assert [json.loads(line)["op"] for line in written] == ["write", "stop"] diff --git a/tests/test_trace_processor.py b/tests/test_trace_processor.py index aed7536f54..43a5b776cc 100644 --- a/tests/test_trace_processor.py +++ b/tests/test_trace_processor.py @@ -721,14 +721,18 @@ def test_backend_span_exporter_deadline_stops_during_5xx_retry_backoff(mock_clie mock_client.return_value.post.return_value = mock_response exporter = BackendSpanExporter(api_key="test_key", max_retries=3, base_delay=1.0) - with patch("agents.tracing.processors.time.sleep") as sleep_for_retry: - exporter._export_with_deadline( - [get_span(mock_processor())], deadline=time.monotonic() + 0.01 - ) + with patch("agents.tracing.processors.time") as clock: + # Spend the deadline budget only during backoff, independent of machine load. + clock.monotonic.return_value = 100.0 + + def advance_clock(delay: float) -> None: + clock.monotonic.return_value += delay + + clock.sleep.side_effect = advance_clock + exporter._export_with_deadline([get_span(mock_processor())], deadline=100.01) assert mock_client.return_value.post.call_count == 1 - sleep_for_retry.assert_called_once() - assert sleep_for_retry.call_args.args[0] <= 0.1 + clock.sleep.assert_called_once_with(pytest.approx(0.01)) exporter.close()