Add: WebSocket endpoint for real-time run status (#11) (#140)

Adds /ws/runs/{run_id} so the dashboard receives status pushes instead of
waiting up to 5s for the next poll.

Because /api/predict runs inference inline, status transitions happen in the
process serving the socket, so a small in-memory pub/sub hub is enough: the
predict handlers publish a terminal event and every socket watching that run
receives it. On connect the socket reads the current status from the database
first, so a client attaching to an already-finished run still gets a terminal
event rather than hanging.

Frontend: useRunPolling now prefers a WebSocket per in-flight run and only
starts the 5s interval when the socket path is unavailable, so behaviour
degrades to exactly what it was before.

Closes #11

Co-authored-by: Ayomipo Solaja <ayomiposolaja@users.noreply.github.com>

Author: Ayomipo18

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)
+  })
+})

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<Map<number, string>>(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(() => {

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)
+  })
+})