919 lines
32 KiB
Python
919 lines
32 KiB
Python
"""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,
|
|
)
|