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.
This commit is contained in:
Alex committed 2026-09-14 10:19:59 +01:00
1 parent 54b540ea42
commit e889392eba
17 files changed
+911 -255

No files matched your search

+92
View File
@@ -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
+60 -34
View File
@@ -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 == {}
+75 -57
View File
@@ -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 ──────────────────────────────────────────────────────
+52
View File
@@ -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