Merge branch 'arc53:main' into main

This commit is contained in:
Manish Madan authored and GitHub committed 2026-09-15 04:18:40 +05:30
commit 2afc711b57
60 files changed
+4696 -1690

No files matched your search

+10 -5
View File
@@ -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/<id>/events`.
- the `/mcp` FastMCP endpoint,
- the chat reconnect reader `GET /api/messages/<id>/events`,
- the notification stream `GET /api/events`,
- the remote-device command stream `GET /api/devices/sessions/<id>/events`, and
- artifact downloads `GET /api/artifacts/<id>/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
+2 -1
View File
@@ -70,7 +70,7 @@ GET /api/messages/<message_id>/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).
<Callout type="warning" emoji="⚠️">
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.
</Callout>
## 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. |
@@ -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/<id>/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/<id>/events`, the notification stream `GET /api/events`, the remote-device command stream `GET /api/devices/sessions/<id>/events`, and artifact downloads `GET /api/artifacts/<id>/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:**
+3 -3
View File
@@ -22,7 +22,7 @@ flowchart TB
Clients["Web UI, widgets and API clients"] --> API
subgraph Runtime["DocsGPT processes"]
API["API process<br/>Flask routes and agent execution<br/>ASGI shell: MCP and chat reconnect"]
API["API process<br/>Flask routes and agent execution<br/>ASGI shell: MCP and streaming routes"]
Redis["Redis<br/>Task broker, results and realtime delivery"]
Worker["Celery workers<br/>Ingestion, parsing, embeddings<br/>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`.
<Callout type="info">
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).
</Callout>
### Agents and retrieval
+9
View File
@@ -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
+30 -31
View File
@@ -39,10 +39,10 @@ redis-cli -n 2 XLEN user:<user_id>:stream
Worker (publish_user_event) Frontend tab
│ ▲
▼ │ GET /api/events SSE
Redis Streams: XADD Flask route
Redis Streams: XADD async route (event loop)
user:<id>:stream ──────────────► replay_backlog (snapshot)
│ +
▼ Topic.subscribe (live tail)
▼ AsyncTopic.subscribe (live tail)
Redis pub/sub: PUBLISH │
user:<id> ────────────────────────────────┘
```
@@ -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:<id>: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:<id>: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:<id>: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:<id>: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=<id>` 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:<id>` 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:<id>:stream MAXLEN 0
# Clear a wedged concurrent-connection counter
redis-cli -n 2 DEL user:<id>:sse_count
# Clear a user's SSE connection leases (they also age out on their own)
redis-cli -n 2 DEL user:<id>: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/<id>/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/<id>/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
+59
View File
@@ -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)
+82
View File
@@ -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)
+22 -112
View File
@@ -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
+66 -28
View File
@@ -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
+1 -5
View File
@@ -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/<session_id>/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/<session_id>/invocations/<invocation_id>/ack",
view_func=ack_invocation, methods=["POST"], endpoint="device_ack",
+5 -79
View File
@@ -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"
)
+153
View File
@@ -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"],
),
]
+208 -275
View File
@@ -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: '(<id>' 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"]),
]
+32 -11
View File
@@ -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.
+212
View File
@@ -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"]),
]
+4 -119
View File
@@ -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/<artifact_id>/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/<artifact_id>/restore")
class RestoreArtifact(Resource):
@api.doc(description="Restore a prior version by appending it as the new current version")
-2
View File
@@ -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)
+13 -5
View File
@@ -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=[
+7 -2
View File
@@ -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.
+88 -7
View File
@@ -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
# ------------------------------------------------------------------
+3 -3
View File
@@ -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:
+31
View File
@@ -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)
+2 -1
View File
@@ -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}"
)
+11
View File
@@ -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."""
+10 -8
View File
@@ -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: "
+56 -5
View File
@@ -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
+60 -7
View File
@@ -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:
+22 -10
View File
@@ -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)
+130
View File
@@ -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()
+69 -2
View File
@@ -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)},
)
+25
View File
@@ -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<RootState>();
@@ -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' && (
@@ -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<string>;
/** 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({
))}
</div>
)}
{unreadable.length > 0 && (
<div
className="flex flex-col gap-0.5 px-2 pb-1 text-xs text-[#B54708] sm:px-3 dark:text-[#FDB022]"
role="status"
>
{unreadable.map((attachment) => (
<span key={attachment.id}>
{t('conversation.attachments.unreadableByModel', {
name: attachment.fileName,
model: modelName,
})}
</span>
))}
</div>
)}
</>
);
}
@@ -0,0 +1,76 @@
import type { Attachment } from '../../upload/uploadSlice';
import { cannotReadAttachment } from './attachmentReadability';
const att = (over: Partial<Attachment> = {}): 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);
});
});
@@ -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;
}
+2 -1
View File
@@ -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",
+18
View File
@@ -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,
+13
View File
@@ -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': {
+116
View File
@@ -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"}
+92
View File
@@ -0,0 +1,92 @@
"""Tests for ``docsgpt/api/asgi_stream.py`` — cleanup that can't stall a worker.
The cleanup runs shielded from the cancellation that ended the response, so
each step needs its own deadline: a close or release waiting on a dead Redis
connection must not hold the request task, or a graceful shutdown, forever.
"""
from __future__ import annotations
import anyio
import pytest
from docsgpt.api import asgi_stream
from docsgpt.api.asgi_stream import ClosingStreamingResponse
from tests.asgi_stream import stream_then_disconnect
class _StuckCloseIterator:
"""Yields one chunk, then ends; its ``aclose`` never returns."""
def __init__(self) -> None:
self._sent = False
self.aclose_started = False
def __aiter__(self):
return self
async def __anext__(self):
if self._sent:
raise StopAsyncIteration
self._sent = True
return b"chunk"
async def aclose(self):
self.aclose_started = True
await anyio.sleep_forever()
async def _one_chunk():
yield b"chunk"
@pytest.mark.unit
@pytest.mark.asyncio
class TestBoundedCleanup:
async def test_stuck_stream_close_is_abandoned_and_on_close_still_runs(self, monkeypatch):
monkeypatch.setattr(asgi_stream, "_CLEANUP_TIMEOUT_SECONDS", 0.05)
body = _StuckCloseIterator()
released = {"ran": False}
async def on_close():
released["ran"] = True
response = ClosingStreamingResponse(body, media_type="text/plain", on_close=on_close)
with anyio.fail_after(2):
status, _, chunks = await stream_then_disconnect(
response, "/", chunks_before_disconnect=10**6
)
assert status == 200
assert chunks == [b"chunk"]
assert body.aclose_started
assert released["ran"] is True
async def test_stuck_on_close_is_abandoned(self, monkeypatch):
monkeypatch.setattr(asgi_stream, "_CLEANUP_TIMEOUT_SECONDS", 0.05)
async def on_close():
await anyio.sleep_forever()
response = ClosingStreamingResponse(_one_chunk(), media_type="text/plain", on_close=on_close)
with anyio.fail_after(2):
status, _, chunks = await stream_then_disconnect(
response, "/", chunks_before_disconnect=10**6
)
assert status == 200
assert chunks == [b"chunk"]
async def test_on_close_runs_when_client_leaves_before_first_frame(self):
released = {"count": 0}
async def on_close():
released["count"] += 1
async def endless():
while True:
await anyio.sleep(0.01)
yield b"tick"
response = ClosingStreamingResponse(endless(), media_type="text/plain", on_close=on_close)
with anyio.fail_after(2):
await stream_then_disconnect(response, "/", chunks_before_disconnect=0)
assert released["count"] == 1
+79 -35
View File
@@ -1,69 +1,84 @@
"""Tests for ``docsgpt/api/async_sse.py``.
Native-async reconnect endpoint: GET /api/messages/<id>/events. Auth gate,
ownership gate, malformed-id rejection, Last-Event-ID normalisation, and the
SSE response shape (headers + ``: connected`` prelude). The route is a
Starlette endpoint, so it's driven through Starlette's TestClient over a
minimal app built from ``async_sse_routes``.
ownership gate, malformed-id rejection, Last-Event-ID normalisation, the
per-user connection lease, and the SSE response shape (headers + ``: connected``
prelude). The route is a Starlette endpoint, so it's driven through Starlette's
TestClient over a minimal app built from ``async_sse_routes``.
"""
from __future__ import annotations
import time
from unittest.mock import AsyncMock, patch
import pytest
from starlette.applications import Starlette
from starlette.testclient import TestClient
from docsgpt.api import async_sse
from docsgpt.api.async_sse import (
_MESSAGE_ID_RE,
_normalise_last_event_id,
async_sse_routes,
)
from docsgpt.core.settings import settings
from tests.asgi_stream import stream_then_disconnect
from tests.fake_async_redis import FakeAsyncRedis
VALID_UUID = "67d65e8f-e7fb-4df1-9e6e-99ea6c830206"
_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 == {}
File diff suppressed because it is too large. Load diff
+395 -215
View File
@@ -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
+88
View File
@@ -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
+46
View File
@@ -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
+502
View File
@@ -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/<session_id>/events" not in rules
assert "/api/devices/poll" in rules
assert "/api/devices/sessions/<session_id>/invocations/<invocation_id>/ack" in rules
assert "/api/devices/sessions/<session_id>/invocations/<invocation_id>/output" in rules
+135 -66
View File
@@ -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"
+153
View File
@@ -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
+68
View File
@@ -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"
+8 -2
View File
@@ -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)
+40
View File
@@ -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
@@ -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
)
+6 -3
View File
@@ -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
+89
View File
@@ -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/<session_id>/events" not in rules
assert "/api/artifacts/<artifact_id>/download" not in rules
# Their siblings stay on Flask.
assert "/api/devices/poll" in rules
assert "/api/artifacts/<artifact_id>" 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"}
+187
View File
@@ -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
+78
View File
@@ -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
+3 -3
View File
@@ -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 ───────────────────────────────────────────────────
+216
View File
@@ -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:<id>: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"]
+9 -8
View File
@@ -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(
+167 -10
View File
@@ -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