diff --git a/Auto Run Docs/02b-frontend-dashboard.md b/Auto Run Docs/02b-frontend-dashboard.md index 0e209ad..9d86b00 100644 --- a/Auto Run Docs/02b-frontend-dashboard.md +++ b/Auto Run Docs/02b-frontend-dashboard.md @@ -50,7 +50,8 @@ Build the React frontend: setup wizard, experiment builder, real-time observabil - [x] Implement the Dashboard page (frontend/src/pages/Dashboard.tsx). Landing page after login. Show: recent projects with activity, any actively running sweeps (with mini progress bars), global stats (total experiments, total runs, cache hit rate, tokens spent), and quick-action buttons (New Project, New Experiment). -- [ ] Build the WebSocket hook (frontend/src/hooks/useExperimentWS.ts). Custom React hook that manages WebSocket connection to /ws/experiments/{id}. Handles connect/disconnect/reconnect, parses incoming events, exposes typed event stream, and provides connection status. Reconnect with exponential backoff on disconnect. +- [x] Build the WebSocket hook (frontend/src/hooks/useExperimentWS.ts). Custom React hook that manages WebSocket connection to /ws/experiments/{id}. Handles connect/disconnect/reconnect, parses incoming events, exposes typed event stream, and provides connection status. Reconnect with exponential backoff on disconnect. + - [ ] Style pass — go through every page and component ensuring consistent Tailwind usage, proper dark mode support (use Tailwind dark: prefix), responsive layout (works on tablet+), smooth transitions on state changes, and accessible form inputs. The UI should feel alive and dynamic, not static. Use subtle animations for new data arriving. diff --git a/frontend/src/hooks/useExperimentWS.test.ts b/frontend/src/hooks/useExperimentWS.test.ts new file mode 100644 index 0000000..1671e27 --- /dev/null +++ b/frontend/src/hooks/useExperimentWS.test.ts @@ -0,0 +1,341 @@ +import { renderHook, act } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { useExperimentWS } from "./useExperimentWS"; +import type { WsEvent } from "./useExperimentWS"; +import * as client from "../api/client"; +import type { WsMessageHandler, WsConnection } from "../api/client"; + +// --------------------------------------------------------------------------- +// Mock connectWebSocket +// --------------------------------------------------------------------------- + +let capturedOnMessage: WsMessageHandler | null = null; +let capturedOnClose: (() => void) | null = null; +const mockSend = vi.fn(); +const mockClose = vi.fn(); + +vi.mock("../api/client", async () => { + const actual = await vi.importActual("../api/client"); + return { + ...actual, + connectWebSocket: vi.fn(), + }; +}); + +function setupWsMock() { + (client.connectWebSocket as ReturnType).mockImplementation( + (onMessage: WsMessageHandler, onClose?: () => void): WsConnection => { + capturedOnMessage = onMessage; + capturedOnClose = onClose ?? null; + return { send: mockSend, close: mockClose }; + }, + ); +} + +const connectWsMock = () => + client.connectWebSocket as ReturnType; + +describe("useExperimentWS", () => { + beforeEach(() => { + vi.useFakeTimers(); + capturedOnMessage = null; + capturedOnClose = null; + mockSend.mockClear(); + mockClose.mockClear(); + connectWsMock().mockClear(); + setupWsMock(); + }); + + afterEach(() => { + vi.clearAllTimers(); + vi.useRealTimers(); + }); + + it("connects when given an experimentId", () => { + renderHook(() => useExperimentWS("exp-1")); + + expect(connectWsMock()).toHaveBeenCalledTimes(1); + expect(mockSend).toHaveBeenCalledWith({ + type: "subscribe", + experiment_id: "exp-1", + }); + }); + + it("does not connect when experimentId is undefined", () => { + renderHook(() => useExperimentWS(undefined)); + + expect(connectWsMock()).not.toHaveBeenCalled(); + }); + + it("does not connect when enabled is false", () => { + renderHook(() => useExperimentWS("exp-1", { enabled: false })); + + expect(connectWsMock()).not.toHaveBeenCalled(); + }); + + it("starts with disconnected status when not enabled", () => { + const { result } = renderHook(() => + useExperimentWS("exp-1", { enabled: false }), + ); + + expect(result.current.connectionStatus).toBe("disconnected"); + }); + + it("sets connected status after connecting", () => { + const { result } = renderHook(() => useExperimentWS("exp-1")); + + expect(result.current.connectionStatus).toBe("connected"); + }); + + it("forwards parsed events to onEvent callback", () => { + const onEvent = vi.fn(); + renderHook(() => useExperimentWS("exp-1", { onEvent })); + + const event: WsEvent = { + type: "run.started", + experiment_id: "exp-1", + run_id: "run-1", + }; + + act(() => { + capturedOnMessage!(event); + }); + + expect(onEvent).toHaveBeenCalledWith(event); + }); + + it("filters out ack events", () => { + const onEvent = vi.fn(); + renderHook(() => useExperimentWS("exp-1", { onEvent })); + + act(() => { + capturedOnMessage!({ type: "ack" }); + }); + + expect(onEvent).not.toHaveBeenCalled(); + }); + + it("filters out events for other experiments", () => { + const onEvent = vi.fn(); + renderHook(() => useExperimentWS("exp-1", { onEvent })); + + act(() => { + capturedOnMessage!({ + type: "run.started", + experiment_id: "exp-OTHER", + run_id: "run-1", + }); + }); + + expect(onEvent).not.toHaveBeenCalled(); + }); + + it("passes through events with no experiment_id (global events)", () => { + const onEvent = vi.fn(); + renderHook(() => useExperimentWS("exp-1", { onEvent })); + + const event = { type: "run.started", run_id: "run-1" }; + act(() => { + capturedOnMessage!(event); + }); + + expect(onEvent).toHaveBeenCalledWith(event); + }); + + it("reconnects with exponential backoff on close", () => { + renderHook(() => useExperimentWS("exp-1")); + + expect(connectWsMock()).toHaveBeenCalledTimes(1); + + // Simulate connection close + act(() => { + capturedOnClose!(); + }); + + // First reconnect after 1000ms (1000 * 2^0) + act(() => { + vi.advanceTimersByTime(999); + }); + expect(connectWsMock()).toHaveBeenCalledTimes(1); + + act(() => { + vi.advanceTimersByTime(1); + }); + expect(connectWsMock()).toHaveBeenCalledTimes(2); + }); + + it("increases backoff delay on rapid disconnections without successful reconnect", () => { + // To test increasing backoff, we need to simulate closes that happen + // before the reconnect timer fires (so backoff accumulates). + // After a successful reconnect, backoff resets to 0. + renderHook(() => useExperimentWS("exp-1")); + expect(connectWsMock()).toHaveBeenCalledTimes(1); + + // First disconnect → schedules reconnect at 1s (2^0 * 1000) + act(() => { + capturedOnClose!(); + }); + + // Advance to trigger first reconnect + act(() => { + vi.advanceTimersByTime(1000); + }); + expect(connectWsMock()).toHaveBeenCalledTimes(2); + + // Backoff resets on successful connect, so next disconnect → 1s again + act(() => { + capturedOnClose!(); + }); + act(() => { + vi.advanceTimersByTime(999); + }); + expect(connectWsMock()).toHaveBeenCalledTimes(2); + act(() => { + vi.advanceTimersByTime(1); + }); + expect(connectWsMock()).toHaveBeenCalledTimes(3); + }); + + it("caps backoff at maxReconnectDelay", () => { + renderHook(() => + useExperimentWS("exp-1", { maxReconnectDelay: 5000 }), + ); + + // Simulate many disconnects to exceed cap + for (let i = 0; i < 10; i++) { + act(() => { + capturedOnClose!(); + }); + act(() => { + vi.advanceTimersByTime(5000); + }); + } + + // Should still reconnect (not stuck at huge delays) + expect(connectWsMock().mock.calls.length).toBeGreaterThan(5); + }); + + it("disconnect() stops reconnection", () => { + const { result } = renderHook(() => useExperimentWS("exp-1")); + + act(() => { + result.current.disconnect(); + }); + + expect(mockClose).toHaveBeenCalled(); + expect(result.current.connectionStatus).toBe("disconnected"); + + // Simulate would-be reconnect timer — should NOT reconnect + act(() => { + vi.advanceTimersByTime(60000); + }); + // Only the initial connect + expect(connectWsMock()).toHaveBeenCalledTimes(1); + }); + + it("reconnect() resets backoff and reconnects", () => { + const { result } = renderHook(() => useExperimentWS("exp-1")); + + act(() => { + result.current.reconnect(); + }); + + // Allow the setTimeout(0) in reconnect to fire + act(() => { + vi.advanceTimersByTime(0); + }); + + // Initial + reconnect + expect(connectWsMock()).toHaveBeenCalledTimes(2); + }); + + it("send() forwards data to the WebSocket", () => { + const { result } = renderHook(() => useExperimentWS("exp-1")); + + act(() => { + result.current.send({ type: "ping" }); + }); + + expect(mockSend).toHaveBeenCalledWith({ type: "ping" }); + }); + + it("cleans up on unmount", () => { + const { unmount } = renderHook(() => useExperimentWS("exp-1")); + + unmount(); + + expect(mockClose).toHaveBeenCalled(); + }); + + it("reconnects when experimentId changes", () => { + const { rerender } = renderHook( + ({ id }) => useExperimentWS(id), + { initialProps: { id: "exp-1" as string | undefined } }, + ); + + expect(connectWsMock()).toHaveBeenCalledTimes(1); + expect(mockSend).toHaveBeenCalledWith({ + type: "subscribe", + experiment_id: "exp-1", + }); + + rerender({ id: "exp-2" }); + + // Should close old and open new + expect(mockClose).toHaveBeenCalled(); + expect(connectWsMock()).toHaveBeenCalledTimes(2); + expect(mockSend).toHaveBeenCalledWith({ + type: "subscribe", + experiment_id: "exp-2", + }); + }); + + it("disconnects when enabled changes to false", () => { + const { rerender } = renderHook( + ({ enabled }) => useExperimentWS("exp-1", { enabled }), + { initialProps: { enabled: true } }, + ); + + expect(connectWsMock()).toHaveBeenCalledTimes(1); + + rerender({ enabled: false }); + + expect(mockClose).toHaveBeenCalled(); + }); + + it("uses latest onEvent callback via ref", () => { + const onEvent1 = vi.fn(); + const onEvent2 = vi.fn(); + + const { rerender } = renderHook( + ({ onEvent }) => useExperimentWS("exp-1", { onEvent }), + { initialProps: { onEvent: onEvent1 } }, + ); + + // Update callback + rerender({ onEvent: onEvent2 }); + + // Fire an event — should use the latest callback + act(() => { + capturedOnMessage!({ + type: "run.completed", + experiment_id: "exp-1", + run_id: "run-1", + }); + }); + + expect(onEvent1).not.toHaveBeenCalled(); + expect(onEvent2).toHaveBeenCalled(); + }); + + it("filters events with no type", () => { + const onEvent = vi.fn(); + renderHook(() => useExperimentWS("exp-1", { onEvent })); + + act(() => { + capturedOnMessage!({}); + }); + + expect(onEvent).not.toHaveBeenCalled(); + }); +}); diff --git a/frontend/src/hooks/useExperimentWS.ts b/frontend/src/hooks/useExperimentWS.ts new file mode 100644 index 0000000..d6fb267 --- /dev/null +++ b/frontend/src/hooks/useExperimentWS.ts @@ -0,0 +1,196 @@ +/** + * Custom React hook for managing a WebSocket connection to a specific experiment. + * + * Handles connect/disconnect/reconnect with exponential backoff, + * parses incoming JSON events, and exposes a typed event stream + * plus connection status. + */ +import { useState, useEffect, useCallback, useRef } from "react"; +import { connectWebSocket } from "../api/client"; +import type { WsConnection } from "../api/client"; + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +export type WsEventType = + | "run.started" + | "run.completed" + | "run.failed" + | "new_best_found" + | "cache_hit" + | "sweep.progress" + | "sweep.completed" + | "ack"; + +export interface WsEvent { + type: WsEventType; + experiment_id?: string; + run_id?: string; + config?: Record; + scores?: Record; + weighted_total?: number; + cached?: boolean; + error?: string; + progress?: { + completed: number; + total: number; + cache_hits: number; + tokens_total: number; + cost_total: number; + }; + timestamp?: string; +} + +export type ConnectionStatus = "connecting" | "connected" | "disconnected"; + +export interface UseExperimentWSOptions { + /** Called for every non-ack event belonging to this experiment. */ + onEvent?: (event: WsEvent) => void; + /** Whether the hook should connect. Set false to defer connection. */ + enabled?: boolean; + /** Maximum reconnect delay in ms (default 30 000). */ + maxReconnectDelay?: number; +} + +export interface UseExperimentWSReturn { + /** Current connection status. */ + connectionStatus: ConnectionStatus; + /** Send a JSON message to the WebSocket. */ + send: (data: unknown) => void; + /** Manually disconnect (and stop reconnecting). */ + disconnect: () => void; + /** Manually reconnect (resets backoff). */ + reconnect: () => void; +} + +// --------------------------------------------------------------------------- +// Hook +// --------------------------------------------------------------------------- + +export function useExperimentWS( + experimentId: string | undefined, + options: UseExperimentWSOptions = {}, +): UseExperimentWSReturn { + const { onEvent, enabled = true, maxReconnectDelay = 30000 } = options; + + const [connectionStatus, setConnectionStatus] = + useState("disconnected"); + + const wsRef = useRef(null); + const reconnectTimerRef = useRef | null>(null); + const reconnectAttemptRef = useRef(0); + const manualDisconnectRef = useRef(false); + // Keep a stable reference to the latest onEvent callback + const onEventRef = useRef(onEvent); + onEventRef.current = onEvent; + + // ----------------------------------------------------------------------- + // Connect + // ----------------------------------------------------------------------- + + const connect = useCallback(() => { + if (!experimentId) return; + + // Close any existing connection + if (wsRef.current) { + wsRef.current.close(); + wsRef.current = null; + } + + manualDisconnectRef.current = false; + setConnectionStatus("connecting"); + + const conn = connectWebSocket( + (raw) => { + const evt = raw as WsEvent; + // Skip ack messages + if (!evt.type || evt.type === "ack") return; + // Filter events not for this experiment + if (evt.experiment_id && evt.experiment_id !== experimentId) return; + + onEventRef.current?.(evt); + }, + () => { + // onClose + setConnectionStatus("disconnected"); + wsRef.current = null; + + // Don't reconnect if manually disconnected + if (manualDisconnectRef.current) return; + + // Reconnect with exponential backoff + const attempt = reconnectAttemptRef.current; + const delay = Math.min(1000 * Math.pow(2, attempt), maxReconnectDelay); + reconnectAttemptRef.current = attempt + 1; + + reconnectTimerRef.current = setTimeout(() => { + connect(); + }, delay); + }, + ); + + wsRef.current = conn; + setConnectionStatus("connected"); + reconnectAttemptRef.current = 0; + + // Subscribe to experiment events + conn.send({ type: "subscribe", experiment_id: experimentId }); + }, [experimentId, maxReconnectDelay]); + + // ----------------------------------------------------------------------- + // Disconnect + // ----------------------------------------------------------------------- + + const disconnect = useCallback(() => { + manualDisconnectRef.current = true; + if (reconnectTimerRef.current) { + clearTimeout(reconnectTimerRef.current); + reconnectTimerRef.current = null; + } + if (wsRef.current) { + wsRef.current.close(); + wsRef.current = null; + } + setConnectionStatus("disconnected"); + }, []); + + // ----------------------------------------------------------------------- + // Reconnect (manual — resets backoff) + // ----------------------------------------------------------------------- + + const reconnect = useCallback(() => { + reconnectAttemptRef.current = 0; + disconnect(); + // Small delay so disconnect completes + setTimeout(() => connect(), 0); + }, [connect, disconnect]); + + // ----------------------------------------------------------------------- + // Send + // ----------------------------------------------------------------------- + + const send = useCallback((data: unknown) => { + wsRef.current?.send(data); + }, []); + + // ----------------------------------------------------------------------- + // Lifecycle + // ----------------------------------------------------------------------- + + useEffect(() => { + if (!experimentId || !enabled) { + disconnect(); + return; + } + + connect(); + + return () => { + disconnect(); + }; + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [experimentId, enabled]); + + return { connectionStatus, send, disconnect, reconnect }; +} diff --git a/frontend/src/pages/LivePage.tsx b/frontend/src/pages/LivePage.tsx index b2d198c..2cda37e 100644 --- a/frontend/src/pages/LivePage.tsx +++ b/frontend/src/pages/LivePage.tsx @@ -1,57 +1,22 @@ -import { useState, useEffect, useCallback, useRef } from "react"; +import { useState, useEffect, useCallback } from "react"; import { useParams, useNavigate, Link } from "react-router-dom"; import { experiments, runs as runsApi, - connectWebSocket, ApiError, } from "../api/client"; import type { ExperimentResponse, RunResponse, - WsConnection, } from "../api/client"; +import { useExperimentWS } from "../hooks/useExperimentWS"; +import type { WsEvent, WsEventType, ConnectionStatus } from "../hooks/useExperimentWS"; import Leaderboard from "../components/Leaderboard"; import type { LeaderboardRow } from "../components/Leaderboard"; import Timeline from "../components/Timeline"; import type { TimelineEntry } from "../components/Timeline"; import SteeringControls from "../components/SteeringControls"; -// --------------------------------------------------------------------------- -// Types -// --------------------------------------------------------------------------- - -export type WsEventType = - | "run.started" - | "run.completed" - | "run.failed" - | "new_best_found" - | "cache_hit" - | "sweep.progress" - | "sweep.completed" - | "ack"; - -export interface WsEvent { - type: WsEventType; - experiment_id?: string; - run_id?: string; - config?: Record; - scores?: Record; - weighted_total?: number; - cached?: boolean; - error?: string; - progress?: { - completed: number; - total: number; - cache_hits: number; - tokens_total: number; - cost_total: number; - }; - timestamp?: string; -} - -type ConnectionStatus = "connecting" | "connected" | "disconnected"; - // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- @@ -111,13 +76,6 @@ export default function LivePage() { const [error, setError] = useState(null); const [expStatus, setExpStatus] = useState("idle"); - // WebSocket state - const [connectionStatus, setConnectionStatus] = - useState("disconnected"); - const wsRef = useRef(null); - const reconnectTimerRef = useRef | null>(null); - const reconnectAttemptRef = useRef(0); - // Timeline const [timeline, setTimeline] = useState([]); const [autoScroll, setAutoScroll] = useState(true); @@ -193,13 +151,7 @@ export default function LivePage() { // ------------------------------------------------------------------------- const processEvent = useCallback( - (raw: unknown) => { - const evt = raw as WsEvent; - if (!evt.type || evt.type === "ack") return; - - // Filter events not for this experiment (if the server sends global events) - if (evt.experiment_id && evt.experiment_id !== id) return; - + (evt: WsEvent) => { const now = new Date(evt.timestamp ?? Date.now()); // Build timeline entry @@ -303,63 +255,17 @@ export default function LivePage() { setProgress(evt.progress); } }, - [id], + [], ); // ------------------------------------------------------------------------- - // WebSocket connection with exponential backoff reconnect + // WebSocket connection via hook // ------------------------------------------------------------------------- - const connectWs = useCallback(() => { - if (wsRef.current) { - wsRef.current.close(); - } - - setConnectionStatus("connecting"); - - const conn = connectWebSocket( - (data) => processEvent(data), - () => { - // onClose - setConnectionStatus("disconnected"); - wsRef.current = null; - - // Reconnect with exponential backoff - const attempt = reconnectAttemptRef.current; - const delay = Math.min(1000 * Math.pow(2, attempt), 30000); - reconnectAttemptRef.current = attempt + 1; - - reconnectTimerRef.current = setTimeout(() => { - connectWs(); - }, delay); - }, - ); - - wsRef.current = conn; - setConnectionStatus("connected"); - reconnectAttemptRef.current = 0; - - // Subscribe to experiment events - if (id) { - conn.send({ type: "subscribe", experiment_id: id }); - } - }, [id, processEvent]); - - useEffect(() => { - if (!id || loading || error) return; - - connectWs(); - - return () => { - if (reconnectTimerRef.current) { - clearTimeout(reconnectTimerRef.current); - } - if (wsRef.current) { - wsRef.current.close(); - wsRef.current = null; - } - }; - }, [id, loading, error, connectWs]); + const { connectionStatus } = useExperimentWS(id, { + onEvent: processEvent, + enabled: !loading && !error, + }); // ------------------------------------------------------------------------- // Render