mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 18:15:02 +00:00
Merge branch 'arc53:main' into main
This commit is contained in:
60 files changed
+4696
-1690
No files matched your search
@@ -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
|
||||
|
||||
@@ -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:**
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
@@ -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
@@ -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"]),
|
||||
]
|
||||
@@ -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.
|
||||
|
||||
@@ -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"]),
|
||||
]
|
||||
@@ -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")
|
||||
|
||||
@@ -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
@@ -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=[
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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: "
|
||||
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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)},
|
||||
)
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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': {
|
||||
|
||||
@@ -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"}
|
||||
@@ -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
|
||||
@@ -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 == {}
|
||||
+525
-625
File diff suppressed because it is too large.
Load diff
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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"}
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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 ───────────────────────────────────────────────────
|
||||
|
||||
@@ -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"]
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in new issue
Block a user