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}>;
}