deepseek-kubernetes-sandbox/examples/dsh-agents-as-a-service/gateway/gateway.py

"""Browser-usable gateway for DeepSeek Harness Web (DSH) sandboxes.

This is the control + data plane in front of the agent-sandbox CRDs.  It is a
single teaching-friendly module with three jobs:

* **Users** -- create a ``SandboxClaim`` over a warm pool, mint a bearer token
  (only its SHA-256 is stored, as a claim annotation), report derived state,
  and delete the claim (which cascades the sandbox and its PVC).
* **Browser sessions** -- exchange a user token for an HttpOnly session cookie
  so the DSH UI can be opened in a browser at ``/u/<user>/``.  Browser requests
  authenticate with that cookie; automation can keep using
  ``Authorization: Bearer <token>``.
* **Reverse proxy** -- HTTP (including SSE) and WebSocket proxying to the
  sandbox's DSH Web server on ``:3080``.  DSH itself is cookie-gated: the
  gateway scrapes the launch token from the pod logs, performs the ``?token=``
  handshake server-side, caches the resulting ``dsh-auth-*`` cookie, and
  injects it upstream.  The browser never sees the launch or bearer token.

A background sweeper suspends sandboxes idle for ``IDLE_TIMEOUT`` seconds.
Suspending deletes the pod only; the PVC (and the user's state) survives.

Deliberately NOT production code: single replica, in-memory state and token
cache, no TLS termination, and launch-token discovery by log scraping.

Run locally with::

    uvicorn gateway:create_app --factory --host 0.0.0.0 --port 8080

The factory signature (``Settings``, ``Kubernetes``, an httpx client and the
websockets connector are all injectable) is what makes the test suite
cluster-free.
"""

from __future__ import annotations

import asyncio
import contextlib
import hashlib
import html
import json
import os
import re
import secrets
import time
from dataclasses import dataclass, field
from urllib.parse import (
    parse_qs,
    parse_qsl,
    quote,
    unquote_plus,
    urlencode,
    urlsplit,
    urlunsplit,
)

import anyio
import httpx
import websockets
from kubernetes import client as k8s_client
from kubernetes import config as k8s_config
from kubernetes.client.exceptions import ApiException
from starlette.applications import Starlette
from starlette.datastructures import Headers
from starlette.requests import Request
from starlette.responses import (
    HTMLResponse,
    JSONResponse,
    RedirectResponse,
    Response,
    StreamingResponse,
)
from starlette.routing import Route, WebSocketRoute
from starlette.websockets import WebSocket

# --- Kubernetes / DSH constants -------------------------------------------

GROUP, VERSION = "extensions.agents.x-k8s.io", "v1beta1"
CORE_GROUP = "agents.x-k8s.io"
TOKEN_ANNOTATION = "aaas.example.com/token-sha256"
# Ownership marker: the gateway shares its namespace with claims applied by
# hand (the resource-layer walkthrough in the example README).  Every claim is
# named "dsh-<user>", so without an explicit marker the sweeper cannot tell its
# own users from a claim someone created manually.
MANAGED_ANNOTATION = "aaas.example.com/managed-by"
MANAGED_BY = "dsh-gateway"
POD_NAME_ANNOTATION = "agents.x-k8s.io/pod-name"
USER_LABEL = "sandbox.users.io/dsh-user"

# RFC 1123 label: lowercase alphanumerics and "-", must start/end alphanumeric.
# The claim name is "dsh-<user>", so cap the user part to keep the claim a
# valid Kubernetes object name.
DNS1123_LABEL = re.compile(r"^[a-z0-9]([-a-z0-9]*[a-z0-9])?$")
MAX_USER_LEN = 63 - len("dsh-")

# Headers a proxy must not forward in either direction (RFC 7230 section 6.1).
HOP_BY_HOP = frozenset(
    {
        "connection",
        "keep-alive",
        "proxy-authenticate",
        "proxy-authorization",
        "te",
        "trailer",
        "trailers",
        "transfer-encoding",
        "upgrade",
    }
)
SAFE_METHODS = frozenset({"GET", "HEAD", "OPTIONS"})
# Headers the gateway owns: never forwarded from the browser, overwritten with
# the gateway's own view of the connection instead.
OWNED_REQUEST_HEADERS = HOP_BY_HOP | {
    "authorization",
    "cookie",
    "host",
    "x-forwarded-for",
    "x-forwarded-host",
    "x-forwarded-port",
    "x-forwarded-proto",
    "x-real-ip",
}

_DSH_WEB_LINE = re.compile(r"dsh web:\s+(?P<url>https?://\S+)")


class UserExists(Exception):
    """The claim already exists (a 409 from the API server)."""


class BodyTooLarge(Exception):
    """The request body exceeded ``MAX_BODY_MB``."""


class HandshakeError(Exception):
    """DSH's launch token or auth cookie could not be obtained."""


class SandboxUnavailable(Exception):
    """The sandbox is not (yet) ready to serve traffic."""


class UpstreamError(Exception):
    """The upstream DSH connection failed."""


# --- Pure helpers (unit-testable without an app) --------------------------


def valid_user(user: str) -> bool:
    """A user is usable as a claim-name suffix (dns-1123 label, length-capped)."""
    return (
        isinstance(user, str)
        and len(user) <= MAX_USER_LEN
        and DNS1123_LABEL.fullmatch(user) is not None
    )


def claim_name(user: str) -> str:
    return f"dsh-{user}"


def claim_user(claim: dict) -> str:
    return claim["metadata"]["name"].removeprefix("dsh-")


def gateway_managed(claim: dict) -> bool:
    """True only for claims this gateway created.

    The sweeper must not suspend claims it does not own: a claim applied from
    a resource-layer manifest (e.g. ``30-claim-alice.yaml``) has the same
    ``dsh-<user>`` name but no ownership marker, and suspending it deletes a
    pod the operator is using.
    """
    annotations = (claim.get("metadata") or {}).get("annotations") or {}
    return annotations.get(MANAGED_ANNOTATION) == MANAGED_BY


def condition(obj: dict, ctype: str) -> bool:
    for c in (obj.get("status") or {}).get("conditions", []):
        if c.get("type") == ctype:
            return c.get("status") == "True"
    return False


def derive_state(sandbox: dict | None) -> str:
    """The state machine a platform shows users (mode x conditions)."""
    if sandbox is None:
        return "Provisioning"
    mode = sandbox.get("spec", {}).get("operatingMode", "Running")
    if mode == "Suspended":
        return "Suspended" if condition(sandbox, "Suspended") else "Suspending"
    return "Ready" if condition(sandbox, "Ready") else "Waking"


def token_matches(claim: dict, token: str | None) -> bool:
    """Timing-safe check of a user token against the stored SHA-256."""
    want = (claim.get("metadata", {}).get("annotations") or {}).get(
        TOKEN_ANNOTATION, ""
    )
    if not token or not want:
        return False
    got = hashlib.sha256(token.encode()).hexdigest()
    return secrets.compare_digest(got, want)


def parse_dsh_launch_token(logs: str) -> str | None:
    """Return the token from the last ``dsh web: http://...?token=...`` line.

    DSH prints the launch URL once at startup.  We scan the tail of the pod log
    and keep the last match so a restarted process (which prints a new token)
    wins over the stale one.
    """
    token = None
    for line in (logs or "").splitlines():
        match = _DSH_WEB_LINE.search(line)
        if not match:
            continue
        values = parse_qs(urlsplit(match.group("url")).query).get("token")
        if values and values[0]:
            token = values[0]
    return token


def extract_dsh_cookies(set_cookie_headers: list[str]) -> str:
    """Join the ``dsh-auth-*`` pairs from Set-Cookie headers into a Cookie value."""
    pairs = []
    for raw in set_cookie_headers:
        first = raw.split(";", 1)[0]
        name, sep, value = first.partition("=")
        name = name.strip()
        if sep and name.startswith("dsh-auth-"):
            pairs.append(f"{name}={value.strip()}")
    return "; ".join(pairs)


def same_origin(headers) -> bool:
    """True when Origin (or Referer) names the same authority as Host."""
    host = (headers.get("host") or "").strip().lower()
    origin = headers.get("origin")
    if origin:
        return urlsplit(origin).netloc.lower() == host
    referer = headers.get("referer")
    if referer:
        return urlsplit(referer).netloc.lower() == host
    return False


def rewrite_location(location: str, user: str) -> str:
    """Prefix absolute-path redirects so the browser stays under ``/u/<user>/``."""
    if not location.startswith("/") or location.startswith("//"):
        return location
    parts = urlsplit(location)
    prefix = f"/u/{user}"
    if parts.path == prefix or parts.path.startswith(f"{prefix}/"):
        return location
    query = [
        (key, value)
        for key, value in parse_qsl(parts.query, keep_blank_values=True)
        if key != "token"
    ]
    return urlunsplit(("", "", prefix + parts.path, urlencode(query), parts.fragment))


def filtered_query(query_string: str) -> str:
    """Drop ``?token=`` before a request is sent upstream.

    The remaining query is forwarded byte-for-byte.  DSH's combined-module
    URLs look like ``plugins/??@deepseek-ai/dsh-client-modules/client.js`` —
    the raw query itself begins with ``?`` and DSH routes it verbatim.
    Re-encoding through ``parse_qsl``/``urlencode`` turns that into ``%3F…``,
    the upstream answers 404, and the client boot dies with "HTML did not
    preload @deepseek-ai/dsh-client-modules/client.js".
    """
    if not query_string:
        return ""
    kept = []
    for part in query_string.split("&"):
        if unquote_plus(part.split("=", 1)[0]) == "token":
            continue
        kept.append(part)
    return "&".join(kept)


def _env_int(name: str, default: int) -> int:
    try:
        return int(os.environ.get(name, default))
    except ValueError:
        return default


def _env_bool(name: str, default: bool) -> bool:
    raw = os.environ.get(name)
    if raw is None:
        return default
    return raw.strip().lower() in {"1", "true", "yes", "on"}


def _env_float(name: str, default: float) -> float:
    try:
        return float(os.environ.get(name, default))
    except ValueError:
        return default


# --- Settings --------------------------------------------------------------


@dataclass(frozen=True)
class Settings:
    """Runtime configuration, read from the environment by default.

    Defaults are the demo values from the example manifests.  Everything is
    injected into ``create_app`` so tests never need the environment or a
    cluster.
    """

    namespace: str = "dsh-demo"
    pool: str = "dsh-pool"
    idle_timeout: int = 60
    wake_timeout: int = 120
    dsh_port: int = 3080
    dsh_container: str = "dsh-web"
    dsh_log_tail_lines: int = 2000
    cookie_name: str = "dsh_gw_session"
    cookie_secure: bool = False
    cookie_max_age: int = 7 * 24 * 60 * 60
    max_body_bytes: int = 32 * 1024 * 1024
    http_timeout: float = 300.0
    sweep_interval: float = 15.0
    # Connection retries while a Ready pod is still binding its port.
    upstream_attempts: int = 3
    upstream_retry_delay: float = 1.0

    @classmethod
    def from_env(cls) -> Settings:
        return cls(
            namespace=os.environ.get("NAMESPACE", "dsh-demo"),
            pool=os.environ.get("POOL", "dsh-pool"),
            idle_timeout=_env_int("IDLE_TIMEOUT", 60),
            wake_timeout=_env_int("WAKE_TIMEOUT", 120),
            dsh_port=_env_int("DSH_PORT", 3080),
            dsh_container=os.environ.get("DSH_CONTAINER", "dsh-web"),
            dsh_log_tail_lines=_env_int("DSH_LOG_TAIL_LINES", 2000),
            cookie_name=os.environ.get("COOKIE_NAME", "dsh_gw_session"),
            cookie_secure=_env_bool("COOKIE_SECURE", False),
            cookie_max_age=_env_int("COOKIE_MAX_AGE", 7 * 24 * 60 * 60),
            max_body_bytes=_env_int("MAX_BODY_MB", 32) * 1024 * 1024,
            http_timeout=_env_float("HTTP_TIMEOUT", 300.0),
            sweep_interval=_env_float("SWEEP_INTERVAL", 15.0),
        )


# --- Kubernetes adapter ----------------------------------------------------
#
# The official Kubernetes client is synchronous.  Every call is pushed onto a
# worker thread with anyio.to_thread so it never blocks the event loop.  The
# client and its credentials are loaded lazily, on the first real call: that
# keeps ``GET /healthz`` (and importing this module) cluster-free.


class Kubernetes:
    """Thin async wrapper over the sync CustomObjects/CoreV1 API clients."""

    def __init__(self, namespace: str = "dsh-demo"):
        self.namespace = namespace
        self._crd: k8s_client.CustomObjectsApi | None = None
        self._core: k8s_client.CoreV1Api | None = None

    def _clients(self):
        if self._crd is None:
            try:
                k8s_config.load_incluster_config()
            except k8s_config.ConfigException:
                k8s_config.load_kube_config()
            self._crd = k8s_client.CustomObjectsApi()
            self._core = k8s_client.CoreV1Api()
        return self._crd, self._core

    async def _call(self, which: str, method: str, *args, **kwargs):
        def op():
            crd, core = self._clients()
            target = crd if which == "crd" else core
            return getattr(target, method)(*args, **kwargs)

        return await anyio.to_thread.run_sync(op)

    # -- claims ------------------------------------------------------------

    async def get_claim(self, user: str) -> dict | None:
        try:
            return await self._call(
                "crd",
                "get_namespaced_custom_object",
                GROUP,
                VERSION,
                self.namespace,
                "sandboxclaims",
                claim_name(user),
            )
        except ApiException as exc:
            if exc.status == 404:
                return None
            raise

    async def create_claim(self, user: str, token_hash: str, pool: str) -> None:
        body = {
            "apiVersion": f"{GROUP}/{VERSION}",
            "kind": "SandboxClaim",
            "metadata": {
                "name": claim_name(user),
                "annotations": {
                    TOKEN_ANNOTATION: token_hash,
                    MANAGED_ANNOTATION: MANAGED_BY,
                },
            },
            "spec": {
                "warmPoolRef": {"name": pool},
                "additionalPodMetadata": {"labels": {USER_LABEL: user}},
            },
        }
        try:
            await self._call(
                "crd",
                "create_namespaced_custom_object",
                GROUP,
                VERSION,
                self.namespace,
                "sandboxclaims",
                body,
            )
        except ApiException as exc:
            if exc.status == 409:
                raise UserExists(user) from exc
            raise

    async def delete_claim(self, user: str) -> None:
        try:
            await self._call(
                "crd",
                "delete_namespaced_custom_object",
                GROUP,
                VERSION,
                self.namespace,
                "sandboxclaims",
                claim_name(user),
            )
        except ApiException as exc:
            if exc.status != 404:  # already gone => idempotent success
                raise

    async def list_claims(self) -> list[dict]:
        result = await self._call(
            "crd",
            "list_namespaced_custom_object",
            GROUP,
            VERSION,
            self.namespace,
            "sandboxclaims",
        )
        return result.get("items", [])

    # -- sandboxes ---------------------------------------------------------

    async def get_sandbox(self, name: str) -> dict | None:
        try:
            return await self._call(
                "crd",
                "get_namespaced_custom_object",
                CORE_GROUP,
                VERSION,
                self.namespace,
                "sandboxes",
                name,
            )
        except ApiException as exc:
            if exc.status == 404:
                return None
            raise

    async def patch_operating_mode(self, name: str, mode: str) -> None:
        await self._call(
            "crd",
            "patch_namespaced_custom_object",
            CORE_GROUP,
            VERSION,
            self.namespace,
            "sandboxes",
            name,
            {"spec": {"operatingMode": mode}},
        )

    # -- pods and logs -----------------------------------------------------

    async def get_pod(self, name: str) -> dict | None:
        try:
            pod = await self._call(
                "core", "read_namespaced_pod", name, self.namespace
            )
        except ApiException as exc:
            if exc.status == 404:
                return None
            raise
        # CoreV1Api returns a V1Pod model while every other call in this module
        # (CRD calls, fakes in tests) yields a dict. Normalise at the boundary
        # so callers can keep using ``pod.get(...)``.
        return k8s_client.ApiClient().sanitize_for_serialization(pod)

    async def read_pod_logs(
        self, name: str, container: str, tail_lines: int
    ) -> str:
        result = await self._call(
            "core",
            "read_namespaced_pod_log",
            name,
            self.namespace,
            container=container,
            tail_lines=tail_lines,
        )
        return result or ""


# --- Gateway ---------------------------------------------------------------


@dataclass
class GatewayState:
    """In-memory, per-process state (single replica by design)."""

    last_activity: dict[str, float] = field(default_factory=dict)
    in_flight: dict[str, int] = field(default_factory=dict)
    maintenance: dict[str, bool] = field(default_factory=dict)
    # (user, pod UID, browser Host) -> "dsh-auth-a=...; dsh-auth-b=..."
    handshake: dict[tuple[str, str, str], str] = field(default_factory=dict)
    lock: asyncio.Lock = field(default_factory=asyncio.Lock)


class Gateway:
    """Route handlers plus the state machine they share."""

    def __init__(
        self,
        settings: Settings,
        k8s,
        http_client: httpx.AsyncClient,
        ws_connect,
    ):
        self.settings = settings
        self.k8s = k8s
        self.http = http_client
        self.ws_connect = ws_connect
        self.state = GatewayState()

    # -- credentials -------------------------------------------------------

    def credentials(self, headers, cookies) -> tuple[str | None, str | None]:
        """Return ``(token, cookie_user)``; bearer wins over the session cookie."""
        auth = headers.get("authorization", "")
        if auth.lower().startswith("bearer "):
            return auth[7:].strip(), None
        raw = cookies.get(self.settings.cookie_name)
        if raw:
            user, sep, token = raw.partition(":")
            if sep and user and token:
                return token, user
        return None, None

    def _authorized(
        self,
        claim: dict,
        token: str | None,
        cookie_user: str | None,
        user: str,
    ) -> bool:
        if cookie_user is not None and cookie_user != user:
            return False  # a session for another user never authorizes this one
        return token_matches(claim, token)

    @staticmethod
    def _mutation_from_cookie(headers, method: str, cookie_user: str | None) -> bool:
        return (
            cookie_user is not None
            and method.upper() not in SAFE_METHODS
            and not same_origin(headers)
        )

    def set_session(self, response: Response, user: str, token: str) -> None:
        response.set_cookie(
            self.settings.cookie_name,
            f"{user}:{token}",
            max_age=self.settings.cookie_max_age,
            path="/",
            httponly=True,
            samesite="lax",
            secure=self.settings.cookie_secure,
        )

    # -- activity / maintenance -------------------------------------------

    async def _touch(self, user: str) -> None:
        async with self.state.lock:
            self.state.last_activity[user] = time.time()

    async def _acquire(self, user: str) -> bool:
        """Reserve an in-flight slot, waiting out a concurrent suspend."""
        deadline = time.monotonic() + self.settings.wake_timeout
        while True:
            async with self.state.lock:
                if not self.state.maintenance.get(user):
                    self.state.in_flight[user] = (
                        self.state.in_flight.get(user, 0) + 1
                    )
                    self.state.last_activity[user] = time.time()
                    return True
            if time.monotonic() > deadline:
                return False
            await asyncio.sleep(0.1)

    async def _release(self, user: str) -> None:
        async with self.state.lock:
            self.state.in_flight[user] = max(
                self.state.in_flight.get(user, 1) - 1, 0
            )
            self.state.last_activity[user] = time.time()

    # -- sandbox / handshake helpers --------------------------------------

    async def _sandbox_of(self, claim: dict) -> dict | None:
        name = (claim.get("status") or {}).get("sandbox", {}).get("name")
        if not name:
            return None
        return await self.k8s.get_sandbox(name)

    async def wait_ready(self, user: str, timeout: int) -> dict | None:
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            claim = await self.k8s.get_claim(user)
            sandbox = await self._sandbox_of(claim) if claim else None
            if sandbox is not None and derive_state(sandbox) == "Ready":
                return sandbox
            await asyncio.sleep(1)
        return None

    async def _resolve_sandbox(self, user: str, claim: dict) -> tuple[dict, str]:
        sandbox = await self._sandbox_of(claim)
        if sandbox is None:
            raise SandboxUnavailable("sandbox not provisioned yet")
        if derive_state(sandbox) != "Ready":
            # Wake-on-connect: safe in ANY state, including mid-suspension.
            self._clear_handshake(user)
            try:
                await self.k8s.patch_operating_mode(
                    sandbox["metadata"]["name"], "Running"
                )
            except Exception as exc:
                raise SandboxUnavailable("could not wake the sandbox") from exc
            sandbox = await self.wait_ready(user, self.settings.wake_timeout)
            if sandbox is None:
                raise SandboxUnavailable("DSH is waking up, retry shortly")
            await self._touch(user)
        fqdn = (sandbox.get("status") or {}).get("serviceFQDN")
        if not fqdn:
            raise SandboxUnavailable("sandbox has no service FQDN yet")
        return sandbox, fqdn

    def _pod_name(self, sandbox: dict) -> str:
        annotations = sandbox.get("metadata", {}).get("annotations") or {}
        return (
            annotations.get(POD_NAME_ANNOTATION)
            or sandbox["metadata"]["name"]
        )

    def _clear_handshake(self, user: str) -> None:
        for key in [k for k in self.state.handshake if k[0] == user]:
            self.state.handshake.pop(key, None)

    async def dsh_cookie(
        self,
        user: str,
        sandbox: dict,
        fqdn: str,
        host: str,
        force: bool = False,
    ) -> tuple[str, tuple[str, str, str], bool]:
        """Return ``(cookie, cache_key, was_cached)``, handshaking on a miss."""
        pod_name = self._pod_name(sandbox)
        try:
            pod = await self.k8s.get_pod(pod_name)
        except ApiException as exc:
            raise HandshakeError(f"could not read pod {pod_name}: {exc}") from exc
        if pod is None:
            raise HandshakeError(f"pod {pod_name} not found")
        uid = (pod.get("metadata") or {}).get("uid") or pod_name
        key = (user, uid, host)
        if not force:
            cached = self.state.handshake.get(key)
            if cached:
                return cached, key, True
        try:
            logs = await self.k8s.read_pod_logs(
                pod_name,
                self.settings.dsh_container,
                self.settings.dsh_log_tail_lines,
            )
        except ApiException as exc:
            raise HandshakeError(
                f"could not read logs for {pod_name}: {exc}"
            ) from exc
        launch_token = parse_dsh_launch_token(logs)
        if not launch_token:
            raise HandshakeError("DSH launch token not found in pod logs")
        url = (
            f"http://{fqdn}:{self.settings.dsh_port}/"
            f"?token={quote(launch_token, safe='')}"
        )
        headers = {
            "Host": host,
            "X-Forwarded-Host": host,
            "X-Forwarded-Proto": "http",
        }
        try:
            response = await self.http.get(
                url,
                headers=headers,
                follow_redirects=False,
                timeout=self.settings.http_timeout,
            )
        except httpx.HTTPError as exc:
            raise HandshakeError(f"DSH handshake failed: {exc}") from exc
        if response.status_code >= 400:
            raise HandshakeError(
                f"DSH handshake returned {response.status_code}"
            )
        cookie = extract_dsh_cookies(response.headers.get_list("set-cookie"))
        if not cookie:
            raise HandshakeError("DSH handshake set no dsh-auth cookie")
        self.state.handshake[key] = cookie
        return cookie, key, False

    # -- landing page / health --------------------------------------------

    async def redirect_to_workspace(self, request: Request) -> Response:
        """``/u/<user>`` -> ``/u/<user>/``; DSH's relative URLs need the slash."""
        user = request.path_params["user"]
        if not valid_user(user):
            return self._error(404, "unknown user")
        target = f"/u/{user}/"
        if request.url.query:
            target = f"{target}?{request.url.query}"
        return RedirectResponse(target, status_code=308)

    async def healthz(self, request: Request) -> JSONResponse:
        return JSONResponse({"status": "ok"})

    async def landing(self, request: Request) -> HTMLResponse:
        token, user = self.credentials(request.headers, request.cookies)
        session = ""
        if user and token:
            safe_user = html.escape(user)
            session = (
                f'<p class="session">Signed in as <strong>{safe_user}</strong>. '
                f'<a href="/u/{quote(user)}/">Open workspace</a></p>'
                '<form method="post" action="/logout">'
                '<button type="submit">Sign out</button></form>'
            )
        body = f"""<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>DSH sandbox gateway</title>
<style>
  body {{ font-family: system-ui, sans-serif; max-width: 42rem;
         margin: 3rem auto; padding: 0 1rem; line-height: 1.5; }}
  form {{ display: flex; gap: .5rem; margin: .5rem 0 1.5rem; }}
  input {{ flex: 1; padding: .4rem; }}
  button {{ padding: .4rem 1rem; }}
  code {{ background: #eee; padding: .1rem .3rem; }}
</style>
</head>
<body>
<h1>DeepSeek Harness sandboxes</h1>
{session}
<section>
<h2>Create a sandbox</h2>
<form method="post" action="/users">
  <input name="user" placeholder="username (dns-1123 label)" required>
  <button type="submit">Create</button>
</form>
</section>
<section>
<h2>Sign in with a token</h2>
<form method="post" action="/login">
  <input name="user" placeholder="username" required>
  <input name="token" placeholder="bearer token" required>
  <button type="submit">Sign in</button>
</form>
</section>
<p>Automation can use <code>Authorization: Bearer &lt;token&gt;</code>
against the JSON API. The browser session is an HttpOnly cookie.</p>
</body>
</html>"""
        return HTMLResponse(body)

    # -- JSON helpers ------------------------------------------------------

    @staticmethod
    def _error(status_code: int, message: str, headers=None) -> JSONResponse:
        return JSONResponse({"error": message}, status_code=status_code, headers=headers)

    @staticmethod
    def _wants_form(request: Request) -> bool:
        ctype = request.headers.get("content-type", "").lower()
        return ctype.startswith(
            ("application/x-www-form-urlencoded", "multipart/form-data")
        )

    async def _read_payload(self, request: Request) -> dict:
        if request.headers.get("content-type", "").lower().startswith(
            "application/json"
        ):
            data = await self._read_body(request)
            try:
                payload = json.loads(data)
            except ValueError:
                payload = None
            return payload if isinstance(payload, dict) else {}
        form = await request.form()
        return {key: form.get(key) for key in form}

    async def _read_body(self, request: Request) -> bytes:
        limit = self.settings.max_body_bytes
        raw_length = request.headers.get("content-length")
        if raw_length:
            try:
                if int(raw_length) > limit:
                    raise BodyTooLarge
            except ValueError:
                pass
        chunks, total = [], 0
        async for chunk in request.stream():
            total += len(chunk)
            if total > limit:
                raise BodyTooLarge
            chunks.append(chunk)
        return b"".join(chunks)

    # -- users -------------------------------------------------------------

    async def signup(self, request: Request) -> Response:
        wants_form = self._wants_form(request)
        try:
            payload = await self._read_payload(request)
        except BodyTooLarge:
            return self._error(413, "request body too large")
        user = payload.get("user", "")
        if not valid_user(user):
            return self._error(
                400, "body must be {'user': '<dns-1123 label>'}"
            )
        if await self.k8s.get_claim(user) is not None:
            return self._error(409, "user exists")
        token = secrets.token_urlsafe(32)
        token_hash = hashlib.sha256(token.encode()).hexdigest()
        try:
            await self.k8s.create_claim(user, token_hash, self.settings.pool)
        except UserExists:
            return self._error(409, "user exists")
        await self._touch(user)
        sandbox = await self.wait_ready(user, self.settings.wake_timeout)
        state = derive_state(sandbox)
        sandbox_name = (sandbox or {}).get("metadata", {}).get("name")
        if wants_form:
            response = RedirectResponse(f"/u/{user}/", status_code=303)
            self.set_session(response, user, token)
            return response
        response = JSONResponse(
            {"user": user, "token": token, "state": state, "sandbox": sandbox_name},
            status_code=201 if state == "Ready" else 202,
        )
        self.set_session(response, user, token)
        return response

    async def login(self, request: Request) -> Response:
        wants_form = self._wants_form(request)
        try:
            payload = await self._read_payload(request)
        except BodyTooLarge:
            return self._error(413, "request body too large")
        user = payload.get("user", "")
        token = payload.get("token", "")
        if not isinstance(user, str) or not isinstance(token, str):
            return self._error(400, "user and token are required")
        claim = await self.k8s.get_claim(user)
        if claim is None:
            return self._error(404, "not found")
        if not token_matches(claim, token):
            return self._error(401, "unauthorized")
        await self._touch(user)
        if wants_form:
            response = RedirectResponse(f"/u/{user}/", status_code=303)
        else:
            response = JSONResponse({"user": user, "state": "ok"})
        self.set_session(response, user, token)
        return response

    async def logout(self, request: Request) -> Response:
        response = RedirectResponse("/", status_code=303)
        response.delete_cookie(self.settings.cookie_name, path="/")
        return response

    async def get_user(self, request: Request) -> Response:
        user = request.path_params["user"]
        if not valid_user(user):
            return self._error(404, "not found")
        claim = await self.k8s.get_claim(user)
        if claim is None:
            return self._error(404, "not found")
        token, cookie_user = self.credentials(request.headers, request.cookies)
        if not self._authorized(claim, token, cookie_user, user):
            return self._error(401, "unauthorized")
        sandbox = await self._sandbox_of(claim)
        return JSONResponse(
            {
                "user": user,
                "state": derive_state(sandbox),
                "sandbox": (sandbox or {}).get("metadata", {}).get("name"),
            }
        )

    async def delete_user(self, request: Request) -> Response:
        user = request.path_params["user"]
        if not valid_user(user):
            return self._error(404, "not found")
        claim = await self.k8s.get_claim(user)
        if claim is None:
            return self._error(404, "not found")
        token, cookie_user = self.credentials(request.headers, request.cookies)
        if not self._authorized(claim, token, cookie_user, user):
            return self._error(401, "unauthorized")
        if self._mutation_from_cookie(request.headers, "DELETE", cookie_user):
            return self._error(403, "cross-origin request rejected")
        await self.k8s.delete_claim(user)
        self._clear_handshake(user)
        async with self.state.lock:
            self.state.last_activity.pop(user, None)
            self.state.in_flight.pop(user, None)
        return JSONResponse(
            {"user": user, "note": "claim deleted; sandbox and PVC cascade"}
        )

    # -- HTTP proxy --------------------------------------------------------

    async def proxy_http(self, request: Request) -> Response:
        user = request.path_params["user"]
        if not valid_user(user):
            return self._error(404, "unknown user")
        claim = await self.k8s.get_claim(user)
        if claim is None:
            return self._error(404, "unknown user")
        token, cookie_user = self.credentials(request.headers, request.cookies)
        if not self._authorized(claim, token, cookie_user, user):
            return self._error(401, "unauthorized")
        if self._mutation_from_cookie(request.headers, request.method, cookie_user):
            return self._error(403, "cross-origin request rejected")

        if not await self._acquire(user):
            return self._error(
                503, "gateway busy, retry shortly", {"Retry-After": "5"}
            )
        streaming = False
        try:
            sandbox, fqdn = await self._resolve_sandbox(user, claim)
            host = request.headers.get("host", "")
            cookie, _, was_cached = await self.dsh_cookie(
                user, sandbox, fqdn, host
            )
            body = await self._read_body(request)

            upstream = await self._send_upstream(
                request, user, fqdn, cookie, body
            )
            if upstream.status_code == 401 and was_cached:
                # The cached DSH cookie expired (pod restart, logout, ...):
                # handshake once more and replay the (already buffered) request.
                await upstream.aclose()
                cookie, _, _ = await self.dsh_cookie(
                    user, sandbox, fqdn, host, force=True
                )
                upstream = await self._send_upstream(
                    request, user, fqdn, cookie, body
                )

            response = StreamingResponse(
                self._relay(user, upstream),
                status_code=upstream.status_code,
                headers=self._response_headers(upstream, user),
            )
            streaming = True
            return response
        except BodyTooLarge:
            return self._error(413, "request body too large")
        except (HandshakeError, SandboxUnavailable) as exc:
            return self._error(
                503,
                f"{exc}; retry shortly",
                {"Retry-After": "10"},
            )
        except UpstreamError:
            return self._error(502, "upstream unavailable")
        finally:
            if not streaming:
                await self._release(user)

    async def _send_upstream(
        self,
        request: Request,
        user: str,
        fqdn: str,
        cookie: str,
        body: bytes,
    ) -> httpx.Response:
        path = request.path_params.get("path", "")
        url = f"http://{fqdn}:{self.settings.dsh_port}/{path}"
        query = filtered_query(request.url.query)
        if query:
            url = f"{url}?{query}"

        headers = {
            key: value
            for key, value in request.headers.items()
            if key.lower() not in OWNED_REQUEST_HEADERS
            and key.lower() != "content-length"
        }
        host = request.headers.get("host", "")
        headers["host"] = host
        headers["cookie"] = cookie  # only the cached DSH auth cookie
        headers["x-forwarded-host"] = host
        headers["x-forwarded-proto"] = request.url.scheme
        if request.client:
            headers["x-forwarded-for"] = request.client.host
        else:
            headers["x-forwarded-for"] = ""
        port = host.rsplit(":", 1)[1] if ":" in host else ""
        headers["x-forwarded-port"] = port or (
            "443" if request.url.scheme == "https" else "80"
        )

        last_error: Exception | None = None
        for attempt in range(self.settings.upstream_attempts):
            try:
                upstream_request = self.http.build_request(
                    request.method, url, headers=headers, content=body
                )
                return await self.http.send(upstream_request, stream=True)
            except (httpx.ConnectError, httpx.ConnectTimeout) as exc:
                # Retry only failures that cannot have reached the upstream;
                # replaying a timed-out POST could duplicate a mutation.
                last_error = exc
                if attempt + 1 < self.settings.upstream_attempts:
                    await asyncio.sleep(self.settings.upstream_retry_delay)
            except httpx.HTTPError as exc:
                raise UpstreamError("upstream request failed") from exc
        raise UpstreamError("upstream unavailable") from last_error

    @staticmethod
    def _response_headers(upstream: httpx.Response, user: str):
        headers = []
        for key, value in upstream.headers.multi_items():
            lowered = key.lower()
            if lowered in HOP_BY_HOP or lowered == "set-cookie":
                continue  # never leak upstream cookies to the browser
            if lowered == "location":
                value = rewrite_location(value, user)
            headers.append((key, value))
        return Headers(raw=[(k.lower().encode("latin-1"), v.encode("latin-1")) for k, v in headers])

    async def _relay(self, user: str, upstream: httpx.Response):
        try:
            async for chunk in upstream.aiter_raw():
                yield chunk
        finally:
            with contextlib.suppress(Exception):
                await upstream.aclose()
            await self._release(user)

    # -- WebSocket proxy ---------------------------------------------------

    async def _ws_reject(self, websocket: WebSocket, code: int) -> None:
        # A close code can only be delivered after the upgrade is accepted;
        # before accept Starlette would turn this into an HTTP 403.
        with contextlib.suppress(Exception):
            await websocket.accept()
        with contextlib.suppress(Exception):
            await websocket.close(code=code)

    async def proxy_ws(self, websocket: WebSocket) -> None:
        user = websocket.path_params["user"]
        if not valid_user(user):
            return await self._ws_reject(websocket, 4401)
        claim = await self.k8s.get_claim(user)
        if claim is None:
            return await self._ws_reject(websocket, 4401)
        token, cookie_user = self.credentials(websocket.headers, websocket.cookies)
        if not self._authorized(claim, token, cookie_user, user):
            return await self._ws_reject(websocket, 4401)
        if cookie_user is not None and not same_origin(websocket.headers):
            return await self._ws_reject(websocket, 4401)

        if not await self._acquire(user):
            return await self._ws_reject(websocket, 1011)
        upstream_cm = None
        try:
            try:
                sandbox, fqdn = await self._resolve_sandbox(user, claim)
                host = websocket.headers.get("host", "")
                cookie, _, _ = await self.dsh_cookie(user, sandbox, fqdn, host)
                path = websocket.path_params.get("path", "")
                url = f"ws://{fqdn}:{self.settings.dsh_port}/{path}"
                query = filtered_query(websocket.url.query)
                if query:
                    url = f"{url}?{query}"
                headers = {
                    key: value
                    for key, value in websocket.headers.items()
                    if key.lower() not in OWNED_REQUEST_HEADERS
                    and not key.lower().startswith("sec-websocket-")
                }
                headers["host"] = host
                headers["cookie"] = cookie
                headers["x-forwarded-host"] = host
                headers["x-forwarded-proto"] = websocket.url.scheme
                if websocket.client:
                    headers["x-forwarded-for"] = websocket.client.host
                subprotocols = [
                    part.strip()
                    for part in websocket.headers.get(
                        "sec-websocket-protocol", ""
                    ).split(",")
                    if part.strip()
                ]
                upstream_cm = self.ws_connect(
                    url,
                    additional_headers=headers,
                    subprotocols=subprotocols or None,
                    open_timeout=self.settings.wake_timeout,
                    max_size=None,
                    proxy=None,
                )
                upstream = await upstream_cm.__aenter__()
            except Exception:
                return await self._ws_reject(websocket, 1011)

            try:
                await websocket.accept(
                    subprotocol=getattr(upstream, "subprotocol", None)
                )
            except RuntimeError:
                pass
            await self._pump(websocket, upstream)
            with contextlib.suppress(Exception):
                await websocket.close(
                    code=self._upstream_close_code(upstream)
                )
        finally:
            if upstream_cm is not None:
                with contextlib.suppress(Exception):
                    await upstream_cm.__aexit__(None, None, None)
            await self._release(user)

    @staticmethod
    def _upstream_close_code(upstream) -> int:
        code = getattr(upstream, "close_code", None)
        if isinstance(code, int) and 1000 <= code <= 4999 and code not in (1005, 1006):
            return code
        return 1000

    async def _pump(self, websocket: WebSocket, upstream) -> None:
        async def client_to_upstream():
            try:
                while True:
                    message = await websocket.receive()
                    if message["type"] == "websocket.disconnect":
                        return
                    if message.get("text") is not None:
                        await upstream.send(message["text"])
                    elif message.get("bytes") is not None:
                        await upstream.send(message["bytes"])
            finally:
                with contextlib.suppress(Exception):
                    await upstream.close()

        async def upstream_to_client():
            async for message in upstream:
                if isinstance(message, str):
                    await websocket.send_text(message)
                else:
                    await websocket.send_bytes(message)

        tasks = [
            asyncio.create_task(client_to_upstream()),
            asyncio.create_task(upstream_to_client()),
        ]
        done, pending = await asyncio.wait(
            tasks, return_when=asyncio.FIRST_COMPLETED
        )
        for task in pending:
            task.cancel()
        await asyncio.gather(*done, *pending, return_exceptions=True)

    # -- idle sweeper ------------------------------------------------------

    async def sweep_once(self) -> None:
        claims = await self.k8s.list_claims()
        for claim in claims:
            # Only suspend claims this gateway created.  Everything else in
            # the namespace (resource-layer claims from the README walkthrough,
            # other controllers' objects) is not ours to manage.
            if not gateway_managed(claim):
                continue
            user = claim_user(claim)
            try:
                sandbox = await self._sandbox_of(claim)
            except ApiException:
                continue  # transient/stale claim; retry next cycle
            if sandbox is None or derive_state(sandbox) != "Ready":
                continue
            async with self.state.lock:
                if self.state.maintenance.get(user):
                    continue
                if self.state.in_flight.get(user, 0) > 0:
                    continue  # requests in flight => the user is active
                idle = time.time() - self.state.last_activity.setdefault(
                    user, time.time()
                )
                if idle < self.settings.idle_timeout:
                    continue
                self.state.maintenance[user] = True
            try:
                print(f"suspending {user} (idle {int(idle)}s)", flush=True)
                await self.k8s.patch_operating_mode(
                    sandbox["metadata"]["name"], "Suspended"
                )
                self._clear_handshake(user)
            finally:
                async with self.state.lock:
                    self.state.maintenance[user] = False

    async def sweeper_loop(self) -> None:
        while True:
            await asyncio.sleep(self.settings.sweep_interval)
            try:
                await self.sweep_once()
            except asyncio.CancelledError:
                raise
            except Exception as exc:
                print(
                    f"sweeper: transient error, retrying next cycle: {exc}",
                    flush=True,
                )


# --- Application factory ---------------------------------------------------


def _lifespan(gateway: Gateway, close_http_client: httpx.AsyncClient | None = None):
    @contextlib.asynccontextmanager
    async def lifespan(app):
        task = asyncio.create_task(gateway.sweeper_loop())
        try:
            yield
        finally:
            task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await task
            if close_http_client is not None:
                with contextlib.suppress(Exception):
                    await close_http_client.aclose()

    return lifespan


def create_app(
    settings: Settings | None = None,
    k8s=None,
    http_client: httpx.AsyncClient | None = None,
    ws_connect=None,
) -> Starlette:
    """Build the ASGI app.  All collaborators are injectable for tests."""
    settings = settings or Settings.from_env()
    k8s = k8s if k8s is not None else Kubernetes(settings.namespace)
    owns_http_client = http_client is None
    if http_client is None:
        http_client = httpx.AsyncClient(
            follow_redirects=False, timeout=settings.http_timeout
        )
    gateway = Gateway(
        settings, k8s, http_client, ws_connect or websockets.connect
    )
    routes = [
        Route("/healthz", gateway.healthz),
        Route("/", gateway.landing),
        Route("/users", gateway.signup, methods=["POST"]),
        Route("/login", gateway.login, methods=["POST"]),
        Route("/logout", gateway.logout, methods=["POST"]),
        Route("/users/{user}", gateway.get_user),
        Route("/users/{user}", gateway.delete_user, methods=["DELETE"]),
        Route("/u/{user}", gateway.redirect_to_workspace),
        Route(
            "/u/{user}/{path:path}",
            gateway.proxy_http,
            methods=["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"],
        ),
        WebSocketRoute("/u/{user}/{path:path}", gateway.proxy_ws),
    ]
    app = Starlette(
        routes=routes,
        lifespan=_lifespan(gateway, http_client if owns_http_client else None),
    )
    app.state.gateway = gateway
    return app