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,147 @@
 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";
+import {
+	type MockWebSocketPublisher,
+	createMockWebSocket,
+} from "testHelpers/websockets";
+import { OneWayWebSocket } from "utils/OneWayWebSocket";
+import { createUseAgentLogs } 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
- */
+const millisecondsInOneMinute = 60_000;
 
-describe.skip("useAgentLogs", () => {
-	afterEach(() => {
-		WS.clean();
+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: "",
+		};
 	});
+}
 
-	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;
+// 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 };
 
-		// Send 3 logs
-		server.send(JSON.stringify(generateLogs(3)));
-		await waitFor(() => {
-			expect(result.current).toHaveLength(3);
+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(initialAgentId: string): 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 mountHook
+	// returns its value.
+	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;
+			},
 		});
+	});
+
+	const { result, rerender } = renderHook(
+		({ agentId, enabled }) => useAgentLogs(agentId, enabled),
+		{ initialProps: { agentId: initialAgentId, enabled: true } },
+	);
+
+	return {
+		rerender,
+		hookResult: result,
+		publisherResult: publisherResult as PublisherResult,
+	};
+}
+
+describe("useAgentLogs", () => {
+	it("Automatically sorts logs that are received out of order", async () => {
+		const { hookResult, publisherResult } = mountHook(MockWorkspaceAgent.id);
+		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));
+	});
 
-		// Disable the hook
-		rerender({ enabled: false });
+	it("Automatically closes the socket connection when the hook is disabled", async () => {
+		const { publisherResult, rerender } = mountHook(MockWorkspaceAgent.id);
+		expect(publisherResult.current.isConnectionOpen()).toBe(true);
+		rerender({ agentId: MockWorkspaceAgent.id, enabled: false });
 		await waitFor(() => {
-			expect(result.current).toHaveLength(0);
+			expect(publisherResult.current.isConnectionOpen()).toBe(false);
 		});
+	});
 
-		// Enable the hook again
-		rerender({ enabled: true });
-		await server.connected;
-		server.send(JSON.stringify(generateLogs(3)));
-		await waitFor(() => {
-			expect(result.current).toHaveLength(3);
+	it("Automatically closes the old connection when the agent ID changes", () => {
+		const { publisherResult, rerender } = mountHook(MockWorkspaceAgent.id);
+		const publisher1 = publisherResult.current;
+		expect(publisher1.isConnectionOpen()).toBe(true);
+
+		const newAgentId = `${MockWorkspaceAgent.id}-2`;
+		rerender({ agentId: newAgentId, enabled: true });
+
+		const publisher2 = publisherResult.current;
+		expect(publisher1.isConnectionOpen()).toBe(false);
+		expect(publisher2.isConnectionOpen()).toBe(true);
+	});
+
+	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(
+			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));
+
+		// Disable the hook (and have the hook close the connection behind the
+		// scenes)
+		rerender({ agentId: MockWorkspaceAgent.id, enabled: false });
+		await waitFor(() => expect(hookResult.current).toHaveLength(0));
+
+		// Re-enable the hook (creating an entirely new connection), and send
+		// new logs
+		rerender({ agentId: MockWorkspaceAgent.id, enabled: true });
+		const newLogs = generateMockLogs(3, new Date("october 3, 2005"));
+		const newEvent = new MessageEvent<string>("message", {
+			data: JSON.stringify(newLogs),
 		});
+		act(() => publisherResult.current.publishMessage(newEvent));
+		await waitFor(() => expect(hookResult.current).toEqual(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: "",
-	}));
-}
@@ -1,47 +1,89 @@
-import { watchWorkspaceAgentLogs } from "api/api";
-import type { WorkspaceAgent, WorkspaceAgentLog } from "api/typesGenerated";
+import {
+	type WatchWorkspaceAgentLogsParams,
+	watchWorkspaceAgentLogs,
+} from "api/api";
+import type { WorkspaceAgentLog } from "api/typesGenerated";
 import { displayError } from "components/GlobalSnackbar/utils";
 import { useEffect, useState } from "react";
+import type { OneWayWebSocket } from "utils/OneWayWebSocket";
 
-export function useAgentLogs(
-	agent: WorkspaceAgent,
-	enabled: boolean,
-): readonly WorkspaceAgentLog[] {
-	const [logs, setLogs] = useState<WorkspaceAgentLog[]>([]);
+type CreateSocket = (
+	agentId: string,
+	params?: WatchWorkspaceAgentLogsParams,
+) => OneWayWebSocket<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: CreateSocket) {
+	return function useAgentLogs(
+		agentId: string,
+		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 createdAtMap = new Map<string, number>();
+			const socket = createSocket(agentId, { after: 0 });
+			socket.addEventListener("message", (e) => {
+				if (e.parseError) {
+					console.warn("Error parsing agent log: ", e.parseError);
+					return;
+				}
+
+				if (e.parsedMessage.length === 0) {
+					return;
+				}
+
+				setLogs((logs) => {
+					const newLogs = [...logs, ...e.parsedMessage];
+					newLogs.sort((l1, l2) => {
+						let d1 = createdAtMap.get(l1.created_at);
+						if (d1 === undefined) {
+							d1 = new Date(l1.created_at).getTime();
+							createdAtMap.set(l1.created_at, d1);
+						}
+						let d2 = createdAtMap.get(l2.created_at);
+						if (d2 === undefined) {
+							d2 = new Date(l2.created_at).getTime();
+							createdAtMap.set(l2.created_at, d2);
+						}
+						return d1 - d2;
+					});
+
+					return newLogs;
+				});
+			});
+
+			socket.addEventListener("error", (e) => {
+				console.error("Error in agent log socket: ", e);
+				displayError(
+					"Unable to watch agent logs",
+					"Please try refreshing the browser",
+				);
+				socket.close();
+			});
+
+			return () => socket.close();
+		}, [createSocket, agentId, enabled]);
+
+		return logs;
+	};
 }
+
+// The baseline implementation to use for production
+export const useAgentLogs = createUseAgentLogs(watchWorkspaceAgentLogs);
@@ -308,12 +308,7 @@ type AgentLogsContentProps = {
 };
 
 const AgentLogsContent: FC<AgentLogsContentProps> = ({ agent }) => {
-	const logs = useAgentLogs(agent, true);
-
-	if (!logs) {
-		return <Loader />;
-	}
-
+	const logs = useAgentLogs(agent.id, true);
 	return (
 		<AgentLogs
 			sources={agent.log_sources}
-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 @@