feat(api): serve long-lived streams on the event loop

Move GET /api/events, the remote-device command stream and artifact
downloads from Flask to Starlette routes mounted ahead of the Flask
catch-all. On Flask each held an a2wsgi threadpool slot for as long as
its response stayed open, and because uvicorn drops writes after a
client disconnects, a closed tab never released it.

- asgi_auth: one JWT/OIDC gate for Starlette routes; the chat reconnect
  reader uses it too
- ClosingStreamingResponse closes the body iterator and releases the
  SSE slot or file handle even when the client leaves before the first
  frame
- AsyncTopic liveness probe replaces the sync client's socket_timeout
  guard against half-open pub/sub sockets
- ASYNC_REDIS_MAX_CONNECTIONS sizes the async Redis pool; every open
  stream holds a connection and redis-py defaults to 100

docsgpt/api/devices/auth.py

@@ -6,7 +6,7 @@ import base64
 import hashlib
 import logging
 import time
-from typing import Optional, Tuple
+from typing import Callable, Mapping, Optional, Tuple
 
 from flask import jsonify, make_response, request
 
@@ -17,6 +17,9 @@ from docsgpt.storage.db.session import db_readonly, db_session
 
 logger = logging.getLogger(__name__)
 
+#: ``(error_code, http_status)`` for a rejected device request.
+DeviceAuthFailure = Tuple[str, int]
+
 
 def hash_session_token(token: str) -> str:
     """SHA-256 over the opaque session token. Stored in ``devices.token_hash``."""
@@ -45,32 +48,46 @@ def _canonical_payload(method: str, path: str, ts: str, body: bytes) -> str:
     return f"{method} {path} {ts} {body_hash}"
 
 
-def verify_device_session(*, touch: bool = True) -> Tuple[Optional[dict], Optional[tuple]]:
-    """Validate the device session token on the current request.
+def authenticate_device(
+    headers: Mapping[str, str],
+    method: str,
+    path: str,
+    get_body: Callable[[], bytes],
+    *,
+    touch: bool = True,
+) -> Tuple[Optional[dict], Optional[DeviceAuthFailure]]:
+    """Validate a device session token from the raw request parts.
 
-    Returns ``(device_row, None)`` on success or ``(None, response)`` on
-    failure, where ``response`` is a ``(json, status)`` Flask tuple.
+    Framework-neutral core shared by the Flask routes and the native-async
+    SSE stream, which calls it from a worker thread since it reads Postgres.
 
     Args:
+        headers: Request headers (a case-insensitive mapping).
+        method: HTTP method, as the CLI signed it.
+        path: Request path without the query string, as the CLI signed it.
+        get_body: Returns the raw body; only called when signatures are required.
         touch: When True, bump ``last_seen_at`` on the device row.
+
+    Returns:
+        tuple: ``(device_row, None)`` on success or ``(None, (code, status))``.
     """
-    auth = request.headers.get("Authorization", "")
+    auth = headers.get("Authorization", "")
     if not auth.lower().startswith("bearer "):
-        return None, _error("missing_token", 401)
+        return None, ("missing_token", 401)
     token = auth.split(" ", 1)[1].strip()
     if not token:
-        return None, _error("missing_token", 401)
+        return None, ("missing_token", 401)
 
     token_hash = hash_session_token(token)
     with db_readonly() as conn:
         device = DevicesRepository(conn).find_by_token_hash(token_hash)
     if device is None:
-        return None, _error("invalid_token", 401)
+        return None, ("invalid_token", 401)
 
     if settings.REMOTE_DEVICE_REQUIRE_SIGNATURE:
-        sig_error = _verify_signature(device)
-        if sig_error is not None:
-            return None, sig_error
+        failure = _verify_signature(device, headers, method, path, get_body)
+        if failure is not None:
+            return None, failure
 
     if touch:
         try:
@@ -82,23 +99,46 @@ def verify_device_session(*, touch: bool = True) -> Tuple[Optional[dict], Option
     return device, None
 
 
-def _verify_signature(device: dict) -> Optional[tuple]:
+def verify_device_session(*, touch: bool = True) -> Tuple[Optional[dict], Optional[tuple]]:
+    """Validate the device session token on the current Flask request.
+
+    Returns ``(device_row, None)`` on success or ``(None, response)`` on
+    failure, where ``response`` is a ready Flask error response.
+
+    Args:
+        touch: When True, bump ``last_seen_at`` on the device row.
+    """
+    device, failure = authenticate_device(
+        request.headers, request.method, request.path, request.get_data, touch=touch
+    )
+    if failure is not None:
+        return None, _error(*failure)
+    return device, None
+
+
+def _verify_signature(
+    device: dict,
+    headers: Mapping[str, str],
+    method: str,
+    path: str,
+    get_body: Callable[[], bytes],
+) -> Optional[DeviceAuthFailure]:
     """Verify ``X-Device-Signature`` against the stored machine pubkey."""
-    sig_b64 = request.headers.get("X-Device-Signature")
-    ts = request.headers.get("X-Device-Timestamp")
-    fp = request.headers.get("X-Device-Machine-Key")
+    sig_b64 = headers.get("X-Device-Signature")
+    ts = headers.get("X-Device-Timestamp")
+    fp = headers.get("X-Device-Machine-Key")
     if not sig_b64 or not ts or not fp:
-        return _error("missing_signature", 401)
+        return ("missing_signature", 401)
 
     try:
         ts_int = int(ts)
     except (TypeError, ValueError):
-        return _error("invalid_timestamp", 401)
+        return ("invalid_timestamp", 401)
     if abs(time.time() - ts_int) > 300:
-        return _error("timestamp_skew", 401)
+        return ("timestamp_skew", 401)
 
     if fp != device.get("machine_pubkey_fingerprint"):
-        return _error("fingerprint_mismatch", 401)
+        return ("fingerprint_mismatch", 401)
 
     # Defer cryptography import to keep cold-start light when signatures
     # are disabled (the dev default).
@@ -111,26 +151,24 @@ def _verify_signature(device: dict) -> Optional[tuple]:
         logger.error(
             "Signature verification requested but ``cryptography`` is not installed."
         )
-        return _error("signature_unsupported", 500)
+        return ("signature_unsupported", 500)
 
     # The full public key isn't stored — only the fingerprint. For signature
     # verification we accept the pubkey in a header too; we trust it iff
     # its fingerprint matches the stored one. This avoids a separate pubkey
     # column for MVP.
-    pubkey_b64 = request.headers.get("X-Device-Machine-Pubkey")
+    pubkey_b64 = headers.get("X-Device-Machine-Pubkey")
     if not pubkey_b64:
-        return _error("missing_pubkey", 401)
+        return ("missing_pubkey", 401)
     if fingerprint_pubkey(pubkey_b64) != device["machine_pubkey_fingerprint"]:
-        return _error("pubkey_fingerprint_mismatch", 401)
+        return ("pubkey_fingerprint_mismatch", 401)
 
-    payload = _canonical_payload(
-        request.method, request.path, ts, request.get_data()
-    ).encode("utf-8")
+    payload = _canonical_payload(method, path, ts, get_body()).encode("utf-8")
     try:
         pubkey = Ed25519PublicKey.from_public_bytes(base64.b64decode(pubkey_b64))
         pubkey.verify(base64.b64decode(sig_b64), payload)
     except (InvalidSignature, ValueError):
-        return _error("invalid_signature", 401)
+        return ("invalid_signature", 401)
 
     return None
 

docsgpt/api/devices/routes.py

@@ -16,7 +16,6 @@ from docsgpt.api.devices.session import (
     ack_invocation,
     me,
     poll,
-    session_events,
     submit_output,
 )
 from docsgpt.devices.normalizer import normalize_command
@@ -313,10 +312,7 @@ def register(bp: Blueprint) -> None:
     bp.add_url_rule(
         "/api/devices/me", view_func=me, methods=["GET"], endpoint="device_me",
     )
-    bp.add_url_rule(
-        "/api/devices/sessions/<session_id>/events", view_func=session_events,
-        methods=["GET"], endpoint="device_session_events",
-    )
+    # The session's SSE stream is native-async (session_events.py, mounted in asgi.py).
     bp.add_url_rule(
         "/api/devices/sessions/<session_id>/invocations/<invocation_id>/ack",
         view_func=ack_invocation, methods=["POST"], endpoint="device_ack",

docsgpt/api/devices/session.py

@@ -1,4 +1,7 @@
-"""Device session endpoints: poll, SSE, ack, output."""
+"""Device session endpoints: poll, me, ack, output.
+
+The SSE command stream itself is a native-async route (``session_events.py``).
+"""
 
 from __future__ import annotations
 
@@ -6,14 +9,10 @@ import gzip
 import io
 import json
 import logging
-import time
-from typing import Iterator
 
-from flask import Response, jsonify, make_response, request, stream_with_context
+from flask import Response, jsonify, make_response, request
 
 from docsgpt.api.devices.auth import verify_device_session
-from docsgpt.core.settings import settings
-from docsgpt.core.shutdown import is_shutting_down
 from docsgpt.devices.broker import get_broker
 from docsgpt.storage.db.repositories.device_audit_log import (
     DeviceAuditLogRepository,
@@ -82,71 +81,6 @@ def me() -> Response:
     return make_response(jsonify(out), 200)
 
 
-def session_events(session_id: str) -> Response:
-    """``GET /api/devices/sessions/{id}/events`` — SSE invocation stream.
-
-    The ``session_id`` must be the ``session_ticket`` the device's own
-    ``/poll`` just issued (the path it was handed as ``session_url``). A
-    stale, mismatched, or fabricated ticket is rejected with ``410 Gone``
-    before any stream is opened.
-    """
-    device, err = verify_device_session()
-    if err is not None:
-        return err
-    broker = get_broker()
-    if not broker.validate_ticket(device["id"], session_id):
-        return make_response(
-            jsonify({"success": False, "error": "session_ticket_invalid"}), 410
-        )
-    sess = broker.register_session(device["id"], device["user_id"])
-
-    keepalive_interval = float(settings.SSE_KEEPALIVE_SECONDS)
-    idle_seconds = float(settings.REMOTE_DEVICE_SESSION_IDLE_SECONDS)
-
-    @stream_with_context
-    def generate() -> Iterator[str]:
-        try:
-            last_keepalive = time.time()
-            while not sess.closed.is_set():
-                # Break promptly on shutdown (see docsgpt/core/shutdown.py).
-                if is_shutting_down():
-                    break
-                now = time.time()
-                if now - sess.last_activity_at > idle_seconds:
-                    yield _sse_event(
-                        "session_end",
-                        {"reason": "inactivity_timeout"},
-                        sess.last_event_id + 1,
-                    )
-                    sess.last_event_id += 1
-                    broker.close_session(sess.session_id, reason="idle")
-                    return
-                envelope = broker.next_command(sess, timeout=1.0)
-                if envelope is None:
-                    if time.time() - last_keepalive >= keepalive_interval:
-                        last_keepalive = time.time()
-                        yield ": heartbeat\n\n"
-                    continue
-                sess.last_event_id += 1
-                sess.last_activity_at = time.time()
-                yield _sse_event("invocation", envelope, sess.last_event_id)
-                last_keepalive = time.time()
-        except GeneratorExit:
-            logger.debug("device SSE generator exiting for session %s", sess.session_id)
-            raise
-        finally:
-            broker.close_session(sess.session_id, reason="generator_exit")
-
-    return Response(
-        generate(),
-        mimetype="text/event-stream",
-        headers={
-            "Cache-Control": "no-cache",
-            "X-Accel-Buffering": "no",
-        },
-    )
-
-
 def ack_invocation(session_id: str, invocation_id: str) -> Response:
     """CLI acks (accepted / denied / auto_approved) the invocation."""
     device, err = verify_device_session()
@@ -262,11 +196,3 @@ def _as_opt_int(value) -> int | None:
         return int(value)
     except (TypeError, ValueError):
         return None
-
-
-def _sse_event(name: str, payload: dict, event_id: int) -> str:
-    return (
-        f"event: {name}\n"
-        f"id: {event_id}\n"
-        f"data: {json.dumps(payload)}\n\n"
-    )

docsgpt/api/devices/session_events.py

@@ -0,0 +1,133 @@
+"""``GET /api/devices/sessions/{session_id}/events`` — a paired device's command stream.
+
+A native-async Starlette route mounted ahead of Flask in ``docsgpt/asgi.py``.
+The CLI holds this stream open while it waits for commands, so it runs on the
+event loop: the token check and ticket claim hop to a worker thread once, then
+commands arrive through the broker's async ``BLPOP`` and an idle session holds
+no thread.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from functools import partial
+from typing import AsyncIterator
+
+import anyio
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+from starlette.routing import Route
+
+from docsgpt.api.asgi_stream import ClosingStreamingResponse
+from docsgpt.api.devices.auth import authenticate_device
+from docsgpt.core.settings import settings
+from docsgpt.core.shutdown import is_shutting_down
+from docsgpt.devices.broker import DeviceBroker, get_broker
+
+logger = logging.getLogger(__name__)
+
+# Upper bound on one broker poll, so a closed session, idle expiry or a
+# server drain is noticed within this many seconds.
+_POLL_TIMEOUT_SECONDS = 1.0
+
+_SSE_HEADERS = {
+    "Cache-Control": "no-cache",
+    "X-Accel-Buffering": "no",
+}
+
+
+def _error(code: str, status_code: int) -> JSONResponse:
+    return JSONResponse({"success": False, "error": code}, status_code=status_code)
+
+
+def _sse_event(name: str, payload: dict, event_id: int) -> str:
+    return (
+        f"event: {name}\n"
+        f"id: {event_id}\n"
+        f"data: {json.dumps(payload)}\n\n"
+    )
+
+
+async def _session_stream(broker: DeviceBroker, device: dict) -> AsyncIterator[str]:
+    """Deliver queued commands until the session closes, idles out, or the server drains."""
+    # Registration consumes the ticket. Shield it so a disconnect can't leave
+    # a registered session that the ``finally`` below never closes.
+    with anyio.CancelScope(shield=True):
+        sess = await anyio.to_thread.run_sync(
+            broker.register_session, device["id"], device["user_id"]
+        )
+    keepalive_interval = float(settings.SSE_KEEPALIVE_SECONDS)
+    idle_seconds = float(settings.REMOTE_DEVICE_SESSION_IDLE_SECONDS)
+    try:
+        last_keepalive = time.time()
+        # ``closed`` is set when the CLI reconnects and a newer session
+        # replaces this one.
+        while not sess.closed.is_set():
+            # Break promptly on shutdown (see docsgpt/core/shutdown.py).
+            if is_shutting_down():
+                break
+            if time.time() - sess.last_activity_at > idle_seconds:
+                yield _sse_event(
+                    "session_end",
+                    {"reason": "inactivity_timeout"},
+                    sess.last_event_id + 1,
+                )
+                sess.last_event_id += 1
+                broker.close_session(sess.session_id, reason="idle")
+                return
+            envelope = await broker.next_command_async(sess, timeout=_POLL_TIMEOUT_SECONDS)
+            if envelope is None:
+                if time.time() - last_keepalive >= keepalive_interval:
+                    last_keepalive = time.time()
+                    yield ": heartbeat\n\n"
+                continue
+            sess.last_event_id += 1
+            sess.last_activity_at = time.time()
+            yield _sse_event("invocation", envelope, sess.last_event_id)
+            last_keepalive = time.time()
+    finally:
+        broker.close_session(sess.session_id, reason="stream_end")
+
+
+async def device_session_events(request: Request) -> Response:
+    """Stream queued invocations to the device holding the poll-issued ticket.
+
+    The ``session_id`` must be the ``session_ticket`` the device's own
+    ``/poll`` just issued (the path it was handed as ``session_url``). A
+    stale, mismatched, or fabricated ticket is rejected with ``410 Gone``
+    before any stream is opened.
+    """
+    session_id = request.path_params["session_id"]
+    body = await request.body()
+    device, failure = await anyio.to_thread.run_sync(
+        partial(
+            authenticate_device,
+            request.headers,
+            request.method,
+            request.url.path,
+            lambda: body,
+        )
+    )
+    if failure is not None:
+        return _error(*failure)
+    broker = get_broker()
+    if not await anyio.to_thread.run_sync(broker.validate_ticket, device["id"], session_id):
+        return _error("session_ticket_invalid", 410)
+    return ClosingStreamingResponse(
+        _session_stream(broker, device),
+        media_type="text/event-stream",
+        headers=_SSE_HEADERS,
+    )
+
+
+# Mounted in ``docsgpt/asgi.py`` ahead of the Flask catch-all. The device's
+# poll, ack and output endpoints stay on Flask (``session.py``).
+device_session_routes = [
+    Route(
+        "/api/devices/sessions/{session_id}/events",
+        device_session_events,
+        methods=["GET"],
+    ),
+]