Add machine connection apis

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",

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."""

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",
+]

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

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
+  )