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

"""Unit tests for the browser-usable DSH gateway.

Everything here runs without a cluster: the Kubernetes adapter, the upstream
httpx client and the WebSocket connector are injected fakes.  The HTTP tests
drive the real ASGI app through ``httpx.ASGITransport``; the WebSocket tests
use Starlette's ``TestClient`` (which runs the lifespan) against a fake
``ws_connect`` async context manager.
"""

import asyncio
import contextlib
import hashlib
import json
import threading
import time
from types import SimpleNamespace

import httpx
import pytest
from kubernetes.client.exceptions import ApiException
from starlette.testclient import TestClient
from starlette.websockets import WebSocketDisconnect

from gateway import (
    Settings,
    UserExists,
    claim_name,
    claim_user,
    condition,
    create_app,
    derive_state,
    extract_dsh_cookies,
    filtered_query,
    gateway_managed,
    parse_dsh_launch_token,
    rewrite_location,
    same_origin,
    token_matches,
)

TOKEN = "s3cret-token"
POD_UID = "pod-uid-1"
POD_IP = "10.0.0.7"
FQDN = "dsh-alice-sbx.test-ns.svc.cluster.local"


def dsh_response(status=200, content=b"", headers=None):
    """A MockTransport response that is a real stream (``aiter_raw`` works)."""
    return httpx.Response(status, headers=headers, stream=httpx.ByteStream(content))


# --- fakes -----------------------------------------------------------------


class FakeK8s:
    def __init__(self):
        self.claims = {}
        self.sandboxes = {}
        self.pods = {}
        self.logs = {}
        self.created = []
        self.deleted = []
        self.patches = []
        self.token_patches = []
        self.conflict_on_create = False
        self.get_claim_error = None
        self.patch_error = None

    async def get_claim(self, user):
        if self.get_claim_error is not None:
            raise self.get_claim_error
        return self.claims.get(user)

    async def create_claim(self, user, token_hash, pool):
        if self.conflict_on_create:
            raise UserExists(user)
        self.created.append((user, token_hash, pool))
        self.claims.setdefault(user, {"metadata": {"name": claim_name(user)}})

    async def delete_claim(self, user):
        self.deleted.append(user)
        self.claims.pop(user, None)

    async def list_claims(self):
        return list(self.claims.values())

    async def get_sandbox(self, name):
        return self.sandboxes.get(name)

    async def patch_operating_mode(self, name, mode):
        if self.patch_error is not None:
            raise self.patch_error
        self.patches.append((name, mode))

    async def patch_token(self, user, token_hash):
        if self.patch_error is not None:
            raise self.patch_error
        self.token_patches.append((user, token_hash))
        claim = self.claims.get(user)
        if claim is not None:
            annotations = claim.setdefault("metadata", {}).setdefault("annotations", {})
            annotations["aaas.example.com/token-sha256"] = token_hash

    async def get_pod(self, name):
        return self.pods.get(name)

    async def read_pod_logs(self, name, container, tail_lines):
        return self.logs.get(name, "")


class FakeUpstream:
    """An httpx.MockTransport handler for both the handshake and the proxy."""

    def __init__(self):
        self.requests = []
        self.handshake_requests = []
        self.handshake_status = 303
        self.handshake_cookies = ["dsh-auth-a=one; Path=/; HttpOnly; SameSite=Strict"]
        self.proxy_handler = None
        self.proxy_requests = []

    def __call__(self, request: httpx.Request) -> httpx.Response:
        self.requests.append(request)
        if request.url.path == "/" and "token=" in str(request.url.query):
            self.handshake_requests.append(request)
            cookies = self.handshake_cookies
            index = min(len(self.handshake_requests) - 1, len(cookies) - 1)
            return dsh_response(
                self.handshake_status,
                headers=[
                    ("location", "/"),
                    ("set-cookie", cookies[index]),
                ],
            )
        self.proxy_requests.append(request)
        if self.proxy_handler is not None:
            return self.proxy_handler(request)
        return dsh_response(200, b"hello from DSH", {"content-type": "text/plain"})


class FakeWS:
    """Async context manager matching websockets' ClientConnection subset."""

    def __init__(self):
        self.sent = []
        self.closed = False
        self.close_code = None
        self.subprotocol = None
        self._stage = 0
        self._sent_event = asyncio.Event()

    async def send(self, data):
        self.sent.append(data)
        self._sent_event.set()

    async def close(self, code=1000):
        self.closed = True
        self.close_code = code

    def __aiter__(self):
        return self

    async def __anext__(self):
        if self._stage == 0:
            self._stage = 1
            return "welcome"
        if self._stage == 1:
            await self._sent_event.wait()
            self._stage = 2
            return "pong"
        raise StopAsyncIteration

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc_info):
        return False


class FakeWSConnect:
    def __init__(self):
        self.calls = []
        self.instances = []
        self.error = None

    def __call__(self, url, **kwargs):
        self.calls.append((url, kwargs))
        if self.error is not None:
            raise self.error
        ws = FakeWS()
        self.instances.append(ws)
        return ws


def add_user(k8s, user="alice", token=TOKEN, ready=True, pod_name=None, suffix=""):
    """Populate the fake API objects for a Ready (or Waking) user."""
    sandbox_name = f"dsh-{user}-sbx{suffix}"
    pod_name = pod_name or sandbox_name
    k8s.claims[user] = {
        "metadata": {
            "name": claim_name(user),
            "annotations": {
                "aaas.example.com/token-sha256": hashlib.sha256(token.encode()).hexdigest()
            },
        },
        "status": {"sandbox": {"name": sandbox_name}},
    }
    k8s.sandboxes[sandbox_name] = {
        "metadata": {"name": sandbox_name, "annotations": {}},
        "spec": {},
        "status": {
            "conditions": [{"type": "Ready", "status": "True"}] if ready else [],
            "serviceFQDN": FQDN,
            "podIPs": [POD_IP],
        },
    }
    k8s.pods[pod_name] = {"metadata": {"name": pod_name, "uid": POD_UID}}
    k8s.logs[pod_name] = f"booting\ndsh web: http://10.0.0.7:3080/?token={TOKEN}-launch\n"
    return sandbox_name, pod_name


@pytest.fixture
def env():
    settings = Settings(
        namespace="test-ns",
        pool="test-pool",
        idle_timeout=60,
        wake_timeout=1,
        sweep_interval=3600,
        upstream_retry_delay=0,
    )
    k8s = FakeK8s()
    upstream = FakeUpstream()
    http_client = httpx.AsyncClient(transport=httpx.MockTransport(upstream), timeout=5)
    ws_connect = FakeWSConnect()
    app = create_app(
        settings=settings,
        k8s=k8s,
        http_client=http_client,
        ws_connect=ws_connect,
    )
    return SimpleNamespace(
        app=app,
        gateway=app.state.gateway,
        settings=settings,
        k8s=k8s,
        upstream=upstream,
        ws_connect=ws_connect,
        http_client=http_client,
    )


def asgi_client(env, cookies=None):
    client = httpx.AsyncClient(
        transport=httpx.ASGITransport(app=env.app),
        base_url="http://gateway.test",
    )
    if cookies:
        client.cookies.update(cookies)
    return client


def stale_handshake(env):
    env.gateway.state.handshake[("alice", POD_UID, "gateway.test")] = "dsh-auth-a=stale"


def raw_scope(method, path, headers=()):
    """An ASGI HTTP scope for driving the app without an httpx client.

    ``spec_version`` matches what uvicorn 0.54 advertises (2.3), which is what
    puts ``StreamingResponse`` on its disconnect-listening path.
    """
    return {
        "type": "http",
        "asgi": {"version": "3.0", "spec_version": "2.3"},
        "http_version": "1.1",
        "method": method,
        "scheme": "http",
        "path": path,
        "raw_path": path.encode(),
        "query_string": b"",
        "root_path": "",
        "server": ("gateway.test", 80),
        "client": ("10.0.0.9", 4321),
        "headers": [(b"host", b"gateway.test"), *headers],
    }


async def asgi_call(env, scope, receive):
    """Drive the app with a scripted ``receive`` and return the sent messages.

    An exception escaping the app fails the test, which is exactly how a
    ``ClientDisconnect``-induced traceback used to surface.
    """
    sent = []

    async def send(message):
        sent.append(message)

    await env.app(scope, receive, send)
    return sent


# --- pure helpers ----------------------------------------------------------


def test_condition_true_when_status_matches():
    obj = {"status": {"conditions": [{"type": "Ready", "status": "True"}]}}
    assert condition(obj, "Ready") is True


def test_condition_false_when_status_is_false():
    obj = {"status": {"conditions": [{"type": "Ready", "status": "False"}]}}
    assert condition(obj, "Ready") is False


def test_condition_false_when_type_absent():
    obj = {"status": {"conditions": [{"type": "Suspended", "status": "True"}]}}
    assert condition(obj, "Ready") is False


def test_condition_false_when_no_status():
    assert condition({}, "Ready") is False


def test_derive_state_provisioning_when_sandbox_missing():
    assert derive_state(None) == "Provisioning"


def test_derive_state_ready():
    sandbox = {"spec": {}, "status": {"conditions": [{"type": "Ready", "status": "True"}]}}
    assert derive_state(sandbox) == "Ready"


def test_derive_state_waking_when_not_yet_ready():
    sandbox = {"spec": {"operatingMode": "Running"}, "status": {"conditions": []}}
    assert derive_state(sandbox) == "Waking"


def test_derive_state_suspended():
    sandbox = {
        "spec": {"operatingMode": "Suspended"},
        "status": {"conditions": [{"type": "Suspended", "status": "True"}]},
    }
    assert derive_state(sandbox) == "Suspended"


def test_derive_state_suspending_while_mode_flipped_but_not_settled():
    sandbox = {"spec": {"operatingMode": "Suspended"}, "status": {"conditions": []}}
    assert derive_state(sandbox) == "Suspending"


def test_claim_name_and_user_roundtrip():
    assert claim_name("alice") == "dsh-alice"
    assert claim_user({"metadata": {"name": "dsh-alice"}}) == "alice"


def test_gateway_managed_requires_ownership_annotation():
    managed = {"metadata": {"annotations": {"aaas.example.com/managed-by": "dsh-gateway"}}}
    manual = {"metadata": {"name": "dsh-alice"}}
    other = {"metadata": {"annotations": {"aaas.example.com/managed-by": "someone"}}}
    assert gateway_managed(managed) is True
    assert gateway_managed(manual) is False
    assert gateway_managed(other) is False


def test_parse_dsh_launch_token_from_real_log_line():
    log = (
        "2026-10-21T10:00:00Z starting\n"
        "dsh web: http://10.244.1.7:3080/?token=abc-_123XYZ\n"
        "some later line\n"
    )
    assert parse_dsh_launch_token(log) == "abc-_123XYZ"


def test_parse_dsh_launch_token_uses_last_match():
    log = (
        "dsh web: http://10.0.0.1:3080/?token=old\n"
        "restart\n"
        "dsh web: http://10.0.0.2:3080/?token=new\n"
    )
    assert parse_dsh_launch_token(log) == "new"


@pytest.mark.parametrize(
    "log",
    ["", "no token here", "dsh web: http://10.0.0.1:3080/?token=", "dsh web: garbage"],
)
def test_parse_dsh_launch_token_returns_none_without_token(log):
    assert parse_dsh_launch_token(log) is None


def test_extract_dsh_cookies_keeps_only_dsh_auth_pairs():
    headers = [
        "dsh-auth-a=one; Path=/; HttpOnly; SameSite=Strict",
        "other=two; Path=/",
        "dsh-auth-b=three; Path=/",
    ]
    assert extract_dsh_cookies(headers) == "dsh-auth-a=one; dsh-auth-b=three"


def test_same_origin_origin_header():
    assert same_origin({"host": "gw.test", "origin": "http://gw.test"}) is True
    assert same_origin({"host": "gw.test", "origin": "http://evil.test"}) is False


def test_same_origin_referer_fallback_and_missing():
    assert same_origin({"host": "gw.test", "referer": "http://gw.test/u/a/"}) is True
    assert same_origin({"host": "gw.test"}) is False


def test_rewrite_location_prefixes_absolute_paths_only():
    assert rewrite_location("/", "alice") == "/u/alice/"
    assert rewrite_location("/assets/app.js", "alice") == "/u/alice/assets/app.js"
    assert rewrite_location("http://x/absolute", "alice") == "http://x/absolute"
    assert rewrite_location("//other/path", "alice") == "//other/path"
    assert rewrite_location("/u/alice/already", "alice") == "/u/alice/already"


def test_rewrite_location_drops_launch_token():
    assert rewrite_location("/?token=secret", "alice") == "/u/alice/"


def test_filtered_query_removes_token_only():
    assert filtered_query("a=1&token=x&b=2&token=y") == "a=1&b=2"


def test_filtered_query_preserves_vite_combined_module_url():
    """DSH's ``plugins/??<specifiers>`` query must reach the upstream verbatim."""
    raw = "?@deepseek-ai/dsh-client-modules/client.js&rev=a1bf7d6a56f9"
    assert filtered_query(raw) == raw
    assert filtered_query(raw + "&token=x") == raw


def test_filtered_query_drops_an_encoded_token_parameter():
    assert filtered_query("a=1&to%6Ben=x&b=2") == "a=1&b=2"


def test_token_matches_is_timing_safe_hash_check():
    claim = {
        "metadata": {
            "annotations": {
                "aaas.example.com/token-sha256": hashlib.sha256(TOKEN.encode()).hexdigest()
            }
        }
    }
    assert token_matches(claim, TOKEN) is True
    assert token_matches(claim, "wrong") is False
    assert token_matches(claim, "") is False
    assert token_matches({"metadata": {}}, TOKEN) is False


def test_settings_from_env(monkeypatch):
    monkeypatch.setenv("NAMESPACE", "ns-x")
    monkeypatch.setenv("POOL", "pool-x")
    monkeypatch.setenv("IDLE_TIMEOUT", "7")
    monkeypatch.setenv("WAKE_TIMEOUT", "8")
    monkeypatch.setenv("DSH_PORT", "9999")
    monkeypatch.setenv("DSH_LOG_TAIL_LINES", "42")
    monkeypatch.setenv("COOKIE_NAME", "sid")
    monkeypatch.setenv("COOKIE_SECURE", "true")
    monkeypatch.setenv("COOKIE_MAX_AGE", "100")
    monkeypatch.setenv("MAX_BODY_MB", "2")
    monkeypatch.setenv("HTTP_TIMEOUT", "12.5")
    monkeypatch.setenv("AUTH_PROXY_HEADER", "X-Forwarded-User")
    settings = Settings.from_env()
    assert settings.namespace == "ns-x"
    assert settings.pool == "pool-x"
    assert settings.idle_timeout == 7
    assert settings.wake_timeout == 8
    assert settings.dsh_port == 9999
    assert settings.dsh_log_tail_lines == 42
    assert settings.cookie_name == "sid"
    assert settings.cookie_secure is True
    assert settings.cookie_max_age == 100
    assert settings.max_body_bytes == 2 * 1024 * 1024
    assert settings.http_timeout == 12.5
    assert settings.auth_proxy_header == "X-Forwarded-User"


def test_settings_auth_proxy_header_defaults_to_none(monkeypatch):
    monkeypatch.delenv("AUTH_PROXY_HEADER", raising=False)
    assert Settings.from_env().auth_proxy_header is None


# --- Kubernetes adapter ----------------------------------------------------


@pytest.mark.asyncio
async def test_kubernetes_create_claim_body(monkeypatch):
    from gateway import Kubernetes

    calls = []

    async def fake_call(which, method, *args):
        calls.append((which, method, args))

    k8s = Kubernetes("test-ns")
    monkeypatch.setattr(k8s, "_call", fake_call)
    await k8s.create_claim("alice", "hash", "test-pool")

    which, method, args = calls[0]
    assert which == "crd"
    assert method == "create_namespaced_custom_object"
    body = args[4]
    assert body["metadata"]["name"] == "dsh-alice"
    assert body["metadata"]["annotations"]["aaas.example.com/token-sha256"] == "hash"
    assert body["metadata"]["annotations"]["aaas.example.com/managed-by"] == "dsh-gateway"
    assert body["spec"]["warmPoolRef"]["name"] == "test-pool"
    assert body["spec"]["additionalPodMetadata"]["labels"]["sandbox.users.io/dsh-user"] == "alice"


@pytest.mark.asyncio
async def test_kubernetes_maps_404_to_none_and_propagates_other_errors(monkeypatch):
    from gateway import Kubernetes

    k8s = Kubernetes("test-ns")
    calls = {"n": 0}

    async def fake_call(which, method, *args):
        calls["n"] += 1
        raise ApiException(status=404 if calls["n"] == 1 else 500)

    monkeypatch.setattr(k8s, "_call", fake_call)
    assert await k8s.get_claim("alice") is None
    with pytest.raises(ApiException):
        await k8s.get_claim("alice")


@pytest.mark.asyncio
async def test_kubernetes_read_pod_logs_forwards_container_and_tail(monkeypatch):
    from gateway import Kubernetes

    calls = []

    class Core:
        def read_namespaced_pod_log(self, *args, **kwargs):
            calls.append((args, kwargs))
            return "log line"

    k8s = Kubernetes("test-ns")
    monkeypatch.setattr(k8s, "_clients", lambda: (None, Core()))
    assert await k8s.read_pod_logs("pod-1", "dsh-web", 200) == "log line"
    assert calls == [(("pod-1", "test-ns"), {"container": "dsh-web", "tail_lines": 200})]


@pytest.mark.asyncio
async def test_kubernetes_get_pod_normalises_the_v1pod_model(monkeypatch):
    """CoreV1Api returns a V1Pod model; callers index it like a dict."""
    from kubernetes.client import V1ObjectMeta, V1Pod

    from gateway import Kubernetes

    async def fake_call(which, method, *args):
        return V1Pod(metadata=V1ObjectMeta(name="dsh-alice", uid="pod-uid-9"))

    k8s = Kubernetes("test-ns")
    monkeypatch.setattr(k8s, "_call", fake_call)

    pod = await k8s.get_pod("dsh-alice")

    assert isinstance(pod, dict)
    assert pod["metadata"]["uid"] == "pod-uid-9"


def test_kubernetes_clients_initialise_once_under_concurrency(monkeypatch):
    """_call runs on a worker pool: the lazy init must not hand out a None client."""
    import gateway as gw
    from gateway import Kubernetes

    created = []

    class Crd:
        def __init__(self):
            created.append("crd")

    class SlowCore:
        def __init__(self):
            time.sleep(0.05)  # widen the window between _crd and _core
            created.append("core")

    monkeypatch.setattr(gw.k8s_client, "CustomObjectsApi", Crd)
    monkeypatch.setattr(gw.k8s_client, "CoreV1Api", SlowCore)
    monkeypatch.setattr(gw.k8s_config, "load_incluster_config", lambda: None)

    k8s = Kubernetes("test-ns")
    results = []

    def worker():
        results.append(k8s._clients())

    threads = [threading.Thread(target=worker) for _ in range(8)]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()

    assert all(core is not None for _, core in results)
    assert created.count("core") == 1


# --- health / landing / redirect ------------------------------------------


@pytest.mark.asyncio
async def test_healthz_needs_no_auth_and_no_cluster(env):
    async with asgi_client(env) as client:
        response = await client.get("/healthz")
    assert response.status_code == 200
    assert response.json() == {"status": "ok"}


@pytest.mark.asyncio
async def test_healthz_does_not_touch_the_cluster_at_startup():
    # No k8s injected: create_app builds the lazy Kubernetes adapter, which
    # must not load credentials until a claim/pod call actually happens.
    app = create_app(
        settings=Settings(namespace="test-ns"),
        http_client=httpx.AsyncClient(transport=httpx.MockTransport(lambda r: dsh_response())),
        ws_connect=lambda *a, **k: None,
    )
    async with httpx.AsyncClient(
        transport=httpx.ASGITransport(app=app),
        base_url="http://gateway.test",
    ) as client:
        response = await client.get("/healthz")
    assert response.status_code == 200


@pytest.mark.asyncio
async def test_landing_page_has_signup_and_login_forms(env):
    async with asgi_client(env) as client:
        response = await client.get("/")
    assert response.status_code == 200
    assert "text/html" in response.headers["content-type"]
    assert 'action="/users"' in response.text
    assert 'action="/login"' in response.text


@pytest.mark.asyncio
async def test_landing_page_shows_session_when_cookie_present(env):
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/")
    assert "Signed in as" in response.text
    assert "/u/alice/" in response.text


@pytest.mark.asyncio
async def test_workspace_redirect_adds_trailing_slash(env):
    async with asgi_client(env) as client:
        response = await client.get("/u/alice?x=1")
    assert response.status_code == 308
    assert response.headers["location"] == "/u/alice/?x=1"


# --- signup / login / logout ----------------------------------------------


@pytest.mark.asyncio
@pytest.mark.parametrize("bad_user", ["Alice", "-alice", "alice-", "", "a_b", 123, None, "a" * 60])
async def test_signup_rejects_invalid_dns1123_label(env, bad_user):
    async with asgi_client(env) as client:
        response = await client.post("/users", json={"user": bad_user})
    assert response.status_code == 400
    assert env.k8s.created == []


@pytest.mark.asyncio
async def test_signup_conflict_when_claim_exists(env):
    add_user(env.k8s)
    async with asgi_client(env) as client:
        response = await client.post("/users", json={"user": "alice"})
    assert response.status_code == 409
    assert env.k8s.created == []


@pytest.mark.asyncio
async def test_signup_conflict_on_concurrent_race(env):
    env.k8s.conflict_on_create = True
    async with asgi_client(env) as client:
        response = await client.post("/users", json={"user": "alice"})
    assert response.status_code == 409


@pytest.mark.asyncio
async def test_signup_json_stores_only_hash_and_returns_token(env, monkeypatch):
    sandbox = {
        "metadata": {"name": "dsh-alice-sbx"},
        "spec": {},
        "status": {"conditions": [{"type": "Ready", "status": "True"}]},
    }

    async def ready(user, timeout):
        return sandbox

    monkeypatch.setattr(env.gateway, "wait_ready", ready)
    async with asgi_client(env) as client:
        response = await client.post("/users", json={"user": "alice"})

    assert response.status_code == 201
    body = response.json()
    assert body["user"] == "alice"
    assert body["state"] == "Ready"
    assert body["sandbox"] == "dsh-alice-sbx"
    assert body["token"]
    user, token_hash, pool = env.k8s.created[0]
    assert user == "alice"
    assert pool == "test-pool"
    assert token_hash == hashlib.sha256(body["token"].encode()).hexdigest()
    assert token_hash != body["token"]
    # The JSON API also gets a browser session cookie (HttpOnly, SameSite=Lax).
    cookie = response.headers["set-cookie"]
    assert "dsh_gw_session=alice:" in cookie
    assert "HttpOnly" in cookie
    assert "SameSite=lax" in cookie


@pytest.mark.asyncio
async def test_signup_returns_202_when_sandbox_not_ready(env, monkeypatch):
    async def not_ready(user, timeout):
        return None

    monkeypatch.setattr(env.gateway, "wait_ready", not_ready)
    async with asgi_client(env) as client:
        response = await client.post("/users", json={"user": "alice"})
    assert response.status_code == 202
    assert response.json()["state"] == "Provisioning"


@pytest.mark.asyncio
async def test_signup_form_redirects_and_sets_session_cookie(env, monkeypatch):
    async def not_ready(user, timeout):
        return None

    monkeypatch.setattr(env.gateway, "wait_ready", not_ready)
    async with asgi_client(env) as client:
        response = await client.post(
            "/users",
            data={"user": "bob"},
            headers={"content-type": "application/x-www-form-urlencoded"},
        )
    assert response.status_code == 303
    assert response.headers["location"] == "/u/bob/"
    cookie = response.headers["set-cookie"]
    assert re_cookie_user(cookie) == "bob"


def re_cookie_user(set_cookie: str) -> str:
    value = set_cookie.split(";", 1)[0].split("=", 1)[1]
    return value.split(":", 1)[0]


@pytest.mark.asyncio
async def test_signup_body_limit_returns_413(env):
    env.settings = Settings(**{**env.settings.__dict__, "max_body_bytes": 8})
    env.gateway.settings = env.settings
    async with asgi_client(env) as client:
        response = await client.post(
            "/users",
            content=json.dumps({"user": "alice"}),
            headers={"content-type": "application/json"},
        )
    assert response.status_code == 413


@pytest.mark.asyncio
async def test_login_form_sets_cookie_and_redirects(env):
    add_user(env.k8s)
    async with asgi_client(env) as client:
        response = await client.post(
            "/login",
            data={"user": "alice", "token": TOKEN},
            headers={"content-type": "application/x-www-form-urlencoded"},
        )
    assert response.status_code == 303
    assert response.headers["location"] == "/u/alice/"
    assert re_cookie_user(response.headers["set-cookie"]) == "alice"


@pytest.mark.asyncio
async def test_login_json_and_failures(env):
    add_user(env.k8s)
    async with asgi_client(env) as client:
        ok = await client.post("/login", json={"user": "alice", "token": TOKEN})
        bad = await client.post("/login", json={"user": "alice", "token": "nope"})
        missing = await client.post("/login", json={"user": "ghost", "token": TOKEN})
    assert ok.status_code == 200
    assert bad.status_code == 401
    assert missing.status_code == 404


@pytest.mark.asyncio
async def test_logout_clears_cookie(env):
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.post("/logout")
    assert response.status_code == 303
    assert response.headers["location"] == "/"
    assert "dsh_gw_session=" in response.headers["set-cookie"]
    assert "Max-Age=0" in response.headers["set-cookie"]


# --- auth ------------------------------------------------------------------


@pytest.mark.asyncio
async def test_get_user_requires_valid_bearer_or_cookie(env):
    add_user(env.k8s)
    async with asgi_client(env) as client:
        forbidden = await client.get("/users/alice", headers={"Authorization": "Bearer wrong"})
        bearer = await client.get("/users/alice", headers={"Authorization": f"Bearer {TOKEN}"})
        cookie = await client.get(
            "/users/alice",
            headers={"Cookie": f"dsh_gw_session=alice:{TOKEN}"},
        )
        wrong_user_cookie = await client.get(
            "/users/alice",
            headers={"Cookie": f"dsh_gw_session=bob:{TOKEN}"},
        )
    assert forbidden.status_code == 401
    assert bearer.status_code == 200
    assert bearer.json()["state"] == "Ready"
    assert cookie.status_code == 200
    assert wrong_user_cookie.status_code == 401


@pytest.mark.asyncio
async def test_get_user_not_found(env):
    async with asgi_client(env) as client:
        response = await client.get("/users/alice")
    assert response.status_code == 404


@pytest.mark.asyncio
@pytest.mark.parametrize("path", ["/users/Alice", "/u/Alice/", "/u/Alice"])
async def test_invalid_user_in_path_is_404(env, path):
    async with asgi_client(env) as client:
        response = await client.get(path)
    assert response.status_code == 404
    assert env.k8s.claims == {}


@pytest.mark.asyncio
async def test_delete_user_requires_auth(env):
    add_user(env.k8s)
    async with asgi_client(env) as client:
        response = await client.delete("/users/alice", headers={"Authorization": "Bearer wrong"})
    assert response.status_code == 401
    assert env.k8s.deleted == []


@pytest.mark.asyncio
async def test_delete_with_cookie_requires_same_origin(env):
    add_user(env.k8s)
    env.gateway.state.handshake[("alice", POD_UID, "gateway.test")] = "dsh-auth-a=x"
    cookies = {"dsh_gw_session": f"alice:{TOKEN}"}
    async with asgi_client(env, cookies=cookies) as client:
        missing = await client.delete("/users/alice")
        cross = await client.delete("/users/alice", headers={"Origin": "http://evil.test"})
        good = await client.delete("/users/alice", headers={"Origin": "http://gateway.test"})
    assert missing.status_code == 403
    assert cross.status_code == 403
    assert good.status_code == 200
    assert env.k8s.deleted == ["alice"]
    assert env.gateway.state.handshake == {}


@pytest.mark.asyncio
async def test_delete_with_bearer_does_not_need_origin(env):
    add_user(env.k8s)
    async with asgi_client(env) as client:
        response = await client.delete("/users/alice", headers={"Authorization": f"Bearer {TOKEN}"})
    assert response.status_code == 200


@pytest.mark.asyncio
async def test_delete_missing_claim_is_404(env):
    add_user(env.k8s)
    env.k8s.claims.clear()
    async with asgi_client(env) as client:
        response = await client.delete("/users/alice", headers={"Authorization": f"Bearer {TOKEN}"})
    assert response.status_code == 404


@pytest.mark.asyncio
async def test_unexpected_api_error_propagates(env):
    add_user(env.k8s)
    env.k8s.get_claim_error = ApiException(status=500)
    async with asgi_client(env) as client:
        with pytest.raises(ApiException):
            await client.get("/users/alice")


# --- DSH handshake ---------------------------------------------------------


@pytest.mark.asyncio
async def test_proxy_handshakes_and_uses_browser_host(env):
    add_user(env.k8s)
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/u/alice/api/state")
    assert response.status_code == 200
    assert response.text == "hello from DSH"
    handshake = env.upstream.handshake_requests[0]
    assert handshake.url.host == FQDN
    assert handshake.url.port == 3080
    assert handshake.url.params["token"] == TOKEN + "-launch"
    assert handshake.headers["host"] == "gateway.test"
    proxy = env.upstream.proxy_requests[0]
    assert proxy.headers["cookie"] == "dsh-auth-a=one"


@pytest.mark.asyncio
async def test_dsh_cookie_is_cached_per_pod_uid_and_host(env):
    add_user(env.k8s)
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        first = await client.get("/u/alice/api/one")
        second = await client.get("/u/alice/api/two")
    assert first.status_code == second.status_code == 200
    assert len(env.upstream.handshake_requests) == 1
    assert len(env.upstream.proxy_requests) == 2


@pytest.mark.asyncio
async def test_missing_launch_token_returns_503_with_retry_after(env):
    _sandbox_name, pod_name = add_user(env.k8s)
    env.k8s.logs[pod_name] = "no token printed yet\n"
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/u/alice/api/state")
    assert response.status_code == 503
    assert response.headers["retry-after"] == "10"
    assert "launch token" in response.json()["error"]


@pytest.mark.asyncio
async def test_handshake_4xx_is_503(env):
    add_user(env.k8s)
    env.upstream.handshake_status = 401
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/u/alice/api/state")
    assert response.status_code == 503


@pytest.mark.asyncio
async def test_missing_cookie_in_handshake_is_503(env):
    add_user(env.k8s)
    env.upstream.handshake_cookies = ["other=1; Path=/"]
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/u/alice/api/state")
    assert response.status_code == 503


@pytest.mark.asyncio
async def test_pod_annotation_overrides_sandbox_name(env):
    sandbox_name, pod_name = add_user(env.k8s)
    env.k8s.sandboxes[sandbox_name]["metadata"]["annotations"]["agents.x-k8s.io/pod-name"] = (
        "custom-pod"
    )
    env.k8s.pods["custom-pod"] = {"metadata": {"uid": "other-uid"}}
    env.k8s.logs["custom-pod"] = env.k8s.logs[pod_name]
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/u/alice/api/state")
    assert response.status_code == 200
    assert ("alice", "other-uid", "gateway.test") in env.gateway.state.handshake


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


@pytest.mark.asyncio
async def test_proxy_strips_gateway_auth_cookie_and_hop_headers(env):
    add_user(env.k8s)

    def handler(request):
        return dsh_response(
            200,
            b"ok",
            {
                "content-type": "text/plain",
                "connection": "keep-alive",
                "set-cookie": "upstream=leak; Path=/",
            },
        )

    env.upstream.proxy_handler = handler
    headers = {
        "Authorization": f"Bearer {TOKEN}",
        "Cookie": "dsh_gw_session=alice:x; browser=1",
        "X-Forwarded-For": "1.2.3.4",
        "X-Forwarded-Proto": "https",
        "Host": "gateway.test",
        "Origin": "http://gateway.test",
    }
    async with asgi_client(env) as client:
        response = await client.get("/u/alice/api/foo?token=secret&x=1", headers=headers)

    assert response.status_code == 200
    proxy = env.upstream.proxy_requests[0]
    assert proxy.url.path == "/api/foo"
    assert proxy.url.query.decode() == "x=1"
    assert "authorization" not in proxy.headers
    assert proxy.headers["cookie"] == "dsh-auth-a=one"
    assert "browser=1" not in proxy.headers["cookie"]
    assert proxy.headers["host"] == "gateway.test"
    assert proxy.headers["origin"] == "http://gateway.test"
    assert proxy.headers["x-forwarded-for"] == "127.0.0.1"
    assert proxy.headers["x-forwarded-proto"] == "http"
    # Upstream Set-Cookie and hop-by-hop headers never reach the browser.
    assert "set-cookie" not in response.headers
    assert "connection" not in response.headers


@pytest.mark.asyncio
async def test_proxy_rewrites_absolute_path_location(env):
    add_user(env.k8s)
    env.upstream.proxy_handler = lambda request: dsh_response(
        303,
        headers=[
            ("location", "/?token=leak&next=1"),
            ("set-cookie", "upstream=leak; Path=/"),
        ],
    )
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/u/alice/api/redirect", follow_redirects=False)
    assert response.status_code == 303
    assert response.headers["location"] == "/u/alice/?next=1"
    assert "set-cookie" not in response.headers


@pytest.mark.asyncio
async def test_proxy_refreshes_cookie_once_on_upstream_401(env):
    add_user(env.k8s)
    env.upstream.handshake_cookies = ["dsh-auth-a=new; Path=/"]
    stale_handshake(env)
    calls = []

    def handler(request):
        calls.append(request)
        if len(calls) == 1:
            return dsh_response(401, b"expired")
        return dsh_response(200, b"fresh")

    env.upstream.proxy_handler = handler
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/u/alice/api/state")
    assert response.status_code == 200
    assert response.text == "fresh"
    assert len(calls) == 2
    assert len(env.upstream.handshake_requests) == 1  # only the forced refresh
    assert calls[1].headers["cookie"] == "dsh-auth-a=new"


@pytest.mark.asyncio
async def test_proxy_does_not_refresh_401_without_cached_cookie(env):
    add_user(env.k8s)
    env.upstream.proxy_handler = lambda request: dsh_response(401, b"no")
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/u/alice/api/state")
    assert response.status_code == 401
    assert len(env.upstream.proxy_requests) == 1


@pytest.mark.asyncio
async def test_proxy_upstream_connection_error_is_502(env):
    add_user(env.k8s)

    def handler(request):
        raise httpx.ConnectError("boom", request=request)

    env.upstream.proxy_handler = handler
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/u/alice/api/state")
    assert response.status_code == 502
    assert env.gateway.state.in_flight["alice"] == 0


@pytest.mark.asyncio
async def test_proxy_body_limit_returns_413(env):
    add_user(env.k8s)
    env.gateway.settings = Settings(**{**env.settings.__dict__, "max_body_bytes": 4})
    async with asgi_client(env) as client:
        response = await client.post(
            "/u/alice/api/echo",
            content=b"too large",
            headers={"Authorization": f"Bearer {TOKEN}"},
        )
    assert response.status_code == 413
    assert env.gateway.state.in_flight["alice"] == 0


@pytest.mark.asyncio
async def test_proxy_client_disconnect_mid_body_returns_499(env):
    """An aborted upload is answered 499 quietly and never reaches DSH."""
    add_user(env.k8s)
    messages = [
        {"type": "http.request", "body": b"partial", "more_body": True},
        {"type": "http.disconnect"},
    ]

    async def receive():
        return messages.pop(0)

    sent = await asgi_call(
        env,
        raw_scope(
            "POST",
            "/u/alice/api/echo",
            headers=[
                (b"origin", b"http://gateway.test"),
                (b"cookie", f"dsh_gw_session=alice:{TOKEN}".encode()),
                (b"content-type", b"application/json"),
            ],
        ),
        receive,
    )
    starts = [message for message in sent if message["type"] == "http.response.start"]
    assert [start["status"] for start in starts] == [499]
    assert env.upstream.proxy_requests == []
    assert env.gateway.state.in_flight["alice"] == 0


@pytest.mark.asyncio
async def test_signup_client_disconnect_mid_body_returns_499(env):
    """Body-reading control routes get the same treatment, not just the proxy."""
    messages = [
        {"type": "http.request", "body": b'{"user": "ali', "more_body": True},
        {"type": "http.disconnect"},
    ]

    async def receive():
        return messages.pop(0)

    sent = await asgi_call(
        env,
        raw_scope("POST", "/users", headers=[(b"content-type", b"application/json")]),
        receive,
    )
    starts = [message for message in sent if message["type"] == "http.response.start"]
    assert [start["status"] for start in starts] == [499]
    assert env.k8s.claims == {}


@pytest.mark.asyncio
async def test_proxy_client_disconnect_mid_stream_frees_slot(env):
    """Closing an open SSE-style stream frees the slot without a traceback."""
    add_user(env.k8s)
    first_chunk = asyncio.Event()

    class SlowStream(httpx.AsyncByteStream):
        async def __aiter__(self):
            yield b"chunk-one"
            first_chunk.set()
            await asyncio.Event().wait()  # the upstream stream stays open

    env.upstream.proxy_handler = lambda request: httpx.Response(
        200, headers={"content-type": "text/event-stream"}, stream=SlowStream()
    )
    messages = [{"type": "http.request", "body": b"", "more_body": False}]

    async def receive():
        if messages:
            return messages.pop(0)
        await first_chunk.wait()
        return {"type": "http.disconnect"}

    sent = await asgi_call(
        env,
        raw_scope(
            "GET",
            "/u/alice/api/stream",
            headers=[(b"cookie", f"dsh_gw_session=alice:{TOKEN}".encode())],
        ),
        receive,
    )
    starts = [message for message in sent if message["type"] == "http.response.start"]
    assert [start["status"] for start in starts] == [200]
    assert env.gateway.state.in_flight["alice"] == 0


@pytest.mark.asyncio
async def test_release_frees_slot_even_when_cancelled(env):
    """A cancelled request must still return its in-flight slot.

    Starlette cancels the streaming task through an anyio cancel scope when the
    client disconnects, and the cancellation can land inside ``_relay``'s
    ``finally``.  Waiting for ``state.lock`` there would strand the slot, which
    keeps the sandbox looking busy forever.
    """
    env.gateway.state.in_flight["alice"] = 1
    await env.gateway.state.lock.acquire()
    try:
        task = asyncio.create_task(env.gateway._release("alice"))
        await asyncio.sleep(0)
        task.cancel()
        with contextlib.suppress(asyncio.CancelledError):
            await task
    finally:
        env.gateway.state.lock.release()
    assert env.gateway.state.in_flight["alice"] == 0


@pytest.mark.asyncio
async def test_proxy_client_disconnect_survives_cancelled_upstream_close(env):
    """A cancellation that lands inside ``upstream.aclose()`` must not skip the release.

    CancelledError is a BaseException, so ``contextlib.suppress(Exception)``
    around the close would let it escape before ``_release`` ran.  A real
    network close awaits, which is where the cancel scope delivers it.
    """
    add_user(env.k8s)
    first_chunk = asyncio.Event()

    class ClosingStream(httpx.AsyncByteStream):
        async def __aiter__(self):
            yield b"chunk-one"
            first_chunk.set()
            await asyncio.Event().wait()

        async def aclose(self):
            await asyncio.sleep(0)  # a checkpoint for the cancel scope

    env.upstream.proxy_handler = lambda request: httpx.Response(
        200, headers={"content-type": "text/event-stream"}, stream=ClosingStream()
    )
    messages = [{"type": "http.request", "body": b"", "more_body": False}]

    async def receive():
        if messages:
            return messages.pop(0)
        await first_chunk.wait()
        return {"type": "http.disconnect"}

    sent = await asgi_call(
        env,
        raw_scope(
            "GET",
            "/u/alice/api/stream",
            headers=[(b"cookie", f"dsh_gw_session=alice:{TOKEN}".encode())],
        ),
        receive,
    )
    starts = [message for message in sent if message["type"] == "http.response.start"]
    assert [start["status"] for start in starts] == [200]
    assert env.gateway.state.in_flight["alice"] == 0


@pytest.mark.asyncio
async def test_proxy_wakes_suspended_sandbox(env, monkeypatch):
    sandbox_name, _ = add_user(env.k8s, ready=False)
    env.k8s.sandboxes[sandbox_name]["spec"]["operatingMode"] = "Suspended"
    env.upstream.proxy_handler = lambda request: dsh_response(200, b"awake")

    async def ready(user, timeout):
        env.k8s.sandboxes[sandbox_name]["spec"]["operatingMode"] = "Running"
        env.k8s.sandboxes[sandbox_name]["status"] = {
            "conditions": [{"type": "Ready", "status": "True"}],
            "serviceFQDN": FQDN,
        }
        return env.k8s.sandboxes[sandbox_name]

    monkeypatch.setattr(env.gateway, "wait_ready", ready)
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/u/alice/api/state")
    assert response.status_code == 200
    assert ("dsh-alice-sbx", "Running") in env.k8s.patches


@pytest.mark.asyncio
async def test_proxy_wake_timeout_returns_503(env, monkeypatch):
    sandbox_name, _ = add_user(env.k8s, ready=False)
    env.k8s.sandboxes[sandbox_name]["spec"]["operatingMode"] = "Suspended"

    async def never(user, timeout):
        return None

    monkeypatch.setattr(env.gateway, "wait_ready", never)
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/u/alice/api/state")
    assert response.status_code == 503
    assert response.headers["retry-after"] == "10"


@pytest.mark.asyncio
async def test_proxy_unknown_user_is_404(env):
    async with asgi_client(env) as client:
        response = await client.get("/u/ghost/api/state")
    assert response.status_code == 404


@pytest.mark.asyncio
async def test_proxy_requires_auth(env):
    add_user(env.k8s)
    async with asgi_client(env) as client:
        response = await client.get("/u/alice/api/state")
    assert response.status_code == 401


@pytest.mark.asyncio
async def test_proxy_cookie_mutation_requires_same_origin(env):
    add_user(env.k8s)
    cookies = {"dsh_gw_session": f"alice:{TOKEN}"}
    async with asgi_client(env, cookies=cookies) as client:
        cross = await client.post(
            "/u/alice/api/mutate", headers={"Origin": "http://evil.test"}, content=b"{}"
        )
        good = await client.post(
            "/u/alice/api/mutate",
            headers={"Origin": "http://gateway.test"},
            content=b"{}",
        )
    assert cross.status_code == 403
    assert good.status_code == 200


def test_stream_timeout_leaves_read_unbounded(env):
    timeout = env.gateway._stream_timeout()
    assert timeout.read is None
    assert timeout.connect == env.settings.http_timeout


@pytest.mark.asyncio
async def test_proxied_requests_use_the_stream_timeout(env, monkeypatch):
    from starlette.requests import Request

    captured = {}
    real_build = env.gateway.http.build_request

    def spy_build(method, url, **kwargs):
        captured.update(kwargs)
        return real_build(method, url, **kwargs)

    monkeypatch.setattr(env.gateway.http, "build_request", spy_build)
    scope = {
        "type": "http",
        "method": "GET",
        "scheme": "http",
        "path": "/u/alice/events",
        "raw_path": b"/u/alice/events",
        "query_string": b"",
        "root_path": "",
        "headers": [(b"host", b"gateway.test")],
        "client": ("127.0.0.1", 4321),
        "server": ("gateway.test", 80),
        "path_params": {"user": "alice", "path": "events"},
    }
    response = await env.gateway._send_upstream(Request(scope), "alice", FQDN, "dsh-auth-a=1", b"")
    await response.aclose()

    assert captured["timeout"].read is None


@pytest.mark.asyncio
async def test_relay_ends_cleanly_when_the_upstream_stream_breaks(env):
    class Broken:
        closed = False

        async def aiter_raw(self):
            yield b"chunk"
            raise httpx.ReadTimeout("upstream went quiet")

        async def aclose(self):
            self.closed = True

    upstream = Broken()
    chunks = [chunk async for chunk in env.gateway._relay("alice", upstream)]

    assert chunks == [b"chunk"]
    assert upstream.closed is True
    assert env.gateway.state.in_flight["alice"] == 0


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


def test_ws_unauthenticated_closes_4401(env):
    with TestClient(env.app) as client:
        with client.websocket_connect("/u/alice/api/remote.mux") as ws:
            with pytest.raises(WebSocketDisconnect) as exc:
                ws.receive_text()
    assert exc.value.code == 4401


def test_ws_cookie_cross_origin_closes_4401(env):
    add_user(env.k8s)
    with TestClient(env.app) as client:
        with client.websocket_connect(
            "/u/alice/api/remote.mux",
            headers={
                "cookie": f"dsh_gw_session=alice:{TOKEN}",
                "origin": "http://evil.test",
            },
        ) as ws:
            with pytest.raises(WebSocketDisconnect) as exc:
                ws.receive_text()
    assert exc.value.code == 4401


def test_ws_bridges_frames_and_injects_cookie(env):
    add_user(env.k8s)
    with TestClient(env.app) as client:
        with client.websocket_connect(
            "/u/alice/api/remote.mux?token=secret&x=1",
            headers={"Authorization": f"Bearer {TOKEN}"},
        ) as ws:
            ws.send_text("hello")
            assert ws.receive_text() == "welcome"
            assert ws.receive_text() == "pong"

    url, kwargs = env.ws_connect.calls[0]
    # DSH's Host fence: dial the pod IP (its bind address) and never forward
    # the browser's Host/Origin, which would make the upgrade 403.
    assert url.startswith(f"ws://{POD_IP}:3080/api/remote.mux")
    assert "token=secret" not in url
    assert "x=1" in url
    headers = kwargs["additional_headers"]
    assert "authorization" not in headers
    assert headers["cookie"] == "dsh-auth-a=one"
    assert "host" not in headers
    assert "origin" not in headers
    assert env.ws_connect.instances[0].sent == ["hello"]


def test_ws_cookie_same_origin_bridges_frames(env):
    add_user(env.k8s)
    with TestClient(env.app) as client:
        with client.websocket_connect(
            "/u/alice/api/remote.mux",
            headers={
                "cookie": f"dsh_gw_session=alice:{TOKEN}",
                "origin": "http://testserver",
            },
        ) as ws:
            ws.send_text("hello")
            assert ws.receive_text() == "welcome"
            assert ws.receive_text() == "pong"

    headers = env.ws_connect.calls[0][1]["additional_headers"]
    assert headers["cookie"] == "dsh-auth-a=one"
    assert "origin" not in headers


def test_ws_handshake_failure_closes_1011(env):
    _, pod_name = add_user(env.k8s)
    env.k8s.logs[pod_name] = "token not ready\n"
    with TestClient(env.app) as client:
        with client.websocket_connect(
            "/u/alice/api/remote.mux",
            headers={"Authorization": f"Bearer {TOKEN}"},
        ) as ws:
            with pytest.raises(WebSocketDisconnect) as exc:
                ws.receive_text()
    assert exc.value.code == 1011


def test_ws_upstream_connect_failure_closes_1011(env):
    add_user(env.k8s)
    env.ws_connect.error = RuntimeError("cannot connect")
    with TestClient(env.app) as client:
        with client.websocket_connect(
            "/u/alice/api/remote.mux",
            headers={"Authorization": f"Bearer {TOKEN}"},
        ) as ws:
            with pytest.raises(WebSocketDisconnect) as exc:
                ws.receive_text()
    assert exc.value.code == 1011


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


def ready_claim(user="alice"):
    return {
        "metadata": {
            "name": claim_name(user),
            "annotations": {"aaas.example.com/managed-by": "dsh-gateway"},
        },
        "status": {"sandbox": {"name": f"dsh-{user}-sbx"}},
    }


def ready_sandbox(user="alice"):
    return {
        "metadata": {"name": f"dsh-{user}-sbx"},
        "spec": {},
        "status": {"conditions": [{"type": "Ready", "status": "True"}]},
    }


@pytest.mark.asyncio
async def test_sweeper_skips_claim_it_does_not_own(env):
    """A hand-applied claim (e.g. README 30-claim-alice.yaml) must not be swept.

    The gateway shares the namespace with the resource-layer walkthrough; those
    claims are named dsh-<user> too, but carry no ownership annotation.  The
    sweeper used to suspend them 60s after first sighting, deleting the pod that
    README step 3 port-forwards to.
    """
    env.k8s.claims["alice"] = {
        "metadata": {"name": claim_name("alice")},
        "status": {"sandbox": {"name": "dsh-alice-sbx"}},
    }
    env.k8s.sandboxes["dsh-alice-sbx"] = ready_sandbox()

    # The sweeper never starts an idle clock for a claim it does not own...
    await env.gateway.sweep_once()
    assert "alice" not in env.gateway.state.last_activity

    # ...and does not suspend it even when it looks long-idle.
    env.gateway.state.last_activity["alice"] = time.time() - env.settings.idle_timeout - 10
    await env.gateway.sweep_once()
    assert env.k8s.patches == []


@pytest.mark.asyncio
async def test_sweeper_suspends_idle_ready_sandbox(env):
    env.k8s.claims["alice"] = ready_claim()
    env.k8s.sandboxes["dsh-alice-sbx"] = ready_sandbox()
    env.gateway.state.last_activity["alice"] = time.time() - env.settings.idle_timeout - 10
    env.gateway.state.handshake[("alice", POD_UID, "gateway.test")] = "dsh-auth-a=x"

    await env.gateway.sweep_once()

    assert env.k8s.patches == [("dsh-alice-sbx", "Suspended")]
    assert env.gateway.state.handshake == {}
    assert env.gateway.state.maintenance["alice"] is False


@pytest.mark.asyncio
async def test_sweeper_skips_in_flight_request(env):
    env.k8s.claims["alice"] = ready_claim()
    env.k8s.sandboxes["dsh-alice-sbx"] = ready_sandbox()
    env.gateway.state.last_activity["alice"] = time.time() - 1000
    env.gateway.state.in_flight["alice"] = 1

    await env.gateway.sweep_once()

    assert env.k8s.patches == []


@pytest.mark.asyncio
async def test_sweeper_skips_recently_active_user(env):
    env.k8s.claims["alice"] = ready_claim()
    env.k8s.sandboxes["dsh-alice-sbx"] = ready_sandbox()
    env.gateway.state.last_activity["alice"] = time.time()

    await env.gateway.sweep_once()

    assert env.k8s.patches == []


@pytest.mark.asyncio
async def test_sweeper_skips_sandbox_not_ready(env):
    env.k8s.claims["alice"] = ready_claim()
    env.k8s.sandboxes["dsh-alice-sbx"] = {
        "metadata": {"name": "dsh-alice-sbx"},
        "spec": {},
        "status": {"conditions": []},
    }
    env.gateway.state.last_activity["alice"] = time.time() - 1000

    await env.gateway.sweep_once()

    assert env.k8s.patches == []


@pytest.mark.asyncio
async def test_sweeper_survives_transient_api_error(env):
    env.k8s.claims["alice"] = ready_claim()
    env.k8s.sandboxes["dsh-alice-sbx"] = ready_sandbox()
    env.gateway.state.last_activity["alice"] = time.time() - 1000
    env.k8s.patch_error = ApiException(status=500)

    with pytest.raises(ApiException):
        await env.gateway.sweep_once()

    # The per-user maintenance flag must be released even on failure.
    assert env.gateway.state.maintenance["alice"] is False


# --- SSO (reverse-proxy identity) ------------------------------------------


def enable_sso(env, header="X-Forwarded-User"):
    """Flip an env into SSO mode: the gateway trusts the named identity header."""
    env.settings = Settings(**{**env.settings.__dict__, "auth_proxy_header": header})
    env.gateway.settings = env.settings
    return env


def test_remote_user_none_when_sso_disabled(env):
    assert env.gateway.remote_user({"X-Forwarded-User": "alice"}) is None


def test_remote_user_returns_valid_and_rejects_invalid(env):
    enable_sso(env)
    g = env.gateway
    assert g.remote_user({"X-Forwarded-User": "alice"}) == "alice"
    assert g.remote_user({"X-Forwarded-User": "  bob "}) == "bob"
    assert g.remote_user({"X-Forwarded-User": "UPPER"}) is None
    assert g.remote_user({"X-Forwarded-User": "bad name"}) is None
    assert g.remote_user({"X-Forwarded-User": ""}) is None
    assert g.remote_user({}) is None


@pytest.mark.asyncio
async def test_sso_landing_new_user_shows_provisioning(env):
    enable_sso(env)
    async with asgi_client(env) as client:
        response = await client.get("/", headers={"X-Forwarded-User": "alice"})
    assert response.status_code == 200
    assert "text/html" in response.headers["content-type"]
    assert "Signed in as <strong>alice</strong>" in response.text
    assert "Provisioning" in response.text
    assert "Recreate sandbox" not in response.text
    assert re_cookie_user(response.headers["set-cookie"]) == "alice"
    assert [c[0] for c in env.k8s.created] == ["alice"]
    assert env.k8s.token_patches == []


@pytest.mark.asyncio
async def test_sso_landing_ready_connects_and_offers_recreate_only(env):
    enable_sso(env)
    add_user(env.k8s)
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.get("/", headers={"X-Forwarded-User": "alice"})
    assert response.status_code == 200
    assert "Signed in as <strong>alice</strong>" in response.text
    assert "Recreate sandbox" in response.text
    # No signup/token forms for SSO users.
    assert 'action="/users"' not in response.text
    assert 'action="/login"' not in response.text
    # Auto-connects to the workspace.
    assert "url=/u/alice/" in response.text
    assert env.k8s.created == []
    assert env.k8s.token_patches == []


@pytest.mark.asyncio
async def test_sso_landing_rotates_token_when_session_missing(env):
    enable_sso(env)
    add_user(env.k8s)
    async with asgi_client(env) as client:
        response = await client.get("/", headers={"X-Forwarded-User": "alice"})
    assert response.status_code == 200
    assert "Recreate sandbox" in response.text
    assert re_cookie_user(response.headers["set-cookie"]) == "alice"
    assert [p[0] for p in env.k8s.token_patches] == ["alice"]
    assert env.k8s.created == []


@pytest.mark.asyncio
async def test_recreate_deletes_and_reprovisions(env):
    enable_sso(env)
    add_user(env.k8s)
    async with asgi_client(env) as client:
        response = await client.post("/recreate", headers={"X-Forwarded-User": "alice"})
    assert response.status_code == 303
    assert response.headers["location"] == "/u/alice/"
    assert env.k8s.deleted == ["alice"]
    assert [c[0] for c in env.k8s.created] == ["alice"]
    assert re_cookie_user(response.headers["set-cookie"]) == "alice"


@pytest.mark.asyncio
async def test_recreate_requires_sso(env):
    add_user(env.k8s)
    async with asgi_client(env) as client:
        response = await client.post("/recreate", headers={"X-Forwarded-User": "alice"})
    assert response.status_code == 401
    assert env.k8s.deleted == []


@pytest.mark.asyncio
async def test_recreate_rejects_cross_origin(env):
    enable_sso(env)
    add_user(env.k8s)
    async with asgi_client(env, cookies={"dsh_gw_session": f"alice:{TOKEN}"}) as client:
        response = await client.post(
            "/recreate",
            headers={"X-Forwarded-User": "alice", "Origin": "http://evil.test"},
        )
    assert response.status_code == 403
    assert env.k8s.deleted == []


@pytest.mark.asyncio
async def test_sso_get_user_authorizes_via_header(env):
    enable_sso(env)
    add_user(env.k8s)
    async with asgi_client(env) as client:
        response = await client.get("/users/alice", headers={"X-Forwarded-User": "alice"})
    assert response.status_code == 200
    assert response.json()["user"] == "alice"


@pytest.mark.asyncio
async def test_sso_proxy_redirects_to_bootstrap_without_session(env):
    enable_sso(env)
    add_user(env.k8s)
    async with asgi_client(env) as client:
        response = await client.get("/u/alice/", headers={"X-Forwarded-User": "alice"})
    assert response.status_code == 303
    assert response.headers["location"] == "/"


@pytest.mark.asyncio
async def test_sso_proxy_redirects_to_bootstrap_when_claim_missing(env):
    enable_sso(env)
    async with asgi_client(env) as client:
        response = await client.get("/u/alice/", headers={"X-Forwarded-User": "alice"})
    assert response.status_code == 303
    assert response.headers["location"] == "/"


@pytest.mark.asyncio
async def test_spoofed_header_ignored_when_sso_disabled(env):
    add_user(env.k8s)
    async with asgi_client(env) as client:
        response = await client.get("/u/alice/", headers={"X-Forwarded-User": "alice"})
    assert response.status_code == 401