From fb66a0b23027ff811d766019c5dda001523d4d80 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 14 Sep 2026 12:52:15 +0100 Subject: [PATCH] fix: stability improvements on event loop --- docsgpt/api/asgi_stream.py | 23 ++++- docsgpt/api/async_sse.py | 45 ++------- docsgpt/api/devices/session_events.py | 109 +++++++++++++-------- docsgpt/api/events/routes.py | 35 ++----- docsgpt/api/user/artifacts/download.py | 4 +- docsgpt/streaming/async_redis.py | 19 +++- docsgpt/streaming/sse_leases.py | 23 ++--- tests/api/test_async_sse_routes.py | 18 ++++ tests/api/test_events_routes.py | 14 ++- tests/asgi_stream.py | 12 ++- tests/devices/test_session_events_async.py | 101 +++++++++++++++++++ tests/fake_async_redis.py | 8 +- tests/test_async_redis.py | 11 +++ tests/test_sse_leases.py | 34 ++++++- 14 files changed, 322 insertions(+), 134 deletions(-) diff --git a/docsgpt/api/asgi_stream.py b/docsgpt/api/asgi_stream.py index 157f2b47..b2ecd38d 100644 --- a/docsgpt/api/asgi_stream.py +++ b/docsgpt/api/asgi_stream.py @@ -3,12 +3,14 @@ from __future__ import annotations import logging -from typing import Any, AsyncIterable, Awaitable, Callable, Optional +from typing import Any, AsyncIterable, AsyncIterator, Awaitable, Callable, Optional import anyio from starlette.responses import StreamingResponse from starlette.types import Receive, Scope, Send +from docsgpt.streaming.sse_leases import StreamLease, hold_lease + logger = logging.getLogger(__name__) # Upper bound on each cleanup step. Cleanup is shielded from the cancellation @@ -16,6 +18,15 @@ logger = logging.getLogger(__name__) # hold the request task, and a graceful shutdown, indefinitely. _CLEANUP_TIMEOUT_SECONDS = 5.0 +# Headers for the browser-facing SSE streams (notifications, chat reconnect). +SSE_HEADERS = { + "Cache-Control": "no-store", + "X-Accel-Buffering": "no", + "Connection": "keep-alive", + # Marks the response as served by the event loop. Purely diagnostic. + "X-SSE-Transport": "async", +} + class ClosingStreamingResponse(StreamingResponse): """``StreamingResponse`` that always finalizes its stream. @@ -50,6 +61,16 @@ class ClosingStreamingResponse(StreamingResponse): await _bounded(self._on_close, "response on_close hook") +def sse_response(stream: AsyncIterator[str], lease: Optional[StreamLease]) -> ClosingStreamingResponse: + """Serve ``stream`` as SSE, refreshing ``lease`` between frames and releasing it however the response ends.""" + return ClosingStreamingResponse( + hold_lease(stream, lease), + media_type="text/event-stream", + headers=SSE_HEADERS, + on_close=lease.release if lease is not None else None, + ) + + async def _bounded(step: Callable[[], Awaitable[None]], what: str) -> None: """Run one cleanup step shielded from cancellation, giving up after the cleanup timeout.""" with anyio.move_on_after(_CLEANUP_TIMEOUT_SECONDS, shield=True) as scope: diff --git a/docsgpt/api/async_sse.py b/docsgpt/api/async_sse.py index 6ba92f07..d3fa3437 100644 --- a/docsgpt/api/async_sse.py +++ b/docsgpt/api/async_sse.py @@ -21,11 +21,11 @@ from typing import Optional import anyio from sqlalchemy import text from starlette.requests import Request -from starlette.responses import JSONResponse, Response +from starlette.responses import Response from starlette.routing import Route -from docsgpt.api.asgi_auth import authenticate -from docsgpt.api.asgi_stream import ClosingStreamingResponse +from docsgpt.api.asgi_auth import authenticate, bind_log_context, json_error +from docsgpt.api.asgi_stream import sse_response from docsgpt.core.settings import settings from docsgpt.storage.db.session import db_readonly from docsgpt.streaming.async_event_replay import ( @@ -36,11 +36,7 @@ from docsgpt.streaming.event_replay import ( DEFAULT_KEEPALIVE_SECONDS, DEFAULT_POLL_TIMEOUT_SECONDS, ) -from docsgpt.streaming.sse_leases import ( - StreamCapExceeded, - acquire_stream_lease, - hold_lease, -) +from docsgpt.streaming.sse_leases import StreamCapExceeded, acquire_stream_lease logger = logging.getLogger(__name__) @@ -89,22 +85,6 @@ def _user_owns_message(message_id: str, user_id: str) -> bool: ) return False -_SSE_HEADERS = { - "Cache-Control": "no-store", - "X-Accel-Buffering": "no", - "Connection": "keep-alive", - # Marks the response as served by the event-loop reader rather than the - # WSGI-threaded Flask fallback. Purely diagnostic — the frontend reads - # the body via fetch+getReader and ignores response headers. - "X-SSE-Transport": "async", -} - - -def _json(message: str, status_code: int) -> JSONResponse: - return JSONResponse( - {"success": False, "message": message}, status_code=status_code - ) - async def stream_message_events(request: Request) -> Response: """GET /api/messages/{message_id}/events — async reconnect tail. @@ -120,17 +100,18 @@ async def stream_message_events(request: Request) -> Response: return error user_id = decoded.get("sub") if isinstance(decoded, dict) else None if not user_id: - return _json("Authentication required", 401) + return json_error("Authentication required", 401) + bind_log_context("message_events", user_id) message_id = request.path_params["message_id"] if not _MESSAGE_ID_RE.match(message_id): - return _json("Invalid message id", 400) + return json_error("Invalid message id", 400) # Ownership check is a sync DB read — push it off the loop. owns = await anyio.to_thread.run_sync(_user_owns_message, message_id, user_id) if not owns: # Same opaque 404 as the Flask route — don't disclose existence. - return _json("Not found", 404) + return json_error("Not found", 404) # Per-user connection cap, shared with /api/events so it bounds a user's # total live SSE footprint. Reserved before the response opens so an @@ -139,7 +120,7 @@ async def stream_message_events(request: Request) -> Response: lease = await acquire_stream_lease(await get_async_redis_instance(), user_id) except StreamCapExceeded: logger.warning("sse.reconnect.rejected user_id=%s (over cap)", user_id) - return _json("Too many concurrent SSE connections", 429) + return json_error("Too many concurrent SSE connections", 429) raw_cursor = request.headers.get("Last-Event-ID") or request.query_params.get( "last_event_id" @@ -163,13 +144,7 @@ async def stream_message_events(request: Request) -> Response: keepalive_seconds=keepalive_seconds, poll_timeout_seconds=DEFAULT_POLL_TIMEOUT_SECONDS, ) - return ClosingStreamingResponse( - hold_lease(stream, lease), - media_type="text/event-stream", - headers=_SSE_HEADERS, - # Released once the response is over, however it ended. - on_close=lease.release if lease is not None else None, - ) + return sse_response(stream, lease) # Mounted in ``docsgpt/asgi.py`` ahead of the Flask catch-all. Keep diff --git a/docsgpt/api/devices/session_events.py b/docsgpt/api/devices/session_events.py index abb1c07f..afb137e3 100644 --- a/docsgpt/api/devices/session_events.py +++ b/docsgpt/api/devices/session_events.py @@ -2,7 +2,7 @@ A native-async Starlette route mounted ahead of Flask in ``docsgpt/asgi.py``. The CLI holds this stream open while it waits for commands, so it runs on the -event loop: the token check and ticket claim hop to a worker thread once, then +event loop: the token check and ticket claim run in worker threads, then commands arrive through the broker's async ``BLPOP`` and an idle session holds no thread. """ @@ -13,18 +13,20 @@ import json import logging import time from functools import partial -from typing import AsyncIterator +from typing import AsyncIterator, Optional import anyio from starlette.requests import Request from starlette.responses import JSONResponse, Response from starlette.routing import Route +from docsgpt.api.asgi_auth import bind_log_context from docsgpt.api.asgi_stream import ClosingStreamingResponse from docsgpt.api.devices.auth import authenticate_device +from docsgpt.core import log_context from docsgpt.core.settings import settings from docsgpt.core.shutdown import is_shutting_down -from docsgpt.devices.broker import DeviceBroker, get_broker +from docsgpt.devices.broker import DeviceBroker, SessionState, get_broker logger = logging.getLogger(__name__) @@ -52,45 +54,57 @@ def _sse_event(name: str, payload: dict, event_id: int) -> str: ) -async def _session_stream(broker: DeviceBroker, device: dict) -> AsyncIterator[str]: +def _signed_path(request: Request) -> str: + """Return the path without the mount prefix, as the Flask device routes see it in ``PATH_INFO``.""" + path = request.scope["path"] + root = request.scope.get("root_path", "") + if root and path.startswith(root): + path = path[len(root):] + return path + + +def _claim_session(broker: DeviceBroker, device: dict, session_id: str) -> Optional[SessionState]: + """Open the device's session iff ``session_id`` is its unexpired ticket; runs in a worker thread.""" + if not broker.validate_ticket(device["id"], session_id): + return None + return broker.register_session(device["id"], device["user_id"]) + + +async def _close_session(broker: DeviceBroker, session_id: str) -> None: + """Close the session once its response is over.""" + broker.close_session(session_id, reason="stream_end") + + +async def _session_stream(broker: DeviceBroker, sess: SessionState) -> AsyncIterator[str]: """Deliver queued commands until the session closes, idles out, or the server drains.""" - # Registration consumes the ticket. Shield it so a disconnect can't leave - # a registered session that the ``finally`` below never closes. - with anyio.CancelScope(shield=True): - sess = await anyio.to_thread.run_sync( - broker.register_session, device["id"], device["user_id"] - ) keepalive_interval = float(settings.SSE_KEEPALIVE_SECONDS) idle_seconds = float(settings.REMOTE_DEVICE_SESSION_IDLE_SECONDS) - try: - last_keepalive = time.time() - # ``closed`` is set when the CLI reconnects and a newer session - # replaces this one. - while not sess.closed.is_set(): - # Break promptly on shutdown (see docsgpt/core/shutdown.py). - if is_shutting_down(): - break - if time.time() - sess.last_activity_at > idle_seconds: - yield _sse_event( - "session_end", - {"reason": "inactivity_timeout"}, - sess.last_event_id + 1, - ) - sess.last_event_id += 1 - broker.close_session(sess.session_id, reason="idle") - return - envelope = await broker.next_command_async(sess, timeout=_POLL_TIMEOUT_SECONDS) - if envelope is None: - if time.time() - last_keepalive >= keepalive_interval: - last_keepalive = time.time() - yield ": heartbeat\n\n" - continue + last_keepalive = time.time() + # ``closed`` is set when the CLI reconnects and a newer session + # replaces this one. + while not sess.closed.is_set(): + # Break promptly on shutdown (see docsgpt/core/shutdown.py). + if is_shutting_down(): + break + if time.time() - sess.last_activity_at > idle_seconds: + yield _sse_event( + "session_end", + {"reason": "inactivity_timeout"}, + sess.last_event_id + 1, + ) sess.last_event_id += 1 - sess.last_activity_at = time.time() - yield _sse_event("invocation", envelope, sess.last_event_id) - last_keepalive = time.time() - finally: - broker.close_session(sess.session_id, reason="stream_end") + broker.close_session(sess.session_id, reason="idle") + return + envelope = await broker.next_command_async(sess, timeout=_POLL_TIMEOUT_SECONDS) + if envelope is None: + if time.time() - last_keepalive >= keepalive_interval: + last_keepalive = time.time() + yield ": heartbeat\n\n" + continue + sess.last_event_id += 1 + sess.last_activity_at = time.time() + yield _sse_event("invocation", envelope, sess.last_event_id) + last_keepalive = time.time() async def device_session_events(request: Request) -> Response: @@ -102,25 +116,34 @@ async def device_session_events(request: Request) -> Response: before any stream is opened. """ session_id = request.path_params["session_id"] - body = await request.body() + bind_log_context("devices.session_events") device, failure = await anyio.to_thread.run_sync( partial( authenticate_device, request.headers, request.method, - request.url.path, - lambda: body, + _signed_path(request), + # Read only when a signature is checked, after the token is known + # good, so an unauthenticated body is never buffered. + lambda: anyio.from_thread.run(request.body), ) ) if failure is not None: return _error(*failure) + log_context.bind(user_id=device["user_id"]) broker = get_broker() - if not await anyio.to_thread.run_sync(broker.validate_ticket, device["id"], session_id): + # Claiming consumes the ticket, so it happens before the response opens and + # ``on_close`` closes the session even if the body never starts. Shielded + # so a cancellation can't leave a registered session with no response. + with anyio.CancelScope(shield=True): + sess = await anyio.to_thread.run_sync(_claim_session, broker, device, session_id) + if sess is None: return _error("session_ticket_invalid", 410) return ClosingStreamingResponse( - _session_stream(broker, device), + _session_stream(broker, sess), media_type="text/event-stream", headers=_SSE_HEADERS, + on_close=partial(_close_session, broker, sess.session_id), ) diff --git a/docsgpt/api/events/routes.py b/docsgpt/api/events/routes.py index 614df456..52bf260a 100644 --- a/docsgpt/api/events/routes.py +++ b/docsgpt/api/events/routes.py @@ -26,7 +26,7 @@ from starlette.responses import Response from starlette.routing import Route from docsgpt.api.asgi_auth import authenticate, bind_log_context, json_error -from docsgpt.api.asgi_stream import ClosingStreamingResponse +from docsgpt.api.asgi_stream import sse_response from docsgpt.cache import PUBSUB_SOCKET_TIMEOUT_SECONDS from docsgpt.core.settings import settings from docsgpt.core.shutdown import is_shutting_down @@ -37,21 +37,13 @@ from docsgpt.events.keys import ( topic_name, ) from docsgpt.streaming.async_broadcast_channel import AsyncTopic -from docsgpt.streaming.async_redis import get_async_redis_instance -from docsgpt.streaming.sse_leases import ( - StreamCapExceeded, - acquire_stream_lease, - hold_lease, -) +from docsgpt.streaming.async_redis import ASYNC_REDIS_OP_TIMEOUT_SECONDS, get_async_redis_instance +from docsgpt.streaming.sse_leases import StreamCapExceeded, acquire_stream_lease logger = logging.getLogger(__name__) SUBSCRIBE_POLL_INTERVAL_SECONDS = 1.0 -# Upper bound on the replay-budget check, so a stalled Redis connection fails -# open instead of holding the request. -_REDIS_OP_TIMEOUT_SECONDS = 5.0 - # WHATWG SSE treats CRLF, CR, and LF equivalently as line terminators. _SSE_LINE_SPLIT = re.compile(r"\r\n|\r|\n") @@ -61,14 +53,6 @@ _SSE_LINE_SPLIT = re.compile(r"\r\n|\r|\n") # would silently fail. _STREAM_ID_RE = re.compile(r"^\d+(-\d+)?$") -_SSE_HEADERS = { - "Cache-Control": "no-store", - "X-Accel-Buffering": "no", - "Connection": "keep-alive", - # Marks the response as served by the event loop. Purely diagnostic. - "X-SSE-Transport": "async", -} - # Only emitted at most once per process so a misconfigured deployment # doesn't drown the logs. _local_user_warned = False @@ -146,7 +130,7 @@ async def _allow_replay( window = max(1, int(settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS)) key = replay_budget_key(user_id) try: - with anyio.fail_after(_REDIS_OP_TIMEOUT_SECONDS): + with anyio.fail_after(ASYNC_REDIS_OP_TIMEOUT_SECONDS): used = int(await redis_client.incr(key)) # Always (re)seed the TTL. Gating on ``used == 1`` would wedge # the counter forever if INCR succeeds but EXPIRE raises on @@ -431,7 +415,8 @@ async def stream_events(request: Request) -> Response: user_id = decoded.get("sub") if isinstance(decoded, dict) else None if not user_id: return json_error("Authentication required", 401) - bind_log_context("event_stream", user_id) + # Same endpoint value the Flask route logged, so saved log queries keep matching. + bind_log_context("event_stream.stream_events", user_id) # In dev deployments without AUTH_TYPE configured, every request # resolves to user_id="local" and shares one stream. Surface this so @@ -476,13 +461,7 @@ async def stream_events(request: Request) -> Response: " (rejected_invalid)" if last_event_id_invalid else "", ) stream = _event_stream(redis_client, user_id, last_event_id, last_event_id_invalid, push_enabled) - return ClosingStreamingResponse( - hold_lease(stream, lease), - media_type="text/event-stream", - headers=_SSE_HEADERS, - # Released once the response is over, however it ended. - on_close=lease.release if lease is not None else None, - ) + return sse_response(stream, lease) # Mounted in ``docsgpt/asgi.py`` ahead of the Flask catch-all. diff --git a/docsgpt/api/user/artifacts/download.py b/docsgpt/api/user/artifacts/download.py index 36802d7b..a57ae727 100644 --- a/docsgpt/api/user/artifacts/download.py +++ b/docsgpt/api/user/artifacts/download.py @@ -51,6 +51,7 @@ def _sanitize_header_filename(filename: Optional[str], fallback: str) -> str: def _ascii_fold(value: str) -> str: + """Drop accents and any other non-ASCII characters from ``value``.""" return unicodedata.normalize("NFKD", value).encode("ascii", "ignore").decode("ascii") @@ -130,7 +131,8 @@ async def download_artifact(request: Request) -> Response: decoded, error = await authenticate(request) if error is not None: return error - bind_log_context("artifact_download", decoded.get("sub") if decoded else None) + # Same endpoint value the Flask-RESTX route logged, so saved log queries keep matching. + bind_log_context("artifacts_download_artifact", decoded.get("sub") if decoded else None) query = request.query_params try: diff --git a/docsgpt/streaming/async_redis.py b/docsgpt/streaming/async_redis.py index 56211eba..c0c60100 100644 --- a/docsgpt/streaming/async_redis.py +++ b/docsgpt/streaming/async_redis.py @@ -15,10 +15,15 @@ from typing import Optional import redis.asyncio as aioredis +from docsgpt.cache import PUBSUB_SOCKET_TIMEOUT_SECONDS, _tcp_keepalive_options from docsgpt.core.settings import settings logger = logging.getLogger(__name__) +# Upper bound on one Redis command made while opening or holding a stream, so +# a stalled connection fails open instead of holding the request or cleanup. +ASYNC_REDIS_OP_TIMEOUT_SECONDS = 5.0 + _async_redis: Optional[aioredis.Redis] = None _creation_failed = False @@ -28,11 +33,12 @@ async def get_async_redis_instance() -> Optional[aioredis.Redis]: ``from_url`` builds the client without opening a socket (connection is lazy), so a transient broker outage surfaces later on the first command - rather than here. Mirrors the sync client's ``socket_connect_timeout`` - and ``health_check_interval`` so a half-open TCP can't wedge the tail - loop past its keepalive cadence. Every open stream holds a pooled - connection, so the pool is sized by ``ASYNC_REDIS_MAX_CONNECTIONS`` - rather than redis-py's default of 100. + rather than here. Uses the sync pub/sub client's timeouts and TCP + keepalive, so a connection NAT/IPVS dropped without a FIN fails a + ``BLPOP`` or ``XRANGE`` instead of hanging it. Pub/sub polls pass their + own read timeout, so ``socket_timeout`` doesn't cut idle subscribers. + Every open stream holds a pooled connection, so the pool is sized by + ``ASYNC_REDIS_MAX_CONNECTIONS`` rather than redis-py's default of 100. """ global _async_redis, _creation_failed if _async_redis is None and not _creation_failed: @@ -40,6 +46,9 @@ async def get_async_redis_instance() -> Optional[aioredis.Redis]: _async_redis = aioredis.Redis.from_url( settings.CACHE_REDIS_URL, socket_connect_timeout=2, + socket_timeout=PUBSUB_SOCKET_TIMEOUT_SECONDS, + socket_keepalive=True, + socket_keepalive_options=_tcp_keepalive_options(), health_check_interval=10, max_connections=int(settings.ASYNC_REDIS_MAX_CONNECTIONS), ) diff --git a/docsgpt/streaming/sse_leases.py b/docsgpt/streaming/sse_leases.py index b32527bd..8a5f3909 100644 --- a/docsgpt/streaming/sse_leases.py +++ b/docsgpt/streaming/sse_leases.py @@ -22,13 +22,10 @@ import anyio from docsgpt.core.settings import settings from docsgpt.events.keys import connection_leases_key +from docsgpt.streaming.async_redis import ASYNC_REDIS_OP_TIMEOUT_SECONDS logger = logging.getLogger(__name__) -# Upper bound on any one lease command, so a stalled Redis connection can't -# hold a request, a stream, or its cleanup. -_REDIS_OP_TIMEOUT_SECONDS = 5.0 - T = TypeVar("T") @@ -63,9 +60,11 @@ class StreamLease: return self.refreshed_at = now try: - with anyio.fail_after(_REDIS_OP_TIMEOUT_SECONDS): + with anyio.fail_after(ASYNC_REDIS_OP_TIMEOUT_SECONDS): async with self.redis.pipeline(transaction=True) as pipe: - pipe.zadd(self.key, {self.lease_id: time.time()}) + # XX: a lease pruned or cleared while refreshes were failing + # stays gone instead of retaking a slot past the cap. + pipe.zadd(self.key, {self.lease_id: time.time()}, xx=True) pipe.expire(self.key, int(self.ttl)) await pipe.execute() except Exception: @@ -73,7 +72,7 @@ class StreamLease: async def release(self) -> None: """Remove this lease; shielded so it completes during cancellation, and bounded.""" - with anyio.move_on_after(_REDIS_OP_TIMEOUT_SECONDS, shield=True): + with anyio.move_on_after(ASYNC_REDIS_OP_TIMEOUT_SECONDS, shield=True): try: await self.redis.zrem(self.key, self.lease_id) except Exception: @@ -99,21 +98,23 @@ async def acquire_stream_lease(redis: Any, user_id: str) -> Optional[StreamLease return None key = connection_leases_key(user_id) ttl = lease_ttl_seconds() - lease_id = uuid.uuid4().hex + lease = StreamLease(redis=redis, key=key, lease_id=uuid.uuid4().hex, ttl=ttl, refreshed_at=time.monotonic()) now = time.time() try: - with anyio.fail_after(_REDIS_OP_TIMEOUT_SECONDS): + with anyio.fail_after(ASYNC_REDIS_OP_TIMEOUT_SECONDS): # One transaction, so concurrent connects are counted one at a time. async with redis.pipeline(transaction=True) as pipe: pipe.zremrangebyscore(key, "-inf", now - ttl) - pipe.zadd(key, {lease_id: now}) + pipe.zadd(key, {lease.lease_id: now}) pipe.zcard(key) pipe.expire(key, int(ttl)) results = await pipe.execute() except Exception: logger.debug("SSE lease acquire failed for user=%s; failing open", user_id, exc_info=True) + # EXEC may have run before its reply was lost; the stream opens without + # a lease, so its member must not keep counting against the cap. + await lease.release() return None - lease = StreamLease(redis=redis, key=key, lease_id=lease_id, ttl=ttl, refreshed_at=time.monotonic()) if int(results[2]) > cap: await lease.release() raise StreamCapExceeded() diff --git a/tests/api/test_async_sse_routes.py b/tests/api/test_async_sse_routes.py index ac45681a..63831a31 100644 --- a/tests/api/test_async_sse_routes.py +++ b/tests/api/test_async_sse_routes.py @@ -184,6 +184,24 @@ class TestOwnershipGate: assert ": connected" in r.text + def test_stream_logs_carry_request_context(self): + from docsgpt.core import log_context + + seen: dict = {} + + async def _gen(message_id, last_event_id=None, **kwargs): + seen.update(log_context.snapshot()) + yield ": connected\n\n" + + with patch(_AUTH, return_value={"sub": "alice"}), patch( + _OWNS, return_value=True + ), patch(_STREAM, _gen): + _client().get(f"/api/messages/{VALID_UUID}/events") + assert seen["endpoint"] == "message_events" + assert seen["user_id"] == "alice" + assert seen["activity_id"] + + # ── Last-Event-ID parsing ─────────────────────────────────────────────────── diff --git a/tests/api/test_events_routes.py b/tests/api/test_events_routes.py index acb46259..b2f188c1 100644 --- a/tests/api/test_events_routes.py +++ b/tests/api/test_events_routes.py @@ -197,6 +197,18 @@ class TestStreamShape: assert record["poll_timeout"] == events_module.SUBSCRIBE_POLL_INTERVAL_SECONDS assert record["liveness_timeout"] == PUBSUB_SOCKET_TIMEOUT_SECONDS + def test_stream_logs_carry_the_flask_endpoint_name(self): + from docsgpt.core import log_context + + seen: dict = {} + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=_redis()) + ), patch(_SUBSCRIBE, _subscribe(probe=lambda: seen.update(log_context.snapshot()))): + _get() + # Same value the Flask route logged, so saved log queries keep matching. + assert seen["endpoint"] == "event_stream.stream_events" + assert seen["user_id"] == "alice" + def test_keepalive_emitted_on_idle_ticks(self, monkeypatch): monkeypatch.setattr(events_module.settings, "SSE_KEEPALIVE_SECONDS", 0) with patch(_AUTH, return_value=ALICE), patch( @@ -474,7 +486,7 @@ class TestReplayBudget: assert await events_module._allow_replay(redis, "alice", OLD_CURSOR) is True async def test_fail_open_when_redis_stalls(self, monkeypatch): - monkeypatch.setattr(events_module, "_REDIS_OP_TIMEOUT_SECONDS", 0.05) + monkeypatch.setattr(events_module, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05) redis = _redis() async def _stalled(_key): diff --git a/tests/asgi_stream.py b/tests/asgi_stream.py index ce8e2957..1daaacc4 100644 --- a/tests/asgi_stream.py +++ b/tests/asgi_stream.py @@ -19,6 +19,7 @@ async def stream_then_disconnect( *, headers: Optional[Mapping[str, str]] = None, query_string: bytes = b"", + root_path: str = "", chunks_before_disconnect: int = 1, timeout: float = 5.0, ) -> tuple[Optional[int], dict[str, str], list[bytes]]: @@ -26,8 +27,9 @@ async def stream_then_disconnect( A response that finishes on its own (a 401, a stream that ends) returns without any disconnect. ``chunks_before_disconnect=0`` disconnects as soon - as the request is read, before any body. ``timeout`` fails the test - instead of hanging it when a stream never ends. + as the request is read, before any body. ``root_path`` mounts the app under + a prefix, which ASGI servers include in ``path``. ``timeout`` fails the + test instead of hanging it when a stream never ends. Returns: tuple: ``(status, headers, body_chunks)`` with lower-cased header names. @@ -46,9 +48,9 @@ async def stream_then_disconnect( "http_version": "1.1", "method": "GET", "scheme": "http", - "path": path, - "raw_path": path.encode("utf-8"), - "root_path": "", + "path": root_path + path, + "raw_path": (root_path + path).encode("utf-8"), + "root_path": root_path, "query_string": query_string, "headers": [ (name.lower().encode("latin-1"), value.encode("latin-1")) diff --git a/tests/devices/test_session_events_async.py b/tests/devices/test_session_events_async.py index c2e72e4f..35fc7c13 100644 --- a/tests/devices/test_session_events_async.py +++ b/tests/devices/test_session_events_async.py @@ -16,13 +16,16 @@ import json import time from types import SimpleNamespace +import anyio import pytest from flask import Flask from starlette.applications import Starlette +from starlette.requests import Request from starlette.testclient import TestClient from docsgpt.api.devices import auth as auth_module from docsgpt.api.devices import session_events as events_module +from docsgpt.core import log_context from docsgpt.core.shutdown import begin_shutdown, reset_shutdown from tests.asgi_stream import stream_then_disconnect @@ -95,6 +98,27 @@ def _path(session_id: str) -> str: return f"/api/devices/sessions/{session_id}/events" +def _scope(session_id: str, headers: dict | None = None) -> dict: + """A raw ASGI GET scope for the stream route.""" + return { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "method": "GET", + "scheme": "http", + "path": _path(session_id), + "raw_path": _path(session_id).encode(), + "root_path": "", + "query_string": b"", + "headers": [ + (name.lower().encode("latin-1"), value.encode("latin-1")) + for name, value in (AUTH if headers is None else headers).items() + ], + "client": ("testclient", 50000), + "server": ("testserver", 80), + } + + def _open(session_id: str, headers: dict | None = None): return TestClient(_app()).get(_path(session_id), headers=headers if headers is not None else AUTH) @@ -150,6 +174,42 @@ class TestAuth: assert _open(ticket).status_code == 200 assert device.touched == [DEVICE["id"]] + @pytest.mark.asyncio + async def test_unauthenticated_request_body_is_never_read(self, device, broker): + # An endless body must not be buffered before the token check rejects it. + reads = 0 + sent: list = [] + + async def receive(): + nonlocal reads + reads += 1 + await anyio.sleep(0) + return {"type": "http.request", "body": b"x" * 65536, "more_body": True} + + async def send(message): + sent.append(message) + + with anyio.fail_after(2): + await _app()(_scope("st_x", headers={}), receive, send) + assert sent[0]["status"] == 401 + assert reads == 0 + + def test_stream_logs_carry_request_context(self, device, broker, monkeypatch): + b, _ = broker + ticket = _ticket(b) + seen: dict = {} + original = b.next_command_async + + async def _capture(sess, timeout=1.0): + seen.update(log_context.snapshot()) + return await original(sess, timeout=timeout) + + monkeypatch.setattr(b, "next_command_async", _capture) + assert _open(ticket).status_code == 200 + assert seen["endpoint"] == "devices.session_events" + assert seen["user_id"] == DEVICE["user_id"] + assert seen["activity_id"] + # ── stream ────────────────────────────────────────────────────────────── @@ -271,6 +331,27 @@ class TestStreamLifecycle: assert b.get_session(ticket) is None assert DEVICE["id"] not in b._sessions_by_device + async def test_ticket_is_claimed_before_the_body_starts(self, device, broker): + b, fake = broker + ticket = _ticket(b) + scope = {**_scope(ticket), "path_params": {"session_id": ticket}} + + async def receive(): + return {"type": "http.disconnect"} + + response = await events_module.device_session_events(Request(scope, receive)) + assert response.status_code == 200 + # A response whose body never runs has still used the ticket up. + assert fake.get(f"dev:ticket:{DEVICE['id']}") is None + assert b.get_session(ticket) is not None + + async def send(_message): + return None + + with anyio.fail_after(2): + await response(scope, receive, send) + assert b.get_session(ticket) is None + # ── machine-key signatures ────────────────────────────────────────────── @@ -320,6 +401,26 @@ class TestSignature: # Rejected before the session opened: the ticket is still unclaimed. assert fake.get(f"dev:ticket:{DEVICE['id']}") == ticket.encode() + @pytest.mark.asyncio + async def test_signed_path_excludes_root_path(self, device, broker, monkeypatch): + # The Flask device routes verify against PATH_INFO, which has the mount + # prefix stripped; the stream must check the same path. + monkeypatch.setattr(auth_module.settings, "REMOTE_DEVICE_REQUIRE_SIGNATURE", True) + b, _ = broker + ticket = _ticket(b) + fingerprint, headers = _signed_headers(_path(ticket)) + device.row["machine_pubkey_fingerprint"] = fingerprint + status, _, chunks = await stream_then_disconnect( + _app(), + _path(ticket), + headers=headers, + root_path="/docsgpt", + chunks_before_disconnect=10**6, + timeout=3, + ) + assert status == 200 + assert b"event: invocation" in b"".join(chunks) + def test_missing_signature_rejected(self, device, broker, monkeypatch): monkeypatch.setattr(auth_module.settings, "REMOTE_DEVICE_REQUIRE_SIGNATURE", True) b, _ = broker diff --git a/tests/fake_async_redis.py b/tests/fake_async_redis.py index 374a82e1..81729c1b 100644 --- a/tests/fake_async_redis.py +++ b/tests/fake_async_redis.py @@ -3,8 +3,9 @@ Sorted sets and transactional pipelines behave like Redis, enough for the per-connection leases. ``incr``, ``expire``, ``xrange`` and ``xinfo_stream`` are ``AsyncMock`` attributes that tests configure directly. Set ``fail`` to -make sorted-set commands and pipelines raise, or ``hang`` to make them never -return (a dead connection). +make sorted-set commands and pipelines raise, ``hang`` to make them never +return (a dead connection), or ``hang_after_execute`` to make a pipeline apply +its commands and then never reply (a reply lost after the server ran EXEC). """ from __future__ import annotations @@ -37,6 +38,7 @@ class FakeAsyncRedis: self.key_ttls: dict[str, int] = {} self.fail: Optional[BaseException] = None self.hang = False + self.hang_after_execute = False self.incr = AsyncMock(return_value=1) self.expire = AsyncMock(return_value=True) self.xrange = AsyncMock(return_value=[]) @@ -146,4 +148,6 @@ class FakePipeline: await self._redis._gate() results = [fn(*args, **kwargs) for fn, args, kwargs in self._ops] self._ops.clear() + if self._redis.hang_after_execute: + await anyio.sleep_forever() return results diff --git a/tests/test_async_redis.py b/tests/test_async_redis.py index 25791c49..07880203 100644 --- a/tests/test_async_redis.py +++ b/tests/test_async_redis.py @@ -40,6 +40,17 @@ class TestGetAsyncRedisInstance: assert kwargs["socket_connect_timeout"] == 2 assert kwargs["health_check_interval"] == 10 + async def test_reads_bounded_like_the_sync_pubsub_client(self): + # A silently dropped connection must fail a BLPOP or XRANGE, not hang it. + from docsgpt.cache import PUBSUB_SOCKET_TIMEOUT_SECONDS + + with patch(_FROM_URL) as from_url: + await async_redis.get_async_redis_instance() + kwargs = from_url.call_args.kwargs + assert kwargs["socket_timeout"] == PUBSUB_SOCKET_TIMEOUT_SECONDS + assert kwargs["socket_keepalive"] is True + assert "socket_keepalive_options" in kwargs + async def test_real_client_pool_honours_setting(self, monkeypatch): # from_url connects lazily, so this never opens a socket. monkeypatch.setattr(async_redis.settings, "ASYNC_REDIS_MAX_CONNECTIONS", 777) diff --git a/tests/test_sse_leases.py b/tests/test_sse_leases.py index 2d393a38..94362526 100644 --- a/tests/test_sse_leases.py +++ b/tests/test_sse_leases.py @@ -89,12 +89,22 @@ class TestAcquire: assert await acquire_stream_lease(redis, "alice") is None async def test_stalled_redis_fails_open_within_timeout(self, monkeypatch): - monkeypatch.setattr(sse_leases, "_REDIS_OP_TIMEOUT_SECONDS", 0.05) + monkeypatch.setattr(sse_leases, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05) redis = FakeAsyncRedis() redis.hang = True with anyio.fail_after(2): assert await acquire_stream_lease(redis, "alice") is None + async def test_timed_out_acquire_leaves_no_orphan_lease(self, monkeypatch): + # EXEC ran on the server but its reply never arrived: the stream opens + # without a lease, so the member it added must not keep counting. + monkeypatch.setattr(sse_leases, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05) + redis = FakeAsyncRedis() + redis.hang_after_execute = True + with anyio.fail_after(2): + assert await acquire_stream_lease(redis, "alice") is None + assert redis.zsets == {} + @pytest.mark.unit @pytest.mark.asyncio @@ -110,7 +120,7 @@ class TestLeaseLifecycle: assert set(redis.zsets[KEY]) == {second.lease_id} async def test_release_is_bounded_when_redis_stalls(self, monkeypatch): - monkeypatch.setattr(sse_leases, "_REDIS_OP_TIMEOUT_SECONDS", 0.05) + monkeypatch.setattr(sse_leases, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05) redis = FakeAsyncRedis() lease = await acquire_stream_lease(redis, "alice") redis.hang = True @@ -144,6 +154,26 @@ class TestLeaseLifecycle: assert long_lived.lease_id not in redis.zsets[KEY] assert len(redis.zsets[KEY]) == 7 + async def test_refresh_does_not_revive_a_pruned_lease(self, clock): + # Refreshes failed past the TTL, a newer connect pruned the lease and + # took the slot. Re-adding it would put the user over the cap. + redis = FakeAsyncRedis() + lease = await acquire_stream_lease(redis, "alice") + redis.zsets[KEY].pop(lease.lease_id) + _fill(redis, 8, clock["t"]) + clock["t"] += 30 + await lease.refresh_if_due() + assert lease.lease_id not in redis.zsets[KEY] + assert len(redis.zsets[KEY]) == 8 + + async def test_refresh_does_not_undo_a_manual_clear(self, clock): + redis = FakeAsyncRedis() + lease = await acquire_stream_lease(redis, "alice") + redis.zsets.clear() # redis-cli DEL user::sse_leases + clock["t"] += 30 + await lease.refresh_if_due() + assert redis.zsets == {} + async def test_refresh_failure_is_swallowed(self, clock): redis = FakeAsyncRedis() lease = await acquire_stream_lease(redis, "alice")