fix: fix duplicated agent logs (#17806)

Fix https://github.com/coder/coder/issues/16355

site/src/api/api.ts

@@ -221,11 +221,11 @@ export const watchBuildLogsByTemplateVersionId = (
 
 export const watchWorkspaceAgentLogs = (
 	agentId: string,
-	{ after, onMessage, onDone, onError }: WatchWorkspaceAgentLogsOptions,
+	params?: WatchWorkspaceAgentLogsParams,
 ) => {
 	const searchParams = new URLSearchParams({
 		follow: "true",
-		after: after.toString(),
+		after: params?.after?.toString() ?? "",
 	});
 
 	/**
@@ -237,32 +237,14 @@ export const watchWorkspaceAgentLogs = (
 		searchParams.set("no_compression", "");
 	}
 
-	const socket = createWebSocket(
-		`/api/v2/workspaceagents/${agentId}/logs`,
+	return new OneWayWebSocket<TypesGen.WorkspaceAgentLog[]>({
+		apiRoute: `/api/v2/workspaceagents/${agentId}/logs`,
 		searchParams,
-	);
-
-	socket.addEventListener("message", (event) => {
-		const logs = JSON.parse(event.data) as TypesGen.WorkspaceAgentLog[];
-		onMessage(logs);
 	});
-
-	socket.addEventListener("error", () => {
-		onError(new Error("socket errored"));
-	});
-
-	socket.addEventListener("close", () => {
-		onDone?.();
-	});
-
-	return socket;
 };
 
-type WatchWorkspaceAgentLogsOptions = {
-	after: number;
-	onMessage: (logs: TypesGen.WorkspaceAgentLog[]) => void;
-	onDone?: () => void;
-	onError: (error: Error) => void;
+type WatchWorkspaceAgentLogsParams = {
+	after?: number;
 };
 
 type WatchBuildLogsByBuildIdOptions = {

site/src/api/queries/workspaces.ts

@@ -5,6 +5,7 @@ import type {
 	ProvisionerLogLevel,
 	UsageAppName,
 	Workspace,
+	WorkspaceAgentLog,
 	WorkspaceBuild,
 	WorkspaceBuildParameter,
 	WorkspacesRequest,
@@ -20,6 +21,7 @@ import type {
 	QueryClient,
 	QueryOptions,
 	UseMutationOptions,
+	UseQueryOptions,
 } from "react-query";
 import { checkAuthorization } from "./authCheck";
 import { disabledRefetchOptions } from "./util";
@@ -342,20 +344,14 @@ export const buildLogs = (workspace: Workspace) => {
 	};
 };
 
-export const agentLogsKey = (workspaceId: string, agentId: string) => [
-	"workspaces",
-	workspaceId,
-	"agents",
-	agentId,
-	"logs",
-];
+export const agentLogsKey = (agentId: string) => ["agents", agentId, "logs"];
 
-export const agentLogs = (workspaceId: string, agentId: string) => {
+export const agentLogs = (agentId: string) => {
 	return {
-		queryKey: agentLogsKey(workspaceId, agentId),
+		queryKey: agentLogsKey(agentId),
 		queryFn: () => API.getWorkspaceAgentLogs(agentId),
 		...disabledRefetchOptions,
-	};
+	} satisfies UseQueryOptions<WorkspaceAgentLog[]>;
 };
 
 // workspace usage options

site/src/modules/resources/AgentLogs/useAgentLogs.test.tsx

@@ -1,142 +0,0 @@
-import { act, renderHook, waitFor } from "@testing-library/react";
-import { API } from "api/api";
-import * as APIModule from "api/api";
-import { agentLogsKey } from "api/queries/workspaces";
-import type { WorkspaceAgentLog } from "api/typesGenerated";
-import WS from "jest-websocket-mock";
-import { type QueryClient, QueryClientProvider } from "react-query";
-import { MockWorkspace, MockWorkspaceAgent } from "testHelpers/entities";
-import { createTestQueryClient } from "testHelpers/renderHelpers";
-import { type UseAgentLogsOptions, useAgentLogs } from "./useAgentLogs";
-
-afterEach(() => {
-	WS.clean();
-});
-
-describe("useAgentLogs", () => {
-	it("should not fetch logs if disabled", async () => {
-		const queryClient = createTestQueryClient();
-		const fetchSpy = jest.spyOn(API, "getWorkspaceAgentLogs");
-		const wsSpy = jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
-		renderUseAgentLogs(queryClient, {
-			workspaceId: MockWorkspace.id,
-			agentId: MockWorkspaceAgent.id,
-			agentLifeCycleState: "ready",
-			enabled: false,
-		});
-		expect(fetchSpy).not.toHaveBeenCalled();
-		expect(wsSpy).not.toHaveBeenCalled();
-	});
-
-	it("should return existing logs without network calls if state is off", async () => {
-		const queryClient = createTestQueryClient();
-		queryClient.setQueryData(
-			agentLogsKey(MockWorkspace.id, MockWorkspaceAgent.id),
-			generateLogs(5),
-		);
-		const fetchSpy = jest.spyOn(API, "getWorkspaceAgentLogs");
-		const wsSpy = jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
-		const { result } = renderUseAgentLogs(queryClient, {
-			workspaceId: MockWorkspace.id,
-			agentId: MockWorkspaceAgent.id,
-			agentLifeCycleState: "off",
-		});
-		await waitFor(() => {
-			expect(result.current).toHaveLength(5);
-		});
-		expect(fetchSpy).not.toHaveBeenCalled();
-		expect(wsSpy).not.toHaveBeenCalled();
-	});
-
-	it("should fetch logs when empty", async () => {
-		const queryClient = createTestQueryClient();
-		const fetchSpy = jest
-			.spyOn(API, "getWorkspaceAgentLogs")
-			.mockResolvedValueOnce(generateLogs(5));
-		jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
-		const { result } = renderUseAgentLogs(queryClient, {
-			workspaceId: MockWorkspace.id,
-			agentId: MockWorkspaceAgent.id,
-			agentLifeCycleState: "ready",
-		});
-		await waitFor(() => {
-			expect(result.current).toHaveLength(5);
-		});
-		expect(fetchSpy).toHaveBeenCalledWith(MockWorkspaceAgent.id);
-	});
-
-	it("should fetch logs and connect to websocket", async () => {
-		const queryClient = createTestQueryClient();
-		const logs = generateLogs(5);
-		const fetchSpy = jest
-			.spyOn(API, "getWorkspaceAgentLogs")
-			.mockResolvedValueOnce(logs);
-		const wsSpy = jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
-		new WS(
-			`ws://localhost/api/v2/workspaceagents/${
-				MockWorkspaceAgent.id
-			}/logs?follow&after=${logs[logs.length - 1].id}`,
-		);
-		const { result } = renderUseAgentLogs(queryClient, {
-			workspaceId: MockWorkspace.id,
-			agentId: MockWorkspaceAgent.id,
-			agentLifeCycleState: "starting",
-		});
-		await waitFor(() => {
-			expect(result.current).toHaveLength(5);
-		});
-		expect(fetchSpy).toHaveBeenCalledWith(MockWorkspaceAgent.id);
-		expect(wsSpy).toHaveBeenCalledWith(MockWorkspaceAgent.id, {
-			after: logs[logs.length - 1].id,
-			onMessage: expect.any(Function),
-			onError: expect.any(Function),
-		});
-	});
-
-	it("update logs from websocket messages", async () => {
-		const queryClient = createTestQueryClient();
-		const logs = generateLogs(5);
-		jest.spyOn(API, "getWorkspaceAgentLogs").mockResolvedValueOnce(logs);
-		const server = new WS(
-			`ws://localhost/api/v2/workspaceagents/${
-				MockWorkspaceAgent.id
-			}/logs?follow&after=${logs[logs.length - 1].id}`,
-		);
-		const { result } = renderUseAgentLogs(queryClient, {
-			workspaceId: MockWorkspace.id,
-			agentId: MockWorkspaceAgent.id,
-			agentLifeCycleState: "starting",
-		});
-		await waitFor(() => {
-			expect(result.current).toHaveLength(5);
-		});
-		await server.connected;
-		act(() => {
-			server.send(JSON.stringify(generateLogs(3)));
-		});
-		await waitFor(() => {
-			expect(result.current).toHaveLength(8);
-		});
-	});
-});
-
-function renderUseAgentLogs(
-	queryClient: QueryClient,
-	options: UseAgentLogsOptions,
-) {
-	return renderHook(() => useAgentLogs(options), {
-		wrapper: ({ children }) => (
-			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
-		),
-	});
-}
-
-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: "",
-	}));
-}

site/src/modules/resources/AgentLogs/useAgentLogs.ts

@@ -1,95 +0,0 @@
-import { watchWorkspaceAgentLogs } from "api/api";
-import { agentLogs } from "api/queries/workspaces";
-import type {
-	WorkspaceAgentLifecycle,
-	WorkspaceAgentLog,
-} from "api/typesGenerated";
-import { useEffectEvent } from "hooks/hookPolyfills";
-import { useEffect, useRef } from "react";
-import { useQuery, useQueryClient } from "react-query";
-
-export type UseAgentLogsOptions = Readonly<{
-	workspaceId: string;
-	agentId: string;
-	agentLifeCycleState: WorkspaceAgentLifecycle;
-	enabled?: boolean;
-}>;
-
-/**
- * Defines a custom hook that gives you all workspace agent logs for a given
- * workspace.Depending on the status of the workspace, all logs may or may not
- * be available.
- */
-export function useAgentLogs(
-	options: UseAgentLogsOptions,
-): readonly WorkspaceAgentLog[] | undefined {
-	const { workspaceId, agentId, agentLifeCycleState, enabled = true } = options;
-	const queryClient = useQueryClient();
-	const queryOptions = agentLogs(workspaceId, agentId);
-	const { data: logs, isFetched } = useQuery({ ...queryOptions, enabled });
-
-	// Track the ID of the last log received when the initial logs response comes
-	// back. If the logs are not complete, the ID will mark the start point of the
-	// Web sockets response so that the remaining logs can be received over time
-	const lastQueriedLogId = useRef(0);
-	useEffect(() => {
-		const isAlreadyTracking = lastQueriedLogId.current !== 0;
-		if (isAlreadyTracking) {
-			return;
-		}
-
-		const lastLog = logs?.at(-1);
-		if (lastLog !== undefined) {
-			lastQueriedLogId.current = lastLog.id;
-		}
-	}, [logs]);
-
-	const addLogs = useEffectEvent((newLogs: WorkspaceAgentLog[]) => {
-		queryClient.setQueryData(
-			queryOptions.queryKey,
-			(oldData: WorkspaceAgentLog[] = []) => [...oldData, ...newLogs],
-		);
-	});
-
-	useEffect(() => {
-		// Stream data only for new logs. Old logs should be loaded beforehand
-		// using a regular fetch to avoid overloading the websocket with all
-		// logs at once.
-		if (!isFetched) {
-			return;
-		}
-
-		// If the agent is off, we don't need to stream logs. This is the only state
-		// where the Coder API can't receive logs for the agent from third-party
-		// apps like envbuilder.
-		if (agentLifeCycleState === "off") {
-			return;
-		}
-
-		const socket = watchWorkspaceAgentLogs(agentId, {
-			after: lastQueriedLogId.current,
-			onMessage: (newLogs) => {
-				// Prevent new logs getting added when a connection is not open
-				if (socket.readyState !== WebSocket.OPEN) {
-					return;
-				}
-				addLogs(newLogs);
-			},
-			onError: (error) => {
-				// For some reason Firefox and Safari throw an error when a websocket
-				// connection is close in the middle of a message and because of that we
-				// can't safely show to the users an error message since most of the
-				// time they are just internal stuff. This does not happen to Chrome at
-				// all and I tried to find better way to "soft close" a WS connection on
-				// those browsers without success.
-				console.error(error);
-			},
-		});
-
-		return () => {
-			socket.close();
-		};
-	}, [addLogs, agentId, agentLifeCycleState, isFetched]);
-
-	return logs;
-}

site/src/modules/resources/AgentRow.tsx

@@ -31,7 +31,6 @@ import { AgentDevcontainerCard } from "./AgentDevcontainerCard";
 import { AgentLatency } from "./AgentLatency";
 import { AGENT_LOG_LINE_HEIGHT } from "./AgentLogs/AgentLogLine";
 import { AgentLogs } from "./AgentLogs/AgentLogs";
-import { useAgentLogs } from "./AgentLogs/useAgentLogs";
 import { AgentMetadata } from "./AgentMetadata";
 import { AgentStatus } from "./AgentStatus";
 import { AgentVersion } from "./AgentVersion";
@@ -41,6 +40,7 @@ import { PortForwardButton } from "./PortForwardButton";
 import { AgentSSHButton } from "./SSHButton/SSHButton";
 import { TerminalLink } from "./TerminalLink/TerminalLink";
 import { VSCodeDesktopButton } from "./VSCodeDesktopButton/VSCodeDesktopButton";
+import { useAgentLogs } from "./useAgentLogs";
 
 export interface AgentRowProps {
 	agent: WorkspaceAgent;
@@ -89,12 +89,7 @@ export const AgentRow: FC<AgentRowProps> = ({
 		["starting", "start_timeout"].includes(agent.lifecycle_state) &&
 			hasStartupFeatures,
 	);
-	const agentLogs = useAgentLogs({
-		workspaceId: workspace.id,
-		agentId: agent.id,
-		agentLifeCycleState: agent.lifecycle_state,
-		enabled: showLogs,
-	});
+	const agentLogs = useAgentLogs(agent, showLogs);
 	const logListRef = useRef<List>(null);
 	const logListDivRef = useRef<HTMLDivElement>(null);
 	const startupLogs = useMemo(() => {

site/src/modules/resources/DownloadAgentLogsButton.stories.tsx

@@ -15,7 +15,7 @@ const meta: Meta<typeof DownloadAgentLogsButton> = {
 	parameters: {
 		queries: [
 			{
-				key: agentLogsKey(MockWorkspace.id, MockWorkspaceAgent.id),
+				key: agentLogsKey(MockWorkspaceAgent.id),
 				data: generateLogs(5),
 			},
 		],

site/src/modules/resources/DownloadAgentLogsButton.tsx

@@ -23,7 +23,7 @@ export const DownloadAgentLogsButton: FC<DownloadAgentLogsButtonProps> = ({
 	const [isDownloading, setIsDownloading] = useState(false);
 
 	const fetchLogs = async () => {
-		const queryOpts = agentLogs(workspaceId, agent.id);
+		const queryOpts = agentLogs(agent.id);
 		let logs = queryClient.getQueryData<WorkspaceAgentLog[]>(
 			queryOpts.queryKey,
 		);

site/src/modules/resources/useAgentLogs.test.ts

@@ -0,0 +1,54 @@
+import { renderHook } 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";
+
+/**
+ * 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("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)));
+		expect(result.current).toHaveLength(3);
+
+		// Disable the hook
+		rerender({ enabled: false });
+		expect(result.current).toHaveLength(0);
+
+		// Enable the hook again
+		rerender({ enabled: true });
+		await server.connected;
+		server.send(JSON.stringify(generateLogs(3)));
+		expect(result.current).toHaveLength(3);
+	});
+});
+
+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: "",
+	}));
+}

site/src/modules/resources/useAgentLogs.ts

@@ -0,0 +1,47 @@
+import { watchWorkspaceAgentLogs } from "api/api";
+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.
+			setLogs([]);
+			return;
+		}
+
+		// 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);
+				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;
+}

site/src/modules/workspaces/WorkspaceMoreActions/DownloadLogsDialog.stories.tsx

@@ -20,7 +20,7 @@ const meta: Meta<typeof DownloadLogsDialog> = {
 				data: generateLogs(200),
 			},
 			{
-				key: agentLogsKey(MockWorkspace.id, MockWorkspaceAgent.id),
+				key: agentLogsKey(MockWorkspaceAgent.id),
 				data: generateLogs(400),
 			},
 		],
@@ -41,7 +41,7 @@ export const Loading: Story = {
 				data: undefined,
 			},
 			{
-				key: agentLogsKey(MockWorkspace.id, MockWorkspaceAgent.id),
+				key: agentLogsKey(MockWorkspaceAgent.id),
 				data: undefined,
 			},
 		],

site/src/modules/workspaces/WorkspaceMoreActions/DownloadLogsDialog.tsx

@@ -53,7 +53,7 @@ export const DownloadLogsDialog: FC<DownloadLogsDialogProps> = ({
 
 	const agentLogQueries = useQueries({
 		queries: allUniqueAgents.map((agent) => ({
-			...agentLogs(workspace.id, agent.id),
+			...agentLogs(agent.id),
 			enabled: open,
 		})),
 	});

site/src/pages/WorkspaceBuildPage/WorkspaceBuildPageView.tsx

@@ -21,7 +21,7 @@ import { useSearchParamsKey } from "hooks/useSearchParamsKey";
 import { BuildAvatar } from "modules/builds/BuildAvatar/BuildAvatar";
 import { DashboardFullPage } from "modules/dashboard/DashboardLayout";
 import { AgentLogs } from "modules/resources/AgentLogs/AgentLogs";
-import { useAgentLogs } from "modules/resources/AgentLogs/useAgentLogs";
+import { useAgentLogs } from "modules/resources/useAgentLogs";
 import {
 	WorkspaceBuildData,
 	WorkspaceBuildDataSkeleton,
@@ -212,13 +212,9 @@ export const WorkspaceBuildPageView: FC<WorkspaceBuildPageViewProps> = ({
 						</Alert>
 					)}
 
-					{tabState.value === "build" ? (
-						<BuildLogsContent logs={logs} />
-					) : (
-						<AgentLogsContent
-							workspaceId={build.workspace_id}
-							agent={selectedAgent!}
-						/>
+					{tabState.value === "build" && <BuildLogsContent logs={logs} />}
+					{tabState.value !== "build" && selectedAgent && (
+						<AgentLogsContent agent={selectedAgent} />
 					)}
 				</ScrollArea>
 			</div>
@@ -286,15 +282,12 @@ const BuildLogsContent: FC<{ logs?: ProvisionerJobLog[] }> = ({ logs }) => {
 	);
 };
 
-const AgentLogsContent: FC<{ workspaceId: string; agent: WorkspaceAgent }> = ({
-	agent,
-	workspaceId,
-}) => {
-	const logs = useAgentLogs({
-		workspaceId,
-		agentId: agent.id,
-		agentLifeCycleState: agent.lifecycle_state,
-	});
+type AgentLogsContentProps = {
+	agent: WorkspaceAgent;
+};
+
+const AgentLogsContent: FC<AgentLogsContentProps> = ({ agent }) => {
+	const logs = useAgentLogs(agent, true);
 
 	if (!logs) {
 		return <Loader />;

site/src/pages/WorkspacePage/WorkspaceActions/WorkspaceActions.stories.tsx

@@ -222,7 +222,7 @@ export const OpenDownloadLogs: Story = {
 				data: generateLogs(200),
 			},
 			{
-				key: agentLogsKey(Mocks.MockWorkspace.id, Mocks.MockWorkspaceAgent.id),
+				key: agentLogsKey(Mocks.MockWorkspaceAgent.id),
 				data: generateLogs(400),
 			},
 		],

site/src/pages/WorkspacePage/WorkspacePage.test.tsx

@@ -50,12 +50,7 @@ const renderWorkspacePage = async (
 	jest
 		.spyOn(API, "getDeploymentConfig")
 		.mockResolvedValueOnce(MockDeploymentConfig);
-	jest
-		.spyOn(apiModule, "watchWorkspaceAgentLogs")
-		.mockImplementation((_, options) => {
-			options.onDone?.();
-			return new WebSocket("");
-		});
+	jest.spyOn(apiModule, "watchWorkspaceAgentLogs");
 
 	renderWithAuth(<WorkspacePage />, {
 		...options,

site/src/testHelpers/storybook.tsx

@@ -75,6 +75,8 @@ export const withWebSocket = (Story: FC, { parameters }: StoryContext) => {
 	let callEventsDelay: number;
 
 	window.WebSocket = class WebSocket {
+		public readyState = 1;
+
 		addEventListener(type: string, callback: CallbackFn) {
 			listeners.set(type, callback);
 
@@ -93,6 +95,8 @@ export const withWebSocket = (Story: FC, { parameters }: StoryContext) => {
 			}, 0);
 		}
 
+		removeEventListener(type: string, callback: CallbackFn) {}
+
 		close() {}
 	} as unknown as typeof window.WebSocket;