chore: import upstream snapshot with attribution

application/api/async_sse.py

@@ -0,0 +1,248 @@
+"""Native-async (ASGI) SSE reader routes, mounted ahead of the Flask app.
+
+These Starlette routes serve the chat-stream *reconnect* path on the event
+loop, so a long-lived, mostly-idle tail costs a coroutine instead of one of
+the 32 a2wsgi threadpool slots (see ``application/asgi.py``). They are the
+sole reconnect reader — the old Flask blueprint has been removed. The heavy
+*producer* (``POST /api/answer/stream`` → agent → LLM) stays on the sync
+path untouched.
+
+Auth, message-id validation, ``Last-Event-ID`` parsing and ownership are
+done here; the snapshot/tail wire format is shared with the producer's
+journal via ``build_message_event_stream_async`` → ``format_sse_event``.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from typing import Optional
+
+import anyio
+from sqlalchemy import text
+from starlette.requests import Request
+from starlette.responses import JSONResponse, StreamingResponse
+from starlette.routing import Route
+
+from application.api.oidc.denylist import is_denied as oidc_session_denied
+from application.auth import handle_auth
+from application.core.settings import settings
+from application.events.keys import connection_counter_key
+from application.storage.db.session import db_readonly
+from application.streaming.async_event_replay import (
+    build_message_event_stream_async,
+)
+from application.streaming.async_redis import get_async_redis_instance
+from application.streaming.event_replay import (
+    DEFAULT_KEEPALIVE_SECONDS,
+    DEFAULT_POLL_TIMEOUT_SECONDS,
+)
+
+logger = logging.getLogger(__name__)
+
+# Per-user concurrent-connection counter TTL (seconds) — orphaned counts
+# from a hard crash self-heal after this window, mirroring the /api/events
+# notification stream. The reconnect reader shares the same counter key, so
+# the cap bounds a user's *total* live SSE footprint.
+_COUNTER_TTL_SECONDS = 3600
+
+# A message_id is the canonical UUID hex format. Reject anything else before
+# the SQL layer so a malformed cookie can't surface as a 500.
+_MESSAGE_ID_RE = re.compile(
+    r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-"
+    r"[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$"
+)
+# ``sequence_no`` is a non-negative decimal integer. Anything else is corrupt
+# client state — fall through to a fresh-replay cursor.
+_SEQUENCE_NO_RE = re.compile(r"^\d+$")
+
+
+def _normalise_last_event_id(raw: Optional[str]) -> Optional[int]:
+    """Parse a ``Last-Event-ID`` cursor; ``None`` for missing/invalid."""
+    if raw is None:
+        return None
+    raw = raw.strip()
+    if not raw or not _SEQUENCE_NO_RE.match(raw):
+        return None
+    return int(raw)
+
+
+def _user_owns_message(message_id: str, user_id: str) -> bool:
+    """Return True iff ``message_id`` belongs to ``user_id``."""
+    try:
+        with db_readonly() as conn:
+            row = conn.execute(
+                text(
+                    """
+                    SELECT 1 FROM conversation_messages
+                    WHERE id = CAST(:id AS uuid)
+                      AND user_id = :u
+                    LIMIT 1
+                    """
+                ),
+                {"id": message_id, "u": user_id},
+            ).first()
+        return row is not None
+    except Exception:
+        logger.exception(
+            "Ownership lookup failed for message_id=%s user_id=%s",
+            message_id,
+            user_id,
+        )
+        return False
+
+_SSE_HEADERS = {
+    "Cache-Control": "no-store",
+    "X-Accel-Buffering": "no",
+    "Connection": "keep-alive",
+    # Marks the response as served by the event-loop reader rather than the
+    # WSGI-threaded Flask fallback. Purely diagnostic — the frontend reads
+    # the body via fetch+getReader and ignores response headers.
+    "X-SSE-Transport": "async",
+}
+
+
+def _json(message: str, status_code: int) -> JSONResponse:
+    return JSONResponse(
+        {"success": False, "message": message}, status_code=status_code
+    )
+
+
+async def _acquire_stream_slot(user_id: str):
+    """Reserve a per-user connection slot; returns ``(redis, key)`` to release.
+
+    Mirrors the ``/api/events`` cap (INCR + safety TTL, reject when the
+    post-increment count exceeds ``SSE_MAX_CONCURRENT_PER_USER``). Returns
+    ``(None, None)`` when the cap is disabled or Redis is unavailable
+    (fail-open, like the notification stream). Raises ``_CapExceeded`` when
+    the user is over the cap so the caller can 429.
+    """
+    cap = int(getattr(settings, "SSE_MAX_CONCURRENT_PER_USER", 0))
+    if cap <= 0:
+        return None, None
+    redis = await get_async_redis_instance()
+    if redis is None:
+        return None, None
+    key = connection_counter_key(user_id)
+    try:
+        current = int(await redis.incr(key))
+    except Exception:
+        logger.debug("async SSE counter INCR failed for user=%s", user_id)
+        return None, None
+    # EXPIRE failure must not bypass the cap, so it's best-effort after INCR.
+    try:
+        await redis.expire(key, _COUNTER_TTL_SECONDS)
+    except Exception:
+        logger.debug("async SSE counter EXPIRE failed for user=%s", user_id)
+    if current > cap:
+        await _release_stream_slot(redis, key)
+        raise _CapExceeded()
+    return redis, key
+
+
+async def _release_stream_slot(redis, key) -> None:
+    if redis is None or key is None:
+        return
+    try:
+        await redis.decr(key)
+    except Exception:
+        logger.debug("async SSE counter DECR failed for key=%s", key)
+
+
+class _CapExceeded(Exception):
+    """Raised when a user is over their concurrent-stream cap."""
+
+
+async def _counted_stream(inner, redis, key):
+    """Wrap the reader so the per-user slot is released when it ends.
+
+    The slot is reserved before the response starts (so over-cap surfaces as
+    HTTP 429, not mid-stream); the release runs in ``finally`` on terminal
+    close, client disconnect, or error. Shielded so a disconnect-cancellation
+    can't skip the DECR and leak the count.
+    """
+    try:
+        async for line in inner:
+            yield line
+    finally:
+        with anyio.CancelScope(shield=True):
+            await _release_stream_slot(redis, key)
+
+
+async def stream_message_events(request: Request) -> JSONResponse | StreamingResponse:
+    """GET /api/messages/{message_id}/events — async reconnect tail.
+
+    Mirrors the Flask handler's gates (auth → id format → ownership →
+    cursor → per-user connection cap) then streams snapshot+tail off the
+    event loop.
+    """
+    # ``handle_auth`` only reads ``request.headers.get("Authorization")``;
+    # Starlette's headers are case-insensitive, so the Flask helper works
+    # verbatim. With AUTH_TYPE unset it returns ``{"sub": "local"}``.
+    decoded = handle_auth(request)
+    if isinstance(decoded, dict) and "error" in decoded:
+        return _json("Authentication error: invalid token", 401)
+    user_id = decoded.get("sub") if isinstance(decoded, dict) else None
+    if not user_id:
+        return _json("Authentication required", 401)
+    if settings.AUTH_TYPE == "oidc" and await anyio.to_thread.run_sync(
+        oidc_session_denied, decoded
+    ):
+        return _json("Authentication error: session revoked", 401)
+
+    message_id = request.path_params["message_id"]
+    if not _MESSAGE_ID_RE.match(message_id):
+        return _json("Invalid message id", 400)
+
+    # Ownership check is a sync DB read — push it off the loop.
+    owns = await anyio.to_thread.run_sync(_user_owns_message, message_id, user_id)
+    if not owns:
+        # Same opaque 404 as the Flask route — don't disclose existence.
+        return _json("Not found", 404)
+
+    # Per-user concurrent-connection cap — reserve before the response opens
+    # so an over-cap caller gets a clean 429 instead of a mid-stream cutoff.
+    try:
+        redis, counter_key = await _acquire_stream_slot(user_id)
+    except _CapExceeded:
+        logger.warning("sse.reconnect.rejected user_id=%s (over cap)", user_id)
+        return _json("Too many concurrent SSE connections", 429)
+
+    raw_cursor = request.headers.get("Last-Event-ID") or request.query_params.get(
+        "last_event_id"
+    )
+    last_event_id = _normalise_last_event_id(raw_cursor)
+    keepalive_seconds = float(
+        getattr(settings, "SSE_KEEPALIVE_SECONDS", DEFAULT_KEEPALIVE_SECONDS)
+    )
+
+    logger.info(
+        "message.event.connect.async message_id=%s user_id=%s last_event_id=%s",
+        message_id,
+        user_id,
+        last_event_id if last_event_id is not None else "-",
+    )
+
+    stream = build_message_event_stream_async(
+        message_id,
+        last_event_id=last_event_id,
+        user_id=user_id,
+        keepalive_seconds=keepalive_seconds,
+        poll_timeout_seconds=DEFAULT_POLL_TIMEOUT_SECONDS,
+    )
+    return StreamingResponse(
+        _counted_stream(stream, redis, counter_key),
+        media_type="text/event-stream",
+        headers=_SSE_HEADERS,
+    )
+
+
+# Mounted in ``application/asgi.py`` ahead of the Flask catch-all. Keep
+# each route's path identical to the Flask blueprint it shadows.
+async_sse_routes = [
+    Route(
+        "/api/messages/{message_id}/events",
+        stream_message_events,
+        methods=["GET"],
+    ),
+]