diff --git a/examples/dsh-agents-as-a-service/gateway/gateway.py b/examples/dsh-agents-as-a-service/gateway/gateway.py index 2865b5c..71b5dba 100644 --- a/examples/dsh-agents-as-a-service/gateway/gateway.py +++ b/examples/dsh-agents-as-a-service/gateway/gateway.py @@ -622,10 +622,12 @@ class Gateway: """Provision ``remote`` idempotently and mint/refresh a session token. The authenticating proxy has already verified ``remote``, so it is - trusted. A brand-new user is provisioned like ``signup`` (claim created, - token minted, sandbox awaited); an existing user whose session cookie is - missing or no longer matches gets a rotated token so a fresh cookie can - be issued. A still-valid session is reused untouched. + trusted. A brand-new user is provisioned (claim created, token minted); + an existing user whose session cookie is missing or no longer matches + gets a rotated token so a fresh cookie can be issued. A still-valid + session is reused untouched. Provisioning is non-blocking: the sandbox + is awaited by the proxy path (``_resolve_sandbox``), so the landing page + can report "Provisioning" instead of stalling for ``WAKE_TIMEOUT``. """ claim = await self.k8s.get_claim(remote) token, cookie_user = self.credentials(headers, cookies) @@ -643,7 +645,6 @@ class Gateway: await self.k8s.patch_token(remote, token_hash) else: await self._touch(remote) - await self.wait_ready(remote, self.settings.wake_timeout) else: await self.k8s.patch_token(remote, token_hash) await self._touch(remote) @@ -835,14 +836,65 @@ class Gateway: async def healthz(self, request: Request) -> JSONResponse: return JSONResponse({"status": "ok"}) + @staticmethod + def _sso_page(user: str, state: str) -> str: + """Minimal SSO landing: connect to the sandbox, or offer to recreate it. + + No signup/token forms — the authenticating proxy has already identified + the user. A "Provisioning" sandbox refreshes itself until a pod exists; + any other state auto-connects to ``/u//`` and offers the single + "Recreate sandbox" action. + """ + safe_user = html.escape(user) + path = quote(user) + if state == "Provisioning": + status = "Provisioning" + body = ( + '' + "

Provisioning your sandbox… this can take a minute.

" + ) + else: + status = html.escape(state) + confirm = ( + 'onsubmit="return confirm(\'Delete this sandbox and start fresh? ' + "Your workspace data will be wiped.')\"" + ) + body = ( + f'' + "

Connecting to your sandbox…

" + f'
' + '
' + ) + return ( + '' + '' + "DSH sandbox" + "

DeepSeek Harness sandbox

" + f"

Signed in as {safe_user}.

" + f"

Sandbox status: {status}.

" + f"{body}" + "" + ) + async def landing(self, request: Request) -> Response: remote = self.remote_user(request.headers) if remote: - response = RedirectResponse(f"/u/{remote}/", status_code=303) - await self._sso_session( + # SSO: auto-provision + connect. A temporary response captures the + # session cookie minted by _sso_session; the real page is rendered + # once the sandbox state is known. + response = HTMLResponse() + claim = await self._sso_session( remote, request.headers, request.cookies, response ) - return response + sandbox = await self._sandbox_of(claim) if claim else None + page = HTMLResponse(self._sso_page(remote, derive_state(sandbox))) + if "set-cookie" in response.headers: + page.headers["set-cookie"] = response.headers["set-cookie"] + return page token, user = self.credentials(request.headers, request.cookies) session = "" if user and token: @@ -1043,6 +1095,38 @@ against the JSON API. The browser session is an HttpOnly cookie.

{"user": user, "note": "claim deleted; sandbox and PVC cascade"} ) + async def recreate(self, request: Request) -> Response: + """Delete and re-provision the SSO user's sandbox for a clean start. + + SSO-only: the identity comes from the trusted proxy header. Deleting the + claim cascades to the sandbox pod and its PVC, wiping user state; we poll + until the claim is fully gone (finalizers) before recreating to avoid a + name-collision 409, then reconnect. + """ + remote = self.remote_user(request.headers) + if not remote: + return self._error(401, "unauthorized") + _token, cookie_user = self.credentials(request.headers, request.cookies) + if cookie_user is not None and cookie_user != remote: + return self._error(401, "unauthorized") + if self._mutation_from_cookie(request.headers, "POST", cookie_user): + return self._error(403, "cross-origin request rejected") + await self.k8s.delete_claim(remote) + self._clear_handshake(remote) + async with self.state.lock: + self.state.last_activity.pop(remote, None) + self.state.in_flight.pop(remote, None) + # Wait out any finalizer so the claim name is reusable. + for _ in range(30): + if await self.k8s.get_claim(remote) is None: + break + await asyncio.sleep(1) + response = RedirectResponse(f"/u/{remote}/", status_code=303) + await self._sso_session( + remote, request.headers, request.cookies, response + ) + return response + # -- HTTP proxy -------------------------------------------------------- async def proxy_http(self, request: Request) -> Response: @@ -1447,6 +1531,7 @@ def create_app( Route("/users", gateway.signup, methods=["POST"]), Route("/login", gateway.login, methods=["POST"]), Route("/logout", gateway.logout, methods=["POST"]), + Route("/recreate", gateway.recreate, methods=["POST"]), Route("/users/{user}", gateway.get_user), Route("/users/{user}", gateway.delete_user, methods=["DELETE"]), Route("/u/{user}", gateway.redirect_to_workspace), diff --git a/examples/dsh-agents-as-a-service/gateway/test_gateway.py b/examples/dsh-agents-as-a-service/gateway/test_gateway.py index 9f68026..f3db972 100644 --- a/examples/dsh-agents-as-a-service/gateway/test_gateway.py +++ b/examples/dsh-agents-as-a-service/gateway/test_gateway.py @@ -1460,32 +1460,36 @@ def test_remote_user_returns_valid_and_rejects_invalid(env): @pytest.mark.asyncio -async def test_sso_landing_provisions_new_user_and_redirects(env, monkeypatch): +async def test_sso_landing_new_user_shows_provisioning(env): enable_sso(env) - - async def ready(user, timeout): - return {"metadata": {"name": f"dsh-{user}-sbx"}} - - monkeypatch.setattr(env.gateway, "wait_ready", ready) async with asgi_client(env) as client: response = await client.get("/", headers={"X-Forwarded-User": "alice"}) - assert response.status_code == 303 - assert response.headers["location"] == "/u/alice/" + assert response.status_code == 200 + assert "text/html" in response.headers["content-type"] + assert "Signed in as alice" 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_reuses_valid_session_without_rotation(env): +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 == 303 - assert response.headers["location"] == "/u/alice/" + assert response.status_code == 200 + assert "Signed in as alice" 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 == [] @@ -1496,13 +1500,54 @@ async def test_sso_landing_rotates_token_when_session_missing(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 == 303 - assert response.headers["location"] == "/u/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)