diff --git a/examples/dsh-agents-as-a-service/docs/architecture.md b/examples/dsh-agents-as-a-service/docs/architecture.md index 5bd5de4..6a07578 100644 --- a/examples/dsh-agents-as-a-service/docs/architecture.md +++ b/examples/dsh-agents-as-a-service/docs/architecture.md @@ -135,6 +135,10 @@ Handshake cache behavior: one entry per `(user, pod UID, Host)`; refreshed once `401`; dropped when the sandbox is suspended or deleted. The launch token and the DSH auth cookie never reach the browser. +Inside `proxy_http` the order is auth, in-flight slot, sandbox resolution/wake, DSH handshake, +body read, upstream send. An upload aborted mid-body may therefore already have paid for a wake +and a handshake, but it never reaches DSH. + ### Wake-on-connect and idle suspend ```mermaid @@ -317,6 +321,7 @@ Notes and limits: | Probe on `/` | Pod `CrashLoopBackOff`, `401` in probe logs | Use `/favicon.svg` for readiness/liveness (shipped templates do) | | Warm pool exhausted | JSON signup returns `202` (state `Waking`); form signup is slow | Increase `replicas` in the warm pool | | Gateway restart | Idle sandboxes stop suspending; handshake cache lost | Rebuilt lazily on next request; consider a shared store for production | +| Browser aborts a request (tab reload, back navigation) during a slow wake or handshake | No response reaches the client; nothing appears in the access log | If the abort lands before the response starts, the body read raises `ClientDisconnect` and the app answers `499` (`CLIENT_CLOSED_REQUEST`), which cannot reach the departed peer. Nothing reached DSH, no ERROR traceback is logged, and the in-flight slot is released. A request already streaming ends with its upstream status. Before the fix every mid-body abort logged a `ClientDisconnect` traceback | | DSH auth cookie expired mid-session | Proxy `401` once, then recovers | Gateway refreshes the handshake once | | Upstream pod name changed after resume | Cached cookie keyed by old pod UID is useless | Cache is invalidated on suspend; next request re-handshakes | | Two gateway replicas | Double suspend/wake races | Keep `replicas: 1` until state is externalized | diff --git a/examples/dsh-agents-as-a-service/gateway/README.md b/examples/dsh-agents-as-a-service/gateway/README.md index e7e6e9c..9185183 100644 --- a/examples/dsh-agents-as-a-service/gateway/README.md +++ b/examples/dsh-agents-as-a-service/gateway/README.md @@ -135,6 +135,8 @@ oauth2-proxy's egress to the gateway, DNS and public HTTPS. | Missing or invalid credentials | `401` | | Cross-origin cookie mutation | `403` | | Request body over `MAX_BODY_MB` | `413` | +| Client aborts mid-body, before the response starts | `499` (`CLIENT_CLOSED_REQUEST`). Nothing is forwarded to DSH and nothing can be delivered; it is not an error, so no traceback is logged. Covers the body-reading control routes (`POST /users`, `POST /login`) as well as the proxy. | +| Client aborts mid-stream (SSE, long response) | The response already started, so its upstream status stands (typically `200`); the stream simply ends. No traceback, and the in-flight slot is released as the relay unwinds. | | Sandbox not provisioned, wake timeout, handshake failure | `503` + `Retry-After: 10` | | Upstream connection/protocol failure | `502` | | WS auth failure (unknown user, bad token, cross-origin) | accept, then close `4401` | @@ -144,6 +146,12 @@ A close code can only be delivered after the WebSocket upgrade is accepted; Starlette would turn an early close into an HTTP 403, which is why WS auth failures are `101`, then close `4401`. +An aborted request is deliberately silent: uvicorn runs with `--no-access-log`, +and the `499` cannot reach a peer that has already left (uvicorn's `send()` is +a no-op once the socket is closed). Sandbox resolution, the wake and the DSH +handshake run before the body is read, so an aborted upload can cost a wake or +a handshake but never reaches DSH. + ## Configuration All values are read from the environment by `Settings.from_env()`; every one diff --git a/examples/dsh-agents-as-a-service/gateway/gateway.py b/examples/dsh-agents-as-a-service/gateway/gateway.py index 27b6ba8..b3f7432 100644 --- a/examples/dsh-agents-as-a-service/gateway/gateway.py +++ b/examples/dsh-agents-as-a-service/gateway/gateway.py @@ -62,7 +62,7 @@ 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.requests import ClientDisconnect, Request from starlette.responses import ( HTMLResponse, JSONResponse, @@ -144,6 +144,11 @@ class UpstreamError(Exception): """The upstream DSH connection failed.""" +# nginx's "client closed request": the peer went away before it could be +# answered, so no status is ever delivered. +CLIENT_CLOSED_REQUEST = 499 + + # --- Pure helpers (unit-testable without an app) -------------------------- @@ -682,9 +687,14 @@ class Gateway: 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() + # Deliberately await-free: this runs from ``finally`` blocks, and the + # slot must be freed even when a client disconnect cancels the request + # mid-stream. Awaiting the lock there can be interrupted before the + # decrement, stranding the slot (the user's sandbox would then never + # look idle again). A plain dict update is atomic because the event + # loop cannot switch tasks between two synchronous statements. + 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 -------------------------------------- @@ -1238,9 +1248,17 @@ against the JSON API. The browser session is an HttpOnly cookie.

print(f"stream for {user} ended early: {exc!r}", flush=True) return finally: - with contextlib.suppress(Exception): - await upstream.aclose() - await self._release(user) + # The slot must be freed even when this ``finally`` runs inside a + # cancelled task: on a client disconnect Starlette cancels the + # streaming task through an anyio cancel scope, and CancelledError + # is a BaseException, so ``suppress(Exception)`` would let a close + # that awaits abort the whole block before ``_release`` ran. A + # stranded slot keeps the user's sandbox looking busy forever. + try: + with contextlib.suppress(Exception): + await upstream.aclose() + finally: + await self._release(user) # -- WebSocket proxy --------------------------------------------------- @@ -1439,6 +1457,20 @@ def _lifespan(gateway: Gateway, close_http_client: httpx.AsyncClient | None = No return lifespan +async def client_gone(request: Request, exc: ClientDisconnect) -> Response: + """Answer a client that vanished mid-request instead of logging a traceback. + + Reading a request body raises ``ClientDisconnect`` when a browser reload, a + back navigation or an ingress timeout aborts the upload. The request never + reached DSH, and nothing can be delivered to the peer that left, so the + honest outcome is 499 (nginx's "client closed request"). Letting it escape + would make Starlette's ``ServerErrorMiddleware`` log one ERROR traceback per + aborted request, which buries real faults. It also covers the control + routes that read a body (signup and login), not just the proxy. + """ + return Response(status_code=CLIENT_CLOSED_REQUEST) + + def create_app( settings: Settings | None = None, k8s=None, @@ -1471,6 +1503,7 @@ def create_app( ] app = Starlette( routes=routes, + exception_handlers={ClientDisconnect: client_gone}, lifespan=_lifespan(gateway, http_client if owns_http_client else None), ) app.state.gateway = gateway diff --git a/examples/dsh-agents-as-a-service/gateway/requirements-dev.txt b/examples/dsh-agents-as-a-service/gateway/requirements-dev.txt index a3d3260..eef29c9 100644 --- a/examples/dsh-agents-as-a-service/gateway/requirements-dev.txt +++ b/examples/dsh-agents-as-a-service/gateway/requirements-dev.txt @@ -2,4 +2,4 @@ pytest==9.1.1 pytest-asyncio==1.4.0 pytest-cov==7.1.0 -ruff==0.17.1 +ruff==0.17.2 diff --git a/examples/dsh-agents-as-a-service/gateway/requirements.txt b/examples/dsh-agents-as-a-service/gateway/requirements.txt index daf3ef1..9b1892a 100644 --- a/examples/dsh-agents-as-a-service/gateway/requirements.txt +++ b/examples/dsh-agents-as-a-service/gateway/requirements.txt @@ -1,4 +1,4 @@ -starlette==1.9.0 +starlette==1.10.0 uvicorn[standard]==0.54.0 httpx==0.28.1 websockets==17.3 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 e10f643..15243af 100644 --- a/examples/dsh-agents-as-a-service/gateway/test_gateway.py +++ b/examples/dsh-agents-as-a-service/gateway/test_gateway.py @@ -8,6 +8,7 @@ use Starlette's ``TestClient`` (which runs the lifespan) against a fake """ import asyncio +import contextlib import hashlib import json import threading @@ -263,6 +264,43 @@ 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 ---------------------------------------------------------- @@ -1049,6 +1087,162 @@ async def test_proxy_body_limit_returns_413(env): 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)