chore: add support for one-way WebSockets to UI (#16855)

Closes #16777 ## Changes made - Added `OneWayWebSocket` utility class to help enforce one-way communication from the server to the client - Updated all client client code to use the new WebSocket-based endpoints made to replace the current SSE-based endpoints - Updated WebSocket event handlers to be aware of new protocols - Refactored existing `useEffect` calls and removed some synchronization bugs - Removed dependencies and types for dealing with SSEs - Addressed some minor Biome warnings
coder · Parkreiner · Mar 28, 2025 · Mar 7, 2025 · Mar 7, 2025 · Mar 18, 2025
commit 8f1eca0c91f38987ed0d18c24291c3d493d1ad0e
diff --git a/site/package.json b/site/package.json
@@ -166,7 +166,6 @@
 		"@vitejs/plugin-react": "4.3.4",
 		"autoprefixer": "10.4.20",
 		"chromatic": "11.25.2",
-		"eventsourcemock": "2.0.0",
 		"express": "4.21.2",
 		"jest": "29.7.0",
 		"jest-canvas-mock": "2.5.2",

diff --git a/site/pnpm-lock.yaml b/site/pnpm-lock.yaml
diff --git a/site/src/@types/eventsourcemock.d.ts b/site/src/@types/eventsourcemock.d.ts
diff --git a/site/src/api/api.ts b/site/src/api/api.ts
@@ -22,9 +22,10 @@
 import globalAxios, { type AxiosInstance, isAxiosError } from "axios";
 import type dayjs from "dayjs";
 import userAgentParser from "ua-parser-js";
+import { OneWayWebSocket } from "utils/OneWayWebSocket";
 import { delay } from "../utils/delay";
-import * as TypesGen from "./typesGenerated";
 import type { PostWorkspaceUsageRequest } from "./typesGenerated";
+import * as TypesGen from "./typesGenerated";
 
 const getMissingParameters = (
 	oldBuildParameters: TypesGen.WorkspaceBuildParameter[],
@@ -101,61 +102,40 @@ const getMissingParameters = (
 };
 
 /**
- *
  * @param agentId
- * @returns An EventSource that emits agent metadata event objects
- * (ServerSentEvent)
+ * @returns {OneWayWebSocket} A OneWayWebSocket that emits Server-Sent Events.
  */
-export const watchAgentMetadata = (agentId: string): EventSource => {
-	return new EventSource(
-		`${location.protocol}//${location.host}/api/v2/workspaceagents/${agentId}/watch-metadata`,
-		{ withCredentials: true },
-	);
+export const watchAgentMetadata = (
+	agentId: string,
+): OneWayWebSocket<TypesGen.ServerSentEvent> => {
+	return new OneWayWebSocket({
+		apiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata-ws`,
+	});
 };
 
 /**
- * @returns {EventSource} An EventSource that emits workspace event objects
- * (ServerSentEvent)
+ * @returns {OneWayWebSocket} A OneWayWebSocket that emits Server-Sent Events.
  */
-export const watchWorkspace = (workspaceId: string): EventSource => {
-	return new EventSource(
-		`${location.protocol}//${location.host}/api/v2/workspaces/${workspaceId}/watch`,
-		{ withCredentials: true },
-	);
+export const watchWorkspace = (
+	workspaceId: string,
+): OneWayWebSocket<TypesGen.ServerSentEvent> => {
+	return new OneWayWebSocket({
+		apiRoute: `/api/v2/workspaces/${workspaceId}/watch-ws`,
+	});
 };
 
-type WatchInboxNotificationsParams = {
+type WatchInboxNotificationsParams = Readonly<{
 	read_status?: "read" | "unread" | "all";
-};
+}>;
 
-export const watchInboxNotifications = (
-	onNewNotification: (res: TypesGen.GetInboxNotificationResponse) => void,
+export function watchInboxNotifications(
 	params?: WatchInboxNotificationsParams,
-) => {
-	const searchParams = new URLSearchParams(params);
-	const socket = createWebSocket(
-		"/api/v2/notifications/inbox/watch",
-		searchParams,
-	);
-
-	socket.addEventListener("message", (event) => {
-		try {
-			const res = JSON.parse(
-				event.data,
-			) as TypesGen.GetInboxNotificationResponse;
-			onNewNotification(res);
-		} catch (error) {
-			console.warn("Error parsing inbox notification: ", error);
-		}
-	});
-
-	socket.addEventListener("error", (event) => {
-		console.warn("Watch inbox notifications error: ", event);
-		socket.close();
+): OneWayWebSocket<TypesGen.GetInboxNotificationResponse> {
+	return new OneWayWebSocket({
+		apiRoute: "/api/v2/notifications/inbox/watch",
+		searchParams: params,
 	});
-
-	return socket;
-};
+}
 
 export const getURLWithSearchParams = (
 	basePath: string,
@@ -1125,7 +1105,7 @@ class ApiMethods {
 	};
 
 	getWorkspaceByOwnerAndName = async (
-		username = "me",
+		username: string,
 		workspaceName: string,
 		params?: TypesGen.WorkspaceOptions,
 	): Promise<TypesGen.Workspace> => {
@@ -1138,7 +1118,7 @@ class ApiMethods {
 	};
 
 	getWorkspaceBuildByNumber = async (
-		username = "me",
+		username: string,
 		workspaceName: string,
 		buildNumber: number,
 	): Promise<TypesGen.WorkspaceBuild> => {
@@ -1324,7 +1304,7 @@ class ApiMethods {
 	};
 
 	createWorkspace = async (
-		userId = "me",
+		userId: string,
 		workspace: TypesGen.CreateWorkspaceRequest,
 	): Promise<TypesGen.Workspace> => {
 		const response = await this.axios.post<TypesGen.Workspace>(
@@ -2542,7 +2522,7 @@ function createWebSocket(
 ) {
 	const protocol = location.protocol === "https:" ? "wss:" : "ws:";
 	const socket = new WebSocket(
-		`${protocol}//${location.host}${path}?${params.toString()}`,
+		`${protocol}//${location.host}${path}?${params}`,
 	);
 	socket.binaryType = "blob";
 	return socket;

diff --git a/site/src/modules/notifications/NotificationsInbox/NotificationsInbox.tsx b/site/src/modules/notifications/NotificationsInbox/NotificationsInbox.tsx
@@ -61,21 +61,31 @@ export const NotificationsInbox: FC<NotificationsInboxProps> = ({
 	);
 
 	useEffect(() => {
-		const socket = watchInboxNotifications(
-			(res) => {
-				updateNotificationsCache((prev) => {
-					return {
-						unread_count: res.unread_count,
-						notifications: [res.notification, ...prev.notifications],
-					};
-				});
-			},
-			{ read_status: "unread" },
-		);
+		const socket = watchInboxNotifications({ read_status: "unread" });
 
-		return () => {
+		socket.addEventListener("message", (e) => {
+			if (e.parseError) {
+				console.warn("Error parsing inbox notification: ", e.parseError);
+				return;
+			}
+
+			const msg = e.parsedMessage;
+			updateNotificationsCache((current) => {
+				return {
+					unread_count: msg.unread_count,
+					notifications: [msg.notification, ...current.notifications],
+				};
+			});
+		});
+
+		socket.addEventListener("error", () => {
+			displayError(
+				"Unable to retrieve latest inbox notifications. Please try refreshing the browser.",
+			);
 			socket.close();
-		};
+		});
+
+		return () => socket.close();
 	}, [updateNotificationsCache]);
 
 	const {

diff --git a/site/src/modules/resources/AgentMetadata.tsx b/site/src/modules/resources/AgentMetadata.tsx
@@ -3,9 +3,11 @@ import Skeleton from "@mui/material/Skeleton";
 import Tooltip from "@mui/material/Tooltip";
 import { watchAgentMetadata } from "api/api";
 import type {
+	ServerSentEvent,
 	WorkspaceAgent,
 	WorkspaceAgentMetadata,
 } from "api/typesGenerated";
+import { displayError } from "components/GlobalSnackbar/utils";
 import { Stack } from "components/Stack/Stack";
 import dayjs from "dayjs";
 import {
@@ -17,6 +19,7 @@ import {
 	useState,
 } from "react";
 import { MONOSPACE_FONT_FAMILY } from "theme/constants";
+import type { OneWayWebSocket } from "utils/OneWayWebSocket";
 
 type ItemStatus = "stale" | "valid" | "loading";
 
@@ -42,58 +45,90 @@ interface AgentMetadataProps {
 	storybookMetadata?: WorkspaceAgentMetadata[];
 }
 
+const maxSocketErrorRetryCount = 3;
+
 export const AgentMetadata: FC<AgentMetadataProps> = ({
 	agent,
 	storybookMetadata,
 }) => {
-	const [metadata, setMetadata] = useState<
-		WorkspaceAgentMetadata[] | undefined
-	>(undefined);
-
+	const [activeMetadata, setActiveMetadata] = useState(storybookMetadata);
 	useEffect(() => {
+		// This is an unfortunate pitfall with this component's testing setup,
+		// but even though we use the value of storybookMetadata as the initial
+		// value of the activeMetadata, we cannot put activeMetadata itself into
+		// the dependency array. If we did, we would destroy and rebuild each
+		// connection every single time a new message comes in from the socket,
+		// because the socket has to be wired up to the state setter
 		if (storybookMetadata !== undefined) {
-			setMetadata(storybookMetadata);
 			return;
 		}
 
-		let timeout: ReturnType<typeof setTimeout> | undefined = undefined;
-
-		const connect = (): (() => void) => {
-			const source = watchAgentMetadata(agent.id);
+		let timeoutId: number | undefined = undefined;
+		let activeSocket: OneWayWebSocket<ServerSentEvent> | null = null;
+		let retries = 0;
+
+		const createNewConnection = () => {
+			const socket = watchAgentMetadata(agent.id);
+			activeSocket = socket;
+
+			socket.addEventListener("error", () => {
+				setActiveMetadata(undefined);
+				window.clearTimeout(timeoutId);
+
+				// The error event is supposed to fire when an error happens
+				// with the connection itself, which implies that the connection
+				// would auto-close. Couldn't find a definitive answer on MDN,
+				// though, so closing it manually just to be safe
+				socket.close();
+				activeSocket = null;
+
+				retries++;
+				if (retries >= maxSocketErrorRetryCount) {
+					displayError(
+						"Unexpected disconnect while watching Metadata changes. Please try refreshing the page.",
+					);
+					return;
+				}
 
-			source.onerror = (e) => {
-				console.error("received error in watch stream", e);
-				setMetadata(undefined);
-				source.close();
+				displayError(
+					"Unexpected disconnect while watching Metadata changes. Creating new connection...",
+				);
+				timeoutId = window.setTimeout(() => {
+					createNewConnection();
+				}, 3_000);
+			});
 
-				timeout = setTimeout(() => {
-					connect();
-				}, 3000);
-			};
+			socket.addEventListener("message", (e) => {
+				if (e.parseError) {
+					displayError(
+						"Unable to process newest response from server. Please try refreshing the page.",
+					);
+					return;
+				}
 
-			source.addEventListener("data", (e) => {
-				const data = JSON.parse(e.data);
-				setMetadata(data);
-			});
-			return () => {
-				if (timeout !== undefined) {
-					clearTimeout(timeout);
+				const msg = e.parsedMessage;
+				if (msg.type === "data") {
+					setActiveMetadata(msg.data as WorkspaceAgentMetadata[]);
 				}
-				source.close();
-			};
+			});
+		};
+
+		createNewConnection();
+		return () => {
+			window.clearTimeout(timeoutId);
+			activeSocket?.close();
 		};
-		return connect();
 	}, [agent.id, storybookMetadata]);
 
-	if (metadata === undefined) {
+	if (activeMetadata === undefined) {
 		return (
 			<section css={styles.root}>
 				<AgentMetadataSkeleton />
 			</section>
 		);
 	}
 
-	return <AgentMetadataView metadata={metadata} />;
+	return <AgentMetadataView metadata={activeMetadata} />;
 };
 
 export const AgentMetadataSkeleton: FC = () => {