stackbox 0.1.0__py3-none-any.whl

stackbox/backends/hosted.py

@@ -0,0 +1,536 @@
+"""Shared logic for hosted providers that do not speak the OpenSandbox API themselves.
+
+The gateway emulates the OpenSandbox lifecycle on top of the provider SDK and runs
+OpenSandbox's ``execd`` inside every sandbox, so clients get the same execd endpoint
+(``GET /sandboxes/{id}/endpoints/44772`` -> public URL + headers) on every provider.
+
+A create here:
+  1. resolves the image to something the provider can boot fast (prepared template or
+     image layer with execd baked in, see ``prepare``),
+  2. creates the provider sandbox tagged with the gateway id,
+  3. starts execd with a per-sandbox access token,
+  4. waits until execd answers on its public URL.
+Any failure after step 2 deletes the provider sandbox before reporting a clear failure;
+if that delete fails, the sandbox still carries the managed tag and the reconciler
+deletes it.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import datetime as dt
+import hashlib
+import hmac
+import logging
+import shlex
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import httpx
+
+from stackbox import ledger as L
+from stackbox.backends.base import (
+    GW_ID_KEY,
+    MANAGED_KEY,
+    Backend,
+    BackendError,
+    Created,
+    CreateFailure,
+    ImageSupport,
+    UpstreamResponse,
+    error_body,
+)
+from stackbox.config import BackendConfig, GatewaySettings
+
+if TYPE_CHECKING:
+    from stackbox.ledger import SandboxRecord
+
+logger = logging.getLogger(__name__)
+
+EXECD_PORT = 44772
+EXECD_PATH = "/opt/opensandbox/execd"
+EXECD_TOKEN_HEADER = "X-EXECD-ACCESS-TOKEN"
+# OpenSandbox's egress sidecar port. The SDK resolves it on every create, but hosted
+# backends run no sidecar, so it gets a placeholder that fails loudly if used.
+EGRESS_PORT = 18080
+DEFAULT_ENTRYPOINT = ["tail", "-f", "/dev/null"]
+_NOOP_ENTRYPOINTS = {"tail -f /dev/null", "sleep infinity", "sleep inf"}
+
+# Provider tag/metadata keys. Providers restrict key characters, so the router's
+# namespaced keys are mapped to plain ones.
+PROVIDER_KEYS = {GW_ID_KEY: "stackbox_gw_id", MANAGED_KEY: "stackbox_managed"}
+_FROM_PROVIDER = {v: k for k, v in PROVIDER_KEYS.items()}
+
+# A failed preparation is retried automatically only after this long.
+PREPARE_RETRY_S = 600
+
+
+def to_provider_tags(metadata: dict[str, str]) -> dict[str, str]:
+    return {PROVIDER_KEYS.get(k, k): v for k, v in metadata.items()}
+
+
+def from_provider_tags(tags: dict[str, str] | None) -> dict[str, str]:
+    return {
+        _FROM_PROVIDER[k]: v for k, v in (tags or {}).items() if k in _FROM_PROVIDER
+    }
+
+
+def image_uri(body: dict[str, Any]) -> str | None:
+    img = body.get("image")
+    if isinstance(img, dict):
+        return img.get("uri")
+    return img if isinstance(img, str) else None
+
+
+def iso(ts: float) -> str:
+    return (
+        dt.datetime.fromtimestamp(ts, dt.UTC)
+        .isoformat(timespec="seconds")
+        .replace("+00:00", "Z")
+    )
+
+
+def parse_iso(value: str) -> float:
+    return dt.datetime.fromisoformat(value).timestamp()
+
+
+def entrypoint_command(entrypoint: list[str] | None) -> str | None:
+    """Shell command for a client entrypoint, or None when it only keeps the box alive."""
+    if not entrypoint or " ".join(entrypoint) in _NOOP_ENTRYPOINTS:
+        return None
+    return shlex.join(entrypoint)
+
+
+def execd_command(token: str, graceful_shutdown: str = "50ms") -> list[str]:
+    return [
+        EXECD_PATH,
+        "--port",
+        str(EXECD_PORT),
+        "--access-token",
+        token,
+        "--graceful-shutdown-timeout",
+        graceful_shutdown,
+    ]
+
+
+@dataclass
+class Provisioned:
+    """A provider sandbox that exists (and must be deleted if anything later fails)."""
+
+    upstream_id: str
+    host: str  # public host for the execd port, no scheme
+    headers: dict[str, str] = field(default_factory=dict)  # provider auth for that host
+    backend_data: dict[str, Any] = field(default_factory=dict)
+    handle: Any = None  # provider SDK object, for the start step
+    scheme: str = "https"  # how the gateway reaches the host (clients choose their own)
+
+
+class PostCreateError(Exception):
+    """The provider sandbox exists, but finishing setup failed."""
+
+
+class HostedBackend(Backend):
+    #: what image_support returns for an image that has not been prepared
+    unprepared_support: ImageSupport = None
+    adopt_ambiguous = False
+
+    def __init__(
+        self,
+        cfg: BackendConfig,
+        *,
+        settings: GatewaySettings,
+        ledger: L.Ledger,
+        token_key: str,
+    ):
+        super().__init__(cfg)
+        self._settings = settings
+        self._ledger = ledger
+        self._token_key = token_key.encode()
+        self._preparing: dict[str, asyncio.Task[dict[str, Any]]] = {}
+        self._http = httpx.AsyncClient(timeout=10)
+        self._execd_digest: str | None = None
+
+    async def aclose(self) -> None:
+        for task in self._preparing.values():
+            task.cancel()
+        await self._http.aclose()
+
+    # ------------------------------------------------------------- helpers
+    def execd_token(self, gw_id: str) -> str:
+        """Per-sandbox execd token, derived so it survives gateway restarts unstored."""
+        return hmac.new(self._token_key, gw_id.encode(), hashlib.sha256).hexdigest()[
+            :40
+        ]
+
+    def execd_digest(self) -> str:
+        if self._execd_digest is None:
+            h = hashlib.sha256()
+            with open(self.execd_file, "rb") as f:
+                for chunk in iter(lambda: f.read(1 << 20), b""):
+                    h.update(chunk)
+            self._execd_digest = h.hexdigest()
+        return self._execd_digest
+
+    @property
+    def execd_file(self) -> Path:
+        """Configured execd_path, else the cached download (fetched on first use)."""
+        from stackbox import execd
+
+        return execd.ensure(self.cfg.execd_path)
+
+    def requested_resources(
+        self, cpu: float | None, memory_gib: float | None
+    ) -> tuple[float, float]:
+        return (
+            cpu if cpu is not None else self.cfg.default_cpu,
+            memory_gib if memory_gib is not None else self.cfg.default_memory_gib,
+        )
+
+    def clamp_timeout(self, requested: Any) -> int:
+        try:
+            seconds = int(requested) if requested is not None else 600
+        except (TypeError, ValueError):
+            seconds = 600
+        return max(60, min(seconds, self.cfg.max_timeout_s))
+
+    # -------------------------------------------------------- image support
+    def image_record(self, image: str) -> L.ImageRecord | None:
+        return self._ledger.get_image(self.name, image)
+
+    def fits(
+        self, rec: L.ImageRecord, cpu: float | None, memory_gib: float | None
+    ) -> bool:
+        """A prepared image that fixes resources (an E2B template) serves only requests
+        within them. Images without fixed resources fit any request."""
+        d = rec.detail or {}
+        if "cpu" not in d:
+            return True
+        return (cpu or 0) <= d["cpu"] + 1e-9 and (memory_gib or 0) <= d.get(
+            "memory_gib", 0
+        ) + 1e-9
+
+    def image_support(
+        self, image: str | None, cpu: float | None, memory_gib: float | None
+    ) -> ImageSupport:
+        if not image:
+            return None
+        rec = self.image_record(image)
+        if rec is not None and rec.status == L.IMAGE_READY:
+            return "ready" if self.fits(rec, cpu, memory_gib) else None
+        if self.cfg.auto_prepare:
+            self.schedule_prepare(image)
+        return self.unprepared_support
+
+    def schedule_prepare(self, image: str) -> None:
+        if image in self._preparing:
+            return
+        rec = self.image_record(image)
+        if (
+            rec is not None
+            and rec.status == L.IMAGE_FAILED
+            and time.time() - rec.updated_at < PREPARE_RETRY_S
+        ):
+            return
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            return
+        task = asyncio.create_task(self._prepare_logged(image))
+        self._preparing[image] = task
+        task.add_done_callback(lambda _t: self._preparing.pop(image, None))
+
+    async def _prepare_logged(self, image: str) -> dict[str, Any]:
+        try:
+            return await self.prepare(image)
+        except Exception:
+            logger.exception("background prepare of %s on %s failed", image, self.name)
+            return {}
+
+    async def prepare(self, image: str) -> dict[str, Any]:
+        t0 = time.monotonic()
+        self._ledger.set_image(self.name, image, L.IMAGE_BUILDING)
+        logger.info("preparing %s on %s", image, self.name)
+        try:
+            ref, detail = await self.build_image(image)
+        except Exception as exc:
+            self._ledger.set_image(
+                self.name,
+                image,
+                L.IMAGE_FAILED,
+                error=f"{type(exc).__name__}: {exc}"[:500],
+            )
+            raise
+        detail = {**detail, "prepare_s": round(time.monotonic() - t0, 1)}
+        self._ledger.set_image(self.name, image, L.IMAGE_READY, ref=ref, detail=detail)
+        logger.info("prepared %s on %s as %s (%s)", image, self.name, ref, detail)
+        return {
+            "backend": self.name,
+            "image": image,
+            "status": L.IMAGE_READY,
+            "ref": ref,
+            **detail,
+        }
+
+    # ------------------------------------------------------ provider hooks
+    async def build_image(self, image: str) -> tuple[str, dict[str, Any]]:
+        raise NotImplementedError
+
+    async def provision(
+        self,
+        *,
+        gw_id: str,
+        image: str,
+        image_rec: L.ImageRecord | None,
+        cpu: float,
+        memory_gib: float,
+        timeout_s: int,
+        env: dict[str, str],
+        entrypoint: list[str] | None,
+        token: str,
+    ) -> Provisioned:
+        """Create the provider sandbox. Raises CreateFailure if it may not exist."""
+        raise NotImplementedError
+
+    async def start(
+        self,
+        prov: Provisioned,
+        *,
+        token: str,
+        env: dict[str, str],
+        entrypoint: list[str] | None,
+    ) -> None:
+        """Start execd (and the client entrypoint) if provisioning did not already."""
+        return
+
+    async def provider_state(self, upstream_id: str) -> tuple[str, float | None] | None:
+        """(OpenSandbox state, expires_at) or None if the sandbox is gone. Raises BackendError."""
+        raise NotImplementedError
+
+    async def kill(self, upstream_id: str) -> bool:
+        """Delete; False if it did not exist."""
+        raise NotImplementedError
+
+    def endpoint_for(
+        self, rec: SandboxRecord, port: int
+    ) -> tuple[str, dict[str, str]] | None:
+        raise NotImplementedError
+
+    async def extend(self, rec: SandboxRecord, expires_at: float) -> UpstreamResponse:
+        return UpstreamResponse(
+            409,
+            error_body(
+                "STACKBOX::RENEW_UNSUPPORTED",
+                f"backend {self.name} ({self.kind}) cannot extend a running sandbox",
+            ),
+        )
+
+    async def list_provider(self, tags: dict[str, str]) -> list[dict[str, Any]]:
+        """Provider sandboxes with these tags as OpenSandbox-shaped dicts (router metadata
+        mapped back). Raises BackendError."""
+        raise NotImplementedError
+
+    # ------------------------------------------------------------ lifecycle
+    async def create(self, body: dict[str, Any], *, gw_id: str) -> Created:
+        from stackbox.router import parse_cpu, parse_memory_gib
+
+        image = image_uri(body)
+        if not image:
+            raise CreateFailure("client", "image.uri is required", 400)
+        limits = body.get("resourceLimits") or {}
+        req_cpu = parse_cpu(limits.get("cpu"))
+        req_mem = parse_memory_gib(limits.get("memory"))
+        support = self.image_support(image, req_cpu, req_mem)
+        rec = self.image_record(image)
+        cpu, memory_gib = self.requested_resources(req_cpu, req_mem)
+        if support is None:
+            raise CreateFailure(
+                "clear", f"image {image} is not prepared on {self.name}"
+            )
+        timeout_s = self.clamp_timeout(body.get("timeout"))
+        env = {str(k): str(v) for k, v in (body.get("env") or {}).items()}
+        entrypoint = body.get("entrypoint")
+        token = self.execd_token(gw_id)
+        t0 = time.time()
+        deadline = time.monotonic() + self._settings.create_timeout_s
+
+        prov = await self.provision(
+            gw_id=gw_id,
+            image=image,
+            image_rec=rec if support == "ready" else None,
+            cpu=cpu,
+            memory_gib=memory_gib,
+            timeout_s=timeout_s,
+            env=env,
+            entrypoint=entrypoint,
+            token=token,
+        )
+        try:
+            await self.start(prov, token=token, env=env, entrypoint=entrypoint)
+            headers = {**prov.headers, EXECD_TOKEN_HEADER: token}
+            await self.wait_execd(f"{prov.scheme}://{prov.host}", headers, deadline)
+        except Exception as exc:
+            await self.discard(prov.upstream_id)
+            raise CreateFailure(
+                "clear",
+                f"sandbox started but execd setup failed: {type(exc).__name__}: {exc}",
+            ) from exc
+
+        alloc_cpu, alloc_mem = self.allocated_resources(image, cpu, memory_gib)
+        expires_at = t0 + timeout_s
+        backend_data = {
+            **prov.backend_data,
+            "host": prov.host,
+            "image": image,
+            "entrypoint": entrypoint or DEFAULT_ENTRYPOINT,
+            "createdAt": iso(t0),
+            "expiresAt": iso(expires_at),
+        }
+        sandbox = self.shape(
+            prov.upstream_id, gw_id, "Running", backend_data=backend_data
+        )
+        return Created(
+            sandbox=sandbox,
+            cpu=alloc_cpu,
+            memory_gib=alloc_mem,
+            backend_data=backend_data,
+            cold=support == "cold",
+        )
+
+    async def wait_execd(
+        self, url: str, headers: dict[str, str], deadline: float
+    ) -> None:
+        last: str = "no response"
+        while time.monotonic() < deadline:
+            try:
+                r = await self._http.get(f"{url}/ping", headers=headers)
+                if r.status_code == 200:
+                    return
+                last = f"HTTP {r.status_code}"
+            except httpx.HTTPError as exc:
+                last = type(exc).__name__
+            await asyncio.sleep(0.25)
+        raise TimeoutError(f"execd at {url} not ready ({last})")
+
+    async def discard(self, upstream_id: str) -> None:
+        try:
+            await self.kill(upstream_id)
+        except Exception as exc:  # noqa: BLE001 - the reconciler deletes it later
+            logger.warning(
+                "could not delete %s on %s after a failed create: %s",
+                upstream_id,
+                self.name,
+                exc,
+            )
+
+    def shape(
+        self,
+        upstream_id: str,
+        gw_id: str | None,
+        state: str,
+        *,
+        backend_data: dict[str, Any] | None = None,
+        expires_at: float | None = None,
+    ) -> dict[str, Any]:
+        bd = backend_data or {}
+        out: dict[str, Any] = {
+            "id": upstream_id,
+            "status": {"state": state, "lastTransitionAt": bd.get("createdAt")},
+            "entrypoint": bd.get("entrypoint") or DEFAULT_ENTRYPOINT,
+            "createdAt": bd.get("createdAt") or iso(time.time()),
+            "metadata": {
+                MANAGED_KEY: self.instance_id,
+                **({GW_ID_KEY: gw_id} if gw_id else {}),
+            },
+        }
+        if bd.get("image"):
+            out["image"] = {"uri": bd["image"]}
+        exp = iso(expires_at) if expires_at else bd.get("expiresAt")
+        if exp:
+            out["expiresAt"] = exp
+        return out
+
+    async def get(self, rec: SandboxRecord) -> UpstreamResponse:
+        assert rec.upstream_id is not None
+        try:
+            found = await self.provider_state(rec.upstream_id)
+        except BackendError as exc:
+            return UpstreamResponse(
+                502, error_body("STACKBOX::BACKEND_ERROR", str(exc))
+            )
+        if found is None:
+            return UpstreamResponse(
+                404, error_body("SANDBOX::NOT_FOUND", "sandbox no longer exists")
+            )
+        state, expires_at = found
+        return UpstreamResponse(
+            200,
+            self.shape(
+                rec.upstream_id,
+                rec.gw_id,
+                state,
+                backend_data=rec.backend_data,
+                expires_at=expires_at,
+            ),
+        )
+
+    async def delete(self, upstream_id: str) -> UpstreamResponse:
+        try:
+            existed = await self.kill(upstream_id)
+        except BackendError as exc:
+            return UpstreamResponse(
+                502, error_body("STACKBOX::BACKEND_ERROR", str(exc))
+            )
+        return UpstreamResponse(204 if existed else 404, None)
+
+    async def subresource(
+        self,
+        rec: SandboxRecord,
+        method: str,
+        rest: str,
+        *,
+        params: Any = None,
+        json: Any = None,
+    ) -> UpstreamResponse:
+        parts = [p for p in rest.split("/") if p]
+        if method == "GET" and len(parts) == 2 and parts[0] == "endpoints":
+            try:
+                port = int(parts[1])
+            except ValueError:
+                return UpstreamResponse(400, error_body("STACKBOX::BAD_PORT", parts[1]))
+            found = self.endpoint_for(rec, port)
+            if found is None and port == EGRESS_PORT:
+                found = (f"no-egress-sidecar-on-{self.kind}.invalid", {})
+            if found is None:
+                return UpstreamResponse(
+                    404,
+                    error_body(
+                        "STACKBOX::PORT_NOT_EXPOSED",
+                        f"port {port} is not exposed on backend {self.name}",
+                    ),
+                )
+            host, headers = found
+            if port == EXECD_PORT:
+                headers = {**headers, EXECD_TOKEN_HEADER: self.execd_token(rec.gw_id)}
+            return UpstreamResponse(200, {"endpoint": host, "headers": headers})
+        if method == "POST" and parts == ["renew-expiration"]:
+            try:
+                expires_at = parse_iso(str((json or {})["expiresAt"]))
+            except (KeyError, ValueError) as exc:
+                return UpstreamResponse(
+                    400, error_body("STACKBOX::BAD_REQUEST", f"expiresAt: {exc}")
+                )
+            return await self.extend(rec, expires_at)
+        return UpstreamResponse(
+            501,
+            error_body(
+                "STACKBOX::UNSUPPORTED",
+                f"{method} /sandboxes/{{id}}/{rest} is not supported on {self.kind} backends",
+            ),
+        )
+
+    async def list_sandboxes(
+        self, metadata: dict[str, str] | None = None
+    ) -> list[dict[str, Any]]:
+        return await self.list_provider(to_provider_tags(metadata or {}))