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.

docs/runbooks/sse-notifications.md

@@ -158,14 +158,16 @@ null clear propagates into the next reconnect.**
 
 The per-user concurrent-connection cap (`SSE_MAX_CONCURRENT_PER_USER`,
 default 8) refused the connection. User has too many tabs open or a
-runaway reconnect loop. `redis-cli -n 2 GET user:<id>:sse_count`
-shows the live counter; the TTL is 1h from the last connection
-attempt (rolling — every INCR re-seeds it), so the key only ages
-out after the user stops reconnecting for a full hour.
+runaway reconnect loop. The cap covers `/api/events` and chat reconnects
+together.
 
-If the count is wedged high without explanation, the
-counter-DECR-in-finally path didn't run (worker SIGKILL, OOM). Wait
-for the TTL or `redis-cli -n 2 DEL user:<id>:sse_count` to reset.
+Each open stream holds a lease in a sorted set scored by its last
+refresh; `redis-cli -n 2 ZRANGE user:<id>:sse_leases 0 -1 WITHSCORES`
+lists them. A live stream refreshes its lease as it sends frames. A
+lease whose stream died without cleanup (worker SIGKILL, OOM) stops
+counting once it is older than the lease TTL: 60s, or 4 ×
+`SSE_KEEPALIVE_SECONDS` if that is larger. The cap heals on its own;
+`redis-cli -n 2 DEL user:<id>:sse_leases` clears it immediately.
 
 ### G. "Replay snapshot stops at 200 events"
 
@@ -305,8 +307,8 @@ redis-cli -n 2 PUBSUB NUMSUB $(redis-cli -n 2 PUBSUB CHANNELS 'user:*')
 # subscribers; OK after explaining to the user)
 redis-cli -n 2 XTRIM user:<id>:stream MAXLEN 0
 
-# Clear a wedged concurrent-connection counter
-redis-cli -n 2 DEL user:<id>:sse_count
+# Clear a user's SSE connection leases (they also age out on their own)
+redis-cli -n 2 DEL user:<id>:sse_leases
 
 # Force-flip every client to re-snapshot (drop the stream key entirely
 # — destroys the backlog; clients reconnect with their last id and

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:

docsgpt/events/keys.py

@@ -18,9 +18,9 @@ def topic_name(user_id: str) -> str:
     return f"user:{user_id}"
 
 
-def connection_counter_key(user_id: str) -> str:
-    """Redis counter tracking active SSE connections for ``user_id``."""
-    return f"user:{user_id}:sse_count"
+def connection_leases_key(user_id: str) -> str:
+    """Redis sorted set of live SSE connection leases for ``user_id``."""
+    return f"user:{user_id}:sse_leases"
 
 
 def replay_budget_key(user_id: str) -> str:

docsgpt/streaming/async_broadcast_channel.py

@@ -31,6 +31,11 @@ OnSubscribe = Callable[[], Union[None, Awaitable[None]]]
 # payload, whose PONG the client swallows before ``get_message`` returns.
 _LIVENESS_PING = "docsgpt-liveness"
 
+# Upper bound on each teardown call (unsubscribe, close). Teardown is shielded
+# from the cancellation that ended the stream, so without a bound a dead
+# connection could hold the stream's task and a graceful shutdown.
+_TEARDOWN_TIMEOUT_SECONDS = 5.0
+
 
 class AsyncTopic:
     """An async pub/sub channel identified by a string name."""
@@ -145,12 +150,11 @@ class AsyncTopic:
                 yield data if isinstance(data, bytes) else str(data).encode("utf-8")
         finally:
             # Client disconnect cancels this generator at the ``await
-            # get_message`` above; without shielding, the cancellation could
-            # re-fire mid-teardown and skip ``aclose()``, leaking the pooled
-            # connection back to nothing. Shield so unsubscribe + aclose
-            # always complete and the connection returns to the pool.
-            with anyio.CancelScope(shield=True):
-                if on_subscribe_fired:
+            # get_message`` above. Each teardown call is shielded so the
+            # cancellation can't skip it and strand the pooled connection, and
+            # bounded separately so a stuck unsubscribe still lets the close run.
+            if on_subscribe_fired:
+                with anyio.move_on_after(_TEARDOWN_TIMEOUT_SECONDS, shield=True):
                     try:
                         await pubsub.unsubscribe(self.name)
                     except Exception:
@@ -159,6 +163,7 @@ class AsyncTopic:
                             self.name,
                             exc_info=True,
                         )
+            with anyio.move_on_after(_TEARDOWN_TIMEOUT_SECONDS, shield=True):
                 try:
                     await pubsub.aclose()
                 except Exception:

docsgpt/streaming/sse_leases.py

@@ -0,0 +1,129 @@
+"""Per-connection leases behind the per-user SSE connection cap.
+
+Each open stream holds a member in the sorted set ``user:{id}:sse_leases``,
+scored by when it last refreshed. Dropping stale members and counting the rest
+enforces ``SSE_MAX_CONCURRENT_PER_USER``. A stream refreshes its lease while it
+runs and removes it when it ends; a lease whose stream died without cleanup
+(SIGKILL, OOM) ages out after the lease TTL. Unlike a shared counter with a
+TTL, an expiring lease can't take another stream's slot with it or drive the
+count negative.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+import uuid
+from contextlib import aclosing
+from dataclasses import dataclass
+from typing import Any, AsyncIterator, Optional, TypeVar
+
+import anyio
+
+from docsgpt.core.settings import settings
+from docsgpt.events.keys import connection_leases_key
+
+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")
+
+
+class StreamCapExceeded(Exception):
+    """Raised when the user already holds ``SSE_MAX_CONCURRENT_PER_USER`` streams."""
+
+
+def lease_ttl_seconds() -> float:
+    """Seconds a lease survives without a refresh.
+
+    A stream yields at least once per keepalive interval and refreshes on the
+    first yield after a third of the TTL, so four keepalives of headroom keep a
+    live stream's lease from expiring.
+    """
+    return max(60.0, 4.0 * float(settings.SSE_KEEPALIVE_SECONDS))
+
+
+@dataclass
+class StreamLease:
+    """One stream's claim on a slot in its user's connection cap."""
+
+    redis: Any
+    key: str
+    lease_id: str
+    ttl: float
+    refreshed_at: float
+
+    async def refresh_if_due(self) -> None:
+        """Re-stamp the lease once a third of its TTL has passed; failures are logged, not raised."""
+        now = time.monotonic()
+        if now - self.refreshed_at < self.ttl / 3:
+            return
+        self.refreshed_at = now
+        try:
+            with anyio.fail_after(_REDIS_OP_TIMEOUT_SECONDS):
+                async with self.redis.pipeline(transaction=True) as pipe:
+                    pipe.zadd(self.key, {self.lease_id: time.time()})
+                    pipe.expire(self.key, int(self.ttl))
+                    await pipe.execute()
+        except Exception:
+            logger.debug("SSE lease refresh failed for key=%s", self.key, exc_info=True)
+
+    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):
+            try:
+                await self.redis.zrem(self.key, self.lease_id)
+            except Exception:
+                logger.debug("SSE lease release failed for key=%s", self.key, exc_info=True)
+
+
+async def acquire_stream_lease(redis: Any, user_id: str) -> Optional[StreamLease]:
+    """Claim a connection slot for ``user_id``.
+
+    Args:
+        redis: The async Redis client, or ``None`` when unavailable.
+        user_id: The user opening the stream.
+
+    Returns:
+        The lease to refresh and release, or ``None`` when the cap is disabled
+        or Redis is unavailable or failing (fail-open).
+
+    Raises:
+        StreamCapExceeded: The user already holds the maximum number of streams.
+    """
+    cap = int(settings.SSE_MAX_CONCURRENT_PER_USER)
+    if cap <= 0 or redis is None:
+        return None
+    key = connection_leases_key(user_id)
+    ttl = lease_ttl_seconds()
+    lease_id = uuid.uuid4().hex
+    now = time.time()
+    try:
+        with anyio.fail_after(_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.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)
+        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()
+    return lease
+
+
+async def hold_lease(stream: AsyncIterator[T], lease: Optional[StreamLease]) -> AsyncIterator[T]:
+    """Yield from ``stream``, refreshing ``lease`` between frames; closing this closes ``stream``."""
+    async with aclosing(stream):
+        async for item in stream:
+            yield item
+            if lease is not None:
+                await lease.refresh_if_due()

tests/api/test_asgi_stream.py

@@ -0,0 +1,92 @@
+"""Tests for ``docsgpt/api/asgi_stream.py`` — cleanup that can't stall a worker.
+
+The cleanup runs shielded from the cancellation that ended the response, so
+each step needs its own deadline: a close or release waiting on a dead Redis
+connection must not hold the request task, or a graceful shutdown, forever.
+"""
+
+from __future__ import annotations
+
+import anyio
+import pytest
+
+from docsgpt.api import asgi_stream
+from docsgpt.api.asgi_stream import ClosingStreamingResponse
+from tests.asgi_stream import stream_then_disconnect
+
+
+class _StuckCloseIterator:
+    """Yields one chunk, then ends; its ``aclose`` never returns."""
+
+    def __init__(self) -> None:
+        self._sent = False
+        self.aclose_started = False
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if self._sent:
+            raise StopAsyncIteration
+        self._sent = True
+        return b"chunk"
+
+    async def aclose(self):
+        self.aclose_started = True
+        await anyio.sleep_forever()
+
+
+async def _one_chunk():
+    yield b"chunk"
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+class TestBoundedCleanup:
+    async def test_stuck_stream_close_is_abandoned_and_on_close_still_runs(self, monkeypatch):
+        monkeypatch.setattr(asgi_stream, "_CLEANUP_TIMEOUT_SECONDS", 0.05)
+        body = _StuckCloseIterator()
+        released = {"ran": False}
+
+        async def on_close():
+            released["ran"] = True
+
+        response = ClosingStreamingResponse(body, media_type="text/plain", on_close=on_close)
+        with anyio.fail_after(2):
+            status, _, chunks = await stream_then_disconnect(
+                response, "/", chunks_before_disconnect=10**6
+            )
+        assert status == 200
+        assert chunks == [b"chunk"]
+        assert body.aclose_started
+        assert released["ran"] is True
+
+    async def test_stuck_on_close_is_abandoned(self, monkeypatch):
+        monkeypatch.setattr(asgi_stream, "_CLEANUP_TIMEOUT_SECONDS", 0.05)
+
+        async def on_close():
+            await anyio.sleep_forever()
+
+        response = ClosingStreamingResponse(_one_chunk(), media_type="text/plain", on_close=on_close)
+        with anyio.fail_after(2):
+            status, _, chunks = await stream_then_disconnect(
+                response, "/", chunks_before_disconnect=10**6
+            )
+        assert status == 200
+        assert chunks == [b"chunk"]
+
+    async def test_on_close_runs_when_client_leaves_before_first_frame(self):
+        released = {"count": 0}
+
+        async def on_close():
+            released["count"] += 1
+
+        async def endless():
+            while True:
+                await anyio.sleep(0.01)
+                yield b"tick"
+
+        response = ClosingStreamingResponse(endless(), media_type="text/plain", on_close=on_close)
+        with anyio.fail_after(2):
+            await stream_then_disconnect(response, "/", chunks_before_disconnect=0)
+        assert released["count"] == 1

tests/api/test_async_sse_routes.py

@@ -1,26 +1,29 @@
 """Tests for ``docsgpt/api/async_sse.py``.
 
 Native-async reconnect endpoint: GET /api/messages/<id>/events. Auth gate,
-ownership gate, malformed-id rejection, Last-Event-ID normalisation, and the
-SSE response shape (headers + ``: connected`` prelude). The route is a
-Starlette endpoint, so it's driven through Starlette's TestClient over a
-minimal app built from ``async_sse_routes``.
+ownership gate, malformed-id rejection, Last-Event-ID normalisation, the
+per-user connection lease, and the SSE response shape (headers + ``: connected``
+prelude). The route is a Starlette endpoint, so it's driven through Starlette's
+TestClient over a minimal app built from ``async_sse_routes``.
 """
 
 from __future__ import annotations
 
+import time
 from unittest.mock import AsyncMock, patch
 
 import pytest
 from starlette.applications import Starlette
 from starlette.testclient import TestClient
 
+from docsgpt.api import async_sse
 from docsgpt.api.async_sse import (
     _MESSAGE_ID_RE,
     _normalise_last_event_id,
     async_sse_routes,
 )
-from docsgpt.core.settings import settings
+from tests.asgi_stream import stream_then_disconnect
+from tests.fake_async_redis import FakeAsyncRedis
 
 VALID_UUID = "67d65e8f-e7fb-4df1-9e6e-99ea6c830206"
 
@@ -28,42 +31,54 @@ _AUTH = "docsgpt.api.asgi_auth.handle_auth"
 _OWNS = "docsgpt.api.async_sse._user_owns_message"
 _STREAM = "docsgpt.api.async_sse.build_message_event_stream_async"
 _AREDIS = "docsgpt.api.async_sse.get_async_redis_instance"
+_LEASES = "user:alice:sse_leases"
+
+
+def _app() -> Starlette:
+    return Starlette(routes=async_sse_routes)
 
 
 def _client() -> TestClient:
-    return TestClient(Starlette(routes=async_sse_routes))
+    return TestClient(_app())
 
 
 @pytest.fixture(autouse=True)
-def _no_redis_by_default():
+def _no_redis_by_default(monkeypatch):
     """Disable the per-user cap (no Redis) so gate tests stay hermetic.
 
-    Cap-specific tests override this with their own mock Redis.
+    Cap-specific tests override this with their own fake Redis.
     """
+    monkeypatch.setattr(async_sse.settings, "SSE_MAX_CONCURRENT_PER_USER", 8)
+    monkeypatch.setattr(async_sse.settings, "SSE_KEEPALIVE_SECONDS", 15)
     with patch(_AREDIS, AsyncMock(return_value=None)):
         yield
 
 
-def _mock_redis(incr_value: int) -> AsyncMock:
-    redis = AsyncMock()
-    redis.incr = AsyncMock(return_value=incr_value)
-    redis.expire = AsyncMock(return_value=True)
-    redis.decr = AsyncMock(return_value=incr_value - 1)
-    return redis
-
-
-def _fake_stream(record: dict | None = None):
-    """Async builder stub: records the cursor, yields the prelude, returns."""
+def _fake_stream(record: dict | None = None, redis: FakeAsyncRedis | None = None):
+    """Async builder stub: records the cursor (and live leases), yields the prelude, returns."""
 
     async def _gen(message_id, last_event_id=None, **kwargs):
         if record is not None:
             record["message_id"] = message_id
             record["last_event_id"] = last_event_id
+            if redis is not None:
+                record["leases_during"] = len(redis.zsets.get(_LEASES, {}))
         yield ": connected\n\n"
 
     return _gen
 
 
+def _endless_stream():
+    async def _gen(message_id, last_event_id=None, **kwargs):
+        import anyio
+
+        while True:
+            await anyio.sleep(0.01)
+            yield ": keepalive\n\n"
+
+    return _gen
+
+
 # ── pure helpers ─────────────────────────────────────────────────────────
 
 
@@ -206,17 +221,14 @@ class TestLastEventIdParsing:
         assert captured["last_event_id"] is None
 
 
-# ── per-user concurrent-connection cap ──────────────────────────────────────
+# ── per-user connection lease ───────────────────────────────────────────────
 
 
 @pytest.mark.unit
 class TestConnectionCap:
-    def _cap(self) -> int:
-        return int(settings.SSE_MAX_CONCURRENT_PER_USER) or 8
-
-    def test_429_when_over_cap(self):
-        cap = self._cap()
-        redis = _mock_redis(incr_value=cap + 1)  # post-incr count exceeds cap
+    def test_429_when_user_already_holds_the_cap(self):
+        redis = FakeAsyncRedis()
+        redis.zsets[_LEASES] = {f"tab-{i}": time.time() for i in range(8)}
         with patch(_AUTH, return_value={"sub": "alice"}), patch(
             _OWNS, return_value=True
         ), patch(_STREAM, _fake_stream()), patch(
@@ -224,23 +236,22 @@ class TestConnectionCap:
         ):
             r = _client().get(f"/api/messages/{VALID_UUID}/events")
         assert r.status_code == 429
-        # The increment is rolled back so a rejected attempt doesn't wedge
-        # the counter at the cap forever.
-        redis.decr.assert_awaited_once()
+        # The rejected attempt removed its own lease and nobody else's.
+        assert set(redis.zsets[_LEASES]) == {f"tab-{i}" for i in range(8)}
 
-    def test_200_and_slot_released_when_under_cap(self):
-        redis = _mock_redis(incr_value=1)
+    def test_200_lease_held_while_streaming_then_released(self):
+        redis = FakeAsyncRedis()
+        captured: dict = {}
         with patch(_AUTH, return_value={"sub": "alice"}), patch(
             _OWNS, return_value=True
-        ), patch(_STREAM, _fake_stream()), patch(
+        ), patch(_STREAM, _fake_stream(captured, redis)), patch(
             _AREDIS, AsyncMock(return_value=redis)
         ):
             r = _client().get(f"/api/messages/{VALID_UUID}/events")
         assert r.status_code == 200
         assert ": connected" in r.text
-        redis.incr.assert_awaited_once()
-        # Slot released when the stream finishes (terminal/close).
-        redis.decr.assert_awaited_once()
+        assert captured["leases_during"] == 1
+        assert redis.zsets == {}
 
     def test_cap_skipped_when_redis_unavailable(self):
         # The autouse fixture already makes get_async_redis_instance -> None;
@@ -250,3 +261,18 @@ class TestConnectionCap:
         ), patch(_STREAM, _fake_stream()):
             r = _client().get(f"/api/messages/{VALID_UUID}/events")
         assert r.status_code == 200
+
+    @pytest.mark.asyncio
+    async def test_disconnect_releases_lease(self):
+        redis = FakeAsyncRedis()
+        with patch(_AUTH, return_value={"sub": "alice"}), patch(
+            _OWNS, return_value=True
+        ), patch(_STREAM, _endless_stream()), patch(
+            _AREDIS, AsyncMock(return_value=redis)
+        ):
+            status, _, chunks = await stream_then_disconnect(
+                _app(), f"/api/messages/{VALID_UUID}/events", chunks_before_disconnect=1
+            )
+        assert status == 200
+        assert chunks[0] == b": keepalive\n\n"
+        assert redis.zsets == {}

tests/api/test_events_routes.py

@@ -1,17 +1,20 @@
 """Tests for docsgpt/api/events/routes.py — the native-async ``GET /api/events``.
 
 The route is a Starlette endpoint served on the event loop. Tests drive it
-through Starlette's ``TestClient`` with the async Redis client and
-``AsyncTopic.subscribe`` faked. Fake subscriptions end after a few frames so
-the stream finishes and the client returns the body; a mid-stream client
-disconnect and the shutdown path go through ``tests/asgi_stream.py``.
+through Starlette's ``TestClient`` with the async Redis client faked
+(``tests/fake_async_redis.py``: real sorted sets for connection leases, mocks
+for streams) and ``AsyncTopic.subscribe`` replaced. Fake subscriptions end
+after a few frames so the stream finishes and the client returns the body; a
+mid-stream client disconnect and the shutdown path go through
+``tests/asgi_stream.py``.
 """
 
 from __future__ import annotations
 
 import inspect
 import json
-from typing import Any
+import time
+from typing import Any, Callable, Optional
 from unittest.mock import AsyncMock, patch
 
 import anyio
@@ -23,11 +26,12 @@ from docsgpt.api.events import routes as events_module
 from docsgpt.cache import PUBSUB_SOCKET_TIMEOUT_SECONDS
 from docsgpt.core.shutdown import begin_shutdown, reset_shutdown
 from tests.asgi_stream import stream_then_disconnect
+from tests.fake_async_redis import FakeAsyncRedis
 
 _AUTH = "docsgpt.api.asgi_auth.handle_auth"
 _AREDIS = "docsgpt.api.events.routes.get_async_redis_instance"
 _SUBSCRIBE = "docsgpt.api.events.routes.AsyncTopic.subscribe"
-_COUNTER = "user:alice:sse_count"
+_LEASES = "user:alice:sse_leases"
 ALICE = {"sub": "alice"}
 OLD_CURSOR = "1735682300000-0"
 
@@ -40,22 +44,23 @@ def _get(headers: dict | None = None, params: dict | None = None):
     return TestClient(_app()).get("/api/events", headers=headers or {}, params=params or {})
 
 
-def _redis(incr: int = 1) -> AsyncMock:
-    redis = AsyncMock()
-    redis.incr = AsyncMock(return_value=incr)
-    redis.expire = AsyncMock(return_value=True)
-    redis.decr = AsyncMock(return_value=0)
-    redis.xinfo_stream = AsyncMock(side_effect=Exception("no such key"))
-    redis.xrange = AsyncMock(return_value=[])
-    return redis
+def _redis() -> FakeAsyncRedis:
+    return FakeAsyncRedis()
 
 
-def _subscribe(*payloads: Any, fire_callback: bool = True, record: dict | None = None):
+def _subscribe(
+    *payloads: Any,
+    fire_callback: bool = True,
+    record: dict | None = None,
+    probe: Optional[Callable[[], None]] = None,
+):
     """``AsyncTopic.subscribe`` stand-in: ack, one idle tick, the payloads, then end."""
 
     async def _impl(self, on_subscribe=None, poll_timeout=1.0, **kwargs):
         if record is not None:
             record.update(kwargs, poll_timeout=poll_timeout, topic=self.name)
+        if probe is not None:
+            probe()
         if fire_callback and on_subscribe is not None:
             result = on_subscribe()
             if inspect.isawaitable(result):
@@ -180,6 +185,7 @@ class TestStreamShape:
         assert r.status_code == 200
         assert r.text == ": connected\n\n: push_disabled\n\n"
         redis.incr.assert_not_awaited()
+        assert redis.zsets == {}
 
     def test_subscribes_to_user_topic_with_liveness_probe(self):
         record: dict = {}
@@ -200,22 +206,34 @@ class TestStreamShape:
         assert ": keepalive\n\n" in r.text
 
 
-# ── concurrency cap ─────────────────────────────────────────────────────
+# ── concurrency cap (per-connection leases) ─────────────────────────────
 
 
 @pytest.mark.unit
 class TestConcurrencyCap:
-    def test_429_when_over_cap_and_increment_rolled_back(self):
-        redis = _redis(incr=9)
+    def test_429_when_user_already_holds_the_cap(self):
+        redis = _redis()
+        redis.zsets[_LEASES] = {f"tab-{i}": time.time() for i in range(8)}
         with patch(_AUTH, return_value=ALICE), patch(
             _AREDIS, AsyncMock(return_value=redis)
         ), patch(_SUBSCRIBE, _subscribe()):
             r = _get()
         assert r.status_code == 429
         assert r.json() == {"success": False, "message": "Too many concurrent SSE connections"}
-        redis.decr.assert_awaited_once_with(_COUNTER)
+        # The rejected attempt removed its own lease and nobody else's.
+        assert set(redis.zsets[_LEASES]) == {f"tab-{i}" for i in range(8)}
 
-    def test_cap_disabled_leaves_counter_alone(self, monkeypatch):
+    def test_leases_left_by_dead_streams_age_out(self):
+        redis = _redis()
+        redis.zsets[_LEASES] = {f"crashed-{i}": time.time() - 3600 for i in range(8)}
+        with patch(_AUTH, return_value=ALICE), patch(
+            _AREDIS, AsyncMock(return_value=redis)
+        ), patch(_SUBSCRIBE, _subscribe()):
+            r = _get()
+        assert r.status_code == 200
+        assert redis.zsets == {}
+
+    def test_cap_disabled_takes_no_lease(self, monkeypatch):
         monkeypatch.setattr(events_module.settings, "SSE_MAX_CONCURRENT_PER_USER", 0)
         redis = _redis()
         with patch(_AUTH, return_value=ALICE), patch(
@@ -223,38 +241,31 @@ class TestConcurrencyCap:
         ), patch(_SUBSCRIBE, _subscribe()):
             r = _get()
         assert r.status_code == 200
-        assert _COUNTER not in [call.args[0] for call in redis.incr.await_args_list]
-        redis.decr.assert_not_awaited()
+        assert redis.zsets == {}
 
-    def test_expire_failure_does_not_bypass_cap(self):
-        redis = _redis(incr=9)
-        redis.expire = AsyncMock(side_effect=Exception("expire blip"))
-        with patch(_AUTH, return_value=ALICE), patch(
-            _AREDIS, AsyncMock(return_value=redis)
-        ), patch(_SUBSCRIBE, _subscribe()):
-            r = _get()
-        assert r.status_code == 429
-
-    def test_incr_failure_fails_open(self):
+    def test_redis_errors_fail_open(self):
         redis = _redis()
-        redis.incr = AsyncMock(side_effect=Exception("redis down"))
+        redis.fail = ConnectionError("reset by peer")
         with patch(_AUTH, return_value=ALICE), patch(
             _AREDIS, AsyncMock(return_value=redis)
         ), patch(_SUBSCRIBE, _subscribe()):
             r = _get()
         assert r.status_code == 200
-        # Nothing was counted, so nothing is released.
-        redis.decr.assert_not_awaited()
 
-    def test_slot_released_when_stream_ends(self):
-        redis = _redis(incr=1)
+    def test_lease_held_while_streaming_and_released_at_end(self):
+        redis = _redis()
+        seen: dict = {}
+
+        def _count_leases():
+            seen["during"] = len(redis.zsets.get(_LEASES, {}))
+
         with patch(_AUTH, return_value=ALICE), patch(
             _AREDIS, AsyncMock(return_value=redis)
-        ), patch(_SUBSCRIBE, _subscribe()):
+        ), patch(_SUBSCRIBE, _subscribe(probe=_count_leases)):
             r = _get()
         assert r.status_code == 200
-        redis.incr.assert_awaited_once_with(_COUNTER)
-        redis.decr.assert_awaited_once_with(_COUNTER)
+        assert seen["during"] == 1
+        assert redis.zsets == {}
 
     def test_redis_unavailable_serves_stream_without_cap(self):
         with patch(_AUTH, return_value=ALICE):
@@ -269,9 +280,9 @@ class TestConcurrencyCap:
 @pytest.mark.unit
 @pytest.mark.asyncio
 class TestStreamLifecycle:
-    async def test_client_disconnect_releases_slot_and_closes_subscription(self, monkeypatch):
+    async def test_client_disconnect_releases_lease_and_closes_subscription(self, monkeypatch):
         monkeypatch.setattr(events_module.settings, "SSE_KEEPALIVE_SECONDS", 0)
-        redis = _redis(incr=1)
+        redis = _redis()
         state: dict = {}
         with patch(_AUTH, return_value=ALICE), patch(
             _AREDIS, AsyncMock(return_value=redis)
@@ -284,12 +295,12 @@ class TestStreamLifecycle:
         assert status == 200
         assert chunks[0] == b": connected\n\n"
         assert state.get("closed") is True
-        redis.decr.assert_awaited_once_with(_COUNTER)
+        assert redis.zsets == {}
 
-    async def test_disconnect_before_first_frame_still_releases_slot_once(self):
+    async def test_disconnect_before_first_frame_still_releases_lease(self):
         # The tab closes as the response starts: the body generator may never
         # run, so its own cleanup can't be what frees the slot.
-        redis = _redis(incr=1)
+        redis = _redis()
         state: dict = {}
         with patch(_AUTH, return_value=ALICE), patch(
             _AREDIS, AsyncMock(return_value=redis)
@@ -298,10 +309,10 @@ class TestStreamLifecycle:
                 _app(), "/api/events", chunks_before_disconnect=0
             )
         assert status == 200
-        redis.decr.assert_awaited_once_with(_COUNTER)
+        assert redis.zsets == {}
 
-    async def test_shutdown_ends_stream_and_releases_slot(self):
-        redis = _redis(incr=1)
+    async def test_shutdown_ends_stream_and_releases_lease(self):
+        redis = _redis()
         state: dict = {}
         begin_shutdown()
         with patch(_AUTH, return_value=ALICE), patch(
@@ -313,7 +324,7 @@ class TestStreamLifecycle:
         assert status == 200
         assert b"".join(chunks) == b": connected\n\n"
         assert state.get("closed") is True
-        redis.decr.assert_awaited_once_with(_COUNTER)
+        assert redis.zsets == {}
 
 
 # ── replay + live tail ──────────────────────────────────────────────────
@@ -457,11 +468,22 @@ class TestReplayBudget:
         assert redis.expire.await_count == 4
         assert all(call.args == ("user:alice:replay_count", 60) for call in redis.expire.await_args_list)
 
-    async def test_fail_open_on_redis_error(self, monkeypatch):
+    async def test_fail_open_on_redis_error(self):
         redis = _redis()
         redis.incr = AsyncMock(side_effect=Exception("redis down"))
         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)
+        redis = _redis()
+
+        async def _stalled(_key):
+            await anyio.sleep_forever()
+
+        redis.incr = AsyncMock(side_effect=_stalled)
+        with anyio.fail_after(2):
+            assert await events_module._allow_replay(redis, "alice", OLD_CURSOR) is True
+
     async def test_recovers_when_seeding_expire_raises(self):
         redis = _redis()
         counter = {"v": 0}
@@ -522,19 +544,15 @@ class TestReplayBudgetRoute:
         live frames would advance the client's cursor past the un-replayed
         window. The 429 keeps the cursor pinned for the next attempt."""
         redis = _redis()
-
-        def _incr(key):
-            return 1 if key == _COUNTER else 31
-
-        redis.incr = AsyncMock(side_effect=_incr)
+        redis.incr = AsyncMock(return_value=31)
         with patch(_AUTH, return_value=ALICE), patch(
             _AREDIS, AsyncMock(return_value=redis)
         ), patch(_SUBSCRIBE, _subscribe()):
             r = _get(headers={"Last-Event-ID": OLD_CURSOR})
         assert r.status_code == 429
         assert r.json() == {"success": False, "message": "Replay budget exhausted"}
-        # The concurrency slot is released so a denied request doesn't hold it.
-        redis.decr.assert_awaited_once_with(_COUNTER)
+        # The connection lease is released so a denied request doesn't hold a slot.
+        assert redis.zsets == {}
 
 
 # ── format helpers ──────────────────────────────────────────────────────

tests/api/user/test_artifacts_routes.py

@@ -738,6 +738,58 @@ class TestDownloadArtifact:
         assert disposition == 'attachment; filename="aInjected: x.txt"'
         assert "injected" not in resp.headers
 
+    def test_unicode_filename_served_with_rfc5987_header(self, _patch_db, token_owner, monkeypatch):
+        # Latin-1 header encoding used to reject this and return 400 after
+        # opening the file.
+        art = _seed_download(_patch_db, filename="报告.pdf", mime_type="application/pdf")
+        storage = _use_storage(monkeypatch, _FakeStorage(b"PDF"))
+        resp = _download(art["id"], token=token_owner)
+        assert resp.status_code == 200
+        assert resp.content == b"PDF"
+        assert resp.headers["content-disposition"] == (
+            f'attachment; filename="artifact-{art["id"]}.pdf"; '
+            "filename*=UTF-8''%E6%8A%A5%E5%91%8A.pdf"
+        )
+        assert storage.opened[0].closed
+
+    def test_accented_filename_keeps_readable_ascii_fallback(self, _patch_db, token_owner, monkeypatch):
+        art = _seed_download(_patch_db, filename="résumé final.pdf")
+        _use_storage(monkeypatch, _FakeStorage(b"CV"))
+        resp = _download(art["id"], token=token_owner)
+        assert resp.status_code == 200
+        assert resp.headers["content-disposition"] == (
+            'attachment; filename="resume final.pdf"; '
+            "filename*=UTF-8''r%C3%A9sum%C3%A9%20final.pdf"
+        )
+
+    def test_disk_file_closed_off_the_event_loop(self, _patch_db, token_owner, monkeypatch):
+        import threading
+
+        class _TrackedHandle:
+            """A non-BytesIO handle, like LocalStorage's open file."""
+
+            def __init__(self, data: bytes):
+                self._buf = io.BytesIO(data)
+                self.close_thread = None
+
+            def read(self, size=-1):
+                return self._buf.read(size)
+
+            def close(self):
+                self.close_thread = threading.current_thread().name
+
+        handle = _TrackedHandle(b"DISK")
+
+        class _HandleStorage(_FakeStorage):
+            def get_file(self, path):
+                return handle
+
+        art = _seed_download(_patch_db)
+        _use_storage(monkeypatch, _HandleStorage())
+        resp = _download(art["id"], token=token_owner)
+        assert resp.content == b"DISK"
+        assert handle.close_thread == "AnyIO worker thread"
+
     def test_non_uuid_id_404(self, _patch_db, token_owner):
         resp = _download("legacy-mongo-objectid-aabbcc", token=token_owner)
         assert resp.status_code == 404

tests/fake_async_redis.py

@@ -0,0 +1,149 @@
+"""In-memory ``redis.asyncio`` stand-in for the SSE lease and event-stream tests.
+
+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).
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+from unittest.mock import AsyncMock
+
+import anyio
+
+
+def _bound(value) -> float:
+    if isinstance(value, bytes):
+        value = value.decode()
+    if value == "-inf":
+        return float("-inf")
+    if value in ("+inf", "inf"):
+        return float("inf")
+    return float(value)
+
+
+def _member(value) -> str:
+    return value.decode() if isinstance(value, bytes) else str(value)
+
+
+class FakeAsyncRedis:
+    """Sorted sets + pipelines in memory; everything else is an ``AsyncMock``."""
+
+    def __init__(self) -> None:
+        self.zsets: dict[str, dict[str, float]] = {}
+        self.key_ttls: dict[str, int] = {}
+        self.fail: Optional[BaseException] = None
+        self.hang = False
+        self.incr = AsyncMock(return_value=1)
+        self.expire = AsyncMock(return_value=True)
+        self.xrange = AsyncMock(return_value=[])
+        self.xinfo_stream = AsyncMock(side_effect=Exception("no such key"))
+
+    async def _gate(self) -> None:
+        if self.hang:
+            await anyio.sleep_forever()
+        if self.fail is not None:
+            raise self.fail
+
+    # -- synchronous cores, shared by direct calls and pipelines -------
+    def _zadd(self, key, mapping, xx=False):
+        zset = self.zsets.setdefault(key, {})
+        added = 0
+        for member, score in mapping.items():
+            member = _member(member)
+            if xx and member not in zset:
+                continue
+            added += member not in zset
+            zset[member] = float(score)
+        if not zset:
+            self.zsets.pop(key, None)
+        return added
+
+    def _zrem(self, key, *members):
+        zset = self.zsets.get(key, {})
+        removed = 0
+        for member in members:
+            if zset.pop(_member(member), None) is not None:
+                removed += 1
+        if key in self.zsets and not zset:
+            self.zsets.pop(key)
+        return removed
+
+    def _zcard(self, key):
+        return len(self.zsets.get(key, {}))
+
+    def _zremrangebyscore(self, key, min, max):
+        lo, hi = _bound(min), _bound(max)
+        zset = self.zsets.get(key, {})
+        stale = [m for m, score in zset.items() if lo <= score <= hi]
+        for member in stale:
+            zset.pop(member)
+        if key in self.zsets and not zset:
+            self.zsets.pop(key)
+        return len(stale)
+
+    def _set_ttl(self, key, seconds):
+        self.key_ttls[key] = int(seconds)
+        return True
+
+    # -- awaitable commands --------------------------------------------
+    async def zadd(self, key, mapping, xx=False):
+        await self._gate()
+        return self._zadd(key, mapping, xx=xx)
+
+    async def zrem(self, key, *members):
+        await self._gate()
+        return self._zrem(key, *members)
+
+    async def zcard(self, key):
+        await self._gate()
+        return self._zcard(key)
+
+    async def zremrangebyscore(self, key, min, max):
+        await self._gate()
+        return self._zremrangebyscore(key, min, max)
+
+    def pipeline(self, transaction=True):
+        return FakePipeline(self)
+
+
+class FakePipeline:
+    """Queues commands and applies them together on ``execute``."""
+
+    def __init__(self, redis: FakeAsyncRedis) -> None:
+        self._redis = redis
+        self._ops: list = []
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *exc):
+        self._ops.clear()
+
+    def _queue(self, fn, *args, **kwargs):
+        self._ops.append((fn, args, kwargs))
+        return self
+
+    def zadd(self, key, mapping, xx=False):
+        return self._queue(self._redis._zadd, key, mapping, xx=xx)
+
+    def zrem(self, key, *members):
+        return self._queue(self._redis._zrem, key, *members)
+
+    def zcard(self, key):
+        return self._queue(self._redis._zcard, key)
+
+    def zremrangebyscore(self, key, min, max):
+        return self._queue(self._redis._zremrangebyscore, key, min, max)
+
+    def expire(self, key, seconds):
+        return self._queue(self._redis._set_ttl, key, seconds)
+
+    async def execute(self):
+        await self._redis._gate()
+        results = [fn(*args, **kwargs) for fn, args, kwargs in self._ops]
+        self._ops.clear()
+        return results

tests/test_async_broadcast_channel.py

@@ -155,6 +155,25 @@ class TestLivenessProbe:
         assert "ASYNC_REDIS_MAX_CONNECTIONS" in str(log.warning.call_args)
         log.exception.assert_not_called()
 
+    async def test_stuck_unsubscribe_does_not_block_teardown(self, monkeypatch):
+        # Teardown is shielded, so each Redis call in it needs its own deadline.
+        monkeypatch.setattr(
+            "docsgpt.streaming.async_broadcast_channel._TEARDOWN_TIMEOUT_SECONDS", 0.05
+        )
+        pubsub = _FakePubSub()
+
+        async def _stuck_unsubscribe(name):
+            await anyio.sleep_forever()
+
+        pubsub.unsubscribe = _stuck_unsubscribe
+        with patch(_AREDIS, _redis_with(pubsub)):
+            agen = AsyncTopic("t").subscribe(poll_timeout=0.01)
+            assert await agen.__anext__() is None
+            with anyio.fail_after(2):
+                await agen.aclose()
+        # The pooled connection is still closed after the stuck unsubscribe.
+        assert pubsub.closed
+
     async def test_no_probe_when_liveness_disabled(self):
         pubsub = _FakePubSub(answer_pings=False)
         with patch(_AREDIS, _redis_with(pubsub)):

tests/test_events_substrate.py

@@ -19,7 +19,7 @@ from docsgpt.api.events.routes import (
     _normalize_last_event_id,
 )
 from docsgpt.events.keys import (
-    connection_counter_key,
+    connection_leases_key,
     stream_id_compare,
     stream_key,
     topic_name,
@@ -39,8 +39,8 @@ class TestKeys:
     def test_topic_name(self):
         assert topic_name("alice") == "user:alice"
 
-    def test_connection_counter_key(self):
-        assert connection_counter_key("alice") == "user:alice:sse_count"
+    def test_connection_leases_key(self):
+        assert connection_leases_key("alice") == "user:alice:sse_leases"
 
 
 # ── stream_id_compare ───────────────────────────────────────────────────

tests/test_sse_leases.py

@@ -0,0 +1,186 @@
+"""Tests for ``docsgpt/streaming/sse_leases.py`` — the per-user SSE connection cap.
+
+Each open stream holds its own lease in a sorted set scored by its last
+refresh. A shared INCR/DECR counter with a TTL could expire under a stream
+that outlived it, after which that stream's DECR took another stream's slot
+or drove the count negative and the cap stopped holding.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import anyio
+import pytest
+
+from docsgpt.streaming import sse_leases
+from docsgpt.streaming.sse_leases import (
+    StreamCapExceeded,
+    acquire_stream_lease,
+    hold_lease,
+)
+from tests.fake_async_redis import FakeAsyncRedis
+
+KEY = "user:alice:sse_leases"
+
+
+@pytest.fixture(autouse=True)
+def _settings(monkeypatch):
+    monkeypatch.setattr(sse_leases.settings, "SSE_MAX_CONCURRENT_PER_USER", 8)
+    monkeypatch.setattr(sse_leases.settings, "SSE_KEEPALIVE_SECONDS", 15)
+
+
+@pytest.fixture
+def clock(monkeypatch):
+    """Drive the module's wall and monotonic clocks by hand."""
+    now = {"t": 1_000_000.0}
+    fake_time = SimpleNamespace(time=lambda: now["t"], monotonic=lambda: now["t"])
+    monkeypatch.setattr(sse_leases, "time", fake_time)
+    return now
+
+
+def _fill(redis: FakeAsyncRedis, count: int, score: float) -> None:
+    redis.zsets[KEY] = {f"other-{i}": score for i in range(count)}
+
+
+@pytest.mark.unit
+def test_lease_ttl_outlasts_keepalive_gaps(monkeypatch):
+    assert sse_leases.lease_ttl_seconds() == 60
+    monkeypatch.setattr(sse_leases.settings, "SSE_KEEPALIVE_SECONDS", 30)
+    # A stream yields at least every keepalive, so refreshing a third of the way
+    # through the TTL always lands before expiry.
+    assert sse_leases.lease_ttl_seconds() == 120
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+class TestAcquire:
+    async def test_adds_own_lease_and_key_ttl(self, clock):
+        redis = FakeAsyncRedis()
+        lease = await acquire_stream_lease(redis, "alice")
+        assert redis.zsets[KEY] == {lease.lease_id: clock["t"]}
+        assert redis.key_ttls[KEY] == 60
+
+    async def test_over_cap_raises_and_leaves_other_leases(self, clock):
+        redis = FakeAsyncRedis()
+        _fill(redis, 8, clock["t"])
+        with pytest.raises(StreamCapExceeded):
+            await acquire_stream_lease(redis, "alice")
+        assert set(redis.zsets[KEY]) == {f"other-{i}" for i in range(8)}
+
+    async def test_stale_leases_from_dead_streams_do_not_count(self, clock):
+        redis = FakeAsyncRedis()
+        _fill(redis, 8, clock["t"] - 61)
+        lease = await acquire_stream_lease(redis, "alice")
+        assert set(redis.zsets[KEY]) == {lease.lease_id}
+
+    async def test_cap_disabled_takes_no_lease(self, monkeypatch):
+        monkeypatch.setattr(sse_leases.settings, "SSE_MAX_CONCURRENT_PER_USER", 0)
+        redis = FakeAsyncRedis()
+        assert await acquire_stream_lease(redis, "alice") is None
+        assert redis.zsets == {}
+
+    async def test_no_redis_takes_no_lease(self):
+        assert await acquire_stream_lease(None, "alice") is None
+
+    async def test_redis_error_fails_open(self):
+        redis = FakeAsyncRedis()
+        redis.fail = ConnectionError("reset by peer")
+        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)
+        redis = FakeAsyncRedis()
+        redis.hang = True
+        with anyio.fail_after(2):
+            assert await acquire_stream_lease(redis, "alice") is None
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+class TestLeaseLifecycle:
+    async def test_release_removes_only_its_own_lease(self, clock):
+        redis = FakeAsyncRedis()
+        first = await acquire_stream_lease(redis, "alice")
+        second = await acquire_stream_lease(redis, "alice")
+        await first.release()
+        assert set(redis.zsets[KEY]) == {second.lease_id}
+        # Releasing twice can't touch anyone else's slot either.
+        await first.release()
+        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)
+        redis = FakeAsyncRedis()
+        lease = await acquire_stream_lease(redis, "alice")
+        redis.hang = True
+        with anyio.fail_after(2):
+            await lease.release()
+
+    async def test_refresh_waits_a_third_of_the_ttl(self, clock):
+        redis = FakeAsyncRedis()
+        lease = await acquire_stream_lease(redis, "alice")
+        start = clock["t"]
+        clock["t"] += 19
+        await lease.refresh_if_due()
+        assert redis.zsets[KEY][lease.lease_id] == start
+        clock["t"] += 1
+        await lease.refresh_if_due()
+        assert redis.zsets[KEY][lease.lease_id] == clock["t"]
+
+    async def test_long_lived_stream_keeps_its_slot(self, clock):
+        # The scenario a shared counter got wrong: a stream open far longer
+        # than the TTL still counts, and newcomers can't exceed the cap.
+        redis = FakeAsyncRedis()
+        long_lived = await acquire_stream_lease(redis, "alice")
+        for _ in range(30):  # 30 x 20s = 10 minutes, refreshing as it goes
+            clock["t"] += 20
+            await long_lived.refresh_if_due()
+        _fill_others = {f"other-{i}": clock["t"] for i in range(7)}
+        redis.zsets[KEY].update(_fill_others)
+        with pytest.raises(StreamCapExceeded):
+            await acquire_stream_lease(redis, "alice")
+        await long_lived.release()
+        assert long_lived.lease_id not in redis.zsets[KEY]
+        assert len(redis.zsets[KEY]) == 7
+
+    async def test_refresh_failure_is_swallowed(self, clock):
+        redis = FakeAsyncRedis()
+        lease = await acquire_stream_lease(redis, "alice")
+        redis.fail = ConnectionError("gone")
+        clock["t"] += 30
+        await lease.refresh_if_due()  # must not raise into the stream
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+class TestHoldLease:
+    async def test_refreshes_between_frames_and_closes_inner_stream(self, clock):
+        redis = FakeAsyncRedis()
+        lease = await acquire_stream_lease(redis, "alice")
+        state = {"closed": False}
+
+        async def inner():
+            try:
+                for i in range(5):
+                    clock["t"] += 25
+                    yield f"frame-{i}"
+            finally:
+                state["closed"] = True
+
+        wrapped = hold_lease(inner(), lease)
+        assert await wrapped.__anext__() == "frame-0"
+        assert redis.zsets[KEY][lease.lease_id] == clock["t"] - 25
+        assert await wrapped.__anext__() == "frame-1"
+        # Resuming after frame-0 refreshed the lease (25s > a third of the
+        # TTL) before the inner stream produced frame-1 another 25s later.
+        assert redis.zsets[KEY][lease.lease_id] == clock["t"] - 25
+        await wrapped.aclose()
+        assert state["closed"] is True
+
+    async def test_passes_through_without_a_lease(self):
+        async def inner():
+            yield "a"
+            yield "b"
+
+        assert [item async for item in hold_lease(inner(), None)] == ["a", "b"]