From 3ba4373e7244ea4d2f51a2df15b4996fceedcebc Mon Sep 17 00:00:00 2001 From: zZebrahz Date: Sat, 6 Jun 2026 19:29:46 -0700 Subject: [PATCH] fix(web): stabilize realtime websocket connections (#58) * fix(web): prevent websocket reconnect loops * fix(api): keep event stream open on snapshot errors * fix(web): delay realtime disconnect indicator * fix(web): keep realtime warning delay stable * fix(web): show delayed realtime warning on mount --- internal/api/handlers/events_ws.go | 16 +- .../RealtimeEventsProvider.test.tsx | 144 +++++++++++++++++- web/src/components/RealtimeEventsProvider.tsx | 47 ++++-- web/src/components/ServerActivity.tsx | 40 ++++- 4 files changed, 225 insertions(+), 22 deletions(-) diff --git a/internal/api/handlers/events_ws.go b/internal/api/handlers/events_ws.go index d0f87b1b..a6b668ff 100644 --- a/internal/api/handlers/events_ws.go +++ b/internal/api/handlers/events_ws.go @@ -391,13 +391,17 @@ func (h *EventsHandler) writeSnapshotFrame( ) error { data, err := h.snapshotForChannel(r, claims, channel) if err != nil { - writeWebSocketError(conn, "internal_error", "Failed to load snapshot") - _ = writeWebSocketControl( - conn, - websocket.CloseMessage, - websocket.FormatCloseMessage(websocket.CloseInternalServerErr, "snapshot failed"), + slog.Error( + "events: failed to build initial snapshot", + "channel", + channel, + "user_id", + claims.UserID, + "error", + err, ) - return err + writeWebSocketError(conn, "internal_error", "Failed to load snapshot") + return nil } return writeWebSocketJSON(conn, evt.EventsSnapshotMessage{ Type: "snapshot", diff --git a/web/src/components/RealtimeEventsProvider.test.tsx b/web/src/components/RealtimeEventsProvider.test.tsx index 2357f87a..a645956b 100644 --- a/web/src/components/RealtimeEventsProvider.test.tsx +++ b/web/src/components/RealtimeEventsProvider.test.tsx @@ -1,5 +1,68 @@ -import { describe, expect, it } from "vitest"; -import { buildEventsUrl } from "./RealtimeEventsProvider"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { act, cleanup, render } from "@testing-library/react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { buildEventsUrl, RealtimeEventsProvider } from "./RealtimeEventsProvider"; + +const mockState = vi.hoisted(() => ({ + user: { + id: 1, + username: "admin", + email: "admin@example.com", + role: "admin", + permissions: [], + download_allowed: true, + }, + pageActivity: { + isVisible: true, + isFocused: true, + isFrozen: false, + canPollDashboard: true, + canApplyRealtimeUpdates: true, + }, +})); + +vi.mock("@/hooks/useAuth", () => ({ + useAuth: () => ({ + user: mockState.user, + profile: null, + }), +})); + +vi.mock("@/hooks/usePageActivity", () => ({ + usePageActivity: () => mockState.pageActivity, +})); + +vi.mock("react-router", () => ({ + useLocation: () => ({ pathname: "/" }), +})); + +class FakeWebSocket { + static CONNECTING = 0; + static OPEN = 1; + static CLOSED = 3; + static instances: FakeWebSocket[] = []; + + onopen: (() => void) | null = null; + onmessage: ((event: MessageEvent) => void) | null = null; + onerror: (() => void) | null = null; + onclose: (() => void) | null = null; + readyState = FakeWebSocket.CONNECTING; + + constructor(public url: string) { + FakeWebSocket.instances.push(this); + } + + send() {} + + close() { + this.readyState = FakeWebSocket.CLOSED; + } + + emitClose() { + this.readyState = FakeWebSocket.CLOSED; + this.onclose?.(); + } +} describe("buildEventsUrl", () => { it("includes auth token and websocket scheme", () => { @@ -20,3 +83,80 @@ describe("buildEventsUrl", () => { ).toBe("ws://localhost:5173/api/v1/events/ws"); }); }); + +describe("RealtimeEventsProvider", () => { + beforeEach(() => { + FakeWebSocket.instances = []; + vi.useFakeTimers(); + vi.stubGlobal("WebSocket", FakeWebSocket); + mockState.pageActivity = { + isVisible: true, + isFocused: true, + isFrozen: false, + canPollDashboard: true, + canApplyRealtimeUpdates: true, + }; + }); + + afterEach(() => { + cleanup(); + vi.useRealTimers(); + vi.unstubAllGlobals(); + }); + + it("ignores stale close events from intentionally closed sockets", () => { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: false }, + mutations: { retry: false }, + }, + }); + const view = render( + + +
+ + , + ); + + expect(FakeWebSocket.instances).toHaveLength(1); + const firstSocket = FakeWebSocket.instances[0]; + + act(() => { + mockState.pageActivity = { + ...mockState.pageActivity, + canApplyRealtimeUpdates: false, + }; + view.rerender( + + +
+ + , + ); + }); + + act(() => { + mockState.pageActivity = { + ...mockState.pageActivity, + canApplyRealtimeUpdates: true, + }; + view.rerender( + + +
+ + , + ); + }); + + expect(FakeWebSocket.instances).toHaveLength(2); + + act(() => { + firstSocket?.emitClose(); + vi.advanceTimersByTime(1_000); + }); + + expect(FakeWebSocket.instances).toHaveLength(2); + }); +}); diff --git a/web/src/components/RealtimeEventsProvider.tsx b/web/src/components/RealtimeEventsProvider.tsx index cca7f065..e2a58389 100644 --- a/web/src/components/RealtimeEventsProvider.tsx +++ b/web/src/components/RealtimeEventsProvider.tsx @@ -354,11 +354,11 @@ export function RealtimeEventsProvider({ children }: { children: ReactNode }) { const { user, profile } = useAuth(); const pageActivity = usePageActivity(); const location = useLocation(); + const authenticatedUserID = user?.id ?? null; const isDashboardRoute = location.pathname === "/admin" || location.pathname === "/admin/"; const allowDashboardRealtimeUpdates = !isDashboardRoute || pageActivity.canPollDashboard; const [connectionState, setConnectionState] = useState("connecting"); const reconnectTimerRef = useRef(undefined); - const closingRef = useRef(false); const socketRef = useRef(null); const helloReceivedRef = useRef(false); const requestCounterRef = useRef(0); @@ -527,7 +527,7 @@ export function RealtimeEventsProvider({ children }: { children: ReactNode }) { } useEffect(() => { - if (!user) { + if (!authenticatedUserID) { return; } if (!pageActivity.canApplyRealtimeUpdates) { @@ -543,15 +543,16 @@ export function RealtimeEventsProvider({ children }: { children: ReactNode }) { type: "active", predicate: (query) => !isDashboardQueryKey(query.queryKey), }); - }, [pageActivity.canApplyRealtimeUpdates, queryClient, user]); + }, [authenticatedUserID, pageActivity.canApplyRealtimeUpdates, queryClient]); useEffect(() => { - if (!user || !pageActivity.canApplyRealtimeUpdates) { + if (!authenticatedUserID || !pageActivity.canApplyRealtimeUpdates) { setConnectionState("disconnected"); return; } - closingRef.current = false; + let closedByEffect = false; + let activeSocket: WebSocket | null = null; const clearReconnect = () => { if (reconnectTimerRef.current !== undefined) { @@ -561,16 +562,22 @@ export function RealtimeEventsProvider({ children }: { children: ReactNode }) { }; const scheduleReconnect = () => { - if (closingRef.current || reconnectTimerRef.current !== undefined) { + if (closedByEffect || reconnectTimerRef.current !== undefined) { return; } reconnectTimerRef.current = window.setTimeout(() => { reconnectTimerRef.current = undefined; + if (closedByEffect) { + return; + } connect(); }, 1_000); }; const connect = () => { + if (closedByEffect) { + return; + } setConnectionState("connecting"); helloReceivedRef.current = false; @@ -583,13 +590,20 @@ export function RealtimeEventsProvider({ children }: { children: ReactNode }) { return; } + activeSocket = socket; socketRef.current = socket; socket.onopen = () => { + if (closedByEffect || socketRef.current !== socket) { + return; + } setConnectionState("live"); }; socket.onmessage = (event) => { + if (closedByEffect || socketRef.current !== socket) { + return; + } if (!canApplyRealtimeUpdatesRef.current) { return; } @@ -617,13 +631,21 @@ export function RealtimeEventsProvider({ children }: { children: ReactNode }) { }; socket.onerror = () => { + if (closedByEffect || socketRef.current !== socket) { + return; + } setConnectionState("disconnected"); }; socket.onclose = () => { + if (socketRef.current !== socket) { + return; + } + socketRef.current = null; + activeSocket = null; helloReceivedRef.current = false; setConnectionState("disconnected"); - if (!closingRef.current) { + if (!closedByEffect) { scheduleReconnect(); } }; @@ -632,15 +654,18 @@ export function RealtimeEventsProvider({ children }: { children: ReactNode }) { connect(); return () => { - closingRef.current = true; + closedByEffect = true; clearReconnect(); for (const [jobId, waiter] of waitersRef.current) { window.clearTimeout(waiter.timeoutId); waiter.reject(new Error(`Realtime events provider closed before job ${jobId} finished`)); } waitersRef.current.clear(); - const socket = socketRef.current; - socketRef.current = null; + const socket = activeSocket; + if (socket && socketRef.current === socket) { + socketRef.current = null; + } + activeSocket = null; if ( socket && (socket.readyState === WebSocket.OPEN || socket.readyState === WebSocket.CONNECTING) @@ -648,7 +673,7 @@ export function RealtimeEventsProvider({ children }: { children: ReactNode }) { socket.close(); } }; - }, [pageActivity.canApplyRealtimeUpdates, queryClient, user]); + }, [authenticatedUserID, pageActivity.canApplyRealtimeUpdates, queryClient]); const value = useMemo( () => ({ diff --git a/web/src/components/ServerActivity.tsx b/web/src/components/ServerActivity.tsx index da2d0d0c..9c926e8e 100644 --- a/web/src/components/ServerActivity.tsx +++ b/web/src/components/ServerActivity.tsx @@ -1,4 +1,4 @@ -import { useState, useMemo } from "react"; +import { useState, useMemo, useEffect, useRef } from "react"; import { Link, useLocation } from "react-router"; import { Popover as PopoverPrimitive } from "radix-ui"; import { Activity, ChevronRight, Loader, ScanLine } from "lucide-react"; @@ -6,9 +6,14 @@ import { useAdminSessions } from "@/hooks/queries/admin/stats"; import { useTasks } from "@/hooks/queries/admin/tasks"; import { useActiveScans } from "@/hooks/queries/admin/scans"; import { useAdminLibraries } from "@/hooks/queries/admin/libraries"; -import { useRealtimeEvents } from "@/components/realtimeEventsContext"; +import { + useRealtimeEvents, + type RealtimeConnectionState, +} from "@/components/realtimeEventsContext"; import type { TaskInfo, ScanRun } from "@/api/types"; +const CONNECTION_PROBLEM_INDICATOR_DELAY_MS = 4_000; + interface ServerActivityProps { /** Hide the trigger button entirely when there is no activity */ hideWhenEmpty?: boolean; @@ -58,6 +63,34 @@ function useServerActivityData() { // ── Main component ─────────────────────────────────────────── +function useDelayedConnectionProblem(connectionState: RealtimeConnectionState) { + const isNonLive = connectionState !== "live"; + const previousIsNonLiveRef = useRef(false); + const [connectionProblemState, setConnectionProblemState] = useState(false); + + useEffect(() => { + const wasNonLive = previousIsNonLiveRef.current; + previousIsNonLiveRef.current = isNonLive; + + if (!isNonLive) { + setConnectionProblemState(false); + return; + } + + if (wasNonLive) { + return; + } + + const timeoutID = window.setTimeout( + () => setConnectionProblemState(true), + CONNECTION_PROBLEM_INDICATOR_DELAY_MS, + ); + return () => window.clearTimeout(timeoutID); + }, [isNonLive]); + + return isNonLive && connectionProblemState; +} + export default function ServerActivity({ hideWhenEmpty = false }: ServerActivityProps) { const [open, setOpen] = useState(false); const location = useLocation(); @@ -72,6 +105,7 @@ export default function ServerActivity({ hideWhenEmpty = false }: ServerActivity connectionState, scansLoaded, } = useServerActivityData(); + const showConnectionProblem = useDelayedConnectionProblem(connectionState); // Keep mounted while popover is open so Radix can animate closed if (hideWhenEmpty && totalActive === 0 && !open) return null; @@ -96,7 +130,7 @@ export default function ServerActivity({ hideWhenEmpty = false }: ServerActivity {totalActive} )} - {connectionState !== "live" && ( + {showConnectionProblem && (