"""Data models for the sandbox client.""" from __future__ import annotations import logging from collections.abc import AsyncIterator, Awaitable, Callable, Iterator from dataclasses import dataclass from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Literal, Optional import httpx from langsmith.sandbox._exceptions import ( SandboxConnectionError, SandboxOperationError, SandboxServerReloadError, ) if TYPE_CHECKING: from langsmith.sandbox._async_sandbox import AsyncSandbox from langsmith.sandbox._sandbox import Sandbox from langsmith.sandbox._ws_execute import ( _AsyncWSStreamControl, _WSStreamControl, ) logger = logging.getLogger(__name__) def _acknowledges_reconnect(msg: dict, command_id: Optional[str]) -> bool: """Whether a 'started' frame proves *this* command was reattached to. A command and its WebSocket are separate things: the command keeps running on the server and the socket is only this client's attachment to it. The server acknowledges every successful reattachment with 'started', which for a command that produces no output is the only evidence it landed. One naming a different command is not evidence, and must not clear the reconnect budget. """ acked = msg.get("command_id") if acked == command_id: return True logger.warning( "Ignoring reconnect acknowledgement for command %r while attached to %r", acked, command_id, ) return False class _StreamEndedBeforeStarted(SandboxOperationError): """A command WebSocket closed before the guest sent its 'started' frame. Internal marker for the idempotently-retryable early close (the proxied tunnel was torn down gracefully mid-handshake), as distinct from a command-level failure. Subclasses SandboxOperationError so callers that catch the public type are unaffected. """ @dataclass class ExecutionResult: """Result of executing a command in a sandbox.""" stdout: str stderr: str exit_code: int @property def success(self) -> bool: """Return True if the command exited with code 0.""" return self.exit_code == 0 @dataclass class ResourceStatus: """Lightweight provisioning status for any async-created resource. Attributes: status: Resource lifecycle status. One of "provisioning", "ready", "failed". status_message: Human-readable details when status is "failed", None otherwise. """ status: str status_message: Optional[str] = None @classmethod def from_dict(cls, data: dict[str, Any]) -> ResourceStatus: """Create a ResourceStatus from API response dict.""" return cls( status=data.get("status", "provisioning"), status_message=data.get("status_message"), ) @dataclass class Snapshot: """Represents a sandbox snapshot. Snapshots are built from Docker images or captured from running sandboxes. They are used to create new sandboxes. Attributes: id: Unique identifier (UUID). name: Display name. status: Build status. One of "building", "ready", "failed". fs_capacity_bytes: Filesystem capacity in bytes. docker_image: Source Docker image (for build snapshots). image_digest: Docker image digest after pull. source_sandbox_id: Source sandbox (for capture snapshots). status_message: Human-readable details when status is "failed". fs_used_bytes: Actual bytes used on the filesystem. created_by: User or service that created the snapshot. registry_id: Private registry ID, if applicable. created_at: Timestamp when the snapshot was created. updated_at: Timestamp when the snapshot was last updated. """ id: str name: str status: str fs_capacity_bytes: int docker_image: Optional[str] = None image_digest: Optional[str] = None source_sandbox_id: Optional[str] = None status_message: Optional[str] = None fs_used_bytes: Optional[int] = None created_by: Optional[str] = None registry_id: Optional[str] = None created_at: Optional[str] = None updated_at: Optional[str] = None @classmethod def from_dict(cls, data: dict[str, Any]) -> Snapshot: """Create a Snapshot from API response dict.""" return cls( id=data.get("id", ""), name=data.get("name", ""), status=data.get("status", "building"), fs_capacity_bytes=data.get("fs_capacity_bytes", 0), docker_image=data.get("docker_image"), image_digest=data.get("image_digest"), source_sandbox_id=data.get("source_sandbox_id"), status_message=data.get("status_message"), fs_used_bytes=data.get("fs_used_bytes"), created_by=data.get("created_by"), registry_id=data.get("registry_id"), created_at=data.get("created_at"), updated_at=data.get("updated_at"), ) # ============================================================================= # Service URL Models # ============================================================================= _AUTH_HEADER = "X-Langsmith-Sandbox-Service-Token" _REFRESH_MARGIN_SECONDS = 30 class ServiceURL: """Authenticated URL for accessing an HTTP service running in a sandbox. Properties auto-refresh the token transparently when it nears expiry. HTTP helper methods (``.get``, ``.post``, etc.) inject the auth header automatically. When constructed by :meth:`SandboxClient.service` or :meth:`Sandbox.service`, the object holds an internal refresher that re-calls the API to obtain a fresh token before the current one expires. Example:: svc = sb.service(port=3000) resp = svc.get("/api/data") # token injected + auto-refreshed print(svc.browser_url) # always-fresh URL """ def __init__( self, browser_url: str, service_url: str, token: str, expires_at: str, *, _refresher: Optional[Callable[[], ServiceURL]] = None, ) -> None: self._browser_url = browser_url self._service_url = service_url self._token = token self._expires_at = expires_at self._refresher = _refresher # -- Auto-refresh logic ------------------------------------------------- def _should_refresh(self) -> bool: if self._refresher is None: return False raw = self._expires_at.replace("Z", "+00:00") expires = datetime.fromisoformat(raw) if expires.tzinfo is None: expires = expires.replace(tzinfo=timezone.utc) remaining = (expires - datetime.now(timezone.utc)).total_seconds() return remaining <= _REFRESH_MARGIN_SECONDS def _maybe_refresh(self) -> None: if self._should_refresh(): fresh = self._refresher() # type: ignore[misc] self._browser_url = fresh._browser_url self._service_url = fresh._service_url self._token = fresh._token self._expires_at = fresh._expires_at # -- Properties (auto-refresh on access) -------------------------------- @property def token(self) -> str: """Return the raw JWT, refreshing if near expiry.""" self._maybe_refresh() return self._token @property def service_url(self) -> str: """Return the base URL, refreshing if near expiry.""" self._maybe_refresh() return self._service_url @property def browser_url(self) -> str: """Return the browser auth URL, refreshing if near expiry.""" self._maybe_refresh() return self._browser_url @property def expires_at(self) -> str: """Return the ISO 8601 expiration, refreshing if near expiry.""" self._maybe_refresh() return self._expires_at # -- HTTP helpers (stateless, one httpx call per request) ---------------- def request(self, method: str, path: str = "/", **kwargs: Any) -> httpx.Response: """Make an HTTP request to the service, injecting the auth header. Args: method: HTTP method (GET, POST, etc.). path: Path relative to the service URL. **kwargs: Forwarded to ``httpx.request``. Returns: httpx.Response. """ url = self.service_url.rstrip("/") + "/" + path.lstrip("/") headers = dict(kwargs.pop("headers", None) or {}) headers[_AUTH_HEADER] = self.token return httpx.request(method, url, headers=headers, **kwargs) def get(self, path: str = "/", **kwargs: Any) -> httpx.Response: """HTTP GET to the service.""" return self.request("GET", path, **kwargs) def post(self, path: str = "/", **kwargs: Any) -> httpx.Response: """HTTP POST to the service.""" return self.request("POST", path, **kwargs) def put(self, path: str = "/", **kwargs: Any) -> httpx.Response: """HTTP PUT to the service.""" return self.request("PUT", path, **kwargs) def patch(self, path: str = "/", **kwargs: Any) -> httpx.Response: """HTTP PATCH to the service.""" return self.request("PATCH", path, **kwargs) def delete(self, path: str = "/", **kwargs: Any) -> httpx.Response: """HTTP DELETE to the service.""" return self.request("DELETE", path, **kwargs) # -- Construction ------------------------------------------------------- @classmethod def from_dict( cls, data: dict[str, Any], *, _refresher: Optional[Callable[[], ServiceURL]] = None, ) -> ServiceURL: """Create a ServiceURL from API response dict.""" return cls( browser_url=data["browser_url"], service_url=data["service_url"], token=data["token"], expires_at=data["expires_at"], _refresher=_refresher, ) def __repr__(self) -> str: return ( f"ServiceURL(service_url={self._service_url!r}, " f"expires_at={self._expires_at!r})" ) class AsyncServiceURL: """Async variant of :class:`ServiceURL` with auto-refreshing token. Properties and HTTP helpers are async. Use with :meth:`AsyncSandboxClient.service` or :meth:`AsyncSandbox.service`. Example:: svc = await sb.service(port=3000) resp = await svc.get("/api/data") print(await svc.get_browser_url()) """ def __init__( self, browser_url: str, service_url: str, token: str, expires_at: str, *, _refresher: Optional[Callable[[], Awaitable[AsyncServiceURL]]] = None, ) -> None: self._browser_url = browser_url self._service_url = service_url self._token = token self._expires_at = expires_at self._refresher = _refresher # -- Auto-refresh logic ------------------------------------------------- def _should_refresh(self) -> bool: if self._refresher is None: return False raw = self._expires_at.replace("Z", "+00:00") expires = datetime.fromisoformat(raw) if expires.tzinfo is None: expires = expires.replace(tzinfo=timezone.utc) remaining = (expires - datetime.now(timezone.utc)).total_seconds() return remaining <= _REFRESH_MARGIN_SECONDS async def _maybe_refresh(self) -> None: if self._should_refresh(): fresh = await self._refresher() # type: ignore[misc] self._browser_url = fresh._browser_url self._service_url = fresh._service_url self._token = fresh._token self._expires_at = fresh._expires_at # -- Async accessors (auto-refresh on access) --------------------------- async def get_token(self) -> str: """Return the raw JWT, refreshing if near expiry.""" await self._maybe_refresh() return self._token async def get_service_url(self) -> str: """Return the base URL, refreshing if near expiry.""" await self._maybe_refresh() return self._service_url async def get_browser_url(self) -> str: """Return the browser auth URL, refreshing if near expiry.""" await self._maybe_refresh() return self._browser_url async def get_expires_at(self) -> str: """Return the ISO 8601 expiration, refreshing if near expiry.""" await self._maybe_refresh() return self._expires_at # -- Sync property access (no refresh, use when token is known-fresh) --- @property def token(self) -> str: """Return the raw JWT without refreshing.""" return self._token @property def service_url(self) -> str: """Return the base URL without refreshing.""" return self._service_url @property def browser_url(self) -> str: """Return the browser auth URL without refreshing.""" return self._browser_url @property def expires_at(self) -> str: """Return the expiration timestamp without refreshing.""" return self._expires_at # -- HTTP helpers (one request per call) -------------------------------- async def request( self, method: str, path: str = "/", **kwargs: Any ) -> httpx.Response: """Make an async HTTP request to the service, injecting the auth header. Args: method: HTTP method (GET, POST, etc.). path: Path relative to the service URL. **kwargs: Forwarded to ``httpx.AsyncClient.request``. Returns: httpx.Response. """ url = (await self.get_service_url()).rstrip("/") + "/" + path.lstrip("/") headers = dict(kwargs.pop("headers", None) or {}) headers[_AUTH_HEADER] = await self.get_token() async with httpx.AsyncClient() as client: return await client.request(method, url, headers=headers, **kwargs) async def get(self, path: str = "/", **kwargs: Any) -> httpx.Response: """Async HTTP GET to the service.""" return await self.request("GET", path, **kwargs) async def post(self, path: str = "/", **kwargs: Any) -> httpx.Response: """Async HTTP POST to the service.""" return await self.request("POST", path, **kwargs) async def put(self, path: str = "/", **kwargs: Any) -> httpx.Response: """Async HTTP PUT to the service.""" return await self.request("PUT", path, **kwargs) async def patch(self, path: str = "/", **kwargs: Any) -> httpx.Response: """Async HTTP PATCH to the service.""" return await self.request("PATCH", path, **kwargs) async def delete(self, path: str = "/", **kwargs: Any) -> httpx.Response: """Async HTTP DELETE to the service.""" return await self.request("DELETE", path, **kwargs) # -- Construction ------------------------------------------------------- @classmethod def from_dict( cls, data: dict[str, Any], *, _refresher: Optional[Callable[[], Awaitable[AsyncServiceURL]]] = None, ) -> AsyncServiceURL: """Create an AsyncServiceURL from API response dict.""" return cls( browser_url=data["browser_url"], service_url=data["service_url"], token=data["token"], expires_at=data["expires_at"], _refresher=_refresher, ) def __repr__(self) -> str: return ( f"AsyncServiceURL(service_url={self._service_url!r}, " f"expires_at={self._expires_at!r})" ) # ============================================================================= # Download URL Models # ============================================================================= DownloadContentDisposition = Literal["attachment", "inline"] """How a download link asks the browser to handle the file.""" @dataclass class DownloadURL: """A link that downloads one sandbox file with no LangSmith credential. The link is pinned to the sandbox, the file path, and the response headers, so it cannot be repointed at another file. It is pinned to the path rather than to a snapshot of the contents, so the file must not be modified while the link is in use. Attributes: download_url: The full URL to fetch. Supports GET, HEAD, and Range. token: The signed token embedded in ``download_url``. expires_at: Expiry timestamp, or None for a link that never expires. """ download_url: str token: str expires_at: Optional[str] = None @classmethod def from_dict(cls, data: dict[str, Any]) -> DownloadURL: """Create a DownloadURL from API response dict.""" return cls( download_url=data.get("download_url", ""), token=data.get("token", ""), expires_at=data.get("expires_at"), ) # ============================================================================= # WebSocket Command Execution Models # ============================================================================= @dataclass class OutputChunk: """A single chunk of streaming output from command execution. Attributes: stream: Either "stdout" or "stderr". data: The text content of this chunk (valid UTF-8, server handles boundary splitting). offset: Byte offset within the stream. Used internally for reconnection; users typically don't need this. """ stream: str data: str offset: int class CommandHandle: """Handle to a running command with streaming output and auto-reconnect. Iterable, yielding OutputChunk objects (stdout and stderr interleaved in arrival order). Access .result after iteration to get the full ExecutionResult. Auto-reconnect behavior: - Server hot-reload (1001 Going Away): reconnect immediately - Network error / unexpected close: reconnect with exponential backoff - User called kill(): do NOT reconnect (propagate error) The auto-reconnect is transparent -- the iterator reconnects and continues yielding chunks without any user intervention. If all reconnect attempts are exhausted, SandboxConnectionError is raised. Construction modes (controlled by ``command_id``): - **New execution** (``command_id=""``, the default): the constructor eagerly reads the server's ``"started"`` message to populate ``command_id`` and ``pid`` before returning. - **Reconnection** (``command_id`` set): skips the started-message read. A reconnect stream's ``"started"`` message is the server's acknowledgement that the reattachment landed, consumed while iterating to clear the reconnect budget — the only such signal for a command that emits no output. Example: handle = sandbox.run("make build", timeout=600, wait=False) for chunk in handle: # auto-reconnects on transient errors print(chunk.data, end="") result = handle.result print(f"Exit code: {result.exit_code}") """ MAX_AUTO_RECONNECTS = 5 _BACKOFF_BASE = 0.5 # seconds _BACKOFF_MAX = 8.0 # seconds def __init__( self, message_stream: Iterator[dict], control: Optional[_WSStreamControl], sandbox: Sandbox, *, command_id: str = "", stdout_offset: int = 0, stderr_offset: int = 0, on_stdout: Optional[Callable[[str], Any]] = None, on_stderr: Optional[Callable[[str], Any]] = None, ) -> None: self._stream = message_stream self._control = control self._sandbox = sandbox self._on_stdout = on_stdout self._on_stderr = on_stderr self._command_id: Optional[str] = None self._pid: Optional[int] = None self._result: Optional[ExecutionResult] = None self._stdout_parts: list[str] = [] self._stderr_parts: list[str] = [] self._exhausted = False self._last_stdout_offset = stdout_offset self._last_stderr_offset = stderr_offset self._reconnect_attempts = 0 # New executions (command_id=""): eager_start reads "started" message. # Reconnections (command_id set): the "started" message is the server's # reattachment acknowledgement, consumed by _iter_stream instead. if command_id: self._command_id = command_id else: self._consume_started() def _consume_started(self) -> None: """Eagerly read the 'started' message to populate command_id and pid. Blocks briefly until the server sends the started message (arrives near-instantly after connection). After this call, command_id and pid are available, and the WebSocket is bound to the control object (so kill() works). """ try: first_msg = next(self._stream) except StopIteration: raise _StreamEndedBeforeStarted( "Command stream ended before 'started' message", operation="command", ) if first_msg.get("type") != "started": raise SandboxOperationError( f"Expected 'started' message, got '{first_msg.get('type')}'", operation="command", ) self._command_id = first_msg.get("command_id") self._pid = first_msg.get("pid") @property def command_id(self) -> Optional[str]: """The server-assigned command ID. Available after construction.""" return self._command_id @property def pid(self) -> Optional[int]: """The process ID on the sandbox. Available after construction.""" return self._pid @property def result(self) -> ExecutionResult: """The final execution result. Blocks until the command completes. Drains the remaining stream if not already exhausted, then returns the ExecutionResult with aggregated stdout, stderr, and exit_code. """ if self._result is None: for _ in self: pass if self._result is None: raise SandboxOperationError( "Command stream ended without exit message", operation="command", ) return self._result def _iter_stream(self) -> Iterator[OutputChunk]: """Iterate over output chunks from the current stream (no reconnect).""" if self._exhausted: return for msg in self._stream: msg_type = msg.get("type") if msg_type == "started": if _acknowledges_reconnect(msg, self._command_id): self._reconnect_attempts = 0 elif msg_type in ("stdout", "stderr"): chunk = OutputChunk( stream=msg_type, data=msg["data"], offset=msg.get("offset", 0), ) if msg_type == "stdout": self._stdout_parts.append(msg["data"]) else: self._stderr_parts.append(msg["data"]) yield chunk elif msg_type == "exit": self._result = ExecutionResult( stdout="".join(self._stdout_parts), stderr="".join(self._stderr_parts), exit_code=msg["exit_code"], ) self._exhausted = True return raise SandboxConnectionError("Command stream ended without exit message") def __iter__(self) -> Iterator[OutputChunk]: """Iterate over output chunks, auto-reconnecting on transient errors. Reconnect strategy: - 1001 Going Away (hot-reload): immediate reconnect, no delay - Other SandboxConnectionError: exponential backoff (0.5s, 1s, 2s...) - After kill(): no reconnect, error propagates """ import time self._reconnect_attempts = 0 while True: try: for chunk in self._iter_stream(): self._reconnect_attempts = 0 # Reset on successful data if chunk.stream == "stdout": self._last_stdout_offset = chunk.offset + len( chunk.data.encode("utf-8") ) if self._on_stdout is not None: self._on_stdout(chunk.data) else: self._last_stderr_offset = chunk.offset + len( chunk.data.encode("utf-8") ) if self._on_stderr is not None: self._on_stderr(chunk.data) yield chunk return # Stream ended normally (exit message received) except SandboxConnectionError as e: if self._control and self._control.killed: raise self._reconnect_attempts += 1 if self._reconnect_attempts > self.MAX_AUTO_RECONNECTS: raise SandboxConnectionError( f"Failed to reattach to the command " f"{self._reconnect_attempts} times in succession, " f"giving up" ) from e is_hot_reload = isinstance(e, SandboxServerReloadError) if not is_hot_reload: delay = min( self._BACKOFF_BASE * (2 ** (self._reconnect_attempts - 1)), self._BACKOFF_MAX, ) time.sleep(delay) assert self._command_id is not None new_handle = self._sandbox.reconnect( self._command_id, stdout_offset=self._last_stdout_offset, stderr_offset=self._last_stderr_offset, ) self._stream = new_handle._stream self._control = new_handle._control self._exhausted = False def kill(self) -> None: """Send a kill signal to the running command (SIGKILL). The server kills the entire process group. The stream will subsequently yield an exit message with a non-zero exit code. Has no effect if the command has already exited or the WebSocket connection is closed. """ if self._control: self._control.send_kill() def send_input(self, data: str) -> None: """Write data to the command's stdin. Args: data: String data to write to stdin. Has no effect if the command has already exited or the WebSocket connection is closed. """ if self._control: self._control.send_input(data) @property def last_stdout_offset(self) -> int: """Last known stdout byte offset (for manual reconnection).""" return self._last_stdout_offset @property def last_stderr_offset(self) -> int: """Last known stderr byte offset (for manual reconnection).""" return self._last_stderr_offset def reconnect(self) -> CommandHandle: """Reconnect to this command from the last known offsets. Returns a new handle that resumes output from where this one left off. Any output produced while disconnected is replayed from the server's ring buffer. Returns: A new CommandHandle. Raises: SandboxOperationError: If command_id is not found or session expired. SandboxConnectionError: If connection to sandbox fails. """ assert self._command_id is not None return self._sandbox.reconnect( self._command_id, stdout_offset=self._last_stdout_offset, stderr_offset=self._last_stderr_offset, ) class AsyncCommandHandle: """Async handle to a running command with streaming output and auto-reconnect. Async iterable, yielding OutputChunk objects (stdout and stderr interleaved in arrival order). Access .result after iteration to get the full ExecutionResult. Auto-reconnect behavior: - Server hot-reload (1001 Going Away): reconnect immediately - Network error / unexpected close: reconnect with exponential backoff - User called kill(): do NOT reconnect (propagate error) Construction modes (controlled by ``command_id``): - **New execution** (``command_id=""``, the default): call ``await handle._ensure_started()`` after construction to read the server's ``"started"`` message and populate ``command_id`` / ``pid``. - **Reconnection** (``command_id`` set): skips the started-message read. A reconnect stream's ``"started"`` message is the server's acknowledgement that the reattachment landed, consumed while iterating to clear the reconnect budget — the only such signal for a command that emits no output. Example: handle = await sandbox.run("make build", timeout=600, wait=False) async for chunk in handle: # auto-reconnects on transient errors print(chunk.data, end="") result = await handle.result print(f"Exit code: {result.exit_code}") """ MAX_AUTO_RECONNECTS = 5 _BACKOFF_BASE = 0.5 # seconds _BACKOFF_MAX = 8.0 # seconds def __init__( self, message_stream: AsyncIterator[dict], control: Optional[_AsyncWSStreamControl], sandbox: AsyncSandbox, *, command_id: str = "", stdout_offset: int = 0, stderr_offset: int = 0, on_stdout: Optional[Callable[[str], Any]] = None, on_stderr: Optional[Callable[[str], Any]] = None, ) -> None: self._stream = message_stream self._control = control self._sandbox = sandbox self._on_stdout = on_stdout self._on_stderr = on_stderr self._command_id: Optional[str] = None self._pid: Optional[int] = None self._result: Optional[ExecutionResult] = None self._stdout_parts: list[str] = [] self._stderr_parts: list[str] = [] self._exhausted = False self._last_stdout_offset = stdout_offset self._last_stderr_offset = stderr_offset self._reconnect_attempts = 0 # New executions (command_id=""): _ensure_started reads "started". # Reconnections (command_id set): the "started" message is the server's # reattachment acknowledgement, consumed by _aiter_stream instead. if command_id: self._command_id = command_id self._started = True else: self._started = False async def _ensure_started(self) -> None: """Read the 'started' message to populate command_id and pid.""" if self._started: return try: first_msg = await self._stream.__anext__() except StopAsyncIteration: raise _StreamEndedBeforeStarted( "Command stream ended before 'started' message", operation="command", ) if first_msg.get("type") != "started": raise SandboxOperationError( f"Expected 'started' message, got '{first_msg.get('type')}'", operation="command", ) self._command_id = first_msg.get("command_id") self._pid = first_msg.get("pid") self._started = True @property def command_id(self) -> Optional[str]: """The server-assigned command ID. Available after _ensure_started.""" return self._command_id @property def pid(self) -> Optional[int]: """The process ID on the sandbox. Available after _ensure_started.""" return self._pid @property async def result(self) -> ExecutionResult: """The final execution result. Awaitable.""" if self._result is None: async for _ in self: pass if self._result is None: raise SandboxOperationError( "Command stream ended without exit message", operation="command", ) return self._result async def _aiter_stream(self) -> AsyncIterator[OutputChunk]: """Iterate over output chunks from the current stream (no reconnect).""" await self._ensure_started() if self._exhausted: return async for msg in self._stream: msg_type = msg.get("type") if msg_type == "started": if _acknowledges_reconnect(msg, self._command_id): self._reconnect_attempts = 0 elif msg_type in ("stdout", "stderr"): chunk = OutputChunk( stream=msg_type, data=msg["data"], offset=msg.get("offset", 0), ) if msg_type == "stdout": self._stdout_parts.append(msg["data"]) else: self._stderr_parts.append(msg["data"]) yield chunk elif msg_type == "exit": self._result = ExecutionResult( stdout="".join(self._stdout_parts), stderr="".join(self._stderr_parts), exit_code=msg["exit_code"], ) self._exhausted = True return raise SandboxConnectionError("Command stream ended without exit message") async def __aiter__(self) -> AsyncIterator[OutputChunk]: """Async iterate with auto-reconnect on transient errors.""" import asyncio self._reconnect_attempts = 0 while True: try: async for chunk in self._aiter_stream(): self._reconnect_attempts = 0 if chunk.stream == "stdout": self._last_stdout_offset = chunk.offset + len( chunk.data.encode("utf-8") ) if self._on_stdout is not None: self._on_stdout(chunk.data) else: self._last_stderr_offset = chunk.offset + len( chunk.data.encode("utf-8") ) if self._on_stderr is not None: self._on_stderr(chunk.data) yield chunk return # Stream ended normally except SandboxConnectionError as e: if self._control and self._control.killed: raise self._reconnect_attempts += 1 if self._reconnect_attempts > self.MAX_AUTO_RECONNECTS: raise SandboxConnectionError( f"Failed to reattach to the command " f"{self._reconnect_attempts} times in succession, " f"giving up" ) from e is_hot_reload = isinstance(e, SandboxServerReloadError) if not is_hot_reload: delay = min( self._BACKOFF_BASE * (2 ** (self._reconnect_attempts - 1)), self._BACKOFF_MAX, ) await asyncio.sleep(delay) assert self._command_id is not None new_handle = await self._sandbox.reconnect( self._command_id, stdout_offset=self._last_stdout_offset, stderr_offset=self._last_stderr_offset, ) self._stream = new_handle._stream self._control = new_handle._control self._exhausted = False async def kill(self) -> None: """Send a kill signal to the running command.""" if self._control: await self._control.send_kill() async def send_input(self, data: str) -> None: """Write data to the command's stdin.""" if self._control: await self._control.send_input(data) @property def last_stdout_offset(self) -> int: """Last known stdout byte offset (for manual reconnection).""" return self._last_stdout_offset @property def last_stderr_offset(self) -> int: """Last known stderr byte offset (for manual reconnection).""" return self._last_stderr_offset async def reconnect(self) -> AsyncCommandHandle: """Reconnect to this command from the last known offsets.""" assert self._command_id is not None return await self._sandbox.reconnect( self._command_id, stdout_offset=self._last_stdout_offset, stderr_offset=self._last_stderr_offset, )