fix(site): only attempt to watch when dev containers enabled (#18892)

DanielleMaywood · web-flow · commit bfb9aa464d4d · 2025-07-17T00:03:59.000+01:00
diff --git a/site/src/modules/resources/useAgentContainers.test.tsx b/site/src/modules/resources/useAgentContainers.test.tsx
@@ -4,23 +4,19 @@ 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 { act } from "react";
+import { QueryClientProvider } from "react-query";
 import {
 	MockWorkspaceAgent,
 	MockWorkspaceAgentDevcontainer,
 } from "testHelpers/entities";
+import { createTestQueryClient } from "testHelpers/renderHelpers";
 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,
-			},
-		},
-	});
+	const queryClient = createTestQueryClient();
 	return ({ children }) => (
 		<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
 	);
@@ -111,22 +107,29 @@ describe("useAgentContainers", () => {
 			),
 		);
 
-		const { unmount } = renderHook(
+		const { result, unmount } = renderHook(
 			() => useAgentContainers(MockWorkspaceAgent),
 			{
 				wrapper: createWrapper(),
 			},
 		);
 
-		// Simulate message event with parsing error
+		// Wait for initial query to complete
+		await waitFor(() => {
+			expect(result.current).toEqual([MockWorkspaceAgentDevcontainer]);
+		});
+
+		// Now 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,
+			act(() => {
+				messageHandler({
+					parseError: new Error("Parse error"),
+					parsedMessage: null,
+				});
 			});
 		}
 
@@ -166,20 +169,27 @@ describe("useAgentContainers", () => {
 			),
 		);
 
-		const { unmount } = renderHook(
+		const { result, unmount } = renderHook(
 			() => useAgentContainers(MockWorkspaceAgent),
 			{
 				wrapper: createWrapper(),
 			},
 		);
 
-		// Simulate error event
+		// Wait for initial query to complete
+		await waitFor(() => {
+			expect(result.current).toEqual([MockWorkspaceAgentDevcontainer]);
+		});
+
+		// Now simulate error event
 		const errorHandler = mockSocket.addEventListener.mock.calls.find(
 			(call) => call[0] === "error",
 		)?.[1];
 
 		if (errorHandler) {
-			errorHandler(new Error("WebSocket error"));
+			act(() => {
+				errorHandler(new Error("WebSocket error"));
+			});
 		}
 
 		await waitFor(() => {
@@ -211,4 +221,36 @@ describe("useAgentContainers", () => {
 
 		watchAgentContainersSpy.mockRestore();
 	});
+
+	it("does not establish WebSocket connection when dev container feature is not enabled", async () => {
+		const watchAgentContainersSpy = jest.spyOn(API, "watchAgentContainers");
+
+		server.use(
+			http.get(
+				`/api/v2/workspaceagents/${MockWorkspaceAgent.id}/containers`,
+				() => {
+					return HttpResponse.json(
+						{ message: "Dev Container feature not enabled." },
+						{ status: 403 },
+					);
+				},
+			),
+		);
+
+		const { result } = renderHook(
+			() => useAgentContainers(MockWorkspaceAgent),
+			{
+				wrapper: createWrapper(),
+			},
+		);
+
+		// Wait for the query to complete and error to be processed
+		await waitFor(() => {
+			expect(result.current).toBeUndefined();
+		});
+
+		expect(watchAgentContainersSpy).not.toHaveBeenCalled();
+
+		watchAgentContainersSpy.mockRestore();
+	});
 });
diff --git a/site/src/modules/resources/useAgentContainers.ts b/site/src/modules/resources/useAgentContainers.ts
@@ -14,7 +14,11 @@ export function useAgentContainers(
 ): readonly WorkspaceAgentDevcontainer[] | undefined {
 	const queryClient = useQueryClient();
 
-	const { data: devcontainers } = useQuery({
+	const {
+		data: devcontainers,
+		error: queryError,
+		isLoading: queryIsLoading,
+	} = useQuery({
 		queryKey: ["agents", agent.id, "containers"],
 		queryFn: () => API.getAgentContainers(agent.id),
 		enabled: agent.status === "connected",
@@ -31,7 +35,7 @@ export function useAgentContainers(
 	);
 
 	useEffect(() => {
-		if (agent.status !== "connected") {
+		if (agent.status !== "connected" || queryIsLoading || queryError) {
 			return;
 		}
 
@@ -57,7 +61,13 @@ export function useAgentContainers(
 		});
 
 		return () => socket.close();
-	}, [agent.id, agent.status, updateDevcontainersCache]);
+	}, [
+		agent.id,
+		agent.status,
+		queryIsLoading,
+		queryError,
+		updateDevcontainersCache,
+	]);
 
 	return devcontainers;
 }
