736 lines
25 KiB
Python

"""Sandbox class for interacting with a specific sandbox instance."""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union, overload
import httpx
from langsmith.sandbox._exceptions import (
DataplaneNotConfiguredError,
ResourceNotFoundError,
SandboxConnectionError,
SandboxNotReadyError,
)
from langsmith.sandbox._helpers import handle_sandbox_http_error
from langsmith.sandbox._models import (
CommandHandle,
ExecutionResult,
ServiceURL,
Snapshot,
)
from langsmith.sandbox._tunnel import Tunnel
if TYPE_CHECKING:
from langsmith.sandbox._client import SandboxClient
RequestHeaders = Optional[Mapping[str, str]]
@dataclass
class Sandbox:
"""Represents an active sandbox for running commands and file operations.
This class is typically obtained from SandboxClient.sandbox() and supports
the context manager protocol for automatic cleanup.
Attributes:
name: Display name (can be updated).
dataplane_url: URL for data plane operations (file I/O, command execution).
Only functional when status is "ready".
id: Unique identifier (UUID). Remains constant even if name changes.
May be None for resources created before ID support was added.
status: Sandbox lifecycle status. One of "provisioning", "ready",
"failed", "stopped".
status_message: Human-readable details when status is "failed", None otherwise.
created_at: Timestamp when the sandbox was created.
updated_at: Timestamp when the sandbox was last updated.
idle_ttl_seconds: Idle timeout TTL in seconds (``0`` means disabled).
Newly-created sandboxes receive a server-side default of ``600``
seconds (10 minutes) when the caller did not set ``idle_ttl_seconds``
explicitly. The launcher stops the sandbox after this many idle
seconds; deletion is anchored to ``stopped_at`` and controlled by
``delete_after_stop_seconds`` (see below).
delete_after_stop_seconds: Seconds after a sandbox enters the
``stopped`` state before it (and its filesystem clone) are
permanently deleted. ``0`` disables stop-anchored deletion;
``None`` falls back to the server default.
stopped_at: Timestamp when the sandbox transitioned to ``stopped``,
or ``None`` while running. The deletion deadline is
``stopped_at + delete_after_stop_seconds``.
snapshot_id: Snapshot ID used to create this sandbox.
vcpus: Number of vCPUs allocated.
mem_bytes: Memory allocation in bytes.
fs_capacity_bytes: Root filesystem capacity in bytes.
Example:
with client.sandbox(snapshot_id="<snapshot-uuid>") as sandbox:
result = sandbox.run("python --version")
print(result.stdout)
"""
# Data fields (from API response)
name: str
dataplane_url: Optional[str] = None
id: Optional[str] = None
status: str = "ready"
status_message: Optional[str] = None
created_at: Optional[str] = None
updated_at: Optional[str] = None
idle_ttl_seconds: Optional[int] = None
delete_after_stop_seconds: Optional[int] = None
stopped_at: Optional[str] = None
snapshot_id: Optional[str] = None
vcpus: Optional[int] = None
mem_bytes: Optional[int] = None
fs_capacity_bytes: Optional[int] = None
# Internal fields (not from API)
_client: SandboxClient = field(repr=False, default=None) # type: ignore
_auto_delete: bool = field(repr=False, default=True)
@classmethod
def from_dict(
cls,
data: dict[str, Any],
client: SandboxClient,
auto_delete: bool = True,
) -> Sandbox:
"""Create a Sandbox from API response dict.
Args:
data: API response dictionary containing sandbox data.
client: Parent SandboxClient for operations.
auto_delete: Whether to delete the sandbox on context exit.
Returns:
Sandbox instance.
"""
return cls(
name=data.get("name", ""),
dataplane_url=data.get("dataplane_url"),
id=data.get("id"),
status=data.get("status", "ready"),
status_message=data.get("status_message"),
created_at=data.get("created_at"),
updated_at=data.get("updated_at"),
idle_ttl_seconds=data.get("idle_ttl_seconds"),
delete_after_stop_seconds=data.get("delete_after_stop_seconds"),
stopped_at=data.get("stopped_at"),
snapshot_id=data.get("snapshot_id"),
vcpus=data.get("vcpus"),
mem_bytes=data.get("mem_bytes"),
fs_capacity_bytes=data.get("fs_capacity_bytes"),
_client=client,
_auto_delete=auto_delete,
)
def __enter__(self) -> Sandbox:
"""Enter context manager."""
return self
def __exit__(
self,
exc_type: Optional[type],
exc_val: Optional[BaseException],
exc_tb: Optional[Any],
) -> None:
"""Exit context manager, optionally deleting the sandbox."""
if self._auto_delete:
try:
self._client.delete_sandbox(self.name)
except Exception:
# Don't raise on cleanup errors
pass
def _require_dataplane_url(self) -> str:
"""Validate and return the dataplane URL.
Returns:
The dataplane URL.
Raises:
SandboxNotReadyError: If sandbox status is not "ready".
DataplaneNotConfiguredError: If dataplane_url is not configured.
"""
if self.status != "ready":
raise SandboxNotReadyError(
f"Sandbox '{self.name}' is not ready (status: {self.status}). "
"Wait for status 'ready' before running operations."
)
if not self.dataplane_url:
raise DataplaneNotConfiguredError(
f"Sandbox '{self.name}' does not have a dataplane_url configured. "
"Runtime operations require a dataplane URL."
)
return self.dataplane_url
@overload
def run(
self,
command: str,
*,
timeout: int = ...,
env: Optional[dict[str, str]] = ...,
cwd: Optional[str] = ...,
shell: str = ...,
on_stdout: Optional[Callable[[str], Any]] = ...,
on_stderr: Optional[Callable[[str], Any]] = ...,
idle_timeout: int = ...,
kill_on_disconnect: bool = ...,
ttl_seconds: int = ...,
pty: bool = ...,
headers: RequestHeaders = ...,
wait: Literal[True] = ...,
) -> ExecutionResult: ...
@overload
def run(
self,
command: str,
*,
timeout: int = ...,
env: Optional[dict[str, str]] = ...,
cwd: Optional[str] = ...,
shell: str = ...,
on_stdout: Optional[Callable[[str], Any]] = ...,
on_stderr: Optional[Callable[[str], Any]] = ...,
idle_timeout: int = ...,
kill_on_disconnect: bool = ...,
ttl_seconds: int = ...,
pty: bool = ...,
headers: RequestHeaders = ...,
wait: Literal[False],
) -> CommandHandle: ...
def run(
self,
command: str,
*,
timeout: int = 60,
env: Optional[dict[str, str]] = None,
cwd: Optional[str] = None,
shell: str = "/bin/bash",
on_stdout: Optional[Callable[[str], Any]] = None,
on_stderr: Optional[Callable[[str], Any]] = None,
idle_timeout: int = 300,
kill_on_disconnect: bool = False,
ttl_seconds: int = 600,
pty: bool = False,
headers: RequestHeaders = None,
wait: bool = True,
) -> Union[ExecutionResult, CommandHandle]:
"""Execute a command in the sandbox.
Args:
command: Shell command to execute.
timeout: Command timeout in seconds.
env: Environment variables to set for the command.
cwd: Working directory for command execution. If None, uses sandbox default.
shell: Shell to use for command execution. Defaults to "/bin/bash".
on_stdout: Callback invoked with each stdout chunk as it arrives.
Blocks until the command completes and returns ExecutionResult.
Cannot be combined with wait=False.
on_stderr: Callback invoked with each stderr chunk as it arrives.
Blocks until the command completes and returns ExecutionResult.
Cannot be combined with wait=False.
idle_timeout: Idle timeout in seconds. If the command has no
connected clients for this duration, it is killed. Defaults
to 300 (5 minutes). Set to -1 for no idle timeout.
Only applies to WebSocket execution.
kill_on_disconnect: If True, kill the command immediately when
the last client disconnects. Defaults to False (command
continues running and can be reconnected to).
ttl_seconds: How long (in seconds) a finished command's session
is kept for reconnection. Defaults to 600 (10 minutes).
Set to -1 to keep indefinitely.
pty: If True, allocate a pseudo-terminal for the command.
Useful for commands that require a TTY (e.g., interactive
programs, commands that use terminal control codes).
Defaults to False.
wait: If True (default), block until the command completes and
return ExecutionResult. If False, return a
CommandHandle immediately for streaming output,
kill, stdin input, and reconnection. Cannot be combined with
on_stdout/on_stderr callbacks.
Returns:
ExecutionResult when wait=True (default).
CommandHandle when wait=False.
Raises:
ValueError: If wait=False is combined with callbacks.
DataplaneNotConfiguredError: If dataplane_url is not configured.
SandboxOperationError: If command execution fails.
CommandTimeoutError: If command exceeds its timeout.
SandboxConnectionError: If connection to sandbox fails after retries.
SandboxNotReadyError: If sandbox is not ready.
SandboxClientError: For other errors.
"""
if not wait and (on_stdout or on_stderr):
raise ValueError(
"Cannot combine wait=False with on_stdout/on_stderr callbacks. "
"Use wait=False and iterate the CommandHandle, or use callbacks."
)
self._require_dataplane_url()
# When not waiting or callbacks are requested, WS is required
use_ws = not wait or on_stdout or on_stderr
if use_ws:
return self._run_ws(
command,
timeout=timeout,
env=env,
cwd=cwd,
shell=shell,
wait=wait,
on_stdout=on_stdout,
on_stderr=on_stderr,
idle_timeout=idle_timeout,
kill_on_disconnect=kill_on_disconnect,
ttl_seconds=ttl_seconds,
pty=pty,
headers=headers,
)
# Default (wait=True, no callbacks): try WS, fall back to HTTP.
# Catch broad exceptions so that unexpected WS failures (e.g. version
# incompatibilities) don't break users who don't need WS features.
try:
return self._run_ws(
command,
timeout=timeout,
env=env,
cwd=cwd,
shell=shell,
wait=True,
on_stdout=None,
on_stderr=None,
idle_timeout=idle_timeout,
kill_on_disconnect=kill_on_disconnect,
ttl_seconds=ttl_seconds,
pty=pty,
headers=headers,
)
except (SandboxConnectionError, ImportError, OSError, TypeError):
return self._run_http(
command,
timeout=timeout,
env=env,
cwd=cwd,
shell=shell,
headers=headers,
)
def _run_ws(
self,
command: str,
*,
timeout: int,
env: Optional[dict[str, str]],
cwd: Optional[str],
shell: str,
wait: bool,
on_stdout: Optional[Callable[[str], Any]],
on_stderr: Optional[Callable[[str], Any]],
idle_timeout: int = 300,
kill_on_disconnect: bool = False,
ttl_seconds: int = 600,
pty: bool = False,
headers: RequestHeaders = None,
) -> Union[ExecutionResult, CommandHandle]:
"""Execute via WebSocket /execute/ws."""
from langsmith.sandbox._ws_execute import run_ws_stream
dataplane_url = self._require_dataplane_url()
api_key = self._client._api_key
ws_kwargs: dict[str, Any] = {
"timeout": timeout,
"env": env,
"cwd": cwd,
"shell": shell,
"on_stdout": on_stdout,
"on_stderr": on_stderr,
"idle_timeout": idle_timeout,
"kill_on_disconnect": kill_on_disconnect,
"ttl_seconds": ttl_seconds,
"pty": pty,
}
if headers is not None:
ws_kwargs["headers"] = headers
msg_stream, control = run_ws_stream(
dataplane_url,
api_key,
command,
**ws_kwargs,
)
handle = CommandHandle(msg_stream, control, self)
if not wait:
return handle
return handle.result # blocks until command completes
def _run_http(
self,
command: str,
*,
timeout: int,
env: Optional[dict[str, str]],
cwd: Optional[str],
shell: str,
headers: RequestHeaders,
) -> ExecutionResult:
"""Execute via HTTP POST /execute (existing implementation)."""
dataplane_url = self._require_dataplane_url()
url = f"{dataplane_url}/execute"
payload: dict[str, Any] = {
"command": command,
"timeout": timeout,
"shell": shell,
}
if env is not None:
payload["env"] = env
if cwd is not None:
payload["cwd"] = cwd
try:
response = self._client._http.post(
url,
json=payload,
timeout=timeout + 10,
headers=self._client._request_headers(headers),
)
response.raise_for_status()
data = response.json()
return ExecutionResult(
stdout=data.get("stdout", ""),
stderr=data.get("stderr", ""),
exit_code=data.get("exit_code", -1),
)
except httpx.HTTPStatusError as e:
handle_sandbox_http_error(e)
raise # pragma: no cover
def reconnect(
self,
command_id: str,
*,
stdout_offset: int = 0,
stderr_offset: int = 0,
headers: RequestHeaders = None,
) -> CommandHandle:
"""Reconnect to a running or recently-finished command.
Resumes output from the given byte offsets. Any output produced while
the client was disconnected is replayed from the server's ring buffer.
Args:
command_id: The command ID from handle.command_id.
stdout_offset: Byte offset to resume stdout from (default: 0).
stderr_offset: Byte offset to resume stderr from (default: 0).
Returns:
A CommandHandle for the command.
Raises:
SandboxOperationError: If command_id is not found or session expired.
SandboxConnectionError: If connection to sandbox fails after retries.
"""
from langsmith.sandbox._ws_execute import reconnect_ws_stream
dataplane_url = self._require_dataplane_url()
api_key = self._client._api_key
reconnect_kwargs: dict[str, Any] = {
"stdout_offset": stdout_offset,
"stderr_offset": stderr_offset,
}
if headers is not None:
reconnect_kwargs["headers"] = headers
msg_stream, control = reconnect_ws_stream(
dataplane_url,
api_key,
command_id,
**reconnect_kwargs,
)
return CommandHandle(
msg_stream,
control,
self,
command_id=command_id,
stdout_offset=stdout_offset,
stderr_offset=stderr_offset,
)
def write(
self,
path: str,
content: Union[str, bytes],
*,
timeout: int = 60,
headers: RequestHeaders = None,
) -> None:
"""Write content to a file in the sandbox.
Args:
path: Target file path in the sandbox.
content: File content (str or bytes).
timeout: Request timeout in seconds.
Raises:
DataplaneNotConfiguredError: If dataplane_url is not configured.
SandboxOperationError: If file write fails.
SandboxConnectionError: If connection to sandbox fails after retries.
SandboxNotReadyError: If sandbox is not ready.
SandboxClientError: For other errors.
"""
dataplane_url = self._require_dataplane_url()
url = f"{dataplane_url}/upload"
# Ensure content is bytes for multipart upload
if isinstance(content, str):
content = content.encode("utf-8")
files = {"file": ("file", content)}
try:
response = self._client._http.post(
url,
params={"path": path},
files=files,
timeout=timeout,
headers=self._client._request_headers(headers),
)
response.raise_for_status()
except httpx.HTTPStatusError as e:
handle_sandbox_http_error(e)
def read(
self, path: str, *, timeout: int = 60, headers: RequestHeaders = None
) -> bytes:
"""Read a file from the sandbox.
Args:
path: File path to read. Supports both absolute paths (e.g., /tmp/file.txt)
and relative paths (resolved from /home/user/).
timeout: Request timeout in seconds.
Returns:
File contents as bytes.
Raises:
DataplaneNotConfiguredError: If dataplane_url is not configured.
ResourceNotFoundError: If the file doesn't exist.
SandboxOperationError: If file read fails.
SandboxConnectionError: If connection to sandbox fails after retries.
SandboxNotReadyError: If sandbox is not ready.
SandboxClientError: For other errors.
"""
dataplane_url = self._require_dataplane_url()
url = f"{dataplane_url}/download"
try:
response = self._client._http.get(
url,
params={"path": path},
timeout=timeout,
headers=self._client._request_headers(headers),
)
response.raise_for_status()
return response.content
except httpx.HTTPStatusError as e:
if e.response.status_code == 404:
raise ResourceNotFoundError(
f"File '{path}' not found in sandbox '{self.name}'",
resource_type="file",
) from e
handle_sandbox_http_error(e)
# This line should never be reached but satisfies type checker
raise # pragma: no cover
def tunnel(
self,
remote_port: int,
*,
local_port: int = 0,
max_reconnects: int = 3,
headers: RequestHeaders = None,
) -> Tunnel:
"""Open a TCP tunnel to a port inside the sandbox.
Creates a local TCP listener that forwards connections through a
yamux-multiplexed WebSocket to the specified port inside the sandbox.
Works with any TCP protocol (databases, Redis, HTTP, etc.).
Use as a context manager for automatic cleanup::
with sandbox.tunnel(remote_port=5432) as t:
conn = psycopg2.connect(host="127.0.0.1", port=t.local_port)
Or manage the lifecycle explicitly::
t = sandbox.tunnel(remote_port=5432)
# ... use tunnel ...
t.close()
Args:
remote_port: TCP port inside the sandbox to tunnel to (1-65535).
local_port: Local port to listen on. Defaults to mirroring
remote_port. Use 0 to let the OS pick an available port.
max_reconnects: Maximum number of automatic reconnect attempts
when the WebSocket session drops. Set to 0 to disable.
Returns:
A Tunnel instance (context manager).
Raises:
ValueError: If port values are out of range.
DataplaneNotConfiguredError: If dataplane_url is not configured.
SandboxNotReadyError: If sandbox is not ready.
"""
if not 1 <= remote_port <= 65535:
raise ValueError(
f"remote_port must be between 1 and 65535 (got {remote_port})"
)
if local_port and not 1 <= local_port <= 65535:
raise ValueError(
f"local_port must be between 1 and 65535 (got {local_port})"
)
dataplane_url = self._require_dataplane_url()
api_key = self._client._api_key
t = Tunnel(
dataplane_url,
api_key,
remote_port,
local_port=local_port,
max_reconnects=max_reconnects,
headers=headers,
)
t._start()
return t
def service(
self,
port: int,
*,
expires_in_seconds: int = 600,
headers: RequestHeaders = None,
) -> ServiceURL:
"""Get an authenticated URL for a service running in this sandbox.
Returns a :class:`ServiceURL` whose properties auto-refresh the
token transparently before it expires.
Args:
port: Port the service is listening on inside the sandbox.
expires_in_seconds: Token TTL in seconds (1--86400, default 600).
headers: Optional per-request header overrides.
Returns:
ServiceURL with auto-refreshing token and HTTP helpers.
Raises:
ResourceNotFoundError: If sandbox not found.
ValueError: If port or expires_in_seconds is out of range.
SandboxClientError: For other errors.
"""
return self._client.service(
self.name,
port,
expires_in_seconds=expires_in_seconds,
headers=headers,
)
def start(
self,
*,
timeout: int = 120,
headers: RequestHeaders = None,
) -> None:
"""Start a stopped sandbox and wait until ready.
After starting, the sandbox's status and dataplane_url are updated
in place.
Args:
timeout: Timeout in seconds when waiting for ready.
headers: Optional per-request header overrides.
Raises:
ResourceNotFoundError: If sandbox not found.
ResourceCreationError: If sandbox fails during startup.
ResourceTimeoutError: If sandbox doesn't become ready within timeout.
SandboxClientError: For other errors.
"""
refreshed = self._client.start_sandbox(
self.name, timeout=timeout, headers=headers
)
self.status = refreshed.status
self.dataplane_url = refreshed.dataplane_url
def stop(self, *, headers: RequestHeaders = None) -> None:
"""Stop a running sandbox (preserves sandbox files for later restart).
Args:
headers: Optional per-request header overrides.
Raises:
ResourceNotFoundError: If sandbox not found.
SandboxClientError: For other errors.
"""
self._client.stop_sandbox(self.name, headers=headers)
self.status = "stopped"
self.dataplane_url = None
def delete(self, *, headers: RequestHeaders = None) -> None:
"""Delete this sandbox.
Args:
headers: Optional per-request header overrides.
Raises:
ResourceNotFoundError: If sandbox not found.
SandboxClientError: For other errors.
"""
self._client.delete_sandbox(self.name, headers=headers)
def capture_snapshot(
self,
name: str,
*,
timeout: int = 60,
headers: RequestHeaders = None,
) -> Snapshot:
"""Capture a snapshot from this sandbox.
Args:
name: Snapshot name.
timeout: Timeout in seconds when waiting for ready.
headers: Optional per-request header overrides.
Returns:
Snapshot in "ready" status.
Raises:
ResourceNotFoundError: If sandbox not found.
ResourceTimeoutError: If snapshot doesn't become ready within timeout.
ResourceCreationError: If snapshot capture fails.
SandboxClientError: For other errors.
"""
return self._client.capture_snapshot(
self.name,
name,
timeout=timeout,
headers=headers,
)