May 23, 2025 · May 23, 2025 · 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
 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 {
 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("useAgentLogs", () => {
 afterEach(() => {
 WS.clean();
 function generateMockLogs(
 logCount: number,
 baseDate = new Date(),
 ): readonly WorkspaceAgentLog[] {
 return Array.from({ length: logCount }, (_, i) => {
 // Make sure that the logs generated each have unique timestamps, so
 // that we can test whether they're being sorted properly before being
 // returned by the hook
 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;

 // Send 3 logs
 server.send(JSON.stringify(generateLogs(3)));
 await waitFor(() => {
 expect(result.current).toHaveLength(3);
 // 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;
 },
 });
 });

 // Disable the hook
 rerender({ enabled: false });
 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, new Date("april 5, 1997"));
 const initialEvent = new MessageEvent<string>("message", {
 data: JSON.stringify(initialLogs),
 });
 publisherResult.current.publishMessage(initialEvent);
 await waitFor(() => {
 expect(result.current).toHaveLength(0);
 // 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));
 });

 // Enable the hook again
 // Disable the hook (and have the hook close the connection behind the
 // scenes)
 rerender({ enabled: false });
 await waitFor(() => expect(hookResult.current).toHaveLength(0));

 // 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, new Date("october 3, 2005"));
 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
 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
 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;
 }
Original file line number	Diff line number	Diff line change
		@@ -1,60 +1,114 @@
		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 {
		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("useAgentLogs", () => {
		afterEach(() => {
		WS.clean();
		function generateMockLogs(
		logCount: number,
		baseDate = new Date(),
		): readonly WorkspaceAgentLog[] {
		return Array.from({ length: logCount }, (_, i) => {
		// Make sure that the logs generated each have unique timestamps, so
		// that we can test whether they're being sorted properly before being
		// returned by the hook
		const logDate = new Date(baseDate.getTime() + i * millisecondsInOneMinute);
		return {
Comment on lines +18 to +22 Copy link MemberAuthor ParkreinerMay 23, 2025• edited Loading Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. Like I said in the PR body, I didn't add a test to assert the sorting functionality yet, but I updated the test helper anyway, so that I can make sure the old/new sets of logs created in the test case are guaranteed to be different, and there's no risk of false positives
		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;

		// Send 3 logs
		server.send(JSON.stringify(generateLogs(3)));
		await waitFor(() => {
		expect(result.current).toHaveLength(3);
		// 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;
		},
		});
		});

		// Disable the hook
		rerender({ enabled: false });
		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, new Date("april 5, 1997"));
		const initialEvent = new MessageEvent<string>("message", {
		data: JSON.stringify(initialLogs),
		});
		publisherResult.current.publishMessage(initialEvent);
		await waitFor(() => {
		expect(result.current).toHaveLength(0);
		// 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));
		});

		// Enable the hook again
		// Disable the hook (and have the hook close the connection behind the
		// scenes)
		rerender({ enabled: false });
		await waitFor(() => expect(hookResult.current).toHaveLength(0));

		// 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, new Date("october 3, 2005"));
		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: "",
		}));
		}
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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);
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,135 @@
		import type { WebSocketEventType } from "utils/OneWayWebSocket";

		export type MockWebSocketPublisher = Readonly<{
Copy link MemberAuthor ParkreinerMay 23, 2025 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. The entire contents of the file were basically copy-pasted from the other test file. The only change was that the publisher type is now exported and was renamed to be a little more clear
		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;
		}