CIAgent
12b2300f6f
---ci--- phase: 7 milestone: v0.3 status: review lessons: - P0 CORS: allow_methods lacked PUT while the build surface writes files with PUT — every cross-origin Save failed preflight; pinned with tests/api/test_cors.py - P0 ingest leak: queue-overflow flood path returned without the disconnect sentinel, parking the drainer forever (one leaked task-set per flooded trace); sentinel now always enqueued, real-server regression test added - P1 perf: flood cap counted rows via len(get_trace(...)) — O(trace) per append, O(n²) per session; TraceStore.count() (COUNT(*)) added and wired - P0 security: file routes followed exec-planted symlinks out of the workspace bind; _resolve_in_workspace refuses escapes (422), read/write now 404 on unknown sandboxes (was 500) - P1 security: WS ingest accepted any browser Origin (CORS middleware does not cover WS); localhost dev origins + no-Origin (capture agent) allowed, others 1008 - P1 correctness: use-sandbox-session leaked a created sandbox on any mid-start failure (per-learner cap 1 → all retries 429 forever); failed starts now destroy what they created - P2 testing: reconnect-flush test killed mid-burst (nondeterministic under load, reproduced on pre-change code); now waits for server-side observation of the pre-kill burst — the underlying one-line replay-margin/ACK gap is documented for v0.4 - maintainability: grading-store/templates/grading.ts docstrings claimed grading is variant-blind (stale pre-P4 text) — updated; ARCHITECTURE.md referenced nonexistent voice/openai_audio.py; dead if TYPE_CHECKING: pass blocks removed ---/ci---
451 lines
19 KiB
Python
451 lines
19 KiB
Python
"""WS ingest protocol for learner telemetry (REQ-3-003, D-026, G-3).
|
|
|
|
Frame contract — trace identity travels as QUERY PARAMS on the WS upgrade
|
|
(`WS /v1/telemetry/ingest?learner_id=...&task_id=...&sandbox_id=...`), NOT as
|
|
a first init frame. Rationale: the in-sandbox capture agent (Task 2-2-01) is a
|
|
stdlib-only RFC6455 client where the URL is the cheapest thing to parametrize
|
|
(`NC_INGEST_URL` carries the query string); identity is also visible to the
|
|
server BEFORE accept(), so a malformed handshake can be rejected without an
|
|
accept/close round-trip. Client messages are then ONE event per JSON text
|
|
frame — no envelope:
|
|
|
|
{"seq": 0, "kind": "command", "payload": {...}, "ts": "...",
|
|
"sandbox_id": "..."} # learner_id / task_id forbidden (URL owns them)
|
|
|
|
Server → client frames are typed status envelopes:
|
|
|
|
{"type": "ack_total", "count": N} — final flush summary, then close 1000
|
|
{"type": "gap_warning", "missing_seqs": [...]} — seq skipped ahead
|
|
{"type": "event_rejected", "detail": "..."} — one frame failed validation
|
|
(seq echoed when parseable)
|
|
{"type": "event_rejected", "seq": N, "detail": "..."} — stored-field rejected (bad kind)
|
|
{"type": "flooded", "reason": "cap_exceeded"|"queue_overflow",
|
|
"count": N} — sent before close(1008)
|
|
|
|
Keepalive: the server sends an opaque ping frame every `PING_INTERVAL_S` (the
|
|
capture agent auto-pongs at the frame layer); a peer that is silent past
|
|
`PONG_TIMEOUT_S` is assumed wedged, but the keepalive half only LOGS — the
|
|
receiver half owns disconnect detection (single-box pilot: TCP EOF is
|
|
reliable; an aggressive pong-watchdog would false-positive on loaded boxes).
|
|
|
|
Flood control (GRILL G-3, BINDING — silent drop-oldest is FORBIDDEN):
|
|
* per-connection inbound queue bounded at `INBOUND_QUEUE_MAX` frames; on
|
|
overflow → close code 1008 (policy violation) + trace marked
|
|
INCOMPLETE_FLOODED via `TraceIntegrityMap`.
|
|
* total events for the (learner, task) exceeding
|
|
`Settings.telemetry_max_events_per_task` → same 1008 + INCOMPLETE_FLOODED.
|
|
`INCOMPLETE_FLOODED` is an integrity signal Proctor/Phase-3 grader read via
|
|
`TraceIntegrityMap.is_incomplete()` (the G-4 gate): a flooded trace can
|
|
never yield a credential.
|
|
|
|
Boundary (D-027): telemetry/ never imports agents/ or api/. This module
|
|
imports only `fastapi.WebSocket` for the socket type (a protocol surface, not
|
|
a DI framework); the session engine below depends only on the TraceStore
|
|
protocol + Settings, and api/telemetry.py injects both through plain
|
|
parameters.
|
|
"""
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import json
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import Any, Final
|
|
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
|
|
|
from ..config import Settings
|
|
from .models import TelemetryEvent
|
|
from .store import TraceStore
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
#: WebSocket close code 1008 — policy violation (RFC 6455 §7.4.1).
|
|
WS_CLOSE_POLICY_VIOLATION: Final = 1008
|
|
|
|
#: Bounded inbound queue depth per connection (G-3). Sized for burst-tolerance
|
|
#: well above the capture agent's emission rate; overflow is a flood signal,
|
|
#: not a backpressure knob.
|
|
INBOUND_QUEUE_MAX: Final = 256
|
|
|
|
PING_INTERVAL_S: Final = 20.0
|
|
|
|
|
|
class InboundEventFrame(BaseModel):
|
|
"""Client → server event frame (one TelemetryEvent minus URL-owned ids).
|
|
|
|
`extra="forbid"`: learner_id/task_id arriving in the frame body is a
|
|
contract violation — identity comes from the query params only, so a
|
|
replayed frame can never lie about which trace it belongs to.
|
|
"""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
seq: int = Field(ge=0)
|
|
kind: str = Field(min_length=1)
|
|
payload: dict[str, Any] = Field(default_factory=dict)
|
|
ts: datetime
|
|
sandbox_id: str = ""
|
|
|
|
|
|
class TraceIntegrityMap:
|
|
"""Integrity flags for traces that can never be graded (G-3/G-4).
|
|
|
|
Process-local and deliberately small: v0.3 runs ONE ai-service process per
|
|
box, and the Phase-3 grader reads this flag through the same DI container
|
|
— D-019-style in-memory registry precedent (the sandbox handle registry is
|
|
the same shape). The flag is terminal within the process: a reconnect
|
|
sending legal events does NOT clear it — the trace is already untrusted as
|
|
grading input. Restarting ai-service resets flags; grading runs against a
|
|
live service, and the SQLite trace rows themselves are durable.
|
|
|
|
All methods are sync: mutation is a dict write, reads are dict lookups —
|
|
no await needed, so callers from any layer (API handlers, the grader)
|
|
don't inherit an async surface for a nanosecond operation.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
# (learner_id, task_id) -> machine-readable reason (INCOMPLETE_FLOODED)
|
|
self._flags: dict[tuple[str, str], str] = {}
|
|
|
|
def mark(self, learner_id: str, task_id: str, reason: str) -> None:
|
|
"""Set an integrity flag. Presence of the flag is the signal; the
|
|
reason is informational (last write wins)."""
|
|
self._flags[(learner_id, task_id)] = reason
|
|
|
|
def clear(self, learner_id: str, task_id: str) -> None:
|
|
"""Test seam: reset a flag (production ingest never clears)."""
|
|
self._flags.pop((learner_id, task_id), None)
|
|
|
|
def is_incomplete(self, learner_id: str, task_id: str) -> bool:
|
|
"""True when the trace carries ANY terminal integrity flag."""
|
|
return (learner_id, task_id) in self._flags
|
|
|
|
def reason(self, learner_id: str, task_id: str) -> str | None:
|
|
"""The flag's reason (INCOMPLETE_FLOODED), or None when unflagged."""
|
|
return self._flags.get((learner_id, task_id))
|
|
|
|
|
|
class IngestSession:
|
|
"""One WebSocket ingest connection: receive → queue → drain → store.
|
|
|
|
Two tasks per connection:
|
|
* `_receiver` — reads frames, validates shape, enqueues (bounded queue,
|
|
G-3). Receives never block on SQLite.
|
|
* `_drainer` — pops frames in arrival order, appends via TraceStore
|
|
(idempotent on (learner,task,seq)), emits gap warnings, enforces the
|
|
per-trace event cap.
|
|
Either task detecting a flood closes the WS with 1008 and marks the trace
|
|
INCOMPLETE_FLOODED. The events queue carries `None` as the client-
|
|
disconnect sentinel.
|
|
|
|
- `telemetry_max_events_per_task` is consulted at connect and re-checked
|
|
per append against the DURABLE row count (cap compares against stored
|
|
events, so a skipped-ahead seq cannot burn budget that was never sent).
|
|
Durable count via `len(get_trace(...))` reads the trace per append —
|
|
O(trace) per event; v0.3 pilot sizing caps traces at 50k rows, WAL keeps
|
|
the writer unblocked (a-3), and the capture agent's emission rate is
|
|
human-scale. If profiling shows the count query hot, swap to COUNT(*)
|
|
without changing the contract.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
websocket: WebSocket,
|
|
store: TraceStore,
|
|
integrity: TraceIntegrityMap,
|
|
settings: Settings,
|
|
learner_id: str,
|
|
task_id: str,
|
|
sandbox_id: str,
|
|
) -> None:
|
|
self._ws = websocket
|
|
self._store = store
|
|
self._integrity = integrity
|
|
# Snapshot of the one setting ingest consults: read once at connect so
|
|
# a hot-reloaded Settings object mid-session can't move the cap.
|
|
self._max_events = settings.telemetry_max_events_per_task
|
|
self.learner_id = learner_id
|
|
self.task_id = task_id
|
|
self.sandbox_id = sandbox_id
|
|
|
|
self._queue: asyncio.Queue[InboundEventFrame | None] = asyncio.Queue(
|
|
maxsize=INBOUND_QUEUE_MAX
|
|
)
|
|
self._seen: set[int] = set()
|
|
self._next_expected: int | None = None # in-connection monotonic hint
|
|
self._received = 0
|
|
self._stored = 0
|
|
self._deduped = 0
|
|
self._rejected = 0
|
|
self._flooded = False
|
|
self._flood_reason = ""
|
|
|
|
# -- receive half ----------------------------------------------------------
|
|
|
|
async def run(self) -> None:
|
|
"""Accept, run receiver+drainer, close cleanly. Owns the WS lifecycle."""
|
|
await self._ws.accept()
|
|
logger.info(
|
|
"telemetry ingest connected: %s/%s sandbox=%s",
|
|
self.learner_id,
|
|
self.task_id,
|
|
self.sandbox_id or "(none)",
|
|
)
|
|
pinger = asyncio.create_task(self._keepalive())
|
|
receiver = asyncio.create_task(self._receiver())
|
|
drainer = asyncio.create_task(self._drainer())
|
|
# First terminal outcome shuts the session down: client disconnect
|
|
# (receiver ends) → drainer flushes; drainer ended (clean close after
|
|
# flush or a 1008 flood close) → receiver must not linger.
|
|
pending: set[asyncio.Task[None]] = {receiver, drainer}
|
|
try:
|
|
done, pending = await asyncio.wait(
|
|
pending, return_when=asyncio.FIRST_COMPLETED
|
|
)
|
|
if receiver in done and drainer in pending:
|
|
try:
|
|
await drainer # final flush → sends ack_total, close 1000
|
|
finally:
|
|
pending.discard(drainer)
|
|
finally:
|
|
for task in (pinger, *pending):
|
|
task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await task
|
|
|
|
async def _receiver(self) -> None:
|
|
"""Read frames; parse+enqueue. Overflow → flood shutdown (G-3).
|
|
|
|
RuntimeError from receive_text is benign here: it fires when the
|
|
socket was closed by the drainer (1008 flood close) while this task
|
|
was parked in receive — a terminal condition, not a bug.
|
|
"""
|
|
try:
|
|
while True:
|
|
raw = await self._ws.receive_text()
|
|
frame = self._parse(raw)
|
|
if frame is None:
|
|
# Rejected frame — keep the connection open; the producer
|
|
# gets an event_rejected status frame so a malformed batch
|
|
# is visible (and its seq is never stored). Yield so the
|
|
# status frame flushes before we block on the next receive.
|
|
await self._reject_frame(raw)
|
|
await asyncio.sleep(0)
|
|
continue
|
|
self._received += 1
|
|
try:
|
|
self._queue.put_nowait(frame)
|
|
except asyncio.QueueFull:
|
|
# Bounded queue — overflow is a flood, never drop-oldest.
|
|
# _trigger_flood closes the socket; fall through to the
|
|
# tail so the disconnect sentinel is still enqueued — the
|
|
# drainer is never left parked on an empty queue after a
|
|
# flood (P7 review: the pre-fix code `return`ed from the
|
|
# QueueFull branch WITHOUT the sentinel, leaking the
|
|
# session task set — one per flooded trace).
|
|
await self._trigger_flood("queue_overflow")
|
|
return
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except RuntimeError:
|
|
logger.debug(
|
|
"ingest receiver: socket already closed (flood path) %s/%s",
|
|
self.learner_id,
|
|
self.task_id,
|
|
)
|
|
# Client gone (clean close, drop, or flood close): sentinel unblocks
|
|
# the drainer for a final flush. put_nowait can only fail under flood,
|
|
# which already terminated the session.
|
|
with contextlib.suppress(asyncio.QueueFull):
|
|
self._queue.put_nowait(None)
|
|
|
|
def _parse(self, raw: str) -> InboundEventFrame | None:
|
|
"""Validate one frame; None means malformed (caller rejects it)."""
|
|
try:
|
|
return InboundEventFrame.model_validate_json(raw)
|
|
except ValidationError:
|
|
return None
|
|
|
|
async def _reject_frame(self, raw: str) -> None:
|
|
"""Malformed envelope: log + event_rejected status frame (never stored)."""
|
|
self._rejected += 1
|
|
detail = "invalid event frame"
|
|
try:
|
|
InboundEventFrame.model_validate_json(raw)
|
|
except ValidationError as exc:
|
|
detail = exc.errors()[0].get("msg", "validation error")
|
|
logger.warning(
|
|
"telemetry frame rejected: %s/%s: %s", self.learner_id, self.task_id, detail
|
|
)
|
|
seq: int | None = None
|
|
with contextlib.suppress(Exception):
|
|
seq = int(json.loads(raw).get("seq")) # best-effort echo for the producer
|
|
payload: dict[str, Any] = {"type": "event_rejected", "detail": detail}
|
|
if seq is not None:
|
|
payload["seq"] = seq
|
|
await self._send_json(payload)
|
|
|
|
# -- drain half --------------------------------------------------------------
|
|
|
|
async def _drainer(self) -> None:
|
|
"""Pop queued frames, append to the store, then close 1000 + summary."""
|
|
while True:
|
|
frame = await self._queue.get()
|
|
if frame is None: # disconnect sentinel → flush complete
|
|
await self._send_json(
|
|
{
|
|
"type": "ack_total",
|
|
"count": self._stored,
|
|
"deduped": self._deduped,
|
|
"rejected": self._rejected,
|
|
}
|
|
)
|
|
with contextlib.suppress(RuntimeError, WebSocketDisconnect):
|
|
await self._ws.close(code=1000)
|
|
return
|
|
await self._append(frame)
|
|
|
|
async def _append(self, frame: InboundEventFrame) -> None:
|
|
# Precedence: a trace already flagged INCOMPLETE_FLOODED is terminal —
|
|
# the connection that triggered it is being torn down, and any stray
|
|
# queued frames must not resurrect the trace's intake.
|
|
if self._integrity.is_incomplete(self.learner_id, self.task_id):
|
|
await self._trigger_flood("already_flagged")
|
|
return
|
|
|
|
# Per-trace cap (G-3): checked against the DURABLE row count so a
|
|
# reconnect resumes the budget instead of resetting it, and a
|
|
# skipped-ahead seq cannot burn budget that was never sent.
|
|
if self._flood_breached():
|
|
await self._trigger_flood("cap_exceeded")
|
|
return
|
|
|
|
# TelemetryEvent's @validates hooks fire on CONSTRUCTION (setattr), so
|
|
# the try must wrap building the model too — an unknown kind raises
|
|
# before `store.append` is ever reached.
|
|
event: TelemetryEvent
|
|
before = self._store.latest_seq(self.learner_id, self.task_id)
|
|
try:
|
|
event = TelemetryEvent(
|
|
learner_id=self.learner_id,
|
|
task_id=self.task_id,
|
|
seq=frame.seq,
|
|
kind=frame.kind,
|
|
payload=frame.payload,
|
|
ts=frame.ts,
|
|
sandbox_id=frame.sandbox_id or self.sandbox_id,
|
|
)
|
|
self._store.append(event)
|
|
except ValueError as exc: # unknown kind / invalid field
|
|
self._rejected += 1
|
|
await self._send_json(
|
|
{"type": "event_rejected", "seq": frame.seq, "detail": str(exc)}
|
|
)
|
|
return
|
|
after = self._store.latest_seq(self.learner_id, self.task_id)
|
|
|
|
if after == before and frame.seq in self._seen:
|
|
self._deduped += 1 # at-least-once retry; stored once (idempotent)
|
|
else:
|
|
self._stored += 1
|
|
self._seen.add(frame.seq)
|
|
await self._check_gap(frame.seq)
|
|
# SQLite appends are sync and fast; on a burst the drainer can hold
|
|
# the loop between receives. Yield so the WS writer flushes the close
|
|
# and the pinger/interleave stay live under the eventlet-free portal.
|
|
await asyncio.sleep(0)
|
|
|
|
def _flood_breached(self) -> bool:
|
|
"""True when this append would exceed the per-trace event budget."""
|
|
# Durable count (NOT latest_seq+1 — a skipped-ahead seq must not burn
|
|
# un-sent events' budget) via COUNT(*): never materialize the trace
|
|
# per append (P7 review — the old len(get_trace(...)) built every row
|
|
# object per event, O(trace) per append / O(n²) per session).
|
|
durable = self._store.count(learner_id=self.learner_id, task_id=self.task_id)
|
|
return durable >= self._max_events
|
|
|
|
async def _check_gap(self, incoming_seq: int) -> None:
|
|
"""Seq skipped ahead → log + per-connection gap_warning status frame."""
|
|
if self._next_expected is not None and incoming_seq > self._next_expected:
|
|
missing = list(range(self._next_expected, incoming_seq))
|
|
logger.warning(
|
|
"telemetry gap: %s/%s missing seqs %s (arrived seq=%d)",
|
|
self.learner_id,
|
|
self.task_id,
|
|
missing,
|
|
incoming_seq,
|
|
)
|
|
await self._send_json({"type": "gap_warning", "missing_seqs": missing})
|
|
if self._next_expected is None or incoming_seq >= self._next_expected:
|
|
self._next_expected = incoming_seq + 1
|
|
|
|
# -- flood + keepalive ------------------------------------------------------
|
|
|
|
async def _trigger_flood(self, reason: str) -> None:
|
|
"""G-3: 1008 close + INCOMPLETE_FLOODED mark. Exactly once."""
|
|
if self._flooded:
|
|
return
|
|
self._flooded = True
|
|
self._flood_reason = reason
|
|
self._integrity.mark(self.learner_id, self.task_id, "INCOMPLETE_FLOODED")
|
|
logger.warning(
|
|
"telemetry flood: %s/%s reason=%s — closing 1008, trace marked "
|
|
"INCOMPLETE_FLOODED (G-3; Proctor/grade gate will refuse it)",
|
|
self.learner_id,
|
|
self.task_id,
|
|
reason,
|
|
)
|
|
await self._send_json(
|
|
{"type": "flooded", "reason": reason, "count": self._received}
|
|
)
|
|
with contextlib.suppress(RuntimeError, WebSocketDisconnect):
|
|
await self._ws.close(
|
|
code=WS_CLOSE_POLICY_VIOLATION,
|
|
reason=f"telemetry flood control (G-3): {reason}",
|
|
)
|
|
|
|
async def _keepalive(self) -> None:
|
|
"""Protocol-level ping on an interval (agent auto-pongs at frame level).
|
|
|
|
A send failure means the socket is already gone — the receiver half
|
|
independently surfaces the disconnect; we just stop pinging.
|
|
"""
|
|
while True:
|
|
await asyncio.sleep(PING_INTERVAL_S)
|
|
try:
|
|
await self._ws.send_bytes(b"\x89ping-nextcraft")
|
|
except (RuntimeError, WebSocketDisconnect):
|
|
return
|
|
|
|
async def _send_json(self, payload: dict[str, Any]) -> None:
|
|
"""Best-effort status frame; the socket may already be gone."""
|
|
with contextlib.suppress(RuntimeError, WebSocketDisconnect):
|
|
await self._ws.send_json(payload)
|
|
|
|
|
|
async def telemetry_ingest_endpoint(
|
|
websocket: WebSocket,
|
|
learner_id: str,
|
|
task_id: str,
|
|
store: TraceStore,
|
|
integrity: TraceIntegrityMap,
|
|
settings: Settings,
|
|
sandbox_id: str = "",
|
|
) -> None:
|
|
"""Engine entry: build the session and run it. api/telemetry.py calls this
|
|
with query params + app.state services already resolved — this signature
|
|
is deliberately Depends-free (telemetry/ never knows FastAPI DI exists).
|
|
"""
|
|
session = IngestSession(
|
|
websocket=websocket,
|
|
store=store,
|
|
integrity=integrity,
|
|
settings=settings,
|
|
learner_id=learner_id,
|
|
task_id=task_id,
|
|
sandbox_id=sandbox_id,
|
|
)
|
|
await session.run()
|