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:

docsgpt/streaming/async_redis.py

@@ -15,10 +15,15 @@ from typing import Optional
 
 import redis.asyncio as aioredis
 
+from docsgpt.cache import PUBSUB_SOCKET_TIMEOUT_SECONDS, _tcp_keepalive_options
 from docsgpt.core.settings import settings
 
 logger = logging.getLogger(__name__)
 
+# Upper bound on one Redis command made while opening or holding a stream, so
+# a stalled connection fails open instead of holding the request or cleanup.
+ASYNC_REDIS_OP_TIMEOUT_SECONDS = 5.0
+
 _async_redis: Optional[aioredis.Redis] = None
 _creation_failed = False
 
@@ -28,11 +33,12 @@ async def get_async_redis_instance() -> Optional[aioredis.Redis]:
 
     ``from_url`` builds the client without opening a socket (connection is
     lazy), so a transient broker outage surfaces later on the first command
-    rather than here. Mirrors the sync client's ``socket_connect_timeout``
-    and ``health_check_interval`` so a half-open TCP can't wedge the tail
-    loop past its keepalive cadence. Every open stream holds a pooled
-    connection, so the pool is sized by ``ASYNC_REDIS_MAX_CONNECTIONS``
-    rather than redis-py's default of 100.
+    rather than here. Uses the sync pub/sub client's timeouts and TCP
+    keepalive, so a connection NAT/IPVS dropped without a FIN fails a
+    ``BLPOP`` or ``XRANGE`` instead of hanging it. Pub/sub polls pass their
+    own read timeout, so ``socket_timeout`` doesn't cut idle subscribers.
+    Every open stream holds a pooled connection, so the pool is sized by
+    ``ASYNC_REDIS_MAX_CONNECTIONS`` rather than redis-py's default of 100.
     """
     global _async_redis, _creation_failed
     if _async_redis is None and not _creation_failed:
@@ -40,6 +46,9 @@ async def get_async_redis_instance() -> Optional[aioredis.Redis]:
             _async_redis = aioredis.Redis.from_url(
                 settings.CACHE_REDIS_URL,
                 socket_connect_timeout=2,
+                socket_timeout=PUBSUB_SOCKET_TIMEOUT_SECONDS,
+                socket_keepalive=True,
+                socket_keepalive_options=_tcp_keepalive_options(),
                 health_check_interval=10,
                 max_connections=int(settings.ASYNC_REDIS_MAX_CONNECTIONS),
             )

docsgpt/streaming/sse_leases.py

@@ -22,13 +22,10 @@ import anyio
 
 from docsgpt.core.settings import settings
 from docsgpt.events.keys import connection_leases_key
+from docsgpt.streaming.async_redis import ASYNC_REDIS_OP_TIMEOUT_SECONDS
 
 logger = logging.getLogger(__name__)
 
-# Upper bound on any one lease command, so a stalled Redis connection can't
-# hold a request, a stream, or its cleanup.
-_REDIS_OP_TIMEOUT_SECONDS = 5.0
-
 T = TypeVar("T")
 
 
@@ -63,9 +60,11 @@ class StreamLease:
             return
         self.refreshed_at = now
         try:
-            with anyio.fail_after(_REDIS_OP_TIMEOUT_SECONDS):
+            with anyio.fail_after(ASYNC_REDIS_OP_TIMEOUT_SECONDS):
                 async with self.redis.pipeline(transaction=True) as pipe:
-                    pipe.zadd(self.key, {self.lease_id: time.time()})
+                    # XX: a lease pruned or cleared while refreshes were failing
+                    # stays gone instead of retaking a slot past the cap.
+                    pipe.zadd(self.key, {self.lease_id: time.time()}, xx=True)
                     pipe.expire(self.key, int(self.ttl))
                     await pipe.execute()
         except Exception:
@@ -73,7 +72,7 @@ class StreamLease:
 
     async def release(self) -> None:
         """Remove this lease; shielded so it completes during cancellation, and bounded."""
-        with anyio.move_on_after(_REDIS_OP_TIMEOUT_SECONDS, shield=True):
+        with anyio.move_on_after(ASYNC_REDIS_OP_TIMEOUT_SECONDS, shield=True):
             try:
                 await self.redis.zrem(self.key, self.lease_id)
             except Exception:
@@ -99,21 +98,23 @@ async def acquire_stream_lease(redis: Any, user_id: str) -> Optional[StreamLease
         return None
     key = connection_leases_key(user_id)
     ttl = lease_ttl_seconds()
-    lease_id = uuid.uuid4().hex
+    lease = StreamLease(redis=redis, key=key, lease_id=uuid.uuid4().hex, ttl=ttl, refreshed_at=time.monotonic())
     now = time.time()
     try:
-        with anyio.fail_after(_REDIS_OP_TIMEOUT_SECONDS):
+        with anyio.fail_after(ASYNC_REDIS_OP_TIMEOUT_SECONDS):
             # One transaction, so concurrent connects are counted one at a time.
             async with redis.pipeline(transaction=True) as pipe:
                 pipe.zremrangebyscore(key, "-inf", now - ttl)
-                pipe.zadd(key, {lease_id: now})
+                pipe.zadd(key, {lease.lease_id: now})
                 pipe.zcard(key)
                 pipe.expire(key, int(ttl))
                 results = await pipe.execute()
     except Exception:
         logger.debug("SSE lease acquire failed for user=%s; failing open", user_id, exc_info=True)
+        # EXEC may have run before its reply was lost; the stream opens without
+        # a lease, so its member must not keep counting against the cap.
+        await lease.release()
         return None
-    lease = StreamLease(redis=redis, key=key, lease_id=lease_id, ttl=ttl, refreshed_at=time.monotonic())
     if int(results[2]) > cap:
         await lease.release()
         raise StreamCapExceeded()

tests/api/test_async_sse_routes.py

@@ -184,6 +184,24 @@ class TestOwnershipGate:
         assert ": connected" in r.text
 
 
+    def test_stream_logs_carry_request_context(self):
+        from docsgpt.core import log_context
+
+        seen: dict = {}
+
+        async def _gen(message_id, last_event_id=None, **kwargs):
+            seen.update(log_context.snapshot())
+            yield ": connected\n\n"
+
+        with patch(_AUTH, return_value={"sub": "alice"}), patch(
+            _OWNS, return_value=True
+        ), patch(_STREAM, _gen):
+            _client().get(f"/api/messages/{VALID_UUID}/events")
+        assert seen["endpoint"] == "message_events"
+        assert seen["user_id"] == "alice"
+        assert seen["activity_id"]
+
+
 # ── Last-Event-ID parsing ───────────────────────────────────────────────────
 
 

tests/api/test_events_routes.py

@@ -197,6 +197,18 @@ class TestStreamShape:
         assert record["poll_timeout"] == events_module.SUBSCRIBE_POLL_INTERVAL_SECONDS
         assert record["liveness_timeout"] == PUBSUB_SOCKET_TIMEOUT_SECONDS
 
+    def test_stream_logs_carry_the_flask_endpoint_name(self):
+        from docsgpt.core import log_context
+
+        seen: dict = {}
+        with patch(_AUTH, return_value=ALICE), patch(
+            _AREDIS, AsyncMock(return_value=_redis())
+        ), patch(_SUBSCRIBE, _subscribe(probe=lambda: seen.update(log_context.snapshot()))):
+            _get()
+        # Same value the Flask route logged, so saved log queries keep matching.
+        assert seen["endpoint"] == "event_stream.stream_events"
+        assert seen["user_id"] == "alice"
+
     def test_keepalive_emitted_on_idle_ticks(self, monkeypatch):
         monkeypatch.setattr(events_module.settings, "SSE_KEEPALIVE_SECONDS", 0)
         with patch(_AUTH, return_value=ALICE), patch(
@@ -474,7 +486,7 @@ class TestReplayBudget:
         assert await events_module._allow_replay(redis, "alice", OLD_CURSOR) is True
 
     async def test_fail_open_when_redis_stalls(self, monkeypatch):
-        monkeypatch.setattr(events_module, "_REDIS_OP_TIMEOUT_SECONDS", 0.05)
+        monkeypatch.setattr(events_module, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05)
         redis = _redis()
 
         async def _stalled(_key):

tests/asgi_stream.py

@@ -19,6 +19,7 @@ async def stream_then_disconnect(
     *,
     headers: Optional[Mapping[str, str]] = None,
     query_string: bytes = b"",
+    root_path: str = "",
     chunks_before_disconnect: int = 1,
     timeout: float = 5.0,
 ) -> tuple[Optional[int], dict[str, str], list[bytes]]:
@@ -26,8 +27,9 @@ async def stream_then_disconnect(
 
     A response that finishes on its own (a 401, a stream that ends) returns
     without any disconnect. ``chunks_before_disconnect=0`` disconnects as soon
-    as the request is read, before any body. ``timeout`` fails the test
-    instead of hanging it when a stream never ends.
+    as the request is read, before any body. ``root_path`` mounts the app under
+    a prefix, which ASGI servers include in ``path``. ``timeout`` fails the
+    test instead of hanging it when a stream never ends.
 
     Returns:
         tuple: ``(status, headers, body_chunks)`` with lower-cased header names.
@@ -46,9 +48,9 @@ async def stream_then_disconnect(
         "http_version": "1.1",
         "method": "GET",
         "scheme": "http",
-        "path": path,
-        "raw_path": path.encode("utf-8"),
-        "root_path": "",
+        "path": root_path + path,
+        "raw_path": (root_path + path).encode("utf-8"),
+        "root_path": root_path,
         "query_string": query_string,
         "headers": [
             (name.lower().encode("latin-1"), value.encode("latin-1"))

tests/devices/test_session_events_async.py

@@ -16,13 +16,16 @@ import json
 import time
 from types import SimpleNamespace
 
+import anyio
 import pytest
 from flask import Flask
 from starlette.applications import Starlette
+from starlette.requests import Request
 from starlette.testclient import TestClient
 
 from docsgpt.api.devices import auth as auth_module
 from docsgpt.api.devices import session_events as events_module
+from docsgpt.core import log_context
 from docsgpt.core.shutdown import begin_shutdown, reset_shutdown
 from tests.asgi_stream import stream_then_disconnect
 
@@ -95,6 +98,27 @@ def _path(session_id: str) -> str:
     return f"/api/devices/sessions/{session_id}/events"
 
 
+def _scope(session_id: str, headers: dict | None = None) -> dict:
+    """A raw ASGI GET scope for the stream route."""
+    return {
+        "type": "http",
+        "asgi": {"version": "3.0", "spec_version": "2.3"},
+        "http_version": "1.1",
+        "method": "GET",
+        "scheme": "http",
+        "path": _path(session_id),
+        "raw_path": _path(session_id).encode(),
+        "root_path": "",
+        "query_string": b"",
+        "headers": [
+            (name.lower().encode("latin-1"), value.encode("latin-1"))
+            for name, value in (AUTH if headers is None else headers).items()
+        ],
+        "client": ("testclient", 50000),
+        "server": ("testserver", 80),
+    }
+
+
 def _open(session_id: str, headers: dict | None = None):
     return TestClient(_app()).get(_path(session_id), headers=headers if headers is not None else AUTH)
 
@@ -150,6 +174,42 @@ class TestAuth:
         assert _open(ticket).status_code == 200
         assert device.touched == [DEVICE["id"]]
 
+    @pytest.mark.asyncio
+    async def test_unauthenticated_request_body_is_never_read(self, device, broker):
+        # An endless body must not be buffered before the token check rejects it.
+        reads = 0
+        sent: list = []
+
+        async def receive():
+            nonlocal reads
+            reads += 1
+            await anyio.sleep(0)
+            return {"type": "http.request", "body": b"x" * 65536, "more_body": True}
+
+        async def send(message):
+            sent.append(message)
+
+        with anyio.fail_after(2):
+            await _app()(_scope("st_x", headers={}), receive, send)
+        assert sent[0]["status"] == 401
+        assert reads == 0
+
+    def test_stream_logs_carry_request_context(self, device, broker, monkeypatch):
+        b, _ = broker
+        ticket = _ticket(b)
+        seen: dict = {}
+        original = b.next_command_async
+
+        async def _capture(sess, timeout=1.0):
+            seen.update(log_context.snapshot())
+            return await original(sess, timeout=timeout)
+
+        monkeypatch.setattr(b, "next_command_async", _capture)
+        assert _open(ticket).status_code == 200
+        assert seen["endpoint"] == "devices.session_events"
+        assert seen["user_id"] == DEVICE["user_id"]
+        assert seen["activity_id"]
+
 
 # ── stream ──────────────────────────────────────────────────────────────
 
@@ -271,6 +331,27 @@ class TestStreamLifecycle:
         assert b.get_session(ticket) is None
         assert DEVICE["id"] not in b._sessions_by_device
 
+    async def test_ticket_is_claimed_before_the_body_starts(self, device, broker):
+        b, fake = broker
+        ticket = _ticket(b)
+        scope = {**_scope(ticket), "path_params": {"session_id": ticket}}
+
+        async def receive():
+            return {"type": "http.disconnect"}
+
+        response = await events_module.device_session_events(Request(scope, receive))
+        assert response.status_code == 200
+        # A response whose body never runs has still used the ticket up.
+        assert fake.get(f"dev:ticket:{DEVICE['id']}") is None
+        assert b.get_session(ticket) is not None
+
+        async def send(_message):
+            return None
+
+        with anyio.fail_after(2):
+            await response(scope, receive, send)
+        assert b.get_session(ticket) is None
+
 
 # ── machine-key signatures ──────────────────────────────────────────────
 
@@ -320,6 +401,26 @@ class TestSignature:
         # Rejected before the session opened: the ticket is still unclaimed.
         assert fake.get(f"dev:ticket:{DEVICE['id']}") == ticket.encode()
 
+    @pytest.mark.asyncio
+    async def test_signed_path_excludes_root_path(self, device, broker, monkeypatch):
+        # The Flask device routes verify against PATH_INFO, which has the mount
+        # prefix stripped; the stream must check the same path.
+        monkeypatch.setattr(auth_module.settings, "REMOTE_DEVICE_REQUIRE_SIGNATURE", True)
+        b, _ = broker
+        ticket = _ticket(b)
+        fingerprint, headers = _signed_headers(_path(ticket))
+        device.row["machine_pubkey_fingerprint"] = fingerprint
+        status, _, chunks = await stream_then_disconnect(
+            _app(),
+            _path(ticket),
+            headers=headers,
+            root_path="/docsgpt",
+            chunks_before_disconnect=10**6,
+            timeout=3,
+        )
+        assert status == 200
+        assert b"event: invocation" in b"".join(chunks)
+
     def test_missing_signature_rejected(self, device, broker, monkeypatch):
         monkeypatch.setattr(auth_module.settings, "REMOTE_DEVICE_REQUIRE_SIGNATURE", True)
         b, _ = broker

tests/fake_async_redis.py

@@ -3,8 +3,9 @@
 Sorted sets and transactional pipelines behave like Redis, enough for the
 per-connection leases. ``incr``, ``expire``, ``xrange`` and ``xinfo_stream``
 are ``AsyncMock`` attributes that tests configure directly. Set ``fail`` to
-make sorted-set commands and pipelines raise, or ``hang`` to make them never
-return (a dead connection).
+make sorted-set commands and pipelines raise, ``hang`` to make them never
+return (a dead connection), or ``hang_after_execute`` to make a pipeline apply
+its commands and then never reply (a reply lost after the server ran EXEC).
 """
 
 from __future__ import annotations
@@ -37,6 +38,7 @@ class FakeAsyncRedis:
         self.key_ttls: dict[str, int] = {}
         self.fail: Optional[BaseException] = None
         self.hang = False
+        self.hang_after_execute = False
         self.incr = AsyncMock(return_value=1)
         self.expire = AsyncMock(return_value=True)
         self.xrange = AsyncMock(return_value=[])
@@ -146,4 +148,6 @@ class FakePipeline:
         await self._redis._gate()
         results = [fn(*args, **kwargs) for fn, args, kwargs in self._ops]
         self._ops.clear()
+        if self._redis.hang_after_execute:
+            await anyio.sleep_forever()
         return results

tests/test_async_redis.py

@@ -40,6 +40,17 @@ class TestGetAsyncRedisInstance:
         assert kwargs["socket_connect_timeout"] == 2
         assert kwargs["health_check_interval"] == 10
 
+    async def test_reads_bounded_like_the_sync_pubsub_client(self):
+        # A silently dropped connection must fail a BLPOP or XRANGE, not hang it.
+        from docsgpt.cache import PUBSUB_SOCKET_TIMEOUT_SECONDS
+
+        with patch(_FROM_URL) as from_url:
+            await async_redis.get_async_redis_instance()
+        kwargs = from_url.call_args.kwargs
+        assert kwargs["socket_timeout"] == PUBSUB_SOCKET_TIMEOUT_SECONDS
+        assert kwargs["socket_keepalive"] is True
+        assert "socket_keepalive_options" in kwargs
+
     async def test_real_client_pool_honours_setting(self, monkeypatch):
         # from_url connects lazily, so this never opens a socket.
         monkeypatch.setattr(async_redis.settings, "ASYNC_REDIS_MAX_CONNECTIONS", 777)

tests/test_sse_leases.py

@@ -89,12 +89,22 @@ class TestAcquire:
         assert await acquire_stream_lease(redis, "alice") is None
 
     async def test_stalled_redis_fails_open_within_timeout(self, monkeypatch):
-        monkeypatch.setattr(sse_leases, "_REDIS_OP_TIMEOUT_SECONDS", 0.05)
+        monkeypatch.setattr(sse_leases, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05)
         redis = FakeAsyncRedis()
         redis.hang = True
         with anyio.fail_after(2):
             assert await acquire_stream_lease(redis, "alice") is None
 
+    async def test_timed_out_acquire_leaves_no_orphan_lease(self, monkeypatch):
+        # EXEC ran on the server but its reply never arrived: the stream opens
+        # without a lease, so the member it added must not keep counting.
+        monkeypatch.setattr(sse_leases, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05)
+        redis = FakeAsyncRedis()
+        redis.hang_after_execute = True
+        with anyio.fail_after(2):
+            assert await acquire_stream_lease(redis, "alice") is None
+        assert redis.zsets == {}
+
 
 @pytest.mark.unit
 @pytest.mark.asyncio
@@ -110,7 +120,7 @@ class TestLeaseLifecycle:
         assert set(redis.zsets[KEY]) == {second.lease_id}
 
     async def test_release_is_bounded_when_redis_stalls(self, monkeypatch):
-        monkeypatch.setattr(sse_leases, "_REDIS_OP_TIMEOUT_SECONDS", 0.05)
+        monkeypatch.setattr(sse_leases, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05)
         redis = FakeAsyncRedis()
         lease = await acquire_stream_lease(redis, "alice")
         redis.hang = True
@@ -144,6 +154,26 @@ class TestLeaseLifecycle:
         assert long_lived.lease_id not in redis.zsets[KEY]
         assert len(redis.zsets[KEY]) == 7
 
+    async def test_refresh_does_not_revive_a_pruned_lease(self, clock):
+        # Refreshes failed past the TTL, a newer connect pruned the lease and
+        # took the slot. Re-adding it would put the user over the cap.
+        redis = FakeAsyncRedis()
+        lease = await acquire_stream_lease(redis, "alice")
+        redis.zsets[KEY].pop(lease.lease_id)
+        _fill(redis, 8, clock["t"])
+        clock["t"] += 30
+        await lease.refresh_if_due()
+        assert lease.lease_id not in redis.zsets[KEY]
+        assert len(redis.zsets[KEY]) == 8
+
+    async def test_refresh_does_not_undo_a_manual_clear(self, clock):
+        redis = FakeAsyncRedis()
+        lease = await acquire_stream_lease(redis, "alice")
+        redis.zsets.clear()  # redis-cli DEL user:<id>:sse_leases
+        clock["t"] += 30
+        await lease.refresh_if_due()
+        assert redis.zsets == {}
+
     async def test_refresh_failure_is_swallowed(self, clock):
         redis = FakeAsyncRedis()
         lease = await acquire_stream_lease(redis, "alice")