diff --git a/ui/index.html b/ui/index.html index 1bb9152e..70a8550e 100644 --- a/ui/index.html +++ b/ui/index.html @@ -4,6 +4,7 @@ + diff --git a/ui/src/context/LiveUpdatesProvider.test.ts b/ui/src/context/LiveUpdatesProvider.test.ts index 3a66cec4..0bbbaabc 100644 --- a/ui/src/context/LiveUpdatesProvider.test.ts +++ b/ui/src/context/LiveUpdatesProvider.test.ts @@ -1,6 +1,6 @@ // @vitest-environment node -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { __liveUpdatesTestUtils } from "./LiveUpdatesProvider"; import { queryKeys } from "../lib/queryKeys"; @@ -117,3 +117,55 @@ describe("LiveUpdatesProvider visible issue toast suppression", () => { ).toBe(true); }); }); + +describe("LiveUpdatesProvider socket helpers", () => { + it("waits for the selected company object to catch up before connecting", () => { + expect(__liveUpdatesTestUtils.resolveLiveCompanyId("company-1", null)).toBeNull(); + expect(__liveUpdatesTestUtils.resolveLiveCompanyId("company-1", "company-2")).toBeNull(); + expect(__liveUpdatesTestUtils.resolveLiveCompanyId("company-1", "company-1")).toBe("company-1"); + }); + + it("defers close until onopen for sockets that are still connecting", () => { + const socket = { + readyState: 0, + onopen: (() => undefined) as (() => void) | null, + onmessage: (() => undefined) as (() => void) | null, + onerror: (() => undefined) as (() => void) | null, + onclose: (() => undefined) as (() => void) | null, + close: vi.fn(), + }; + + __liveUpdatesTestUtils.closeSocketQuietly(socket as never, "provider_unmount"); + + expect(socket.close).not.toHaveBeenCalled(); + expect(socket.onmessage).toBeNull(); + expect(socket.onclose).toBeNull(); + expect(socket.onopen).toBeTypeOf("function"); + expect(socket.onerror).toBeTypeOf("function"); + + socket.onopen?.(); + + expect(socket.close).toHaveBeenCalledWith(1000, "provider_unmount"); + expect(socket.onopen).toBeNull(); + expect(socket.onerror).toBeNull(); + }); + + it("closes open sockets immediately without leaving handlers behind", () => { + const socket = { + readyState: 1, + onopen: (() => undefined) as (() => void) | null, + onmessage: (() => undefined) as (() => void) | null, + onerror: (() => undefined) as (() => void) | null, + onclose: (() => undefined) as (() => void) | null, + close: vi.fn(), + }; + + __liveUpdatesTestUtils.closeSocketQuietly(socket as never, "stale_connection"); + + expect(socket.close).toHaveBeenCalledWith(1000, "stale_connection"); + expect(socket.onopen).toBeNull(); + expect(socket.onmessage).toBeNull(); + expect(socket.onerror).toBeNull(); + expect(socket.onclose).toBeNull(); + }); +}); diff --git a/ui/src/context/LiveUpdatesProvider.tsx b/ui/src/context/LiveUpdatesProvider.tsx index ec593b3b..b4f83388 100644 --- a/ui/src/context/LiveUpdatesProvider.tsx +++ b/ui/src/context/LiveUpdatesProvider.tsx @@ -14,6 +14,17 @@ import { useLocation } from "../lib/router"; const TOAST_COOLDOWN_WINDOW_MS = 10_000; const TOAST_COOLDOWN_MAX = 3; const RECONNECT_SUPPRESS_MS = 2000; +const SOCKET_CONNECTING = 0; +const SOCKET_OPEN = 1; + +type LiveUpdatesSocketLike = { + readyState: number; + onopen: ((this: WebSocket, ev: Event) => unknown) | null; + onmessage: ((this: WebSocket, ev: MessageEvent) => unknown) | null; + onerror: ((this: WebSocket, ev: Event) => unknown) | null; + onclose: ((this: WebSocket, ev: CloseEvent) => unknown) | null; + close: (code?: number, reason?: string) => void; +}; function readString(value: unknown): string | null { return typeof value === "string" && value.length > 0 ? value : null; @@ -652,33 +663,88 @@ function handleLiveEvent( } } +function resolveLiveCompanyId( + selectedCompanyId: string | null, + selectedCompanyLiveId: string | null, +): string | null { + return selectedCompanyId && selectedCompanyId === selectedCompanyLiveId + ? selectedCompanyId + : null; +} + +function resetSocketHandlers(target: LiveUpdatesSocketLike) { + target.onopen = null; + target.onmessage = null; + target.onerror = null; + target.onclose = null; +} + +function closeSocketQuietly(target: LiveUpdatesSocketLike | null, reason: string) { + if (!target) return; + + if (target.readyState === SOCKET_CONNECTING) { + // Let the handshake complete and then close. Calling close() while the + // socket is still CONNECTING is what triggers the noisy browser error. + target.onopen = () => { + resetSocketHandlers(target); + target.close(1000, reason); + }; + target.onmessage = null; + target.onerror = () => undefined; + target.onclose = null; + return; + } + + resetSocketHandlers(target); + + if (target.readyState === SOCKET_OPEN) { + target.close(1000, reason); + } +} + export const __liveUpdatesTestUtils = { + closeSocketQuietly, invalidateActivityQueries, + resolveLiveCompanyId, shouldSuppressActivityToastForVisibleIssue, shouldSuppressRunStatusToastForVisibleIssue, shouldSuppressAgentStatusToastForVisibleIssue, }; export function LiveUpdatesProvider({ children }: { children: ReactNode }) { - const { selectedCompanyId } = useCompany(); + const { selectedCompanyId, selectedCompany } = useCompany(); const queryClient = useQueryClient(); const { pushToast } = useToast(); const location = useLocation(); const gateRef = useRef({ cooldownHits: new Map(), suppressUntil: 0 }); const pathnameRef = useRef(location.pathname); - const { data: session } = useQuery({ + const { data: session, status: sessionStatus } = useQuery({ queryKey: queryKeys.auth.session, queryFn: () => authApi.getSession(), retry: false, }); const currentUserId = session?.user?.id ?? session?.session?.userId ?? null; + const socketAuthKey = session?.session?.id ?? currentUserId ?? "signed_out"; + const liveCompanyId = resolveLiveCompanyId(selectedCompanyId, selectedCompany?.id ?? null); + const canConnectSocket = sessionStatus === "success" && session !== null && liveCompanyId !== null; + const currentActorRef = useRef<{ userId: string | null; agentId: string | null }>({ + userId: currentUserId, + agentId: null, + }); useEffect(() => { pathnameRef.current = location.pathname; }, [location.pathname]); useEffect(() => { - if (!selectedCompanyId) return; + currentActorRef.current = { + userId: currentUserId, + agentId: null, + }; + }, [currentUserId]); + + useEffect(() => { + if (!canConnectSocket || !liveCompanyId) return; let closed = false; let reconnectAttempt = 0; @@ -705,55 +771,63 @@ export function LiveUpdatesProvider({ children }: { children: ReactNode }) { const connect = () => { if (closed) return; const protocol = window.location.protocol === "https:" ? "wss" : "ws"; - const url = `${protocol}://${window.location.host}/api/companies/${encodeURIComponent(selectedCompanyId)}/events/ws`; - socket = new WebSocket(url); + const url = `${protocol}://${window.location.host}/api/companies/${encodeURIComponent(liveCompanyId)}/events/ws`; + const nextSocket = new WebSocket(url); + socket = nextSocket; - socket.onopen = () => { + nextSocket.onopen = () => { + if (closed || socket !== nextSocket) { + closeSocketQuietly(nextSocket, "stale_connection"); + return; + } if (reconnectAttempt > 0) { gateRef.current.suppressUntil = Date.now() + RECONNECT_SUPPRESS_MS; } reconnectAttempt = 0; }; - socket.onmessage = (message) => { + nextSocket.onmessage = (message) => { const raw = typeof message.data === "string" ? message.data : ""; if (!raw) return; try { const parsed = JSON.parse(raw) as LiveEvent; - handleLiveEvent(queryClient, selectedCompanyId, pathnameRef.current, parsed, pushToast, gateRef.current, { - userId: currentUserId, - agentId: null, + handleLiveEvent(queryClient, liveCompanyId, pathnameRef.current, parsed, pushToast, gateRef.current, { + userId: currentActorRef.current.userId, + agentId: currentActorRef.current.agentId, }); } catch { // Ignore non-JSON payloads. } }; - socket.onerror = () => { - socket?.close(); + nextSocket.onerror = () => { + // Wait for onclose to drive the reconnect. Self-closing here is what + // produces the "closed before connection established" browser noise. }; - socket.onclose = () => { + nextSocket.onclose = () => { + if (socket !== nextSocket) return; + socket = null; if (closed) return; scheduleReconnect(); }; }; - connect(); + // Delay initial connect slightly so React StrictMode's double-invoke + // cleanup fires before the WebSocket is created, avoiding the + // "WebSocket closed before connection established" dev-mode error. + const connectTimer = window.setTimeout(connect, 0); return () => { closed = true; + window.clearTimeout(connectTimer); clearReconnect(); - if (socket) { - socket.onopen = null; - socket.onmessage = null; - socket.onerror = null; - socket.onclose = null; - socket.close(1000, "provider_unmount"); - } + const activeSocket = socket; + socket = null; + closeSocketQuietly(activeSocket, "provider_unmount"); }; - }, [queryClient, selectedCompanyId, pushToast, currentUserId]); + }, [queryClient, liveCompanyId, pushToast, canConnectSocket, socketAuthKey]); return <>{children}; }