"""``/query`` POST + ``/query/runs/{runId}/events`` SSE proxy. CONTRACT.md §6 + §7. ``POST /query`` returns **202** with ``{jobId, sseUrl}`` — never the runId, which arrives in the terminal ``done`` event payload once the run commits. ``GET /query/runs/{runId}/events`` is a thin alias that resolves runId → jobId via ``jobs.subject_id`` and proxies the underlying ``/jobs/{jobId}/events`` stream. ``GET /query/runs/{runId}`` for committed runs already exists in ``src/api/routers/history.py`` (Stage 3); we don't touch it. """ from __future__ import annotations import logging import uuid from fastapi import APIRouter, Depends, Header, Path, Query, Request from fastapi.responses import StreamingResponse from src.api.deps import forbid_readers, require_query_access from src.api.dto.common import JobStartResponse from src.api.dto.query import QueryRequest from src.api.errors import QueryRunNotFound from src.api.jobs import runner, store from src.api.limits import enforce_paid_request from src.api.jobs.sse import event_stream from src.api.validation import assert_known_model from src.lib.auth.models import User # Side-effect import: registers the "query" job handler with the registry. # Without this import, runner.enqueue would raise KeyError("No job handler # registered for type 'query'"). Stage-5 will add a sibling import for # "ingest"; we keep it here (not in app.py) so the router file is # self-contained. from src.api.jobs.types import query as _query_job # noqa: F401 logger = logging.getLogger(__name__) router = APIRouter(tags=["query"]) def _validate_model_overrides(models, *, allow_premium: bool) -> None: """Reject unknown per-stage model overrides before any paid work. ``models.{generation,classify,judge}`` are free-form strings that flow straight to the provider, so without this a query-capable user could pin every stage (incl. ~50 parallel CONCEPTUAL judge calls) to the most expensive model and burn the daily cap far faster. We constrain each override to a model that actually exists in the tools registry's ``llm`` category (the same set the FE offers via ``GET /tools?category=llm``). ``allow_premium`` is the caller's admin flag: a non-admin query user may NOT pin premium-priced (claude-opus-class) models, so one leaked/low-trust login can't drain the shared daily spend cap in a few queries. """ if models is None: return for field in ("generation", "classify", "judge"): assert_known_model( getattr(models, field, None), field=field, allow_premium=allow_premium ) @router.post("/query", status_code=202, response_model=JobStartResponse) def post_query( payload: QueryRequest, user: User = Depends(require_query_access) ) -> JobStartResponse: """Kick off a query job. Returns 202 + ``{jobId, sseUrl}``. No ``runId`` in this response: live queries are addressed by ``jobId`` until the terminal ``done`` event commits the run, at which point ``done.result.runId`` arrives and the FE swaps its URL. See CONTRACT.md §7 for the locked design. The per-subject lock is keyed ``query:{username}`` so a single user's sequential queries serialize (avoiding accidental parallel paid calls); two users hit the global ``max_concurrent_jobs`` semaphore independently. Access + cost guardrails run *before* anything is enqueued: :func:`require_query_access` blocks read-only viewers, and :func:`enforce_paid_request` enforces the per-user rate limit and the daily USD spend caps so no single login can drive a runaway bill. """ enforce_paid_request(user) _validate_model_overrides(payload.models, allow_premium=user.is_admin) job_id = uuid.uuid4().hex # Persist the request body so a startup-after-crash reconciler can # see what was running; also gives ``GET /jobs/{id}`` a meaningful # params payload before the runner even starts. params = payload.model_dump() # Stash the caller's username for the runner — query_history wants it # but the wire DTO doesn't carry it. Prefix with `_` so we strip it # back out before persisting (see query.py:_dump_persistable). params["_username"] = user.username store.insert_job( id=job_id, type="query", params_json=params, username=user.username, ) runner.enqueue( job_id, type_="query", params=params, subject_key=f"user:{user.username}", ) return JobStartResponse(job_id=job_id, sse_url=f"/jobs/{job_id}/events") @router.get("/query/runs/{runId}/events") async def stream_query_run_events( request: Request, run_id: int = Path(..., alias="runId", ge=1), last_event_id_header: str | None = Header(default=None, alias="Last-Event-ID"), last_event_id_query: int | None = Query(default=None, alias="lastEventId"), user: User = Depends(forbid_readers), ) -> StreamingResponse: """SSE alias that resolves ``runId`` → ``jobId`` and proxies the stream. The FE flips its URL from ``?job=...`` to ``?run={runId}`` after the terminal ``done`` event fires. On a deep-link refresh the FE hits this endpoint with the runId; we look up the underlying job by ``jobs.subject_id = str(runId)`` and reuse the same ``event_stream`` generator the ``/jobs/{id}/events`` endpoint uses. 404 ``QUERY_RUN_NOT_FOUND`` when the runId has no associated job — which can mean either "the run never existed" or "the job's subject_id was never bound" (the runner only writes subject_id on successful persist_run, so failed live runs land here too). """ job_row = store.fetch_job_by_subject(type="query", subject_id=str(run_id)) if job_row is None: raise QueryRunNotFound(run_id=run_id) # Tenancy: a non-owner can't tail another user's run stream (same 404 as # a missing run — no oracle). Admins may watch any run. if not user.is_admin and (getattr(job_row, "username", None) or "") != user.username: raise QueryRunNotFound(run_id=run_id) # Header takes precedence over the query param fallback. last_id = 0 if last_event_id_header: try: last_id = int(last_event_id_header) except ValueError: last_id = 0 elif last_event_id_query is not None: last_id = max(0, int(last_event_id_query)) job_id = job_row.id async def _gen(): async for frame in event_stream(job_id, last_event_id=last_id): if await request.is_disconnected(): return yield frame return StreamingResponse( _gen(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache, no-transform", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, )