coder · Parkreiner · May 23, 2025 · May 23, 2025 · May 23, 2025 · May 23, 2025
diff --git a/site/src/api/api.ts b/site/src/api/api.ts
@@ -264,7 +264,7 @@ export const watchWorkspaceAgentLogs = (
 	});
 };
 
-type WatchWorkspaceAgentLogsParams = {
+export type WatchWorkspaceAgentLogsParams = {
 	after?: number;
 };
 

@@ -229,13 +229,12 @@ export function makeUseEmbeddedMetadata(
 			manager.getMetadata,
 		);
 
-		// biome-ignore lint/correctness/useExhaustiveDependencies(manager.clearMetadataByKey): baked into containing hook
 		const stableMetadataResult = useMemo<UseEmbeddedMetadataResult>(() => {
 			return {
 				metadata,
 				clearMetadataByKey: manager.clearMetadataByKey,
 			};
-		}, [metadata]);
+		}, [manager, metadata]);
 
 		return stableMetadataResult;
 	};

@@ -6,6 +6,7 @@ import type {
 	Template,
 	Workspace,
 	WorkspaceAgent,
+	WorkspaceAgentLog,
 	WorkspaceAgentMetadata,
 } from "api/typesGenerated";
 import { Button } from "components/Button/Button";
@@ -76,24 +77,24 @@ export const AgentRow: FC<AgentRowProps> = ({
 		["starting", "start_timeout"].includes(agent.lifecycle_state) &&
 			hasStartupFeatures,
 	);
-	const agentLogs = useAgentLogs(agent, showLogs);
+	const agentLogs = useAgentLogs(agent.id, showLogs);
 	const logListRef = useRef<List>(null);
 	const logListDivRef = useRef<HTMLDivElement>(null);
-	const startupLogs = useMemo(() => {
-		const allLogs = agentLogs || [];
-
-		const logs = [...allLogs];
-		if (agent.logs_overflowed) {
-			logs.push({
+	const startupLogs = useMemo<readonly WorkspaceAgentLog[]>(() => {
+		if (!agent.logs_overflowed) {
+			return agentLogs;
+		}
+		return [
+			...agentLogs,
+			{
 				id: -1,
 				level: "error",
 				output:
 					"Startup logs exceeded the max size of 1MB, and will not continue to be written to the database! Logs will continue to be written to the /tmp/coder-startup-script.log file in the workspace.",
 				created_at: new Date().toISOString(),
-				created_at: new Date().toISOString(),
+				created_at: agentLogs.at(-1).created_at,
-				created_at: new Date().toISOString(),
+				created_at: agentLogs.at(-1).created_at,
 				source_id: "",
-			});
-		}
-		return logs;
+			},
+		];
 	}, [agentLogs, agent.logs_overflowed]);
 	const [bottomOfLogs, setBottomOfLogs] = useState(true);
 

@@ -1,60 +1,201 @@
 import { renderHook, waitFor } from "@testing-library/react";
 import type { WorkspaceAgentLog } from "api/typesGenerated";
-import WS from "jest-websocket-mock";
+import { act } from "react";
 import { MockWorkspaceAgent } from "testHelpers/entities";
-import { useAgentLogs } from "./useAgentLogs";
-
-/**
- * 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
- */
-
-describe.skip("useAgentLogs", () => {
-	afterEach(() => {
-		WS.clean();
-	});
-
-	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)));
-		await waitFor(() => {
-			expect(result.current).toHaveLength(3);
+import {
+	type MockWebSocketPublisher,
+	createMockWebSocket,
+} from "testHelpers/websockets";
+import { OneWayWebSocket } from "utils/OneWayWebSocket";
+import {
+	type CreateUseAgentLogsOptions,
+	createUseAgentLogs,
+} from "./useAgentLogs";
+
+const millisecondsInOneMinute = 60_000;
+
+function generateMockLogs(
+	logCount: number,
+	baseDate = new Date(),
-	baseDate = new Date(),
+	baseDate = new Date("Thu, 07 Aug 2025 21:38:12 GMT"),
-	baseDate = new Date(),
+	baseDate = new Date("Thu, 07 Aug 2025 21:38:12 GMT"),
+): readonly WorkspaceAgentLog[] {
+	return Array.from({ length: logCount }, (_, i) => {
+		// Make sure that the logs generated each have unique timestamps, so
+		// that we can test whether the hook is sorting them properly as it's
+		// receiving them over time
+		const logDate = new Date(baseDate.getTime() + i * millisecondsInOneMinute);
+		return {
+			id: i,
+			created_at: logDate.toISOString(),
+			level: "info",
+			output: `Log ${i}`,
+			source_id: "",
+		};
+	});
+}
+
+// A mutable object holding the most recent mock WebSocket publisher that was
+// created when initializing a mock WebSocket. Inner value will be undefined if
+// the hook is disabled on mount, but will always be defined otherwise
+type PublisherResult = { current: MockWebSocketPublisher | undefined };
+
+type MountHookOptions = Readonly<{
+	initialAgentId: string;
+	enabled?: boolean;
+	onError?: CreateUseAgentLogsOptions["onError"];
+}>;
+
+type MountHookResult = Readonly<{
+	rerender: (props: { agentId: string; enabled: boolean }) => void;
+	publisherResult: PublisherResult;
+
+	// Note: the `current` property is only "halfway" readonly; the value is
+	// readonly, but the key is still mutable
+	hookResult: { current: readonly WorkspaceAgentLog[] };
+}>;
+
+function mountHook(options: MountHookOptions): MountHookResult {
+	const { initialAgentId, enabled = true, onError = jest.fn() } = options;
+
+	const publisherResult: PublisherResult = { current: undefined };
+	const useAgentLogs = createUseAgentLogs({
+		onError,
+		createSocket: (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;
+				},
+			});
+		},
+	});
+
+	const { result, rerender } = renderHook(
+		({ agentId, enabled }) => useAgentLogs(agentId, enabled),
+		{ initialProps: { agentId: initialAgentId, enabled: enabled } },
+	);
+
+	return {
+		rerender,
+		hookResult: result,
+		publisherResult,
+	};
+}
+
+describe("useAgentLogs", () => {
+	it("Automatically sorts logs that are received out of order", async () => {
+		const { hookResult, publisherResult } = mountHook({
+			initialAgentId: MockWorkspaceAgent.id,
 		});
 
-		// Disable the hook
-		rerender({ enabled: false });
-		await waitFor(() => {
-			expect(result.current).toHaveLength(0);
+		const logs = generateMockLogs(10, new Date("september 9, 1999"));
+		const reversed = logs.toReversed();
+
+		for (const log of reversed) {
+			act(() => {
+				publisherResult.current?.publishMessage(
+					new MessageEvent<string>("message", {
+						data: JSON.stringify([log]),
+					}),
+				);
+			});
+		}
+		await waitFor(() => expect(hookResult.current).toEqual(logs));
+	});
+
+	it("Never creates a connection if hook is disabled on mount", () => {
+		const { publisherResult } = mountHook({
+			initialAgentId: MockWorkspaceAgent.id,
+			enabled: false,
 		});
 
-		// Enable the hook again
-		rerender({ enabled: true });
-		await server.connected;
-		server.send(JSON.stringify(generateLogs(3)));
+		expect(publisherResult.current).toBe(undefined);
+	});
+
+	it("Automatically closes the socket connection when the hook is disabled", async () => {
+		const { publisherResult, rerender } = mountHook({
+			initialAgentId: MockWorkspaceAgent.id,
+		});
+
+		expect(publisherResult.current?.isConnectionOpen).toBe(true);
+		rerender({ agentId: MockWorkspaceAgent.id, enabled: false });
 		await waitFor(() => {
-			expect(result.current).toHaveLength(3);
+			expect(publisherResult.current?.isConnectionOpen).toBe(false);
 		});
 	});
-});
 
-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: "",
-	}));
-}
+	it("Automatically closes the old connection when the agent ID changes", () => {
+		const { publisherResult, rerender } = mountHook({
+			initialAgentId: MockWorkspaceAgent.id,
+		});
+
+		const publisher1 = publisherResult.current;
+		expect(publisher1?.isConnectionOpen).toBe(true);
+
+		rerender({
+			enabled: true,
+			agentId: `${MockWorkspaceAgent.id}-new-value`,
+		});
+
+		const publisher2 = publisherResult.current;
+		expect(publisher1).not.toBe(publisher2);
+		expect(publisher1?.isConnectionOpen).toBe(false);
+		expect(publisher2?.isConnectionOpen).toBe(true);
+	});
+
+	it("Calls error callback when error is received (but only while hook is enabled)", async () => {
+		const onError = jest.fn();
+		const { publisherResult, rerender } = mountHook({
+			initialAgentId: MockWorkspaceAgent.id,
+			// Start off disabled so that we can check that the callback is
+			// never called when there is no connection
+			enabled: false,
+			onError,
+		});
+
+		const errorEvent = new Event("error");
+		act(() => publisherResult.current?.publishError(errorEvent));
+		expect(onError).not.toHaveBeenCalled();
+
+		rerender({ agentId: MockWorkspaceAgent.id, enabled: true });
+		act(() => publisherResult.current?.publishError(errorEvent));
+		expect(onError).toHaveBeenCalledTimes(1);
+	});
+
+	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({
+			initialAgentId: MockWorkspaceAgent.id,
+		});
+
+		// Send initial logs so that we have something to clear out later
+		const initialLogs = generateMockLogs(3, new Date("april 5, 1997"));
+		const initialEvent = new MessageEvent<string>("message", {
+			data: JSON.stringify(initialLogs),
+		});
+		act(() => publisherResult.current?.publishMessage(initialEvent));
+		await waitFor(() => expect(hookResult.current).toEqual(initialLogs));
+
+		// Need to do the following steps multiple times to make sure that we
+		// don't break anything after the first disable
+		const mockDates: readonly string[] = ["october 3, 2005", "august 1, 2025"];
+		for (const md of mockDates) {
+			// Disable the hook to clear current logs
+			rerender({ agentId: MockWorkspaceAgent.id, enabled: false });
+			await waitFor(() => expect(hookResult.current).toHaveLength(0));
+
+			// Re-enable the hook and send new logs
+			rerender({ agentId: MockWorkspaceAgent.id, enabled: true });
+			const newLogs = generateMockLogs(3, new Date(md));
+			const newEvent = new MessageEvent<string>("message", {
+				data: JSON.stringify(newLogs),
+			});
+			act(() => publisherResult.current?.publishMessage(newEvent));
+			await waitFor(() => expect(hookResult.current).toEqual(newLogs));
+		}
+	});
+});
-Original file line number
+Diff line change
@@ Expand Up / @@ -264,7 +264,7 @@ export const watchWorkspaceAgentLogs = ( @@
     	});
     };
-    type WatchWorkspaceAgentLogsParams = {
+    export type WatchWorkspaceAgentLogsParams = {
     	after?: number;
     };
@@ Expand Down @@