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/content/Guides/ocr.mdx b/docs/content/Guides/ocr.mdx index b49c28fa..1a79eb84 100644 --- a/docs/content/Guides/ocr.mdx +++ b/docs/content/Guides/ocr.mdx @@ -150,6 +150,15 @@ 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. 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 OCR on, DocsGPT probes every page's text layer, OCRs the pages that have diff --git a/docs/runbooks/sse-notifications.md b/docs/runbooks/sse-notifications.md index e8ee55f5..1d100c43 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: ────────────────────────────────┘ ``` @@ -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" @@ -204,11 +206,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. --- @@ -303,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 @@ -324,6 +328,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 +365,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..b2ecd38d --- /dev/null +++ b/docsgpt/api/asgi_stream.py @@ -0,0 +1,82 @@ +"""Streaming response for native-async routes whose streams hold resources.""" + +from __future__ import annotations + +import logging +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 +# 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 + +# 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. + + 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. 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: AsyncIterable[Any], + *, + on_close: Optional[Callable[[], Awaitable[None]]] = None, + **kwargs: Any, + ) -> 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: + 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") + + +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: + 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 fc6493bf..d3fa3437 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. @@ -21,13 +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 Response 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, bind_log_context, json_error +from docsgpt.api.asgi_stream import sse_response 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, @@ -37,15 +36,10 @@ from docsgpt.streaming.event_replay import ( DEFAULT_KEEPALIVE_SECONDS, DEFAULT_POLL_TIMEOUT_SECONDS, ) +from docsgpt.streaming.sse_leases import StreamCapExceeded, acquire_stream_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( @@ -91,122 +85,42 @@ 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 _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 → 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) + 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 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) + 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" @@ -230,11 +144,7 @@ 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), - media_type="text/event-stream", - headers=_SSE_HEADERS, - ) + return sse_response(stream, lease) # Mounted in ``docsgpt/asgi.py`` ahead of the Flask catch-all. Keep 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..65b47f2f --- /dev/null +++ b/docsgpt/api/devices/session_events.py @@ -0,0 +1,153 @@ +"""``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 run in worker threads, 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_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, SessionState, 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 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" + f"data: {json.dumps(payload)}\n\n" + ) + + +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 + + +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.""" + keepalive_interval = float(settings.SSE_KEEPALIVE_SECONDS) + idle_seconds = float(settings.REMOTE_DEVICE_SESSION_IDLE_SECONDS) + 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() + + +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"] + bind_log_context("devices.session_events") + device, failure = await anyio.to_thread.run_sync( + partial( + authenticate_device, + request.headers, + request.method, + _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() + # 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( + broker.redeem_ticket, device["id"], device["user_id"], session_id + ) + if sess is None: + return _error("session_ticket_invalid", 410) + return ClosingStreamingResponse( + _session_stream(broker, sess), + media_type="text/event-stream", + headers=_SSE_HEADERS, + on_close=partial(_close_session, broker, sess.session_id), + ) + + +# 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..52bf260a 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,26 +16,32 @@ import json import logging import re import time -from typing import Iterator, Optional +from contextlib import aclosing +from typing import AsyncIterator, Optional -from flask import Blueprint, Response, jsonify, make_response, request, stream_with_context +import anyio +from redis.asyncio import Redis +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 sse_response +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, topic_name, ) -from docsgpt.streaming.broadcast_channel import Topic +from docsgpt.streaming.async_broadcast_channel import AsyncTopic +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__) -events = Blueprint("event_stream", __name__) - SUBSCRIBE_POLL_INTERVAL_SECONDS = 1.0 # WHATWG SSE treats CRLF, CR, and LF equivalently as line terminators. @@ -74,19 +84,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: 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 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,15 +108,16 @@ def _oldest_retained_id(redis_client, user_id: str) -> Optional[str]: return None -def _allow_replay( - redis_client, user_id: str, last_event_id: Optional[str] +async def _allow_replay( + 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: @@ -119,12 +130,13 @@ 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)) - # 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) + 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 + # 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", @@ -166,10 +178,10 @@ def _replay_floor_id() -> Optional[str]: return f"{floor_ms}-0" -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``. +async def _replay_backlog( + 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``. Capped at ``max_count`` rows; clients catch up across reconnects. Parse failures are skipped; the Streams id is injected into the @@ -186,7 +198,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 +208,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 +236,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 +252,95 @@ 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, - ) - - # 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: Optional[Redis], + 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 +366,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 +385,86 @@ 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 - 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, - ) + # 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) + # 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 + # 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 + redis_client = await get_async_redis_instance() if push_enabled else None + + # 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 lease is not None: + await lease.release() + 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 + stream = _event_stream(redis_client, user_id, last_event_id, last_event_id_invalid, push_enabled) + return sse_response(stream, lease) + + +# 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..a57ae727 --- /dev/null +++ b/docsgpt/api/user/artifacts/download.py @@ -0,0 +1,212 @@ +"""``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 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 +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 _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") + + +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]]]: + """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 the handle; closing anything but an in-memory buffer runs in a worker thread.""" + close = getattr(file_obj, "close", None) + 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: + """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 + # 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: + 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) + + 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) + + # 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. + file_obj = await anyio.to_thread.run_sync(storage.get_file, storage_path) + return ClosingStreamingResponse( + _read_chunks(file_obj), + media_type=mime_type, + headers=headers, + 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..09d55360 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__) @@ -58,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}" @@ -139,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 ) @@ -191,13 +229,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 +242,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/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/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..4a7b2942 100644 --- a/docsgpt/parser/file/image_parser.py +++ b/docsgpt/parser/file/image_parser.py @@ -1,14 +1,65 @@ """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"}) + +# 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. + + 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, + 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 + 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, Image.DecompressionBombError) as exc: + raise DocumentParseError(f"Could not read the image: {exc}") from exc + return out.getvalue(), frames class ImageParser(BaseParser): @@ -24,8 +75,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/streaming/async_broadcast_channel.py b/docsgpt/streaming/async_broadcast_channel.py index f18e3c42..57e05f87 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,15 @@ 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" + +# 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.""" @@ -36,6 +47,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 +57,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 +77,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 +92,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 +102,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: @@ -97,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: @@ -111,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/async_redis.py b/docsgpt/streaming/async_redis.py index 4dadf6a6..c0c60100 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 @@ -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,9 +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. + 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: @@ -38,7 +46,11 @@ 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), ) except ValueError as e: logger.error("Invalid Redis URL for async client: %s", e) diff --git a/docsgpt/streaming/sse_leases.py b/docsgpt/streaming/sse_leases.py new file mode 100644 index 00000000..8a5f3909 --- /dev/null +++ b/docsgpt/streaming/sse_leases.py @@ -0,0 +1,130 @@ +"""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 +from docsgpt.streaming.async_redis import ASYNC_REDIS_OP_TIMEOUT_SECONDS + +logger = logging.getLogger(__name__) + +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(ASYNC_REDIS_OP_TIMEOUT_SECONDS): + async with self.redis.pipeline(transaction=True) as pipe: + # 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: + 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(ASYNC_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 = StreamLease(redis=redis, key=key, lease_id=uuid.uuid4().hex, ttl=ttl, refreshed_at=time.monotonic()) + now = time.time() + try: + 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.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 + 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/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/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_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 ac1c055c..63831a31 100644 --- a/tests/api/test_async_sse_routes.py +++ b/tests/api/test_async_sse_routes.py @@ -1,69 +1,84 @@ """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" -_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" +_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 ───────────────────────────────────────────────────────── @@ -169,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 ─────────────────────────────────────────────────── @@ -206,17 +239,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 +254,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 +279,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 d7f05480..b2f188c1 100644 --- a/tests/api/test_events_routes.py +++ b/tests/api/test_events_routes.py @@ -1,615 +1,581 @@ -"""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 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 -import threading -from typing import Any -from unittest.mock import MagicMock, patch +import time +from typing import Any, Callable, Optional +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 +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" +_LEASES = "user:alice:sse_leases" +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() -> FakeAsyncRedis: + return FakeAsyncRedis() - 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, + 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): + 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() + assert redis.zsets == {} - 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_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( + _AREDIS, AsyncMock(return_value=_redis()) + ), patch(_SUBSCRIBE, _subscribe()): + r = _get() + assert ": keepalive\n\n" in r.text -# ── concurrency cap ───────────────────────────────────────────────────── +# ── concurrency cap (per-connection leases) ───────────────────────────── +@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_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 - # 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"} + # 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_skips_cap_when_zero_disabled(self): - from docsgpt.api.events import routes as events_module + 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 == {} - app = _make_app() - redis_client = MagicMock() + 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( + _AREDIS, AsyncMock(return_value=redis) + ), patch(_SUBSCRIBE, _subscribe()): + r = _get() + assert r.status_code == 200 + assert redis.zsets == {} - 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_redis_errors_fail_open(self): + redis = _redis() + 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 + + 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(probe=_count_leases)): + r = _get() + assert r.status_code == 200 + assert seen["during"] == 1 + assert redis.zsets == {} + + 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_lease_and_closes_subscription(self, monkeypatch): + monkeypatch.setattr(events_module.settings, "SSE_KEEPALIVE_SECONDS", 0) + redis = _redis() + 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 + assert redis.zsets == {} + + 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() + 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 + assert redis.zsets == {} + + async def test_shutdown_ends_stream_and_releases_lease(self): + redis = _redis() + 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 + assert redis.zsets == {} # ── 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): + 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_fail_open_when_redis_stalls(self, monkeypatch): + monkeypatch.setattr(events_module, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05) + redis = _redis() - redis.incr.assert_not_called() - redis.xlen.assert_not_called() + async def _stalled(_key): + await anyio.sleep_forever() - def test_allow_replay_passes_until_budget_exhausted(self): - from docsgpt.api.events.routes import _allow_replay + redis.incr = AsyncMock(side_effect=_stalled) + with anyio.fail_after(2): + assert await events_module._allow_replay(redis, "alice", OLD_CURSOR) is True - 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_recovers_when_seeding_expire_raises(self): + redis = _redis() + counter = {"v": 0} - def _incr(_key): - counter["v"] += 1 - return counter["v"] + def _incr(_key): + counter["v"] += 1 + return counter["v"] - redis.incr.side_effect = _incr + 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 - # 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_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 test_allow_replay_fail_open_on_redis_error(self): - from docsgpt.api.events.routes import _allow_replay + 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" - 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_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) == [] - def test_allow_replay_recovers_when_seeding_expire_raises(self): - """Regression: INCR=1 then EXPIRE raising must not wedge the key. + 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" - 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 + 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() - counter = {"v": 0} + 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 _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() - - def _incr(key): - if key == "user:alice:sse_count": - return 1 - # Budget counter: report over-limit. - return 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", - }, - ) + """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() + 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 - # 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 connection lease is released so a denied request doesn't hold a slot. + assert redis.zsets == {} +# ── 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 +591,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..2e0e6028 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,131 @@ 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_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 + 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 +1030,6 @@ class TestMalformedArtifactId: [ ("GetArtifact", (), "get", None), ("GetArtifactVersion", (1,), "get", None), - ("DownloadArtifact", (), "get", None), ("RestoreArtifact", (), "post", {"version": 1}), ], ) @@ -930,13 +1109,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 +1119,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..1daaacc4 --- /dev/null +++ b/tests/asgi_stream.py @@ -0,0 +1,88 @@ +"""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"", + root_path: str = "", + 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. ``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. + """ + 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": 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")) + 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..c0a3f316 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 @@ -83,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: @@ -187,6 +198,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 +228,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..35fc7c13 --- /dev/null +++ b/tests/devices/test_session_events_async.py @@ -0,0 +1,502 @@ +"""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 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 + +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 _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) + + +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"]] + + @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 ────────────────────────────────────────────────────────────── + + +@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 + + 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 ────────────────────────────────────────────── + + +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() + + @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 + 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..98f1381e 100644 --- a/tests/devices/test_session_ticket.py +++ b/tests/devices/test_session_ticket.py @@ -2,25 +2,28 @@ 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. +* 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. """ 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 # --------------------------------------------------------------------------- @@ -96,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) # --------------------------------------------------------------------------- @@ -104,6 +146,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 +207,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/fake_async_redis.py b/tests/fake_async_redis.py new file mode 100644 index 00000000..81729c1b --- /dev/null +++ b/tests/fake_async_redis.py @@ -0,0 +1,153 @@ +"""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, ``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 + +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.hang_after_execute = 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() + if self._redis.hang_after_execute: + await anyio.sleep_forever() + return results 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..d0225b22 --- /dev/null +++ b/tests/parser/file/test_image_vision_conversion.py @@ -0,0 +1,99 @@ +"""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_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 ( + {"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_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..0b888c89 --- /dev/null +++ b/tests/test_async_broadcast_channel.py @@ -0,0 +1,187 @@ +"""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_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)): + 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..07880203 --- /dev/null +++ b/tests/test_async_redis.py @@ -0,0 +1,78 @@ +"""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_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) + 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 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..94362526 --- /dev/null +++ b/tests/test_sse_leases.py @@ -0,0 +1,216 @@ +"""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, "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 +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, "ASYNC_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_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") + 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"] 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