fix: apply feedback

Parkreiner · Parkreiner · commit 6cdfc21ed97a · 2025-03-12T20:04:49.000Z
diff --git a/site/src/api/api.ts b/site/src/api/api.ts
@@ -107,7 +107,9 @@ const getMissingParameters = (
  * @returns An EventSource that emits agent metadata event objects
  * (ServerSentEvent)
  */
-export const watchAgentMetadata = (agentId: string): OneWayWebSocket => {
+export const watchAgentMetadata = (
+	agentId: string,
+): OneWayWebSocket<TypesGen.ServerSentEvent> => {
 	return new OneWayWebSocket({
 		apiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata-ws`,
 	});
@@ -117,7 +119,9 @@ export const watchAgentMetadata = (agentId: string): OneWayWebSocket => {
  * @returns {EventSource} An EventSource that emits workspace event objects
  * (ServerSentEvent)
  */
-export const watchWorkspace = (workspaceId: string): OneWayWebSocket => {
+export const watchWorkspace = (
+	workspaceId: string,
+): OneWayWebSocket<TypesGen.ServerSentEvent> => {
 	return new OneWayWebSocket({
 		apiRoute: `/api/v2/workspaces/${workspaceId}/watch-ws`,
 	});
diff --git a/site/src/modules/resources/AgentMetadata.tsx b/site/src/modules/resources/AgentMetadata.tsx
@@ -64,7 +64,7 @@ export const AgentMetadata: FC<AgentMetadataProps> = ({
 		}
 
 		let timeoutId: number | undefined = undefined;
-		let activeSocket: OneWayWebSocket | undefined = undefined;
+		let activeSocket: OneWayWebSocket<ServerSentEvent> | null = null;
 		let retries = 0;
 
 		const createNewConnection = () => {
@@ -80,7 +80,7 @@ export const AgentMetadata: FC<AgentMetadataProps> = ({
 				// would auto-close. Couldn't find a definitive answer on MDN,
 				// though, so closing it manually just to be safe
 				socket.close();
-				activeSocket = undefined;
+				activeSocket = null;
 
 				retries++;
 				if (retries >= maxSocketErrorRetryCount) {
@@ -99,15 +99,16 @@ export const AgentMetadata: FC<AgentMetadataProps> = ({
 			});
 
 			socket.addEventListener("message", (e) => {
-				try {
-					const payload = JSON.parse(e.data) as ServerSentEvent;
-					if (payload.type === "data") {
-						setActiveMetadata(payload.data as WorkspaceAgentMetadata[]);
-					}
-				} catch {
+				if (e.parseError) {
 					displayError(
 						"Unable to process newest response from server. Please try refreshing the page.",
 					);
+					return;
+				}
+
+				const msg = e.parsedMessage;
+				if (msg.type === "data") {
+					setActiveMetadata(msg.data as WorkspaceAgentMetadata[]);
 				}
 			});
 		};
diff --git a/site/src/pages/WorkspacePage/WorkspacePage.tsx b/site/src/pages/WorkspacePage/WorkspacePage.tsx
@@ -85,15 +85,15 @@ export const WorkspacePage: FC = () => {
 
 		const socket = watchWorkspace(workspaceId);
 		socket.addEventListener("message", (event) => {
-			try {
-				const sse = JSON.parse(event.data);
-				if (sse.type === "data") {
-					updateWorkspaceData(sse.data as Workspace);
-				}
-			} catch {
+			if (event.parseError) {
 				displayError(
 					"Unable to process latest data from the server. Please try refreshing the page.",
 				);
+				return;
+			}
+
+			if (event.parsedMessage.type === "data") {
+				updateWorkspaceData(event.parsedMessage.data as Workspace);
 			}
 		});
 		socket.addEventListener("error", () => {
diff --git a/site/src/utils/OneWayWebSocket.ts b/site/src/utils/OneWayWebSocket.ts
@@ -1,6 +1,6 @@
 /**
- * @file A WebSocket that can only receive messages from the server, and cannot
- * ever send anything.
+ * @file A wrapper over WebSockets that (1) enforces one-way communication, and
+ * (2) supports automatically parsing JSON messages as they come in.
  *
  * This should ALWAYS be favored in favor of using Server-Sent Events and the
  * built-in EventSource class for doing one-way communication. SSEs have a hard
@@ -12,76 +12,155 @@
  * without it being clear why.
  *
  * WebSockets do not have this limitation, even on HTTP/1.1 – all modern
- * browsers implement at least some degree of multiplexing for them. This file
- * just provides a wrapper to make it harder to use WebSockets for two-way
- * communication by accident.
+ * browsers implement at least some degree of multiplexing for them.
  */
 
 // Not bothering with trying to borrow methods from the base WebSocket type
-// because it's a mess of inheritance and generics.
+// because it's already a mess of inheritance and generics, and we're going to
+// have to add a few more
 type WebSocketEventType = "close" | "error" | "message" | "open";
 
-type WebSocketEventPayloadMap = {
+type OneWayMessageEvent<TData> = Readonly<
+	| {
+			sourceEvent: MessageEvent<string>;
+			parsedMessage: TData;
+			parseError: undefined;
+	  }
+	| {
+			sourceEvent: MessageEvent<string>;
+			parsedMessage: undefined;
+			parseError: Error;
+	  }
+>;
+
+type OneWayEventPayloadMap<TData> = {
 	close: CloseEvent;
 	error: Event;
-	message: MessageEvent<string>;
+	message: OneWayMessageEvent<TData>;
 	open: Event;
 };
 
-// The generics need to apply on a per-method-invocation basis; they cannot be
-// generalized to the top-level interface definition
-interface OneWayWebSocketApi {
-	addEventListener: <T extends WebSocketEventType>(
-		eventType: T,
-		callback: (payload: WebSocketEventPayloadMap[T]) => void,
-	) => void;
+type WebSocketMessageCallback = (payload: MessageEvent<string>) => void;
 
-	removeEventListener: <T extends WebSocketEventType>(
-		eventType: T,
-		callback: (payload: WebSocketEventPayloadMap[T]) => void,
-	) => void;
-
-	close: (closeCode?: number, reason?: string) => void;
-}
+type OneWayEventCallback<TData, TEvent extends WebSocketEventType> = (
+	payload: OneWayEventPayloadMap<TData>[TEvent],
+) => void;
 
 type OneWayWebSocketInit = Readonly<{
-	apiRoute: `/${string}`;
+	apiRoute: string;
 	location?: Location;
 	protocols?: string | string[];
 }>;
 
-// Implementing wrapper around the base WebSocket class instead of doing fancy
-// compile-time type-casts so that we have more runtime assurance that we won't
-// accidentally send a message from the client to the server
-export class OneWayWebSocket implements OneWayWebSocketApi {
+interface OneWayWebSocketApi<TData> {
+	addEventListener: <TEvent extends WebSocketEventType>(
+		eventType: TEvent,
+		callback: OneWayEventCallback<TData, TEvent>,
+	) => void;
+
+	removeEventListener: <TEvent extends WebSocketEventType>(
+		eventType: TEvent,
+		callback: OneWayEventCallback<TData, TEvent>,
+	) => void;
+
+	close: (closeCode?: number, reason?: string) => void;
+}
+
+export class OneWayWebSocket<TData = unknown>
+	implements OneWayWebSocketApi<TData>
+{
 	#socket: WebSocket;
+	#messageCallbackWrappers = new Map<
+		OneWayEventCallback<TData, "message">,
+		WebSocketMessageCallback
+	>();
 
 	constructor(init: OneWayWebSocketInit) {
 		const { apiRoute, protocols, location = window.location } = init;
+		if (apiRoute.at(0) !== "/") {
+			throw new Error(`API route ${apiRoute} does not begin with a space`);
+		}
 
 		const protocol = location.protocol === "https:" ? "wss:" : "ws:";
 		const url = `${protocol}//${location.host}${apiRoute}`;
 		this.#socket = new WebSocket(url, protocols);
 	}
 
-	// Just because this is a React project, all public methods are defined as
-	// arrow functions so they can be passed around as values without losing
-	// their `this` context
-	addEventListener = <T extends WebSocketEventType>(
-		message: T,
-		callback: (payload: WebSocketEventPayloadMap[T]) => void,
-	): void => {
-		this.#socket.addEventListener(message, callback);
-	};
-
-	removeEventListener = <T extends WebSocketEventType>(
-		message: T,
-		callback: (payload: WebSocketEventPayloadMap[T]) => void,
-	): void => {
-		this.#socket.removeEventListener(message, callback);
-	};
-
-	close = (closeCode?: number, reason?: string): void => {
+	addEventListener<TEvent extends WebSocketEventType>(
+		event: TEvent,
+		callback: OneWayEventCallback<TData, TEvent>,
+	): void {
+		// Not happy about all the type assertions, but there are some nasty
+		// type contravariance issues if you try to resolve the function types
+		// properly. This is actually the lesser of two evils
+		const looseCallback = callback as OneWayEventCallback<
+			TData,
+			WebSocketEventType
+		>;
+
+		if (this.#messageCallbackWrappers.has(looseCallback)) {
+			return;
+		}
+		if (event !== "message") {
+			this.#socket.addEventListener(event, looseCallback);
+			return;
+		}
+
+		const wrapped = (event: MessageEvent<string>): void => {
+			const messageCallback = looseCallback as OneWayEventCallback<
+				TData,
+				"message"
+			>;
+
+			try {
+				const message = JSON.parse(event.data) as TData;
+				messageCallback({
+					sourceEvent: event,
+					parseError: undefined,
+					parsedMessage: message,
+				});
+			} catch (err) {
+				messageCallback({
+					sourceEvent: event,
+					parseError: err as Error,
+					parsedMessage: undefined,
+				});
+			}
+		};
+
+		this.#socket.addEventListener(event as "message", wrapped);
+		this.#messageCallbackWrappers.set(looseCallback, wrapped);
+	}
+
+	removeEventListener<TEvent extends WebSocketEventType>(
+		event: TEvent,
+		callback: OneWayEventCallback<TData, TEvent>,
+	): void {
+		const looseCallback = callback as OneWayEventCallback<
+			TData,
+			WebSocketEventType
+		>;
+
+		if (!this.#messageCallbackWrappers.has(looseCallback)) {
+			return;
+		}
+		if (event !== "message") {
+			this.#socket.removeEventListener(event, looseCallback);
+			return;
+		}
+
+		const wrapper = this.#messageCallbackWrappers.get(looseCallback);
+		if (wrapper === undefined) {
+			throw new Error(
+				`Cannot unregister callback for event ${event}. This is likely an issue with the browser itself.`,
+			);
+		}
+
+		this.#socket.removeEventListener(event as "message", wrapper);
+		this.#messageCallbackWrappers.delete(looseCallback);
+	}
+
+	close(closeCode?: number, reason?: string): void {
 		this.#socket.close(closeCode, reason);
-	};
+	}
 }
