"""Data models for the sandbox client.""" from __future__ import annotations from collections.abc import AsyncIterator, Awaitable, Callable, Iterator from dataclasses import dataclass from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, 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, ) @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})" ) # ============================================================================= # 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, since reconnect streams don't emit one. 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, ) -> None: self._stream = message_stream self._control = control self._sandbox = sandbox 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 # New executions (command_id=""): eager_start reads "started" message. # Reconnections (command_id set): skip eager_start since reconnect # streams don't send a "started" message. 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 SandboxOperationError( "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 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 self._exhausted = True 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 reconnect_attempts = 0 while True: try: for chunk in self._iter_stream(): reconnect_attempts = 0 # Reset on successful data if chunk.stream == "stdout": self._last_stdout_offset = chunk.offset + len( chunk.data.encode("utf-8") ) else: self._last_stderr_offset = chunk.offset + len( chunk.data.encode("utf-8") ) yield chunk return # Stream ended normally (exit message received) except SandboxConnectionError as e: if self._control and self._control.killed: raise reconnect_attempts += 1 if reconnect_attempts > self.MAX_AUTO_RECONNECTS: raise SandboxConnectionError( f"Lost connection {reconnect_attempts} times in " f"succession, giving up" ) from e is_hot_reload = isinstance(e, SandboxServerReloadError) if not is_hot_reload: delay = min( self._BACKOFF_BASE * (2 ** (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, since reconnect streams don't emit one. 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, ) -> None: self._stream = message_stream self._control = control self._sandbox = sandbox 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 # New executions (command_id=""): _ensure_started reads "started". # Reconnections (command_id set): skip since reconnect streams # don't send a "started" message. 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 SandboxOperationError( "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 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 self._exhausted = True async def __aiter__(self) -> AsyncIterator[OutputChunk]: """Async iterate with auto-reconnect on transient errors.""" import asyncio reconnect_attempts = 0 while True: try: async for chunk in self._aiter_stream(): reconnect_attempts = 0 if chunk.stream == "stdout": self._last_stdout_offset = chunk.offset + len( chunk.data.encode("utf-8") ) else: self._last_stderr_offset = chunk.offset + len( chunk.data.encode("utf-8") ) yield chunk return # Stream ended normally except SandboxConnectionError as e: if self._control and self._control.killed: raise reconnect_attempts += 1 if reconnect_attempts > self.MAX_AUTO_RECONNECTS: raise SandboxConnectionError( f"Lost connection {reconnect_attempts} times " f"in succession, giving up" ) from e is_hot_reload = isinstance(e, SandboxServerReloadError) if not is_hot_reload: delay = min( self._BACKOFF_BASE * (2 ** (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, )