fix: stability improvements on event loop

docsgpt/api/asgi_stream.py

@@ -3,12 +3,14 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, AsyncIterable, Awaitable, Callable, Optional
+from typing import Any, AsyncIterable, AsyncIterator, Awaitable, Callable, Optional
 
 import anyio
 from starlette.responses import StreamingResponse
 from starlette.types import Receive, Scope, Send
 
+from docsgpt.streaming.sse_leases import StreamLease, hold_lease
+
 logger = logging.getLogger(__name__)
 
 # Upper bound on each cleanup step. Cleanup is shielded from the cancellation
@@ -16,6 +18,15 @@ logger = logging.getLogger(__name__)
 # hold the request task, and a graceful shutdown, indefinitely.
 _CLEANUP_TIMEOUT_SECONDS = 5.0
 
+# Headers for the browser-facing SSE streams (notifications, chat reconnect).
+SSE_HEADERS = {
+    "Cache-Control": "no-store",
+    "X-Accel-Buffering": "no",
+    "Connection": "keep-alive",
+    # Marks the response as served by the event loop. Purely diagnostic.
+    "X-SSE-Transport": "async",
+}
+
 
 class ClosingStreamingResponse(StreamingResponse):
     """``StreamingResponse`` that always finalizes its stream.
@@ -50,6 +61,16 @@ class ClosingStreamingResponse(StreamingResponse):
                 await _bounded(self._on_close, "response on_close hook")
 
 
+def sse_response(stream: AsyncIterator[str], lease: Optional[StreamLease]) -> ClosingStreamingResponse:
+    """Serve ``stream`` as SSE, refreshing ``lease`` between frames and releasing it however the response ends."""
+    return ClosingStreamingResponse(
+        hold_lease(stream, lease),
+        media_type="text/event-stream",
+        headers=SSE_HEADERS,
+        on_close=lease.release if lease is not None else None,
+    )
+
+
 async def _bounded(step: Callable[[], Awaitable[None]], what: str) -> None:
     """Run one cleanup step shielded from cancellation, giving up after the cleanup timeout."""
     with anyio.move_on_after(_CLEANUP_TIMEOUT_SECONDS, shield=True) as scope:

docsgpt/api/async_sse.py

@@ -21,11 +21,11 @@ from typing import Optional
 import anyio
 from sqlalchemy import text
 from starlette.requests import Request
-from starlette.responses import JSONResponse, Response
+from starlette.responses import Response
 from starlette.routing import Route
 
-from docsgpt.api.asgi_auth import authenticate
-from docsgpt.api.asgi_stream import ClosingStreamingResponse
+from docsgpt.api.asgi_auth import authenticate, bind_log_context, json_error
+from docsgpt.api.asgi_stream import sse_response
 from docsgpt.core.settings import settings
 from docsgpt.storage.db.session import db_readonly
 from docsgpt.streaming.async_event_replay import (
@@ -36,11 +36,7 @@ from docsgpt.streaming.event_replay import (
     DEFAULT_KEEPALIVE_SECONDS,
     DEFAULT_POLL_TIMEOUT_SECONDS,
 )
-from docsgpt.streaming.sse_leases import (
-    StreamCapExceeded,
-    acquire_stream_lease,
-    hold_lease,
-)
+from docsgpt.streaming.sse_leases import StreamCapExceeded, acquire_stream_lease
 
 logger = logging.getLogger(__name__)
 
@@ -89,22 +85,6 @@ def _user_owns_message(message_id: str, user_id: str) -> bool:
         )
         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 stream_message_events(request: Request) -> Response:
     """GET /api/messages/{message_id}/events — async reconnect tail.
@@ -120,17 +100,18 @@ async def stream_message_events(request: Request) -> Response:
         return error
     user_id = decoded.get("sub") if isinstance(decoded, dict) else None
     if not user_id:
-        return _json("Authentication required", 401)
+        return json_error("Authentication required", 401)
+    bind_log_context("message_events", user_id)
 
     message_id = request.path_params["message_id"]
     if not _MESSAGE_ID_RE.match(message_id):
-        return _json("Invalid message id", 400)
+        return json_error("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)
+        return json_error("Not found", 404)
 
     # Per-user connection cap, shared with /api/events so it bounds a user's
     # total live SSE footprint. Reserved before the response opens so an
@@ -139,7 +120,7 @@ async def stream_message_events(request: Request) -> Response:
         lease = await acquire_stream_lease(await get_async_redis_instance(), user_id)
     except StreamCapExceeded:
         logger.warning("sse.reconnect.rejected user_id=%s (over cap)", user_id)
-        return _json("Too many concurrent SSE connections", 429)
+        return json_error("Too many concurrent SSE connections", 429)
 
     raw_cursor = request.headers.get("Last-Event-ID") or request.query_params.get(
         "last_event_id"
@@ -163,13 +144,7 @@ async def stream_message_events(request: Request) -> Response:
         keepalive_seconds=keepalive_seconds,
         poll_timeout_seconds=DEFAULT_POLL_TIMEOUT_SECONDS,
     )
-    return ClosingStreamingResponse(
-        hold_lease(stream, lease),
-        media_type="text/event-stream",
-        headers=_SSE_HEADERS,
-        # Released once the response is over, however it ended.
-        on_close=lease.release if lease is not None else None,
-    )
+    return sse_response(stream, lease)
 
 
 # Mounted in ``docsgpt/asgi.py`` ahead of the Flask catch-all. Keep

docsgpt/api/devices/session_events.py

@@ -2,7 +2,7 @@
 
 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
+event loop: the token check and ticket claim run in worker threads, then
 commands arrive through the broker's async ``BLPOP`` and an idle session holds
 no thread.
 """
@@ -13,18 +13,20 @@ import json
 import logging
 import time
 from functools import partial
-from typing import AsyncIterator
+from typing import AsyncIterator, Optional
 
 import anyio
 from starlette.requests import Request
 from starlette.responses import JSONResponse, Response
 from starlette.routing import Route
 
+from docsgpt.api.asgi_auth import bind_log_context
 from docsgpt.api.asgi_stream import ClosingStreamingResponse
 from docsgpt.api.devices.auth import authenticate_device
+from docsgpt.core import log_context
 from docsgpt.core.settings import settings
 from docsgpt.core.shutdown import is_shutting_down
-from docsgpt.devices.broker import DeviceBroker, get_broker
+from docsgpt.devices.broker import DeviceBroker, SessionState, get_broker
 
 logger = logging.getLogger(__name__)
 
@@ -52,45 +54,57 @@ def _sse_event(name: str, payload: dict, event_id: int) -> str:
     )
 
 
-async def _session_stream(broker: DeviceBroker, device: dict) -> AsyncIterator[str]:
+def _signed_path(request: Request) -> str:
+    """Return the path without the mount prefix, as the Flask device routes see it in ``PATH_INFO``."""
+    path = request.scope["path"]
+    root = request.scope.get("root_path", "")
+    if root and path.startswith(root):
+        path = path[len(root):]
+    return path
+
+
+def _claim_session(broker: DeviceBroker, device: dict, session_id: str) -> Optional[SessionState]:
+    """Open the device's session iff ``session_id`` is its unexpired ticket; runs in a worker thread."""
+    if not broker.validate_ticket(device["id"], session_id):
+        return None
+    return broker.register_session(device["id"], device["user_id"])
+
+
+async def _close_session(broker: DeviceBroker, session_id: str) -> None:
+    """Close the session once its response is over."""
+    broker.close_session(session_id, reason="stream_end")
+
+
+async def _session_stream(broker: DeviceBroker, sess: SessionState) -> 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
+    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
-            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")
+            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()
 
 
 async def device_session_events(request: Request) -> Response:
@@ -102,25 +116,34 @@ async def device_session_events(request: Request) -> Response:
     before any stream is opened.
     """
     session_id = request.path_params["session_id"]
-    body = await request.body()
+    bind_log_context("devices.session_events")
     device, failure = await anyio.to_thread.run_sync(
         partial(
             authenticate_device,
             request.headers,
             request.method,
-            request.url.path,
-            lambda: body,
+            _signed_path(request),
+            # Read only when a signature is checked, after the token is known
+            # good, so an unauthenticated body is never buffered.
+            lambda: anyio.from_thread.run(request.body),
         )
     )
     if failure is not None:
         return _error(*failure)
+    log_context.bind(user_id=device["user_id"])
     broker = get_broker()
-    if not await anyio.to_thread.run_sync(broker.validate_ticket, device["id"], session_id):
+    # Claiming consumes the ticket, so it happens before the response opens and
+    # ``on_close`` closes the session even if the body never starts. Shielded
+    # so a cancellation can't leave a registered session with no response.
+    with anyio.CancelScope(shield=True):
+        sess = await anyio.to_thread.run_sync(_claim_session, broker, device, session_id)
+    if sess is None:
         return _error("session_ticket_invalid", 410)
     return ClosingStreamingResponse(
-        _session_stream(broker, device),
+        _session_stream(broker, sess),
         media_type="text/event-stream",
         headers=_SSE_HEADERS,
+        on_close=partial(_close_session, broker, sess.session_id),
     )
 
 

docsgpt/api/events/routes.py

@@ -26,7 +26,7 @@ from starlette.responses import Response
 from starlette.routing import Route
 
 from docsgpt.api.asgi_auth import authenticate, bind_log_context, json_error
-from docsgpt.api.asgi_stream import ClosingStreamingResponse
+from docsgpt.api.asgi_stream import sse_response
 from docsgpt.cache import PUBSUB_SOCKET_TIMEOUT_SECONDS
 from docsgpt.core.settings import settings
 from docsgpt.core.shutdown import is_shutting_down
@@ -37,21 +37,13 @@ from docsgpt.events.keys import (
     topic_name,
 )
 from docsgpt.streaming.async_broadcast_channel import AsyncTopic
-from docsgpt.streaming.async_redis import get_async_redis_instance
-from docsgpt.streaming.sse_leases import (
-    StreamCapExceeded,
-    acquire_stream_lease,
-    hold_lease,
-)
+from docsgpt.streaming.async_redis import ASYNC_REDIS_OP_TIMEOUT_SECONDS, get_async_redis_instance
+from docsgpt.streaming.sse_leases import StreamCapExceeded, acquire_stream_lease
 
 logger = logging.getLogger(__name__)
 
 SUBSCRIBE_POLL_INTERVAL_SECONDS = 1.0
 
-# Upper bound on the replay-budget check, so a stalled Redis connection fails
-# open instead of holding the request.
-_REDIS_OP_TIMEOUT_SECONDS = 5.0
-
 # WHATWG SSE treats CRLF, CR, and LF equivalently as line terminators.
 _SSE_LINE_SPLIT = re.compile(r"\r\n|\r|\n")
 
@@ -61,14 +53,6 @@ _SSE_LINE_SPLIT = re.compile(r"\r\n|\r|\n")
 # would silently fail.
 _STREAM_ID_RE = re.compile(r"^\d+(-\d+)?$")
 
-_SSE_HEADERS = {
-    "Cache-Control": "no-store",
-    "X-Accel-Buffering": "no",
-    "Connection": "keep-alive",
-    # Marks the response as served by the event loop. Purely diagnostic.
-    "X-SSE-Transport": "async",
-}
-
 # Only emitted at most once per process so a misconfigured deployment
 # doesn't drown the logs.
 _local_user_warned = False
@@ -146,7 +130,7 @@ async def _allow_replay(
     window = max(1, int(settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS))
     key = replay_budget_key(user_id)
     try:
-        with anyio.fail_after(_REDIS_OP_TIMEOUT_SECONDS):
+        with anyio.fail_after(ASYNC_REDIS_OP_TIMEOUT_SECONDS):
             used = int(await redis_client.incr(key))
             # Always (re)seed the TTL. Gating on ``used == 1`` would wedge
             # the counter forever if INCR succeeds but EXPIRE raises on
@@ -431,7 +415,8 @@ async def stream_events(request: Request) -> Response:
     user_id = decoded.get("sub") if isinstance(decoded, dict) else None
     if not user_id:
         return json_error("Authentication required", 401)
-    bind_log_context("event_stream", user_id)
+    # Same endpoint value the Flask route logged, so saved log queries keep matching.
+    bind_log_context("event_stream.stream_events", user_id)
 
     # In dev deployments without AUTH_TYPE configured, every request
     # resolves to user_id="local" and shares one stream. Surface this so
@@ -476,13 +461,7 @@ async def stream_events(request: Request) -> Response:
         " (rejected_invalid)" if last_event_id_invalid else "",
     )
     stream = _event_stream(redis_client, user_id, last_event_id, last_event_id_invalid, push_enabled)
-    return ClosingStreamingResponse(
-        hold_lease(stream, lease),
-        media_type="text/event-stream",
-        headers=_SSE_HEADERS,
-        # Released once the response is over, however it ended.
-        on_close=lease.release if lease is not None else None,
-    )
+    return sse_response(stream, lease)
 
 
 # Mounted in ``docsgpt/asgi.py`` ahead of the Flask catch-all.

docsgpt/api/user/artifacts/download.py

@@ -51,6 +51,7 @@ def _sanitize_header_filename(filename: Optional[str], fallback: str) -> str:
 
 
 def _ascii_fold(value: str) -> str:
+    """Drop accents and any other non-ASCII characters from ``value``."""
     return unicodedata.normalize("NFKD", value).encode("ascii", "ignore").decode("ascii")
 
 
@@ -130,7 +131,8 @@ async def download_artifact(request: Request) -> Response:
     decoded, error = await authenticate(request)
     if error is not None:
         return error
-    bind_log_context("artifact_download", decoded.get("sub") if decoded else None)
+    # Same endpoint value the Flask-RESTX route logged, so saved log queries keep matching.
+    bind_log_context("artifacts_download_artifact", decoded.get("sub") if decoded else None)
     query = request.query_params
 
     try: