From e889392eba6494366f7f0666ea1c1e07db06a25e Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 14 Sep 2026 10:19:59 +0100 Subject: [PATCH] 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 | 20 +- docsgpt/api/asgi_stream.py | 49 +++-- docsgpt/api/async_sse.py | 93 ++-------- docsgpt/api/devices/session_events.py | 2 + docsgpt/api/events/routes.py | 85 ++++----- docsgpt/api/user/artifacts/download.py | 35 +++- docsgpt/events/keys.py | 6 +- docsgpt/streaming/async_broadcast_channel.py | 17 +- docsgpt/streaming/sse_leases.py | 129 +++++++++++++ tests/api/test_asgi_stream.py | 92 +++++++++ tests/api/test_async_sse_routes.py | 94 ++++++---- tests/api/test_events_routes.py | 132 +++++++------ tests/api/user/test_artifacts_routes.py | 52 ++++++ tests/fake_async_redis.py | 149 +++++++++++++++ tests/test_async_broadcast_channel.py | 19 ++ tests/test_events_substrate.py | 6 +- tests/test_sse_leases.py | 186 +++++++++++++++++++ 17 files changed, 911 insertions(+), 255 deletions(-) create mode 100644 docsgpt/streaming/sse_leases.py create mode 100644 tests/api/test_asgi_stream.py create mode 100644 tests/fake_async_redis.py create mode 100644 tests/test_sse_leases.py diff --git a/docs/runbooks/sse-notifications.md b/docs/runbooks/sse-notifications.md index 70047ec9..1d100c43 100644 --- a/docs/runbooks/sse-notifications.md +++ b/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::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::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::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::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::stream MAXLEN 0 -# Clear a wedged concurrent-connection counter -redis-cli -n 2 DEL user::sse_count +# Clear a user's SSE connection leases (they also age out on their own) +redis-cli -n 2 DEL user::sse_leases # Force-flip every client to re-snapshot (drop the stream key entirely # — destroys the backlog; clients reconnect with their last id and diff --git a/docsgpt/api/asgi_stream.py b/docsgpt/api/asgi_stream.py index d2de6629..157f2b47 100644 --- a/docsgpt/api/asgi_stream.py +++ b/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) diff --git a/docsgpt/api/async_sse.py b/docsgpt/api/async_sse.py index ed293891..6ba92f07 100644 --- a/docsgpt/api/async_sse.py +++ b/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, ) diff --git a/docsgpt/api/devices/session_events.py b/docsgpt/api/devices/session_events.py index ffd9231c..abb1c07f 100644 --- a/docsgpt/api/devices/session_events.py +++ b/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" diff --git a/docsgpt/api/events/routes.py b/docsgpt/api/events/routes.py index e4d634cd..614df456 100644 --- a/docsgpt/api/events/routes.py +++ b/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, ) diff --git a/docsgpt/api/user/artifacts/download.py b/docsgpt/api/user/artifacts/download.py index 4547c298..36802d7b 100644 --- a/docsgpt/api/user/artifacts/download.py +++ b/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: diff --git a/docsgpt/events/keys.py b/docsgpt/events/keys.py index d72777eb..041c3db2 100644 --- a/docsgpt/events/keys.py +++ b/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: diff --git a/docsgpt/streaming/async_broadcast_channel.py b/docsgpt/streaming/async_broadcast_channel.py index 1c83f187..57e05f87 100644 --- a/docsgpt/streaming/async_broadcast_channel.py +++ b/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: diff --git a/docsgpt/streaming/sse_leases.py b/docsgpt/streaming/sse_leases.py new file mode 100644 index 00000000..b32527bd --- /dev/null +++ b/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() diff --git a/tests/api/test_asgi_stream.py b/tests/api/test_asgi_stream.py new file mode 100644 index 00000000..c574e7c4 --- /dev/null +++ b/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 diff --git a/tests/api/test_async_sse_routes.py b/tests/api/test_async_sse_routes.py index 290fe4a0..ac45681a 100644 --- a/tests/api/test_async_sse_routes.py +++ b/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//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 == {} diff --git a/tests/api/test_events_routes.py b/tests/api/test_events_routes.py index a36bbff5..acb46259 100644 --- a/tests/api/test_events_routes.py +++ b/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 ────────────────────────────────────────────────────── diff --git a/tests/api/user/test_artifacts_routes.py b/tests/api/user/test_artifacts_routes.py index f77d8ba9..2e0e6028 100644 --- a/tests/api/user/test_artifacts_routes.py +++ b/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 diff --git a/tests/fake_async_redis.py b/tests/fake_async_redis.py new file mode 100644 index 00000000..374a82e1 --- /dev/null +++ b/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 diff --git a/tests/test_async_broadcast_channel.py b/tests/test_async_broadcast_channel.py index 49865cf1..0b888c89 100644 --- a/tests/test_async_broadcast_channel.py +++ b/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)): diff --git a/tests/test_events_substrate.py b/tests/test_events_substrate.py index 5d728094..7b3578e3 100644 --- a/tests/test_events_substrate.py +++ b/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 ─────────────────────────────────────────────────── diff --git a/tests/test_sse_leases.py b/tests/test_sse_leases.py new file mode 100644 index 00000000..2d393a38 --- /dev/null +++ b/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"]