src/api/routers/query.py

"""``/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",
        },
    )