fix(api): address review on the event-loop streams

- Per-user SSE cap uses per-connection leases in a sorted set instead of
  a shared INCR/DECR counter with a TTL. A stream that outlived the TTL
  could let the counter expire, then decrement another stream's slot or
  drive the count negative. Leases refresh while a stream sends frames
  and age out when a stream dies without cleanup.
- Shielded cleanup is bounded per step: stream close and on_close in
  ClosingStreamingResponse, unsubscribe and close in AsyncTopic, lease
  release, and the replay-budget check, so a dead Redis connection can't
  hold a request or a graceful shutdown.
- Artifact downloads send an ASCII filename plus an RFC 5987 filename*
  for non-ASCII names, build headers before opening the file, and close
  disk handles off the event loop.
- Type hints and docstrings on the new helpers.

docsgpt/api/asgi_stream.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Awaitable, Callable, Optional
+from typing import Any, AsyncIterable, Awaitable, Callable, Optional
 
 import anyio
 from starlette.responses import StreamingResponse
@@ -11,6 +11,11 @@ from starlette.types import Receive, Scope, Send
 
 logger = logging.getLogger(__name__)
 
+# Upper bound on each cleanup step. Cleanup is shielded from the cancellation
+# that ended the response, so a step stuck on a dead connection would otherwise
+# hold the request task, and a graceful shutdown, indefinitely.
+_CLEANUP_TIMEOUT_SECONDS = 5.0
+
 
 class ClosingStreamingResponse(StreamingResponse):
     """``StreamingResponse`` that always finalizes its stream.
@@ -18,11 +23,19 @@ class ClosingStreamingResponse(StreamingResponse):
     When a client disconnects, Starlette cancels the send loop and drops the
     body iterator: a generator paused between frames, or one that never
     started, would leave its cleanup to garbage collection. This closes the
-    iterator and then runs ``on_close`` once the response is over, shielded
-    from the cancellation that ended it.
+    iterator and then runs ``on_close`` once the response is over. Each step is
+    shielded from the cancellation that ended the response and bounded by
+    ``_CLEANUP_TIMEOUT_SECONDS``, so ``on_close`` still runs if closing the
+    iterator hangs.
     """
 
-    def __init__(self, content, *, on_close: Optional[Callable[[], Awaitable[None]]] = None, **kwargs) -> None:
+    def __init__(
+        self,
+        content: AsyncIterable[Any],
+        *,
+        on_close: Optional[Callable[[], Awaitable[None]]] = None,
+        **kwargs: Any,
+    ) -> None:
         super().__init__(content, **kwargs)
         self._on_close = on_close
 
@@ -30,15 +43,19 @@ class ClosingStreamingResponse(StreamingResponse):
         try:
             await super().__call__(scope, receive, send)
         finally:
-            with anyio.CancelScope(shield=True):
-                aclose = getattr(self.body_iterator, "aclose", None)
-                if aclose is not None:
-                    try:
-                        await aclose()
-                    except Exception:
-                        logger.exception("closing the response stream failed")
-                if self._on_close is not None:
-                    try:
-                        await self._on_close()
-                    except Exception:
-                        logger.exception("response on_close hook failed")
+            aclose = getattr(self.body_iterator, "aclose", None)
+            if aclose is not None:
+                await _bounded(aclose, "closing the response stream")
+            if self._on_close is not None:
+                await _bounded(self._on_close, "response on_close hook")
+
+
+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:
+        try:
+            await step()
+        except Exception:
+            logger.exception("%s failed", what)
+    if scope.cancelled_caught:
+        logger.warning("%s timed out after %ss", what, _CLEANUP_TIMEOUT_SECONDS)

docsgpt/api/async_sse.py

@@ -21,12 +21,12 @@ from typing import Optional
 import anyio
 from sqlalchemy import text
 from starlette.requests import Request
-from starlette.responses import JSONResponse, StreamingResponse
+from starlette.responses import JSONResponse, Response
 from starlette.routing import Route
 
 from docsgpt.api.asgi_auth import authenticate
+from docsgpt.api.asgi_stream import ClosingStreamingResponse
 from docsgpt.core.settings import settings
-from docsgpt.events.keys import connection_counter_key
 from docsgpt.storage.db.session import db_readonly
 from docsgpt.streaming.async_event_replay import (
     build_message_event_stream_async,
@@ -36,15 +36,14 @@ 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,
+)
 
 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(
@@ -107,68 +106,7 @@ def _json(message: str, status_code: int) -> JSONResponse:
     )
 
 
-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:
+async def stream_message_events(request: Request) -> Response:
     """GET /api/messages/{message_id}/events — async reconnect tail.
 
     Mirrors the Flask handler's gates (auth → id format → ownership →
@@ -194,11 +132,12 @@ async def stream_message_events(request: Request) -> JSONResponse | StreamingRes
         # 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.
+    # 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
+    # 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:
+        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)
 
@@ -224,10 +163,12 @@ async def stream_message_events(request: Request) -> JSONResponse | StreamingRes
         keepalive_seconds=keepalive_seconds,
         poll_timeout_seconds=DEFAULT_POLL_TIMEOUT_SECONDS,
     )
-    return StreamingResponse(
-        _counted_stream(stream, redis, counter_key),
+    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,
     )
 
 

docsgpt/api/devices/session_events.py

@@ -39,10 +39,12 @@ _SSE_HEADERS = {
 
 
 def _error(code: str, status_code: int) -> JSONResponse:
+    """Return the ``{"success": false, "error": code}`` body the Flask device routes use."""
     return JSONResponse({"success": False, "error": code}, status_code=status_code)
 
 
 def _sse_event(name: str, payload: dict, event_id: int) -> str:
+    """Encode one SSE record with an event name, an id and a JSON data line."""
     return (
         f"event: {name}\n"
         f"id: {event_id}\n"

docsgpt/api/events/routes.py

@@ -17,9 +17,10 @@ import logging
 import re
 import time
 from contextlib import aclosing
-from functools import partial
 from typing import AsyncIterator, Optional
 
+import anyio
+from redis.asyncio import Redis
 from starlette.requests import Request
 from starlette.responses import Response
 from starlette.routing import Route
@@ -30,7 +31,6 @@ from docsgpt.cache import PUBSUB_SOCKET_TIMEOUT_SECONDS
 from docsgpt.core.settings import settings
 from docsgpt.core.shutdown import is_shutting_down
 from docsgpt.events.keys import (
-    connection_counter_key,
     replay_budget_key,
     stream_id_compare,
     stream_key,
@@ -38,14 +38,19 @@ from docsgpt.events.keys import (
 )
 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,
+)
 
 logger = logging.getLogger(__name__)
 
 SUBSCRIBE_POLL_INTERVAL_SECONDS = 1.0
 
-# Safety TTL on the per-user connection counter: counts orphaned by a hard
-# crash self-heal after this window.
-_COUNTER_TTL_SECONDS = 3600
+# 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")
@@ -95,7 +100,7 @@ def _decode(value) -> Optional[str]:
     return str(value)
 
 
-async def _oldest_retained_id(redis_client, user_id: str) -> Optional[str]:
+async def _oldest_retained_id(redis_client: Redis, user_id: str) -> Optional[str]:
     """Return the id of the oldest entry still in the stream, or ``None``.
 
     Used to detect ``Last-Event-ID`` having slid off the back of the
@@ -120,14 +125,15 @@ async def _oldest_retained_id(redis_client, user_id: str) -> Optional[str]:
 
 
 async def _allow_replay(
-    redis_client, user_id: str, last_event_id: Optional[str]
+    redis_client: Optional[Redis], user_id: str, last_event_id: Optional[str]
 ) -> bool:
     """Per-user sliding-window snapshot-replay budget.
 
-    Fails open on Redis errors or when the budget is disabled. No-cursor
-    connects never consume budget: fresh sessions start live and do no
-    snapshot work, so counting them only starved real reconnects (a burst
-    of fresh tabs could 429 a user's cursor-bearing reconnect).
+    Fails open on Redis errors, a stalled connection, or when the budget is
+    disabled. No-cursor connects never consume budget: fresh sessions start
+    live and do no snapshot work, so counting them only starved real
+    reconnects (a burst of fresh tabs could 429 a user's cursor-bearing
+    reconnect).
     """
     budget = int(settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW)
     if budget <= 0:
@@ -140,12 +146,13 @@ async def _allow_replay(
     window = max(1, int(settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS))
     key = replay_budget_key(user_id)
     try:
-        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
-        # the seeding call. EXPIRE on an existing key resets the TTL
-        # to ``window`` — within ±1s of the per-window budget semantic.
-        await redis_client.expire(key, window)
+        with anyio.fail_after(_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
+            # the seeding call. EXPIRE on an existing key resets the TTL
+            # to ``window`` — within ±1s of the per-window budget semantic.
+            await redis_client.expire(key, window)
     except Exception:
         logger.debug(
             "replay budget probe failed for user=%s; failing open",
@@ -188,7 +195,7 @@ def _replay_floor_id() -> Optional[str]:
 
 
 async def _replay_backlog(
-    redis_client, user_id: str, last_event_id: Optional[str], max_count: int
+    redis_client: Redis, user_id: str, last_event_id: Optional[str], max_count: int
 ) -> list[tuple[str, str]]:
     """Return ``(entry_id, sse_line)`` pairs for backlog entries past ``last_event_id``.
 
@@ -261,15 +268,8 @@ def _truncation_notice_line(oldest_id: str) -> str:
     )
 
 
-async def _release_slot(redis_client, counter_key: str) -> None:
-    try:
-        await redis_client.decr(counter_key)
-    except Exception:
-        logger.debug("SSE connection counter DECR failed for key=%s", counter_key)
-
-
 async def _event_stream(
-    redis_client,
+    redis_client: Optional[Redis],
     user_id: str,
     last_event_id: Optional[str],
     last_event_id_invalid: bool,
@@ -451,34 +451,22 @@ async def stream_events(request: Request) -> Response:
     last_event_id_invalid = raw_last_event_id is not None and last_event_id is None
 
     push_enabled = settings.ENABLE_SSE_PUSH
-    cap = int(settings.SSE_MAX_CONCURRENT_PER_USER)
     redis_client = await get_async_redis_instance() if push_enabled else None
-    counter_key = connection_counter_key(user_id)
-    counted = False
 
-    if redis_client is not None and cap > 0:
-        try:
-            current = int(await redis_client.incr(counter_key))
-            counted = True
-        except Exception:
-            logger.debug("SSE connection counter INCR failed for user=%s", user_id)
-        if counted:
-            # EXPIRE failure must NOT bypass the cap.
-            try:
-                await redis_client.expire(counter_key, _COUNTER_TTL_SECONDS)
-            except Exception:
-                logger.debug("SSE connection counter EXPIRE failed for user=%s", user_id)
-            if current > cap:
-                await _release_slot(redis_client, counter_key)
-                return json_error("Too many concurrent SSE connections", 429)
+    # Reserve a per-user connection slot before the response opens, so an
+    # over-cap caller gets a clean 429 instead of a mid-stream cutoff.
+    try:
+        lease = await acquire_stream_lease(redis_client, user_id)
+    except StreamCapExceeded:
+        return json_error("Too many concurrent SSE connections", 429)
 
     # Replay budget is checked before the stream opens so a denial surfaces
     # as HTTP 429 instead of a silent snapshot skip: a live tail carrying
     # ``id:`` headers would advance the client's cursor past the entries it
     # never received. 429 keeps the cursor pinned and the frontend backs off.
     if redis_client is not None and not await _allow_replay(redis_client, user_id, last_event_id):
-        if counted:
-            await _release_slot(redis_client, counter_key)
+        if lease is not None:
+            await lease.release()
         return json_error("Replay budget exhausted", 429)
 
     logger.info(
@@ -487,12 +475,13 @@ async def stream_events(request: Request) -> Response:
         last_event_id or "-",
         " (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(
-        _event_stream(redis_client, user_id, last_event_id, last_event_id_invalid, push_enabled),
+        hold_lease(stream, lease),
         media_type="text/event-stream",
         headers=_SSE_HEADERS,
         # Released once the response is over, however it ended.
-        on_close=partial(_release_slot, redis_client, counter_key) if counted else None,
+        on_close=lease.release if lease is not None else None,
     )
 
 

docsgpt/api/user/artifacts/download.py

@@ -11,9 +11,12 @@ from __future__ import annotations
 
 import io
 import logging
+import os
 import re
+import unicodedata
 from functools import partial
 from typing import AsyncIterator, BinaryIO, Mapping, Optional, Tuple
+from urllib.parse import quote
 
 import anyio
 from starlette.requests import Request
@@ -47,6 +50,25 @@ def _sanitize_header_filename(filename: Optional[str], fallback: str) -> str:
     return cleaned or fallback
 
 
+def _ascii_fold(value: str) -> str:
+    return unicodedata.normalize("NFKD", value).encode("ascii", "ignore").decode("ascii")
+
+
+def _content_disposition(filename: Optional[str], artifact_id: str) -> str:
+    """Build an ``attachment`` Content-Disposition value that any filename can travel in.
+
+    Header values are Latin-1, so a non-ASCII name gets an ASCII ``filename``
+    fallback plus an RFC 5987 ``filename*`` carrying the real name.
+    """
+    name = _sanitize_header_filename(filename, f"artifact-{artifact_id}")
+    if name.isascii():
+        return f'attachment; filename="{name}"'
+    stem, ext = os.path.splitext(name)
+    ascii_stem = _ascii_fold(stem).strip()
+    fallback = f"{ascii_stem or f'artifact-{artifact_id}'}{_ascii_fold(ext)}"
+    return f"attachment; filename=\"{fallback}\"; filename*=UTF-8''{quote(name, safe='')}"
+
+
 def _load_version(
     artifact_id: str, decoded_token: Optional[dict], query: Mapping[str, str]
 ) -> Tuple[Optional[dict], Optional[Tuple[str, int]]]:
@@ -90,9 +112,14 @@ async def _read_chunks(file_obj: BinaryIO) -> AsyncIterator[bytes]:
 
 
 async def _close_file(file_obj: BinaryIO) -> None:
+    """Close the handle; closing anything but an in-memory buffer runs in a worker thread."""
     close = getattr(file_obj, "close", None)
-    if callable(close):
+    if not callable(close):
+        return
+    if isinstance(file_obj, io.BytesIO):
         close()
+    else:
+        await anyio.to_thread.run_sync(close)
 
 
 async def download_artifact(request: Request) -> Response:
@@ -118,7 +145,6 @@ async def download_artifact(request: Request) -> Response:
         if not storage_path:
             return json_error("No file for this version", 404)
 
-        filename = _sanitize_header_filename(version_row.get("filename"), f"artifact-{artifact_id}")
         mime_type = version_row.get("mime_type") or "application/octet-stream"
         storage = await anyio.to_thread.run_sync(StorageCreator.get_storage)
 
@@ -157,6 +183,9 @@ async def download_artifact(request: Request) -> Response:
                 )
             return RedirectResponse(url, status_code=302)
 
+        # Build the headers before opening the file, so nothing between the
+        # open and the response can fail and leave the handle unclosed.
+        headers = {"Content-Disposition": _content_disposition(version_row.get("filename"), artifact_id)}
         # Stream the bytes in chunks instead of buffering the whole object in
         # memory (artifacts can be many MB); the handle closes when the
         # response ends, including on a client disconnect.
@@ -164,7 +193,7 @@ async def download_artifact(request: Request) -> Response:
         return ClosingStreamingResponse(
             _read_chunks(file_obj),
             media_type=mime_type,
-            headers={"Content-Disposition": f'attachment; filename="{filename}"'},
+            headers=headers,
             on_close=partial(_close_file, file_obj),
         )
     except FileNotFoundError: