feat(site): use websocket connection for devcontainer updates (#18808)

Instead of polling every 10 seconds, we instead use a WebSocket
connection for more timely updates.

site/src/api/api.ts

@@ -129,6 +129,14 @@ export const watchWorkspace = (
 	});
 };
 
+export const watchAgentContainers = (
+	agentId: string,
+): OneWayWebSocket<TypesGen.WorkspaceAgentListContainersResponse> => {
+	return new OneWayWebSocket({
+		apiRoute: `/api/v2/workspaceagents/${agentId}/containers/watch`,
+	});
+};
+
 type WatchInboxNotificationsParams = Readonly<{
 	read_status?: "read" | "unread" | "all";
 }>;

site/src/modules/resources/AgentDevcontainerCard.tsx

@@ -130,12 +130,6 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
 
 			return { previousData };
 		},
-		onSuccess: async () => {
-			// Invalidate the containers query to refetch updated data.
-			await queryClient.invalidateQueries({
-				queryKey: ["agents", parentAgent.id, "containers"],
-			});
-		},
 		onError: (error, _, context) => {
 			// If the mutation fails, use the context returned from
 			// onMutate to roll back.

site/src/modules/resources/AgentRow.tsx

@@ -2,14 +2,12 @@ import type { Interpolation, Theme } from "@emotion/react";
 import Collapse from "@mui/material/Collapse";
 import Divider from "@mui/material/Divider";
 import Skeleton from "@mui/material/Skeleton";
-import { API } from "api/api";
 import type {
 	Template,
 	Workspace,
 	WorkspaceAgent,
 	WorkspaceAgentMetadata,
 } from "api/typesGenerated";
-import { isAxiosError } from "axios";
 import { Button } from "components/Button/Button";
 import { DropdownArrow } from "components/DropdownArrow/DropdownArrow";
 import { Stack } from "components/Stack/Stack";
@@ -25,7 +23,6 @@ import {
 	useRef,
 	useState,
 } from "react";
-import { useQuery } from "react-query";
 import AutoSizer from "react-virtualized-auto-sizer";
 import type { FixedSizeList as List, ListOnScrollProps } from "react-window";
 import { AgentApps, organizeAgentApps } from "./AgentApps/AgentApps";
@@ -41,6 +38,7 @@ import { PortForwardButton } from "./PortForwardButton";
 import { AgentSSHButton } from "./SSHButton/SSHButton";
 import { TerminalLink } from "./TerminalLink/TerminalLink";
 import { VSCodeDesktopButton } from "./VSCodeDesktopButton/VSCodeDesktopButton";
+import { useAgentContainers } from "./useAgentContainers";
 import { useAgentLogs } from "./useAgentLogs";
 
 interface AgentRowProps {
@@ -133,20 +131,7 @@ export const AgentRow: FC<AgentRowProps> = ({
 		setBottomOfLogs(distanceFromBottom < AGENT_LOG_LINE_HEIGHT);
 	}, []);
 
-	const { data: devcontainers } = useQuery({
-		queryKey: ["agents", agent.id, "containers"],
-		queryFn: () => API.getAgentContainers(agent.id),
-		enabled: agent.status === "connected",
-		select: (res) => res.devcontainers,
-		// TODO: Implement a websocket connection to get updates on containers
-		// without having to poll.
-		refetchInterval: ({ state }) => {
-			const { error } = state;
-			return isAxiosError(error) && error.response?.status === 403
-				? false
-				: 10_000;
-		},
-	});
+	const devcontainers = useAgentContainers(agent);
 
 	// This is used to show the parent apps of the devcontainer.
 	const [showParentApps, setShowParentApps] = useState(false);

site/src/modules/resources/useAgentContainers.test.tsx

@@ -0,0 +1,196 @@
+import { renderHook, waitFor } from "@testing-library/react";
+import * as API from "api/api";
+import type { WorkspaceAgentListContainersResponse } from "api/typesGenerated";
+import * as GlobalSnackbar from "components/GlobalSnackbar/utils";
+import { http, HttpResponse } from "msw";
+import type { FC, PropsWithChildren } from "react";
+import { QueryClient, QueryClientProvider } from "react-query";
+import {
+	MockWorkspaceAgent,
+	MockWorkspaceAgentDevcontainer,
+} from "testHelpers/entities";
+import { server } from "testHelpers/server";
+import type { OneWayWebSocket } from "utils/OneWayWebSocket";
+import { useAgentContainers } from "./useAgentContainers";
+
+const createWrapper = (): FC<PropsWithChildren> => {
+	const queryClient = new QueryClient({
+		defaultOptions: {
+			queries: {
+				retry: false,
+			},
+		},
+	});
+	return ({ children }) => (
+		<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+	);
+};
+
+describe("useAgentContainers", () => {
+	it("returns containers when agent is connected", async () => {
+		server.use(
+			http.get(
+				`/api/v2/workspaceagents/${MockWorkspaceAgent.id}/containers`,
+				() => {
+					return HttpResponse.json({
+						devcontainers: [MockWorkspaceAgentDevcontainer],
+						containers: [],
+					});
+				},
+			),
+		);
+
+		const { result } = renderHook(
+			() => useAgentContainers(MockWorkspaceAgent),
+			{
+				wrapper: createWrapper(),
+			},
+		);
+
+		await waitFor(() => {
+			expect(result.current).toEqual([MockWorkspaceAgentDevcontainer]);
+		});
+	});
+
+	it("returns undefined when agent is not connected", () => {
+		const disconnectedAgent = {
+			...MockWorkspaceAgent,
+			status: "disconnected" as const,
+		};
+
+		const { result } = renderHook(() => useAgentContainers(disconnectedAgent), {
+			wrapper: createWrapper(),
+		});
+
+		expect(result.current).toBeUndefined();
+	});
+
+	it("handles API errors gracefully", async () => {
+		server.use(
+			http.get(
+				`/api/v2/workspaceagents/${MockWorkspaceAgent.id}/containers`,
+				() => {
+					return HttpResponse.error();
+				},
+			),
+		);
+
+		const { result } = renderHook(
+			() => useAgentContainers(MockWorkspaceAgent),
+			{
+				wrapper: createWrapper(),
+			},
+		);
+
+		await waitFor(() => {
+			expect(result.current).toBeUndefined();
+		});
+	});
+
+	it("handles parsing errors from WebSocket", async () => {
+		const displayErrorSpy = jest.spyOn(GlobalSnackbar, "displayError");
+		const watchAgentContainersSpy = jest.spyOn(API, "watchAgentContainers");
+
+		const mockSocket = {
+			addEventListener: jest.fn(),
+			close: jest.fn(),
+		};
+		watchAgentContainersSpy.mockReturnValue(
+			mockSocket as unknown as OneWayWebSocket<WorkspaceAgentListContainersResponse>,
+		);
+
+		server.use(
+			http.get(
+				`/api/v2/workspaceagents/${MockWorkspaceAgent.id}/containers`,
+				() => {
+					return HttpResponse.json({
+						devcontainers: [MockWorkspaceAgentDevcontainer],
+						containers: [],
+					});
+				},
+			),
+		);
+
+		const { unmount } = renderHook(
+			() => useAgentContainers(MockWorkspaceAgent),
+			{
+				wrapper: createWrapper(),
+			},
+		);
+
+		// Simulate message event with parsing error
+		const messageHandler = mockSocket.addEventListener.mock.calls.find(
+			(call) => call[0] === "message",
+		)?.[1];
+
+		if (messageHandler) {
+			messageHandler({
+				parseError: new Error("Parse error"),
+				parsedMessage: null,
+			});
+		}
+
+		await waitFor(() => {
+			expect(displayErrorSpy).toHaveBeenCalledWith(
+				"Failed to update containers",
+				"Please try refreshing the page",
+			);
+		});
+
+		unmount();
+		displayErrorSpy.mockRestore();
+		watchAgentContainersSpy.mockRestore();
+	});
+
+	it("handles WebSocket errors", async () => {
+		const displayErrorSpy = jest.spyOn(GlobalSnackbar, "displayError");
+		const watchAgentContainersSpy = jest.spyOn(API, "watchAgentContainers");
+
+		const mockSocket = {
+			addEventListener: jest.fn(),
+			close: jest.fn(),
+		};
+		watchAgentContainersSpy.mockReturnValue(
+			mockSocket as unknown as OneWayWebSocket<WorkspaceAgentListContainersResponse>,
+		);
+
+		server.use(
+			http.get(
+				`/api/v2/workspaceagents/${MockWorkspaceAgent.id}/containers`,
+				() => {
+					return HttpResponse.json({
+						devcontainers: [MockWorkspaceAgentDevcontainer],
+						containers: [],
+					});
+				},
+			),
+		);
+
+		const { unmount } = renderHook(
+			() => useAgentContainers(MockWorkspaceAgent),
+			{
+				wrapper: createWrapper(),
+			},
+		);
+
+		// Simulate error event
+		const errorHandler = mockSocket.addEventListener.mock.calls.find(
+			(call) => call[0] === "error",
+		)?.[1];
+
+		if (errorHandler) {
+			errorHandler(new Error("WebSocket error"));
+		}
+
+		await waitFor(() => {
+			expect(displayErrorSpy).toHaveBeenCalledWith(
+				"Failed to load containers",
+				"Please try refreshing the page",
+			);
+		});
+
+		unmount();
+		displayErrorSpy.mockRestore();
+		watchAgentContainersSpy.mockRestore();
+	});
+});

site/src/modules/resources/useAgentContainers.ts

@@ -0,0 +1,59 @@
+import { API, watchAgentContainers } from "api/api";
+import type {
+	WorkspaceAgent,
+	WorkspaceAgentDevcontainer,
+	WorkspaceAgentListContainersResponse,
+} from "api/typesGenerated";
+import { displayError } from "components/GlobalSnackbar/utils";
+import { useEffectEvent } from "hooks/hookPolyfills";
+import { useEffect } from "react";
+import { useQuery, useQueryClient } from "react-query";
+
+export function useAgentContainers(
+	agent: WorkspaceAgent,
+): readonly WorkspaceAgentDevcontainer[] | undefined {
+	const queryClient = useQueryClient();
+
+	const { data: devcontainers } = useQuery({
+		queryKey: ["agents", agent.id, "containers"],
+		queryFn: () => API.getAgentContainers(agent.id),
+		enabled: agent.status === "connected",
+		select: (res) => res.devcontainers,
+		staleTime: Number.POSITIVE_INFINITY,
+	});
+
+	const updateDevcontainersCache = useEffectEvent(
+		async (data: WorkspaceAgentListContainersResponse) => {
+			const queryKey = ["agents", agent.id, "containers"];
+
+			queryClient.setQueryData(queryKey, data);
+		},
+	);
+
+	useEffect(() => {
+		const socket = watchAgentContainers(agent.id);
+
+		socket.addEventListener("message", (event) => {
+			if (event.parseError) {
+				displayError(
+					"Failed to update containers",
+					"Please try refreshing the page",
+				);
+				return;
+			}
+
+			updateDevcontainersCache(event.parsedMessage);
+		});
+
+		socket.addEventListener("error", () => {
+			displayError(
+				"Failed to load containers",
+				"Please try refreshing the page",
+			);
+		});
+
+		return () => socket.close();
+	}, [agent.id, updateDevcontainersCache]);
+
+	return devcontainers;
+}