nextcraft/apps/ai-service/ai_service/api/telemetry.py

"""/v1/telemetry — WS ingest + trace read endpoints (REQ-3-003, D-026, G-3).

The router composes the telemetry engine via DI: `telemetry/ingest.py` owns
the WS protocol (frame contract + flood control + keepalive) and this module
only wires `app.state.trace_store` / `app.state.trace_integrity` /
`app.state.settings` into it, plus the two HTTP read faces:

    WS   /v1/telemetry/ingest?learner_id&task_id[&sandbox_id]   (D-026)
    GET  /v1/telemetry/traces/{learner_id}/{task_id}            ordered trace; 404 unknown
    GET  /v1/telemetry/gaps/{learner_id}/{task_id}              missing seqs   ; 404 unknown

The WS route is a thin DI shell: it validates the query-param identity and
the Origin (browser pages are gated to the localhost dev origins — CORS
middleware does not cover WS upgrades; the stdlib capture agent sends no
Origin and is unaffected), pulls store/integrity/settings from `app.state`,
and calls `telemetry_ingest_endpoint(...)` — the engine stays
FastAPI-DI-free so it's testable without a router and the api/ layer owns
all composition.

Unknown-trace contract: a trace is KNOWN when it has >=1 stored event OR
carries an integrity flag — a flooded trace with zero stored rows still 200s
so Proctor/grader can read WHY it's unusable (G-4 consumes
`integrity_reason`). `TraceResponse.incomplete` / `.integrity_reason` mirror
the map so HTTP consumers never touch process internals.
"""

from __future__ import annotations

from fastapi import APIRouter, Depends, HTTPException, WebSocket
from pydantic import BaseModel

from ..telemetry.ingest import (
    TraceIntegrityMap,
    telemetry_ingest_endpoint,
)
from ..telemetry.models import TelemetryEvent
from ..telemetry.store import TraceStore
from .deps import get_trace_integrity, get_trace_store

router = APIRouter(prefix="/v1/telemetry", tags=["telemetry"])

#: Browser Origins allowed to open the ingest socket (A-008 mirror). The
#: stdlib capture agent sends NO Origin header (it is not a browser) and
#: stays allowed; a malicious page loaded in the learner's browser would
#: carry an Origin and must not be able to poison/flood the trace. CORS
#: middleware does NOT cover WebSocket upgrades, so this gate is explicit.
_ALLOWED_WS_ORIGINS = frozenset(
    {"http://localhost:3000", "http://127.0.0.1:3000", "http://localhost:8420"}
)


# --- WS ingest (D-026) ---------------------------------------------------------


@router.websocket("/ingest")
async def telemetry_ingest_ws(websocket: WebSocket) -> None:
    """DI shell: resolve app.state services, then hand the socket to the engine.

    The engine's session + flood logic is fully typed and testable without
    FastAPI; this shim is the only place the two layers meet.
    """
    origin = (websocket.headers.get("origin") or "").strip()
    if origin and origin not in _ALLOWED_WS_ORIGINS:
        # Same-origin dev pages (Next.js on :3000, the service itself on
        # :8420) pass; anything else is refused pre-accept. Non-browser
        # producers (the capture agent, tests) send no Origin and pass.
        await websocket.close(
            code=1008, reason=f"origin {origin!r} not allowed for telemetry ingest"
        )
        return
    query = websocket.query_params
    learner_id = query.get("learner_id", "")
    task_id = query.get("task_id", "")
    if not learner_id or not task_id:
        # Reject BEFORE accept: closing the pre-accept handshake is the
        # cheapest denial and unambiguous for the stdlib capture agent.
        await websocket.close(code=1008, reason="learner_id/task_id query params required")
        return
    await telemetry_ingest_endpoint(
        websocket=websocket,
        learner_id=learner_id,
        task_id=task_id,
        sandbox_id=query.get("sandbox_id", ""),
        store=websocket.app.state.trace_store,
        integrity=websocket.app.state.trace_integrity,
        settings=websocket.app.state.settings,
    )


# --- HTTP reads -----------------------------------------------------------------


class TraceResponse(BaseModel):
    """Ordered trace + integrity signal (G-4 reads incomplete/reason)."""

    learner_id: str
    task_id: str
    events: list[TelemetryEvent]
    incomplete: bool
    integrity_reason: str | None


class GapsResponse(BaseModel):
    """Missing seqs + integrity signal."""

    learner_id: str
    task_id: str
    gaps: list[int]
    incomplete: bool
    integrity_reason: str | None


def _is_known_trace(
    store: TraceStore, integrity: TraceIntegrityMap, learner_id: str, task_id: str
) -> bool:
    """Known = has stored events OR carries an integrity flag (a flooded trace
    with zero rows must still be readable — Proctor needs the reason)."""
    return (
        store.latest_seq(learner_id, task_id) >= 0
        or integrity.is_incomplete(learner_id, task_id)
    )


@router.get("/traces/{learner_id}/{task_id}", response_model=TraceResponse)
async def get_trace(
    learner_id: str,
    task_id: str,
    store: TraceStore = Depends(get_trace_store),
    integrity: TraceIntegrityMap = Depends(get_trace_integrity),
) -> TraceResponse:
    if not _is_known_trace(store, integrity, learner_id, task_id):
        raise HTTPException(
            status_code=404, detail=f"unknown trace {learner_id!r}/{task_id!r}"
        )
    return TraceResponse(
        learner_id=learner_id,
        task_id=task_id,
        events=store.get_trace(learner_id, task_id),
        incomplete=integrity.is_incomplete(learner_id, task_id),
        integrity_reason=integrity.reason(learner_id, task_id),
    )


@router.get("/gaps/{learner_id}/{task_id}", response_model=GapsResponse)
async def get_gaps(
    learner_id: str,
    task_id: str,
    store: TraceStore = Depends(get_trace_store),
    integrity: TraceIntegrityMap = Depends(get_trace_integrity),
) -> GapsResponse:
    if not _is_known_trace(store, integrity, learner_id, task_id):
        raise HTTPException(
            status_code=404, detail=f"unknown trace {learner_id!r}/{task_id!r}"
        )
    return GapsResponse(
        learner_id=learner_id,
        task_id=task_id,
        gaps=store.gaps(learner_id, task_id),
        incomplete=integrity.is_incomplete(learner_id, task_id),
        integrity_reason=integrity.reason(learner_id, task_id),
    )