From 5c357756da6b7303a3f19b7f63540c35737ec671 Mon Sep 17 00:00:00 2001 From: Tobias Herber <22559657+herber@users.noreply.github.com> Date: Fri, 16 Oct 2026 13:23:20 +0200 Subject: [PATCH] Add machine connection apis --- src/metorial/__init__.py | 25 ++ src/metorial/_magnetar_sdk.py | 76 +++++++ src/metorial/sandbox/__init__.py | 28 +++ src/metorial/sandbox/_access.py | 52 +++++ src/metorial/sandbox/_errors.py | 90 ++++++++ src/metorial/sandbox/_mux.py | 305 +++++++++++++++++++++++++ src/metorial/sandbox/files.py | 379 +++++++++++++++++++++++++++++++ src/metorial/sandbox/port.py | 117 ++++++++++ src/metorial/sandbox/shell.py | 305 +++++++++++++++++++++++++ tests/test_sandbox.py | 375 ++++++++++++++++++++++++++++++ 10 files changed, 1752 insertions(+) create mode 100644 src/metorial/sandbox/__init__.py create mode 100644 src/metorial/sandbox/_access.py create mode 100644 src/metorial/sandbox/_errors.py create mode 100644 src/metorial/sandbox/_mux.py create mode 100644 src/metorial/sandbox/files.py create mode 100644 src/metorial/sandbox/port.py create mode 100644 src/metorial/sandbox/shell.py create mode 100644 tests/test_sandbox.py diff --git a/src/metorial/__init__.py b/src/metorial/__init__.py index 29de7d03..324d9f64 100644 --- a/src/metorial/__init__.py +++ b/src/metorial/__init__.py @@ -113,6 +113,19 @@ build_openai_tools, call_openai_tools, ) +from .sandbox import ( + SandboxAccessError, + SandboxCommandError, + SandboxCommandResult, + SandboxConnectionNotice, + SandboxFileEntry, + SandboxFileList, + SandboxFiles, + SandboxPort, + SandboxProcess, + SandboxProcessExit, + SandboxShell, +) def _configure_sdk_logging() -> None: @@ -215,6 +228,18 @@ def _configure_sdk_logging() -> None: "MagnetarMagicMcpGroup", "MagnetarPortalsGroup", "create_magnetar_sdk", + # Sandbox access + "SandboxAccessError", + "SandboxCommandError", + "SandboxCommandResult", + "SandboxConnectionNotice", + "SandboxFileEntry", + "SandboxFileList", + "SandboxFiles", + "SandboxPort", + "SandboxProcess", + "SandboxProcessExit", + "SandboxShell", # MCP "MetorialMagnetarMcpSession", "MagnetarMcpSessionInit", diff --git a/src/metorial/_magnetar_sdk.py b/src/metorial/_magnetar_sdk.py index 06aff8f7..300712d6 100644 --- a/src/metorial/_magnetar_sdk.py +++ b/src/metorial/_magnetar_sdk.py @@ -10,9 +10,13 @@ from __future__ import annotations +import asyncio +from collections.abc import Callable from dataclasses import dataclass from typing import Any +import httpx + from metorial._endpoint import MetorialEndpointManager from metorial._generated.magnetar.endpoints.callback_events import ( MetorialCallbackEventsEndpoint, @@ -334,6 +338,15 @@ ) from metorial._sdk_builder import MetorialSDKBuilder from metorial._sdk_shared import SDKConfig, get_api_host, get_headers +from metorial.sandbox import ( + SandboxAccessCache, + SandboxAccessError, + SandboxAccessGrant, + SandboxConnectionNotice, + SandboxFiles, + SandboxPort, + SandboxShell, +) # ── Nested endpoint groups ──────────────────────────────────────────────────── @@ -489,6 +502,69 @@ def __init__(self, manager: MetorialEndpointManager) -> None: self.errors = MetorialSandboxMachinesErrorsEndpoint(manager) self.usage_history = MetorialSandboxMachinesUsageEndpoint(manager) + def _access_cache( + self, kind: str, machine_id: str, ttl_seconds: float | None, port: int = 0 + ) -> SandboxAccessCache: + async def provider() -> SandboxAccessGrant: + if kind == "port": + grant: Any = await asyncio.to_thread( + self.access.ports, machine_id, ports=[port], ttl_seconds=ttl_seconds + ) + url = (grant.urls.ports or {}).get(str(port)) + label = f"Port {port}" + elif kind == "files": + grant = await asyncio.to_thread( + self.access.files, machine_id, ttl_seconds=ttl_seconds + ) + url, label = grant.urls.files, "File" + else: + grant = await asyncio.to_thread( + self.access.shell, machine_id, ttl_seconds=ttl_seconds + ) + url, label = grant.urls.exec_, "Command" + + if not url: + raise SandboxAccessError( + "unavailable", f"{label} access is not available for this machine." + ) + return SandboxAccessGrant(url=url, token=grant.token, expires_at=grant.expires_at) + + return SandboxAccessCache(provider) + + async def files( + self, + machine_id: str, + *, + ttl_seconds: float | None = None, + on_notice: Callable[[SandboxConnectionNotice], None] | None = None, + ) -> SandboxFiles: + return await SandboxFiles.connect( + self._access_cache("files", machine_id, ttl_seconds), on_notice=on_notice + ) + + async def shell( + self, + machine_id: str, + *, + ttl_seconds: float | None = None, + on_notice: Callable[[SandboxConnectionNotice], None] | None = None, + ) -> SandboxShell: + return await SandboxShell.connect( + self._access_cache("shell", machine_id, ttl_seconds), on_notice=on_notice + ) + + async def connect( + self, + machine_id: str, + port: int, + *, + ttl_seconds: float | None = None, + client: httpx.AsyncClient | None = None, + ) -> SandboxPort: + return await SandboxPort.connect( + port, self._access_cache("port", machine_id, ttl_seconds, port), client=client + ) + class MagnetarSandboxImagesGroup(MetorialSandboxImagesEndpoint): """Sandbox images with versions/builds sub-endpoints.""" diff --git a/src/metorial/sandbox/__init__.py b/src/metorial/sandbox/__init__.py new file mode 100644 index 00000000..75754807 --- /dev/null +++ b/src/metorial/sandbox/__init__.py @@ -0,0 +1,28 @@ +from ._access import SandboxAccessCache, SandboxAccessGrant, SandboxAccessProvider +from ._errors import SandboxAccessError, SandboxCommandError +from ._mux import SandboxConnectionNotice +from .files import SandboxFileEntry, SandboxFileList, SandboxFiles +from .port import SandboxPort +from .shell import ( + SandboxCommandResult, + SandboxProcess, + SandboxProcessExit, + SandboxShell, +) + +__all__ = [ + "SandboxAccessCache", + "SandboxAccessError", + "SandboxAccessGrant", + "SandboxAccessProvider", + "SandboxCommandError", + "SandboxCommandResult", + "SandboxConnectionNotice", + "SandboxFileEntry", + "SandboxFileList", + "SandboxFiles", + "SandboxPort", + "SandboxProcess", + "SandboxProcessExit", + "SandboxShell", +] diff --git a/src/metorial/sandbox/_access.py b/src/metorial/sandbox/_access.py new file mode 100644 index 00000000..8ddd0d2e --- /dev/null +++ b/src/metorial/sandbox/_access.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import datetime, timezone + +_EXPIRY_MARGIN_SECONDS = 30 + + +@dataclass +class SandboxAccessGrant: + url: str + token: str + expires_at: datetime | None = None + + +SandboxAccessProvider = Callable[[], Awaitable[SandboxAccessGrant]] + + +class SandboxAccessCache: + def __init__(self, provider: SandboxAccessProvider) -> None: + self._provider = provider + self._grant: SandboxAccessGrant | None = None + self._lock = asyncio.Lock() + + def _is_usable(self, grant: SandboxAccessGrant) -> bool: + if grant.expires_at is None: + return True + + expires_at = grant.expires_at + if expires_at.tzinfo is None: + expires_at = expires_at.replace(tzinfo=timezone.utc) + + remaining = (expires_at - datetime.now(timezone.utc)).total_seconds() + return remaining > _EXPIRY_MARGIN_SECONDS + + async def get(self) -> SandboxAccessGrant: + grant = self._grant + if grant is not None and self._is_usable(grant): + return grant + + async with self._lock: + grant = self._grant + if grant is not None and self._is_usable(grant): + return grant + + self._grant = await self._provider() + return self._grant + + def invalidate(self) -> None: + self._grant = None diff --git a/src/metorial/sandbox/_errors.py b/src/metorial/sandbox/_errors.py new file mode 100644 index 00000000..b4717ddb --- /dev/null +++ b/src/metorial/sandbox/_errors.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from typing import Any + + +class SandboxAccessError(Exception): + def __init__( + self, + code: str, + message: str | None = None, + *, + retryable: bool = False, + path: str | None = None, + ) -> None: + super().__init__(message or code.replace("_", " ")) + self.code = code + self.retryable = retryable + self.path = path + + +class SandboxCommandError(SandboxAccessError): + pass + + +def closed_error() -> SandboxAccessError: + return SandboxAccessError("closed", "The sandbox connection was closed.") + + +def server_error( + raw: Any, error_class: type[SandboxAccessError] = SandboxAccessError +) -> SandboxAccessError: + data = raw if isinstance(raw, dict) else {} + code = data.get("code") + message = data.get("message") + path = data.get("path") + + return error_class( + code if isinstance(code, str) else "unknown", + message if isinstance(message, str) else None, + retryable=data.get("retryable") is True, + path=path if isinstance(path, str) else None, + ) + + +_CLOSE_REASONS = { + "machine_paused": "The machine paused, so this connection was closed.", + "machine_restarting": "The machine is restarting.", + "machine_terminated": "The machine was terminated.", + "idle_timeout": "This connection was closed after being idle.", +} + + +def describe_socket_close(code: int | None, reason: str) -> SandboxAccessError: + if code == 4401: + return SandboxAccessError( + "unauthorized", "Access to the machine was rejected.", retryable=True + ) + if code == 4404: + return SandboxAccessError("not_found", "This machine could not be found.") + if code == 4410: + return SandboxAccessError( + "machine_terminated", _CLOSE_REASONS["machine_terminated"] + ) + if code == 4429: + return SandboxAccessError( + "rate_limited", "Too many connections to the machine.", retryable=True + ) + if code == 4503: + if reason == "machine_failed": + return SandboxAccessError("machine_failed", "The machine has failed.") + return SandboxAccessError( + "machine_unavailable", "The machine is not available.", retryable=True + ) + if code == 1012: + kind = ( + reason + if reason in ("machine_terminated", "machine_restarting") + else "machine_paused" + ) + return SandboxAccessError( + kind, _CLOSE_REASONS[kind], retryable=kind != "machine_terminated" + ) + if code == 1000 and reason == "idle_timeout": + return SandboxAccessError( + "idle_timeout", _CLOSE_REASONS["idle_timeout"], retryable=True + ) + + return SandboxAccessError( + "disconnected", "The connection to the machine was lost.", retryable=True + ) diff --git a/src/metorial/sandbox/_mux.py b/src/metorial/sandbox/_mux.py new file mode 100644 index 00000000..ebacff14 --- /dev/null +++ b/src/metorial/sandbox/_mux.py @@ -0,0 +1,305 @@ +from __future__ import annotations + +import asyncio +import contextlib +import itertools +import json +import secrets +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import aiohttp + +from ._access import SandboxAccessCache +from ._errors import SandboxAccessError, closed_error, describe_socket_close + +_MAX_CONNECT_ATTEMPTS = 3 +_CONNECT_BASE_DELAY = 0.3 +_MAX_RECONNECT_DELAY = 10.0 +_MAX_MESSAGE_SIZE = 16 * 1024 * 1024 + +_RECONNECT_CODES = frozenset( + { + "disconnected", + "unauthorized", + "rate_limited", + "machine_restarting", + "machine_unavailable", + } +) + +_id_counter = itertools.count(1) +_id_prefix = secrets.token_hex(3) + + +def create_request_id() -> str: + return f"{_id_prefix}_{next(_id_counter):x}" + + +def socket_protocols(token: str) -> list[str]: + return ["metorial-sandbox-v1", f"metorial-sandbox-token.{token}"] + + +def encode_frame(request_id: str, payload: bytes, tag: int | None = None) -> bytes: + id_bytes = request_id.encode() + header = bytes([len(id_bytes)]) + id_bytes + if tag is not None: + header += bytes([tag]) + return header + payload + + +def decode_frame( + data: bytes, has_tag: bool = False +) -> tuple[str, int | None, bytes] | None: + if not data: + return None + + id_length = data[0] + offset = 1 + id_length + if len(data) < offset + (1 if has_tag else 0): + return None + + request_id = data[1:offset].decode(errors="replace") + tag: int | None = None + if has_tag: + tag = data[offset] + offset += 1 + + return request_id, tag, data[offset:] + + +@dataclass +class SandboxConnectionNotice: + notice: str + reason: str | None = None + idle_in_seconds: float | None = None + + +NoticeHandler = Callable[[SandboxConnectionNotice], None] + + +def _handshake_error(error: aiohttp.WSServerHandshakeError) -> SandboxAccessError: + if error.status in (401, 403): + return SandboxAccessError( + "unauthorized", "Access to the machine was rejected.", retryable=True + ) + if error.status == 404: + return SandboxAccessError("not_found", "This machine could not be found.") + if error.status == 429: + return SandboxAccessError( + "rate_limited", "Too many connections to the machine.", retryable=True + ) + return SandboxAccessError( + "disconnected", f"The connection failed ({error.status}).", retryable=True + ) + + +class MuxSocket: + def __init__( + self, + access: SandboxAccessCache, + *, + on_text: Callable[[dict[str, Any]], None], + on_binary: Callable[[bytes], None], + on_disconnect: Callable[[SandboxAccessError], None], + on_notice: NoticeHandler | None = None, + ) -> None: + self._access = access + self._on_text = on_text + self._on_binary = on_binary + self._on_disconnect = on_disconnect + self._on_notice = on_notice + + self._session: aiohttp.ClientSession | None = None + self._ws: aiohttp.ClientWebSocketResponse[bool] | None = None + self._opening: asyncio.Task[aiohttp.ClientWebSocketResponse[bool]] | None = None + self._reader: asyncio.Task[None] | None = None + self._reconnect: asyncio.Task[None] | None = None + self._reconnect_attempt = 0 + self._closed = False + + @property + def is_closed(self) -> bool: + return self._closed + + @property + def is_open(self) -> bool: + return self._ws is not None and not self._ws.closed + + async def ensure_open(self) -> aiohttp.ClientWebSocketResponse[bool]: + if self._closed: + raise closed_error() + if self._ws is not None and not self._ws.closed: + return self._ws + + if self._opening is None or self._opening.done(): + self._opening = asyncio.ensure_future(self._connect()) + + return await asyncio.shield(self._opening) + + async def send(self, data: str | bytes) -> bool: + ws = self._ws + if ws is None or ws.closed: + return False + + try: + if isinstance(data, str): + await ws.send_str(data) + else: + await ws.send_bytes(data) + return True + except (ConnectionError, RuntimeError, aiohttp.ClientError): + return False + + async def close(self) -> None: + if self._closed: + return + self._closed = True + + for task in (self._reconnect, self._opening): + if task is not None and not task.done(): + task.cancel() + + ws, self._ws = self._ws, None + if ws is not None: + with contextlib.suppress(Exception): + await ws.close(code=1000, message=b"closed") + + if self._reader is not None and not self._reader.done(): + self._reader.cancel() + + if self._session is not None: + await self._session.close() + self._session = None + + self._on_disconnect(closed_error()) + + async def _connect(self) -> aiohttp.ClientWebSocketResponse[bool]: + last_error: SandboxAccessError | None = None + + for attempt in range(1, _MAX_CONNECT_ATTEMPTS + 1): + if attempt > 1: + await asyncio.sleep(_CONNECT_BASE_DELAY * 2 ** (attempt - 2)) + if self._closed: + raise closed_error() + + grant = await self._access.get() + if self._session is None or self._session.closed: + self._session = aiohttp.ClientSession() + + try: + ws = await self._session.ws_connect( + grant.url, + protocols=socket_protocols(grant.token), + max_msg_size=_MAX_MESSAGE_SIZE, + ) + except aiohttp.WSServerHandshakeError as error: + last_error = _handshake_error(error) + except (aiohttp.ClientError, OSError, asyncio.TimeoutError) as error: + last_error = SandboxAccessError( + "disconnected", str(error) or None, retryable=True + ) + else: + if self.is_closed: + await ws.close(code=1000, message=b"closed") + raise closed_error() + + self._ws = ws + self._reader = asyncio.ensure_future(self._read(ws)) + return ws + + if not last_error.retryable: + raise last_error + self._access.invalidate() + + assert last_error is not None + raise last_error + + async def _read(self, ws: aiohttp.ClientWebSocketResponse[bool]) -> None: + close_code: int | None = None + close_reason = "" + + try: + while True: + message = await ws.receive() + + if message.type == aiohttp.WSMsgType.TEXT: + self._handle_text(message.data) + elif message.type == aiohttp.WSMsgType.BINARY: + self._on_binary(message.data) + elif message.type == aiohttp.WSMsgType.CLOSE: + close_code = message.data + close_reason = message.extra or "" + break + elif message.type in ( + aiohttp.WSMsgType.CLOSING, + aiohttp.WSMsgType.CLOSED, + aiohttp.WSMsgType.ERROR, + ): + break + except asyncio.CancelledError: + return + except Exception: + pass + + if close_code is None: + close_code = ws.close_code + + if self._closed or self._ws is not ws: + return + + error = describe_socket_close(close_code, close_reason) + if error.code == "unauthorized": + self._access.invalidate() + + self._ws = None + self._on_disconnect(error) + + if error.code in _RECONNECT_CODES: + self._schedule_reconnect() + + def _schedule_reconnect(self) -> None: + if self._closed or (self._reconnect is not None and not self._reconnect.done()): + return + + delay = min(_CONNECT_BASE_DELAY * 2**self._reconnect_attempt, _MAX_RECONNECT_DELAY) + self._reconnect_attempt += 1 + self._reconnect = asyncio.ensure_future(self._reconnect_after(delay)) + + async def _reconnect_after(self, delay: float) -> None: + await asyncio.sleep(delay) + + try: + await self.ensure_open() + self._reconnect_attempt = 0 + except SandboxAccessError as error: + if error.code in _RECONNECT_CODES: + self._reconnect = None + self._schedule_reconnect() + except Exception: + pass + + def _handle_text(self, data: str) -> None: + try: + message = json.loads(data) + except ValueError: + return + + if not isinstance(message, dict): + return + + if message.get("type") == "connection_notice": + if self._on_notice is not None: + reason = message.get("reason") + idle = message.get("idleInSeconds") + self._on_notice( + SandboxConnectionNotice( + notice=str(message.get("notice") or ""), + reason=reason if isinstance(reason, str) else None, + idle_in_seconds=idle if isinstance(idle, (int, float)) else None, + ) + ) + return + + self._on_text(message) diff --git a/src/metorial/sandbox/files.py b/src/metorial/sandbox/files.py new file mode 100644 index 00000000..9ff44e00 --- /dev/null +++ b/src/metorial/sandbox/files.py @@ -0,0 +1,379 @@ +from __future__ import annotations + +import asyncio +import hashlib +import json +import weakref +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from datetime import datetime, timezone +from types import TracebackType +from typing import Any, TypeVar + +from ._access import SandboxAccessCache +from ._errors import SandboxAccessError, server_error +from ._mux import ( + MuxSocket, + NoticeHandler, + create_request_id, + decode_frame, + encode_frame, +) + +T = TypeVar("T") +_EntriesHandler = Callable[[list[Any]], None] + +_CHUNK_SIZE = 64 * 1024 +_DEFAULT_LIST_LIMIT = 10000 +_MAX_ATTEMPTS = 3 +_RESEND_DELAY = 0.2 + +_unsent_errors: weakref.WeakSet[SandboxAccessError] = weakref.WeakSet() + + +@dataclass +class SandboxFileEntry: + name: str + path: str + type: str + size: int + mode: int + modified_at: datetime + target: str | None = None + + +@dataclass +class SandboxFileList: + entries: list[SandboxFileEntry] + truncated: bool + + +@dataclass +class _PendingOperation: + future: asyncio.Future[dict[str, Any]] + on_entries: _EntriesHandler | None = None + on_data: Callable[[bytes], None] | None = None + chunks: list[bytes] = field(default_factory=list) + + +def _join_path(parent: str, child: str) -> str: + return f"{parent.rstrip('/')}/{child.lstrip('/')}" + + +def _base_name(path: str) -> str: + return path.rstrip("/").split("/")[-1] or "/" + + +def _to_entry(raw: dict[str, Any], parent: str | None = None) -> SandboxFileEntry: + name = raw.get("name") + if parent is None: + path = str(raw.get("path")) + else: + path = _join_path(parent, str(name if isinstance(name, str) else raw.get("path"))) + + target = raw.get("target") + return SandboxFileEntry( + name=name if isinstance(name, str) else _base_name(path), + path=path, + type=str(raw.get("type")), + size=int(raw.get("size") or 0), + mode=int(raw.get("mode") or 0), + modified_at=datetime.fromtimestamp( + int(raw.get("mtime") or 0) / 1000, tz=timezone.utc + ), + target=target if isinstance(target, str) else None, + ) + + +class SandboxFiles: + def __init__( + self, access: SandboxAccessCache, *, on_notice: NoticeHandler | None = None + ) -> None: + self._pending: dict[str, _PendingOperation] = {} + self._socket = MuxSocket( + access, + on_text=self._handle_text, + on_binary=self._handle_binary, + on_disconnect=self._handle_disconnect, + on_notice=on_notice, + ) + + @classmethod + async def connect( + cls, access: SandboxAccessCache, *, on_notice: NoticeHandler | None = None + ) -> SandboxFiles: + files = cls(access, on_notice=on_notice) + await files._socket.ensure_open() + return files + + @property + def is_closed(self) -> bool: + return self._socket.is_closed + + async def __aenter__(self) -> SandboxFiles: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + await self.close() + + async def list( + self, + path: str, + *, + recursive: bool = False, + max_depth: int | None = None, + limit: int | None = None, + ) -> SandboxFileList: + async def run() -> SandboxFileList: + entries: list[SandboxFileEntry] = [] + fields: dict[str, Any] = { + "path": path, + "limit": limit if limit is not None else _DEFAULT_LIST_LIMIT, + } + if recursive: + fields["recursive"] = True + if max_depth is not None: + fields["maxDepth"] = max_depth + + result = await self._run( + "list", + fields, + on_entries=lambda batch: entries.extend(_to_entry(raw, path) for raw in batch), + ) + return SandboxFileList(entries=entries, truncated=result.get("truncated") is True) + + return await self._with_retry(run) + + async def stat(self, path: str, *, follow_symlinks: bool = False) -> SandboxFileEntry: + async def run() -> SandboxFileEntry: + result = await self._run( + "stat", {"path": path, "followSymlinks": follow_symlinks} + ) + return _to_entry(result.get("entry") or {}) + + return await self._with_retry(run) + + async def exists(self, path: str) -> bool: + try: + await self.stat(path) + return True + except SandboxAccessError as error: + if error.code == "not_found": + return False + raise + + async def read( + self, path: str, *, offset: int | None = None, length: int | None = None + ) -> bytes: + async def run() -> bytes: + chunks: list[bytes] = [] + fields: dict[str, Any] = {"path": path} + if offset is not None: + fields["offset"] = offset + if length is not None: + fields["length"] = length + + await self._run("read", fields, on_data=chunks.append) + return b"".join(chunks) + + return await self._with_retry(run) + + async def read_text( + self, + path: str, + *, + encoding: str = "utf-8", + offset: int | None = None, + length: int | None = None, + ) -> str: + return (await self.read(path, offset=offset, length=length)).decode(encoding) + + async def write( + self, + path: str, + data: str | bytes, + *, + overwrite: bool = True, + parents: bool = False, + mode: int | str | None = None, + ) -> int: + content = data.encode() if isinstance(data, str) else bytes(data) + checksum = hashlib.sha256(content).hexdigest() + + fields: dict[str, Any] = { + "path": path, + "size": len(content), + "overwrite": overwrite, + } + if parents: + fields["parents"] = True + if mode is not None: + fields["mode"] = mode + + async def upload(request_id: str, is_active: Callable[[], bool]) -> None: + for start in range(0, len(content), _CHUNK_SIZE): + if not is_active(): + return + await self._socket.send( + encode_frame(request_id, content[start : start + _CHUNK_SIZE]) + ) + + if is_active(): + await self._socket.send( + json.dumps({"id": request_id, "type": "end", "sha256": checksum}) + ) + + async def run() -> int: + result = await self._run("write", fields, upload=upload) + return int(result.get("size") or len(content)) + + return await self._with_retry(run, can_resend=overwrite) + + async def mkdir( + self, path: str, *, parents: bool = False, mode: int | str | None = None + ) -> None: + fields: dict[str, Any] = {"path": path, "parents": parents} + if mode is not None: + fields["mode"] = mode + + await self._with_retry(lambda: self._run("mkdir", fields), can_resend=parents) + + async def move( + self, source: str, destination: str, *, overwrite: bool = False + ) -> None: + await self._with_retry( + lambda: self._run( + "move", {"from": source, "to": destination, "overwrite": overwrite} + ), + can_resend=False, + ) + + async def delete( + self, path: str, *, recursive: bool = False, missing_ok: bool = False + ) -> None: + fields: dict[str, Any] = {"path": path, "recursive": recursive} + if missing_ok: + fields["missingOk"] = True + + await self._with_retry(lambda: self._run("delete", fields), can_resend=missing_ok) + + async def close(self) -> None: + await self._socket.close() + + def _handle_text(self, message: dict[str, Any]) -> None: + request_id = message.get("id") + if not isinstance(request_id, str): + return + + operation = self._pending.get(request_id) + if operation is None: + return + + kind = message.get("type") + if kind == "entries" and operation.on_entries is not None: + entries = message.get("entries") + operation.on_entries(entries if isinstance(entries, list) else []) + elif kind == "result": + fields = {k: v for k, v in message.items() if k not in ("id", "type")} + self._settle(request_id, result=fields) + elif kind == "error": + self._settle(request_id, error=server_error(message.get("error"))) + + def _handle_binary(self, data: bytes) -> None: + frame = decode_frame(data) + if frame is None: + return + + operation = self._pending.get(frame[0]) + if operation is not None and operation.on_data is not None: + operation.on_data(frame[2]) + + def _handle_disconnect(self, error: SandboxAccessError) -> None: + for request_id in list(self._pending): + self._settle(request_id, error=error) + + def _settle( + self, + request_id: str, + *, + result: dict[str, Any] | None = None, + error: SandboxAccessError | None = None, + ) -> None: + operation = self._pending.pop(request_id, None) + if operation is None or operation.future.done(): + return + + if error is not None: + operation.future.set_exception(error) + else: + operation.future.set_result(result or {}) + + async def _with_retry( + self, run: Callable[[], Awaitable[T]], *, can_resend: bool = True + ) -> T: + attempt = 1 + while True: + try: + return await run() + except SandboxAccessError as error: + should_resend = ( + (can_resend or error in _unsent_errors) + and attempt < _MAX_ATTEMPTS + and not self.is_closed + and error.retryable + ) + if not should_resend: + raise + + await asyncio.sleep(_RESEND_DELAY * attempt) + attempt += 1 + + async def _run( + self, + op: str, + fields: dict[str, Any], + *, + on_entries: _EntriesHandler | None = None, + on_data: Callable[[bytes], None] | None = None, + upload: Callable[[str, Callable[[], bool]], Awaitable[None]] | None = None, + ) -> dict[str, Any]: + try: + await self._socket.ensure_open() + except SandboxAccessError as error: + _unsent_errors.add(error) + raise + + request_id = create_request_id() + future: asyncio.Future[dict[str, Any]] = asyncio.get_running_loop().create_future() + self._pending[request_id] = _PendingOperation( + future=future, on_entries=on_entries, on_data=on_data + ) + + if not await self._socket.send(json.dumps({"id": request_id, "op": op, **fields})): + lost = SandboxAccessError( + "disconnected", "The file connection was lost.", retryable=True + ) + _unsent_errors.add(lost) + self._settle(request_id, error=lost) + return await future + + upload_task: asyncio.Task[None] | None = None + if upload is not None: + upload_task = asyncio.ensure_future(upload(request_id, lambda: not future.done())) + + try: + return await future + except asyncio.CancelledError: + self._pending.pop(request_id, None) + await self._socket.send( + json.dumps({"id": create_request_id(), "op": "cancel", "target": request_id}) + ) + raise + finally: + if upload_task is not None and not upload_task.done(): + upload_task.cancel() diff --git a/src/metorial/sandbox/port.py b/src/metorial/sandbox/port.py new file mode 100644 index 00000000..419b7c5d --- /dev/null +++ b/src/metorial/sandbox/port.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +from types import TracebackType +from typing import Any +from urllib.parse import urlsplit + +import httpx + +from ._access import SandboxAccessCache +from ._errors import closed_error + +_TOKEN_HEADER = "Metorial-Sandbox-Token" + + +class SandboxPort: + def __init__( + self, + port: int, + access: SandboxAccessCache, + *, + client: httpx.AsyncClient | None = None, + ) -> None: + self.port = port + self._access = access + self._client = client or httpx.AsyncClient() + self._owns_client = client is None + self._closed = False + self._url = "" + + @classmethod + async def connect( + cls, + port: int, + access: SandboxAccessCache, + *, + client: httpx.AsyncClient | None = None, + ) -> SandboxPort: + connection = cls(port, access, client=client) + connection._url = (await access.get()).url + return connection + + @property + def url(self) -> str: + return self._url + + @property + def is_closed(self) -> bool: + return self._closed + + async def __aenter__(self) -> SandboxPort: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + await self.close() + + async def get_token(self) -> str: + return (await self._access.get()).token + + async def request(self, method: str, path: str, **kwargs: Any) -> httpx.Response: + if self._closed: + raise closed_error() + + headers = httpx.Headers(kwargs.pop("headers", None)) + + async def send() -> httpx.Response: + grant = await self._access.get() + self._url = grant.url + headers[_TOKEN_HEADER] = grant.token + return await self._client.request( + method, self._resolve(grant.url, path), headers=headers, **kwargs + ) + + response = await send() + if response.status_code == 401: + self._access.invalidate() + response = await send() + + return response + + async def get(self, path: str, **kwargs: Any) -> httpx.Response: + return await self.request("GET", path, **kwargs) + + async def head(self, path: str, **kwargs: Any) -> httpx.Response: + return await self.request("HEAD", path, **kwargs) + + async def post(self, path: str, **kwargs: Any) -> httpx.Response: + return await self.request("POST", path, **kwargs) + + async def put(self, path: str, **kwargs: Any) -> httpx.Response: + return await self.request("PUT", path, **kwargs) + + async def patch(self, path: str, **kwargs: Any) -> httpx.Response: + return await self.request("PATCH", path, **kwargs) + + async def delete(self, path: str, **kwargs: Any) -> httpx.Response: + return await self.request("DELETE", path, **kwargs) + + async def close(self) -> None: + if self._closed: + return + self._closed = True + if self._owns_client: + await self._client.aclose() + + def _resolve(self, base: str, path: str) -> str: + if path.startswith(("http://", "https://")): + target, origin = urlsplit(path), urlsplit(base) + if (target.scheme, target.netloc) != (origin.scheme, origin.netloc): + raise ValueError(f"{path} is not on the sandbox port {self.port}.") + return path + + return f"{base.rstrip('/')}/{path.lstrip('/')}" diff --git a/src/metorial/sandbox/shell.py b/src/metorial/sandbox/shell.py new file mode 100644 index 00000000..947bb0a4 --- /dev/null +++ b/src/metorial/sandbox/shell.py @@ -0,0 +1,305 @@ +from __future__ import annotations + +import asyncio +import json +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from types import TracebackType +from typing import Any + +from ._access import SandboxAccessCache +from ._errors import SandboxAccessError, SandboxCommandError, server_error +from ._mux import ( + MuxSocket, + NoticeHandler, + create_request_id, + decode_frame, + encode_frame, +) + +_STDERR = 2 + +OutputHandler = Callable[[bytes], None] +SandboxCommand = str | Sequence[str] + + +@dataclass +class SandboxProcessExit: + exit_code: int | None + signal: str | None + timed_out: bool + + +@dataclass +class SandboxCommandResult(SandboxProcessExit): + stdout_bytes: bytes + stderr_bytes: bytes + + @property + def stdout(self) -> str: + return self.stdout_bytes.decode(errors="replace") + + @property + def stderr(self) -> str: + return self.stderr_bytes.decode(errors="replace") + + +def _to_argv(command: SandboxCommand) -> list[str]: + if isinstance(command, str): + return ["/bin/sh", "-c", command] + return list(command) + + +def _lost_error() -> SandboxCommandError: + return SandboxCommandError("connection_interrupted", "The shell connection was lost.") + + +class SandboxProcess: + def __init__(self, request_id: str, socket: MuxSocket) -> None: + loop = asyncio.get_running_loop() + + self.id = request_id + self.pid: int | None = None + self._socket = socket + self._started: asyncio.Future[int] = loop.create_future() + self._exit: asyncio.Future[SandboxProcessExit] = loop.create_future() + self._stdout_handlers: list[OutputHandler] = [] + self._stderr_handlers: list[OutputHandler] = [] + + @property + def is_done(self) -> bool: + return self._exit.done() + + async def started(self) -> int: + return await asyncio.shield(self._started) + + async def wait(self) -> SandboxProcessExit: + return await asyncio.shield(self._exit) + + def on_stdout(self, handler: OutputHandler) -> None: + self._stdout_handlers.append(handler) + + def on_stderr(self, handler: OutputHandler) -> None: + self._stderr_handlers.append(handler) + + async def write(self, data: str | bytes) -> None: + self._ensure_running() + payload = data.encode() if isinstance(data, str) else data + if not await self._socket.send(encode_frame(self.id, payload)): + raise _lost_error() + + async def close_stdin(self) -> None: + await self._send_control({"type": "stdin_eof"}) + + async def signal(self, signal: str) -> None: + name = signal.upper() + if name.startswith("SIG"): + name = name[3:] + await self._send_control({"type": "signal", "signal": name}) + + async def kill(self) -> None: + await self.signal("KILL") + + def _handle_output(self, stream: int | None, data: bytes) -> None: + handlers = self._stderr_handlers if stream == _STDERR else self._stdout_handlers + for handler in handlers: + handler(data) + + def _handle_started(self, pid: int) -> None: + self.pid = pid + if not self._started.done(): + self._started.set_result(pid) + + def _handle_exit(self, exit: SandboxProcessExit) -> None: + if not self._started.done(): + self._started.set_result(self.pid if self.pid is not None else -1) + if not self._exit.done(): + self._exit.set_result(exit) + + def _handle_error(self, error: SandboxAccessError) -> None: + for future in (self._started, self._exit): + if not future.done(): + future.set_exception(error) + future.exception() + + async def _send_control(self, fields: dict[str, Any]) -> None: + self._ensure_running() + if not await self._socket.send(json.dumps({"id": self.id, **fields})): + raise _lost_error() + + def _ensure_running(self) -> None: + if self.is_done: + raise SandboxCommandError("unknown_command", "The command has already finished.") + + +class SandboxShell: + def __init__( + self, access: SandboxAccessCache, *, on_notice: NoticeHandler | None = None + ) -> None: + self._processes: dict[str, SandboxProcess] = {} + self._socket = MuxSocket( + access, + on_text=self._handle_text, + on_binary=self._handle_binary, + on_disconnect=self._handle_disconnect, + on_notice=on_notice, + ) + + @classmethod + async def connect( + cls, access: SandboxAccessCache, *, on_notice: NoticeHandler | None = None + ) -> SandboxShell: + shell = cls(access, on_notice=on_notice) + await shell._socket.ensure_open() + return shell + + @property + def is_closed(self) -> bool: + return self._socket.is_closed + + async def __aenter__(self) -> SandboxShell: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + await self.close() + + async def spawn( + self, + command: SandboxCommand, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_ms: int | None = None, + stdin: bool = False, + on_stdout: OutputHandler | None = None, + on_stderr: OutputHandler | None = None, + ) -> SandboxProcess: + await self._socket.ensure_open() + + request_id = create_request_id() + process = SandboxProcess(request_id, self._socket) + if on_stdout is not None: + process.on_stdout(on_stdout) + if on_stderr is not None: + process.on_stderr(on_stderr) + + self._processes[request_id] = process + + request: dict[str, Any] = { + "id": request_id, + "op": "run", + "v": 1, + "command": _to_argv(command), + } + if cwd is not None: + request["cwd"] = cwd + if env is not None: + request["env"] = env + if stdin: + request["stdin"] = True + if timeout_ms is not None: + request["timeoutMs"] = timeout_ms + + if not await self._socket.send(json.dumps(request)): + self._processes.pop(request_id, None) + process._handle_error(_lost_error()) + + return process + + async def command( + self, + command: SandboxCommand, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_ms: int | None = None, + stdin: str | bytes | None = None, + on_stdout: OutputHandler | None = None, + on_stderr: OutputHandler | None = None, + ) -> SandboxCommandResult: + stdout: list[bytes] = [] + stderr: list[bytes] = [] + + process = await self.spawn( + command, + cwd=cwd, + env=env, + timeout_ms=timeout_ms, + stdin=stdin is not None, + on_stdout=stdout.append, + on_stderr=stderr.append, + ) + if on_stdout is not None: + process.on_stdout(on_stdout) + if on_stderr is not None: + process.on_stderr(on_stderr) + + if stdin is not None and not process.is_done: + if stdin: + await process.write(stdin) + await process.close_stdin() + + exit = await process.wait() + return SandboxCommandResult( + exit_code=exit.exit_code, + signal=exit.signal, + timed_out=exit.timed_out, + stdout_bytes=b"".join(stdout), + stderr_bytes=b"".join(stderr), + ) + + async def close(self) -> None: + await self._socket.close() + + def _handle_text(self, message: dict[str, Any]) -> None: + request_id = message.get("id") + if not isinstance(request_id, str): + return + + process = self._processes.get(request_id) + if process is None: + return + + kind = message.get("type") + if kind == "started": + process._handle_started(int(message.get("pid") or 0)) + elif kind == "exit": + self._processes.pop(request_id, None) + code = message.get("code") + signal = message.get("signal") + process._handle_exit( + SandboxProcessExit( + exit_code=code if isinstance(code, int) else None, + signal=signal if isinstance(signal, str) else None, + timed_out=message.get("timedOut") is True, + ) + ) + elif kind == "error": + self._processes.pop(request_id, None) + process._handle_error(server_error(message.get("error"), SandboxCommandError)) + + def _handle_binary(self, data: bytes) -> None: + frame = decode_frame(data, has_tag=True) + if frame is None: + return + + process = self._processes.get(frame[0]) + if process is not None: + process._handle_output(frame[1], frame[2]) + + def _handle_disconnect(self, error: SandboxAccessError) -> None: + interrupted = ( + error + if error.code == "closed" + else SandboxCommandError("connection_interrupted", str(error)) + ) + + processes = list(self._processes.values()) + self._processes.clear() + for process in processes: + process._handle_error(interrupted) diff --git a/tests/test_sandbox.py b/tests/test_sandbox.py new file mode 100644 index 00000000..5de0d0ce --- /dev/null +++ b/tests/test_sandbox.py @@ -0,0 +1,375 @@ +from __future__ import annotations + +import asyncio +import hashlib +import json +from collections.abc import AsyncIterator +from datetime import datetime, timedelta, timezone +from typing import Any + +import pytest +from aiohttp import WSMsgType, web + +from metorial.sandbox import ( + SandboxAccessCache, + SandboxAccessError, + SandboxAccessGrant, + SandboxCommandError, + SandboxFiles, + SandboxPort, + SandboxShell, +) +from metorial.sandbox._mux import decode_frame, encode_frame + + +class FakeSandbox: + def __init__(self) -> None: + self.files: dict[str, bytes] = {} + self.valid_tokens = {"good"} + self.connections = 0 + self.flaky_drops = 0 + self.move_requests = 0 + self.token_requests = 0 + self.base_url = "" + + def _token(self, request: web.Request) -> str | None: + protocols = request.headers.get("Sec-WebSocket-Protocol", "").split(",") + for protocol in (p.strip() for p in protocols): + if protocol.startswith("metorial-sandbox-token."): + return protocol[len("metorial-sandbox-token.") :] + return None + + async def socket(self, request: web.Request) -> web.StreamResponse: + if self._token(request) not in self.valid_tokens: + return web.Response(status=401, text="token_invalid") + + ws = web.WebSocketResponse(protocols=["metorial-sandbox-v1"]) + await ws.prepare(request) + self.connections += 1 + + if request.path == "/files": + await self._files(ws) + else: + await self._exec(ws) + return ws + + async def _files(self, ws: web.WebSocketResponse) -> None: + uploads: dict[str, tuple[str, list[bytes]]] = {} + + async for message in ws: + if message.type == WSMsgType.BINARY: + frame = decode_frame(message.data) + assert frame is not None + if frame[0] in uploads: + uploads[frame[0]][1].append(frame[2]) + continue + + if message.type != WSMsgType.TEXT: + continue + + data = json.loads(message.data) + request_id = data["id"] + + async def reply(fields: dict[str, Any], request_id: str = request_id) -> None: + await ws.send_str(json.dumps({"id": request_id, **fields})) + + if data.get("type") == "end": + upload = uploads.pop(request_id, None) + if upload is None: + continue + content = b"".join(upload[1]) + if data.get("sha256") != hashlib.sha256(content).hexdigest(): + await reply({"type": "error", "error": {"code": "checksum_mismatch"}}) + continue + self.files[upload[0]] = content + await reply({"type": "result", "size": len(content)}) + continue + + op, path = data.get("op"), data.get("path") + + if op == "write": + if not data["overwrite"] and path in self.files: + error = {"code": "already_exists", "path": path} + await reply({"type": "error", "error": error}) + else: + uploads[request_id] = (path, []) + elif op == "read": + if path not in self.files: + error = {"code": "not_found", "message": "missing", "path": path} + await reply({"type": "error", "error": error}) + continue + content = self.files[path] + await reply({"type": "start", "size": len(content), "mtime": 0, "offset": 0}) + for start in range(0, len(content), 5): + await ws.send_bytes(encode_frame(request_id, content[start : start + 5])) + await reply({"type": "result"}) + elif op == "stat": + if path not in self.files: + await reply({"type": "error", "error": {"code": "not_found"}}) + continue + entry = { + "path": path, + "type": "file", + "size": len(self.files[path]), + "mode": 420, + "mtime": 1000, + } + await reply({"type": "result", "entry": entry}) + elif op == "list": + if path == "/flaky" and self.flaky_drops > 0: + self.flaky_drops -= 1 + await reply({"type": "entries", "entries": [{"name": "partial"}]}) + await ws.close(code=1011) + return + entries = [ + {"name": name[len(path) + 1 :], "type": "file", "size": len(c), "mtime": 0} + for name, c in self.files.items() + if name.startswith(f"{path}/") + ] + await reply({"type": "entries", "entries": entries[:1]}) + await reply({"type": "entries", "entries": entries[1:]}) + await reply({"type": "result", "count": len(entries), "truncated": False}) + elif op == "move": + self.move_requests += 1 + await ws.close(code=1011) + return + elif op == "delete": + self.files.pop(path, None) + await reply({"type": "result"}) + + async def _exec(self, ws: web.WebSocketResponse) -> None: + stdin: dict[str, list[bytes]] = {} + + async def output(request_id: str, stream: int, text: str) -> None: + await ws.send_bytes(encode_frame(request_id, text.encode(), stream)) + + async for message in ws: + if message.type == WSMsgType.BINARY: + frame = decode_frame(message.data) + assert frame is not None + stdin.setdefault(frame[0], []).append(frame[2]) + continue + + if message.type != WSMsgType.TEXT: + continue + + data = json.loads(message.data) + request_id = data["id"] + + async def send(fields: dict[str, Any], request_id: str = request_id) -> None: + await ws.send_str(json.dumps({"id": request_id, **fields})) + + if data.get("type") == "stdin_eof": + text = b"".join(stdin.pop(request_id, [])).decode().upper() + await output(request_id, 1, text) + await send({"type": "exit", "code": 0, "signal": None, "timedOut": False}) + continue + + if data.get("type") == "signal": + exit = {"code": None, "signal": data["signal"], "timedOut": False} + await send({"type": "exit", **exit}) + continue + + program, *args = data["command"] + await send({"type": "started", "pid": 42}) + + if program == "missing": + error = {"code": "spawn_failed", "message": "not found"} + await send({"type": "error", "error": error}) + elif program == "drop": + await ws.close(code=1011) + return + elif data.get("stdin") or program == "sleep": + stdin.setdefault(request_id, []) + else: + await output(request_id, 1, f"ran {' '.join(args)}") + await output(request_id, 2, f"cwd={data.get('cwd', '')}") + await send({"type": "exit", "code": 3, "signal": None, "timedOut": False}) + + async def http(self, request: web.Request) -> web.Response: + token = request.headers.get("Metorial-Sandbox-Token") + if token not in self.valid_tokens: + return web.Response(status=401, text="token_invalid") + + body = await request.text() + return web.json_response( + { + "method": request.method, + "url": str(request.rel_url), + "body": body, + "test": request.headers.get("x-test"), + } + ) + + def access(self, path: str, tokens: list[str] | None = None) -> SandboxAccessCache: + queue = list(tokens or ["good"]) + + async def provider() -> SandboxAccessGrant: + self.token_requests += 1 + token = queue.pop(0) if len(queue) > 1 else queue[0] + base = self.base_url + url = f"{base}/app" if path == "/port" else base.replace("http", "ws") + path + expires_at = datetime.now(timezone.utc) + timedelta(hours=1) + return SandboxAccessGrant(url=url, token=token, expires_at=expires_at) + + return SandboxAccessCache(provider) + + +@pytest.fixture +async def sandbox() -> AsyncIterator[FakeSandbox]: + fake = FakeSandbox() + app = web.Application() + app.router.add_get("/files", fake.socket) + app.router.add_get("/exec", fake.socket) + app.router.add_route("*", "/app/{tail:.*}", fake.http) + + runner = web.AppRunner(app) + await runner.setup() + site = web.TCPSite(runner, "127.0.0.1", 0) + await site.start() + + port = runner.addresses[0][1] + fake.base_url = f"http://127.0.0.1:{port}" + + yield fake + await runner.cleanup() + + +async def test_files_roundtrip(sandbox: FakeSandbox) -> None: + async with await SandboxFiles.connect(sandbox.access("/files")) as files: + big = b"x" * (200 * 1024) + await files.write("/work/a.txt", "hello world") + await files.write("/work/b.bin", big) + + assert await files.read_text("/work/a.txt") == "hello world" + assert await files.read("/work/b.bin") == big + + entry = await files.stat("/work/a.txt") + assert (entry.name, entry.path, entry.type, entry.size) == ( + "a.txt", + "/work/a.txt", + "file", + 11, + ) + + listing = await files.list("/work") + assert not listing.truncated + assert sorted(e.path for e in listing.entries) == ["/work/a.txt", "/work/b.bin"] + + await files.delete("/work/a.txt") + assert not await files.exists("/work/a.txt") + + +async def test_files_errors(sandbox: FakeSandbox) -> None: + async with await SandboxFiles.connect(sandbox.access("/files")) as files: + with pytest.raises(SandboxAccessError) as missing: + await files.read("/nope") + assert (missing.value.code, missing.value.path) == ("not_found", "/nope") + + await files.write("/exists", "a") + with pytest.raises(SandboxAccessError) as exists: + await files.write("/exists", "b", overwrite=False) + assert exists.value.code == "already_exists" + + +async def test_files_refreshes_rejected_token(sandbox: FakeSandbox) -> None: + files = await SandboxFiles.connect(sandbox.access("/files", ["stale", "good"])) + assert sandbox.token_requests == 2 + await files.write("/refresh", "ok") + assert await files.read_text("/refresh") == "ok" + await files.close() + + +async def test_files_resends_pending_list(sandbox: FakeSandbox) -> None: + async with await SandboxFiles.connect(sandbox.access("/files")) as files: + await files.write("/flaky/one", "1") + sandbox.flaky_drops = 2 + + listing = await files.list("/flaky") + assert sandbox.flaky_drops == 0 + assert [e.path for e in listing.entries] == ["/flaky/one"] + + +async def test_files_does_not_resend_move(sandbox: FakeSandbox) -> None: + async with await SandboxFiles.connect(sandbox.access("/files")) as files: + with pytest.raises(SandboxAccessError) as error: + await files.move("/a", "/b") + assert error.value.code == "disconnected" + assert sandbox.move_requests == 1 + + +async def test_files_closed(sandbox: FakeSandbox) -> None: + files = await SandboxFiles.connect(sandbox.access("/files")) + await files.close() + with pytest.raises(SandboxAccessError) as error: + await files.stat("/x") + assert error.value.code == "closed" + + +async def test_shell_commands(sandbox: FakeSandbox) -> None: + async with await SandboxShell.connect(sandbox.access("/exec")) as shell: + result = await shell.command(["echo", "hi", "there"], cwd="/tmp") + assert (result.exit_code, result.signal, result.timed_out) == (3, None, False) + assert (result.stdout, result.stderr) == ("ran hi there", "cwd=/tmp") + + via_shell = await shell.command("ls -la") + assert via_shell.stdout == "ran -c ls -la" + + piped = await shell.command(["cat"], stdin="shout") + assert piped.stdout == "SHOUT" + + +async def test_shell_spawn_signal(sandbox: FakeSandbox) -> None: + async with await SandboxShell.connect(sandbox.access("/exec")) as shell: + process = await shell.spawn(["sleep", "100"]) + assert await process.started() == 42 + + await process.signal("SIGTERM") + exit = await process.wait() + assert (exit.exit_code, exit.signal) == (None, "TERM") + + +async def test_shell_errors_and_reconnect(sandbox: FakeSandbox) -> None: + async with await SandboxShell.connect(sandbox.access("/exec")) as shell: + with pytest.raises(SandboxCommandError) as spawn_error: + await shell.command(["missing"]) + assert spawn_error.value.code == "spawn_failed" + + before = sandbox.connections + with pytest.raises(SandboxCommandError) as dropped: + await shell.command(["drop"]) + assert dropped.value.code == "connection_interrupted" + + await asyncio.sleep(0.6) + assert sandbox.connections - before == 1 + + result = await shell.command(["echo", "again"]) + assert result.stdout == "ran again" + + +async def test_port_requests(sandbox: FakeSandbox) -> None: + async with await SandboxPort.connect(3000, sandbox.access("/port")) as port: + response = await port.get("/testing?x=1", headers={"x-test": "yes"}) + assert response.json() == { + "method": "GET", + "url": "/app/testing?x=1", + "body": "", + "test": "yes", + } + + posted = await port.post("/items", json={"name": "a"}) + assert posted.json()["body"] == '{"name":"a"}' + assert port.url == f"{sandbox.base_url}/app" + + with pytest.raises(ValueError): + await port.get("https://example.com/") + + +async def test_port_retries_401(sandbox: FakeSandbox) -> None: + port = await SandboxPort.connect(3000, sandbox.access("/port", ["stale", "good"])) + response = await port.get("/") + assert response.status_code == 200 + + await port.close() + with pytest.raises(SandboxAccessError): + await port.get("/")