coder · Parkreiner · May 23, 2025 · May 23, 2025 · May 23, 2025 · May 23, 2025
diff --git a/site/src/modules/resources/useAgentLogs.test.ts b/site/src/modules/resources/useAgentLogs.test.ts
@@ -1,60 +1,103 @@
 import { renderHook, waitFor } from "@testing-library/react";
 import type { WorkspaceAgentLog } from "api/typesGenerated";
-import WS from "jest-websocket-mock";
 import { MockWorkspaceAgent } from "testHelpers/entities";
-import { useAgentLogs } from "./useAgentLogs";
+import { createUseAgentLogs } from "./useAgentLogs";
+import {
+	createMockWebSocket,
+	type MockWebSocketPublisher,
+} from "testHelpers/websockets";
+import { OneWayWebSocket } from "utils/OneWayWebSocket";
 
-/**
- * TODO: WS does not support multiple tests running at once in isolation so we
- * have one single test that test the most common scenario.
- * Issue: https://github.com/romgain/jest-websocket-mock/issues/172
- */
+function generateMockLogs(count: number): WorkspaceAgentLog[] {
+	return Array.from({ length: count }, (_, i) => ({
+		id: i,
+		created_at: new Date().toISOString(),
+		level: "info",
+		output: `Log ${i}`,
+		source_id: "",
+	}));
+}
 
-describe("useAgentLogs", () => {
-	afterEach(() => {
-		WS.clean();
+// A mutable object holding the most recent mock WebSocket publisher. The inner
+// value will change as the hook opens/closes new connections
+type PublisherResult = {
+	current: MockWebSocketPublisher;
+};
+
+type MountHookResult = Readonly<{
+	// Note: the value of `current` should be readonly, but the `current`
+	// property itself should be mutable
+	hookResult: {
+		current: readonly WorkspaceAgentLog[];
+	};
+	rerender: (props: { enabled: boolean }) => void;
+	publisherResult: PublisherResult;
+}>;
+
+function mountHook(): MountHookResult {
+	// Have to cheat the types a little bit to avoid a chicken-and-the-egg
+	// scenario. publisherResult will be initialized with an undefined current
+	// value, but it'll be guaranteed not to be undefined by the time this
+	// function returns.
+	const publisherResult: Partial<PublisherResult> = { current: undefined };
+	const useAgentLogs = createUseAgentLogs((agentId, params) => {
+		return new OneWayWebSocket({
+			apiRoute: `/api/v2/workspaceagents/${agentId}/logs`,
+			searchParams: new URLSearchParams({
+				follow: "true",
+				after: params?.after?.toString() || "0",
+			}),
+			websocketInit: (url) => {
+				const [mockSocket, mockPublisher] = createMockWebSocket(url);
+				publisherResult.current = mockPublisher;
+				return mockSocket;
+			},
+		});
 	});
 
-	it("clear logs when disabled to avoid duplicates", async () => {
-		const server = new WS(
-			`ws://localhost/api/v2/workspaceagents/${
-				MockWorkspaceAgent.id
-			}/logs?follow&after=0`,
-		);
-		const { result, rerender } = renderHook(
-			({ enabled }) => useAgentLogs(MockWorkspaceAgent, enabled),
-			{ initialProps: { enabled: true } },
-		);
-		await server.connected;
-
-		// Send 3 logs
-		server.send(JSON.stringify(generateLogs(3)));
+	const { result, rerender } = renderHook(
+		({ enabled }) => useAgentLogs(MockWorkspaceAgent, enabled),
+		{ initialProps: { enabled: true } },
+	);
+
+	return {
+		rerender,
+		hookResult: result,
+		publisherResult: publisherResult as PublisherResult,
+	};
+}
+
+describe("useAgentLogs", () => {
+	it("clears logs when hook becomes disabled (protection to avoid duplicate logs when hook goes back to being re-enabled)", async () => {
+		const { hookResult, publisherResult, rerender } = mountHook();
+
+		// Verify that logs can be received after mount
+		const initialLogs = generateMockLogs(3);
+		const initialEvent = new MessageEvent<string>("message", {
+			data: JSON.stringify(initialLogs),
+		});
+		publisherResult.current.publishMessage(initialEvent);
 		await waitFor(() => {
-			expect(result.current).toHaveLength(3);
+			// Using expect.arrayContaining to account for the fact that we're
+			// not guaranteed to receive WebSocket events in order
+			expect(hookResult.current).toEqual(expect.arrayContaining(initialLogs));
 		});
 
-		// Disable the hook
+		// Disable the hook (and have the hook close the connection behind the
+		// scenes)
 		rerender({ enabled: false });
-		await waitFor(() => {
-			expect(result.current).toHaveLength(0);
-		});
+		await waitFor(() => expect(hookResult.current).toHaveLength(0));
 
-		// Enable the hook again
+		// Re-enable the hook (creating an entirely new connection), and send
+		// new logs
 		rerender({ enabled: true });
-		await server.connected;
-		server.send(JSON.stringify(generateLogs(3)));
+		const newLogs = generateMockLogs(3);
+		const newEvent = new MessageEvent<string>("message", {
+			data: JSON.stringify(newLogs),
+		});
+		publisherResult.current.publishMessage(newEvent);
 		await waitFor(() => {
-			expect(result.current).toHaveLength(3);
+			expect(hookResult.current).toEqual(expect.arrayContaining(newLogs));
 		});
 	});
 });
-
-function generateLogs(count: number): WorkspaceAgentLog[] {
-	return Array.from({ length: count }, (_, i) => ({
-		id: i,
-		created_at: new Date().toISOString(),
-		level: "info",
-		output: `Log ${i}`,
-		source_id: "",
-	}));
-}
diff --git a/site/src/modules/resources/useAgentLogs.ts b/site/src/modules/resources/useAgentLogs.ts
@@ -3,45 +3,63 @@ import type { WorkspaceAgent, WorkspaceAgentLog } from "api/typesGenerated";
 import { displayError } from "components/GlobalSnackbar/utils";
 import { useEffect, useState } from "react";
 
-export function useAgentLogs(
-	agent: WorkspaceAgent,
-	enabled: boolean,
-): readonly WorkspaceAgentLog[] {
-	const [logs, setLogs] = useState<WorkspaceAgentLog[]>([]);
-
-	useEffect(() => {
-		if (!enabled) {
-			// Clean up the logs when the agent is not enabled. So it can receive logs
-			// from the beginning without duplicating the logs.
+export function createUseAgentLogs(
+	createSocket: typeof watchWorkspaceAgentLogs,
+) {
+	return function useAgentLogs(
+		agent: WorkspaceAgent,
+		enabled: boolean,
+	): readonly WorkspaceAgentLog[] {
+		const [logs, setLogs] = useState<readonly WorkspaceAgentLog[]>([]);
+
+		// Clean up the logs when the agent is not enabled, using a mid-render
+		// sync to remove any risk of screen flickering. Clearing the logs helps
+		// ensure that if the hook flips back to being enabled, we can receive a
+		// fresh set of logs from the beginning with zero risk of duplicates.
+		const [prevEnabled, setPrevEnabled] = useState(enabled);
+		if (!enabled && prevEnabled) {
 			setLogs([]);
-			return;
+			setPrevEnabled(false);
 		}
 
-		// Always fetch the logs from the beginning. We may want to optimize this in
-		// the future, but it would add some complexity in the code that maybe does
-		// not worth it.
-		const socket = watchWorkspaceAgentLogs(agent.id, { after: 0 });
-		socket.addEventListener("message", (e) => {
-			if (e.parseError) {
-				console.warn("Error parsing agent log: ", e.parseError);
+		useEffect(() => {
+			if (!enabled) {
 				return;
 			}
-			setLogs((logs) => [...logs, ...e.parsedMessage]);
-		});
-
-		socket.addEventListener("error", (e) => {
-			console.error("Error in agent log socket: ", e);
-			displayError(
-				"Unable to watch the agent logs",
-				"Please try refreshing the browser",
-			);
-			socket.close();
-		});
-
-		return () => {
-			socket.close();
-		};
-	}, [agent.id, enabled]);
-
-	return logs;
+
+			// Always fetch the logs from the beginning. We may want to optimize
+			// this in the future, but it would add some complexity in the code
+			// that might not be worth it.
+			const socket = createSocket(agent.id, { after: 0 });
+			socket.addEventListener("message", (e) => {
+				if (e.parseError) {
+					console.warn("Error parsing agent log: ", e.parseError);
+					return;
+				}
+				setLogs((logs) => [...logs, ...e.parsedMessage]);
+			});
+
+			socket.addEventListener("error", (e) => {
+				console.error("Error in agent log socket: ", e);
+				displayError(
+					"Unable to watch the agent logs",
+					"Please try refreshing the browser",
+				);
+				socket.close();
+			});
+
+			return () => socket.close();
+
+			// createSocket shouldn't ever change for the lifecycle of the hook,
+			// but Biome isn't smart enough to detect constant dependencies for
+			// higher-order hooks. Adding it to the array (even though it
+			// shouldn't ever be needed) seemed like the least fragile way to
+			// resolve the warning.
+		}, [createSocket, agent.id, enabled]);
+
+		return logs;
+	};
 }
+
+// The baseline implementation to use for production
+export const useAgentLogs = createUseAgentLogs(watchWorkspaceAgentLogs);
diff --git a/site/src/testHelpers/websockets.ts b/site/src/testHelpers/websockets.ts
@@ -0,0 +1,135 @@
+import type { WebSocketEventType } from "utils/OneWayWebSocket";
+
+export type MockWebSocketPublisher = Readonly<{
+	publishMessage: (event: MessageEvent<string>) => void;
+	publishError: (event: ErrorEvent) => void;
+	publishClose: (event: CloseEvent) => void;
+	publishOpen: (event: Event) => void;
+}>;
+
+export function createMockWebSocket(
+	url: string,
+	protocols?: string | string[],
+): readonly [WebSocket, MockWebSocketPublisher] {
+	type EventMap = {
+		message: MessageEvent<string>;
+		error: ErrorEvent;
+		close: CloseEvent;
+		open: Event;
+	};
+	type CallbackStore = {
+		[K in keyof EventMap]: ((event: EventMap[K]) => void)[];
+	};
+
+	let activeProtocol: string;
+	if (Array.isArray(protocols)) {
+		activeProtocol = protocols[0] ?? "";
+	} else if (typeof protocols === "string") {
+		activeProtocol = protocols;
+	} else {
+		activeProtocol = "";
+	}
+
+	let closed = false;
+	const store: CallbackStore = {
+		message: [],
+		error: [],
+		close: [],
+		open: [],
+	};
+
+	const mockSocket: WebSocket = {
+		CONNECTING: 0,
+		OPEN: 1,
+		CLOSING: 2,
+		CLOSED: 3,
+
+		url,
+		protocol: activeProtocol,
+		readyState: 1,
+		binaryType: "blob",
+		bufferedAmount: 0,
+		extensions: "",
+		onclose: null,
+		onerror: null,
+		onmessage: null,
+		onopen: null,
+		send: jest.fn(),
+		dispatchEvent: jest.fn(),
+
+		addEventListener: <E extends WebSocketEventType>(
+			eventType: E,
+			callback: WebSocketEventMap[E],
+		) => {
+			if (closed) {
+				return;
+			}
+
+			const subscribers = store[eventType];
+			const cb = callback as unknown as CallbackStore[E][0];
+			if (!subscribers.includes(cb)) {
+				subscribers.push(cb);
+			}
+		},
+
+		removeEventListener: <E extends WebSocketEventType>(
+			eventType: E,
+			callback: WebSocketEventMap[E],
+		) => {
+			if (closed) {
+				return;
+			}
+
+			const subscribers = store[eventType];
+			const cb = callback as unknown as CallbackStore[E][0];
+			if (subscribers.includes(cb)) {
+				const updated = store[eventType].filter((c) => c !== cb);
+				store[eventType] = updated as unknown as CallbackStore[E];
+			}
+		},
+
+		close: () => {
+			closed = true;
+		},
+	};
+
+	const publisher: MockWebSocketPublisher = {
+		publishOpen: (event) => {
+			if (closed) {
+				return;
+			}
+			for (const sub of store.open) {
+				sub(event);
+			}
+		},
+
+		publishError: (event) => {
+			if (closed) {
+				return;
+			}
+			for (const sub of store.error) {
+				sub(event);
+			}
+		},
+
+		publishMessage: (event) => {
+			if (closed) {
+				return;
+			}
+			for (const sub of store.message) {
+				sub(event);
+			}
+		},
+
+		publishClose: (event) => {
+			if (closed) {
+				return;
+			}
+			for (const sub of store.close) {
+				sub(event);
+			}
+		},
+	};
+
+	return [mockSocket, publisher] as const;
+}