diff --git a/frontend/src/hooks/useRunPolling.test.tsx b/frontend/src/hooks/useRunPolling.test.tsx new file mode 100644 index 0000000..9753fb6 --- /dev/null +++ b/frontend/src/hooks/useRunPolling.test.tsx @@ -0,0 +1,121 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' +import { renderHook, act } from '@testing-library/react' +import { useRunPolling, POLL_INTERVAL_MS } from './useRunPolling' +import type { Run } from '../api' + +class FakeWebSocket { + static instances: FakeWebSocket[] = [] + static readonly OPEN = 1 + static readonly CLOSED = 3 + + readyState = FakeWebSocket.OPEN + onmessage: ((event: { data: string }) => void) | null = null + onerror: (() => void) | null = null + + constructor(public url: string) { + FakeWebSocket.instances.push(this) + } + + close() { + this.readyState = FakeWebSocket.CLOSED + } + + emit(payload: unknown) { + this.onmessage?.({ data: JSON.stringify(payload) }) + } + + fail() { + this.onerror?.() + } +} + +const originalWebSocket = globalThis.WebSocket + +function run(id: number, status: Run['status']): Run { + return { + id, + kind: 'demo', + status, + analysis_type: 'deforestation', + created_at: '2026-01-01T00:00:00Z', + updated_at: '2026-01-01T00:00:00Z', + } +} + +beforeEach(() => { + FakeWebSocket.instances = [] + vi.stubGlobal('WebSocket', FakeWebSocket) + vi.useFakeTimers() +}) + +afterEach(() => { + vi.useRealTimers() + vi.stubGlobal('WebSocket', originalWebSocket) +}) + +describe('useRunPolling', () => { + it('does not poll while the WebSocket path is healthy', () => { + const fetchRuns = vi.fn() + renderHook(() => useRunPolling([run(1, 'running')], fetchRuns)) + + act(() => { + vi.advanceTimersByTime(POLL_INTERVAL_MS * 3) + }) + + expect(fetchRuns).not.toHaveBeenCalled() + }) + + it('refetches as soon as a run reports a terminal status', () => { + const fetchRuns = vi.fn() + renderHook(() => useRunPolling([run(1, 'running')], fetchRuns)) + + act(() => { + FakeWebSocket.instances[0].emit({ type: 'status', run_id: 1, status: 'completed' }) + }) + + expect(fetchRuns).toHaveBeenCalledTimes(1) + }) + + it('falls back to polling once a socket errors', () => { + const fetchRuns = vi.fn() + renderHook(() => useRunPolling([run(1, 'running')], fetchRuns)) + + // The error flips the hook to unhealthy, which starts the interval. + // `act` flushes that state update synchronously, so no `waitFor` is + // needed — and `waitFor` would deadlock against the fake timers here. + act(() => { + FakeWebSocket.instances[0].fail() + }) + + act(() => { + vi.advanceTimersByTime(POLL_INTERVAL_MS * 2) + }) + + expect(fetchRuns).toHaveBeenCalledTimes(2) + }) + + it('neither polls nor connects when nothing is running', () => { + const fetchRuns = vi.fn() + renderHook(() => useRunPolling([run(1, 'completed')], fetchRuns)) + + act(() => { + vi.advanceTimersByTime(POLL_INTERVAL_MS * 3) + }) + + expect(FakeWebSocket.instances).toHaveLength(0) + expect(fetchRuns).not.toHaveBeenCalled() + }) + + it('still reports the running → completed transition', () => { + const onCompleted = vi.fn() + const { rerender } = renderHook( + ({ runs }) => useRunPolling(runs, vi.fn(), onCompleted), + { initialProps: { runs: [run(1, 'running')] } }, + ) + + rerender({ runs: [run(1, 'completed')] }) + + expect(onCompleted).toHaveBeenCalledTimes(1) + expect(onCompleted.mock.calls[0][0].id).toBe(1) + }) +}) diff --git a/frontend/src/hooks/useRunPolling.ts b/frontend/src/hooks/useRunPolling.ts index f076408..0ab2b82 100644 --- a/frontend/src/hooks/useRunPolling.ts +++ b/frontend/src/hooks/useRunPolling.ts @@ -1,23 +1,52 @@ -import { useEffect, useRef } from 'react' +import { useCallback, useEffect, useMemo, useRef } from 'react' import type { Run } from '../api' +import { useRunWebSocket } from './useRunWebSocket' +/** Interval used when falling back to polling, in milliseconds. */ +export const POLL_INTERVAL_MS = 5000 + +/** + * Keep `runs` fresh while any of them is still running. + * + * Prefers a WebSocket per in-flight run and refetches as soon as one reports a + * terminal status. Polling is only started when the WebSocket path is + * unavailable — no WebSocket implementation, or a socket that errored — so the + * hook degrades to its previous behaviour rather than losing updates. + */ export function useRunPolling( runs: Run[], fetchRuns: () => void, onCompleted?: (run: Run) => void, ) { const prevRunsRef = useRef>(new Map()) + const fetchRunsRef = useRef(fetchRuns) + + useEffect(() => { + fetchRunsRef.current = fetchRuns + }, [fetchRuns]) + + const runningIds = useMemo( + () => runs.filter((r) => r.status === 'running').map((r) => r.id), + [runs], + ) + + const handleTerminal = useCallback(() => { + fetchRunsRef.current() + }, []) + + const { healthy } = useRunWebSocket(runningIds, handleTerminal) useEffect(() => { - const hasRunning = runs.some((r) => r.status === 'running') - if (!hasRunning) return + if (runningIds.length === 0) return + // A healthy socket already pushes updates, so polling would be redundant. + if (healthy) return const interval = setInterval(() => { - fetchRuns() - }, 5000) + fetchRunsRef.current() + }, POLL_INTERVAL_MS) return () => clearInterval(interval) - }, [runs, fetchRuns]) + }, [runningIds, healthy]) // Detect transitions from running → completed useEffect(() => { diff --git a/frontend/src/hooks/useRunWebSocket.test.tsx b/frontend/src/hooks/useRunWebSocket.test.tsx new file mode 100644 index 0000000..1d8c3ca --- /dev/null +++ b/frontend/src/hooks/useRunWebSocket.test.tsx @@ -0,0 +1,173 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' +import { renderHook, act, waitFor } from '@testing-library/react' +import { + buildRunWebSocketUrl, + isWebSocketSupported, + useRunWebSocket, +} from './useRunWebSocket' + +/** Minimal scriptable WebSocket stand-in. */ +class FakeWebSocket { + static instances: FakeWebSocket[] = [] + static readonly CONNECTING = 0 + static readonly OPEN = 1 + static readonly CLOSING = 2 + static readonly CLOSED = 3 + + readyState = FakeWebSocket.OPEN + onmessage: ((event: { data: string }) => void) | null = null + onerror: (() => void) | null = null + closed = false + + constructor(public url: string) { + FakeWebSocket.instances.push(this) + } + + close() { + this.closed = true + this.readyState = FakeWebSocket.CLOSED + } + + emit(payload: unknown) { + this.onmessage?.({ data: JSON.stringify(payload) }) + } + + emitRaw(data: string) { + this.onmessage?.({ data }) + } + + fail() { + this.onerror?.() + } +} + +const originalWebSocket = globalThis.WebSocket + +beforeEach(() => { + FakeWebSocket.instances = [] + vi.stubGlobal('WebSocket', FakeWebSocket) +}) + +afterEach(() => { + vi.stubGlobal('WebSocket', originalWebSocket) + vi.unstubAllEnvs() +}) + +describe('buildRunWebSocketUrl', () => { + it('derives a ws:// URL from the page origin when no API base is set', () => { + expect(buildRunWebSocketUrl(7)).toBe(`${window.location.origin.replace(/^http/, 'ws')}/ws/runs/7`) + }) + + it('upgrades an https API base to wss', () => { + vi.stubEnv('VITE_API_BASE_URL', 'https://api.example.com') + expect(buildRunWebSocketUrl(3)).toBe('wss://api.example.com/ws/runs/3') + }) + + it('does not double up the slash when the base has a trailing one', () => { + vi.stubEnv('VITE_API_BASE_URL', 'https://api.example.com/') + expect(buildRunWebSocketUrl(3)).toBe('wss://api.example.com/ws/runs/3') + }) +}) + +describe('useRunWebSocket', () => { + it('opens one socket per running run', () => { + renderHook(() => useRunWebSocket([1, 2], vi.fn())) + expect(FakeWebSocket.instances).toHaveLength(2) + expect(FakeWebSocket.instances.map((s) => s.url)).toEqual([ + buildRunWebSocketUrl(1), + buildRunWebSocketUrl(2), + ]) + }) + + it('opens nothing when there are no running runs', () => { + renderHook(() => useRunWebSocket([], vi.fn())) + expect(FakeWebSocket.instances).toHaveLength(0) + }) + + it('invokes the callback on a terminal status', () => { + const onTerminal = vi.fn() + renderHook(() => useRunWebSocket([1], onTerminal)) + + act(() => { + FakeWebSocket.instances[0].emit({ type: 'status', run_id: 1, status: 'completed' }) + }) + + expect(onTerminal).toHaveBeenCalledTimes(1) + expect(onTerminal.mock.calls[0][0]).toMatchObject({ run_id: 1, status: 'completed' }) + }) + + it('treats a failed run as terminal', () => { + const onTerminal = vi.fn() + renderHook(() => useRunWebSocket([1], onTerminal)) + + act(() => { + FakeWebSocket.instances[0].emit({ + type: 'status', + run_id: 1, + status: 'failed', + error: 'boom', + }) + }) + + expect(onTerminal).toHaveBeenCalledTimes(1) + }) + + it('ignores a non-terminal status', () => { + const onTerminal = vi.fn() + renderHook(() => useRunWebSocket([1], onTerminal)) + + act(() => { + FakeWebSocket.instances[0].emit({ type: 'status', run_id: 1, status: 'running' }) + }) + + expect(onTerminal).not.toHaveBeenCalled() + }) + + it('survives an unparseable frame', () => { + const onTerminal = vi.fn() + renderHook(() => useRunWebSocket([1], onTerminal)) + + expect(() => { + act(() => { + FakeWebSocket.instances[0].emitRaw('not json') + }) + }).not.toThrow() + expect(onTerminal).not.toHaveBeenCalled() + }) + + it('reports unhealthy once a socket errors', async () => { + const { result } = renderHook(() => useRunWebSocket([1], vi.fn())) + expect(result.current.healthy).toBe(true) + + act(() => { + FakeWebSocket.instances[0].fail() + }) + + await waitFor(() => expect(result.current.healthy).toBe(false)) + }) + + it('reports unhealthy when the environment has no WebSocket', async () => { + vi.stubGlobal('WebSocket', undefined) + expect(isWebSocketSupported()).toBe(false) + + const { result } = renderHook(() => useRunWebSocket([1], vi.fn())) + await waitFor(() => expect(result.current.healthy).toBe(false)) + }) + + it('closes its sockets on unmount', () => { + const { unmount } = renderHook(() => useRunWebSocket([1, 2], vi.fn())) + unmount() + expect(FakeWebSocket.instances.every((s) => s.closed)).toBe(true) + }) + + it('does not reopen sockets when the id list is unchanged', () => { + const { rerender } = renderHook(({ ids }) => useRunWebSocket(ids, vi.fn()), { + initialProps: { ids: [1] }, + }) + expect(FakeWebSocket.instances).toHaveLength(1) + + // A fresh array with equal contents must not churn the connection. + rerender({ ids: [1] }) + expect(FakeWebSocket.instances).toHaveLength(1) + }) +}) diff --git a/frontend/src/hooks/useRunWebSocket.ts b/frontend/src/hooks/useRunWebSocket.ts new file mode 100644 index 0000000..dfd26f3 --- /dev/null +++ b/frontend/src/hooks/useRunWebSocket.ts @@ -0,0 +1,117 @@ +import { useEffect, useRef, useState } from 'react' +import { getApiBaseUrl } from '../api' +import type { RunStatus } from '../api' + +export interface RunStatusEvent { + type: 'status' | 'error' + run_id: number + status?: RunStatus + result?: Record + error?: string +} + +const TERMINAL_STATUSES: ReadonlySet = new Set(['completed', 'failed']) + +/** + * Resolve the WebSocket URL for a run. + * + * `getApiBaseUrl()` is empty by default, meaning "same origin", so fall back to + * the current page's host and swap http(s) for ws(s). + */ +export function buildRunWebSocketUrl(runId: number): string { + const base = getApiBaseUrl() + const origin = base !== '' ? base : window.location.origin + const wsOrigin = origin.replace(/^http/, 'ws') + return `${wsOrigin.replace(/\/$/, '')}/ws/runs/${runId}` +} + +/** True when the environment provides a WebSocket implementation. */ +export function isWebSocketSupported(): boolean { + return typeof WebSocket !== 'undefined' +} + +/** + * Watch a set of runs over WebSockets. + * + * Opens one socket per run id and calls `onTerminal` when a run reports + * `completed` or `failed`. `healthy` is false when the browser has no + * WebSocket implementation, or when any socket errored — the caller uses that + * to decide whether polling still needs to run. + */ +export function useRunWebSocket( + runIds: number[], + onTerminal: (event: RunStatusEvent) => void, +): { healthy: boolean } { + const [failed, setFailed] = useState(false) + const onTerminalRef = useRef(onTerminal) + + // Keep the latest callback without making it a effect dependency, so a new + // inline function on each render does not tear down and reopen the sockets. + useEffect(() => { + onTerminalRef.current = onTerminal + }, [onTerminal]) + + // Stable key so re-renders with an equal id list do not reopen sockets. + const key = runIds.join(',') + + useEffect(() => { + if (!isWebSocketSupported()) { + setFailed(true) + return + } + const ids = key === '' ? [] : key.split(',').map(Number) + if (ids.length === 0) return + + const sockets: WebSocket[] = [] + let cancelled = false + + for (const runId of ids) { + let socket: WebSocket + try { + socket = new WebSocket(buildRunWebSocketUrl(runId)) + } catch { + // Constructor throws on a malformed URL or a blocked scheme. + setFailed(true) + continue + } + + socket.onmessage = (message) => { + if (cancelled) return + let event: RunStatusEvent + try { + event = JSON.parse(String(message.data)) as RunStatusEvent + } catch { + // A frame we cannot parse is not fatal; polling still covers us. + return + } + if (event.status && TERMINAL_STATUSES.has(event.status)) { + onTerminalRef.current(event) + } + } + + socket.onerror = () => { + if (!cancelled) setFailed(true) + } + + sockets.push(socket) + } + + return () => { + cancelled = true + for (const socket of sockets) { + // Detach handlers first so a close-triggered error cannot flip + // `failed` after this effect has been torn down. + socket.onmessage = null + socket.onerror = null + if ( + socket.readyState === WebSocket.OPEN || + socket.readyState === WebSocket.CONNECTING + ) { + socket.close() + } + } + } + }, [key]) + + return { healthy: !failed } +} diff --git a/src/climatevision/api/main.py b/src/climatevision/api/main.py index 23d12e1..04fb92f 100644 --- a/src/climatevision/api/main.py +++ b/src/climatevision/api/main.py @@ -23,7 +23,19 @@ from pydantic import field_validator -from fastapi import FastAPI, File, Form, HTTPException, UploadFile, Header, Query, Depends, Request +from fastapi import ( + FastAPI, + File, + Form, + HTTPException, + UploadFile, + Header, + Query, + Depends, + Request, + WebSocket, + WebSocketDisconnect, +) from fastapi.responses import FileResponse, RedirectResponse from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles @@ -52,6 +64,7 @@ from climatevision.governance import explain_prediction, SHAPExplainer from climatevision.security.api_security import SecurityMiddleware from climatevision.workers.alert_delivery import AlertDeliveryWorker +from climatevision.api.run_events import build_status_event, run_event_hub logger = logging.getLogger(__name__) @@ -119,6 +132,71 @@ def _utc_now_iso() -> str: return datetime.now(timezone.utc).isoformat() +# ===== Run Status Streaming ===== + +# Statuses after which no further updates are expected, so the socket closes. +TERMINAL_RUN_STATUSES = frozenset({"completed", "failed"}) + +# Application-defined close code for "the requested run does not exist". +# The 4000-4999 range is reserved for application use by RFC 6455. +WS_RUN_NOT_FOUND = 4404 + + +def _read_run_status( + run_id: int, +) -> Optional[tuple[str, Optional[dict[str, Any]]]]: + """Read a run's current status and latest result payload. + + Args: + run_id: The run to read. + + Returns: + A ``(status, result_payload)`` pair, or ``None`` if no such run exists. + The payload is ``None`` when the run has not produced a result yet. + """ + with get_connection() as conn: + run = conn.execute( + "SELECT status FROM runs WHERE id = ?", (run_id,) + ).fetchone() + if run is None: + return None + result = conn.execute( + "SELECT payload_json FROM results WHERE run_id = ? ORDER BY id DESC LIMIT 1", + (run_id,), + ).fetchone() + + payload: Optional[dict[str, Any]] = None + if result is not None: + payload = json.loads(result["payload_json"]) + return run["status"], payload + + +def _status_event_for( + run_id: int, status: str, result_payload: Optional[dict[str, Any]] +) -> dict[str, Any]: + """Build a status event, routing the payload to ``result`` or ``error``. + + A failed run stores its message under ``error`` inside the result payload, + so surface that as the event's ``error`` field rather than as a result. + + Args: + run_id: The run the event describes. + status: The run's current status. + result_payload: The stored result payload, if any. + + Returns: + The event dictionary to send over the WebSocket. + """ + if status == "failed": + error = None + if result_payload is not None: + error = result_payload.get("error") + return build_status_event(run_id, status, error=error or "Inference failed") + if status == "completed": + return build_status_event(run_id, status, result=result_payload) + return build_status_event(run_id, status) + + # ===== Request/Response Models ===== class PredictRequest(BaseModel): @@ -778,6 +856,51 @@ def get_run(run_id: int) -> dict[str, Any]: }, } + @app.websocket("/ws/runs/{run_id}") + async def watch_run(websocket: WebSocket, run_id: int) -> None: + """Stream status updates for a single run. + + On connect the current status is read from the database and sent + immediately, so a client attaching to an already-finished run still + receives a terminal event. If the run is still in progress the socket + then streams transitions until it reaches ``completed``/``failed``, + and closes once a terminal status has been delivered. + """ + await websocket.accept() + + # Subscribe before the initial read, otherwise a transition happening + # between the read and the subscribe would be missed entirely. + async with run_event_hub.subscribe(run_id) as queue: + snapshot = _read_run_status(run_id) + if snapshot is None: + await websocket.send_json( + {"type": "error", "run_id": run_id, "error": "Run not found"} + ) + await websocket.close(code=WS_RUN_NOT_FOUND) + return + + status, result_payload = snapshot + await websocket.send_json( + _status_event_for(run_id, status, result_payload) + ) + if status in TERMINAL_RUN_STATUSES: + await websocket.close() + return + + try: + while True: + event = await queue.get() + await websocket.send_json(event) + if event.get("status") in TERMINAL_RUN_STATUSES: + break + except WebSocketDisconnect: + # Client went away mid-run; the subscription is released by + # the context manager, nothing else to clean up. + logger.debug("WebSocket client disconnected from run %s", run_id) + return + + await websocket.close() + @app.get("/api/reports/{run_id}") def get_report(run_id: int) -> dict[str, Any]: """Structured carbon impact report for a completed deforestation run. @@ -943,6 +1066,10 @@ async def predict_json( (run_id, json.dumps(result_payload), None, result_created_at), ) + await run_event_hub.publish( + run_id, _status_event_for(run_id, status, result_payload) + ) + return {"run_id": run_id, "result": result_payload} @app.post("/api/predict/upload") @@ -1027,6 +1154,10 @@ async def predict_upload( (run_id, json.dumps(result_payload), None, result_created_at), ) + await run_event_hub.publish( + run_id, _status_event_for(run_id, status, result_payload) + ) + return {"run_id": run_id, "result": result_payload} # ===== Explainability Endpoints ===== diff --git a/src/climatevision/api/run_events.py b/src/climatevision/api/run_events.py new file mode 100644 index 0000000..ba3db7d --- /dev/null +++ b/src/climatevision/api/run_events.py @@ -0,0 +1,119 @@ +"""In-process publish/subscribe hub for run status events. + +`/api/predict` and `/api/predict/upload` execute inference inline, so a run's +status transitions happen inside the same process that serves the WebSocket. +That makes a small in-memory hub sufficient: the predict handlers publish a +terminal event, and every socket currently watching that run receives it. + +The hub deliberately does not persist anything. The WebSocket handler reads the +authoritative current state from the database when a client connects, and only +uses the hub for changes that happen *while* the client is attached. A client +that connects after a run has already finished still gets a terminal event from +that initial database read. + +Note that this is per-process state. Under a multi-worker deployment a socket +served by worker A will not observe a transition published by worker B; such a +client falls back to the initial snapshot plus the frontend's polling path. A +cross-process broker (Redis pub/sub or Postgres LISTEN/NOTIFY) would be the +natural upgrade and is tracked separately rather than folded in here. +""" + +from __future__ import annotations + +import asyncio +import logging +from contextlib import asynccontextmanager +from typing import Any, AsyncIterator + +logger = logging.getLogger(__name__) + +# Bounded so a client that stops reading cannot grow the queue without limit. +_MAX_QUEUED_EVENTS = 32 + + +class RunEventHub: + """Fan-out of run status events to the sockets watching each run.""" + + def __init__(self) -> None: + self._subscribers: dict[int, set[asyncio.Queue[dict[str, Any]]]] = {} + self._lock = asyncio.Lock() + + async def publish(self, run_id: int, event: dict[str, Any]) -> None: + """Deliver an event to every subscriber of ``run_id``. + + Args: + run_id: The run the event belongs to. + event: The payload to deliver. Delivered as-is. + """ + async with self._lock: + queues = list(self._subscribers.get(run_id, ())) + + for queue in queues: + try: + queue.put_nowait(event) + except asyncio.QueueFull: + # A subscriber that cannot keep up is skipped rather than + # allowed to block inference from completing. + logger.warning( + "Dropping run event for run %s: subscriber queue is full", run_id + ) + + @asynccontextmanager + async def subscribe( + self, run_id: int + ) -> AsyncIterator[asyncio.Queue[dict[str, Any]]]: + """Subscribe to ``run_id`` for the duration of the context. + + Args: + run_id: The run to watch. + + Yields: + A queue that receives every event published for the run. + """ + queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=_MAX_QUEUED_EVENTS) + async with self._lock: + self._subscribers.setdefault(run_id, set()).add(queue) + try: + yield queue + finally: + async with self._lock: + subscribers = self._subscribers.get(run_id) + if subscribers is not None: + subscribers.discard(queue) + # Keep the map from growing once a run has no watchers. + if not subscribers: + del self._subscribers[run_id] + + def subscriber_count(self, run_id: int) -> int: + """Return how many sockets are currently watching ``run_id``.""" + return len(self._subscribers.get(run_id, ())) + + +# Module-level hub shared by the predict handlers and the WebSocket endpoint. +run_event_hub = RunEventHub() + + +def build_status_event( + run_id: int, + status: str, + *, + result: dict[str, Any] | None = None, + error: str | None = None, +) -> dict[str, Any]: + """Build the wire payload for a run status event. + + Args: + run_id: The run the event describes. + status: The run's status, e.g. ``running``/``completed``/``failed``. + result: The result payload, for a ``completed`` run. + error: The failure message, for a ``failed`` run. + + Returns: + The event dictionary sent over the WebSocket. + """ + event: dict[str, Any] = {"type": "status", "run_id": run_id, "status": status} + if result is not None: + event["result"] = result + if error is not None: + event["error"] = error + return event diff --git a/tests/test_ws_run_status.py b/tests/test_ws_run_status.py new file mode 100644 index 0000000..7bb48f2 --- /dev/null +++ b/tests/test_ws_run_status.py @@ -0,0 +1,327 @@ +"""Tests for the /ws/runs/{run_id} run status WebSocket.""" + +import asyncio +import json +from unittest.mock import patch + +import pytest +from fastapi.testclient import TestClient + +from climatevision.api.main import WS_RUN_NOT_FOUND +from climatevision.api.run_events import ( + RunEventHub, + build_status_event, + run_event_hub, +) +from climatevision.db import get_connection + + +def _insert_run(status: str) -> int: + """Insert a run with the given status and return its id.""" + with get_connection() as conn: + cur = conn.execute( + """ + INSERT INTO runs (kind, status, analysis_type, created_at, updated_at) + VALUES (?, ?, ?, ?, ?) + """, + ( + "test", + status, + "deforestation", + "2026-01-01T00:00:00Z", + "2026-01-01T00:00:00Z", + ), + ) + return int(cur.lastrowid) + + +def _insert_result(run_id: int, payload: dict) -> None: + """Attach a stored result payload to a run.""" + with get_connection() as conn: + conn.execute( + """ + INSERT INTO results (run_id, payload_json, mask_path, created_at) + VALUES (?, ?, ?, ?) + """, + (run_id, json.dumps(payload), None, "2026-01-01T00:00:00Z"), + ) + + +def test_unknown_run_is_reported_and_closed(client: TestClient) -> None: + """A socket for a nonexistent run gets an error event, then a close.""" + with client.websocket_connect("/ws/runs/99999999") as ws: + event = ws.receive_json() + assert event["type"] == "error" + assert event["run_id"] == 99999999 + assert "not found" in event["error"].lower() + + # Starlette surfaces the server-side close as a websocket.close message. + closed = ws.receive() + assert closed["type"] == "websocket.close" + assert closed["code"] == WS_RUN_NOT_FOUND + + +def test_already_completed_run_sends_terminal_snapshot(client: TestClient) -> None: + """Connecting to a finished run still yields a terminal event immediately. + + This is the late-attach case: the transition happened before the client + connected, so it can only come from the initial database read. + """ + run_id = _insert_run("completed") + _insert_result(run_id, {"analysis_type": "deforestation", "deforested_pixels": 42}) + + with client.websocket_connect(f"/ws/runs/{run_id}") as ws: + event = ws.receive_json() + + assert event["type"] == "status" + assert event["run_id"] == run_id + assert event["status"] == "completed" + assert event["result"]["deforested_pixels"] == 42 + + +def test_failed_run_surfaces_error_not_result(client: TestClient) -> None: + """A failed run reports its message under `error`, with no `result`.""" + run_id = _insert_run("failed") + _insert_result(run_id, {"analysis_type": "deforestation", "error": "GEE timeout"}) + + with client.websocket_connect(f"/ws/runs/{run_id}") as ws: + event = ws.receive_json() + + assert event["status"] == "failed" + assert event["error"] == "GEE timeout" + assert "result" not in event + + +def test_failed_run_without_stored_message_still_reports_an_error( + client: TestClient, +) -> None: + """A failed run with no `error` in its payload gets a fallback message.""" + run_id = _insert_run("failed") + _insert_result(run_id, {"analysis_type": "deforestation"}) + + with client.websocket_connect(f"/ws/runs/{run_id}") as ws: + event = ws.receive_json() + + assert event["status"] == "failed" + assert event["error"] == "Inference failed" + + +def test_running_run_streams_transition_then_closes(client: TestClient) -> None: + """A run still in progress streams its snapshot, then the terminal event.""" + run_id = _insert_run("running") + + with client.websocket_connect(f"/ws/runs/{run_id}") as ws: + snapshot = ws.receive_json() + assert snapshot["status"] == "running" + assert "result" not in snapshot + + # Publishing on the app's event loop mirrors what the predict handler + # does once inference finishes. + portal = ws.portal # type: ignore[attr-defined] + portal.call( + run_event_hub.publish, + run_id, + build_status_event(run_id, "completed", result={"deforested_pixels": 7}), + ) + + terminal = ws.receive_json() + assert terminal["status"] == "completed" + assert terminal["result"]["deforested_pixels"] == 7 + + closed = ws.receive() + assert closed["type"] == "websocket.close" + + +def test_running_run_ignores_events_for_other_runs(client: TestClient) -> None: + """Events published for a different run are not delivered to this socket.""" + watched_id = _insert_run("running") + other_id = _insert_run("running") + + with client.websocket_connect(f"/ws/runs/{watched_id}") as ws: + assert ws.receive_json()["status"] == "running" + + portal = ws.portal # type: ignore[attr-defined] + portal.call( + run_event_hub.publish, + other_id, + build_status_event(other_id, "completed", result={}), + ) + portal.call( + run_event_hub.publish, + watched_id, + build_status_event(watched_id, "failed", error="boom"), + ) + + event = ws.receive_json() + + # The first event received is the one for the watched run, proving the + # other run's event was never queued onto this subscription. + assert event["run_id"] == watched_id + assert event["status"] == "failed" + assert event["error"] == "boom" + + +def test_subscription_is_released_when_the_client_disconnects( + client: TestClient, +) -> None: + """Leaving the socket context drops the subscription, avoiding a leak.""" + run_id = _insert_run("running") + + with client.websocket_connect(f"/ws/runs/{run_id}") as ws: + ws.receive_json() + assert run_event_hub.subscriber_count(run_id) == 1 + + # Give the server task a moment to unwind the context manager. + for _ in range(50): + if run_event_hub.subscriber_count(run_id) == 0: + break + import time + + time.sleep(0.01) + + assert run_event_hub.subscriber_count(run_id) == 0 + + +# ===== Predict-to-hub wiring ===== + + +def test_predict_publishes_a_completed_event( + client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + """A successful /api/predict publishes the run's terminal event.""" + monkeypatch.setenv("CLIMATEVISION_ALLOW_DEV_KEY", "1") + published: list[tuple[int, dict]] = [] + + async def record(run_id: int, event: dict) -> None: + published.append((run_id, event)) + + monkeypatch.setattr(run_event_hub, "publish", record) + + fake_result = {"inference": {"forest_percentage": 72.3}} + with patch( + "climatevision.api.main.run_inference_from_gee", return_value=fake_result + ): + response = client.post( + "/api/predict", + json={ + "bbox": [-60.0, -15.0, -45.0, -5.0], + "start_date": "2023-01-01", + "end_date": "2023-06-30", + "analysis_type": "deforestation", + }, + headers={"X-API-Key": "cv_dev"}, + ) + + assert response.status_code == 200 + run_id = response.json()["run_id"] + assert len(published) == 1 + published_run_id, event = published[0] + assert published_run_id == run_id + assert event["status"] == "completed" + assert event["result"]["inference"]["forest_percentage"] == 72.3 + + +def test_predict_publishes_a_failed_event_when_inference_raises( + client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + """A failed /api/predict publishes `failed` with the error message.""" + monkeypatch.setenv("CLIMATEVISION_ALLOW_DEV_KEY", "1") + published: list[tuple[int, dict]] = [] + + async def record(run_id: int, event: dict) -> None: + published.append((run_id, event)) + + monkeypatch.setattr(run_event_hub, "publish", record) + + with patch( + "climatevision.api.main.run_inference_from_gee", + side_effect=RuntimeError("GEE unavailable"), + ): + response = client.post( + "/api/predict", + json={ + "bbox": [-60.0, -15.0, -45.0, -5.0], + "start_date": "2023-01-01", + "end_date": "2023-06-30", + "analysis_type": "deforestation", + }, + headers={"X-API-Key": "cv_dev"}, + ) + + assert response.status_code == 200 + assert len(published) == 1 + _, event = published[0] + assert event["status"] == "failed" + assert event["error"] == "GEE unavailable" + assert "result" not in event + + +# ===== RunEventHub unit tests ===== + + +def test_hub_delivers_to_every_subscriber_of_a_run() -> None: + """All sockets watching the same run receive the event.""" + + async def scenario() -> None: + hub = RunEventHub() + async with hub.subscribe(1) as first, hub.subscribe(1) as second: + await hub.publish(1, {"type": "status", "status": "completed"}) + assert (await first.get())["status"] == "completed" + assert (await second.get())["status"] == "completed" + + asyncio.run(scenario()) + + +def test_hub_publish_to_a_run_with_no_subscribers_is_a_no_op() -> None: + """Publishing to an unwatched run neither raises nor retains state.""" + + async def scenario() -> None: + hub = RunEventHub() + await hub.publish(404, {"type": "status", "status": "completed"}) + assert hub.subscriber_count(404) == 0 + + asyncio.run(scenario()) + + +def test_hub_drops_events_for_a_subscriber_that_stops_reading() -> None: + """A full queue is skipped so publishing never blocks inference.""" + + async def scenario() -> None: + hub = RunEventHub() + async with hub.subscribe(1) as queue: + # Fill well past the bound without ever draining the queue. + for i in range(100): + await hub.publish(1, {"type": "status", "seq": i}) + assert queue.full() + # Publishing still returns rather than blocking forever. + await asyncio.wait_for(hub.publish(1, {"type": "status", "seq": 100}), 1) + + asyncio.run(scenario()) + + +def test_hub_forgets_a_run_once_its_last_subscriber_leaves() -> None: + """The subscriber map does not grow once a run has no watchers.""" + + async def scenario() -> None: + hub = RunEventHub() + async with hub.subscribe(7): + assert hub.subscriber_count(7) == 1 + assert hub.subscriber_count(7) == 0 + + asyncio.run(scenario()) + + +@pytest.mark.parametrize( + "kwargs,expected_keys", + [ + ({}, {"type", "run_id", "status"}), + ({"result": {"a": 1}}, {"type", "run_id", "status", "result"}), + ({"error": "boom"}, {"type", "run_id", "status", "error"}), + ], +) +def test_build_status_event_only_includes_provided_fields( + kwargs: dict, expected_keys: set +) -> None: + """Optional fields stay out of the payload unless explicitly supplied.""" + event = build_status_event(1, "completed", **kwargs) + assert set(event) == expected_keys