From 54b540ea42794e3c1e7bbe930850aa5fb68746fe Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 14 Sep 2026 08:33:46 +0100 Subject: [PATCH 1/7] feat(api): serve long-lived streams on the event loop Move GET /api/events, the remote-device command stream and artifact downloads from Flask to Starlette routes mounted ahead of the Flask catch-all. On Flask each held an a2wsgi threadpool slot for as long as its response stayed open, and because uvicorn drops writes after a client disconnects, a closed tab never released it. - asgi_auth: one JWT/OIDC gate for Starlette routes; the chat reconnect reader uses it too - ClosingStreamingResponse closes the body iterator and releases the SSE slot or file handle even when the client leaves before the first frame - AsyncTopic liveness probe replaces the sync client's socket_timeout guard against half-open pub/sub sockets - ASYNC_REDIS_MAX_CONNECTIONS sizes the async Redis pool; every open stream holds a connection and redis-py defaults to 100 --- AGENTS.md | 15 +- docs/content/Agents/notifications.mdx | 3 +- .../Deploying/Development-Environment.mdx | 2 +- docs/content/Guides/Architecture.mdx | 6 +- docs/runbooks/sse-notifications.md | 41 +- docsgpt/api/asgi_auth.py | 59 + docsgpt/api/asgi_stream.py | 44 + docsgpt/api/async_sse.py | 20 +- docsgpt/api/devices/auth.py | 94 +- docsgpt/api/devices/routes.py | 6 +- docsgpt/api/devices/session.py | 84 +- docsgpt/api/devices/session_events.py | 133 ++ docsgpt/api/events/routes.py | 487 ++++--- docsgpt/api/user/artifacts/authz.py | 43 +- docsgpt/api/user/artifacts/download.py | 181 +++ docsgpt/api/user/artifacts/routes.py | 123 +- docsgpt/app.py | 2 - docsgpt/asgi.py | 18 +- docsgpt/core/settings.py | 9 +- docsgpt/devices/broker.py | 60 +- docsgpt/streaming/async_broadcast_channel.py | 50 +- docsgpt/streaming/async_redis.py | 19 +- tests/api/test_asgi_auth.py | 116 ++ tests/api/test_async_sse_routes.py | 2 +- tests/api/test_events_routes.py | 1116 ++++++++--------- tests/api/user/test_artifacts_routes.py | 558 +++++---- tests/asgi_stream.py | 86 ++ tests/devices/conftest.py | 38 + tests/devices/test_session_events_async.py | 401 ++++++ tests/devices/test_session_ticket.py | 158 ++- tests/test_asgi.py | 89 ++ tests/test_async_broadcast_channel.py | 168 +++ tests/test_async_redis.py | 67 + 33 files changed, 2822 insertions(+), 1476 deletions(-) create mode 100644 docsgpt/api/asgi_auth.py create mode 100644 docsgpt/api/asgi_stream.py create mode 100644 docsgpt/api/devices/session_events.py create mode 100644 docsgpt/api/user/artifacts/download.py create mode 100644 tests/api/test_asgi_auth.py create mode 100644 tests/asgi_stream.py create mode 100644 tests/devices/test_session_events_async.py create mode 100644 tests/test_async_broadcast_channel.py create mode 100644 tests/test_async_redis.py diff --git a/AGENTS.md b/AGENTS.md index b8aeb247..90356fc4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -66,12 +66,17 @@ inner loop (quick startup, the Werkzeug interactive debugger), but it serves **only** the WSGI Flask app and omits the routes mounted on the ASGI shell in `docsgpt/asgi.py`: -- the `/mcp` FastMCP endpoint, and -- the native-async SSE reconnect reader `GET /api/messages//events`. +- the `/mcp` FastMCP endpoint, +- the chat reconnect reader `GET /api/messages//events`, +- the notification stream `GET /api/events`, +- the remote-device command stream `GET /api/devices/sessions//events`, and +- artifact downloads `GET /api/artifacts//download`. -Under `flask run` those paths 404. Chat still works (`POST /stream` is a -Flask route), but a stream interrupted by a disconnect won't auto-resume on -reconnect. Use `flask run` only when you don't need those routes. +These are native-async Starlette routes because they hold a response open +for a long time; on Flask each would pin a WSGI threadpool slot. Under +`flask run` those paths 404: chat still works (`POST /stream` is a Flask +route), but live notifications, stream auto-resume, paired devices and +artifact downloads don't. Use `flask run` only when you don't need them. Production uses `gunicorn -k uvicorn_worker.UvicornWorker` against the same `docsgpt.asgi:asgi_app` target; see `docsgpt/Dockerfile` for the diff --git a/docs/content/Agents/notifications.mdx b/docs/content/Agents/notifications.mdx index 7039c0b5..c6633556 100644 --- a/docs/content/Agents/notifications.mdx +++ b/docs/content/Agents/notifications.mdx @@ -70,7 +70,7 @@ GET /api/messages//events This replays the message's events past your last-seen sequence number and tails the rest live. It is backed by the Postgres `message_events` journal (retained for `MESSAGE_EVENTS_RETENTION_DAYS`, default 14). -The chat reconnect endpoint is a native-async route served by the ASGI entrypoint. Under a plain `flask run` dev server it returns `404`; run the backend via the ASGI app (`uvicorn docsgpt.asgi:asgi_app`) or the production gunicorn uvicorn worker to use it. See the [Development Environment](/Deploying/Development-Environment) guide. +Both channels are native-async routes served by the ASGI entrypoint. Under a plain `flask run` dev server they return `404`; run the backend via the ASGI app (`uvicorn docsgpt.asgi:asgi_app`) or the production gunicorn uvicorn worker to use them. See the [Development Environment](/Deploying/Development-Environment) guide. ## Settings @@ -81,6 +81,7 @@ The chat reconnect endpoint is a native-async route served by the ASGI entrypoin | `EVENTS_STREAM_MAXLEN` | `1000` | Per-user backlog cap (approximate). | | `SSE_KEEPALIVE_SECONDS` | `15` | Keepalive comment-frame cadence (keep below your proxy's idle timeout). | | `SSE_MAX_CONCURRENT_PER_USER` | `8` | Max simultaneous SSE connections per user (`0` disables the cap). | +| `ASYNC_REDIS_MAX_CONNECTIONS` | `2000` | Redis connection pool per API worker. Each open SSE connection holds one, so this caps concurrent streams per worker. | | `EVENTS_REPLAY_MAX_PER_REQUEST` | `200` | Max backlog entries per replay request. | | `EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW` | `30` | Per-user replay requests per window (`0` disables). | | `EVENTS_REPLAY_BUDGET_WINDOW_SECONDS` | `60` | Replay budget window length. | diff --git a/docs/content/Deploying/Development-Environment.mdx b/docs/content/Deploying/Development-Environment.mdx index fd2cc8c2..e7e877c4 100644 --- a/docs/content/Deploying/Development-Environment.mdx +++ b/docs/content/Deploying/Development-Environment.mdx @@ -133,7 +133,7 @@ To run the DocsGPT backend locally, you'll need to set up a Python environment a flask --app docsgpt/app.py run --host=0.0.0.0 --port=7091 ``` - But it serves **only** the WSGI Flask app and omits the routes mounted on the ASGI shell in `docsgpt/asgi.py`: the `/mcp` FastMCP endpoint and the native-async SSE reconnect reader `GET /api/messages//events`. Under `flask run` those paths return 404 — chat still works (`POST /stream` is a Flask route), but a stream interrupted by a disconnect won't auto-resume on reconnect. Use `flask run` only when you don't need those routes. + But it serves **only** the WSGI Flask app and omits the native-async routes mounted on the ASGI shell in `docsgpt/asgi.py`: the `/mcp` FastMCP endpoint, the chat reconnect reader `GET /api/messages//events`, the notification stream `GET /api/events`, the remote-device command stream `GET /api/devices/sessions//events`, and artifact downloads `GET /api/artifacts//download`. Under `flask run` those paths return 404 — chat still works (`POST /stream` is a Flask route), but live notifications, stream auto-resume, paired devices and artifact downloads don't. Use `flask run` only when you don't need them. 6. **Start the Celery Worker:** diff --git a/docs/content/Guides/Architecture.mdx b/docs/content/Guides/Architecture.mdx index de76c087..a65f3dcb 100644 --- a/docs/content/Guides/Architecture.mdx +++ b/docs/content/Guides/Architecture.mdx @@ -22,7 +22,7 @@ flowchart TB Clients["Web UI, widgets and API clients"] --> API subgraph Runtime["DocsGPT processes"] - API["API process
Flask routes and agent execution
ASGI shell: MCP and chat reconnect"] + API["API process
Flask routes and agent execution
ASGI shell: MCP and streaming routes"] Redis["Redis
Task broker, results and realtime delivery"] Worker["Celery workers
Ingestion, parsing, embeddings
and background agent runs"] API -->|"Enqueue work"| Redis @@ -63,7 +63,7 @@ The storage boxes describe responsibilities, not a required number of servers. F ### API entry points -[The ASGI application](https://github.com/arc53/DocsGPT/blob/main/docsgpt/asgi.py) mounts Flask through a WSGI adapter, the FastMCP server, and native asynchronous chat-reconnect readers in one process. +[The ASGI application](https://github.com/arc53/DocsGPT/blob/main/docsgpt/asgi.py) mounts Flask through a WSGI adapter, the FastMCP server, and native asynchronous routes in one process. The asynchronous routes are the ones that hold a response open: the notification stream, chat reconnect, the remote-device command stream, and artifact downloads. They run on the event loop, so an idle connection doesn't occupy a Flask worker thread. | Interface | Purpose | | --- | --- | @@ -76,7 +76,7 @@ The storage boxes describe responsibilities, not a required number of servers. F DocsGPT also acts as an **MCP client**: agents can call external MCP servers through the [MCP tool integration](/Guides/Integrations/mcp-tool-integration). This is a separate direction from exposing DocsGPT search at `/mcp`. -Use the ASGI entry point for the complete application. A plain `flask run` server omits `/mcp` and the native async chat-reconnect routes. See [Development Environment](/Deploying/Development-Environment). +Use the ASGI entry point for the complete application. A plain `flask run` server omits `/mcp` and the native asynchronous routes: notifications, chat reconnect, the remote-device command stream, and artifact downloads. See [Development Environment](/Deploying/Development-Environment). ### Agents and retrieval diff --git a/docs/runbooks/sse-notifications.md b/docs/runbooks/sse-notifications.md index e8ee55f5..70047ec9 100644 --- a/docs/runbooks/sse-notifications.md +++ b/docs/runbooks/sse-notifications.md @@ -39,10 +39,10 @@ redis-cli -n 2 XLEN user::stream Worker (publish_user_event) Frontend tab │ ▲ ▼ │ GET /api/events SSE - Redis Streams: XADD Flask route + Redis Streams: XADD async route (event loop) user::stream ──────────────► replay_backlog (snapshot) │ + - ▼ Topic.subscribe (live tail) + ▼ AsyncTopic.subscribe (live tail) Redis pub/sub: PUBLISH │ user: ────────────────────────────────┘ ``` @@ -204,11 +204,13 @@ the snapshot would re-serve from the start, the cap would re-trip). logs to confirm the publisher is still firing. - Pull `event.connect user=` from the API logs to confirm the client is reconnecting. -- Check the gunicorn worker count and `WSGIMiddleware(workers=32)` — - if the deploy reduced worker count, the per-user cap is still 8 - but total concurrent SSE connections are bounded by `gunicorn - workers × 32`. A capacity miss looks like users randomly getting - 429'd. +- Check the async Redis pool. Every open `/api/events` stream holds + one connection from it, so concurrent streams per worker are bounded + by `ASYNC_REDIS_MAX_CONNECTIONS` (default 2000). When the pool is + full the API logs `async Redis pool exhausted subscribing to + user:` and the stream closes at once, so clients reconnect in a + loop. Raise the setting (keeping the total across workers under the + Redis server's `maxclients`) or add workers. --- @@ -324,6 +326,7 @@ Everything in `docsgpt/core/settings.py`: | `EVENTS_STREAM_MAXLEN` | `1000` | Per-user backlog cap. Approximate via `XADD MAXLEN ~`. | | `SSE_KEEPALIVE_SECONDS` | `15` | Comment-frame cadence. Must sit under reverse-proxy idle close. | | `SSE_MAX_CONCURRENT_PER_USER` | `8` | Cap on simultaneous SSE connections per user. 0 = disabled. | +| `ASYNC_REDIS_MAX_CONNECTIONS` | `2000` | Async Redis pool per API worker; each open stream holds one connection. | | `EVENTS_REPLAY_MAX_PER_REQUEST` | `200` | Hard cap on snapshot rows per request. | | `EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW` | `30` | Per-user replays per window. 0 = disabled. | | `EVENTS_REPLAY_BUDGET_WINDOW_SECONDS` | `60` | Window length. | @@ -360,22 +363,16 @@ running task and the terminal SSE arrives later, the toast pops back. Intentional ("notify the user it's done"); revisit if the re-surface UX is too aggressive for v2. -### Reconnect reader needs the ASGI entrypoint +### Both SSE channels need the ASGI entrypoint -The chat-stream reconnect reader `GET /api/messages//events` -is a native-async Starlette route mounted in -`docsgpt/asgi.py`, not a Flask route. Plain `flask run` -serves only the WSGI Flask app, so under it that endpoint 404s -and reconnect-after-disconnect can't resume. Run the backend via -`uvicorn docsgpt.asgi:asgi_app --reload` (or the production -gunicorn uvicorn-worker) to exercise it. - -### Werkzeug doesn't auto-reload route files - -The dev server (`flask run`) doesn't watch -`docsgpt/api/events/routes.py` for changes by default. -After editing the route, restart Flask manually — `--reload` -isn't on. (Production gunicorn reloads via deploy.) +`GET /api/events` and the chat-stream reconnect reader +`GET /api/messages//events` are native-async Starlette routes +mounted in `docsgpt/asgi.py`, not Flask routes. Plain `flask run` +serves only the WSGI Flask app, so under it both endpoints 404: +no live notifications, and reconnect-after-disconnect can't resume. +Run the backend via `uvicorn docsgpt.asgi:asgi_app --reload` (or the +production gunicorn uvicorn-worker) to exercise them; `--reload` +picks up edits to `docsgpt/api/events/routes.py`. ### MCP OAuth completion can fall outside the user stream's MAXLEN window diff --git a/docsgpt/api/asgi_auth.py b/docsgpt/api/asgi_auth.py new file mode 100644 index 00000000..93f08a23 --- /dev/null +++ b/docsgpt/api/asgi_auth.py @@ -0,0 +1,59 @@ +"""Request authentication for the native-async (Starlette) routes. + +Mirrors ``authenticate_request`` and the log-context hooks in +``docsgpt/app.py`` so a route moved off Flask keeps the same gates: the same +JWT decoder and 401 payloads, and the OIDC revocation check. +""" + +from __future__ import annotations + +import uuid +from contextvars import Token +from typing import Optional, Tuple + +import anyio +from starlette.requests import Request +from starlette.responses import JSONResponse + +from docsgpt.api.oidc.denylist import is_denied as oidc_session_denied +from docsgpt.auth import handle_auth +from docsgpt.core import log_context +from docsgpt.core.settings import settings + + +async def authenticate(request: Request) -> Tuple[Optional[dict], Optional[JSONResponse]]: + """Decode the caller's JWT the way Flask's ``authenticate_request`` does. + + Args: + request: The incoming Starlette request. + + Returns: + tuple: ``(claims, None)`` for an authenticated caller, ``(None, None)`` + when no token was sent (the route decides whether that is allowed), or + ``(None, response)`` carrying the 401 to return. + """ + decoded = handle_auth(request) + if not decoded: + return None, None + if "error" in decoded: + return None, JSONResponse(decoded, status_code=401) + # The denylist is a sync Redis read; keep it off the event loop. + if settings.AUTH_TYPE == "oidc" and await anyio.to_thread.run_sync(oidc_session_denied, decoded): + return None, JSONResponse( + {"message": "Authentication error: session revoked", "error": "token_revoked"}, + status_code=401, + ) + return decoded, None + + +def bind_log_context(endpoint: str, user_id: Optional[str] = None) -> Token: + """Stamp this request's log records with a fresh activity id, the endpoint and user. + + Each request runs in its own task, so the binding ends with the request. + """ + return log_context.bind(activity_id=str(uuid.uuid4()), endpoint=endpoint, user_id=user_id) + + +def json_error(message: str, status_code: int) -> JSONResponse: + """Return the ``{"success": false, "message": ...}`` body the Flask routes use.""" + return JSONResponse({"success": False, "message": message}, status_code=status_code) diff --git a/docsgpt/api/asgi_stream.py b/docsgpt/api/asgi_stream.py new file mode 100644 index 00000000..d2de6629 --- /dev/null +++ b/docsgpt/api/asgi_stream.py @@ -0,0 +1,44 @@ +"""Streaming response for native-async routes whose streams hold resources.""" + +from __future__ import annotations + +import logging +from typing import Awaitable, Callable, Optional + +import anyio +from starlette.responses import StreamingResponse +from starlette.types import Receive, Scope, Send + +logger = logging.getLogger(__name__) + + +class ClosingStreamingResponse(StreamingResponse): + """``StreamingResponse`` that always finalizes its stream. + + 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. + """ + + def __init__(self, content, *, on_close: Optional[Callable[[], Awaitable[None]]] = None, **kwargs) -> None: + super().__init__(content, **kwargs) + self._on_close = on_close + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + 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") diff --git a/docsgpt/api/async_sse.py b/docsgpt/api/async_sse.py index fc6493bf..ed293891 100644 --- a/docsgpt/api/async_sse.py +++ b/docsgpt/api/async_sse.py @@ -2,7 +2,7 @@ These Starlette routes serve the chat-stream *reconnect* path on the event loop, so a long-lived, mostly-idle tail costs a coroutine instead of one of -the 32 a2wsgi threadpool slots (see ``docsgpt/asgi.py``). They are the +the a2wsgi threadpool slots (see ``docsgpt/asgi.py``). They are the sole reconnect reader — the old Flask blueprint has been removed. The heavy *producer* (``POST /api/answer/stream`` → agent → LLM) stays on the sync path untouched. @@ -24,8 +24,7 @@ from starlette.requests import Request from starlette.responses import JSONResponse, StreamingResponse from starlette.routing import Route -from docsgpt.api.oidc.denylist import is_denied as oidc_session_denied -from docsgpt.auth import handle_auth +from docsgpt.api.asgi_auth import authenticate from docsgpt.core.settings import settings from docsgpt.events.keys import connection_counter_key from docsgpt.storage.db.session import db_readonly @@ -176,19 +175,14 @@ async def stream_message_events(request: Request) -> JSONResponse | StreamingRes cursor → per-user connection cap) then streams snapshot+tail off the event loop. """ - # ``handle_auth`` only reads ``request.headers.get("Authorization")``; - # Starlette's headers are case-insensitive, so the Flask helper works - # verbatim. With AUTH_TYPE unset it returns ``{"sub": "local"}``. - decoded = handle_auth(request) - if isinstance(decoded, dict) and "error" in decoded: - return _json("Authentication error: invalid token", 401) + # Same JWT decoder and OIDC revocation check as the Flask routes. With + # AUTH_TYPE unset the caller resolves to ``{"sub": "local"}``. + decoded, error = await authenticate(request) + if error is not None: + return error user_id = decoded.get("sub") if isinstance(decoded, dict) else None if not user_id: return _json("Authentication required", 401) - if settings.AUTH_TYPE == "oidc" and await anyio.to_thread.run_sync( - oidc_session_denied, decoded - ): - return _json("Authentication error: session revoked", 401) message_id = request.path_params["message_id"] if not _MESSAGE_ID_RE.match(message_id): diff --git a/docsgpt/api/devices/auth.py b/docsgpt/api/devices/auth.py index 89a1b419..72bcde03 100644 --- a/docsgpt/api/devices/auth.py +++ b/docsgpt/api/devices/auth.py @@ -6,7 +6,7 @@ import base64 import hashlib import logging import time -from typing import Optional, Tuple +from typing import Callable, Mapping, Optional, Tuple from flask import jsonify, make_response, request @@ -17,6 +17,9 @@ from docsgpt.storage.db.session import db_readonly, db_session logger = logging.getLogger(__name__) +#: ``(error_code, http_status)`` for a rejected device request. +DeviceAuthFailure = Tuple[str, int] + def hash_session_token(token: str) -> str: """SHA-256 over the opaque session token. Stored in ``devices.token_hash``.""" @@ -45,32 +48,46 @@ def _canonical_payload(method: str, path: str, ts: str, body: bytes) -> str: return f"{method} {path} {ts} {body_hash}" -def verify_device_session(*, touch: bool = True) -> Tuple[Optional[dict], Optional[tuple]]: - """Validate the device session token on the current request. +def authenticate_device( + headers: Mapping[str, str], + method: str, + path: str, + get_body: Callable[[], bytes], + *, + touch: bool = True, +) -> Tuple[Optional[dict], Optional[DeviceAuthFailure]]: + """Validate a device session token from the raw request parts. - Returns ``(device_row, None)`` on success or ``(None, response)`` on - failure, where ``response`` is a ``(json, status)`` Flask tuple. + Framework-neutral core shared by the Flask routes and the native-async + SSE stream, which calls it from a worker thread since it reads Postgres. Args: + headers: Request headers (a case-insensitive mapping). + method: HTTP method, as the CLI signed it. + path: Request path without the query string, as the CLI signed it. + get_body: Returns the raw body; only called when signatures are required. touch: When True, bump ``last_seen_at`` on the device row. + + Returns: + tuple: ``(device_row, None)`` on success or ``(None, (code, status))``. """ - auth = request.headers.get("Authorization", "") + auth = headers.get("Authorization", "") if not auth.lower().startswith("bearer "): - return None, _error("missing_token", 401) + return None, ("missing_token", 401) token = auth.split(" ", 1)[1].strip() if not token: - return None, _error("missing_token", 401) + return None, ("missing_token", 401) token_hash = hash_session_token(token) with db_readonly() as conn: device = DevicesRepository(conn).find_by_token_hash(token_hash) if device is None: - return None, _error("invalid_token", 401) + return None, ("invalid_token", 401) if settings.REMOTE_DEVICE_REQUIRE_SIGNATURE: - sig_error = _verify_signature(device) - if sig_error is not None: - return None, sig_error + failure = _verify_signature(device, headers, method, path, get_body) + if failure is not None: + return None, failure if touch: try: @@ -82,23 +99,46 @@ def verify_device_session(*, touch: bool = True) -> Tuple[Optional[dict], Option return device, None -def _verify_signature(device: dict) -> Optional[tuple]: +def verify_device_session(*, touch: bool = True) -> Tuple[Optional[dict], Optional[tuple]]: + """Validate the device session token on the current Flask request. + + Returns ``(device_row, None)`` on success or ``(None, response)`` on + failure, where ``response`` is a ready Flask error response. + + Args: + touch: When True, bump ``last_seen_at`` on the device row. + """ + device, failure = authenticate_device( + request.headers, request.method, request.path, request.get_data, touch=touch + ) + if failure is not None: + return None, _error(*failure) + return device, None + + +def _verify_signature( + device: dict, + headers: Mapping[str, str], + method: str, + path: str, + get_body: Callable[[], bytes], +) -> Optional[DeviceAuthFailure]: """Verify ``X-Device-Signature`` against the stored machine pubkey.""" - sig_b64 = request.headers.get("X-Device-Signature") - ts = request.headers.get("X-Device-Timestamp") - fp = request.headers.get("X-Device-Machine-Key") + sig_b64 = headers.get("X-Device-Signature") + ts = headers.get("X-Device-Timestamp") + fp = headers.get("X-Device-Machine-Key") if not sig_b64 or not ts or not fp: - return _error("missing_signature", 401) + return ("missing_signature", 401) try: ts_int = int(ts) except (TypeError, ValueError): - return _error("invalid_timestamp", 401) + return ("invalid_timestamp", 401) if abs(time.time() - ts_int) > 300: - return _error("timestamp_skew", 401) + return ("timestamp_skew", 401) if fp != device.get("machine_pubkey_fingerprint"): - return _error("fingerprint_mismatch", 401) + return ("fingerprint_mismatch", 401) # Defer cryptography import to keep cold-start light when signatures # are disabled (the dev default). @@ -111,26 +151,24 @@ def _verify_signature(device: dict) -> Optional[tuple]: logger.error( "Signature verification requested but ``cryptography`` is not installed." ) - return _error("signature_unsupported", 500) + return ("signature_unsupported", 500) # The full public key isn't stored — only the fingerprint. For signature # verification we accept the pubkey in a header too; we trust it iff # its fingerprint matches the stored one. This avoids a separate pubkey # column for MVP. - pubkey_b64 = request.headers.get("X-Device-Machine-Pubkey") + pubkey_b64 = headers.get("X-Device-Machine-Pubkey") if not pubkey_b64: - return _error("missing_pubkey", 401) + return ("missing_pubkey", 401) if fingerprint_pubkey(pubkey_b64) != device["machine_pubkey_fingerprint"]: - return _error("pubkey_fingerprint_mismatch", 401) + return ("pubkey_fingerprint_mismatch", 401) - payload = _canonical_payload( - request.method, request.path, ts, request.get_data() - ).encode("utf-8") + payload = _canonical_payload(method, path, ts, get_body()).encode("utf-8") try: pubkey = Ed25519PublicKey.from_public_bytes(base64.b64decode(pubkey_b64)) pubkey.verify(base64.b64decode(sig_b64), payload) except (InvalidSignature, ValueError): - return _error("invalid_signature", 401) + return ("invalid_signature", 401) return None diff --git a/docsgpt/api/devices/routes.py b/docsgpt/api/devices/routes.py index ba4569c1..0ea2d6d7 100644 --- a/docsgpt/api/devices/routes.py +++ b/docsgpt/api/devices/routes.py @@ -16,7 +16,6 @@ from docsgpt.api.devices.session import ( ack_invocation, me, poll, - session_events, submit_output, ) from docsgpt.devices.normalizer import normalize_command @@ -313,10 +312,7 @@ def register(bp: Blueprint) -> None: bp.add_url_rule( "/api/devices/me", view_func=me, methods=["GET"], endpoint="device_me", ) - bp.add_url_rule( - "/api/devices/sessions//events", view_func=session_events, - methods=["GET"], endpoint="device_session_events", - ) + # The session's SSE stream is native-async (session_events.py, mounted in asgi.py). bp.add_url_rule( "/api/devices/sessions//invocations//ack", view_func=ack_invocation, methods=["POST"], endpoint="device_ack", diff --git a/docsgpt/api/devices/session.py b/docsgpt/api/devices/session.py index 4343c965..36482e23 100644 --- a/docsgpt/api/devices/session.py +++ b/docsgpt/api/devices/session.py @@ -1,4 +1,7 @@ -"""Device session endpoints: poll, SSE, ack, output.""" +"""Device session endpoints: poll, me, ack, output. + +The SSE command stream itself is a native-async route (``session_events.py``). +""" from __future__ import annotations @@ -6,14 +9,10 @@ import gzip import io import json import logging -import time -from typing import Iterator -from flask import Response, jsonify, make_response, request, stream_with_context +from flask import Response, jsonify, make_response, request from docsgpt.api.devices.auth import verify_device_session -from docsgpt.core.settings import settings -from docsgpt.core.shutdown import is_shutting_down from docsgpt.devices.broker import get_broker from docsgpt.storage.db.repositories.device_audit_log import ( DeviceAuditLogRepository, @@ -82,71 +81,6 @@ def me() -> Response: return make_response(jsonify(out), 200) -def session_events(session_id: str) -> Response: - """``GET /api/devices/sessions/{id}/events`` — SSE invocation stream. - - The ``session_id`` must be the ``session_ticket`` the device's own - ``/poll`` just issued (the path it was handed as ``session_url``). A - stale, mismatched, or fabricated ticket is rejected with ``410 Gone`` - before any stream is opened. - """ - device, err = verify_device_session() - if err is not None: - return err - broker = get_broker() - if not broker.validate_ticket(device["id"], session_id): - return make_response( - jsonify({"success": False, "error": "session_ticket_invalid"}), 410 - ) - sess = broker.register_session(device["id"], device["user_id"]) - - keepalive_interval = float(settings.SSE_KEEPALIVE_SECONDS) - idle_seconds = float(settings.REMOTE_DEVICE_SESSION_IDLE_SECONDS) - - @stream_with_context - def generate() -> Iterator[str]: - try: - last_keepalive = time.time() - while not sess.closed.is_set(): - # Break promptly on shutdown (see docsgpt/core/shutdown.py). - if is_shutting_down(): - break - now = time.time() - if now - 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 = broker.next_command(sess, timeout=1.0) - 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() - except GeneratorExit: - logger.debug("device SSE generator exiting for session %s", sess.session_id) - raise - finally: - broker.close_session(sess.session_id, reason="generator_exit") - - return Response( - generate(), - mimetype="text/event-stream", - headers={ - "Cache-Control": "no-cache", - "X-Accel-Buffering": "no", - }, - ) - - def ack_invocation(session_id: str, invocation_id: str) -> Response: """CLI acks (accepted / denied / auto_approved) the invocation.""" device, err = verify_device_session() @@ -262,11 +196,3 @@ def _as_opt_int(value) -> int | None: return int(value) except (TypeError, ValueError): return None - - -def _sse_event(name: str, payload: dict, event_id: int) -> str: - return ( - f"event: {name}\n" - f"id: {event_id}\n" - f"data: {json.dumps(payload)}\n\n" - ) diff --git a/docsgpt/api/devices/session_events.py b/docsgpt/api/devices/session_events.py new file mode 100644 index 00000000..ffd9231c --- /dev/null +++ b/docsgpt/api/devices/session_events.py @@ -0,0 +1,133 @@ +"""``GET /api/devices/sessions/{session_id}/events`` — a paired device's command stream. + +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 +commands arrive through the broker's async ``BLPOP`` and an idle session holds +no thread. +""" + +from __future__ import annotations + +import json +import logging +import time +from functools import partial +from typing import AsyncIterator + +import anyio +from starlette.requests import Request +from starlette.responses import JSONResponse, Response +from starlette.routing import Route + +from docsgpt.api.asgi_stream import ClosingStreamingResponse +from docsgpt.api.devices.auth import authenticate_device +from docsgpt.core.settings import settings +from docsgpt.core.shutdown import is_shutting_down +from docsgpt.devices.broker import DeviceBroker, get_broker + +logger = logging.getLogger(__name__) + +# Upper bound on one broker poll, so a closed session, idle expiry or a +# server drain is noticed within this many seconds. +_POLL_TIMEOUT_SECONDS = 1.0 + +_SSE_HEADERS = { + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", +} + + +def _error(code: str, status_code: int) -> JSONResponse: + return JSONResponse({"success": False, "error": code}, status_code=status_code) + + +def _sse_event(name: str, payload: dict, event_id: int) -> str: + return ( + f"event: {name}\n" + f"id: {event_id}\n" + f"data: {json.dumps(payload)}\n\n" + ) + + +async def _session_stream(broker: DeviceBroker, device: dict) -> 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 + 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") + + +async def device_session_events(request: Request) -> Response: + """Stream queued invocations to the device holding the poll-issued ticket. + + The ``session_id`` must be the ``session_ticket`` the device's own + ``/poll`` just issued (the path it was handed as ``session_url``). A + stale, mismatched, or fabricated ticket is rejected with ``410 Gone`` + before any stream is opened. + """ + session_id = request.path_params["session_id"] + body = await request.body() + device, failure = await anyio.to_thread.run_sync( + partial( + authenticate_device, + request.headers, + request.method, + request.url.path, + lambda: body, + ) + ) + if failure is not None: + return _error(*failure) + broker = get_broker() + if not await anyio.to_thread.run_sync(broker.validate_ticket, device["id"], session_id): + return _error("session_ticket_invalid", 410) + return ClosingStreamingResponse( + _session_stream(broker, device), + media_type="text/event-stream", + headers=_SSE_HEADERS, + ) + + +# Mounted in ``docsgpt/asgi.py`` ahead of the Flask catch-all. The device's +# poll, ack and output endpoints stay on Flask (``session.py``). +device_session_routes = [ + Route( + "/api/devices/sessions/{session_id}/events", + device_session_events, + methods=["GET"], + ), +] diff --git a/docsgpt/api/events/routes.py b/docsgpt/api/events/routes.py index 7b45926b..e4d634cd 100644 --- a/docsgpt/api/events/routes.py +++ b/docsgpt/api/events/routes.py @@ -4,6 +4,10 @@ Subscribe-then-snapshot pattern: subscribe to ``user:{user_id}`` pub/sub, snapshot the Redis Streams backlog past ``Last-Event-ID`` inside the SUBSCRIBE-ack callback, flush snapshot, then tail live events (dedup'd by stream id). See ``docs/runbooks/sse-notifications.md``. + +A native-async Starlette route mounted ahead of Flask in ``docsgpt/asgi.py``. +Every open browser tab holds this stream, so it runs on the event loop with +the async Redis client instead of pinning a WSGI threadpool slot per tab. """ from __future__ import annotations @@ -12,11 +16,17 @@ import json import logging import re import time -from typing import Iterator, Optional +from contextlib import aclosing +from functools import partial +from typing import AsyncIterator, Optional -from flask import Blueprint, Response, jsonify, make_response, request, stream_with_context +from starlette.requests import Request +from starlette.responses import Response +from starlette.routing import Route -from docsgpt.cache import get_redis_instance +from docsgpt.api.asgi_auth import authenticate, bind_log_context, json_error +from docsgpt.api.asgi_stream import ClosingStreamingResponse +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 ( @@ -26,14 +36,17 @@ from docsgpt.events.keys import ( stream_key, topic_name, ) -from docsgpt.streaming.broadcast_channel import Topic +from docsgpt.streaming.async_broadcast_channel import AsyncTopic +from docsgpt.streaming.async_redis import get_async_redis_instance logger = logging.getLogger(__name__) -events = Blueprint("event_stream", __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 + # WHATWG SSE treats CRLF, CR, and LF equivalently as line terminators. _SSE_LINE_SPLIT = re.compile(r"\r\n|\r|\n") @@ -43,6 +56,14 @@ _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 @@ -74,19 +95,19 @@ def _decode(value) -> Optional[str]: return str(value) -def _oldest_retained_id(redis_client, user_id: str) -> Optional[str]: +async def _oldest_retained_id(redis_client, 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 MAXLEN'd window. """ try: - info = redis_client.xinfo_stream(stream_key(user_id)) + info = await redis_client.xinfo_stream(stream_key(user_id)) except Exception: return None if not isinstance(info, dict): return None - # redis-py 7.4 returns str-keyed dicts here; the bytes-key probe is + # redis-py returns str-keyed dicts here; the bytes-key probe is # defence in depth in case ``decode_responses`` is ever flipped. first_entry = info.get("first-entry") or info.get(b"first-entry") if not first_entry: @@ -98,7 +119,7 @@ def _oldest_retained_id(redis_client, user_id: str) -> Optional[str]: return None -def _allow_replay( +async def _allow_replay( redis_client, user_id: str, last_event_id: Optional[str] ) -> bool: """Per-user sliding-window snapshot-replay budget. @@ -119,12 +140,12 @@ def _allow_replay( window = max(1, int(settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS)) key = replay_budget_key(user_id) try: - used = int(redis_client.incr(key)) + 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. - redis_client.expire(key, window) + await redis_client.expire(key, window) except Exception: logger.debug( "replay budget probe failed for user=%s; failing open", @@ -166,10 +187,10 @@ def _replay_floor_id() -> Optional[str]: return f"{floor_ms}-0" -def _replay_backlog( +async def _replay_backlog( redis_client, user_id: str, last_event_id: Optional[str], max_count: int -) -> Iterator[tuple[str, str]]: - """Yield ``(entry_id, sse_line)`` for backlog entries past ``last_event_id``. +) -> list[tuple[str, str]]: + """Return ``(entry_id, sse_line)`` pairs for backlog entries past ``last_event_id``. Capped at ``max_count`` rows; clients catch up across reconnects. Parse failures are skipped; the Streams id is injected into the @@ -186,7 +207,7 @@ def _replay_backlog( # Exclusive start: '(' skips the already-delivered entry. start = f"({last_event_id}" try: - entries = redis_client.xrange( + entries = await redis_client.xrange( stream_key(user_id), min=start, max="+", count=max_count ) except Exception as exc: @@ -196,14 +217,15 @@ def _replay_backlog( last_event_id or "-", exc, ) - return + return [] + lines: list[tuple[str, str]] = [] for entry_id, fields in entries: entry_id_str = _decode(entry_id) if not entry_id_str: continue - # decode_responses=False on the cache client ⇒ field keys/values - # are bytes. The string-key fallback covers a future flip of that + # decode_responses=False on the client ⇒ field keys/values are + # bytes. The string-key fallback covers a future flip of that # default without a forced refactor here. raw_event = None if isinstance(fields, dict): @@ -223,7 +245,8 @@ def _replay_backlog( "Replay envelope parse failed for entry %s; passing through raw", entry_id_str, ) - yield entry_id_str, _format_sse(event_str, event_id=entry_id_str) + lines.append((entry_id_str, _format_sse(event_str, event_id=entry_id_str))) + return lines def _truncation_notice_line(oldest_id: str) -> str: @@ -238,209 +261,102 @@ def _truncation_notice_line(oldest_id: str) -> str: ) -@events.route("/api/events", methods=["GET"]) -def stream_events() -> Response: - decoded = getattr(request, "decoded_token", None) - user_id = decoded.get("sub") if isinstance(decoded, dict) else None - if not user_id: - return make_response( - jsonify({"success": False, "message": "Authentication required"}), - 401, - ) +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) - # In dev deployments without AUTH_TYPE configured, every request - # resolves to user_id="local" and shares one stream. Surface this so - # an accidentally-multi-user dev box doesn't silently cross-stream. - global _local_user_warned - if user_id == "local" and not _local_user_warned: - logger.warning( - "SSE serving user_id='local' (AUTH_TYPE not set). " - "All clients on this deployment will share one event stream." - ) - _local_user_warned = True - - raw_last_event_id = request.headers.get("Last-Event-ID") or request.args.get( - "last_event_id" - ) - last_event_id = _normalize_last_event_id(raw_last_event_id) - last_event_id_invalid = raw_last_event_id is not None and last_event_id is None +async def _event_stream( + redis_client, + user_id: str, + last_event_id: Optional[str], + last_event_id_invalid: bool, + push_enabled: bool, +) -> AsyncIterator[str]: + """Yield the prelude, the replayed backlog, then live events until the stream ends.""" + connect_ts = time.monotonic() + replayed_count = 0 keepalive_seconds = float(settings.SSE_KEEPALIVE_SECONDS) - push_enabled = settings.ENABLE_SSE_PUSH - cap = int(settings.SSE_MAX_CONCURRENT_PER_USER) + try: + # First frame primes intermediaries (Cloudflare, nginx) so they + # don't sit on a buffer waiting for body bytes. + yield ": connected\n\n" - redis_client = get_redis_instance() - counter_key = connection_counter_key(user_id) - counted = False + if not push_enabled: + yield ": push_disabled\n\n" + return - if push_enabled and redis_client is not None and cap > 0: - try: - current = int(redis_client.incr(counter_key)) - counted = True - except Exception: - current = 0 - logger.debug( - "SSE connection counter INCR failed for user=%s", user_id - ) - if counted: - # 1h safety TTL — orphaned counts from hard crashes self-heal. - # EXPIRE failure must NOT clobber ``current`` and bypass the cap. + replay_lines: list[str] = [] + max_replayed_id: Optional[str] = None + replay_done = False + + # If the client sent a malformed Last-Event-ID, surface the + # truncation notice before the subscribe loop. Buffering it into + # ``replay_lines`` would lose it when ``subscribe`` returns + # immediately (Redis down) and the loop body never runs. + if last_event_id_invalid: + yield _truncation_notice_line("") + replayed_count += 1 + + async def _on_subscribe() -> None: + # Runs inside AsyncTopic.subscribe after the SUBSCRIBE is + # acked. Reading the backlog here means any publisher firing + # between SUBSCRIBE-send and SUBSCRIBE-ack has its XADD captured + # by XRANGE *and* its PUBLISH buffered at the connection layer + # until we read it — closing the replay/subscribe race. + # + # Truncation contract: ``backlog.truncated`` is emitted ONLY + # when the client's ``Last-Event-ID`` has slid off the MAXLEN'd + # window or the age floor — the journal is genuinely gone past + # the cursor, so the frontend should clear its cursor and + # refetch state. Cap-hit skips silently: the cursor advances via + # the per-entry ``id:`` headers. Budget exhaustion never reaches + # this callback — the route 429s before opening the stream. + nonlocal max_replayed_id, replay_done try: - redis_client.expire(counter_key, 3600) - except Exception: - logger.debug( - "SSE connection counter EXPIRE failed for user=%s", user_id - ) - if current > cap: - try: - redis_client.decr(counter_key) - except Exception: - logger.debug( - "SSE connection counter DECR failed for user=%s", - user_id, - ) - return make_response( - jsonify( - { - "success": False, - "message": "Too many concurrent SSE connections", - } - ), - 429, - ) + if redis_client is None or last_event_id is None: + # Fresh session: start live. Replaying the whole + # retained window shipped weeks-old entries on every + # tab-open (MAXLEN caps by count, not age). + return + oldest = await _oldest_retained_id(redis_client, user_id) + floor = _replay_floor_id() + # The snapshot can't reach past whichever is newer: the + # MAXLEN'd window edge or the age floor. + effective_oldest = oldest + if floor and ( + effective_oldest is None + or stream_id_compare(floor, effective_oldest) > 0 + ): + effective_oldest = floor + if ( + effective_oldest + and stream_id_compare(last_event_id, effective_oldest) < 0 + ): + replay_lines.append(_truncation_notice_line(effective_oldest)) + replay_cap = int(settings.EVENTS_REPLAY_MAX_PER_REQUEST) + for entry_id, sse_line in await _replay_backlog( + redis_client, user_id, last_event_id, replay_cap + ): + replay_lines.append(sse_line) + max_replayed_id = entry_id + finally: + # Always flip the flag — even on partial-replay failure the + # outer loop must reach the flush step so we don't silently + # strand whatever entries did land. + replay_done = True - # Replay budget is checked here, before the generator opens the - # stream, so a denial can surface as HTTP 429 instead of a silent - # snapshot skip. The earlier in-generator skip lost events between - # the client's cursor and the first live-tailed entry: the live - # tail still carried ``id:`` headers, the frontend advanced - # ``lastEventId`` to one of those ids, and the events in between - # were never reachable on the next reconnect. 429 keeps the - # cursor pinned and lets the frontend back off until the window - # slides (eventStreamClient.ts treats 429 as escalated backoff). - if push_enabled and redis_client is not None and not _allow_replay( - redis_client, user_id, last_event_id - ): - if counted: - try: - redis_client.decr(counter_key) - except Exception: - logger.debug( - "SSE connection counter DECR failed for user=%s", - user_id, - ) - return make_response( - jsonify( - { - "success": False, - "message": "Replay budget exhausted", - } - ), - 429, + topic = AsyncTopic(topic_name(user_id)) + last_keepalive = time.monotonic() + subscription = topic.subscribe( + on_subscribe=_on_subscribe, + poll_timeout=SUBSCRIBE_POLL_INTERVAL_SECONDS, + liveness_timeout=PUBSUB_SOCKET_TIMEOUT_SECONDS, ) - - @stream_with_context - def generate() -> Iterator[str]: - connect_ts = time.monotonic() - replayed_count = 0 - try: - # First frame primes intermediaries (Cloudflare, nginx) so they - # don't sit on a buffer waiting for body bytes. - yield ": connected\n\n" - - if not push_enabled: - yield ": push_disabled\n\n" - return - - replay_lines: list[str] = [] - max_replayed_id: Optional[str] = None - replay_done = False - - # If the client sent a malformed Last-Event-ID, surface the - # truncation notice synchronously *before* the subscribe - # loop. Buffering it into ``replay_lines`` would lose it - # when ``Topic.subscribe`` returns immediately (Redis down) - # — the loop body never runs, and the flush at line ~335 - # never fires. - if last_event_id_invalid: - yield _truncation_notice_line("") - replayed_count += 1 - - def _on_subscribe_callback() -> None: - # Runs synchronously inside Topic.subscribe after the - # SUBSCRIBE is acked. By doing XRANGE here, any publisher - # firing between SUBSCRIBE-send and SUBSCRIBE-ack has its - # XADD captured by XRANGE *and* its PUBLISH buffered at - # the connection layer until we read it — closing the - # replay/subscribe race the design doc warns about. - # - # Truncation contract: ``backlog.truncated`` is emitted - # ONLY when the client's ``Last-Event-ID`` has slid off - # the MAXLEN'd window — that's the case where the - # journal is genuinely gone past the cursor and the - # frontend should clear its slice cursor and refetch - # state. Cap-hit skips the snapshot silently: the - # cursor advances via the per-entry ``id:`` headers - # and the frontend's slice keeps the latest id so the - # next reconnect resumes from there. Budget-exhausted - # never reaches this callback — the route 429s before - # opening the stream, keeping the cursor pinned. - # Conflating these with stale-cursor truncation would - # tell the client to clear its cursor and re-receive - # the same oldest-N entries on every reconnect — - # locking the user out of entries past N. - nonlocal max_replayed_id, replay_done - try: - if redis_client is None: - return - if last_event_id is None: - # Fresh session: start live. A no-cursor connect has - # no state to catch up on — replaying the whole - # retained window shipped weeks-old entries on every - # tab-open (MAXLEN caps by count, not age). Clients - # that had state present a cursor; the malformed- - # cursor case already got its truncation notice - # before the subscribe loop. - return - oldest = _oldest_retained_id(redis_client, user_id) - floor = _replay_floor_id() - # The snapshot can't reach past whichever is newer: - # the MAXLEN'd window edge or the age floor. - effective_oldest = oldest - if floor and ( - effective_oldest is None - or stream_id_compare(floor, effective_oldest) > 0 - ): - effective_oldest = floor - if ( - effective_oldest - and stream_id_compare(last_event_id, effective_oldest) < 0 - ): - # The Last-Event-ID has slid off the replayable - # window. Tell the client so it can fetch full state. - replay_lines.append( - _truncation_notice_line(effective_oldest) - ) - replay_cap = int(settings.EVENTS_REPLAY_MAX_PER_REQUEST) - for entry_id, sse_line in _replay_backlog( - redis_client, user_id, last_event_id, replay_cap - ): - replay_lines.append(sse_line) - max_replayed_id = entry_id - finally: - # Always flip the flag — even on partial-replay failure - # the outer loop must reach the flush step so we don't - # silently strand whatever entries did land. - replay_done = True - - topic = Topic(topic_name(user_id)) - last_keepalive = time.monotonic() - for payload in topic.subscribe( - on_subscribe=_on_subscribe_callback, - poll_timeout=SUBSCRIBE_POLL_INTERVAL_SECONDS, - ): - # Break promptly on shutdown — this a2wsgi thread can't be - # cancelled by asyncio (see docsgpt/core/shutdown.py). + async with aclosing(subscription): + async for payload in subscription: if is_shutting_down(): break @@ -466,14 +382,10 @@ def stream_events() -> Response: if isinstance(envelope, dict): candidate = envelope.get("id") # Only trust ids that look like real Redis Streams - # ids (``ms`` or ``ms-seq``). A malformed or - # adversarial publisher could otherwise pin - # dedupe forever — a lex-greater bogus id would - # make every legitimate later id compare ``<=`` - # and get dropped silently. - if isinstance(candidate, str) and _STREAM_ID_RE.match( - candidate - ): + # ids. A malformed or adversarial publisher could + # otherwise pin dedupe forever — a lex-greater bogus + # id would make every later id compare ``<=``. + if isinstance(candidate, str) and _STREAM_ID_RE.match(candidate): event_id = candidate except Exception: pass @@ -489,49 +401,102 @@ def stream_events() -> Response: yield _format_sse(event_str, event_id=event_id) last_keepalive = now - # Topic.subscribe exited before the first yield (transient - # Redis hiccup between SUBSCRIBE-ack and the first poll, or - # an immediate Redis-down return). The callback may already - # have populated the snapshot — flush it so the client gets - # the backlog instead of a silent drop. Safe no-op when the - # in-loop flush ran (it clear()'d the buffer) and when the - # callback never fired (replay_done stays False). - if replay_done and replay_lines: - for line in replay_lines: - yield line - replayed_count += 1 - replay_lines.clear() - except GeneratorExit: - return + # The subscription exited before the first yield (transient Redis + # hiccup between SUBSCRIBE-ack and the first poll, or an immediate + # Redis-down return). The callback may already have populated the + # snapshot — flush it so the client gets the backlog. + if replay_done and replay_lines: + for line in replay_lines: + yield line + replayed_count += 1 + replay_lines.clear() + except Exception: + # Client disconnect (cancellation / GeneratorExit) is a BaseException + # and passes straight through; only genuine bugs land here. + logger.exception("SSE event-stream generator crashed for user=%s", user_id) + finally: + logger.info( + "event.disconnect user=%s duration_s=%.1f replayed=%d", + user_id, + time.monotonic() - connect_ts, + replayed_count, + ) + + +async def stream_events(request: Request) -> Response: + """``GET /api/events`` — the caller's live notification stream.""" + decoded, error = await authenticate(request) + if error is not None: + return error + 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) + + # In dev deployments without AUTH_TYPE configured, every request + # resolves to user_id="local" and shares one stream. Surface this so + # an accidentally-multi-user dev box doesn't silently cross-stream. + global _local_user_warned + if user_id == "local" and not _local_user_warned: + logger.warning( + "SSE serving user_id='local' (AUTH_TYPE not set). " + "All clients on this deployment will share one event stream." + ) + _local_user_warned = True + + raw_last_event_id = request.headers.get("Last-Event-ID") or request.query_params.get( + "last_event_id" + ) + last_event_id = _normalize_last_event_id(raw_last_event_id) + 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.exception( - "SSE event-stream generator crashed for user=%s", user_id - ) - finally: - duration_s = time.monotonic() - connect_ts - logger.info( - "event.disconnect user=%s duration_s=%.1f replayed=%d", - user_id, - duration_s, - replayed_count, - ) - if counted and redis_client is not None: - try: - redis_client.decr(counter_key) - except Exception: - logger.debug( - "SSE connection counter DECR failed for user=%s on disconnect", - user_id, - ) + 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) + + # 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) + return json_error("Replay budget exhausted", 429) - response = Response(generate(), mimetype="text/event-stream") - response.headers["Cache-Control"] = "no-store" - response.headers["X-Accel-Buffering"] = "no" - response.headers["Connection"] = "keep-alive" logger.info( "event.connect user=%s last_event_id=%s%s", user_id, last_event_id or "-", " (rejected_invalid)" if last_event_id_invalid else "", ) - return response + return ClosingStreamingResponse( + _event_stream(redis_client, user_id, last_event_id, last_event_id_invalid, push_enabled), + 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, + ) + + +# Mounted in ``docsgpt/asgi.py`` ahead of the Flask catch-all. +event_stream_routes = [ + Route("/api/events", stream_events, methods=["GET"]), +] diff --git a/docsgpt/api/user/artifacts/authz.py b/docsgpt/api/user/artifacts/authz.py index 429df071..5c778f6c 100644 --- a/docsgpt/api/user/artifacts/authz.py +++ b/docsgpt/api/user/artifacts/authz.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Optional +from typing import Mapping, Optional from flask import request @@ -40,18 +40,16 @@ class Principal: return self.agent_id is not None -def resolve_principal() -> Principal: - """Resolve the caller to a :class:`Principal`. +def principal_for(decoded_token: Optional[dict], api_key: Optional[str]) -> Principal: + """Resolve a caller from its decoded JWT or ``api_key``, independent of the web framework. - Priority: a decoded JWT (owner session) first, then an ``api_key`` query/form - param (agent-scoped). An unresolvable caller yields an anonymous principal - (both fields ``None``) that only a valid ``share_token`` can authorize. + Priority: a decoded JWT (owner session) first, then the ``api_key`` + (agent-scoped). An unresolvable caller yields an anonymous principal (both + fields ``None``) that only a valid ``share_token`` can authorize. """ - decoded_token = getattr(request, "decoded_token", None) if decoded_token: return Principal(user_id=decoded_token.get("sub")) - api_key = request.args.get("api_key") or request.form.get("api_key") if api_key: with db_readonly() as conn: agent = AgentsRepository(conn).find_by_key(api_key) @@ -60,6 +58,18 @@ def resolve_principal() -> Principal: return Principal() +def resolve_principal() -> Principal: + """Resolve the current Flask request's caller to a :class:`Principal`. + + Reads the decoded JWT, then an ``api_key`` query/form param; see + :func:`principal_for`. + """ + return principal_for( + getattr(request, "decoded_token", None), + request.args.get("api_key") or request.form.get("api_key"), + ) + + def _shared_row_for(conn, conversation_id, share_token): """Return the shared_conversations row iff share_token grants this conversation, else None.""" if not share_token: @@ -87,7 +97,12 @@ def user_can_access_conversation( return _shared_row_for(conn, conversation_id, share_token) is not None -def authorize_artifact(conn, artifact: dict, principal: Principal) -> bool: +def authorize_artifact( + conn, + artifact: dict, + principal: Principal, + args: Optional[Mapping[str, str]] = None, +) -> bool: """Authorize a READ of ``artifact`` for ``principal``; missing parent fails closed. A low-trust agent api_key is confined to a single conversation it proves by @@ -96,7 +111,13 @@ def authorize_artifact(conn, artifact: dict, principal: Principal) -> bool: JWT owner or ``shared_with`` collaborator gets full access to every artifact of the parent; a share-token holder is confined to the shared ``first_n_queries`` snapshot (an artifact whose ``message_id`` is outside it, or NULL, is denied). + + ``args`` are the request's query params; they default to the current Flask + request's, and a native-async route passes its own. """ + if args is None: + args = request.args + if principal.is_agent_scoped: # An agent key is not the owner's session and is embedded in public widget # JS, so require the artifact to belong to the owner AND the request to @@ -104,7 +125,7 @@ def authorize_artifact(conn, artifact: dict, principal: Principal) -> bool: # the key alone, without that id, cannot download a known artifact. if str(artifact.get("user_id")) != str(principal.user_id): return False - req_conv = request.args.get("conversation_id") + req_conv = args.get("conversation_id") parent = artifact.get("conversation_id") if not req_conv or parent is None or str(parent) != str(req_conv): return False @@ -114,7 +135,7 @@ def authorize_artifact(conn, artifact: dict, principal: Principal) -> bool: conversation_id = artifact.get("conversation_id") workflow_run_id = artifact.get("workflow_run_id") - share_token = request.args.get("share_token") + share_token = args.get("share_token") if conversation_id is not None: # Owner or shared_with collaborator: full access to every artifact. diff --git a/docsgpt/api/user/artifacts/download.py b/docsgpt/api/user/artifacts/download.py new file mode 100644 index 00000000..4547c298 --- /dev/null +++ b/docsgpt/api/user/artifacts/download.py @@ -0,0 +1,181 @@ +"""``GET /api/artifacts/{artifact_id}/download`` — an artifact version's bytes. + +A native-async Starlette route mounted ahead of Flask in ``docsgpt/asgi.py``, +so a large download to a slow client streams from the event loop instead of +pinning a WSGI threadpool slot for the whole transfer. Authorization is the +same parent-derived check the Flask artifact routes use; the database reads +and storage calls run in worker threads. +""" + +from __future__ import annotations + +import io +import logging +import re +from functools import partial +from typing import AsyncIterator, BinaryIO, Mapping, Optional, Tuple + +import anyio +from starlette.requests import Request +from starlette.responses import JSONResponse, RedirectResponse, 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.user.artifacts.authz import authorize_artifact, principal_for +from docsgpt.core.settings import settings +from docsgpt.storage.db.base_repository import looks_like_uuid +from docsgpt.storage.db.repositories.artifacts import ArtifactsRepository +from docsgpt.storage.db.session import db_readonly +from docsgpt.storage.storage_creator import StorageCreator + +logger = logging.getLogger(__name__) + +# Presigned-URL TTL for private S3 artifact downloads (seconds). +_PRESIGNED_URL_TTL = 300 + +_ARTIFACT_URL_ENVELOPE_MIME = "application/vnd.docsgpt.artifact-url+json" + +_CHUNK_SIZE = 65536 + + +def _sanitize_header_filename(filename: Optional[str], fallback: str) -> str: + """Strip CRLF / quotes from a display filename for a Content-Disposition header.""" + if not filename: + return fallback + cleaned = re.sub(r'[\r\n"]', "", str(filename)).strip() + return cleaned or fallback + + +def _load_version( + artifact_id: str, decoded_token: Optional[dict], query: Mapping[str, str] +) -> Tuple[Optional[dict], Optional[Tuple[str, int]]]: + """Authorize the caller and fetch the requested version row; runs in a worker thread. + + Returns: + tuple: ``(version_row, None)`` or ``(None, (message, status))``. + """ + principal = principal_for(decoded_token, query.get("api_key")) + with db_readonly() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.get_artifact(artifact_id) + if artifact is None: + return None, ("Artifact not found", 404) + if not authorize_artifact(conn, artifact, principal, args=query): + return None, ("Forbidden", 403) + version = artifact.get("current_version") + version_arg = query.get("version") + if version_arg is not None: + try: + version = int(version_arg) + except ValueError: + return None, ("Invalid version", 400) + version_row = repo.get_version(artifact_id, version) + if version_row is None: + return None, ("Version not found", 404) + return version_row, None + + +async def _read_chunks(file_obj: BinaryIO) -> AsyncIterator[bytes]: + """Yield the file in chunks; reads from disk or a socket hop to a worker thread.""" + in_memory = isinstance(file_obj, io.BytesIO) + while True: + if in_memory: + chunk = file_obj.read(_CHUNK_SIZE) + else: + chunk = await anyio.to_thread.run_sync(file_obj.read, _CHUNK_SIZE) + if not chunk: + return + yield chunk + + +async def _close_file(file_obj: BinaryIO) -> None: + close = getattr(file_obj, "close", None) + if callable(close): + close() + + +async def download_artifact(request: Request) -> Response: + """Serve an artifact version: a 302 (or a JSON URL envelope) on S3, streamed bytes otherwise.""" + artifact_id = request.path_params["artifact_id"] + if not looks_like_uuid(artifact_id): + return json_error("Artifact not found", 404) + decoded, error = await authenticate(request) + if error is not None: + return error + bind_log_context("artifact_download", decoded.get("sub") if decoded else None) + query = request.query_params + + try: + version_row, failure = await anyio.to_thread.run_sync( + _load_version, artifact_id, decoded, query + ) + if failure is not None: + return json_error(*failure) + + # The object key is derived only from the stored path, never client input. + storage_path = version_row.get("storage_path") + 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) + + # With URL_STRATEGY=="s3" the contract is to hand back a presigned + # URL. If the active backend can't mint one, that's a config error: + # surface a 500 rather than silently proxying bytes from a backend + # the operator expected to be off the hot path. + if getattr(settings, "URL_STRATEGY", "backend") == "s3": + try: + url = await anyio.to_thread.run_sync( + partial(storage.generate_presigned_url, storage_path, expires_in=_PRESIGNED_URL_TTL) + ) + except NotImplementedError: + logger.error( + "URL_STRATEGY=s3 but %s cannot mint presigned URLs", + type(storage).__name__, + ) + return json_error("Storage misconfigured", 500) + # A 302 to a cross-origin S3 URL can't be read by the app's authed + # fetch (the bucket has no CORS grant for the app origin). When the + # client opts in via ?disposition=url (or Accept: application/json) + # hand the presigned URL back as JSON so it can navigate to it + # top-level (no CORS). Default stays a 302 so nothing else breaks. + if query.get("disposition") == "url" or ( + "application/json" in request.headers.get("Accept", "") + ): + # Tag the envelope with a distinctive media type so the client + # keys off a server-set signal, not the JSON body shape. A + # ``data`` artifact whose bytes happen to be + # ``{"success":true,"url":"..."}`` is streamed under the backend + # strategy with its own content-type and can never carry this + # vendor type, so it can't be mistaken for a redirect + # (open-redirect gadget). Content-Type is CORS-safelisted. + return JSONResponse( + {"success": True, "url": url}, media_type=_ARTIFACT_URL_ENVELOPE_MIME + ) + return RedirectResponse(url, status_code=302) + + # 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. + file_obj = await anyio.to_thread.run_sync(storage.get_file, storage_path) + return ClosingStreamingResponse( + _read_chunks(file_obj), + media_type=mime_type, + headers={"Content-Disposition": f'attachment; filename="{filename}"'}, + on_close=partial(_close_file, file_obj), + ) + except FileNotFoundError: + return json_error("File not found", 404) + except Exception: + logger.error("Error downloading artifact %s", artifact_id, exc_info=True) + return JSONResponse({"success": False}, status_code=400) + + +# Mounted in ``docsgpt/asgi.py`` ahead of the Flask catch-all. The other +# artifact routes stay on Flask (``routes.py``). +artifact_download_routes = [ + Route("/api/artifacts/{artifact_id}/download", download_artifact, methods=["GET"]), +] diff --git a/docsgpt/api/user/artifacts/routes.py b/docsgpt/api/user/artifacts/routes.py index e426288a..9ab0b501 100644 --- a/docsgpt/api/user/artifacts/routes.py +++ b/docsgpt/api/user/artifacts/routes.py @@ -1,18 +1,15 @@ -"""Artifact metadata and download routes (parent-derived authz).""" +"""Artifact metadata, restore and delete routes (parent-derived authz). + +Downloads are served by the native-async route in ``download.py``. +""" from __future__ import annotations -import re -from typing import Optional - from flask import ( - Response, current_app, jsonify, make_response, - redirect, request, - stream_with_context, ) from flask_restx import Namespace, Resource @@ -24,7 +21,6 @@ from docsgpt.api.user.artifacts.authz import ( resolve_principal, user_can_access_conversation, ) -from docsgpt.core.settings import settings from docsgpt.storage.db.base_repository import looks_like_uuid from docsgpt.storage.db.repositories.artifacts import ArtifactsRepository from docsgpt.storage.db.repositories.conversations import ConversationsRepository @@ -34,20 +30,6 @@ from docsgpt.storage.storage_creator import StorageCreator artifacts_ns = Namespace("artifacts", description="Artifact operations", path="/api") -# Presigned-URL TTL for private S3 artifact downloads (seconds). -_PRESIGNED_URL_TTL = 300 - - -_ARTIFACT_URL_ENVELOPE_MIME = "application/vnd.docsgpt.artifact-url+json" - - -def _sanitize_header_filename(filename: Optional[str], fallback: str) -> str: - """Strip CRLF / quotes from a display filename for a Content-Disposition header.""" - if not filename: - return fallback - cleaned = re.sub(r'[\r\n"]', "", str(filename)).strip() - return cleaned or fallback - def _artifact_summary(artifact: dict) -> dict: """Project an artifact identity row to its API metadata shape (owner id withheld).""" @@ -255,103 +237,6 @@ class GetArtifactVersion(Resource): return make_response(jsonify({"success": False}), 400) -@artifacts_ns.route("/artifacts//download") -class DownloadArtifact(Resource): - @api.doc(description="Download an artifact's bytes (302 to a presigned URL on S3)") - def get(self, artifact_id: str): - if not looks_like_uuid(artifact_id): - return make_response(jsonify({"success": False, "message": "Artifact not found"}), 404) - principal = resolve_principal() - version_arg = request.args.get("version") - try: - with db_readonly() as conn: - repo = ArtifactsRepository(conn) - artifact = repo.get_artifact(artifact_id) - if artifact is None: - return make_response(jsonify({"success": False, "message": "Artifact not found"}), 404) - if not authorize_artifact(conn, artifact, principal): - return make_response(jsonify({"success": False, "message": "Forbidden"}), 403) - version = artifact.get("current_version") - if version_arg is not None: - try: - version = int(version_arg) - except ValueError: - return make_response(jsonify({"success": False, "message": "Invalid version"}), 400) - version_row = repo.get_version(artifact_id, version) - - if version_row is None: - return make_response(jsonify({"success": False, "message": "Version not found"}), 404) - # The object key is derived only from the stored path, never client input. - storage_path = version_row.get("storage_path") - if not storage_path: - return make_response(jsonify({"success": False, "message": "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 = StorageCreator.get_storage() - - # With URL_STRATEGY=="s3" the contract is to hand back a presigned - # URL. If the active backend can't mint one, that's a config error: - # surface a 500 rather than silently proxying bytes from a backend - # the operator expected to be off the hot path. - if getattr(settings, "URL_STRATEGY", "backend") == "s3": - try: - url = storage.generate_presigned_url(storage_path, expires_in=_PRESIGNED_URL_TTL) - except NotImplementedError: - current_app.logger.error( - "URL_STRATEGY=s3 but %s cannot mint presigned URLs", - type(storage).__name__, - ) - return make_response( - jsonify({"success": False, "message": "Storage misconfigured"}), - 500, - ) - # A 302 to a cross-origin S3 URL can't be read by the app's authed - # fetch (the bucket has no CORS grant for the app origin). When the - # client opts in via ?disposition=url (or Accept: application/json) - # hand the presigned URL back as JSON so it can navigate to it - # top-level (no CORS). Default stays a 302 so nothing else breaks. - if request.args.get("disposition") == "url" or ( - "application/json" in request.headers.get("Accept", "") - ): - resp = make_response(jsonify({"success": True, "url": url}), 200) - # Tag the envelope with a distinctive media type so the client - # keys off a server-set signal, not the JSON body shape. A - # ``data`` artifact whose bytes happen to be - # ``{"success":true,"url":"..."}`` is streamed under the backend - # strategy with its own (``application/json``) content-type and - # can never carry this vendor type, so it can't be mistaken for a - # redirect (open-redirect gadget). Content-Type is CORS-safelisted, - # so this stays readable cross-origin without expose-headers. - resp.headers["Content-Type"] = _ARTIFACT_URL_ENVELOPE_MIME - return resp - return redirect(url, code=302) - - # Stream the bytes in chunks instead of buffering the whole object in - # worker memory (artifacts can be many MB); close the handle when done. - file_obj = storage.get_file(storage_path) - - def _stream(): - try: - for chunk in iter(lambda: file_obj.read(65536), b""): - yield chunk - finally: - close = getattr(file_obj, "close", None) - if callable(close): - close() - - return Response( - stream_with_context(_stream()), - mimetype=mime_type, - headers={"Content-Disposition": f'attachment; filename="{filename}"'}, - ) - except FileNotFoundError: - return make_response(jsonify({"success": False, "message": "File not found"}), 404) - except Exception as err: - current_app.logger.error(f"Error downloading artifact: {err}", exc_info=True) - return make_response(jsonify({"success": False}), 400) - - @artifacts_ns.route("/artifacts//restore") class RestoreArtifact(Resource): @api.doc(description="Restore a prior version by appending it as the new current version") diff --git a/docsgpt/app.py b/docsgpt/app.py index 81c303dc..db3dceee 100644 --- a/docsgpt/app.py +++ b/docsgpt/app.py @@ -20,7 +20,6 @@ from docsgpt.api import api # noqa: E402 from docsgpt.api.admin import admin_ns # noqa: E402 from docsgpt.api.answer import answer # noqa: E402 from docsgpt.api.devices import devices_bp # noqa: E402 -from docsgpt.api.events.routes import events # noqa: E402 from docsgpt.api.internal.routes import internal # noqa: E402 from docsgpt.api.oidc import oidc_bp # noqa: E402 from docsgpt.api.oidc.denylist import is_denied as oidc_session_denied # noqa: E402 @@ -99,7 +98,6 @@ validate_default_chat_tools() app = Flask(__name__) app.register_blueprint(user) app.register_blueprint(answer) -app.register_blueprint(events) app.register_blueprint(internal) app.register_blueprint(connector) app.register_blueprint(devices_bp) diff --git a/docsgpt/asgi.py b/docsgpt/asgi.py index e694fefc..4de1603c 100644 --- a/docsgpt/asgi.py +++ b/docsgpt/asgi.py @@ -9,6 +9,9 @@ from starlette.middleware.cors import CORSMiddleware from starlette.routing import Mount from docsgpt.api.async_sse import async_sse_routes +from docsgpt.api.devices.session_events import device_session_routes +from docsgpt.api.events.routes import event_stream_routes +from docsgpt.api.user.artifacts.download import artifact_download_routes from docsgpt.app import app as flask_app from docsgpt.core.settings import settings from docsgpt.mcp_server import mcp @@ -28,12 +31,17 @@ _backend = StaticUI.wrap( asgi_app = Starlette( routes=[ Mount("/mcp", app=mcp_app), - # Native-async SSE readers intercept their exact paths before the - # Flask catch-all, so a mostly-idle reconnect tail rides the event - # loop instead of pinning a WSGI threadpool slot. Order matters: - # Starlette matches routes top-to-bottom, so these must precede the - # Mount("/") that hands everything else to Flask. + # Native-async routes intercept their exact paths before the Flask + # catch-all. Each holds its response open for a long time (the chat + # reconnect tail, the notification stream, a device's command stream, + # large artifact downloads), so it rides the event loop instead of + # pinning a WSGI threadpool slot. Order matters: Starlette matches + # routes top-to-bottom, so these must precede the Mount("/") that + # hands everything else to Flask. *async_sse_routes, + *event_stream_routes, + *device_session_routes, + *artifact_download_routes, Mount("/", app=_backend), ], middleware=[ diff --git a/docsgpt/core/settings.py b/docsgpt/core/settings.py index 14c834b2..a04861b3 100644 --- a/docsgpt/core/settings.py +++ b/docsgpt/core/settings.py @@ -444,9 +444,14 @@ class Settings(BaseSettings): GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS: int = 30 WSGI_THREADPOOL_WORKERS: int = 96 SSE_KEEPALIVE_SECONDS: int = Field(default=15, ge=1) - # Simultaneous SSE connections per user; each holds a WSGI thread and a Redis pub/sub - # connection. 8 covers multi-tab use without one user starving the pool. 0 disables. + # Simultaneous SSE connections per user; each holds a pooled async Redis connection for + # its lifetime. 8 covers multi-tab use without one user starving the pool. 0 disables. SSE_MAX_CONCURRENT_PER_USER: int = 8 + # Pool size of the async Redis client behind the event-loop routes, per process. Every + # open notification tab, chat reconnect and device session holds one connection, so this + # caps concurrent streams per worker (redis-py's own default is 100). Keep the total + # across workers below the Redis server's maxclients (10000 by default). + ASYNC_REDIS_MAX_CONNECTIONS: int = Field(default=2000, ge=1) # Backlog entries XRANGE returns per /api/events snapshot. Bounds what one replay moves # from Redis to the wire: a client looping Last-Event-ID reconnects enumerates at most # this many per round-trip, and the budget below bounds total throughput. diff --git a/docsgpt/devices/broker.py b/docsgpt/devices/broker.py index e9eff343..781df24b 100644 --- a/docsgpt/devices/broker.py +++ b/docsgpt/devices/broker.py @@ -38,8 +38,11 @@ import uuid from dataclasses import dataclass, field from typing import Any, Dict, Iterator, Optional +import anyio + from docsgpt.cache import get_redis_instance from docsgpt.core.settings import settings +from docsgpt.streaming.async_redis import get_async_redis_instance logger = logging.getLogger(__name__) @@ -191,13 +194,8 @@ class DeviceBroker: return None if not popped: return None - _key, raw = popped - try: - envelope = json.loads(_as_str(raw)) - except (TypeError, ValueError): - logger.warning("dropping malformed command envelope") - return None - if not isinstance(envelope, dict): + envelope = self._decode_envelope(popped[1]) + if envelope is None: return None # Drop an envelope whose invocation was already reaped (timed out / # cleaned up) after it was queued, so a command the user already saw @@ -209,6 +207,54 @@ class DeviceBroker: return None return envelope + async def next_command_async( + self, session: SessionState, timeout: float = 1.0 + ) -> Optional[Dict[str, Any]]: + """Await up to ``timeout`` for the next queued command envelope. + + Event-loop twin of :meth:`next_command` for the native-async SSE + stream: the same envelope checks over the async Redis client, so an + idle device session holds no thread. + """ + redis = await get_async_redis_instance() + if redis is None: + await anyio.sleep(timeout) + return None + try: + popped = await redis.blpop(_cmd_key(session.device_id), timeout=timeout) + except Exception: + logger.exception("async blpop failed for %s", session.device_id) + await anyio.sleep(timeout) + return None + if not popped: + return None + envelope = self._decode_envelope(popped[1]) + if envelope is None: + return None + # Same reaped-invocation guard as ``next_command``; a failed lookup + # drops the envelope too, as ``get_invocation`` returning None does. + inv_id = envelope.get("invocation_id") + if inv_id: + try: + reaped = not await redis.exists(_inv_key(inv_id)) + except Exception: + logger.exception("invocation lookup failed for %s", inv_id) + reaped = True + if reaped: + logger.debug("dropping reaped invocation %s", inv_id) + return None + return envelope + + @staticmethod + def _decode_envelope(raw: Any) -> Optional[Dict[str, Any]]: + """Parse a queued command envelope; ``None`` for anything malformed.""" + try: + envelope = json.loads(_as_str(raw)) + except (TypeError, ValueError): + logger.warning("dropping malformed command envelope") + return None + return envelope if isinstance(envelope, dict) else None + # ------------------------------------------------------------------ # Polling / tickets # ------------------------------------------------------------------ diff --git a/docsgpt/streaming/async_broadcast_channel.py b/docsgpt/streaming/async_broadcast_channel.py index f18e3c42..1c83f187 100644 --- a/docsgpt/streaming/async_broadcast_channel.py +++ b/docsgpt/streaming/async_broadcast_channel.py @@ -1,4 +1,4 @@ -"""Async Redis pub/sub Topic for the native-async SSE reader. +"""Async Redis pub/sub Topic for the native-async SSE readers. Event-loop twin of :class:`docsgpt.streaming.broadcast_channel.Topic`. Same contract — ``subscribe`` yields ``None`` on poll timeout (so the @@ -15,9 +15,11 @@ from __future__ import annotations import inspect import logging +import time from typing import AsyncIterator, Awaitable, Callable, Optional, Union import anyio +from redis.exceptions import MaxConnectionsError from docsgpt.streaming.async_redis import get_async_redis_instance @@ -25,6 +27,10 @@ logger = logging.getLogger(__name__) OnSubscribe = Callable[[], Union[None, Awaitable[None]]] +# Payload of the liveness PING. It differs from redis-py's own health-check +# payload, whose PONG the client swallows before ``get_message`` returns. +_LIVENESS_PING = "docsgpt-liveness" + class AsyncTopic: """An async pub/sub channel identified by a string name.""" @@ -36,6 +42,7 @@ class AsyncTopic: self, on_subscribe: Optional[OnSubscribe] = None, poll_timeout: float = 1.0, + liveness_timeout: Optional[float] = None, ) -> AsyncIterator[Optional[bytes]]: """Subscribe to the topic; yield raw payloads or ``None`` on tick. @@ -45,6 +52,13 @@ class AsyncTopic: live message is processed. If Redis is unavailable, returns immediately without yielding so the caller can fall back to a direct snapshot read. Cleanly unsubscribes on close / disconnect. + + ``liveness_timeout`` protects long-lived subscribers from a + connection NAT/IPVS dropped without a FIN. Such a socket never + errors; ``get_message`` just keeps timing out. After that many + seconds with nothing received a PING goes out, and if nothing comes + back within the same window the subscription ends, so the client + reconnects and replays from its cursor. """ redis = await get_async_redis_instance() if redis is None: @@ -58,6 +72,14 @@ class AsyncTopic: try: try: await pubsub.subscribe(self.name) + except MaxConnectionsError: + # Every pooled connection is held by another open stream. + logger.warning( + "async Redis pool exhausted subscribing to %s; raise " + "ASYNC_REDIS_MAX_CONNECTIONS if this worker should hold more streams", + self.name, + ) + return except Exception: # Transient subscribe failure is treated like "Redis # unavailable": yield nothing, let the caller fall back to @@ -65,6 +87,8 @@ class AsyncTopic: # pubsub down cleanly. logger.exception("async pubsub.subscribe failed for %s", self.name) return + last_received = time.monotonic() + ping_sent_at: Optional[float] = None while True: try: msg = await pubsub.get_message(timeout=poll_timeout) @@ -73,10 +97,34 @@ class AsyncTopic: "async pubsub.get_message failed for %s", self.name ) return + now = time.monotonic() if msg is None: + if liveness_timeout is not None: + if ping_sent_at is not None: + if now - ping_sent_at >= liveness_timeout: + logger.info( + "pubsub liveness probe unanswered for %s; closing subscriber", + self.name, + ) + return + elif now - last_received >= liveness_timeout: + try: + await pubsub.ping(_LIVENESS_PING) + except Exception: + logger.info( + "pubsub liveness probe failed for %s; closing subscriber", + self.name, + exc_info=True, + ) + return + ping_sent_at = now yield None continue + last_received = now + ping_sent_at = None msg_type = msg.get("type") + if msg_type == "pong": + continue if msg_type == "subscribe": if not on_subscribe_fired and on_subscribe is not None: try: diff --git a/docsgpt/streaming/async_redis.py b/docsgpt/streaming/async_redis.py index 4dadf6a6..56211eba 100644 --- a/docsgpt/streaming/async_redis.py +++ b/docsgpt/streaming/async_redis.py @@ -1,11 +1,11 @@ -"""Lazy async Redis client for the native-async SSE reader. +"""Lazy async Redis client for the native-async (event-loop) routes. -Async twin of :func:`docsgpt.cache.get_redis_instance`. The -Starlette-mounted reader (``docsgpt.api.async_sse``) tails pub/sub on -the event loop, so it needs a ``redis.asyncio`` client rather than the -sync one used by the producer side. The app runs a single ASGI worker / -event loop, so a module-level singleton is sufficient and avoids -reconnecting per request. +Async twin of :func:`docsgpt.cache.get_redis_instance`. The Starlette-mounted +routes — the chat reconnect reader, ``/api/events`` and the device command +stream — wait on Redis from the event loop, so they need a ``redis.asyncio`` +client rather than the sync one used by the producer side. The app runs a +single ASGI worker / event loop, so a module-level singleton is sufficient +and avoids reconnecting per request. """ from __future__ import annotations @@ -30,7 +30,9 @@ async def get_async_redis_instance() -> Optional[aioredis.Redis]: 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. + 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. """ global _async_redis, _creation_failed if _async_redis is None and not _creation_failed: @@ -39,6 +41,7 @@ async def get_async_redis_instance() -> Optional[aioredis.Redis]: settings.CACHE_REDIS_URL, socket_connect_timeout=2, health_check_interval=10, + max_connections=int(settings.ASYNC_REDIS_MAX_CONNECTIONS), ) except ValueError as e: logger.error("Invalid Redis URL for async client: %s", e) diff --git a/tests/api/test_asgi_auth.py b/tests/api/test_asgi_auth.py new file mode 100644 index 00000000..cf72eb43 --- /dev/null +++ b/tests/api/test_asgi_auth.py @@ -0,0 +1,116 @@ +"""Tests for ``docsgpt/api/asgi_auth.py`` — auth shared by the Starlette routes. + +The helper must gate a native-async route exactly like ``authenticate_request`` +gates a Flask one: same decoder, same 401 payloads, same OIDC revocation check. +""" + +from __future__ import annotations + +import json +from unittest.mock import patch + +import pytest +from jose import jwt +from starlette.requests import Request + +from docsgpt.api import asgi_auth +from docsgpt.core import log_context + + +def _request(headers: dict | None = None) -> Request: + raw = [(k.lower().encode(), v.encode()) for k, v in (headers or {}).items()] + return Request( + {"type": "http", "method": "GET", "path": "/", "headers": raw, "query_string": b""} + ) + + +@pytest.mark.unit +@pytest.mark.asyncio +class TestAuthenticate: + async def test_local_mode_resolves_local_user(self, monkeypatch): + monkeypatch.setattr(asgi_auth.settings, "AUTH_TYPE", None) + decoded, error = await asgi_auth.authenticate(_request()) + assert error is None + assert decoded == {"sub": "local"} + + async def test_missing_token_is_anonymous_not_rejected(self, monkeypatch): + # Flask leaves ``decoded_token`` None here and lets the route decide + # (a share token can still authorize an artifact download). + monkeypatch.setattr(asgi_auth.settings, "AUTH_TYPE", "simple_jwt") + decoded, error = await asgi_auth.authenticate(_request()) + assert decoded is None + assert error is None + + async def test_invalid_token_rejected_with_decoder_payload(self, monkeypatch): + monkeypatch.setattr(asgi_auth.settings, "AUTH_TYPE", "simple_jwt") + monkeypatch.setattr(asgi_auth.settings, "JWT_SECRET_KEY", "test-secret") + decoded, error = await asgi_auth.authenticate( + _request({"Authorization": "Bearer not-a-jwt"}) + ) + assert decoded is None + assert error.status_code == 401 + assert json.loads(error.body) == { + "message": "Authentication error: invalid token", + "error": "invalid_token", + } + + async def test_valid_token_decoded(self, monkeypatch): + monkeypatch.setattr(asgi_auth.settings, "AUTH_TYPE", "simple_jwt") + monkeypatch.setattr(asgi_auth.settings, "JWT_SECRET_KEY", "test-secret") + token = jwt.encode({"sub": "alice"}, "test-secret", algorithm="HS256") + decoded, error = await asgi_auth.authenticate( + _request({"Authorization": f"Bearer {token}"}) + ) + assert error is None + assert decoded["sub"] == "alice" + + async def test_revoked_oidc_session_rejected(self, monkeypatch): + monkeypatch.setattr(asgi_auth.settings, "AUTH_TYPE", "oidc") + claims = {"sub": "alice", "iat": 1} + with patch.object(asgi_auth, "handle_auth", return_value=claims), patch.object( + asgi_auth, "oidc_session_denied", return_value=True + ) as denied: + decoded, error = await asgi_auth.authenticate(_request()) + assert decoded is None + assert error.status_code == 401 + assert json.loads(error.body) == { + "message": "Authentication error: session revoked", + "error": "token_revoked", + } + denied.assert_called_once_with(claims) + + async def test_live_oidc_session_passes(self, monkeypatch): + monkeypatch.setattr(asgi_auth.settings, "AUTH_TYPE", "oidc") + with patch.object( + asgi_auth, "handle_auth", return_value={"sub": "alice", "iat": 1} + ), patch.object(asgi_auth, "oidc_session_denied", return_value=False): + decoded, error = await asgi_auth.authenticate(_request()) + assert error is None + assert decoded["sub"] == "alice" + + async def test_denylist_only_consulted_for_oidc(self, monkeypatch): + monkeypatch.setattr(asgi_auth.settings, "AUTH_TYPE", "simple_jwt") + with patch.object( + asgi_auth, "handle_auth", return_value={"sub": "alice"} + ), patch.object(asgi_auth, "oidc_session_denied") as denied: + await asgi_auth.authenticate(_request()) + denied.assert_not_called() + + +@pytest.mark.unit +def test_bind_log_context_stamps_request_keys(): + token = asgi_auth.bind_log_context("event_stream", "alice") + try: + snap = log_context.snapshot() + assert snap["endpoint"] == "event_stream" + assert snap["user_id"] == "alice" + assert snap["activity_id"] + finally: + log_context.reset(token) + + +@pytest.mark.unit +def test_json_error_shape(): + response = asgi_auth.json_error("Forbidden", 403) + assert response.status_code == 403 + assert json.loads(response.body) == {"success": False, "message": "Forbidden"} diff --git a/tests/api/test_async_sse_routes.py b/tests/api/test_async_sse_routes.py index ac1c055c..290fe4a0 100644 --- a/tests/api/test_async_sse_routes.py +++ b/tests/api/test_async_sse_routes.py @@ -24,7 +24,7 @@ from docsgpt.core.settings import settings VALID_UUID = "67d65e8f-e7fb-4df1-9e6e-99ea6c830206" -_AUTH = "docsgpt.api.async_sse.handle_auth" +_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" diff --git a/tests/api/test_events_routes.py b/tests/api/test_events_routes.py index d7f05480..a36bbff5 100644 --- a/tests/api/test_events_routes.py +++ b/tests/api/test_events_routes.py @@ -1,615 +1,551 @@ -"""Tests for docsgpt/api/events/routes.py — the SSE endpoint. +"""Tests for docsgpt/api/events/routes.py — the native-async ``GET /api/events``. -The SSE generator runs in a separate thread under the WSGI test client; -we drive it with mocked Redis (the ``pubsub.get_message`` and ``xrange`` -sequences) and read the response body until we have enough records to -assert on, then close the response to terminate the generator. +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``. """ from __future__ import annotations +import inspect import json -import threading from typing import Any -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, patch +import anyio import pytest -from flask import Flask, request +from starlette.applications import Starlette +from starlette.testclient import TestClient + +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 + +_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" +ALICE = {"sub": "alice"} +OLD_CURSOR = "1735682300000-0" -def _make_app(): - """Mount the events blueprint on a bare Flask app + JWT shim. - - The shim mimics ``docsgpt/app.py`` populating - ``request.decoded_token`` so the SSE handler's auth gate sees a - user-id without requiring the full app stack. - """ - from docsgpt.api.events.routes import events - - app = Flask(__name__) - app.register_blueprint(events) - app.config["TESTING"] = True - - @app.before_request - def _shim_auth(): # noqa: D401 - header = request.headers.get("X-Test-Sub") - request.decoded_token = {"sub": header} if header else None - - return app +def _app() -> Starlette: + return Starlette(routes=events_module.event_stream_routes) -class _FakePubSub: - """Minimal Redis pub/sub stand-in for the SSE handler. - - ``messages`` is a list of message dicts the generator should see in - order. After exhausting it, ``get_message`` returns ``None`` (poll - timeout) so the generator stays alive emitting keepalives until the - test closes the response. - """ - - def __init__(self, messages: list[dict[str, Any]]): - self._messages = list(messages) - self.subscribed: list[str] = [] - self.unsubscribed: list[str] = [] - self.closed = False - self._lock = threading.Lock() - - def subscribe(self, name: str): - self.subscribed.append(name) - - def unsubscribe(self, name: str): - self.unsubscribed.append(name) - - def close(self): - self.closed = True - - def get_message(self, timeout: float = 0): - with self._lock: - if self._messages: - return self._messages.pop(0) - return None +def _get(headers: dict | None = None, params: dict | None = None): + return TestClient(_app()).get("/api/events", headers=headers or {}, params=params or {}) -def _drain_until(response, predicate, max_chunks: int = 200) -> bytes: - """Consume the streamed response until ``predicate(buf)`` is true. +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 - Returns the accumulated bytes. Closes the response so the generator - exits cleanly via GeneratorExit. - """ - buf = b"" - iterator = response.iter_encoded() - for _ in range(max_chunks): + +def _subscribe(*payloads: Any, fire_callback: bool = True, record: dict | 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 fire_callback and on_subscribe is not None: + result = on_subscribe() + if inspect.isawaitable(result): + await result + yield None + for payload in payloads: + yield payload + + return _impl + + +def _subscribe_dies_after_callback(): + """SUBSCRIBE acks and the callback runs, then the connection drops before any poll.""" + + async def _impl(self, on_subscribe=None, poll_timeout=1.0, **kwargs): + if on_subscribe is not None: + result = on_subscribe() + if inspect.isawaitable(result): + await result + return + yield # pragma: no cover — make the function an async generator + + return _impl + + +def _endless_subscribe(state: dict): + """A subscription that idles until cancelled; records its own teardown.""" + + async def _impl(self, on_subscribe=None, poll_timeout=1.0, **kwargs): + if on_subscribe is not None: + result = on_subscribe() + if inspect.isawaitable(result): + await result try: - chunk = next(iterator) - except StopIteration: - break - if not chunk: - continue - buf += chunk - if predicate(buf): - break - response.close() - return buf + while True: + await anyio.sleep(0.01) + yield None + finally: + state["closed"] = True + + return _impl + + +def _stored(payload: dict) -> bytes: + return json.dumps(payload).encode() + + +@pytest.fixture(autouse=True) +def _baseline(monkeypatch): + """Hermetic defaults: push on, cap 8, no async Redis, clean shutdown flag.""" + monkeypatch.setattr(events_module.settings, "AUTH_TYPE", None) + monkeypatch.setattr(events_module.settings, "ENABLE_SSE_PUSH", True) + monkeypatch.setattr(events_module.settings, "SSE_MAX_CONCURRENT_PER_USER", 8) + monkeypatch.setattr(events_module.settings, "SSE_KEEPALIVE_SECONDS", 15) + monkeypatch.setattr(events_module.settings, "EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW", 30) + monkeypatch.setattr(events_module.settings, "EVENTS_REPLAY_BUDGET_WINDOW_SECONDS", 60) + monkeypatch.setattr(events_module.settings, "EVENTS_REPLAY_MAX_AGE_HOURS", 0) + monkeypatch.setattr(events_module.settings, "EVENTS_REPLAY_MAX_PER_REQUEST", 200) + monkeypatch.setattr( + "docsgpt.streaming.async_broadcast_channel.get_async_redis_instance", + AsyncMock(return_value=None), + ) + reset_shutdown() + with patch(_AREDIS, AsyncMock(return_value=None)): + yield + reset_shutdown() # ── auth gate ─────────────────────────────────────────────────────────── +@pytest.mark.unit class TestAuthGate: - def test_rejects_when_no_decoded_token(self): - app = _make_app() - with app.test_client() as c: - r = c.get("/api/events") + def test_401_when_no_token(self): + with patch(_AUTH, return_value=None): + r = _get() + assert r.status_code == 401 + assert r.json() == {"success": False, "message": "Authentication required"} + + def test_401_when_token_has_no_sub(self): + with patch(_AUTH, return_value={"email": "x@y.z"}): + r = _get() assert r.status_code == 401 - def test_rejects_when_decoded_token_missing_sub(self): - from docsgpt.api.events import routes as events_module - - app = _make_app() - - # Clear the shim's behavior — supply a decoded_token without sub. - @app.before_request - def _override(): - request.decoded_token = {"email": "x@y.z"} - - with patch.object(events_module, "get_redis_instance", return_value=None): - with app.test_client() as c: - r = c.get("/api/events") + def test_401_passes_decoder_error_through(self): + err = {"message": "Authentication error: token expired", "error": "token_expired"} + with patch(_AUTH, return_value=err): + r = _get() assert r.status_code == 401 + assert r.json() == err + + def test_401_when_oidc_session_revoked(self, monkeypatch): + monkeypatch.setattr(events_module.settings, "AUTH_TYPE", "oidc") + with patch(_AUTH, return_value={"sub": "alice", "iat": 1}), patch( + "docsgpt.api.asgi_auth.oidc_session_denied", return_value=True + ): + r = _get() + assert r.status_code == 401 + assert r.json()["error"] == "token_revoked" # ── streaming response shape ──────────────────────────────────────────── +@pytest.mark.unit class TestStreamShape: - def test_returns_event_stream_mimetype_and_no_buffering_header(self): - from docsgpt.api.events import routes as events_module + def test_event_stream_headers_and_prelude(self): + with patch(_AUTH, return_value=ALICE): + r = _get() + assert r.status_code == 200 + assert r.headers["content-type"].startswith("text/event-stream") + assert r.headers["cache-control"] == "no-store" + assert r.headers["x-accel-buffering"] == "no" + assert r.headers["x-sse-transport"] == "async" + assert r.text.startswith(": connected\n\n") - app = _make_app() - with patch.object(events_module, "get_redis_instance", return_value=None): - with app.test_client() as c: - r = c.get("/api/events", headers={"X-Test-Sub": "alice"}) - assert r.status_code == 200 - assert r.mimetype == "text/event-stream" - assert r.headers.get("Cache-Control") == "no-store" - assert r.headers.get("X-Accel-Buffering") == "no" - # Drain enough to see the prelude comment then close. - body = _drain_until(r, lambda b: b": connected" in b) - assert b": connected" in body + def test_push_disabled_skips_redis_entirely(self, monkeypatch): + monkeypatch.setattr(events_module.settings, "ENABLE_SSE_PUSH", False) + redis = _redis() + with patch(_AUTH, return_value=ALICE), patch(_AREDIS, AsyncMock(return_value=redis)): + r = _get() + assert r.status_code == 200 + assert r.text == ": connected\n\n: push_disabled\n\n" + redis.incr.assert_not_awaited() - def test_emits_push_disabled_when_setting_off(self): - from docsgpt.api.events import routes as events_module + def test_subscribes_to_user_topic_with_liveness_probe(self): + record: dict = {} + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=_redis()) + ), patch(_SUBSCRIBE, _subscribe(record=record)): + _get() + assert record["topic"] == "user:alice" + assert record["poll_timeout"] == events_module.SUBSCRIBE_POLL_INTERVAL_SECONDS + assert record["liveness_timeout"] == PUBSUB_SOCKET_TIMEOUT_SECONDS - app = _make_app() - with patch.object(events_module, "get_redis_instance", return_value=None), \ - patch.object(events_module.settings, "ENABLE_SSE_PUSH", False): - with app.test_client() as c: - r = c.get("/api/events", headers={"X-Test-Sub": "alice"}) - body = _drain_until(r, lambda b: b": push_disabled" in b) - assert b": push_disabled" in body - assert b": connected" in body # prelude still emitted + 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( + _AREDIS, AsyncMock(return_value=_redis()) + ), patch(_SUBSCRIBE, _subscribe()): + r = _get() + assert ": keepalive\n\n" in r.text # ── concurrency cap ───────────────────────────────────────────────────── +@pytest.mark.unit class TestConcurrencyCap: - def test_returns_429_when_user_over_cap(self): - from docsgpt.api.events import routes as events_module - - app = _make_app() - redis_client = MagicMock() - # First INCR returns 9 (over cap of 8). - redis_client.incr.return_value = 9 - - with patch.object(events_module, "get_redis_instance", return_value=redis_client), \ - patch.object(events_module.settings, "SSE_MAX_CONCURRENT_PER_USER", 8): - with app.test_client() as c: - r = c.get("/api/events", headers={"X-Test-Sub": "alice"}) + def test_429_when_over_cap_and_increment_rolled_back(self): + redis = _redis(incr=9) + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _subscribe()): + r = _get() assert r.status_code == 429 - # DECR fired to release the over-cap increment. - redis_client.decr.assert_called_once_with("user:alice:sse_count") + assert r.json() == {"success": False, "message": "Too many concurrent SSE connections"} + redis.decr.assert_awaited_once_with(_COUNTER) - def test_skips_cap_when_zero_disabled(self): - from docsgpt.api.events import routes as events_module + def test_cap_disabled_leaves_counter_alone(self, monkeypatch): + monkeypatch.setattr(events_module.settings, "SSE_MAX_CONCURRENT_PER_USER", 0) + redis = _redis() + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), 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() - app = _make_app() - redis_client = MagicMock() + 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 - with patch.object(events_module, "get_redis_instance", return_value=redis_client), \ - patch.object(events_module.settings, "SSE_MAX_CONCURRENT_PER_USER", 0), \ - patch.object(events_module, "Topic") as mock_topic_cls: - mock_topic = MagicMock() - mock_topic.subscribe.return_value = iter([]) - mock_topic_cls.return_value = mock_topic - redis_client.xinfo_stream.side_effect = Exception("no stream") - redis_client.xrange.return_value = [] - with app.test_client() as c: - r = c.get("/api/events", headers={"X-Test-Sub": "alice"}) - assert r.status_code == 200 - # Concurrency counter not touched when cap is 0. The - # replay-budget INCR is unrelated and may still fire. - incr_keys = [ - call.args[0] for call in redis_client.incr.call_args_list - ] - assert "user:alice:sse_count" not in incr_keys - _drain_until(r, lambda b: b": connected" in b) + def test_incr_failure_fails_open(self): + redis = _redis() + redis.incr = AsyncMock(side_effect=Exception("redis down")) + 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) + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _subscribe()): + r = _get() + assert r.status_code == 200 + redis.incr.assert_awaited_once_with(_COUNTER) + redis.decr.assert_awaited_once_with(_COUNTER) + + def test_redis_unavailable_serves_stream_without_cap(self): + with patch(_AUTH, return_value=ALICE): + r = _get() + assert r.status_code == 200 + assert r.text.startswith(": connected") + + +# ── disconnect and shutdown (raw ASGI) ────────────────────────────────── + + +@pytest.mark.unit +@pytest.mark.asyncio +class TestStreamLifecycle: + async def test_client_disconnect_releases_slot_and_closes_subscription(self, monkeypatch): + monkeypatch.setattr(events_module.settings, "SSE_KEEPALIVE_SECONDS", 0) + redis = _redis(incr=1) + state: dict = {} + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _endless_subscribe(state)): + # Prelude + one keepalive proves the subscription is live before + # the tab closes. + status, _, chunks = await stream_then_disconnect( + _app(), "/api/events", chunks_before_disconnect=2 + ) + assert status == 200 + assert chunks[0] == b": connected\n\n" + assert state.get("closed") is True + redis.decr.assert_awaited_once_with(_COUNTER) + + async def test_disconnect_before_first_frame_still_releases_slot_once(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) + state: dict = {} + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _endless_subscribe(state)): + status, _, _ = await stream_then_disconnect( + _app(), "/api/events", chunks_before_disconnect=0 + ) + assert status == 200 + redis.decr.assert_awaited_once_with(_COUNTER) + + async def test_shutdown_ends_stream_and_releases_slot(self): + redis = _redis(incr=1) + state: dict = {} + begin_shutdown() + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _endless_subscribe(state)): + status, _, chunks = await stream_then_disconnect( + _app(), "/api/events", chunks_before_disconnect=10**6, timeout=3 + ) + assert status == 200 + assert b"".join(chunks) == b": connected\n\n" + assert state.get("closed") is True + redis.decr.assert_awaited_once_with(_COUNTER) # ── replay + live tail ────────────────────────────────────────────────── +@pytest.mark.unit class TestReplayAndTail: def test_replay_yields_xrange_entries_with_injected_id(self): - from docsgpt.api.events import routes as events_module + redis = _redis() + redis.xrange = AsyncMock( + return_value=[ + ( + b"1735682400000-0", + {b"event": _stored({"type": "source.ingest.progress", "payload": {"current": 25}})}, + ), + ] + ) + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _subscribe()): + r = _get(headers={"Last-Event-ID": OLD_CURSOR}) + body = r.text + assert body.startswith(": connected\n\n") + assert "id: 1735682400000-0" in body + data_line = next(line for line in body.split("\n") if line.startswith("data: ")) + envelope = json.loads(data_line[len("data: "):]) + assert envelope["id"] == "1735682400000-0" + assert envelope["payload"] == {"current": 25} + assert "backlog.truncated" not in body + assert redis.xrange.await_args.kwargs["min"] == f"({OLD_CURSOR}" - app = _make_app() - redis_client = MagicMock() - redis_client.incr.return_value = 1 - # Empty stream (no truncation). - redis_client.xinfo_stream.side_effect = Exception("nope") - # XRANGE returns one stored envelope (without ``id``); the route - # injects the entry id on the way out. - stored_event = json.dumps( - { - "type": "source.ingest.progress", - "ts": "2026-04-28T00:00:00.000Z", - "user_id": "alice", - "topic": "user:alice", - "scope": {"kind": "source", "id": "src-1"}, - "payload": {"current": 25, "total": 100}, - } - ).encode() - redis_client.xrange.return_value = [ - (b"1735682400000-0", {b"event": stored_event}), - ] - - # Topic.subscribe yields an immediate timeout so the generator - # keeps running long enough to flush replay; subsequent calls - # also return None. - from docsgpt.api.events.routes import _SSE_LINE_SPLIT # noqa: F401 - - # Fake the broadcast Topic to invoke on_subscribe immediately - # then yield None ticks until close. - def _fake_subscribe(self, on_subscribe=None, poll_timeout=1.0): - if on_subscribe is not None: - on_subscribe() - while True: - yield None - - with patch.object(events_module, "get_redis_instance", return_value=redis_client), \ - patch.object( - events_module.Topic, "subscribe", _fake_subscribe, create=False - ): - with app.test_client() as c: - r = c.get( - "/api/events", - headers={"X-Test-Sub": "alice", "Last-Event-ID": "1735682300000-0"}, - ) - body = _drain_until( - r, - lambda b: b'"current": 25' in b or b'"current":25' in b, - max_chunks=80, - ) - # Replay yields the entry id as the SSE id field. - assert b"id: 1735682400000-0" in body - # Envelope was rewritten to include the injected id. - assert b'"id": "1735682400000-0"' in body or b'"id":"1735682400000-0"' in body - # The connect log fires before replay. - assert b": connected" in body + def test_query_param_cursor_is_honoured(self): + redis = _redis() + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _subscribe()): + _get(params={"last_event_id": OLD_CURSOR}) + assert redis.xrange.await_args.kwargs["min"] == f"({OLD_CURSOR}" def test_snapshot_flushed_when_subscribe_dies_after_callback(self): - """Regression: if ``on_subscribe`` populated ``replay_lines`` but - ``Topic.subscribe`` exits before yielding once (transient Redis - hiccup between SUBSCRIBE-ack and the first poll), the snapshot - must still reach the client. Prior to the fix the in-loop flush - was the only flush, so the backlog was silently dropped. - """ - from docsgpt.api.events import routes as events_module - - app = _make_app() - redis_client = MagicMock() - redis_client.incr.return_value = 1 - redis_client.xinfo_stream.side_effect = Exception("nope") - stored_event = json.dumps( - { - "type": "notification", - "payload": {"text": "from snapshot"}, - } - ).encode() - redis_client.xrange.return_value = [ - (b"1735682400000-0", {b"event": stored_event}), - ] - - # Mimic the broadcast_channel race: SUBSCRIBE acks, on_subscribe - # runs, then the next get_message raises and the generator - # returns without ever yielding. - def _subscribe_dies_after_callback( - self, on_subscribe=None, poll_timeout=1.0 - ): - if on_subscribe is not None: - on_subscribe() - return - yield # pragma: no cover (make the function a generator) - - with patch.object(events_module, "get_redis_instance", return_value=redis_client), \ - patch.object( - events_module.Topic, - "subscribe", - _subscribe_dies_after_callback, - create=False, - ): - with app.test_client() as c: - r = c.get( - "/api/events", - headers={ - "X-Test-Sub": "alice", - "Last-Event-ID": "1735682300000-0", - }, - ) - body = _drain_until( - r, - lambda b: b"from snapshot" in b, - max_chunks=80, - ) - # Snapshot frame must have been flushed via the post-loop - # safety net even though Topic.subscribe exited before - # the in-loop flush could fire. - assert b"id: 1735682400000-0" in body - assert b"from snapshot" in body - # XRANGE was issued exactly once (no double-flush). - redis_client.xrange.assert_called_once() + redis = _redis() + redis.xrange = AsyncMock( + return_value=[ + (b"1735682400000-0", {b"event": _stored({"type": "notification", "payload": {"text": "from snapshot"}})}), + ] + ) + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _subscribe_dies_after_callback()): + r = _get(headers={"Last-Event-ID": OLD_CURSOR}) + assert "id: 1735682400000-0" in r.text + assert "from snapshot" in r.text + redis.xrange.assert_awaited_once() def test_invalid_last_event_id_emits_truncation_notice(self): - from docsgpt.api.events import routes as events_module - - app = _make_app() - redis_client = MagicMock() - redis_client.incr.return_value = 1 - redis_client.xinfo_stream.return_value = {"first-entry": [b"1-0", []]} - redis_client.xrange.return_value = [] - - def _fake_subscribe(self, on_subscribe=None, poll_timeout=1.0): - if on_subscribe is not None: - on_subscribe() - while True: - yield None - - with patch.object(events_module, "get_redis_instance", return_value=redis_client), \ - patch.object(events_module.Topic, "subscribe", _fake_subscribe, create=False): - with app.test_client() as c: - r = c.get( - "/api/events", - headers={"X-Test-Sub": "alice", "Last-Event-ID": "definitely-not-an-id"}, - ) - body = _drain_until( - r, lambda b: b"backlog.truncated" in b, max_chunks=80 - ) - assert b"backlog.truncated" in body + redis = _redis() + redis.xinfo_stream = AsyncMock(return_value={"first-entry": [b"1-0", []]}) + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _subscribe()): + r = _get(headers={"Last-Event-ID": "definitely-not-an-id"}) + assert "backlog.truncated" in r.text + # A corrupt cursor is a fresh session: nothing to replay. + redis.xrange.assert_not_awaited() def test_live_tail_rejects_malformed_event_id_for_dedupe(self): - """A pub/sub envelope carrying a non-Redis-Streams ``id`` must not - seed the dedup floor. Otherwise an adversarial or buggy publisher - could ship ``id="9999999999999-9"`` (lex-greater than any real - id) and pin every subsequent legitimate event below the floor, - silently dropping the user's notifications. + """A pub/sub envelope with a non-Streams ``id`` must not seed the dedup + floor, or a bogus lex-greater id would pin every later event below it. + The event itself still ships, just without an ``id:`` line.""" + redis = _redis() + redis.xrange = AsyncMock( + return_value=[(b"1735682400000-0", {b"event": _stored({"type": "x", "payload": {"step": "replay"}})})] + ) + live_bogus = _stored({"id": "definitely-not-an-id", "type": "x", "payload": {"step": "live-bogus"}}) + live_real = _stored({"id": "1735682500000-0", "type": "x", "payload": {"step": "live-real"}}) + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _subscribe(live_bogus, live_real)): + r = _get(headers={"Last-Event-ID": OLD_CURSOR}) + body = r.text + assert "live-real" in body + assert "id: 1735682500000-0" in body + assert "live-bogus" in body + assert "id: definitely-not-an-id" not in body - The event itself should still be delivered to the client — we - just refuse to use the bogus id for ordering, so it ships - without an SSE ``id:`` header and ``max_replayed_id`` stays put. - """ - from docsgpt.api.events import routes as events_module + def test_live_event_already_covered_by_snapshot_is_dropped(self): + redis = _redis() + redis.xrange = AsyncMock( + return_value=[(b"1735682400000-0", {b"event": _stored({"type": "x", "payload": {"step": "replay"}})})] + ) + duplicate = _stored({"id": "1735682400000-0", "type": "x", "payload": {"step": "replay"}}) + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _subscribe(duplicate)): + r = _get(headers={"Last-Event-ID": OLD_CURSOR}) + assert r.text.count("id: 1735682400000-0") == 1 - app = _make_app() - redis_client = MagicMock() - redis_client.incr.return_value = 1 - redis_client.xinfo_stream.side_effect = Exception("nope") - # Snapshot covers ids up to 1735682400000-0; max_replayed_id - # becomes that value after the in-loop flush. - replay_event = json.dumps({ - "type": "source.ingest.progress", - "payload": {"step": "replay"}, - }).encode() - redis_client.xrange.return_value = [ - (b"1735682400000-0", {b"event": replay_event}), - ] - - live_bogus = json.dumps({ - "id": "definitely-not-an-id", - "type": "source.ingest.completed", - "payload": {"step": "live-bogus"}, - }) - live_real = json.dumps({ - "id": "1735682500000-0", - "type": "source.ingest.completed", - "payload": {"step": "live-real"}, - }) - - def _fake_subscribe(self, on_subscribe=None, poll_timeout=1.0): - # ``Topic.subscribe`` already unpacks redis-py pubsub dicts - # and yields the raw ``data`` bytes (or ``None`` on poll - # timeout). Mirror that contract. - if on_subscribe is not None: - on_subscribe() - yield live_bogus.encode() - yield live_real.encode() - while True: - yield None - - with patch.object( - events_module, "get_redis_instance", return_value=redis_client - ), patch.object( - events_module.Topic, "subscribe", _fake_subscribe, create=False - ): - with app.test_client() as c: - r = c.get( - "/api/events", - headers={ - "X-Test-Sub": "alice", - "Last-Event-ID": "1735682300000-0", - }, - ) - body = _drain_until( - r, lambda b: b"live-real" in b, max_chunks=80 - ) - - # Live-real arrived (its id is strictly greater than the - # replayed snapshot's id), with its valid id surfaced as - # the SSE ``id:`` header so the frontend can advance. - assert b"live-real" in body - assert b"id: 1735682500000-0" in body - - # The bogus-id event was still delivered to the client, - # but no ``id: definitely-not-an-id`` line was emitted — - # the malformed id never reached the SSE wire and so - # could not pin the dedup floor. - assert b"live-bogus" in body - assert b"id: definitely-not-an-id" not in body + def test_non_json_live_payload_passes_through_without_id(self): + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=_redis()) + ), patch(_SUBSCRIBE, _subscribe(b"plain text")): + r = _get() + assert "data: plain text\n\n" in r.text -# ── format helpers (already covered in test_events_substrate but -# duplicated here as a smoke for the route's surface) ───────────────── +# ── replay budget and helpers ─────────────────────────────────────────── -class TestReplayRateLimit: - """Enumeration defenses on the per-user backlog.""" +@pytest.mark.unit +@pytest.mark.asyncio +class TestReplayBudget: + async def test_budget_disabled(self, monkeypatch): + monkeypatch.setattr(events_module.settings, "EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW", 0) + assert await events_module._allow_replay(_redis(), "alice", OLD_CURSOR) is True - def test_allow_replay_returns_true_when_budget_disabled(self): - from docsgpt.api.events.routes import _allow_replay + async def test_redis_unavailable_fails_open(self): + assert await events_module._allow_replay(None, "alice", OLD_CURSOR) is True - with patch("docsgpt.api.events.routes.settings") as mock_settings: - mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 0 - mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60 - assert _allow_replay(MagicMock(), "alice", "1735682400000-0") is True + async def test_no_cursor_never_consumes_budget(self, monkeypatch): + """Fresh sessions start live and must never 429 on the replay budget, + no matter how many tabs open at once.""" + monkeypatch.setattr(events_module.settings, "EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW", 3) + redis = _redis() + for _ in range(50): + assert await events_module._allow_replay(redis, "alice", None) is True + redis.incr.assert_not_awaited() - def test_allow_replay_returns_true_when_redis_unavailable(self): - from docsgpt.api.events.routes import _allow_replay + async def test_passes_until_budget_exhausted(self, monkeypatch): + monkeypatch.setattr(events_module.settings, "EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW", 3) + redis = _redis() + counter = {"v": 0} - with patch("docsgpt.api.events.routes.settings") as mock_settings: - mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 5 - mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60 - assert _allow_replay(None, "alice", "1735682400000-0") is True + def _incr(_key): + counter["v"] += 1 + return counter["v"] - def test_allow_replay_never_consumes_budget_without_cursor(self): - """No cursor ⇒ no replay work ⇒ no budget. Fresh sessions start - live and must never 429 on the replay budget, no matter how many - tabs open at once (30 fresh connects/min used to exhaust it). - """ - from docsgpt.api.events.routes import _allow_replay + redis.incr = AsyncMock(side_effect=_incr) + results = [await events_module._allow_replay(redis, "alice", OLD_CURSOR) for _ in range(4)] + assert results == [True, True, True, False] + # TTL re-seeded on every INCR so a failed seeding EXPIRE can't wedge the key. + assert redis.expire.await_count == 4 + assert all(call.args == ("user:alice:replay_count", 60) for call in redis.expire.await_args_list) - with patch("docsgpt.api.events.routes.settings") as mock_settings: - mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 3 - mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60 - redis = MagicMock() - redis.xlen.return_value = 42 + async def test_fail_open_on_redis_error(self, monkeypatch): + redis = _redis() + redis.incr = AsyncMock(side_effect=Exception("redis down")) + assert await events_module._allow_replay(redis, "alice", OLD_CURSOR) is True - for _ in range(50): - assert _allow_replay(redis, "alice", None) is True + async def test_recovers_when_seeding_expire_raises(self): + redis = _redis() + counter = {"v": 0} - redis.incr.assert_not_called() - redis.xlen.assert_not_called() + def _incr(_key): + counter["v"] += 1 + return counter["v"] - def test_allow_replay_passes_until_budget_exhausted(self): - from docsgpt.api.events.routes import _allow_replay + redis.incr = AsyncMock(side_effect=_incr) + redis.expire = AsyncMock(side_effect=[Exception("expire blip"), True]) + assert await events_module._allow_replay(redis, "alice", OLD_CURSOR) is True + assert await events_module._allow_replay(redis, "alice", OLD_CURSOR) is True + assert redis.expire.await_count == 2 - with patch("docsgpt.api.events.routes.settings") as mock_settings: - mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 3 - mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60 - redis = MagicMock() - counter = {"v": 0} + async def test_replay_backlog_passes_count_to_xrange(self): + redis = _redis() + assert await events_module._replay_backlog(redis, "alice", None, 200) == [] + assert redis.xrange.await_args.kwargs["count"] == 200 - def _incr(_key): - counter["v"] += 1 - return counter["v"] + async def test_replay_backlog_skips_entries_without_event(self): + redis = _redis() + redis.xrange = AsyncMock( + return_value=[ + (b"1-0", {b"other": b"x"}), + (b"2-0", {b"event": b"not json"}), + ] + ) + lines = await events_module._replay_backlog(redis, "alice", "0-0", 200) + assert [entry_id for entry_id, _ in lines] == ["2-0"] + assert lines[0][1] == "id: 2-0\ndata: not json\n\n" - redis.incr.side_effect = _incr + async def test_replay_backlog_returns_nothing_on_xrange_error(self): + redis = _redis() + redis.xrange = AsyncMock(side_effect=Exception("WRONGTYPE")) + assert await events_module._replay_backlog(redis, "alice", "0-0", 200) == [] - # Cursor set → XLEN short-circuit doesn't fire, INCR always runs. - cursor = "1735682400000-0" - # First three pass. - assert _allow_replay(redis, "alice", cursor) is True - assert _allow_replay(redis, "alice", cursor) is True - assert _allow_replay(redis, "alice", cursor) is True - # Fourth refused. - assert _allow_replay(redis, "alice", cursor) is False - # TTL re-seeded on every successful INCR so a transient - # EXPIRE failure on the seeding call can't wedge the key. - assert redis.expire.call_count == 4 - for call in redis.expire.call_args_list: - assert call.args[1] == 60 + async def test_oldest_retained_id(self): + redis = _redis() + redis.xinfo_stream = AsyncMock(return_value={"first-entry": [b"5-0", [b"event", b"{}"]]}) + assert await events_module._oldest_retained_id(redis, "alice") == "5-0" - def test_allow_replay_fail_open_on_redis_error(self): - from docsgpt.api.events.routes import _allow_replay + async def test_oldest_retained_id_bytes_keys(self): + redis = _redis() + redis.xinfo_stream = AsyncMock(return_value={b"first-entry": [b"6-0", []]}) + assert await events_module._oldest_retained_id(redis, "alice") == "6-0" - with patch("docsgpt.api.events.routes.settings") as mock_settings: - mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 5 - mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60 - redis = MagicMock() - redis.incr.side_effect = Exception("redis down") - assert _allow_replay(redis, "alice", "1735682400000-0") is True + async def test_oldest_retained_id_none_on_error_or_empty(self): + redis = _redis() + assert await events_module._oldest_retained_id(redis, "alice") is None + redis.xinfo_stream = AsyncMock(return_value={"first-entry": None}) + assert await events_module._oldest_retained_id(redis, "alice") is None - def test_allow_replay_recovers_when_seeding_expire_raises(self): - """Regression: INCR=1 then EXPIRE raising must not wedge the key. - - Earlier code only called EXPIRE when ``used == 1``. If that EXPIRE - raised, the counter persisted with no TTL and every subsequent - call hit ``used > 1`` without re-seeding — locking the user out - until an operator DEL'd the key. The fix calls EXPIRE on every - successful INCR so the next call still re-seeds the TTL. - """ - from docsgpt.api.events.routes import _allow_replay - - with patch("docsgpt.api.events.routes.settings") as mock_settings: - mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 5 - mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60 - redis = MagicMock() - counter = {"v": 0} - - def _incr(_key): - counter["v"] += 1 - return counter["v"] - - redis.incr.side_effect = _incr - # First EXPIRE raises (the seeding call that would have - # wedged the key under the old gated logic). Second EXPIRE - # succeeds — and crucially, must still run. - redis.expire.side_effect = [Exception("expire blip"), True] - - cursor = "1735682400000-0" - # First call: INCR=1 succeeds, EXPIRE raises -> outer except - # returns True (fail-open) for this call. - assert _allow_replay(redis, "alice", cursor) is True - # Second call: INCR=2, EXPIRE succeeds -> still under budget, - # and the TTL is now seeded (no permanent lockout). - assert _allow_replay(redis, "alice", cursor) is True - - assert redis.expire.call_count == 2 - # Both EXPIRE calls used the configured window. - for call in redis.expire.call_args_list: - assert call.args[1] == 60 - - def test_replay_backlog_passes_count_to_xrange(self): - from docsgpt.api.events.routes import _replay_backlog - - redis = MagicMock() - redis.xrange.return_value = [] - # Drain the iterator so xrange is actually called. - list(_replay_backlog(redis, "alice", None, 200)) - redis.xrange.assert_called_once() - kwargs = redis.xrange.call_args.kwargs - assert kwargs.get("count") == 200 +@pytest.mark.unit +class TestReplayBudgetRoute: def test_returns_429_when_replay_budget_exhausted(self): - """Route refuses the connection rather than serving live tail - only. Earlier behavior silently skipped replay and let the - client advance ``lastEventId`` via id-bearing live frames, - permanently stranding the un-replayed window. The 429 keeps - the cursor pinned so the next reconnect (after the budget - window slides) can replay normally. - """ - from docsgpt.api.events import routes as events_module - - app = _make_app() - redis_client = MagicMock() + """The route refuses rather than serving live-tail only: id-bearing + 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): - if key == "user:alice:sse_count": - return 1 - # Budget counter: report over-limit. - return 31 + return 1 if key == _COUNTER else 31 - redis_client.incr.side_effect = _incr - - with patch.object( - events_module, "get_redis_instance", return_value=redis_client - ), patch.object( - events_module.settings, - "EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW", - 30, - ): - with app.test_client() as c: - r = c.get( - "/api/events", - headers={ - "X-Test-Sub": "alice", - "Last-Event-ID": "1735682300000-0", - }, - ) + redis.incr = AsyncMock(side_effect=_incr) + 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 - # Concurrency slot is released so a budget-denied request - # doesn't permanently consume a connection from the cap. - redis_client.decr.assert_called_once_with("user:alice:sse_count") + 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) +# ── format helpers ────────────────────────────────────────────────────── + + +@pytest.mark.unit class TestFormatHelpers: def test_format_sse_two_terminating_newlines(self): - from docsgpt.api.events.routes import _format_sse - - out = _format_sse("hello", event_id="1-0") + out = events_module._format_sse("hello", event_id="1-0") assert out.endswith("\n\n") - # Exactly one ``id:`` and one ``data:``. - lines = out.rstrip("\n").split("\n") - assert lines == ["id: 1-0", "data: hello"] + assert out.rstrip("\n").split("\n") == ["id: 1-0", "data: hello"] @pytest.mark.parametrize( "candidate, expected", @@ -625,122 +561,56 @@ class TestFormatHelpers: ], ) def test_normalize_last_event_id(self, candidate, expected): - from docsgpt.api.events.routes import _normalize_last_event_id - - assert _normalize_last_event_id(candidate) == expected + assert events_module._normalize_last_event_id(candidate) == expected # ── replay policy: fresh sessions start live, snapshots are age-capped ── +@pytest.mark.unit class TestReplayPolicy: - def test_replay_floor_id_disabled_when_setting_zero(self): - from docsgpt.api.events.routes import _replay_floor_id + def test_replay_floor_id_disabled_when_setting_zero(self, monkeypatch): + monkeypatch.setattr(events_module.settings, "EVENTS_REPLAY_MAX_AGE_HOURS", 0) + assert events_module._replay_floor_id() is None - with patch("docsgpt.api.events.routes.settings") as mock_settings: - mock_settings.EVENTS_REPLAY_MAX_AGE_HOURS = 0 - assert _replay_floor_id() is None - - def test_replay_floor_id_is_ms_stream_id(self): - from docsgpt.api.events.routes import _replay_floor_id - - with patch("docsgpt.api.events.routes.settings") as mock_settings: - mock_settings.EVENTS_REPLAY_MAX_AGE_HOURS = 48 - with patch( - "docsgpt.api.events.routes.time.time", - return_value=1_800_000_000.0, - ): - floor = _replay_floor_id() + def test_replay_floor_id_is_ms_stream_id(self, monkeypatch): + monkeypatch.setattr(events_module.settings, "EVENTS_REPLAY_MAX_AGE_HOURS", 48) + with patch("docsgpt.api.events.routes.time.time", return_value=1_800_000_000.0): + floor = events_module._replay_floor_id() assert floor == f"{(1_800_000_000 - 48 * 3600) * 1000}-0" - def test_replay_backlog_uses_cursor_when_newer_than_floor(self): - from docsgpt.api.events.routes import _replay_backlog + @pytest.mark.asyncio + async def test_replay_backlog_uses_cursor_when_newer_than_floor(self): + redis = _redis() + with patch("docsgpt.api.events.routes._replay_floor_id", return_value="1000-0"): + await events_module._replay_backlog(redis, "alice", "2000-0", 200) + assert redis.xrange.await_args.kwargs["min"] == "(2000-0" - redis = MagicMock() - redis.xrange.return_value = [] - with patch( - "docsgpt.api.events.routes._replay_floor_id", - return_value="1000-0", - ): - list(_replay_backlog(redis, "alice", "2000-0", 200)) - assert redis.xrange.call_args.kwargs["min"] == "(2000-0" - - def test_replay_backlog_clamps_start_to_age_floor(self): - from docsgpt.api.events.routes import _replay_backlog - - redis = MagicMock() - redis.xrange.return_value = [] - with patch( - "docsgpt.api.events.routes._replay_floor_id", - return_value="5000-0", - ): - list(_replay_backlog(redis, "alice", "2000-0", 200)) - # Cursor is older than the floor: replay starts at the floor - # (inclusive), not the cursor — entries past their shelf life - # are never shipped. - assert redis.xrange.call_args.kwargs["min"] == "5000-0" + @pytest.mark.asyncio + async def test_replay_backlog_clamps_start_to_age_floor(self): + redis = _redis() + with patch("docsgpt.api.events.routes._replay_floor_id", return_value="5000-0"): + await events_module._replay_backlog(redis, "alice", "2000-0", 200) + # Entries past their shelf life are never shipped. + assert redis.xrange.await_args.kwargs["min"] == "5000-0" def test_no_cursor_connect_never_replays(self): - """A fresh session (no Last-Event-ID) starts live: no XRANGE, no - weeks-old backlog on every tab-open.""" - from docsgpt.api.events import routes as events_module - - app = _make_app() - redis_client = MagicMock() - redis_client.incr.return_value = 1 - - def _fake_subscribe(self, on_subscribe=None, poll_timeout=1.0): - if on_subscribe is not None: - on_subscribe() - yield None - - with patch.object( - events_module, "get_redis_instance", return_value=redis_client - ), patch.object( - events_module.Topic, "subscribe", _fake_subscribe, create=False - ): - with app.test_client() as c: - r = c.get("/api/events", headers={"X-Test-Sub": "alice"}) - body = _drain_until( - r, lambda b: b": connected" in b, max_chunks=40 - ) - assert b": connected" in body - redis_client.xrange.assert_not_called() + redis = _redis() + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _subscribe()): + r = _get() + assert r.text.startswith(": connected") + redis.xrange.assert_not_awaited() def test_cursor_older_than_floor_emits_truncation_notice(self): - """An age-clamped snapshot has a gap the client can't see from - entry ids alone — it must get the truncation notice so it - refetches full state instead of trusting a partial replay.""" - from docsgpt.api.events import routes as events_module - - app = _make_app() - redis_client = MagicMock() - redis_client.incr.return_value = 1 - # No retained-oldest available: the floor alone must trigger it. - redis_client.xinfo_stream.side_effect = Exception("nope") - redis_client.xrange.return_value = [] - - def _fake_subscribe(self, on_subscribe=None, poll_timeout=1.0): - if on_subscribe is not None: - on_subscribe() - yield None - - with patch.object( - events_module, "get_redis_instance", return_value=redis_client - ), patch.object( - events_module.Topic, "subscribe", _fake_subscribe, create=False - ), patch.object( + """An age-clamped snapshot has a gap the client can't see from entry + ids alone, so it must get the truncation notice and refetch state.""" + redis = _redis() + with patch(_AUTH, return_value=ALICE), patch( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _subscribe()), patch.object( events_module, "_replay_floor_id", return_value="9999999999999-0" ): - with app.test_client() as c: - r = c.get( - "/api/events", - headers={ - "X-Test-Sub": "alice", - "Last-Event-ID": "1735682400000-0", - }, - ) - body = _drain_until( - r, lambda b: b"backlog.truncated" in b, max_chunks=40 - ) - assert b"backlog.truncated" in body + r = _get(headers={"Last-Event-ID": "1735682400000-0"}) + assert "backlog.truncated" in r.text diff --git a/tests/api/user/test_artifacts_routes.py b/tests/api/user/test_artifacts_routes.py index 34370c0a..f77d8ba9 100644 --- a/tests/api/user/test_artifacts_routes.py +++ b/tests/api/user/test_artifacts_routes.py @@ -1,16 +1,28 @@ -"""Route-layer tests for the artifacts API (parent-derived authz).""" +"""Route-layer tests for the artifacts API (parent-derived authz). + +Metadata, restore and delete are Flask-RESTX resources driven through a test +request context. The download is a native-async Starlette route +(``docsgpt/api/user/artifacts/download.py``) driven through Starlette's +``TestClient``, with the JWT decoder patched to the caller's claims. +""" from __future__ import annotations import io import uuid from contextlib import contextmanager +from unittest.mock import patch import pytest from flask import request from sqlalchemy import text +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import StreamingResponse +from starlette.testclient import TestClient from docsgpt.api.user.artifacts import authz +from docsgpt.api.user.artifacts.download import artifact_download_routes, download_artifact from docsgpt.storage.db.repositories.agents import AgentsRepository from docsgpt.storage.db.repositories.artifacts import ArtifactsRepository from docsgpt.storage.db.repositories.conversations import ConversationsRepository @@ -18,6 +30,7 @@ from docsgpt.storage.db.repositories.shared_conversations import ( SharedConversationsRepository, ) from docsgpt.storage.db.repositories.workflow_runs import WorkflowRunsRepository +from tests.asgi_stream import stream_then_disconnect OWNER = "artifact_owner" @@ -35,6 +48,7 @@ def _patch_db(pg_conn, monkeypatch): monkeypatch.setattr("docsgpt.api.user.artifacts.routes.db_readonly", _use_conn) monkeypatch.setattr("docsgpt.api.user.artifacts.routes.db_session", _use_conn) + monkeypatch.setattr("docsgpt.api.user.artifacts.download.db_readonly", _use_conn) return pg_conn @@ -102,6 +116,51 @@ def _call(flask_app, resource_cls, *args, token=None, query=None, json_body=None return getattr(resource, method)(*args) +def _download(artifact_id, *, token=None, query=None, headers=None): + """GET the native-async download route; ``token`` is what the JWT decoder returns.""" + client = TestClient(Starlette(routes=artifact_download_routes)) + with patch("docsgpt.api.asgi_auth.handle_auth", return_value=token): + return client.get( + f"/api/artifacts/{artifact_id}/download", + params=query or {}, + headers=headers or {}, + follow_redirects=False, + ) + + +class _FakeStorage: + def __init__(self, data: bytes = b""): + self._data = data + self.deleted: list = [] + self.opened: list = [] + + def get_file(self, path): + handle = io.BytesIO(self._data) + self.opened.append(handle) + return handle + + def generate_presigned_url(self, path, expires_in=300): + return f"https://signed.example/{path}?exp={expires_in}" + + def delete_file(self, path): + self.deleted.append(path) + return True + + +def _use_storage(monkeypatch, storage, strategy="backend"): + """Serve downloads from ``storage`` under the given ``URL_STRATEGY``.""" + monkeypatch.setattr( + "docsgpt.api.user.artifacts.download.StorageCreator.get_storage", + lambda: storage, + ) + monkeypatch.setattr( + "docsgpt.api.user.artifacts.download.settings.URL_STRATEGY", + strategy, + raising=False, + ) + return storage + + # --------------------------------------------------------------------------- # List # --------------------------------------------------------------------------- @@ -291,9 +350,20 @@ class TestSharedAccess: resp = _call(flask_app, GetArtifact, art["id"], token={"sub": SHARED_USER}) assert resp.status_code == 200 - def test_share_token_holder_can_download(self, _patch_db, flask_app, monkeypatch): - from docsgpt.api.user.artifacts.routes import DownloadArtifact + def test_shared_with_user_can_download(self, _patch_db, monkeypatch): + conv = _make_conversation(_patch_db) + ConversationsRepository(_patch_db).add_shared_user(str(conv["id"]), SHARED_USER) + art = _make_artifact( + _patch_db, conversation_id=str(conv["id"]), + filename="f.bin", storage_path="inputs/owner/artifacts/x/v1/f.bin", + ) + _use_storage(monkeypatch, _FakeStorage(b"SHARED")) + resp = _download(art["id"], token={"sub": SHARED_USER}) + assert resp.status_code == 200 + assert resp.content == b"SHARED" + + def test_share_token_holder_can_download(self, _patch_db, monkeypatch): conv = _make_conversation(_patch_db) # Attach to the first message so it falls inside the first_n_queries snapshot. msg = _make_message(_patch_db, str(conv["id"])) @@ -308,25 +378,23 @@ class TestSharedAccess: share = SharedConversationsRepository(_patch_db).create( str(conv["id"]), OWNER, first_n_queries=1 ) - - storage = _FakeStorage(b"PDFDATA") - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.StorageCreator.get_storage", - lambda: storage, - ) - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.settings.URL_STRATEGY", - "backend", - raising=False, - ) + _use_storage(monkeypatch, _FakeStorage(b"PDFDATA")) # Anonymous link holder (no JWT) supplies the share token. - resp = _call( - flask_app, DownloadArtifact, art["id"], token=None, - query={"share_token": str(share["uuid"])}, - ) + resp = _download(art["id"], token=None, query={"share_token": str(share["uuid"])}) assert resp.status_code == 200 - assert resp.data == b"PDFDATA" + assert resp.content == b"PDFDATA" + + def test_anonymous_without_share_token_denied(self, _patch_db, monkeypatch): + conv = _make_conversation(_patch_db) + art = _make_artifact( + _patch_db, conversation_id=str(conv["id"]), + storage_path="inputs/owner/artifacts/x/v1/f.bin", + ) + _use_storage(monkeypatch, _FakeStorage(b"SECRET")) + + resp = _download(art["id"], token=None) + assert resp.status_code == 403 # --------------------------------------------------------------------------- @@ -361,19 +429,6 @@ class TestShareTokenSnapshotScope: ) return conv_id, in_art, out_art, null_art, str(share["uuid"]) - @staticmethod - def _mock_storage(monkeypatch, data=b"BYTES"): - storage = _FakeStorage(data) - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.StorageCreator.get_storage", - lambda: storage, - ) - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.settings.URL_STRATEGY", - "backend", raising=False, - ) - return storage - def test_share_token_list_only_snapshot(self, _patch_db, flask_app): from docsgpt.api.user.artifacts.routes import ListArtifacts @@ -418,31 +473,23 @@ class TestShareTokenSnapshotScope: ) assert resp.status_code == 403 - def test_share_token_download_in_snapshot_200( - self, _patch_db, flask_app, monkeypatch - ): - from docsgpt.api.user.artifacts.routes import DownloadArtifact - + def test_share_token_download_in_snapshot_200(self, _patch_db, monkeypatch): _, in_art, _out, _null, token = self._seed(_patch_db) - self._mock_storage(monkeypatch, b"INDATA") - resp = _call( - flask_app, DownloadArtifact, in_art["id"], token=None, - query={"share_token": token}, - ) + _use_storage(monkeypatch, _FakeStorage(b"INDATA")) + resp = _download(in_art["id"], token=None, query={"share_token": token}) assert resp.status_code == 200 - assert resp.data == b"INDATA" - - def test_share_token_download_out_of_snapshot_403( - self, _patch_db, flask_app, monkeypatch - ): - from docsgpt.api.user.artifacts.routes import DownloadArtifact + assert resp.content == b"INDATA" + def test_share_token_download_out_of_snapshot_403(self, _patch_db, monkeypatch): _, _in, out_art, _null, token = self._seed(_patch_db) - self._mock_storage(monkeypatch, b"OUTDATA") - resp = _call( - flask_app, DownloadArtifact, out_art["id"], token=None, - query={"share_token": token}, - ) + _use_storage(monkeypatch, _FakeStorage(b"OUTDATA")) + resp = _download(out_art["id"], token=None, query={"share_token": token}) + assert resp.status_code == 403 + + def test_share_token_download_null_message_id_403(self, _patch_db, monkeypatch): + _, _in, _out, null_art, token = self._seed(_patch_db) + _use_storage(monkeypatch, _FakeStorage(b"NULLDATA")) + resp = _download(null_art["id"], token=None, query={"share_token": token}) assert resp.status_code == 403 def test_owner_sees_all_artifacts(self, _patch_db, flask_app, token_owner): @@ -484,149 +531,190 @@ class TestShareTokenSnapshotScope: # --------------------------------------------------------------------------- -# Download +# Download (native-async route) # --------------------------------------------------------------------------- -class _FakeStorage: - def __init__(self, data: bytes = b""): - self._data = data - self.deleted: list = [] - - def get_file(self, path): - return io.BytesIO(self._data) - - def generate_presigned_url(self, path, expires_in=300): - return f"https://signed.example/{path}?exp={expires_in}" - - def delete_file(self, path): - self.deleted.append(path) - return True +def _seed_download(conn, **over): + conv = _make_conversation(conn) + defaults = dict( + conversation_id=str(conv["id"]), + filename="deck.pptx", + mime_type="application/vnd.ms-powerpoint", + storage_path="inputs/owner/artifacts/x/v1/deck.pptx", + ) + defaults.update(over) + return _make_artifact(conn, **defaults) @pytest.mark.unit class TestDownloadArtifact: - def _seed(self, conn, **over): - conv = _make_conversation(conn) - defaults = dict( - conversation_id=str(conv["id"]), - filename="deck.pptx", - mime_type="application/vnd.ms-powerpoint", - storage_path="inputs/owner/artifacts/x/v1/deck.pptx", - ) - defaults.update(over) - return _make_artifact(conn, **defaults) - def test_local_streams_bytes_with_content_disposition( - self, _patch_db, flask_app, token_owner, monkeypatch + self, _patch_db, token_owner, monkeypatch ): - from docsgpt.api.user.artifacts.routes import DownloadArtifact + art = _seed_download(_patch_db) + storage = _use_storage(monkeypatch, _FakeStorage(b"BINARY")) - art = self._seed(_patch_db) - storage = _FakeStorage(b"BINARY") - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.StorageCreator.get_storage", - lambda: storage, - ) - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.settings.URL_STRATEGY", - "backend", - raising=False, - ) - - resp = _call(flask_app, DownloadArtifact, art["id"], token=token_owner) + resp = _download(art["id"], token=token_owner) assert resp.status_code == 200 - assert resp.data == b"BINARY" - assert resp.headers["Content-Disposition"] == 'attachment; filename="deck.pptx"' - assert resp.headers["Content-Type"] == "application/vnd.ms-powerpoint" + assert resp.content == b"BINARY" + assert resp.headers["content-disposition"] == 'attachment; filename="deck.pptx"' + assert resp.headers["content-type"] == "application/vnd.ms-powerpoint" + assert storage.opened[0].closed - def test_local_download_is_streamed_not_buffered( - self, _patch_db, flask_app, token_owner, monkeypatch + def test_missing_mime_type_defaults_to_octet_stream( + self, _patch_db, token_owner, monkeypatch ): - # The response body must be a stream (generator), not the whole object - # buffered into memory via make_response(file_obj.read()). - from docsgpt.api.user.artifacts.routes import DownloadArtifact + art = _seed_download(_patch_db, mime_type=None) + _use_storage(monkeypatch, _FakeStorage(b"X")) + resp = _download(art["id"], token=token_owner) + assert resp.headers["content-type"] == "application/octet-stream" - art = self._seed(_patch_db) - storage = _FakeStorage(b"Z" * 200_000) - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.StorageCreator.get_storage", - lambda: storage, - ) - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.settings.URL_STRATEGY", - "backend", raising=False, - ) - resp = _call(flask_app, DownloadArtifact, art["id"], token=token_owner) - assert resp.is_streamed - assert resp.get_data() == b"Z" * 200_000 - - def test_s3_strategy_redirects_to_presigned( - self, _patch_db, flask_app, token_owner, monkeypatch + def test_missing_filename_falls_back_to_artifact_id( + self, _patch_db, token_owner, monkeypatch ): - from docsgpt.api.user.artifacts.routes import DownloadArtifact + art = _seed_download(_patch_db, filename=None) + _use_storage(monkeypatch, _FakeStorage(b"X")) + resp = _download(art["id"], token=token_owner) + assert resp.headers["content-disposition"] == f'attachment; filename="artifact-{art["id"]}"' - art = self._seed(_patch_db) - storage = _FakeStorage(b"unused") - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.StorageCreator.get_storage", - lambda: storage, - ) - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.settings.URL_STRATEGY", - "s3", - raising=False, - ) + def test_disk_file_streamed_in_chunks_and_closed( + self, _patch_db, token_owner, monkeypatch, tmp_path + ): + # A real file handle is read off the event loop, chunk by chunk. + payload = bytes(range(256)) * 800 # 204_800 bytes -> 4 chunks + path = tmp_path / "deck.pptx" + path.write_bytes(payload) + handles = [] - resp = _call(flask_app, DownloadArtifact, art["id"], token=token_owner) + class _DiskStorage: + def get_file(self, _path): + handle = open(path, "rb") + handles.append(handle) + return handle + + art = _seed_download(_patch_db) + _use_storage(monkeypatch, _DiskStorage()) + resp = _download(art["id"], token=token_owner) + assert resp.status_code == 200 + assert resp.content == payload + assert handles[0].closed + + def test_explicit_version_served(self, _patch_db, token_owner, monkeypatch): + art = _seed_download(_patch_db) + ArtifactsRepository(_patch_db).append_version( + art["id"], filename="deck-v2.pptx", storage_path="inputs/owner/artifacts/x/v2/deck.pptx", + ) + seen_paths = [] + + class _PathStorage(_FakeStorage): + def get_file(self, path): + seen_paths.append(path) + return super().get_file(path) + + _use_storage(monkeypatch, _PathStorage(b"V1")) + resp = _download(art["id"], token=token_owner, query={"version": "1"}) + assert resp.status_code == 200 + assert seen_paths == ["inputs/owner/artifacts/x/v1/deck.pptx"] + + _download(art["id"], token=token_owner) + assert seen_paths[-1] == "inputs/owner/artifacts/x/v2/deck.pptx" + + def test_invalid_version_400(self, _patch_db, token_owner, monkeypatch): + art = _seed_download(_patch_db) + _use_storage(monkeypatch, _FakeStorage(b"X")) + resp = _download(art["id"], token=token_owner, query={"version": "latest"}) + assert resp.status_code == 400 + assert resp.json() == {"success": False, "message": "Invalid version"} + + def test_unknown_version_404(self, _patch_db, token_owner, monkeypatch): + art = _seed_download(_patch_db) + _use_storage(monkeypatch, _FakeStorage(b"X")) + resp = _download(art["id"], token=token_owner, query={"version": "7"}) + assert resp.status_code == 404 + assert resp.json() == {"success": False, "message": "Version not found"} + + def test_version_without_file_404(self, _patch_db, token_owner, monkeypatch): + art = _seed_download(_patch_db, storage_path=None) + _use_storage(monkeypatch, _FakeStorage(b"X")) + resp = _download(art["id"], token=token_owner) + assert resp.status_code == 404 + assert resp.json() == {"success": False, "message": "No file for this version"} + + def test_missing_bytes_404(self, _patch_db, token_owner, monkeypatch): + class _GoneStorage(_FakeStorage): + def get_file(self, path): + raise FileNotFoundError(path) + + art = _seed_download(_patch_db) + _use_storage(monkeypatch, _GoneStorage()) + resp = _download(art["id"], token=token_owner) + assert resp.status_code == 404 + assert resp.json() == {"success": False, "message": "File not found"} + + def test_storage_failure_400(self, _patch_db, token_owner, monkeypatch): + class _BrokenStorage(_FakeStorage): + def get_file(self, path): + raise RuntimeError("backend exploded") + + art = _seed_download(_patch_db) + _use_storage(monkeypatch, _BrokenStorage()) + resp = _download(art["id"], token=token_owner) + assert resp.status_code == 400 + assert resp.json() == {"success": False} + + def test_s3_strategy_redirects_to_presigned(self, _patch_db, token_owner, monkeypatch): + art = _seed_download(_patch_db) + storage = _use_storage(monkeypatch, _FakeStorage(b"unused"), strategy="s3") + + resp = _download(art["id"], token=token_owner) assert resp.status_code == 302 - assert resp.headers["Location"].startswith("https://signed.example/") + assert resp.headers["location"] == ( + "https://signed.example/inputs/owner/artifacts/x/v1/deck.pptx?exp=300" + ) + assert storage.opened == [] def test_s3_strategy_disposition_url_returns_json( - self, _patch_db, flask_app, token_owner, monkeypatch + self, _patch_db, token_owner, monkeypatch ): # ?disposition=url opts into a JSON envelope (for a top-level browser # navigation) instead of the CORS-blocked cross-origin 302. - from docsgpt.api.user.artifacts.routes import DownloadArtifact + art = _seed_download(_patch_db) + _use_storage(monkeypatch, _FakeStorage(b"unused"), strategy="s3") - art = self._seed(_patch_db) - storage = _FakeStorage(b"unused") - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.StorageCreator.get_storage", - lambda: storage, - ) - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.settings.URL_STRATEGY", - "s3", - raising=False, - ) - - resp = _call( - flask_app, - DownloadArtifact, - art["id"], - token=token_owner, - query={"disposition": "url"}, - ) + resp = _download(art["id"], token=token_owner, query={"disposition": "url"}) assert resp.status_code == 200 - assert resp.json["success"] is True - assert resp.json["url"].startswith("https://signed.example/") + assert resp.json()["success"] is True + assert resp.json()["url"].startswith("https://signed.example/") # The envelope carries a distinctive media type so the client keys off # a server signal, not the JSON body shape (open-redirect gadget guard). - assert resp.mimetype == "application/vnd.docsgpt.artifact-url+json" + assert resp.headers["content-type"] == "application/vnd.docsgpt.artifact-url+json" - def test_stranger_denied(self, _patch_db, flask_app, monkeypatch): - from docsgpt.api.user.artifacts.routes import DownloadArtifact + def test_s3_strategy_accept_json_returns_json(self, _patch_db, token_owner, monkeypatch): + art = _seed_download(_patch_db) + _use_storage(monkeypatch, _FakeStorage(b"unused"), strategy="s3") - art = self._seed(_patch_db) - resp = _call(flask_app, DownloadArtifact, art["id"], token={"sub": STRANGER}) + resp = _download(art["id"], token=token_owner, headers={"Accept": "application/json"}) + assert resp.status_code == 200 + assert resp.headers["content-type"] == "application/vnd.docsgpt.artifact-url+json" + + def test_stranger_denied(self, _patch_db, monkeypatch): + art = _seed_download(_patch_db) + storage = _use_storage(monkeypatch, _FakeStorage(b"SECRET")) + resp = _download(art["id"], token={"sub": STRANGER}) assert resp.status_code == 403 + assert resp.json() == {"success": False, "message": "Forbidden"} + assert storage.opened == [] + + def test_invalid_bearer_401_with_decoder_payload(self, _patch_db, monkeypatch): + art = _seed_download(_patch_db) + _use_storage(monkeypatch, _FakeStorage(b"X")) + err = {"message": "Authentication error: invalid token", "error": "invalid_token"} + resp = _download(art["id"], token=err) + assert resp.status_code == 401 + assert resp.json() == err def test_s3_strategy_misconfigured_backend_500( - self, _patch_db, flask_app, token_owner, monkeypatch + self, _patch_db, token_owner, monkeypatch ): - from docsgpt.api.user.artifacts.routes import DownloadArtifact - - art = self._seed(_patch_db) - class _NoPresignStorage: def get_file(self, path): return io.BytesIO(b"unused") @@ -634,39 +722,79 @@ class TestDownloadArtifact: def generate_presigned_url(self, path, expires_in=300): raise NotImplementedError("backend cannot mint presigned URLs") - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.StorageCreator.get_storage", - lambda: _NoPresignStorage(), - ) - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.settings.URL_STRATEGY", - "s3", - raising=False, - ) + art = _seed_download(_patch_db) + _use_storage(monkeypatch, _NoPresignStorage(), strategy="s3") - resp = _call(flask_app, DownloadArtifact, art["id"], token=token_owner) + resp = _download(art["id"], token=token_owner) assert resp.status_code == 500 + assert resp.json() == {"success": False, "message": "Storage misconfigured"} - def test_crlf_filename_sanitized_in_header( - self, _patch_db, flask_app, token_owner, monkeypatch - ): - from docsgpt.api.user.artifacts.routes import DownloadArtifact - - art = self._seed(_patch_db, filename='a"\r\nInjected: x.txt') - storage = _FakeStorage(b"X") - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.StorageCreator.get_storage", - lambda: storage, - ) - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.settings.URL_STRATEGY", - "backend", - raising=False, - ) - resp = _call(flask_app, DownloadArtifact, art["id"], token=token_owner) - disposition = resp.headers["Content-Disposition"] + def test_crlf_filename_sanitized_in_header(self, _patch_db, token_owner, monkeypatch): + art = _seed_download(_patch_db, filename='a"\r\nInjected: x.txt') + _use_storage(monkeypatch, _FakeStorage(b"X")) + resp = _download(art["id"], token=token_owner) + disposition = resp.headers["content-disposition"] assert "\r" not in disposition and "\n" not in disposition assert disposition == 'attachment; filename="aInjected: x.txt"' + assert "injected" not in resp.headers + + 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 + assert resp.json() == {"success": False, "message": "Artifact not found"} + + def test_unknown_artifact_404(self, _patch_db, token_owner): + resp = _download(str(uuid.uuid4()), token=token_owner) + assert resp.status_code == 404 + + +@pytest.mark.unit +@pytest.mark.asyncio +class TestDownloadStreaming: + async def test_download_is_streamed_not_buffered(self, _patch_db, token_owner, monkeypatch): + # The body must be a chunked stream, not the whole object read into memory. + art = _seed_download(_patch_db) + _use_storage(monkeypatch, _FakeStorage(b"Z" * 200_000)) + request = Request( + { + "type": "http", + "method": "GET", + "path": f"/api/artifacts/{art['id']}/download", + "path_params": {"artifact_id": str(art["id"])}, + "headers": [], + "query_string": b"", + } + ) + with patch("docsgpt.api.asgi_auth.handle_auth", return_value=token_owner): + resp = await download_artifact(request) + assert isinstance(resp, StreamingResponse) + chunks = [chunk async for chunk in resp.body_iterator] + assert [len(chunk) for chunk in chunks] == [65536, 65536, 65536, 3392] + + async def test_client_disconnect_closes_file( + self, _patch_db, token_owner, monkeypatch, tmp_path + ): + path = tmp_path / "big.bin" + path.write_bytes(b"B" * (1024 * 1024)) + handles = [] + + class _DiskStorage: + def get_file(self, _path): + handle = open(path, "rb") + handles.append(handle) + return handle + + art = _seed_download(_patch_db) + _use_storage(monkeypatch, _DiskStorage()) + with patch("docsgpt.api.asgi_auth.handle_auth", return_value=token_owner): + status, _, chunks = await stream_then_disconnect( + Starlette(routes=artifact_download_routes), + f"/api/artifacts/{art['id']}/download", + chunks_before_disconnect=1, + ) + assert status == 200 + assert chunks[0] == b"B" * 65536 + assert handles[0].closed # --------------------------------------------------------------------------- @@ -850,7 +978,6 @@ class TestMalformedArtifactId: [ ("GetArtifact", (), "get", None), ("GetArtifactVersion", (1,), "get", None), - ("DownloadArtifact", (), "get", None), ("RestoreArtifact", (), "post", {"version": 1}), ], ) @@ -930,13 +1057,9 @@ class TestApiKeyPrincipal: ) assert resp.status_code == 403 - def test_api_key_download_requires_conversation_id( - self, _patch_db, flask_app, monkeypatch - ): + def test_api_key_download_requires_conversation_id(self, _patch_db, monkeypatch): # Download follows the same bearer-capability rule as get: the key alone is # denied; the key + matching conversation_id is served. - from docsgpt.api.user.artifacts.routes import DownloadArtifact - agent = _make_agent(_patch_db) conv = _make_agent_conversation(_patch_db, agent["id"]) art = _make_artifact( @@ -944,28 +1067,33 @@ class TestApiKeyPrincipal: filename="f.bin", storage_path="inputs/owner/artifacts/x/v1/f.bin", ) _wire_api_key(monkeypatch, _patch_db) - storage = _FakeStorage(b"BYTES") - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.StorageCreator.get_storage", - lambda: storage, - ) - monkeypatch.setattr( - "docsgpt.api.user.artifacts.routes.settings.URL_STRATEGY", - "backend", raising=False, - ) + _use_storage(monkeypatch, _FakeStorage(b"BYTES")) - denied = _call( - flask_app, DownloadArtifact, art["id"], token=None, - query={"api_key": "secret-key"}, - ) + denied = _download(art["id"], token=None, query={"api_key": "secret-key"}) assert denied.status_code == 403 - ok = _call( - flask_app, DownloadArtifact, art["id"], token=None, + ok = _download( + art["id"], token=None, query={"api_key": "secret-key", "conversation_id": str(conv["id"])}, ) assert ok.status_code == 200 - assert ok.data == b"BYTES" + assert ok.content == b"BYTES" + + def test_api_key_download_outside_agent_scope_denied(self, _patch_db, monkeypatch): + _make_agent(_patch_db) + other_conv = _make_conversation(_patch_db, user_id=OWNER) # no agent_id + art = _make_artifact( + _patch_db, conversation_id=str(other_conv["id"]), + storage_path="inputs/owner/artifacts/x/v1/f.bin", + ) + _wire_api_key(monkeypatch, _patch_db) + _use_storage(monkeypatch, _FakeStorage(b"BYTES")) + + resp = _download( + art["id"], token=None, + query={"api_key": "secret-key", "conversation_id": str(other_conv["id"])}, + ) + assert resp.status_code == 403 def test_api_key_list_without_conversation_id_forbidden( self, _patch_db, flask_app, monkeypatch diff --git a/tests/asgi_stream.py b/tests/asgi_stream.py new file mode 100644 index 00000000..ce8e2957 --- /dev/null +++ b/tests/asgi_stream.py @@ -0,0 +1,86 @@ +"""Drive a streaming ASGI response and disconnect the client mid-stream. + +Starlette's ``TestClient`` buffers a response until the app returns, so it +can't exercise what happens when a browser tab closes on an endless SSE +stream. This harness speaks raw ASGI: it hands the app one GET, collects the +response, and delivers ``http.disconnect`` once enough body has arrived. +""" + +from __future__ import annotations + +from typing import Mapping, Optional + +import anyio + + +async def stream_then_disconnect( + app, + path: str, + *, + headers: Optional[Mapping[str, str]] = None, + query_string: bytes = b"", + chunks_before_disconnect: int = 1, + timeout: float = 5.0, +) -> tuple[Optional[int], dict[str, str], list[bytes]]: + """Run ``app`` for one GET, disconnecting after ``chunks_before_disconnect`` body chunks. + + 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. + + Returns: + tuple: ``(status, headers, body_chunks)`` with lower-cased header names. + """ + disconnected = anyio.Event() + if chunks_before_disconnect <= 0: + disconnected.set() + request_delivered = False + status: Optional[int] = None + response_headers: dict[str, str] = {} + chunks: list[bytes] = [] + + scope = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "method": "GET", + "scheme": "http", + "path": path, + "raw_path": path.encode("utf-8"), + "root_path": "", + "query_string": query_string, + "headers": [ + (name.lower().encode("latin-1"), value.encode("latin-1")) + for name, value in (headers or {}).items() + ], + "client": ("testclient", 50000), + "server": ("testserver", 80), + } + + async def receive(): + nonlocal request_delivered + if not request_delivered: + request_delivered = True + return {"type": "http.request", "body": b"", "more_body": False} + await disconnected.wait() + return {"type": "http.disconnect"} + + async def send(message): + nonlocal status + if message["type"] == "http.response.start": + status = message["status"] + for name, value in message.get("headers", []): + response_headers[name.decode("latin-1")] = value.decode("latin-1") + elif message["type"] == "http.response.body": + body = message.get("body", b"") + if body: + chunks.append(body) + if len(chunks) >= chunks_before_disconnect: + disconnected.set() + if not message.get("more_body", False): + disconnected.set() + + with anyio.fail_after(timeout): + await app(scope, receive, send) + return status, response_headers, chunks diff --git a/tests/devices/conftest.py b/tests/devices/conftest.py index d573263c..bf2663bc 100644 --- a/tests/devices/conftest.py +++ b/tests/devices/conftest.py @@ -11,6 +11,8 @@ deliberately *not* faithful in several ways, so tests must not lean on them: real Redis with ``approximate=True`` only trims past a slack. One ``FakeRedis`` shared by two ``DeviceBroker`` instances models two processes (e.g. a Celery worker and the web tier) talking through one Redis. +``AsyncFakeRedis`` puts an awaitable face on the same store for the broker's +async session path, so sync writers and the async reader share state. """ from __future__ import annotations @@ -18,6 +20,7 @@ from __future__ import annotations import threading import time +import anyio import pytest from docsgpt.devices.broker import DeviceBroker @@ -187,6 +190,24 @@ class FakeRedis: return out or None +class AsyncFakeRedis: + """Awaitable facade over a ``FakeRedis`` for the broker's async session path.""" + + def __init__(self, sync: FakeRedis) -> None: + self.sync = sync + + async def blpop(self, key, timeout=0): + popped = self.sync.blpop(key, timeout=timeout) + if popped is None: + # Real BLPOP waits out ``timeout``; a short sleep keeps a polling + # loop from spinning without slowing tests down. + await anyio.sleep(min(float(timeout or 0), 0.01)) + return popped + + async def exists(self, key): + return self.sync.exists(key) + + @pytest.fixture def fake_redis() -> FakeRedis: return FakeRedis() @@ -199,3 +220,20 @@ def broker_env(monkeypatch, fake_redis): "docsgpt.devices.broker.get_redis_instance", lambda: fake_redis ) return DeviceBroker(), fake_redis + + +@pytest.fixture +def async_broker_env(monkeypatch, fake_redis): + """Like ``broker_env``, with the async client backed by the same store.""" + async_redis = AsyncFakeRedis(fake_redis) + + async def _get_async_redis(): + return async_redis + + monkeypatch.setattr( + "docsgpt.devices.broker.get_redis_instance", lambda: fake_redis + ) + monkeypatch.setattr( + "docsgpt.devices.broker.get_async_redis_instance", _get_async_redis + ) + return DeviceBroker(), fake_redis diff --git a/tests/devices/test_session_events_async.py b/tests/devices/test_session_events_async.py new file mode 100644 index 00000000..c2e72e4f --- /dev/null +++ b/tests/devices/test_session_events_async.py @@ -0,0 +1,401 @@ +"""Tests for the native-async device SSE stream. + +``GET /api/devices/sessions/{session_id}/events`` runs on the event loop: the +device-token check runs in a worker thread, commands arrive through the +broker's async ``BLPOP``, and an idle stream holds no thread. Streams here end +on their own (idle timeout, superseded session, shutdown) so the client +returns; a mid-stream disconnect goes through ``tests/asgi_stream.py``. +Ticket rejection (410) and the poll→stream loop live in test_session_ticket.py. +""" + +from __future__ import annotations + +import base64 +import hashlib +import json +import time +from types import SimpleNamespace + +import pytest +from flask import Flask +from starlette.applications import Starlette +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.shutdown import begin_shutdown, reset_shutdown +from tests.asgi_stream import stream_then_disconnect + +from .conftest import AsyncFakeRedis + +TOKEN = "tok_good" +DEVICE = { + "id": "dev_async", + "user_id": "user_async", + "name": "laptop", + "status": "active", + "approval_mode": "ask", + "machine_pubkey_fingerprint": "fp", +} +AUTH = {"Authorization": f"Bearer {TOKEN}"} + + +class _Ctx: + def __enter__(self): + return None + + def __exit__(self, *exc): + return False + + +@pytest.fixture(autouse=True) +def _baseline(monkeypatch): + monkeypatch.setattr(events_module.settings, "SSE_KEEPALIVE_SECONDS", 15) + monkeypatch.setattr(events_module.settings, "REMOTE_DEVICE_SESSION_IDLE_SECONDS", 0.2) + monkeypatch.setattr(auth_module.settings, "REMOTE_DEVICE_REQUIRE_SIGNATURE", False) + reset_shutdown() + yield + reset_shutdown() + + +@pytest.fixture +def device(monkeypatch): + """Resolve ``TOKEN`` to a device row without Postgres; other tokens are unknown.""" + row = dict(DEVICE) + touched: list = [] + + class _Repo: + def __init__(self, _conn): + pass + + def find_by_token_hash(self, token_hash): + return row if token_hash == auth_module.hash_session_token(TOKEN) else None + + def touch_last_seen(self, device_id): + touched.append(device_id) + + monkeypatch.setattr(auth_module, "DevicesRepository", _Repo) + monkeypatch.setattr(auth_module, "db_readonly", _Ctx) + monkeypatch.setattr(auth_module, "db_session", _Ctx) + return SimpleNamespace(row=row, touched=touched) + + +@pytest.fixture +def broker(async_broker_env, monkeypatch): + broker, fake = async_broker_env + monkeypatch.setattr(events_module, "get_broker", lambda: broker) + return broker, fake + + +def _app() -> Starlette: + return Starlette(routes=events_module.device_session_routes) + + +def _path(session_id: str) -> str: + return f"/api/devices/sessions/{session_id}/events" + + +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) + + +def _queue(broker, invocation_id: str = "inv_1") -> None: + broker.dispatch_invocation( + DEVICE["id"], + DEVICE["user_id"], + {"invocation_id": invocation_id, "action": "run_command", "command": "ls"}, + ) + + +def _ticket(broker, invocation_id: str = "inv_1") -> str: + _queue(broker, invocation_id) + return broker.claim_ticket(DEVICE["id"], 30) + + +def _records(body: str) -> list[dict]: + """Parse SSE records into ``{event, id, data}`` dicts, dropping comments.""" + out = [] + for block in body.split("\n\n"): + fields = {} + for line in block.split("\n"): + if not line or line.startswith(":"): + continue + name, _, value = line.partition(": ") + fields[name] = value + if fields: + if "data" in fields: + fields["data"] = json.loads(fields["data"]) + out.append(fields) + return out + + +# ── auth ──────────────────────────────────────────────────────────────── + + +@pytest.mark.unit +class TestAuth: + def test_401_without_token(self, device, broker): + r = _open("st_x", headers={}) + assert r.status_code == 401 + assert r.json() == {"success": False, "error": "missing_token"} + + def test_401_for_unknown_token(self, device, broker): + r = _open("st_x", headers={"Authorization": "Bearer tok_other"}) + assert r.status_code == 401 + assert r.json() == {"success": False, "error": "invalid_token"} + + def test_valid_token_touches_last_seen(self, device, broker): + b, _ = broker + ticket = _ticket(b) + assert _open(ticket).status_code == 200 + assert device.touched == [DEVICE["id"]] + + +# ── stream ────────────────────────────────────────────────────────────── + + +@pytest.mark.unit +class TestStream: + def test_delivers_queued_invocation_then_ends_idle_session(self, device, broker): + b, fake = broker + ticket = _ticket(b) + r = _open(ticket) + assert r.status_code == 200 + assert r.headers["content-type"].startswith("text/event-stream") + assert r.headers["cache-control"] == "no-cache" + assert r.headers["x-accel-buffering"] == "no" + records = _records(r.text) + assert [rec["event"] for rec in records] == ["invocation", "session_end"] + assert records[0]["id"] == "1" + assert records[0]["data"]["invocation_id"] == "inv_1" + assert records[0]["data"]["command"] == "ls" + assert records[1]["id"] == "2" + assert records[1]["data"] == {"reason": "inactivity_timeout"} + # The ticket became the session id and the session is gone once the stream ends. + assert fake.get(f"dev:ticket:{DEVICE['id']}") is None + assert b.get_session(ticket) is None + + def test_delivers_commands_in_queue_order(self, device, broker): + b, _ = broker + ticket = _ticket(b, "inv_1") + _queue(b, "inv_2") + records = _records(_open(ticket).text) + invocations = [rec for rec in records if rec["event"] == "invocation"] + assert [rec["data"]["invocation_id"] for rec in invocations] == ["inv_1", "inv_2"] + assert [rec["id"] for rec in records] == ["1", "2", "3"] + + def test_heartbeat_comment_while_idle(self, device, broker, monkeypatch): + monkeypatch.setattr(events_module.settings, "SSE_KEEPALIVE_SECONDS", 0.05) + monkeypatch.setattr(events_module.settings, "REMOTE_DEVICE_SESSION_IDLE_SECONDS", 0.3) + b, _ = broker + r = _open(_ticket(b)) + assert ": heartbeat\n\n" in r.text + + def test_reaped_invocation_is_not_delivered(self, device, broker): + b, fake = broker + ticket = _ticket(b) + fake.delete("dev:inv:inv_1") + records = _records(_open(ticket).text) + assert [rec["event"] for rec in records] == ["session_end"] + assert records[0]["id"] == "1" + + def test_output_activity_keeps_session_alive(self, device, broker, monkeypatch): + # submit_output_chunk (a Flask thread) bumps last_activity_at; the async + # idle check must read it rather than a stale copy. + monkeypatch.setattr(events_module.settings, "REMOTE_DEVICE_SESSION_IDLE_SECONDS", 0.3) + b, _ = broker + ticket = _ticket(b) + original = b.next_command_async + started = time.time() + + async def _poll_while_output_flows(sess, timeout=1.0): + if time.time() - started < 0.6: + sess.last_activity_at = time.time() + return await original(sess, timeout=timeout) + + monkeypatch.setattr(b, "next_command_async", _poll_while_output_flows) + r = _open(ticket) + assert time.time() - started >= 0.6 + assert [rec["event"] for rec in _records(r.text)] == ["invocation", "session_end"] + + +@pytest.mark.unit +@pytest.mark.asyncio +class TestStreamLifecycle: + async def test_reconnect_ends_superseded_stream_and_keeps_new_session( + self, device, broker, monkeypatch + ): + monkeypatch.setattr(events_module.settings, "REMOTE_DEVICE_SESSION_IDLE_SECONDS", 30) + b, _ = broker + ticket = _ticket(b) + replacement: dict = {} + original = b.next_command_async + + async def _reconnect_then_poll(sess, timeout=1.0): + envelope = await original(sess, timeout=timeout) + if not replacement: + # The CLI reconnected on another request: its session replaces this one. + replacement["sess"] = b.register_session(DEVICE["id"], DEVICE["user_id"]) + return envelope + + monkeypatch.setattr(b, "next_command_async", _reconnect_then_poll) + status, _, chunks = await stream_then_disconnect( + _app(), _path(ticket), headers=AUTH, chunks_before_disconnect=10**6, timeout=3 + ) + assert status == 200 + assert [rec["event"] for rec in _records(b"".join(chunks).decode())] == ["invocation"] + assert b.get_session(ticket) is None + assert b.get_session(replacement["sess"].session_id) is replacement["sess"] + + async def test_shutdown_ends_stream(self, device, broker, monkeypatch): + monkeypatch.setattr(events_module.settings, "REMOTE_DEVICE_SESSION_IDLE_SECONDS", 30) + b, _ = broker + ticket = _ticket(b) + begin_shutdown() + status, _, chunks = await stream_then_disconnect( + _app(), _path(ticket), headers=AUTH, chunks_before_disconnect=10**6, timeout=3 + ) + assert status == 200 + assert b"".join(chunks) == b"" + assert b.get_session(ticket) is None + + async def test_client_disconnect_closes_session(self, device, broker, monkeypatch): + monkeypatch.setattr(events_module.settings, "REMOTE_DEVICE_SESSION_IDLE_SECONDS", 30) + b, _ = broker + ticket = _ticket(b) + status, _, chunks = await stream_then_disconnect( + _app(), _path(ticket), headers=AUTH, chunks_before_disconnect=1, timeout=3 + ) + assert status == 200 + assert chunks[0].startswith(b"event: invocation\n") + assert b.get_session(ticket) is None + assert DEVICE["id"] not in b._sessions_by_device + + +# ── machine-key signatures ────────────────────────────────────────────── + + +def _signed_headers(path: str, *, sign_path: str | None = None) -> tuple[str, dict]: + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey + + key = Ed25519PrivateKey.generate() + raw_pub = key.public_key().public_bytes( + serialization.Encoding.Raw, serialization.PublicFormat.Raw + ) + fingerprint = hashlib.sha256(raw_pub).hexdigest() + ts = str(int(time.time())) + payload = auth_module._canonical_payload("GET", sign_path or path, ts, b"") + headers = { + **AUTH, + "X-Device-Signature": base64.b64encode(key.sign(payload.encode("utf-8"))).decode(), + "X-Device-Timestamp": ts, + "X-Device-Machine-Key": fingerprint, + "X-Device-Machine-Pubkey": base64.b64encode(raw_pub).decode(), + } + return fingerprint, headers + + +@pytest.mark.unit +class TestSignature: + def test_signed_stream_request_accepted(self, device, broker, monkeypatch): + 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 + r = _open(ticket, headers=headers) + assert r.status_code == 200 + assert "event: invocation" in r.text + + def test_signature_for_another_path_rejected(self, device, broker, monkeypatch): + monkeypatch.setattr(auth_module.settings, "REMOTE_DEVICE_REQUIRE_SIGNATURE", True) + b, fake = broker + ticket = _ticket(b) + fingerprint, headers = _signed_headers(_path(ticket), sign_path="/api/devices/poll") + device.row["machine_pubkey_fingerprint"] = fingerprint + r = _open(ticket, headers=headers) + assert r.status_code == 401 + assert r.json() == {"success": False, "error": "invalid_signature"} + # Rejected before the session opened: the ticket is still unclaimed. + assert fake.get(f"dev:ticket:{DEVICE['id']}") == ticket.encode() + + def test_missing_signature_rejected(self, device, broker, monkeypatch): + monkeypatch.setattr(auth_module.settings, "REMOTE_DEVICE_REQUIRE_SIGNATURE", True) + b, _ = broker + r = _open(_ticket(b)) + assert r.status_code == 401 + assert r.json() == {"success": False, "error": "missing_signature"} + + +# ── broker async poll ─────────────────────────────────────────────────── + + +@pytest.mark.unit +@pytest.mark.asyncio +class TestNextCommandAsync: + async def test_returns_queued_envelope(self, async_broker_env): + b, _ = async_broker_env + _queue(b) + sess = b.register_session(DEVICE["id"], DEVICE["user_id"]) + envelope = await b.next_command_async(sess, timeout=0.01) + assert envelope["invocation_id"] == "inv_1" + + async def test_none_when_queue_empty(self, async_broker_env): + b, _ = async_broker_env + sess = b.register_session(DEVICE["id"], DEVICE["user_id"]) + assert await b.next_command_async(sess, timeout=0.01) is None + + @pytest.mark.parametrize("raw", [b"not json", b'["not", "a", "dict"]']) + async def test_drops_malformed_envelope(self, async_broker_env, raw): + b, fake = async_broker_env + fake.rpush(f"dev:cmd:{DEVICE['id']}", raw) + sess = b.register_session(DEVICE["id"], DEVICE["user_id"]) + assert await b.next_command_async(sess, timeout=0.01) is None + assert fake.llen(f"dev:cmd:{DEVICE['id']}") == 0 + + async def test_drops_reaped_invocation(self, async_broker_env): + b, fake = async_broker_env + _queue(b) + fake.delete("dev:inv:inv_1") + sess = b.register_session(DEVICE["id"], DEVICE["user_id"]) + assert await b.next_command_async(sess, timeout=0.01) is None + + async def test_none_when_async_redis_unavailable(self, async_broker_env, monkeypatch): + b, _ = async_broker_env + + async def _no_redis(): + return None + + monkeypatch.setattr("docsgpt.devices.broker.get_async_redis_instance", _no_redis) + sess = b.register_session(DEVICE["id"], DEVICE["user_id"]) + assert await b.next_command_async(sess, timeout=0.01) is None + + async def test_none_when_blpop_fails(self, async_broker_env, fake_redis, monkeypatch): + b, _ = async_broker_env + + class _Broken(AsyncFakeRedis): + async def blpop(self, key, timeout=0): + raise ConnectionError("reset by peer") + + broken = _Broken(fake_redis) + + async def _get(): + return broken + + monkeypatch.setattr("docsgpt.devices.broker.get_async_redis_instance", _get) + sess = b.register_session(DEVICE["id"], DEVICE["user_id"]) + assert await b.next_command_async(sess, timeout=0.01) is None + + +@pytest.mark.unit +def test_flask_blueprint_leaves_session_stream_to_asgi(): + from docsgpt.api.devices import devices_bp + + app = Flask(__name__) + app.register_blueprint(devices_bp) + rules = {rule.rule for rule in app.url_map.iter_rules()} + assert "/api/devices/sessions//events" not in rules + assert "/api/devices/poll" in rules + assert "/api/devices/sessions//invocations//ack" in rules + assert "/api/devices/sessions//invocations//output" in rules diff --git a/tests/devices/test_session_ticket.py b/tests/devices/test_session_ticket.py index c2778e93..ac4d38c9 100644 --- a/tests/devices/test_session_ticket.py +++ b/tests/devices/test_session_ticket.py @@ -4,23 +4,26 @@ Two layers: * Broker unit tests for ``claim_ticket`` / ``validate_ticket`` (issue, match, mismatch, eviction, absence). -* Route tests proving the real CLI loop still works: ``/poll`` issues a - ticket and ``session_events`` accepts *that* ticket, while a mismatched - ticket is rejected with ``410`` before any stream opens. +* Route tests proving the real CLI loop still works: ``/poll`` (Flask) issues + a ticket and the native-async stream accepts *that* ticket, while a + mismatched ticket is rejected with ``410`` before any stream opens. """ from __future__ import annotations -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from flask import Flask +from starlette.applications import Starlette +from starlette.testclient import TestClient from docsgpt.api.devices import auth as auth_module from docsgpt.api.devices import session as session_module +from docsgpt.api.devices import session_events as session_events_module from docsgpt.devices.broker import DeviceBroker -from .conftest import FakeRedis +from .conftest import AsyncFakeRedis, FakeRedis # --------------------------------------------------------------------------- @@ -104,6 +107,14 @@ def app(): return Flask(__name__) +@pytest.fixture(autouse=True) +def _short_idle(monkeypatch): + # An idle stream ends on its own, so the test client returns. + monkeypatch.setattr( + session_events_module.settings, "REMOTE_DEVICE_SESSION_IDLE_SECONDS", 0.1 + ) + + def _device_row() -> dict: return { "id": "dev_route", @@ -157,72 +168,91 @@ def _call(app: Flask, view, path: str, *args): c.stop() -def test_poll_to_sse_upgrade_with_issued_ticket_works(app): - """The legitimate loop: /poll issues a ticket, session_events accepts it.""" +def _open_stream(session_id: str): + client = TestClient(Starlette(routes=session_events_module.device_session_routes)) + ctxs = _patched_auth() + for c in ctxs: + c.start() + try: + return client.get( + f"/api/devices/sessions/{session_id}/events", + headers={"Authorization": "Bearer tok_good"}, + ) + finally: + for c in ctxs: + c.stop() + + +def _broker_patches(fake: FakeRedis, broker: DeviceBroker): + return [ + patch("docsgpt.devices.broker.get_redis_instance", return_value=fake), + patch( + "docsgpt.devices.broker.get_async_redis_instance", + AsyncMock(return_value=AsyncFakeRedis(fake)), + ), + patch.object(session_module, "get_broker", return_value=broker), + patch.object(session_events_module, "get_broker", return_value=broker), + ] + + +@pytest.fixture +def wired_broker(): fake = FakeRedis() broker = DeviceBroker() - with patch("docsgpt.devices.broker.get_redis_instance", return_value=fake): - # Queue work so /poll returns a ticket rather than 202. - broker.dispatch_invocation( - "dev_route", "user_route", - {"invocation_id": "inv_a", "action": "run_command"}, - ) - with patch.object(session_module, "get_broker", return_value=broker): - poll_resp = _call(app, session_module.poll, "/api/devices/poll") - assert poll_resp.status_code == 200 - payload = poll_resp.get_json() - ticket = payload["session_ticket"] - assert payload["session_url"] == f"/api/devices/sessions/{ticket}/events" - assert payload["expires_in"] == 30 - - # CLI opens the exact session_url it was handed. - sse_resp = _call( - app, - session_module.session_events, - f"/api/devices/sessions/{ticket}/events", - ticket, - ) - try: - assert sse_resp.status_code == 200 - assert sse_resp.mimetype == "text/event-stream" - finally: - sse_resp.close() + patches = _broker_patches(fake, broker) + for p in patches: + p.start() + try: + yield broker, fake + finally: + for p in patches: + p.stop() -def test_session_events_rejects_mismatched_ticket(app): +def test_poll_to_sse_upgrade_with_issued_ticket_works(app, wired_broker): + """The legitimate loop: /poll issues a ticket, the stream accepts it.""" + broker, _fake = wired_broker + # Queue work so /poll returns a ticket rather than 202. + broker.dispatch_invocation( + "dev_route", "user_route", + {"invocation_id": "inv_a", "action": "run_command"}, + ) + poll_resp = _call(app, session_module.poll, "/api/devices/poll") + assert poll_resp.status_code == 200 + payload = poll_resp.get_json() + ticket = payload["session_ticket"] + assert payload["session_url"] == f"/api/devices/sessions/{ticket}/events" + assert payload["expires_in"] == 30 + + # CLI opens the exact session_url it was handed. + sse_resp = _open_stream(ticket) + assert sse_resp.status_code == 200 + assert sse_resp.headers["content-type"].startswith("text/event-stream") + assert "event: invocation" in sse_resp.text + assert '"invocation_id": "inv_a"' in sse_resp.text + + +def test_session_events_rejects_mismatched_ticket(app, wired_broker): """A fabricated/mismatched session_id is 410 Gone, no stream opened.""" - fake = FakeRedis() - broker = DeviceBroker() - with patch("docsgpt.devices.broker.get_redis_instance", return_value=fake): - broker.dispatch_invocation( - "dev_route", "user_route", - {"invocation_id": "inv_b", "action": "run_command"}, - ) - with patch.object(session_module, "get_broker", return_value=broker): - # Poll issues the real ticket... - _call(app, session_module.poll, "/api/devices/poll") - # ...but the client opens a different one. - resp = _call( - app, - session_module.session_events, - "/api/devices/sessions/st_bogus/events", - "st_bogus", - ) + broker, fake = wired_broker + broker.dispatch_invocation( + "dev_route", "user_route", + {"invocation_id": "inv_b", "action": "run_command"}, + ) + # Poll issues the real ticket... + _call(app, session_module.poll, "/api/devices/poll") + issued = fake.get("dev:ticket:dev_route") + # ...but the client opens a different one. + resp = _open_stream("st_bogus") assert resp.status_code == 410 - assert resp.get_json()["error"] == "session_ticket_invalid" + assert resp.json()["error"] == "session_ticket_invalid" + # No session registered, so the real ticket is still unclaimed. + assert fake.get("dev:ticket:dev_route") == issued + assert "dev_route" not in broker._sessions_by_device -def test_session_events_rejects_when_never_polled(app): +def test_session_events_rejects_when_never_polled(wired_broker): """Opening the SSE stream without a prior poll is rejected (410).""" - fake = FakeRedis() - broker = DeviceBroker() - with patch("docsgpt.devices.broker.get_redis_instance", return_value=fake): - with patch.object(session_module, "get_broker", return_value=broker): - resp = _call( - app, - session_module.session_events, - "/api/devices/sessions/st_anything/events", - "st_anything", - ) + resp = _open_stream("st_anything") assert resp.status_code == 410 - assert resp.get_json()["error"] == "session_ticket_invalid" + assert resp.json()["error"] == "session_ticket_invalid" diff --git a/tests/test_asgi.py b/tests/test_asgi.py index 44c47c86..aa45b95c 100644 --- a/tests/test_asgi.py +++ b/tests/test_asgi.py @@ -168,3 +168,92 @@ def test_wsgi_threadpool_sized_from_settings(): assert asgi._WSGI_THREADPOOL == int(settings.WSGI_THREADPOOL_WORKERS) assert asgi._WSGI_THREADPOOL >= 64 + + +_UUID = "67d65e8f-e7fb-4df1-9e6e-99ea6c830206" + + +@pytest.mark.unit +@pytest.mark.parametrize( + "path, endpoint_name", + [ + ("/api/events", "stream_events"), + ("/api/devices/sessions/st_abc/events", "device_session_events"), + (f"/api/artifacts/{_UUID}/download", "download_artifact"), + (f"/api/messages/{_UUID}/events", "stream_message_events"), + ], +) +def test_long_lived_routes_served_on_event_loop_not_flask(path, endpoint_name): + """Held-open GETs must resolve to a Starlette route ahead of the Flask + mount; otherwise each one pins a WSGI threadpool slot.""" + from starlette.routing import Match + + from docsgpt.asgi import asgi_app + + scope = { + "type": "http", + "method": "GET", + "path": path, + "root_path": "", + "headers": [], + "query_string": b"", + } + for route in asgi_app.routes: + match, _ = route.matches(scope) + if match is Match.FULL: + endpoint = getattr(route, "endpoint", None) + assert endpoint is not None, f"{path} fell through to {route!r}" + assert endpoint.__name__ == endpoint_name + return + pytest.fail(f"no route matched {path}") + + +@pytest.mark.unit +def test_flask_no_longer_registers_moved_routes(): + from docsgpt.app import app as flask_app + + rules = {rule.rule for rule in flask_app.url_map.iter_rules()} + assert "/api/events" not in rules + assert "/api/devices/sessions//events" not in rules + assert "/api/artifacts//download" not in rules + # Their siblings stay on Flask. + assert "/api/devices/poll" in rules + assert "/api/artifacts/" in rules + + +@pytest.mark.unit +def test_event_stream_served_through_full_asgi_stack(monkeypatch): + """Auth, CORS and the SSE route work together in the real app.""" + from unittest.mock import AsyncMock + + from starlette.testclient import TestClient + + from docsgpt.asgi import asgi_app + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "AUTH_TYPE", None) + monkeypatch.setattr( + "docsgpt.api.events.routes.get_async_redis_instance", AsyncMock(return_value=None) + ) + monkeypatch.setattr( + "docsgpt.streaming.async_broadcast_channel.get_async_redis_instance", + AsyncMock(return_value=None), + ) + with TestClient(asgi_app) as client: + r = client.get("/api/events", headers={"Origin": "http://example.com"}) + assert r.status_code == 200 + assert r.headers["x-sse-transport"] == "async" + assert r.headers.get("access-control-allow-origin") == "*" + assert r.text.startswith(": connected") + + +@pytest.mark.unit +def test_artifact_download_served_through_full_asgi_stack(): + from starlette.testclient import TestClient + + from docsgpt.asgi import asgi_app + + with TestClient(asgi_app) as client: + r = client.get("/api/artifacts/not-a-uuid/download") + assert r.status_code == 404 + assert r.json() == {"success": False, "message": "Artifact not found"} diff --git a/tests/test_async_broadcast_channel.py b/tests/test_async_broadcast_channel.py new file mode 100644 index 00000000..49865cf1 --- /dev/null +++ b/tests/test_async_broadcast_channel.py @@ -0,0 +1,168 @@ +"""Tests for ``AsyncTopic.subscribe``'s liveness probe. + +A pub/sub socket that NAT/IPVS dropped silently never errors: ``get_message`` +keeps timing out exactly as if the channel were idle. The sync subscriber +bounds this with ``socket_timeout`` (see tests/test_pubsub_timeouts.py). The +async subscriber sends a PING after a quiet spell and ends the subscription +when nothing comes back, so the SSE client reconnects and replays from its +cursor instead of silently missing events. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import anyio +import pytest + +from docsgpt.streaming.async_broadcast_channel import AsyncTopic + +_AREDIS = "docsgpt.streaming.async_broadcast_channel.get_async_redis_instance" + +IDLE = object() + + +class _FakePubSub: + """Scripted ``redis.asyncio`` pubsub. + + ``script`` entries are message dicts (returned at once) or ``IDLE`` (a + ``get_message`` poll that times out). After the script, every poll is idle. + A PING queues a pong reply unless ``answer_pings`` is off. + """ + + def __init__(self, script=(), *, answer_pings=True, ping_error=None): + self._script = [{"type": "subscribe", "channel": b"t", "data": 1}, *script] + self.answer_pings = answer_pings + self.ping_error = ping_error + self.pings = 0 + self._pong_pending = False + self.unsubscribed = False + self.closed = False + + async def subscribe(self, name): + return None + + async def get_message(self, timeout=0.0): + if self._pong_pending: + self._pong_pending = False + return {"type": "pong", "pattern": None, "channel": None, "data": b"liveness"} + if self._script: + item = self._script.pop(0) + if item is not IDLE: + return item + await anyio.sleep(timeout) + return None + + async def ping(self, message=None): + self.pings += 1 + if self.ping_error is not None: + raise self.ping_error + if self.answer_pings: + self._pong_pending = True + return True + + async def unsubscribe(self, name): + self.unsubscribed = True + + async def aclose(self): + self.closed = True + + +def _redis_with(pubsub: _FakePubSub) -> AsyncMock: + client = MagicMock() + client.pubsub.return_value = pubsub + return AsyncMock(return_value=client) + + +@pytest.mark.unit +@pytest.mark.asyncio +class TestLivenessProbe: + async def test_unanswered_ping_ends_subscription(self): + pubsub = _FakePubSub(answer_pings=False) + ticks = 0 + with patch(_AREDIS, _redis_with(pubsub)): + with anyio.fail_after(2): + async for payload in AsyncTopic("t").subscribe( + poll_timeout=0.01, liveness_timeout=0.05 + ): + assert payload is None + ticks += 1 + assert ticks > 0 + assert pubsub.pings == 1 + assert pubsub.unsubscribed + assert pubsub.closed + + async def test_answered_pings_keep_subscription_open(self): + pubsub = _FakePubSub(answer_pings=True) + with patch(_AREDIS, _redis_with(pubsub)): + agen = AsyncTopic("t").subscribe(poll_timeout=0.01, liveness_timeout=0.03) + try: + with anyio.fail_after(2): + # Roughly 0.4s of idle polls: many probe windows, and the + # pong replies are consumed rather than yielded. + for _ in range(40): + assert await agen.__anext__() is None + finally: + await agen.aclose() + assert pubsub.pings >= 2 + assert pubsub.closed + + async def test_published_traffic_counts_as_liveness(self): + message = {"type": "message", "channel": b"t", "data": b"event"} + # Traffic every 3 polls (0.03s) never leaves a 0.05s quiet window. + script = [message, IDLE, IDLE] * 10 + pubsub = _FakePubSub(script, answer_pings=False) + received = [] + with patch(_AREDIS, _redis_with(pubsub)): + with anyio.fail_after(3): + async for payload in AsyncTopic("t").subscribe( + poll_timeout=0.01, liveness_timeout=0.05 + ): + if payload is not None: + received.append(payload) + assert pubsub.pings == 0 + assert received == [b"event"] * 10 + # The only probe went out after the traffic stopped, and it ended the stream. + assert pubsub.pings == 1 + + async def test_ping_failure_ends_subscription(self): + pubsub = _FakePubSub(ping_error=ConnectionError("broken pipe")) + with patch(_AREDIS, _redis_with(pubsub)): + with anyio.fail_after(2): + async for _payload in AsyncTopic("t").subscribe( + poll_timeout=0.01, liveness_timeout=0.03 + ): + pass + assert pubsub.pings == 1 + assert pubsub.closed + + async def test_pool_exhaustion_warns_with_setting_name(self): + from redis.exceptions import MaxConnectionsError + + pubsub = _FakePubSub() + + async def _exhausted(name): + raise MaxConnectionsError("Too many connections") + + pubsub.subscribe = _exhausted + with patch(_AREDIS, _redis_with(pubsub)), patch( + "docsgpt.streaming.async_broadcast_channel.logger" + ) as log: + items = [payload async for payload in AsyncTopic("t").subscribe(poll_timeout=0.01)] + assert items == [] + assert pubsub.closed + assert log.warning.called + assert "ASYNC_REDIS_MAX_CONNECTIONS" in str(log.warning.call_args) + log.exception.assert_not_called() + + async def test_no_probe_when_liveness_disabled(self): + pubsub = _FakePubSub(answer_pings=False) + with patch(_AREDIS, _redis_with(pubsub)): + agen = AsyncTopic("t").subscribe(poll_timeout=0.005) + try: + with anyio.fail_after(2): + for _ in range(30): + assert await agen.__anext__() is None + finally: + await agen.aclose() + assert pubsub.pings == 0 diff --git a/tests/test_async_redis.py b/tests/test_async_redis.py new file mode 100644 index 00000000..25791c49 --- /dev/null +++ b/tests/test_async_redis.py @@ -0,0 +1,67 @@ +"""Tests for the async Redis client behind the event-loop routes. + +Each open event-loop stream (a notification tab, a chat reconnect, a device +session) holds one pooled connection for its whole life. redis-py's pool +defaults to 100 connections, which one worker outgrows long before its event +loop does, so the pool size must come from settings. +""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +import docsgpt.streaming.async_redis as async_redis +from docsgpt.core.settings import Settings + +_FROM_URL = "docsgpt.streaming.async_redis.aioredis.Redis.from_url" + + +@pytest.fixture(autouse=True) +def _reset_singleton(): + async_redis._async_redis = None + async_redis._creation_failed = False + yield + async_redis._async_redis = None + async_redis._creation_failed = False + + +@pytest.mark.unit +@pytest.mark.asyncio +class TestGetAsyncRedisInstance: + async def test_pool_sized_from_settings(self, monkeypatch): + monkeypatch.setattr(async_redis.settings, "ASYNC_REDIS_MAX_CONNECTIONS", 4321) + with patch(_FROM_URL) as from_url: + client = await async_redis.get_async_redis_instance() + assert client is from_url.return_value + kwargs = from_url.call_args.kwargs + assert kwargs["max_connections"] == 4321 + assert kwargs["socket_connect_timeout"] == 2 + assert kwargs["health_check_interval"] == 10 + + 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) + monkeypatch.setattr(async_redis.settings, "CACHE_REDIS_URL", "redis://127.0.0.1:1/0") + client = await async_redis.get_async_redis_instance() + try: + assert client.connection_pool.max_connections == 777 + finally: + await client.aclose() + + async def test_default_pool_outgrows_redis_py_default(self): + assert Settings.model_fields["ASYNC_REDIS_MAX_CONNECTIONS"].default > 100 + + async def test_singleton(self): + with patch(_FROM_URL) as from_url: + first = await async_redis.get_async_redis_instance() + second = await async_redis.get_async_redis_instance() + assert first is second + assert from_url.call_count == 1 + + async def test_invalid_url_marks_failed_and_stops_retrying(self): + with patch(_FROM_URL, side_effect=ValueError("bad url")) as from_url: + assert await async_redis.get_async_redis_instance() is None + assert await async_redis.get_async_redis_instance() is None + assert from_url.call_count == 1 From e889392eba6494366f7f0666ea1c1e07db06a25e Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 14 Sep 2026 10:19:59 +0100 Subject: [PATCH 2/7] 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"] From fb66a0b23027ff811d766019c5dda001523d4d80 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 14 Sep 2026 12:52:15 +0100 Subject: [PATCH 3/7] 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") From 72904b0b9598d0ea1527c4dab4d221e866da4172 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 14 Sep 2026 14:06:59 +0100 Subject: [PATCH 4/7] fix: Make ticket claim atomic. --- docsgpt/api/devices/session_events.py | 15 ++++-------- docsgpt/devices/broker.py | 35 +++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 10 deletions(-) diff --git a/docsgpt/api/devices/session_events.py b/docsgpt/api/devices/session_events.py index afb137e3..65b47f2f 100644 --- a/docsgpt/api/devices/session_events.py +++ b/docsgpt/api/devices/session_events.py @@ -13,7 +13,7 @@ import json import logging import time from functools import partial -from typing import AsyncIterator, Optional +from typing import AsyncIterator import anyio from starlette.requests import Request @@ -63,13 +63,6 @@ def _signed_path(request: Request) -> str: 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") @@ -132,11 +125,13 @@ async def device_session_events(request: Request) -> Response: return _error(*failure) log_context.bind(user_id=device["user_id"]) broker = get_broker() - # Claiming consumes the ticket, so it happens before the response opens and + # Redeeming 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) + sess = await anyio.to_thread.run_sync( + broker.redeem_ticket, device["id"], device["user_id"], session_id + ) if sess is None: return _error("session_ticket_invalid", 410) return ClosingStreamingResponse( diff --git a/docsgpt/devices/broker.py b/docsgpt/devices/broker.py index 781df24b..09d55360 100644 --- a/docsgpt/devices/broker.py +++ b/docsgpt/devices/broker.py @@ -61,6 +61,16 @@ def _ticket_key(device_id: str) -> str: return f"dev:ticket:{device_id}" +# Deletes the device's ticket iff it still holds the presented one. Runs +# server-side so two requests racing with one ticket can't both redeem it. +_REDEEM_TICKET_LUA = """ +if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('DEL', KEYS[1]) +end +return 0 +""" + + def _inv_key(invocation_id: str) -> str: return f"dev:inv:{invocation_id}" @@ -142,6 +152,31 @@ class DeviceBroker: except Exception: logger.exception("ticket lookup failed for %s", device_id) session_id = _as_str(issued) if issued else f"st_{uuid.uuid4().hex}" + return self._adopt_session(session_id, device_id, user_id) + + def redeem_ticket(self, device_id: str, user_id: str, ticket: str) -> Optional[SessionState]: + """Consume the device's poll-issued ``ticket`` and open its session under it. + + Returns ``None``, and opens nothing, unless ``ticket`` is still the + device's unexpired ticket. The compare and delete are one Redis step, + so a replay racing the first redeem can't replace its session. + """ + if not ticket: + return None + redis = get_redis_instance() + if redis is None: + return None + try: + redeemed = redis.eval(_REDEEM_TICKET_LUA, 1, _ticket_key(device_id), ticket) + except Exception: + logger.exception("ticket redeem failed for %s", device_id) + return None + if not redeemed: + return None + return self._adopt_session(ticket, device_id, user_id) + + def _adopt_session(self, session_id: str, device_id: str, user_id: str) -> SessionState: + """Make a new session the device's live one, closing the one it replaces.""" sess = SessionState( session_id=session_id, device_id=device_id, user_id=user_id ) From 1bf7380aa5df0875aa926d0f3fdae3fd42dabbf4 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 14 Sep 2026 14:12:21 +0100 Subject: [PATCH 5/7] fix: mini test --- tests/devices/conftest.py | 8 ++++++ tests/devices/test_session_ticket.py | 43 ++++++++++++++++++++++++++-- 2 files changed, 49 insertions(+), 2 deletions(-) diff --git a/tests/devices/conftest.py b/tests/devices/conftest.py index bf2663bc..c0a3f316 100644 --- a/tests/devices/conftest.py +++ b/tests/devices/conftest.py @@ -86,6 +86,14 @@ class FakeRedis: def expire(self, key, ttl): # TTL is not simulated. return True + def eval(self, _script, _numkeys, key, expected): + """Model the broker's one Lua script: delete ``key`` iff it holds ``expected``.""" + with self._lock: + if self.kv.get(key) != _b(expected): + return 0 + del self.kv[key] + return 1 + # -- lists -------------------------------------------------------- def llen(self, key): with self._lock: diff --git a/tests/devices/test_session_ticket.py b/tests/devices/test_session_ticket.py index ac4d38c9..98f1381e 100644 --- a/tests/devices/test_session_ticket.py +++ b/tests/devices/test_session_ticket.py @@ -2,8 +2,8 @@ Two layers: -* Broker unit tests for ``claim_ticket`` / ``validate_ticket`` (issue, match, - mismatch, eviction, absence). +* Broker unit tests for ``claim_ticket`` / ``validate_ticket`` / + ``redeem_ticket`` (issue, match, mismatch, eviction, absence, single use). * Route tests proving the real CLI loop still works: ``/poll`` (Flask) issues a ticket and the native-async stream accepts *that* ticket, while a mismatched ticket is rejected with ``410`` before any stream opens. @@ -99,6 +99,45 @@ def test_register_session_consumes_issued_ticket_as_session_id(broker_env): assert fake.get("dev:ticket:dev_x") is None +def test_redeem_ticket_opens_session_under_the_ticket(broker_env): + broker, fake = broker_env + _queue_work(broker) + ticket = broker.claim_ticket("dev_x", 30) + sess = broker.redeem_ticket("dev_x", "user_x", ticket) + assert sess.session_id == ticket + assert broker.get_session(ticket) is sess + assert fake.get("dev:ticket:dev_x") is None + + +def test_redeem_ticket_succeeds_only_once(broker_env): + # Two requests racing with one ticket: the loser must not open a session + # that replaces the winner's live stream. + broker, _fake = broker_env + _queue_work(broker) + ticket = broker.claim_ticket("dev_x", 30) + first = broker.redeem_ticket("dev_x", "user_x", ticket) + assert broker.redeem_ticket("dev_x", "user_x", ticket) is None + assert not first.closed.is_set() + assert broker._sessions_by_device["dev_x"] is first + + +def test_redeem_ticket_rejects_other_ticket_and_keeps_issued_one(broker_env): + broker, fake = broker_env + _queue_work(broker) + ticket = broker.claim_ticket("dev_x", 30) + assert broker.redeem_ticket("dev_x", "user_x", "st_not_the_one") is None + assert broker.redeem_ticket("dev_x", "user_x", "") is None + assert fake.get("dev:ticket:dev_x") == ticket.encode() + assert "dev_x" not in broker._sessions_by_device + + +def test_redeem_ticket_rejects_when_redis_unavailable(monkeypatch): + monkeypatch.setattr("docsgpt.devices.broker.get_redis_instance", lambda: None) + broker = DeviceBroker() + assert broker.redeem_ticket("dev_x", "user_x", "st_anything") is None + assert "dev_x" not in broker._sessions_by_device + + # --------------------------------------------------------------------------- # Route tests (poll -> SSE upgrade) # --------------------------------------------------------------------------- From 3030aba39d4580c6a0d65dd809d75e9575738d76 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 14 Sep 2026 17:27:13 +0100 Subject: [PATCH 6/7] fix: handle non text uploads more carefully --- docs/content/Guides/ocr.mdx | 8 + docsgpt/llm/handlers/base.py | 31 +++ docsgpt/parser/file/anydoc_parser.py | 3 +- docsgpt/parser/file/base_parser.py | 11 ++ docsgpt/parser/file/bulk.py | 18 +- docsgpt/parser/file/image_parser.py | 47 ++++- docsgpt/worker.py | 71 ++++++- frontend/src/components/MessageInput.tsx | 25 +++ .../message-input/AttachmentChipList.tsx | 27 +++ .../attachmentReadability.test.ts | 76 ++++++++ .../message-input/attachmentReadability.ts | 33 ++++ frontend/src/locale/en.json | 3 +- frontend/src/upload/uploadSlice.test.ts | 18 ++ frontend/src/upload/uploadSlice.ts | 13 ++ tests/llm/handlers/test_llm_handlers.py | 68 +++++++ tests/parser/file/test_anydoc_parser.py | 10 +- tests/parser/file/test_bulk.py | 40 ++++ .../file/test_image_vision_conversion.py | 64 +++++++ tests/parser/file/test_ocr_parser.py | 9 +- tests/test_upload_limits.py | 17 +- tests/worker/test_attachment_worker.py | 177 +++++++++++++++++- 21 files changed, 729 insertions(+), 40 deletions(-) create mode 100644 frontend/src/components/message-input/attachmentReadability.test.ts create mode 100644 frontend/src/components/message-input/attachmentReadability.ts create mode 100644 tests/parser/file/test_image_vision_conversion.py diff --git a/docs/content/Guides/ocr.mdx b/docs/content/Guides/ocr.mdx index b49c28fa..2b185141 100644 --- a/docs/content/Guides/ocr.mdx +++ b/docs/content/Guides/ocr.mdx @@ -150,6 +150,14 @@ takes over as the fallback parser. If that fallback also extracts almost nothing — OCR off, or no engine available — the upload fails with a clear message instead of silently indexing an empty document. +Chat attachments are the exception: a scanned PDF attached from the message +box is kept with no text (`extraction.status: no_text`), because the file +itself is what a model reads. Models that take PDFs get the document; models +that take only images get its pages as images; a text-only model is told it +cannot read the file, and the message box warns before sending. Images work +the same way with OCR off, and TIFF and BMP attachments are stored as PNG, +since model providers do not accept those formats. + Mixed documents — text pages with scanned pages among them — convert through anydoc, which reads the text pages and skips the scanned ones. With OCR on, DocsGPT probes every page's text layer, OCRs the pages that have diff --git a/docsgpt/llm/handlers/base.py b/docsgpt/llm/handlers/base.py index ec80ab3d..0cd4c84c 100644 --- a/docsgpt/llm/handlers/base.py +++ b/docsgpt/llm/handlers/base.py @@ -324,6 +324,25 @@ class LLMHandler(ABC): return images_data + @staticmethod + def _is_unreadable_visual(attachment: Dict) -> bool: + """Whether an unsupported attachment is an image or PDF with no text to use. + + Images and scanned PDFs (``extraction.status == "no_text"``) reach this + path only on a model that takes neither format natively. They carry no + text, so dropping them would let the model answer as if nothing had + been attached. + """ + mime_type = attachment.get("mime_type") or "" + if not (mime_type.startswith("image/") or mime_type == "application/pdf"): + return False + extraction = (attachment.get("metadata") or {}).get("extraction") or {} + status = extraction.get("status") + if status == "no_text": + return True + content = attachment.get("content") + return status in (None, "ok") and content is not None and not str(content).strip() + def _append_unsupported_attachments( self, messages: List[Dict], attachments: List[Dict] ) -> List[Dict]: @@ -339,8 +358,12 @@ class LLMHandler(ABC): """ prepared_messages = messages.copy() attachment_texts = [] + unreadable_names = [] for attachment in attachments: + if self._is_unreadable_visual(attachment): + unreadable_names.append(attachment.get("filename") or "an attached file") + continue # ``metadata.extraction`` records what parsing actually did. # Rows predating it have no ``extraction`` key and pass; rows # whose extraction failed must not reach the prompt (a PG row @@ -374,6 +397,14 @@ class LLMHandler(ABC): "Scope any whole-document claims to this portion.]\n\n" ) attachment_texts.append(f"Attached file content:\n\n{note}{content}") + if unreadable_names: + listed = ", ".join(f'"{name}"' for name in unreadable_names) + attachment_texts.append( + f"[NOTE: The user attached {listed}, but the current model cannot read " + "images or scanned PDFs, so the file contents are not available to you. " + "Tell the user you cannot see the file and suggest switching to a model " + "that supports images or PDFs. Do not guess what it contains.]" + ) if attachment_texts: combined_text = "\n\n".join(attachment_texts) diff --git a/docsgpt/parser/file/anydoc_parser.py b/docsgpt/parser/file/anydoc_parser.py index 22b5271f..3e964671 100644 --- a/docsgpt/parser/file/anydoc_parser.py +++ b/docsgpt/parser/file/anydoc_parser.py @@ -23,6 +23,7 @@ from docsgpt.core.settings import settings from docsgpt.parser.file.base_parser import ( BaseParser, DocumentParseError, + NoTextLayerError, delegate_parse, module_available, ) @@ -424,7 +425,7 @@ class AnydocParser(BaseParser): "tesseract binary or a DeepSeek-OCR endpoint (OCR_ENGINE), or " "through the optional docling extra when installed." ) - raise DocumentParseError( + raise NoTextLayerError( f"{path.name} appears to be a scanned PDF (no text layer), and " f"{type(fallback).__name__} extracted almost nothing{hint}" ) diff --git a/docsgpt/parser/file/base_parser.py b/docsgpt/parser/file/base_parser.py index 7b8907d1..32c81a53 100644 --- a/docsgpt/parser/file/base_parser.py +++ b/docsgpt/parser/file/base_parser.py @@ -26,6 +26,17 @@ class DocumentParseError(Exception): """ +class NoTextLayerError(DocumentParseError): + """A readable document with no text to extract, such as a scanned PDF. + + Still a failed parse wherever text is the point (source ingestion needs + something to embed), so it subclasses ``DocumentParseError`` and callers + that do not know about it keep failing loudly. The file itself is intact, + though: a chat attachment can still go to a model that reads the format + natively, which is why the attachment worker catches this type. + """ + + class BaseParser: """Base class for all parsers.""" diff --git a/docsgpt/parser/file/bulk.py b/docsgpt/parser/file/bulk.py index 1332a9c5..4828f893 100644 --- a/docsgpt/parser/file/bulk.py +++ b/docsgpt/parser/file/bulk.py @@ -105,9 +105,6 @@ def _image_entries(factory: Callable[[], BaseParser], suffixes) -> Dict[str, Bas return {suffix: factory() for suffix in suffixes} -_LEGACY_IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg") - - def _legacy_file_extractor(pdf_text_fast_path: bool = False, ocr_enabled: bool = False) -> Dict[str, BaseParser]: """Parser map that needs neither docling nor anydoc. @@ -124,7 +121,10 @@ def _legacy_file_extractor(pdf_text_fast_path: bool = False, ocr_enabled: bool = from docsgpt.parser.file.ocr_parser import IMAGE_SUFFIXES pdf_parser: BaseParser = PDFParser() - images: Dict[str, BaseParser] = _image_entries(ImageParser, _LEGACY_IMAGE_SUFFIXES) + # Every image suffix gets an entry even without OCR: the image itself is + # what a vision model is sent, so a missing entry only means the upload is + # refused before any model sees it (.webp was, in production). + images: Dict[str, BaseParser] = _image_entries(ImageParser, sorted(IMAGE_SUFFIXES)) if ocr_enabled: native = _native_ocr_parsers() if native is not None: @@ -490,6 +490,7 @@ class SimpleDirectoryReader(BaseReader): data: Union[str, List[str]] = "" data_list: List[str] = [] metadata_list = [] + first_error: Optional[DocumentParseError] = None self.file_token_counts = {} self.failed_files = [] @@ -529,6 +530,7 @@ class SimpleDirectoryReader(BaseReader): except DocumentParseError as e: logging.warning(f"Skipping unreadable file {input_file.name}: {e}") self.failed_files.append((input_file, str(e))) + first_error = first_error or e report_progress(file_index + 1) continue @@ -579,10 +581,10 @@ class SimpleDirectoryReader(BaseReader): # ingest tasks) treat it as terminal and tell the user. if self.failed_files and not data_list: if len(self.failed_files) == 1: - # Single-file read (the attachment path): the parser's own - # message reaches the user verbatim, so don't wrap it in - # "None of the 1 file(s)…". - raise DocumentParseError(self.failed_files[0][1]) + # Single-file read (the attachment path): re-raise the parser's + # own exception, so its message reaches the user verbatim and + # its type (``NoTextLayerError`` for a scan) reaches the caller. + raise first_error names = ", ".join(p.name for p, _ in self.failed_files[:5]) raise DocumentParseError( f"None of the {len(self.failed_files)} files could be parsed: " diff --git a/docsgpt/parser/file/image_parser.py b/docsgpt/parser/file/image_parser.py index f4dad9ef..e3382682 100644 --- a/docsgpt/parser/file/image_parser.py +++ b/docsgpt/parser/file/image_parser.py @@ -1,14 +1,51 @@ """Image parser. -Contains parser for .png, .jpg, .jpeg files. +Contains the parser for image files (.png, .jpg, .jpeg, .tiff, .tif, .bmp, +.webp) and the PNG re-encoding for the formats model providers reject. """ +import io from pathlib import Path +from typing import BinaryIO, Dict, Tuple, Union + import requests -from typing import Dict, Union -from docsgpt.parser.file.base_parser import BaseParser from docsgpt.core.settings import settings +from docsgpt.parser.file.base_parser import BaseParser, DocumentParseError + +# Image types the vision APIs refuse: OpenAI and Anthropic take png, jpeg, +# webp and gif only. A chat attachment in one of these formats is re-encoded +# to PNG before it is stored for the model. +VISION_CONVERTIBLE_MIME_TYPES = frozenset({"image/tiff", "image/bmp", "image/x-ms-bmp"}) + + +def convert_image_to_png(file_obj: BinaryIO) -> Tuple[bytes, int]: + """Re-encode the first frame of an image as PNG. + + Args: + file_obj: Readable binary stream of the source image. + + Returns: + Tuple[bytes, int]: The PNG bytes, and how many frames (pages) the + source had. Only the first is kept. + + Raises: + DocumentParseError: If the bytes are not an image Pillow can decode. + """ + from PIL import Image, UnidentifiedImageError + + try: + with Image.open(file_obj) as image: + frames = getattr(image, "n_frames", 1) + image.seek(0) + frame = image + if frame.mode not in ("RGB", "RGBA", "L", "LA"): + frame = frame.convert("RGBA" if "A" in frame.getbands() else "RGB") + out = io.BytesIO() + frame.save(out, format="PNG") + except (UnidentifiedImageError, OSError) as exc: + raise DocumentParseError(f"Could not read the image: {exc}") from exc + return out.getvalue(), frames class ImageParser(BaseParser): @@ -24,8 +61,8 @@ class ImageParser(BaseParser): # alternatively you can use local vision capable LLM with open(file, "rb") as file_loaded: files = {'file': file_loaded} - response = requests.post(doc2md_service, files=files, timeout=100) - data = response.json()["markdown"] + response = requests.post(doc2md_service, files=files, timeout=100) + data = response.json()["markdown"] else: data = "" return data diff --git a/docsgpt/worker.py b/docsgpt/worker.py index b300f1a3..3ce7f15e 100755 --- a/docsgpt/worker.py +++ b/docsgpt/worker.py @@ -1,4 +1,5 @@ import datetime +import io import json import logging import mimetypes @@ -23,8 +24,13 @@ from docsgpt.parser.embedding_pipeline import ( assert_index_complete, embed_and_store_documents, ) +from docsgpt.parser.file.base_parser import NoTextLayerError from docsgpt.parser.file.bulk import SimpleDirectoryReader, get_default_file_extractor from docsgpt.parser.file.constants import SUPPORTED_SOURCE_EXTENSIONS +from docsgpt.parser.file.image_parser import ( + VISION_CONVERTIBLE_MIME_TYPES, + convert_image_to_png, +) from docsgpt.parser.remote.remote_creator import ( RemoteCreator, normalize_remote_data, @@ -1598,6 +1604,44 @@ def _reject_attachment_zip_bomb(local_path: str) -> None: raise AttachmentRejectedError(reason) +def _readable_without_text(filename: str) -> bool: + """Whether a model can read an attachment from its original file alone. + + PDFs and images are sent to models as the file itself (a PDF as page + images on image-only models); every other format needs extracted text. + + Args: + filename: The upload's original filename. + + Returns: + bool: True for PDFs and images. + """ + mime_type = mimetypes.guess_type(filename)[0] or "" + return mime_type == "application/pdf" or mime_type.startswith("image/") + + +def _store_png_copy(storage, relative_path: str, mime_type: str) -> tuple[str, dict]: + """Store a PNG copy, beside the original, of an image the providers reject. + + Args: + storage: The storage backend holding the upload. + relative_path: Storage path of the original image. + mime_type: The original's mime type, e.g. ``image/tiff``. + + Returns: + tuple[str, dict]: The PNG's storage path, and the conversion record + kept in the attachment's metadata. + + Raises: + DocumentParseError: If the image cannot be decoded. + """ + with storage.get_file(relative_path) as source: + png_bytes, frames = convert_image_to_png(source) + png_path = f"{os.path.splitext(relative_path)[0]}.png" + storage.save_file(io.BytesIO(png_bytes), png_path) + return png_path, {"from": mime_type, "to": "image/png", "frames": frames} + + def _bounded_attachment_copy(local_path: str) -> tuple[str, bool]: """Bound how much of a text attachment reaches the parser. @@ -1794,7 +1838,24 @@ def attachment_worker(self, file_info, user): except OSError: pass - attachment_document = storage.process_file(relative_path, _parse_local_file) + extraction_status = "ok" + no_text_reason = None + try: + attachment_document = storage.process_file(relative_path, _parse_local_file) + except NoTextLayerError as exc: + # A PDF with no text layer is still a usable attachment: models + # that read PDFs natively, or as page images, are sent the stored + # file rather than its text. Keep it with nothing to inline; the + # LLM handler names it to any model that cannot read it. + if not _readable_without_text(filename): + raise + logging.info( + f"Attachment {filename} has no text layer; keeping it for native reading", + extra={"user": user}, + ) + attachment_document = Document(text="", extra_info={}) + extraction_status = "no_text" + no_text_reason = str(exc)[:1024] content = attachment_document.text # A fast-path parser may have delegated to its fallback for this file, # so record the engine that actually ran rather than the one selected. @@ -1823,11 +1884,12 @@ def attachment_worker(self, file_info, user): metadata = { **metadata, "extraction": { - "status": "ok", + "status": extraction_status, "parser": parser_name, "truncated": truncated, "original_tokens": original_tokens, "stored_tokens": token_count, + **({"reason": no_text_reason} if no_text_reason else {}), }, } @@ -1847,6 +1909,10 @@ def attachment_worker(self, file_info, user): ) mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream" + if mime_type in VISION_CONVERTIBLE_MIME_TYPES: + relative_path, conversion = _store_png_copy(storage, relative_path, mime_type) + mime_type = conversion["to"] + metadata = {**metadata, "image_conversion": conversion} _upsert_attachment_row( user, @@ -1873,6 +1939,7 @@ def attachment_worker(self, file_info, user): "filename": filename, "token_count": token_count, "mime_type": mime_type, + "extraction_status": extraction_status, }, scope={"kind": "attachment", "id": str(attachment_id)}, ) diff --git a/frontend/src/components/MessageInput.tsx b/frontend/src/components/MessageInput.tsx index ddff0c3f..c7c29171 100644 --- a/frontend/src/components/MessageInput.tsx +++ b/frontend/src/components/MessageInput.tsx @@ -28,6 +28,7 @@ import { import { ActiveState, Doc } from '../models/misc'; import { selectSelectedDocs, + selectSelectedModel, selectSourceDocs, selectToken, setSelectedDocs, @@ -47,6 +48,7 @@ import { ToolsTrigger, } from './message-input'; import { useArmedSend } from './message-input/armedSend'; +import { cannotReadAttachment } from './message-input/attachmentReadability'; import { handleAbort } from '../conversation/conversationSlice'; import { AUDIO_FILE_ACCEPT_ATTR, @@ -345,6 +347,21 @@ export default function MessageInput({ const sourceDocs = useSelector(selectSourceDocs); const token = useSelector(selectToken); const attachments = useSelector(selectAttachments); + const selectedModel = useSelector(selectSelectedModel); + const unreadableAttachmentIds = useMemo( + () => + new Set( + attachments + .filter((attachment) => + cannotReadAttachment( + attachment, + selectedModel?.supported_attachment_types, + ), + ) + .map((attachment) => attachment.id), + ), + [attachments, selectedModel], + ); const dispatch = useDispatch(); const store = useStore(); @@ -453,6 +470,12 @@ export default function MessageInput({ ...(Number.isFinite(tokenCount) ? { token_count: tokenCount } : {}), + ...(typeof payload.mime_type === 'string' + ? { mimeType: payload.mime_type } + : {}), + ...(typeof payload.extraction_status === 'string' + ? { extractionStatus: payload.extraction_status } + : {}), }, }), ); @@ -1662,6 +1685,8 @@ export default function MessageInput({ onDragStart={handleDragStart} onDragOver={handleDragOver} onDropOn={handleDropOn} + unreadableIds={unreadableAttachmentIds} + modelName={selectedModel?.display_name} /> {sendArmed && sendReadiness.state === 'waiting' && ( diff --git a/frontend/src/components/message-input/AttachmentChipList.tsx b/frontend/src/components/message-input/AttachmentChipList.tsx index c90e7faf..4d1a9b3e 100644 --- a/frontend/src/components/message-input/AttachmentChipList.tsx +++ b/frontend/src/components/message-input/AttachmentChipList.tsx @@ -13,6 +13,10 @@ type AttachmentChipListProps = { onDragStart: (e: React.DragEvent, id: string) => void; onDragOver: (e: React.DragEvent) => void; onDropOn: (e: React.DragEvent, targetId: string) => void; + /** Completed attachments the selected model has no way to read. */ + unreadableIds?: Set; + /** Display name of the selected model, for the unreadable-file warning. */ + modelName?: string; }; export default function AttachmentChipList({ @@ -22,6 +26,8 @@ export default function AttachmentChipList({ onDragStart, onDragOver, onDropOn, + unreadableIds, + modelName, }: AttachmentChipListProps) { const { t } = useTranslation(); @@ -31,6 +37,11 @@ export default function AttachmentChipList({ const failures = attachments.filter( (attachment) => attachment.status === 'failed' && attachment.errorMessage, ); + // Not a failure: the file is kept and still sends. It only warns that the + // model picked right now would receive nothing it can read. + const unreadable = modelName + ? attachments.filter((attachment) => unreadableIds?.has(attachment.id)) + : []; return ( <> @@ -141,6 +152,22 @@ export default function AttachmentChipList({ ))} )} + + {unreadable.length > 0 && ( +
+ {unreadable.map((attachment) => ( + + {t('conversation.attachments.unreadableByModel', { + name: attachment.fileName, + model: modelName, + })} + + ))} +
+ )} ); } diff --git a/frontend/src/components/message-input/attachmentReadability.test.ts b/frontend/src/components/message-input/attachmentReadability.test.ts new file mode 100644 index 00000000..c05cc7fd --- /dev/null +++ b/frontend/src/components/message-input/attachmentReadability.test.ts @@ -0,0 +1,76 @@ +import type { Attachment } from '../../upload/uploadSlice'; +import { cannotReadAttachment } from './attachmentReadability'; + +const att = (over: Partial = {}): Attachment => ({ + id: 'a1', + fileName: 'scan.pdf', + progress: 100, + status: 'completed', + taskId: 't1', + ...over, +}); + +const PDF_AND_IMAGES = [ + 'application/pdf', + 'image/png', + 'image/jpeg', + 'image/webp', +]; +const IMAGES_ONLY = ['image/png', 'image/jpeg', 'image/webp']; +const TEXT_ONLY: string[] = []; + +const scannedPdf = att({ + mimeType: 'application/pdf', + extractionStatus: 'no_text', + token_count: 0, +}); +const photo = att({ + fileName: 'photo.webp', + mimeType: 'image/webp', + extractionStatus: 'ok', + token_count: 0, +}); + +describe('cannotReadAttachment', () => { + it('flags a scanned PDF on a text-only model', () => { + expect(cannotReadAttachment(scannedPdf, TEXT_ONLY)).toBe(true); + }); + + it('lets a model that reads PDFs natively take a scanned PDF', () => { + expect(cannotReadAttachment(scannedPdf, PDF_AND_IMAGES)).toBe(false); + }); + + it('lets an image-only model take a scanned PDF as page images', () => { + expect(cannotReadAttachment(scannedPdf, IMAGES_ONLY)).toBe(false); + }); + + it('flags an image on a text-only model', () => { + expect(cannotReadAttachment(photo, TEXT_ONLY)).toBe(true); + }); + + it('lets a vision model take the image', () => { + expect(cannotReadAttachment(photo, IMAGES_ONLY)).toBe(false); + }); + + it('does not flag a PDF whose text was extracted', () => { + const textPdf = att({ + mimeType: 'application/pdf', + extractionStatus: 'ok', + token_count: 1200, + }); + expect(cannotReadAttachment(textPdf, TEXT_ONLY)).toBe(false); + }); + + it('stays quiet when the model or the file type is unknown', () => { + expect(cannotReadAttachment(scannedPdf, undefined)).toBe(false); + expect( + cannotReadAttachment(att({ extractionStatus: 'no_text' }), TEXT_ONLY), + ).toBe(false); + }); + + it('ignores attachments that have not finished processing', () => { + expect( + cannotReadAttachment({ ...scannedPdf, status: 'processing' }, TEXT_ONLY), + ).toBe(false); + }); +}); diff --git a/frontend/src/components/message-input/attachmentReadability.ts b/frontend/src/components/message-input/attachmentReadability.ts new file mode 100644 index 00000000..9c768edd --- /dev/null +++ b/frontend/src/components/message-input/attachmentReadability.ts @@ -0,0 +1,33 @@ +import type { Attachment } from '../../upload/uploadSlice'; + +/** + * Whether the selected model has no way to read a completed attachment. + * + * A file with extracted text reaches any model as text. A file without any + * (an image, or a scanned PDF stored with `extraction_status: 'no_text'`) + * only reaches a model that takes that format natively; a PDF also reaches a + * model that takes images, which is sent its pages as images. When the model + * or the file type is unknown, stay quiet rather than warn on a guess. + */ +export function cannotReadAttachment( + attachment: Attachment, + supportedTypes: string[] | undefined, +): boolean { + if (attachment.status !== 'completed') return false; + const mime = attachment.mimeType; + if (!mime || !supportedTypes) return false; + + const isPdf = mime === 'application/pdf'; + if (!isPdf && !mime.startsWith('image/')) return false; + + const hasText = + attachment.extractionStatus !== 'no_text' && + (attachment.token_count ?? 0) > 0; + if (hasText) return false; + + if (supportedTypes.includes(mime)) return false; + if (isPdf && supportedTypes.some((type) => type.startsWith('image/'))) { + return false; + } + return true; +} diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index af03d4a2..d433f191 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -1105,7 +1105,8 @@ "waitingToSend_other": "Will send when {{count}} files finish processing…", "cancelQueuedSend": "Cancel", "sendBlockedByFailed": "{{names}} couldn't be processed — remove to send", - "unsupportedType": "Not a supported file type ({{extension}})" + "unsupportedType": "Not a supported file type ({{extension}})", + "unreadableByModel": "{{model}} can't read {{name}} — switch to a model that supports images or PDFs" }, "retry": "Retry", "reasoning": "Reasoning", diff --git a/frontend/src/upload/uploadSlice.test.ts b/frontend/src/upload/uploadSlice.test.ts index 365aa727..aa56ffb7 100644 --- a/frontend/src/upload/uploadSlice.test.ts +++ b/frontend/src/upload/uploadSlice.test.ts @@ -401,6 +401,24 @@ describe('attachment race recovery', () => { ...overrides, }); + it('keeps the mime type and extraction status a completed attachment reports', () => { + let state = reducer(undefined, addAttachment(makeAttachment())); + state = reducer( + state, + sseEventReceived( + attEvent('attachment.completed', { + token_count: 0, + mime_type: 'application/pdf', + extraction_status: 'no_text', + }), + ), + ); + + expect(state.attachments[0].status).toBe('completed'); + expect(state.attachments[0].mimeType).toBe('application/pdf'); + expect(state.attachments[0].extractionStatus).toBe('no_text'); + }); + it('drops attachment.completed silently when no row matches attachmentId', () => { const state = reducer( undefined, diff --git a/frontend/src/upload/uploadSlice.ts b/frontend/src/upload/uploadSlice.ts index f512e385..db5a00fa 100644 --- a/frontend/src/upload/uploadSlice.ts +++ b/frontend/src/upload/uploadSlice.ts @@ -46,6 +46,13 @@ export interface Attachment { */ attachmentId?: string; token_count?: number; + /** Mime type the server stored (a converted TIFF/BMP reports ``image/png``). */ + mimeType?: string; + /** + * ``metadata.extraction.status`` of a completed attachment: ``ok``, or + * ``no_text`` for a scanned PDF kept for models that read it natively. + */ + extractionStatus?: string; /** Why a ``failed`` attachment failed, when known (server or client gate). */ errorMessage?: string; } @@ -299,6 +306,12 @@ export const uploadSlice = createSlice({ if (Number.isFinite(tokenCount)) { attachment.token_count = tokenCount; } + if (typeof payload.mime_type === 'string') { + attachment.mimeType = payload.mime_type; + } + if (typeof payload.extraction_status === 'string') { + attachment.extractionStatus = payload.extraction_status; + } break; } case 'attachment.failed': { diff --git a/tests/llm/handlers/test_llm_handlers.py b/tests/llm/handlers/test_llm_handlers.py index bd544a05..53a47610 100644 --- a/tests/llm/handlers/test_llm_handlers.py +++ b/tests/llm/handlers/test_llm_handlers.py @@ -2394,3 +2394,71 @@ class TestAttachmentExtractionGate: assert "truncated" in prompt assert "100,000" in prompt assert "250,000" in prompt + + def test_scanned_pdf_row_tells_the_model_it_cannot_read_the_file(self): + """A no-text PDF reaches this path only on a model that reads neither + PDFs nor images. Dropping it silently lets the model answer as if + nothing were attached; naming it lets the model tell the user.""" + handler = ConcreteHandler() + messages = [{"role": "system", "content": "sys"}] + attachments = [{ + "id": "a1", + "filename": "bylaws.pdf", + "mime_type": "application/pdf", + "content": "", + "metadata": {"extraction": {"status": "no_text"}}, + }] + + prompt = handler._append_unsupported_attachments(messages, attachments)[0]["content"] + + assert "bylaws.pdf" in prompt + assert "cannot read" in prompt + assert "Attached file content" not in prompt + + def test_image_without_text_is_named_instead_of_an_empty_block(self): + handler = ConcreteHandler() + messages = [{"role": "system", "content": "sys"}] + attachments = [{ + "id": "a1", + "filename": "photo.webp", + "mime_type": "image/webp", + "content": "", + "metadata": {"extraction": {"status": "ok"}}, + }] + + prompt = handler._append_unsupported_attachments(messages, attachments)[0]["content"] + + assert "photo.webp" in prompt + assert "cannot read" in prompt + assert "Attached file content" not in prompt + + def test_image_with_extracted_text_is_still_appended(self): + handler = ConcreteHandler() + messages = [{"role": "system", "content": "sys"}] + attachments = [{ + "id": "a1", + "filename": "receipt.png", + "mime_type": "image/png", + "content": "Total due: 42 EUR", + "metadata": {"extraction": {"status": "ok"}}, + }] + + prompt = handler._append_unsupported_attachments(messages, attachments)[0]["content"] + + assert "Total due: 42 EUR" in prompt + assert "cannot read" not in prompt + + def test_failed_scan_row_is_still_skipped_silently(self): + handler = ConcreteHandler() + messages = [{"role": "system", "content": "sys"}] + attachments = [{ + "id": "a1", + "filename": "broken.pdf", + "mime_type": "application/pdf", + "content": None, + "metadata": {"extraction": {"status": "failed", "error": "boom"}}, + }] + + result = handler._append_unsupported_attachments(messages, attachments) + + assert result[0]["content"] == "sys" diff --git a/tests/parser/file/test_anydoc_parser.py b/tests/parser/file/test_anydoc_parser.py index 4d387800..638d3327 100644 --- a/tests/parser/file/test_anydoc_parser.py +++ b/tests/parser/file/test_anydoc_parser.py @@ -527,22 +527,28 @@ def test_tableize_never_touches_docling_reroute(monkeypatch): def test_scanned_pdf_with_near_empty_fallback_fails_loudly_when_ocr_off(tmp_path): """A scan whose fallback (OCR off) extracts almost nothing must fail with an actionable message, not be stored as an empty document.""" + from docsgpt.parser.file.base_parser import NoTextLayerError + path = _scanned_pdf(tmp_path / "scan.pdf") fallback = _RecordingFallback(result=" ") fallback.ocr_enabled = False parser = AnydocParser(fallback_parser=fallback) - with pytest.raises(DocumentParseError, match="OCR_ENABLED"): + # Still a DocumentParseError for source ingestion, but typed, so the + # attachment worker can keep the file for models that read it natively. + with pytest.raises(NoTextLayerError, match="OCR_ENABLED"): parser.parse_file(path) assert parser.last_engine is None def test_scanned_pdf_with_near_empty_fallback_and_ocr_on_reports_it(tmp_path): + from docsgpt.parser.file.base_parser import NoTextLayerError + path = _scanned_pdf(tmp_path / "scan.pdf") fallback = _RecordingFallback(result="x") fallback.ocr_enabled = True - with pytest.raises(DocumentParseError, match="even with OCR enabled"): + with pytest.raises(NoTextLayerError, match="even with OCR enabled"): AnydocParser(fallback_parser=fallback).parse_file(path) diff --git a/tests/parser/file/test_bulk.py b/tests/parser/file/test_bulk.py index f77abbe4..99bf02cf 100644 --- a/tests/parser/file/test_bulk.py +++ b/tests/parser/file/test_bulk.py @@ -222,6 +222,31 @@ class TestSimpleDirectoryReaderLoadData: with pytest.raises(DocumentParseError, match="InvalidCxxCompiler"): reader.load_data() + def test_single_unparseable_file_keeps_the_parsers_error_type(self, tmp_path): + """The attachment worker tells a scan from a broken file by the error's type. + + Re-wrapping the parser's exception in a plain ``DocumentParseError`` + would erase ``NoTextLayerError`` and turn every scanned PDF back into + a failed upload. + """ + from docsgpt.parser.file.base_parser import NoTextLayerError + from docsgpt.parser.file.bulk import SimpleDirectoryReader + + (tmp_path / "scan.pdf").write_bytes(b"%PDF-1.7") + + mock_parser = MagicMock() + mock_parser.parser_config_set = True + mock_parser.parse_file.side_effect = NoTextLayerError( + "scan.pdf appears to be a scanned PDF (no text layer)" + ) + + reader = SimpleDirectoryReader( + input_files=[str(tmp_path / "scan.pdf")], + file_extractor={".pdf": mock_parser}, + ) + with pytest.raises(NoTextLayerError, match="scanned PDF"): + reader.load_data() + def test_all_files_unparseable_raises(self, tmp_path): """Nothing parsed is a failed ingest, not an empty success.""" from docsgpt.parser.file.base_parser import DocumentParseError @@ -588,6 +613,21 @@ class TestParserEngineSwitch: assert type(extractor[".xlsx"].fallback_parser).__name__ == "ExcelParser" assert type(extractor[".png"]).__name__ == "ImageParser" + def test_anydoc_without_docling_admits_every_image_format(self, monkeypatch): + """The production image ships neither docling nor OCR. An image is + sent to the model as an image, so .webp/.tiff/.bmp need a parser + entry or the worker refuses them before any model sees them.""" + pytest.importorskip("anydoc") + import sys + + from docsgpt.parser.file.bulk import get_default_file_extractor + + monkeypatch.setitem(sys.modules, "docling", None) + extractor = get_default_file_extractor(engine="anydoc") + + for suffix in (".webp", ".tiff", ".tif", ".bmp"): + assert type(extractor[suffix]).__name__ == "ImageParser", suffix + def test_docling_engine_keeps_docling_map(self): pytest.importorskip("docling") from docsgpt.parser.file.bulk import get_default_file_extractor diff --git a/tests/parser/file/test_image_vision_conversion.py b/tests/parser/file/test_image_vision_conversion.py new file mode 100644 index 00000000..738636fa --- /dev/null +++ b/tests/parser/file/test_image_vision_conversion.py @@ -0,0 +1,64 @@ +"""Re-encoding images the model providers reject (TIFF, BMP) as PNG. + +Chat attachments reach vision models as the stored file itself. OpenAI and +Anthropic accept png/jpeg/webp/gif only, so a TIFF fax or a BMP screenshot has +to be re-encoded before it can be sent at all. +""" + +import io + +import pytest +from PIL import Image + +from docsgpt.parser.file.base_parser import DocumentParseError +from docsgpt.parser.file.image_parser import ( + VISION_CONVERTIBLE_MIME_TYPES, + convert_image_to_png, +) + + +def _encoded(fmt, frames=1, size=(40, 24)): + """An in-memory image of ``frames`` pages, each a different red level.""" + images = [Image.new("RGB", size, color=(i * 60, 20, 200)) for i in range(frames)] + buf = io.BytesIO() + if frames > 1: + images[0].save(buf, format=fmt, save_all=True, append_images=images[1:]) + else: + images[0].save(buf, format=fmt) + buf.seek(0) + return buf + + +@pytest.mark.unit +class TestConvertImageToPng: + def test_tiff_becomes_png(self): + png, frames = convert_image_to_png(_encoded("TIFF")) + + decoded = Image.open(io.BytesIO(png)) + assert decoded.format == "PNG" + assert decoded.size == (40, 24) + assert frames == 1 + + def test_multi_page_tiff_keeps_the_first_page_and_reports_the_count(self): + png, frames = convert_image_to_png(_encoded("TIFF", frames=3)) + + decoded = Image.open(io.BytesIO(png)).convert("RGB") + assert frames == 3 + assert decoded.getpixel((0, 0))[0] == 0 # page one's red level + + def test_bmp_becomes_png(self): + png, frames = convert_image_to_png(_encoded("BMP")) + + assert Image.open(io.BytesIO(png)).format == "PNG" + assert frames == 1 + + def test_unreadable_bytes_raise_a_parse_error(self): + with pytest.raises(DocumentParseError): + convert_image_to_png(io.BytesIO(b"not an image")) + + def test_only_formats_the_providers_reject_are_converted(self): + assert {"image/tiff", "image/bmp"} <= VISION_CONVERTIBLE_MIME_TYPES + assert not ( + {"image/png", "image/jpeg", "image/webp", "image/gif"} + & VISION_CONVERTIBLE_MIME_TYPES + ) diff --git a/tests/parser/file/test_ocr_parser.py b/tests/parser/file/test_ocr_parser.py index de850bc0..a182d3ca 100644 --- a/tests/parser/file/test_ocr_parser.py +++ b/tests/parser/file/test_ocr_parser.py @@ -536,13 +536,16 @@ class TestNativeOcrImageParser: @pytest.mark.unit class TestExtractorWiring: - def test_legacy_map_without_ocr_is_unchanged(self): + def test_legacy_map_without_ocr_hands_every_image_to_image_parser(self): + """Without OCR an image still reaches the model as an image, so every + image suffix the upload gate admits needs an entry. Only png/jpg/jpeg + had one, and a .webp was refused outright (prod 2026-09-10).""" from docsgpt.parser.file.bulk import _legacy_file_extractor extractor = _legacy_file_extractor() assert type(extractor[".pdf"]).__name__ == "PDFParser" - assert type(extractor[".png"]).__name__ == "ImageParser" - assert ".tiff" not in extractor + for suffix in op.IMAGE_SUFFIXES: + assert type(extractor[suffix]).__name__ == "ImageParser", suffix def test_legacy_map_with_ocr_uses_native_parsers(self): from docsgpt.parser.file.bulk import _legacy_file_extractor diff --git a/tests/test_upload_limits.py b/tests/test_upload_limits.py index 180f271b..cf15cc6f 100644 --- a/tests/test_upload_limits.py +++ b/tests/test_upload_limits.py @@ -145,23 +145,24 @@ def test_enforce_parseable_attachment_rejects_binary_behind_a_bom(bom, tmp_path) def test_enforce_parseable_attachment_uses_the_extractor_it_is_given(tmp_path): """The worker holds the live parser table; a trimmed install must not admit on trust. - Without docling the fallback extractor has no .webp handler, so a .webp - would otherwise skip the content check and be read as plain text. + Without docling the fallback extractor has no .vtt handler, so a binary + file named .vtt would otherwise skip the content check and be read as + plain text. """ from docsgpt.upload_limits import ( enforce_parseable_attachment, UnsupportedUploadTypeError, ) - path = tmp_path / "scan.webp" - path.write_bytes(b"RIFF\x00\x00\x00\x00WEBPVP8 " + bytes(range(256))) + path = tmp_path / "subs.vtt" + path.write_bytes(b"\x00\x01\x02\x03" + bytes(range(256))) - # Default list: .webp is parser-backed, admitted on its name. - enforce_parseable_attachment(path, "scan.webp") + # Default list: .vtt is parser-backed, admitted on its name. + enforce_parseable_attachment(path, "subs.vtt") - # The extractor actually loaded has no .webp parser. + # The extractor actually loaded has no .vtt parser. with pytest.raises(UnsupportedUploadTypeError): - enforce_parseable_attachment(path, "scan.webp", {".pdf", ".docx"}) + enforce_parseable_attachment(path, "subs.vtt", {".pdf", ".docx"}) @pytest.mark.parametrize( diff --git a/tests/worker/test_attachment_worker.py b/tests/worker/test_attachment_worker.py index 882a8976..ac218e39 100644 --- a/tests/worker/test_attachment_worker.py +++ b/tests/worker/test_attachment_worker.py @@ -384,14 +384,14 @@ class TestAttachmentTypeGuard: ): """The guard judges against the parser table actually loaded. - Without docling the fallback extractor has no .webp parser, so a - .webp the route admitted on its name would otherwise be opened as - plain text here. + The route admits ``.vtt`` by name, but a docling-less install has no + VTT parser, so a binary file named ``.vtt`` would otherwise be opened + as plain text here. """ from docsgpt import worker - local_path = tmp_path / "scan.webp" - local_path.write_bytes(b"RIFF\x00\x00\x00\x00WEBPVP8 " + bytes(range(256))) + local_path = tmp_path / "subs.vtt" + local_path.write_bytes(b"\x00\x01\x02\x03" + bytes(range(256))) events = [] fake_storage = MagicMock(name="storage") @@ -399,7 +399,7 @@ class TestAttachmentTypeGuard: str(local_path) ) monkeypatch.setattr(worker.StorageCreator, "get_storage", lambda: fake_storage) - # The docling-less fallback table: images are handled, .webp is not. + # A trimmed parser table: .png is handled, .vtt is not. monkeypatch.setattr( worker, "get_default_file_extractor", @@ -412,19 +412,176 @@ class TestAttachmentTypeGuard: ) file_info = { - "filename": "scan.webp", + "filename": "subs.vtt", "attachment_id": "507f1f77bcf86cd799439014", - "path": "uploads/user1/attachments/scan.webp", + "path": "uploads/user1/attachments/subs.vtt", "metadata": {"source": "chat"}, } with pytest.raises( - worker.AttachmentRejectedError, match=r"Unsupported file type: \.webp" + worker.AttachmentRejectedError, match=r"Unsupported file type: \.vtt" ): worker.attachment_worker(task_self, file_info, "user1") failed = [payload for name, payload in events if name == "attachment.failed"] - assert failed and failed[0]["error"] == "Unsupported file type: .webp" + assert failed and failed[0]["error"] == "Unsupported file type: .vtt" + + def test_scanned_pdf_completes_without_text_for_models_that_read_it_natively( + self, pg_conn, patch_worker_db, task_self, monkeypatch + ): + """A PDF with no text layer is still a usable attachment. + + Models that read PDFs natively (or as page images) are sent the + stored file, not its extracted text, so failing the upload only kept + the file from a model that could read it (prod 2026-09-12: a paying + user's scanned club bylaws). The row is kept with ``status: no_text`` + — text inlining still skips it — and the user is told it completed. + """ + from docsgpt import worker + from docsgpt.parser.file.base_parser import NoTextLayerError + + published: list[tuple[str, dict]] = [] + monkeypatch.setattr( + worker, + "publish_user_event", + lambda user, event, payload, **kw: published.append((event, payload)), + ) + fake_storage = MagicMock(name="storage") + fake_storage.process_file.side_effect = NoTextLayerError( + "bylaws.pdf appears to be a scanned PDF (no text layer), and " + "PDFParser extracted almost nothing" + ) + monkeypatch.setattr(worker.StorageCreator, "get_storage", lambda: fake_storage) + monkeypatch.setattr( + worker, "get_default_file_extractor", lambda ocr_enabled=False, pdf_text_fast_path=False: {} + ) + + file_info = { + "filename": "bylaws.pdf", + "attachment_id": "507f1f77bcf86cd799439021", + "path": "uploads/user1/attachments/bylaws.pdf", + "metadata": {"source": "chat"}, + } + + result = worker.attachment_worker(task_self, file_info, "user1") + + assert result["token_count"] == 0 + assert result["mime_type"] == "application/pdf" + extraction = result["metadata"]["extraction"] + assert extraction["status"] == "no_text" + assert "scanned PDF" in extraction["reason"] + + row = AttachmentsRepository(pg_conn).get_by_legacy_id( + file_info["attachment_id"], "user1" + ) + assert row["upload_path"] == file_info["path"] + assert row["content"] == "" + assert row["metadata"]["extraction"]["status"] == "no_text" + + events = [event for event, _ in published] + assert "attachment.failed" not in events + completed = [payload for event, payload in published if event == "attachment.completed"] + assert completed[0]["extraction_status"] == "no_text" + assert completed[0]["mime_type"] == "application/pdf" + + def test_no_text_on_a_type_models_cannot_read_natively_still_fails( + self, pg_conn, patch_worker_db, task_self, monkeypatch + ): + """Only PDFs and images have a native path; anything else without text stays a failure.""" + from docsgpt import worker + from docsgpt.parser.file.base_parser import NoTextLayerError + + published: list[tuple[str, dict]] = [] + monkeypatch.setattr( + worker, + "publish_user_event", + lambda user, event, payload, **kw: published.append((event, payload)), + ) + fake_storage = MagicMock(name="storage") + fake_storage.process_file.side_effect = NoTextLayerError("notes.docx has no text") + monkeypatch.setattr(worker.StorageCreator, "get_storage", lambda: fake_storage) + monkeypatch.setattr( + worker, "get_default_file_extractor", lambda ocr_enabled=False, pdf_text_fast_path=False: {} + ) + + file_info = { + "filename": "notes.docx", + "attachment_id": "507f1f77bcf86cd799439022", + "path": "uploads/user1/attachments/notes.docx", + "metadata": {"source": "chat"}, + } + + with pytest.raises(NoTextLayerError): + worker.attachment_worker(task_self, file_info, "user1") + + events = [event for event, _ in published] + assert "attachment.failed" in events + assert "attachment.completed" not in events + + def test_tiff_is_stored_as_png_so_providers_accept_it( + self, pg_conn, patch_worker_db, task_self, monkeypatch + ): + """OpenAI and Anthropic reject TIFF, so the row must point at a PNG copy.""" + import io + + from PIL import Image + + from docsgpt import worker + + buf = io.BytesIO() + Image.new("RGB", (32, 20), color=(10, 120, 200)).save(buf, format="TIFF") + tiff_bytes = buf.getvalue() + + published: list[tuple[str, dict]] = [] + saved: dict[str, bytes] = {} + + def _save(data, path, **kwargs): + saved[path] = data.read() + return {"storage_type": "local"} + + fake_storage = MagicMock(name="storage") + fake_storage.process_file.return_value = Document(text="", extra_info={}) + fake_storage.get_file.side_effect = lambda path: io.BytesIO(tiff_bytes) + fake_storage.save_file.side_effect = _save + monkeypatch.setattr(worker.StorageCreator, "get_storage", lambda: fake_storage) + monkeypatch.setattr( + worker, "get_default_file_extractor", lambda ocr_enabled=False, pdf_text_fast_path=False: {} + ) + monkeypatch.setattr( + worker, + "publish_user_event", + lambda user, event, payload, **kw: published.append((event, payload)), + ) + + file_info = { + "filename": "fax.tiff", + "attachment_id": "507f1f77bcf86cd799439023", + "path": "uploads/user1/attachments/abc/fax.tiff", + "metadata": {"source": "chat"}, + } + + result = worker.attachment_worker(task_self, file_info, "user1") + + png_path = "uploads/user1/attachments/abc/fax.png" + assert list(saved) == [png_path] + assert Image.open(io.BytesIO(saved[png_path])).format == "PNG" + assert result["path"] == png_path + assert result["mime_type"] == "image/png" + assert result["metadata"]["image_conversion"] == { + "from": "image/tiff", + "to": "image/png", + "frames": 1, + } + + row = AttachmentsRepository(pg_conn).get_by_legacy_id( + file_info["attachment_id"], "user1" + ) + assert row["filename"] == "fax.tiff" + assert row["upload_path"] == png_path + assert row["mime_type"] == "image/png" + + completed = [payload for event, payload in published if event == "attachment.completed"] + assert completed[0]["mime_type"] == "image/png" def test_text_without_a_parser_is_parsed( self, pg_conn, patch_worker_db, task_self, monkeypatch, tmp_path From ad201f83184f48f8561198812e6f76e17cad99f1 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 14 Sep 2026 17:45:43 +0100 Subject: [PATCH 7/7] fix: refuse oversized TIFF/BMP attachments before converting them A deflate-compressed TIFF under 1 MB can declare 144 million pixels and take 1.2 GB to convert to PNG, and Pillow only warns below 179 million. Read the dimensions from the header and refuse images over 40 million pixels before any pixel data is decoded. Pillow's DecompressionBombError is now raised as DocumentParseError, so the upload fails once instead of being retried. --- docs/content/Guides/ocr.mdx | 3 +- docsgpt/parser/file/image_parser.py | 18 ++++++++-- .../file/test_image_vision_conversion.py | 35 +++++++++++++++++++ 3 files changed, 53 insertions(+), 3 deletions(-) diff --git a/docs/content/Guides/ocr.mdx b/docs/content/Guides/ocr.mdx index 2b185141..1a79eb84 100644 --- a/docs/content/Guides/ocr.mdx +++ b/docs/content/Guides/ocr.mdx @@ -156,7 +156,8 @@ itself is what a model reads. Models that take PDFs get the document; models that take only images get its pages as images; a text-only model is told it cannot read the file, and the message box warns before sending. Images work the same way with OCR off, and TIFF and BMP attachments are stored as PNG, -since model providers do not accept those formats. +since model providers do not accept those formats. A TIFF or BMP larger than +40 million pixels is refused instead of converted. Mixed documents — text pages with scanned pages among them — convert through anydoc, which reads the text pages and skips the scanned ones. With diff --git a/docsgpt/parser/file/image_parser.py b/docsgpt/parser/file/image_parser.py index e3382682..4a7b2942 100644 --- a/docsgpt/parser/file/image_parser.py +++ b/docsgpt/parser/file/image_parser.py @@ -18,6 +18,12 @@ from docsgpt.parser.file.base_parser import BaseParser, DocumentParseError # to PNG before it is stored for the model. VISION_CONVERTIBLE_MIME_TYPES = frozenset({"image/tiff", "image/bmp", "image/x-ms-bmp"}) +# Largest image re-encoded to PNG. A deflate TIFF under 1 MB can declare +# 144 million pixels and take 1.2 GB to convert, while Pillow only warns below +# 179 million. Vision models downscale far below this cap (Anthropic refuses +# more than 8000 px per side), so a larger image gains nothing. +MAX_CONVERTIBLE_PIXELS = 40_000_000 + def convert_image_to_png(file_obj: BinaryIO) -> Tuple[bytes, int]: """Re-encode the first frame of an image as PNG. @@ -30,12 +36,20 @@ def convert_image_to_png(file_obj: BinaryIO) -> Tuple[bytes, int]: source had. Only the first is kept. Raises: - DocumentParseError: If the bytes are not an image Pillow can decode. + DocumentParseError: If the bytes are not an image Pillow can decode, + or the image is larger than ``MAX_CONVERTIBLE_PIXELS``. """ from PIL import Image, UnidentifiedImageError try: with Image.open(file_obj) as image: + # Opening reads only the header; refuse before any pixels decode. + width, height = image.size + if width * height > MAX_CONVERTIBLE_PIXELS: + raise DocumentParseError( + f"The image is too large to convert ({width}×{height} pixels; the limit is " + f"{MAX_CONVERTIBLE_PIXELS // 1_000_000} million pixels). Resize it and upload it again." + ) frames = getattr(image, "n_frames", 1) image.seek(0) frame = image @@ -43,7 +57,7 @@ def convert_image_to_png(file_obj: BinaryIO) -> Tuple[bytes, int]: frame = frame.convert("RGBA" if "A" in frame.getbands() else "RGB") out = io.BytesIO() frame.save(out, format="PNG") - except (UnidentifiedImageError, OSError) as exc: + except (UnidentifiedImageError, OSError, Image.DecompressionBombError) as exc: raise DocumentParseError(f"Could not read the image: {exc}") from exc return out.getvalue(), frames diff --git a/tests/parser/file/test_image_vision_conversion.py b/tests/parser/file/test_image_vision_conversion.py index 738636fa..d0225b22 100644 --- a/tests/parser/file/test_image_vision_conversion.py +++ b/tests/parser/file/test_image_vision_conversion.py @@ -52,10 +52,45 @@ class TestConvertImageToPng: assert Image.open(io.BytesIO(png)).format == "PNG" assert frames == 1 + def test_palette_tiff_is_stored_as_rgb_png(self): + buf = io.BytesIO() + Image.new("RGB", (40, 24), color=(0, 204, 51)).convert("P").save(buf, format="TIFF") + buf.seek(0) + + png, _ = convert_image_to_png(buf) + + decoded = Image.open(io.BytesIO(png)) + assert decoded.mode == "RGB" + assert decoded.getpixel((0, 0)) == (0, 204, 51) + def test_unreadable_bytes_raise_a_parse_error(self): with pytest.raises(DocumentParseError): convert_image_to_png(io.BytesIO(b"not an image")) + def test_image_over_the_pixel_limit_is_refused_before_decoding(self, monkeypatch): + # A deflate TIFF under 1 MB can decode to over a gigabyte, so the + # dimensions in the header are checked before any pixel data is read. + from PIL import TiffImagePlugin + + from docsgpt.parser.file import image_parser + + def decode(*args, **kwargs): + raise AssertionError("pixel data was decoded") + + monkeypatch.setattr(image_parser, "MAX_CONVERTIBLE_PIXELS", 500) + monkeypatch.setattr(TiffImagePlugin.TiffImageFile, "load", decode) + + with pytest.raises(DocumentParseError, match="40×24"): + convert_image_to_png(_encoded("TIFF")) + + def test_decompression_bomb_refusal_is_a_parse_error(self, monkeypatch): + # Pillow refuses images over twice MAX_IMAGE_PIXELS with an error that + # is not an OSError; it must fail the upload like any unreadable image. + monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 100) + + with pytest.raises(DocumentParseError): + convert_image_to_png(_encoded("TIFF")) + def test_only_formats_the_providers_reject_are_converted(self): assert {"image/tiff", "image/bmp"} <= VISION_CONVERTIBLE_MIME_TYPES assert not (