From b8cfe76b0e8211e73ba8777a5ce40d44eade8321 Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Fri, 7 Mar 2025 23:10:36 +0000 Subject: [PATCH 01/19] feat: add support for one-way websockets to UI --- site/package.json | 7 +- site/pnpm-lock.yaml | 9 +-- site/src/@types/eventsourcemock.d.ts | 1 - site/src/api/api.ts | 34 ++++---- site/src/modules/resources/AgentMetadata.tsx | 62 ++++++++------- .../modules/templates/useWatchVersionLogs.ts | 42 ++++------ .../WorkspacePage/WorkspacePage.test.tsx | 20 +---- .../src/pages/WorkspacePage/WorkspacePage.tsx | 29 ++++--- site/src/utils/OneWayWebSocket.ts | 77 +++++++++++++++++++ 9 files changed, 177 insertions(+), 104 deletions(-) delete mode 100644 site/src/@types/eventsourcemock.d.ts create mode 100644 site/src/utils/OneWayWebSocket.ts diff --git a/site/package.json b/site/package.json index 892e1d50a005f..e7ec0fc8aeb41 100644 --- a/site/package.json +++ b/site/package.json @@ -164,7 +164,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", @@ -189,7 +188,11 @@ "vite-plugin-checker": "0.8.0", "vite-plugin-turbosnap": "1.0.3" }, - "browserslist": ["chrome 110", "firefox 111", "safari 16.0"], + "browserslist": [ + "chrome 110", + "firefox 111", + "safari 16.0" + ], "resolutions": { "optionator": "0.9.3", "semver": "7.6.2" diff --git a/site/pnpm-lock.yaml b/site/pnpm-lock.yaml index 62ae51082e96a..4306e34c37a96 100644 --- a/site/pnpm-lock.yaml +++ b/site/pnpm-lock.yaml @@ -397,9 +397,6 @@ importers: chromatic: specifier: 11.25.2 version: 11.25.2 - eventsourcemock: - specifier: 2.0.0 - version: 2.0.0 express: specifier: 4.21.2 version: 4.21.2 @@ -3694,6 +3691,7 @@ packages: eslint@8.52.0: resolution: {integrity: sha512-zh/JHnaixqHZsolRB/w9/02akBk9EPrOs9JwcTP2ek7yL5bVvXuRariiaAjjoJ5DvuwQ1WAE/HsMz+w17YgBCg==, tarball: https://registry.npmjs.org/eslint/-/eslint-8.52.0.tgz} engines: {node: ^12.22.0 || ^14.17.0 || >=16.0.0} + deprecated: This version is no longer supported. Please see https://eslint.org/version-support for other options. hasBin: true espree@9.6.1: @@ -3737,9 +3735,6 @@ packages: eventemitter3@4.0.7: resolution: {integrity: sha512-8guHBZCwKnFhYdHr2ysuRWErTwhoN2X8XELRlrRwpmfeY2jjuUN4taQMsULKUVo1K4DvZl+0pgfyoysHxvmvEw==, tarball: https://registry.npmjs.org/eventemitter3/-/eventemitter3-4.0.7.tgz} - eventsourcemock@2.0.0: - resolution: {integrity: sha512-tSmJnuE+h6A8/hLRg0usf1yL+Q8w01RQtmg0Uzgoxk/HIPZrIUeAr/A4es/8h1wNsoG8RdiESNQLTKiNwbSC3Q==, tarball: https://registry.npmjs.org/eventsourcemock/-/eventsourcemock-2.0.0.tgz} - execa@5.1.1: resolution: {integrity: sha512-8uSpZZocAZRBAPIEINJj3Lo9HyGitllczc27Eh5YYojjMFMn8yHMDMaUHE2Jqfq05D/wucwI4JGURyXt1vchyg==, tarball: https://registry.npmjs.org/execa/-/execa-5.1.1.tgz} engines: {node: '>=10'} @@ -9908,8 +9903,6 @@ snapshots: eventemitter3@4.0.7: {} - eventsourcemock@2.0.0: {} - execa@5.1.1: dependencies: cross-spawn: 7.0.6 diff --git a/site/src/@types/eventsourcemock.d.ts b/site/src/@types/eventsourcemock.d.ts deleted file mode 100644 index 296c4f19c33ce..0000000000000 --- a/site/src/@types/eventsourcemock.d.ts +++ /dev/null @@ -1 +0,0 @@ -declare module "eventsourcemock"; diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 627ede80976c6..3fc6ca29f900d 100644 --- 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[], @@ -106,10 +107,10 @@ const getMissingParameters = ( * @returns An EventSource that emits agent metadata event objects * (ServerSentEvent) */ -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 => { + const protocol = location.protocol === "https:" ? "wss:" : "ws:"; + return new OneWayWebSocket( + `${protocol}//${location.host}/api/v2/workspaceagents/${agentId}/watch-metadata-ws`, ); }; @@ -117,10 +118,10 @@ export const watchAgentMetadata = (agentId: string): EventSource => { * @returns {EventSource} An EventSource that emits workspace event objects * (ServerSentEvent) */ -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 => { + const protocol = location.protocol === "https:" ? "wss:" : "ws:"; + return new OneWayWebSocket( + `${protocol}//${location.host}/api/v2/workspaces/${workspaceId}/watch-ws`, ); }; @@ -1080,12 +1081,13 @@ class ApiMethods { }; getWorkspaceByOwnerAndName = async ( - username = "me", + username: string | undefined, workspaceName: string, params?: TypesGen.WorkspaceOptions, ): Promise => { + const user = username || "me"; const response = await this.axios.get( - `/api/v2/users/${username}/workspace/${workspaceName}`, + `/api/v2/users/${user}/workspace/${workspaceName}`, { params }, ); @@ -1093,12 +1095,13 @@ class ApiMethods { }; getWorkspaceBuildByNumber = async ( - username = "me", + username: string | undefined, workspaceName: string, buildNumber: number, ): Promise => { + const name = username || "me"; const response = await this.axios.get( - `/api/v2/users/${username}/workspace/${workspaceName}/builds/${buildNumber}`, + `/api/v2/users/${name}/workspace/${workspaceName}/builds/${buildNumber}`, ); return response.data; @@ -1279,11 +1282,12 @@ class ApiMethods { }; createWorkspace = async ( - userId = "me", + userId: string | undefined, workspace: TypesGen.CreateWorkspaceRequest, ): Promise => { + const id = userId || "me"; const response = await this.axios.post( - `/api/v2/users/${userId}/workspaces`, + `/api/v2/users/${id}/workspaces`, workspace, ); diff --git a/site/src/modules/resources/AgentMetadata.tsx b/site/src/modules/resources/AgentMetadata.tsx index 81b5a14994e81..98818abae359a 100644 --- 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"; @@ -46,43 +49,50 @@ export const AgentMetadata: FC = ({ agent, storybookMetadata, }) => { - const [metadata, setMetadata] = useState< - WorkspaceAgentMetadata[] | undefined - >(undefined); - + const [metadata, setMetadata] = useState(); useEffect(() => { + // Even though we're using storybookMetadata as the initial value of the + // `metadata` state, we can't sync on `metadata` itself. If we did, the + // moment we update the state with a new event, we would re-trigger the + // effect and immediately destroy the connection if (storybookMetadata !== undefined) { - setMetadata(storybookMetadata); return; } - let timeout: ReturnType | undefined = undefined; + let timeoutId: number | undefined = undefined; + let latestSocket: OneWayWebSocket | undefined = undefined; - const connect = (): (() => void) => { - const source = watchAgentMetadata(agent.id); + const createNewConnection = () => { + const socket = watchAgentMetadata(agent.id); + latestSocket = socket; - source.onerror = (e) => { - console.error("received error in watch stream", e); + socket.addEventListener("error", () => { + displayError("Socket closed unexpectedly. Creating new connection..."); setMetadata(undefined); - source.close(); - - timeout = setTimeout(() => { - connect(); - }, 3000); - }; - - source.addEventListener("data", (e) => { - const data = JSON.parse(e.data); - setMetadata(data); + timeoutId = window.setTimeout(() => { + createNewConnection(); + }, 3_000); }); - return () => { - if (timeout !== undefined) { - clearTimeout(timeout); + + socket.addEventListener("message", (e) => { + try { + const payload = JSON.parse(e.data) as ServerSentEvent; + if (payload.type === "data") { + setMetadata(payload.data as WorkspaceAgentMetadata[]); + } + } catch { + displayError( + "Unable to process newest response from server. Please try refreshing the page.", + ); } - source.close(); - }; + }); + }; + + createNewConnection(); + return () => { + window.clearTimeout(timeoutId); + latestSocket?.close(); }; - return connect(); }, [agent.id, storybookMetadata]); if (metadata === undefined) { diff --git a/site/src/modules/templates/useWatchVersionLogs.ts b/site/src/modules/templates/useWatchVersionLogs.ts index 5574e083a9849..1e77b0eb1b073 100644 --- a/site/src/modules/templates/useWatchVersionLogs.ts +++ b/site/src/modules/templates/useWatchVersionLogs.ts @@ -1,46 +1,38 @@ import { watchBuildLogsByTemplateVersionId } from "api/api"; import type { ProvisionerJobLog, TemplateVersion } from "api/typesGenerated"; +import { useEffectEvent } from "hooks/hookPolyfills"; import { useEffect, useState } from "react"; export const useWatchVersionLogs = ( templateVersion: TemplateVersion | undefined, options?: { onDone: () => Promise }, ) => { - const [logs, setLogs] = useState(); + const [logs, setLogs] = useState(); const templateVersionId = templateVersion?.id; - const templateVersionStatus = templateVersion?.job.status; + const [cachedVersionId, setCachedVersionId] = useState(templateVersionId); + if (cachedVersionId !== templateVersionId) { + setCachedVersionId(templateVersionId); + setLogs([]); + } - // biome-ignore lint/correctness/useExhaustiveDependencies: consider refactoring + const stableOnDone = useEffectEvent(() => options?.onDone()); + const status = templateVersion?.job.status; + const canWatch = status === "running" || status === "pending"; useEffect(() => { - setLogs(undefined); - }, [templateVersionId]); - - useEffect(() => { - if (!templateVersionId || !templateVersionStatus) { - return; - } - - if ( - templateVersionStatus !== "running" && - templateVersionStatus !== "pending" - ) { + if (!templateVersionId || !canWatch) { return; } const socket = watchBuildLogsByTemplateVersionId(templateVersionId, { - onMessage: (log) => { - setLogs((logs) => (logs ? [...logs, log] : [log])); - }, - onDone: options?.onDone, - onError: (error) => { - console.error(error); + onError: (error) => console.error(error), + onDone: stableOnDone, + onMessage: (newLog) => { + setLogs((current) => [...(current ?? []), newLog]); }, }); - return () => { - socket.close(); - }; - }, [options?.onDone, templateVersionId, templateVersionStatus]); + return () => socket.close(); + }, [stableOnDone, canWatch, templateVersionId]); return logs; }; diff --git a/site/src/pages/WorkspacePage/WorkspacePage.test.tsx b/site/src/pages/WorkspacePage/WorkspacePage.test.tsx index 50f47a4721320..d120ad5546c17 100644 --- a/site/src/pages/WorkspacePage/WorkspacePage.test.tsx +++ b/site/src/pages/WorkspacePage/WorkspacePage.test.tsx @@ -2,7 +2,7 @@ import { screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import * as apiModule from "api/api"; import type { TemplateVersionParameter, Workspace } from "api/typesGenerated"; -import EventSourceMock from "eventsourcemock"; +import MockServerSocket from "jest-websocket-mock"; import { DashboardContext, type DashboardProvider, @@ -84,23 +84,11 @@ const testButton = async ( const user = userEvent.setup(); await user.click(button); - expect(actionMock).toBeCalled(); + expect(actionMock).toHaveBeenCalled(); }; -let originalEventSource: typeof window.EventSource; - -beforeAll(() => { - originalEventSource = window.EventSource; - // mocking out EventSource for SSE - window.EventSource = EventSourceMock; -}); - -beforeEach(() => { - jest.resetAllMocks(); -}); - -afterAll(() => { - window.EventSource = originalEventSource; +afterEach(() => { + MockServerSocket.clean(); }); describe("WorkspacePage", () => { diff --git a/site/src/pages/WorkspacePage/WorkspacePage.tsx b/site/src/pages/WorkspacePage/WorkspacePage.tsx index cd2b5f48cb6d3..578bce9958ab7 100644 --- a/site/src/pages/WorkspacePage/WorkspacePage.tsx +++ b/site/src/pages/WorkspacePage/WorkspacePage.tsx @@ -15,6 +15,7 @@ import { useQuery, useQueryClient } from "react-query"; import { useParams } from "react-router-dom"; import { WorkspaceReadyPage } from "./WorkspaceReadyPage"; import { type WorkspacePermissions, workspaceChecks } from "./permissions"; +import { displayError } from "components/GlobalSnackbar/utils"; export const WorkspacePage: FC = () => { const queryClient = useQueryClient(); @@ -82,20 +83,26 @@ export const WorkspacePage: FC = () => { return; } - const eventSource = watchWorkspace(workspaceId); - - eventSource.addEventListener("data", async (event) => { - const newWorkspaceData = JSON.parse(event.data) as Workspace; - await updateWorkspaceData(newWorkspaceData); + 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 { + displayError( + "Unable to process latest data from the server. Please try refreshing the page.", + ); + } }); - - eventSource.addEventListener("error", (event) => { - console.error("Error on getting workspace changes.", event); + socket.addEventListener("error", () => { + displayError( + "Unable to get workspace changes. Connection has been closed.", + ); }); - return () => { - eventSource.close(); - }; + return () => socket.close(); }, [updateWorkspaceData, workspaceId]); // Page statuses diff --git a/site/src/utils/OneWayWebSocket.ts b/site/src/utils/OneWayWebSocket.ts new file mode 100644 index 0000000000000..87dca9954ea64 --- /dev/null +++ b/site/src/utils/OneWayWebSocket.ts @@ -0,0 +1,77 @@ +/** + * @file A WebSocket that can only receive messages from the server, and cannot + * ever send anything. + * + * 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 + * limitation on HTTP/1.1 and below where there is a maximum number of 6 ports + * that can ever be used for a domain (sometimes less depending on the browser). + * Not only is this limit shared with short-lived REST requests, but it also + * applies across tabs and windows. So if a user opens Coder in multiple tabs, + * there is a very real possibility that parts of the app will start to lock up + * 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. + */ + +// Not bothering with trying to borrow methods from the base WebSocket type +// because it's a mess of inheritance and generics. +type WebSocketEventType = "close" | "error" | "message" | "open"; + +type WebSocketEventPayloadMap = { + close: CloseEvent; + error: Event; + message: MessageEvent; + 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: ( + eventType: T, + callback: (payload: WebSocketEventPayloadMap[T]) => void, + ) => void; + + removeEventListener: ( + eventType: T, + callback: (payload: WebSocketEventPayloadMap[T]) => void, + ) => void; + + close: (closeCode?: number, reason?: string) => void; +} + +// 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 { + #socket: WebSocket; + + constructor(url: string | URL, protocols?: string | string[]) { + 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 = ( + message: T, + callback: (payload: WebSocketEventPayloadMap[T]) => void, + ): void => { + this.#socket.addEventListener(message, callback); + }; + + removeEventListener = ( + message: T, + callback: (payload: WebSocketEventPayloadMap[T]) => void, + ): void => { + this.#socket.removeEventListener(message, callback); + }; + + close = (closeCode?: number, reason?: string): void => { + this.#socket.close(closeCode, reason); + }; +} From 367906de28fbdb2fec5eb99e3a96be03550002e8 Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Fri, 7 Mar 2025 23:15:25 +0000 Subject: [PATCH 02/19] fix: apply formatting --- site/package.json | 6 +----- site/src/pages/WorkspacePage/WorkspacePage.tsx | 2 +- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/site/package.json b/site/package.json index e7ec0fc8aeb41..507469fd4f8b2 100644 --- a/site/package.json +++ b/site/package.json @@ -188,11 +188,7 @@ "vite-plugin-checker": "0.8.0", "vite-plugin-turbosnap": "1.0.3" }, - "browserslist": [ - "chrome 110", - "firefox 111", - "safari 16.0" - ], + "browserslist": ["chrome 110", "firefox 111", "safari 16.0"], "resolutions": { "optionator": "0.9.3", "semver": "7.6.2" diff --git a/site/src/pages/WorkspacePage/WorkspacePage.tsx b/site/src/pages/WorkspacePage/WorkspacePage.tsx index 578bce9958ab7..82b0d2b3dc67f 100644 --- a/site/src/pages/WorkspacePage/WorkspacePage.tsx +++ b/site/src/pages/WorkspacePage/WorkspacePage.tsx @@ -5,6 +5,7 @@ import { workspaceBuildsKey } from "api/queries/workspaceBuilds"; import { workspaceByOwnerAndName } from "api/queries/workspaces"; import type { Workspace } from "api/typesGenerated"; import { ErrorAlert } from "components/Alert/ErrorAlert"; +import { displayError } from "components/GlobalSnackbar/utils"; import { Loader } from "components/Loader/Loader"; import { Margins } from "components/Margins/Margins"; import { useEffectEvent } from "hooks/hookPolyfills"; @@ -15,7 +16,6 @@ import { useQuery, useQueryClient } from "react-query"; import { useParams } from "react-router-dom"; import { WorkspaceReadyPage } from "./WorkspaceReadyPage"; import { type WorkspacePermissions, workspaceChecks } from "./permissions"; -import { displayError } from "components/GlobalSnackbar/utils"; export const WorkspacePage: FC = () => { const queryClient = useQueryClient(); From 09f5e95c876eda1ab5a3b0f9b8d919b9f073b651 Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Fri, 7 Mar 2025 23:21:24 +0000 Subject: [PATCH 03/19] docs: remove outdated comment --- site/src/modules/resources/AgentMetadata.tsx | 4 ---- 1 file changed, 4 deletions(-) diff --git a/site/src/modules/resources/AgentMetadata.tsx b/site/src/modules/resources/AgentMetadata.tsx index 98818abae359a..f245e67943ace 100644 --- a/site/src/modules/resources/AgentMetadata.tsx +++ b/site/src/modules/resources/AgentMetadata.tsx @@ -51,10 +51,6 @@ export const AgentMetadata: FC = ({ }) => { const [metadata, setMetadata] = useState(); useEffect(() => { - // Even though we're using storybookMetadata as the initial value of the - // `metadata` state, we can't sync on `metadata` itself. If we did, the - // moment we update the state with a new event, we would re-trigger the - // effect and immediately destroy the connection if (storybookMetadata !== undefined) { return; } From 04a384654f30b542a0ad9bd572651cbfef356ec2 Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Fri, 7 Mar 2025 23:23:13 +0000 Subject: [PATCH 04/19] fix: add missing clear call --- site/src/modules/resources/AgentMetadata.tsx | 1 + 1 file changed, 1 insertion(+) diff --git a/site/src/modules/resources/AgentMetadata.tsx b/site/src/modules/resources/AgentMetadata.tsx index f245e67943ace..99253e49503ec 100644 --- a/site/src/modules/resources/AgentMetadata.tsx +++ b/site/src/modules/resources/AgentMetadata.tsx @@ -65,6 +65,7 @@ export const AgentMetadata: FC = ({ socket.addEventListener("error", () => { displayError("Socket closed unexpectedly. Creating new connection..."); setMetadata(undefined); + window.clearTimeout(timeoutId); timeoutId = window.setTimeout(() => { createNewConnection(); }, 3_000); From ca8e94f0bfa360f32d66684dda5f901c71b8716d Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Fri, 7 Mar 2025 23:43:36 +0000 Subject: [PATCH 05/19] fix: streamline biome fixes --- site/src/api/api.ts | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 3fc6ca29f900d..550570f51cc9c 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -1081,13 +1081,12 @@ class ApiMethods { }; getWorkspaceByOwnerAndName = async ( - username: string | undefined, + username: string, workspaceName: string, params?: TypesGen.WorkspaceOptions, ): Promise => { - const user = username || "me"; const response = await this.axios.get( - `/api/v2/users/${user}/workspace/${workspaceName}`, + `/api/v2/users/${username}/workspace/${workspaceName}`, { params }, ); @@ -1095,13 +1094,12 @@ class ApiMethods { }; getWorkspaceBuildByNumber = async ( - username: string | undefined, + username: string, workspaceName: string, buildNumber: number, ): Promise => { - const name = username || "me"; const response = await this.axios.get( - `/api/v2/users/${name}/workspace/${workspaceName}/builds/${buildNumber}`, + `/api/v2/users/${username}/workspace/${workspaceName}/builds/${buildNumber}`, ); return response.data; @@ -1282,12 +1280,11 @@ class ApiMethods { }; createWorkspace = async ( - userId: string | undefined, + userId: string, workspace: TypesGen.CreateWorkspaceRequest, ): Promise => { - const id = userId || "me"; const response = await this.axios.post( - `/api/v2/users/${id}/workspaces`, + `/api/v2/users/${userId}/workspaces`, workspace, ); From 81a723a05c8df83d967844cfb22911f81a2ee0cb Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Wed, 12 Mar 2025 16:13:17 +0000 Subject: [PATCH 06/19] fix: resolve Storybook metadata setup bug --- site/src/modules/resources/AgentMetadata.tsx | 43 ++++++++++++++++---- 1 file changed, 34 insertions(+), 9 deletions(-) diff --git a/site/src/modules/resources/AgentMetadata.tsx b/site/src/modules/resources/AgentMetadata.tsx index 99253e49503ec..e87916508b15e 100644 --- a/site/src/modules/resources/AgentMetadata.tsx +++ b/site/src/modules/resources/AgentMetadata.tsx @@ -45,11 +45,13 @@ interface AgentMetadataProps { storybookMetadata?: WorkspaceAgentMetadata[]; } +const maxSocketErrorRetryCount = 3; + export const AgentMetadata: FC = ({ agent, storybookMetadata, }) => { - const [metadata, setMetadata] = useState(); + const [activeMetadata, setActiveMetadata] = useState(storybookMetadata); useEffect(() => { if (storybookMetadata !== undefined) { return; @@ -57,25 +59,41 @@ export const AgentMetadata: FC = ({ let timeoutId: number | undefined = undefined; let latestSocket: OneWayWebSocket | undefined = undefined; + let retries = 0; const createNewConnection = () => { const socket = watchAgentMetadata(agent.id); latestSocket = socket; socket.addEventListener("error", () => { - displayError("Socket closed unexpectedly. Creating new connection..."); - setMetadata(undefined); + setActiveMetadata(undefined); window.clearTimeout(timeoutId); - timeoutId = window.setTimeout(() => { - createNewConnection(); - }, 3_000); + + retries++; + if (retries < maxSocketErrorRetryCount) { + displayError( + "Unexpected disconnect while watching Metadata changes. Creating new connection...", + ); + timeoutId = window.setTimeout(() => { + createNewConnection(); + }, 3_000); + return; + } + + displayError( + "Unexpected disconnect while watching Metadata changes. Cannot connect to server", + ); + // The socket should already be closed by this point, but doing + // this just to be thorough + socket.close(); + latestSocket = undefined; }); socket.addEventListener("message", (e) => { try { const payload = JSON.parse(e.data) as ServerSentEvent; if (payload.type === "data") { - setMetadata(payload.data as WorkspaceAgentMetadata[]); + setActiveMetadata(payload.data as WorkspaceAgentMetadata[]); } } catch { displayError( @@ -90,9 +108,16 @@ export const AgentMetadata: FC = ({ window.clearTimeout(timeoutId); latestSocket?.close(); }; + + // 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 }, [agent.id, storybookMetadata]); - if (metadata === undefined) { + if (activeMetadata === undefined) { return (
@@ -100,7 +125,7 @@ export const AgentMetadata: FC = ({ ); } - return ; + return ; }; export const AgentMetadataSkeleton: FC = () => { From 4364a3d7bf5adcc250b6718c10b8adc680e04f0c Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Wed, 12 Mar 2025 16:14:46 +0000 Subject: [PATCH 07/19] docs: make warning more obvious --- site/src/modules/resources/AgentMetadata.tsx | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/site/src/modules/resources/AgentMetadata.tsx b/site/src/modules/resources/AgentMetadata.tsx index e87916508b15e..43429d1de2853 100644 --- a/site/src/modules/resources/AgentMetadata.tsx +++ b/site/src/modules/resources/AgentMetadata.tsx @@ -53,6 +53,12 @@ export const AgentMetadata: FC = ({ }) => { 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) { return; } @@ -108,13 +114,6 @@ export const AgentMetadata: FC = ({ window.clearTimeout(timeoutId); latestSocket?.close(); }; - - // 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 }, [agent.id, storybookMetadata]); if (activeMetadata === undefined) { From ecb29401b5f9547031fa84d49b0b25b2fccb949a Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Wed, 12 Mar 2025 16:20:18 +0000 Subject: [PATCH 08/19] fix: beef up socket retry logic --- site/src/modules/resources/AgentMetadata.tsx | 29 +++++++++++--------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/site/src/modules/resources/AgentMetadata.tsx b/site/src/modules/resources/AgentMetadata.tsx index 43429d1de2853..2f586b266c0fc 100644 --- a/site/src/modules/resources/AgentMetadata.tsx +++ b/site/src/modules/resources/AgentMetadata.tsx @@ -64,35 +64,38 @@ export const AgentMetadata: FC = ({ } let timeoutId: number | undefined = undefined; - let latestSocket: OneWayWebSocket | undefined = undefined; + let activeSocket: OneWayWebSocket | undefined = undefined; let retries = 0; const createNewConnection = () => { const socket = watchAgentMetadata(agent.id); - latestSocket = socket; + 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 = undefined; + retries++; - if (retries < maxSocketErrorRetryCount) { + if (retries >= maxSocketErrorRetryCount) { displayError( - "Unexpected disconnect while watching Metadata changes. Creating new connection...", + "Unexpected disconnect while watching Metadata changes. Please try refreshing the page.", ); - timeoutId = window.setTimeout(() => { - createNewConnection(); - }, 3_000); return; } displayError( - "Unexpected disconnect while watching Metadata changes. Cannot connect to server", + "Unexpected disconnect while watching Metadata changes. Creating new connection...", ); - // The socket should already be closed by this point, but doing - // this just to be thorough - socket.close(); - latestSocket = undefined; + timeoutId = window.setTimeout(() => { + createNewConnection(); + }, 3_000); }); socket.addEventListener("message", (e) => { @@ -112,7 +115,7 @@ export const AgentMetadata: FC = ({ createNewConnection(); return () => { window.clearTimeout(timeoutId); - latestSocket?.close(); + activeSocket?.close(); }; }, [agent.id, storybookMetadata]); From bfe4d9f5ad79959d0ce18b178786add144f9f3ee Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Wed, 12 Mar 2025 18:16:40 +0000 Subject: [PATCH 09/19] fix: make it harder to initialize OWWS --- site/src/api/api.ts | 14 ++++++-------- site/src/utils/OneWayWebSocket.ts | 12 +++++++++++- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 550570f51cc9c..24d918325ec54 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -108,10 +108,9 @@ const getMissingParameters = ( * (ServerSentEvent) */ export const watchAgentMetadata = (agentId: string): OneWayWebSocket => { - const protocol = location.protocol === "https:" ? "wss:" : "ws:"; - return new OneWayWebSocket( - `${protocol}//${location.host}/api/v2/workspaceagents/${agentId}/watch-metadata-ws`, - ); + return new OneWayWebSocket({ + apiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata-ws`, + }); }; /** @@ -119,10 +118,9 @@ export const watchAgentMetadata = (agentId: string): OneWayWebSocket => { * (ServerSentEvent) */ export const watchWorkspace = (workspaceId: string): OneWayWebSocket => { - const protocol = location.protocol === "https:" ? "wss:" : "ws:"; - return new OneWayWebSocket( - `${protocol}//${location.host}/api/v2/workspaces/${workspaceId}/watch-ws`, - ); + return new OneWayWebSocket({ + apiRoute: `/api/v2/workspaces/${workspaceId}/watch-ws`, + }); }; export const getURLWithSearchParams = ( diff --git a/site/src/utils/OneWayWebSocket.ts b/site/src/utils/OneWayWebSocket.ts index 87dca9954ea64..29eaa6b801562 100644 --- a/site/src/utils/OneWayWebSocket.ts +++ b/site/src/utils/OneWayWebSocket.ts @@ -44,13 +44,23 @@ interface OneWayWebSocketApi { close: (closeCode?: number, reason?: string) => void; } +type OneWayWebSocketInit = Readonly<{ + 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 { #socket: WebSocket; - constructor(url: string | URL, protocols?: string | string[]) { + constructor(init: OneWayWebSocketInit) { + const { apiRoute, protocols, location = window.location } = init; + + const protocol = location.protocol === "https:" ? "wss:" : "ws:"; + const url = `${protocol}//${location.host}${apiRoute}`; this.#socket = new WebSocket(url, protocols); } From 6cdfc21ed97a48d0ccaa28541b9154bab9370c2b Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Wed, 12 Mar 2025 20:04:49 +0000 Subject: [PATCH 10/19] fix: apply feedback --- site/src/api/api.ts | 8 +- site/src/modules/resources/AgentMetadata.tsx | 17 +- .../src/pages/WorkspacePage/WorkspacePage.tsx | 12 +- site/src/utils/OneWayWebSocket.ts | 171 +++++++++++++----- 4 files changed, 146 insertions(+), 62 deletions(-) diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 24d918325ec54..b64844e47327b 100644 --- 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 => { 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 => { 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 index 2f586b266c0fc..5e5501809ee49 100644 --- a/site/src/modules/resources/AgentMetadata.tsx +++ b/site/src/modules/resources/AgentMetadata.tsx @@ -64,7 +64,7 @@ export const AgentMetadata: FC = ({ } let timeoutId: number | undefined = undefined; - let activeSocket: OneWayWebSocket | undefined = undefined; + let activeSocket: OneWayWebSocket | null = null; let retries = 0; const createNewConnection = () => { @@ -80,7 +80,7 @@ export const AgentMetadata: FC = ({ // 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 = ({ }); 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 index 82b0d2b3dc67f..a55971abfb576 100644 --- 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 index 29eaa6b801562..572433d2c8ea5 100644 --- 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 = Readonly< + | { + sourceEvent: MessageEvent; + parsedMessage: TData; + parseError: undefined; + } + | { + sourceEvent: MessageEvent; + parsedMessage: undefined; + parseError: Error; + } +>; + +type OneWayEventPayloadMap = { close: CloseEvent; error: Event; - message: MessageEvent; + message: OneWayMessageEvent; 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: ( - eventType: T, - callback: (payload: WebSocketEventPayloadMap[T]) => void, - ) => void; +type WebSocketMessageCallback = (payload: MessageEvent) => void; - removeEventListener: ( - eventType: T, - callback: (payload: WebSocketEventPayloadMap[T]) => void, - ) => void; - - close: (closeCode?: number, reason?: string) => void; -} +type OneWayEventCallback = ( + payload: OneWayEventPayloadMap[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 { + addEventListener: ( + eventType: TEvent, + callback: OneWayEventCallback, + ) => void; + + removeEventListener: ( + eventType: TEvent, + callback: OneWayEventCallback, + ) => void; + + close: (closeCode?: number, reason?: string) => void; +} + +export class OneWayWebSocket + implements OneWayWebSocketApi +{ #socket: WebSocket; + #messageCallbackWrappers = new Map< + OneWayEventCallback, + 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 = ( - message: T, - callback: (payload: WebSocketEventPayloadMap[T]) => void, - ): void => { - this.#socket.addEventListener(message, callback); - }; - - removeEventListener = ( - message: T, - callback: (payload: WebSocketEventPayloadMap[T]) => void, - ): void => { - this.#socket.removeEventListener(message, callback); - }; - - close = (closeCode?: number, reason?: string): void => { + addEventListener( + event: TEvent, + callback: OneWayEventCallback, + ): 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): 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( + event: TEvent, + callback: OneWayEventCallback, + ): 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); - }; + } } From 682e2f481844453915758d70d7059ffd540a44a7 Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Wed, 12 Mar 2025 20:07:44 +0000 Subject: [PATCH 11/19] fix: update JSDoc --- site/src/api/api.ts | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/site/src/api/api.ts b/site/src/api/api.ts index b64844e47327b..426d3a8539d7b 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -104,8 +104,7 @@ 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, @@ -116,8 +115,7 @@ export const watchAgentMetadata = ( }; /** - * @returns {EventSource} An EventSource that emits workspace event objects - * (ServerSentEvent) + * @returns {OneWayWebSocket} A OneWayWebSocket that emits Server-Sent Events. */ export const watchWorkspace = ( workspaceId: string, From 9b19ceb66365047e0cec7701b9833ba93953d079 Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Tue, 18 Mar 2025 23:19:59 +0000 Subject: [PATCH 12/19] chore: add missing socket unit tests --- site/src/utils/OneWayWebSocket.test.ts | 467 +++++++++++++++++++++++++ site/src/utils/OneWayWebSocket.ts | 66 +++- 2 files changed, 516 insertions(+), 17 deletions(-) create mode 100644 site/src/utils/OneWayWebSocket.test.ts diff --git a/site/src/utils/OneWayWebSocket.test.ts b/site/src/utils/OneWayWebSocket.test.ts new file mode 100644 index 0000000000000..753f8451708d3 --- /dev/null +++ b/site/src/utils/OneWayWebSocket.test.ts @@ -0,0 +1,467 @@ +/** + * @file Sets up unit tests for OneWayWebSocket. + * + * 2025-03-18 - Really wanted to define these as integration tests with MSW, but + * getting it set up correctly for Jest and JSDOM got a little screwy. That can + * be revisited in the future, but in the meantime, we're assuming that the base + * WebSocket class doesn't have any bugs, and can safely be mocked out. + */ + +import { + type OneWayMessageEvent, + type WebSocketEventType, + OneWayWebSocket, +} from "./OneWayWebSocket"; + +type MockSocket = WebSocket & { + publishMessage: (event: MessageEvent) => void; + publishError: (event: ErrorEvent) => void; + publishClose: (event: CloseEvent) => void; + publishOpen: (event: Event) => void; +}; + +function createMockWebSocket( + url: string, + protocols?: string | string[], +): MockSocket { + type EventMap = { + message: MessageEvent; + error: ErrorEvent; + close: CloseEvent; + open: Event; + }; + type CallbackStore = { + [K in keyof EventMap]: ((event: EventMap[K]) => void)[]; + }; + + let activeProtocol: string; + if (Array.isArray(protocols)) { + activeProtocol = protocols[0] ?? ""; + } else if (typeof protocols === "string") { + activeProtocol = protocols; + } else { + activeProtocol = ""; + } + + let closed = false; + const store: CallbackStore = { + message: [], + error: [], + close: [], + open: [], + }; + + return { + CONNECTING: 0, + OPEN: 1, + CLOSING: 2, + CLOSED: 3, + + url, + protocol: activeProtocol, + readyState: 1, + binaryType: "blob", + bufferedAmount: 0, + extensions: "", + onclose: null, + onerror: null, + onmessage: null, + onopen: null, + send: jest.fn(), + dispatchEvent: jest.fn(), + + addEventListener: ( + eventType: E, + callback: WebSocketEventMap[E], + ) => { + if (closed) { + return; + } + + const subscribers = store[eventType]; + const cb = callback as unknown as CallbackStore[E][0]; + if (!subscribers.includes(cb)) { + subscribers.push(cb); + } + }, + + removeEventListener: ( + eventType: E, + callback: WebSocketEventMap[E], + ) => { + if (closed) { + return; + } + + const subscribers = store[eventType]; + const cb = callback as unknown as CallbackStore[E][0]; + if (subscribers.includes(cb)) { + const updated = store[eventType].filter((c) => c !== cb); + store[eventType] = updated as unknown as CallbackStore[E]; + } + }, + + close: () => { + closed = true; + }, + + publishOpen: (event) => { + for (const sub of store.open) { + sub(event); + } + }, + + publishError: (event) => { + for (const sub of store.error) { + sub(event); + } + }, + + publishMessage: (event) => { + for (const sub of store.message) { + sub(event); + } + }, + + publishClose: (event) => { + for (const sub of store.close) { + sub(event); + } + }, + }; +} + +/** + * Qualities to test: + * 1. Lets a consumer to add an event listener of each type + * 2. Lets a consumer to add multiple event listeners of each type + * 3. Registering a callback that is already registered will do nothing + * 4. Surfaces a parse error if a message is not formatted as JSON + * 4. Lets a consumer remove event listeners of each type + * 5. Closing the socket renders it inert + */ +describe(OneWayWebSocket.name, () => { + const dummyRoute = "/api/v2/blah"; + + it("Errors out if API route does not start with '/api/v2/'", () => { + const testRoutes: string[] = ["blah", "", "/", "/api", "/api/v225"]; + + for (const r of testRoutes) { + expect(() => { + new OneWayWebSocket({ + apiRoute: r, + websocketInit: createMockWebSocket, + }); + }).toThrow(Error); + } + }); + + it("Lets a consumer add an event listener of each type", () => { + let mock!: MockSocket; + const oneWay = new OneWayWebSocket({ + apiRoute: dummyRoute, + websocketInit: (url, protocols) => { + const socket = createMockWebSocket(url, protocols); + mock = socket; + return socket; + }, + }); + + const onOpen = jest.fn(); + const onClose = jest.fn(); + const onError = jest.fn(); + const onMessage = jest.fn(); + + oneWay.addEventListener("open", onOpen); + oneWay.addEventListener("close", onClose); + oneWay.addEventListener("error", onError); + oneWay.addEventListener("message", onMessage); + + mock.publishOpen(new Event("open")); + mock.publishClose(new CloseEvent("close")); + mock.publishError( + new ErrorEvent("error", { + error: new Error("Whoops - connection broke"), + }), + ); + mock.publishMessage( + new MessageEvent("message", { + data: "null", + }), + ); + + expect(onOpen).toHaveBeenCalledTimes(1); + expect(onClose).toHaveBeenCalledTimes(1); + expect(onError).toHaveBeenCalledTimes(1); + expect(onMessage).toHaveBeenCalledTimes(1); + }); + + it("Lets a consumer remove an event listener of each type", () => { + let mock!: MockSocket; + const oneWay = new OneWayWebSocket({ + apiRoute: dummyRoute, + websocketInit: (url, protocols) => { + const socket = createMockWebSocket(url, protocols); + mock = socket; + return socket; + }, + }); + + const onOpen = jest.fn(); + const onClose = jest.fn(); + const onError = jest.fn(); + const onMessage = jest.fn(); + + oneWay.addEventListener("open", onOpen); + oneWay.addEventListener("close", onClose); + oneWay.addEventListener("error", onError); + oneWay.addEventListener("message", onMessage); + + oneWay.removeEventListener("open", onOpen); + oneWay.removeEventListener("close", onClose); + oneWay.removeEventListener("error", onError); + oneWay.removeEventListener("message", onMessage); + + mock.publishOpen(new Event("open")); + mock.publishClose(new CloseEvent("close")); + mock.publishError( + new ErrorEvent("error", { + error: new Error("Whoops - connection broke"), + }), + ); + mock.publishMessage( + new MessageEvent("message", { + data: "null", + }), + ); + + expect(onOpen).toHaveBeenCalledTimes(0); + expect(onClose).toHaveBeenCalledTimes(0); + expect(onError).toHaveBeenCalledTimes(0); + expect(onMessage).toHaveBeenCalledTimes(0); + }); + + it("Only calls each callback once if callback is added multiple times", () => { + let mock!: MockSocket; + const oneWay = new OneWayWebSocket({ + apiRoute: dummyRoute, + websocketInit: (url, protocols) => { + const socket = createMockWebSocket(url, protocols); + mock = socket; + return socket; + }, + }); + + const onOpen = jest.fn(); + const onClose = jest.fn(); + const onError = jest.fn(); + const onMessage = jest.fn(); + + for (let i = 0; i < 10; i++) { + oneWay.addEventListener("open", onOpen); + oneWay.addEventListener("close", onClose); + oneWay.addEventListener("error", onError); + oneWay.addEventListener("message", onMessage); + } + + mock.publishOpen(new Event("open")); + mock.publishClose(new CloseEvent("close")); + mock.publishError( + new ErrorEvent("error", { + error: new Error("Whoops - connection broke"), + }), + ); + mock.publishMessage( + new MessageEvent("message", { + data: "null", + }), + ); + + expect(onOpen).toHaveBeenCalledTimes(1); + expect(onClose).toHaveBeenCalledTimes(1); + expect(onError).toHaveBeenCalledTimes(1); + expect(onMessage).toHaveBeenCalledTimes(1); + }); + + it("Lets consumers register multiple callbacks for each event type", () => { + let mock!: MockSocket; + const oneWay = new OneWayWebSocket({ + apiRoute: dummyRoute, + websocketInit: (url, protocols) => { + const socket = createMockWebSocket(url, protocols); + mock = socket; + return socket; + }, + }); + + const onOpen1 = jest.fn(); + const onClose1 = jest.fn(); + const onError1 = jest.fn(); + const onMessage1 = jest.fn(); + oneWay.addEventListener("open", onOpen1); + oneWay.addEventListener("close", onClose1); + oneWay.addEventListener("error", onError1); + oneWay.addEventListener("message", onMessage1); + + const onOpen2 = jest.fn(); + const onClose2 = jest.fn(); + const onError2 = jest.fn(); + const onMessage2 = jest.fn(); + oneWay.addEventListener("open", onOpen2); + oneWay.addEventListener("close", onClose2); + oneWay.addEventListener("error", onError2); + oneWay.addEventListener("message", onMessage2); + + mock.publishOpen(new Event("open")); + mock.publishClose(new CloseEvent("close")); + mock.publishError( + new ErrorEvent("error", { + error: new Error("Whoops - connection broke"), + }), + ); + mock.publishMessage( + new MessageEvent("message", { + data: "null", + }), + ); + + expect(onOpen1).toHaveBeenCalledTimes(1); + expect(onClose1).toHaveBeenCalledTimes(1); + expect(onError1).toHaveBeenCalledTimes(1); + expect(onMessage1).toHaveBeenCalledTimes(1); + + expect(onOpen2).toHaveBeenCalledTimes(1); + expect(onClose2).toHaveBeenCalledTimes(1); + expect(onError2).toHaveBeenCalledTimes(1); + expect(onMessage2).toHaveBeenCalledTimes(1); + }); + + it("Computes the socket protocol based on the browser location protocol", () => { + const oneWay1 = new OneWayWebSocket({ + apiRoute: dummyRoute, + websocketInit: createMockWebSocket, + location: { + protocol: "https:", + host: "www.cool.com", + }, + }); + const oneWay2 = new OneWayWebSocket({ + apiRoute: dummyRoute, + websocketInit: createMockWebSocket, + location: { + protocol: "http:", + host: "www.cool.com", + }, + }); + + expect(oneWay1.url).toMatch(/^wss:\/\//); + expect(oneWay2.url).toMatch(/^ws:\/\//); + }); + + it("Gives consumers pre-parsed versions of message events", () => { + let mock!: MockSocket; + const oneWay = new OneWayWebSocket({ + apiRoute: dummyRoute, + websocketInit: (url, protocols) => { + const socket = createMockWebSocket(url, protocols); + mock = socket; + return socket; + }, + }); + + const onMessage = jest.fn(); + oneWay.addEventListener("message", onMessage); + + const payload = { + value: 5, + cool: "yes", + }; + const event = new MessageEvent("message", { + data: JSON.stringify(payload), + }); + + mock.publishMessage(event); + expect(onMessage).toHaveBeenCalledWith({ + sourceEvent: event, + parsedMessage: payload, + parseError: undefined, + }); + }); + + it("Exposes parsing error if message payload could not be parsed as JSON", () => { + let mock!: MockSocket; + const oneWay = new OneWayWebSocket({ + apiRoute: dummyRoute, + websocketInit: (url, protocols) => { + const socket = createMockWebSocket(url, protocols); + mock = socket; + return socket; + }, + }); + + const onMessage = jest.fn(); + oneWay.addEventListener("message", onMessage); + + const payload = "definitely not valid JSON"; + const event = new MessageEvent("message", { + data: payload, + }); + mock.publishMessage(event); + + const arg: OneWayMessageEvent = onMessage.mock.lastCall[0]; + expect(arg.sourceEvent).toEqual(event); + expect(arg.parsedMessage).toEqual(undefined); + expect(arg.parseError).toBeInstanceOf(Error); + }); + + it("Passes all search param values through Websocket URL", () => { + const input1: Record = { + cool: "yeah", + yeah: "cool", + blah: "5", + }; + const oneWay1 = new OneWayWebSocket({ + apiRoute: dummyRoute, + websocketInit: createMockWebSocket, + searchParams: input1, + location: { + protocol: "https:", + host: "www.blah.com", + }, + }); + let [base, params] = oneWay1.url.split("?"); + expect(base).toBe("wss://www.blah.com/api/v2/blah"); + for (const [key, value] of Object.entries(input1)) { + expect(params).toContain(`${key}=${value}`); + } + + const input2 = new URLSearchParams(input1); + const oneWay2 = new OneWayWebSocket({ + apiRoute: dummyRoute, + websocketInit: createMockWebSocket, + searchParams: input2, + location: { + protocol: "https:", + host: "www.blah.com", + }, + }); + [base, params] = oneWay2.url.split("?"); + expect(base).toBe("wss://www.blah.com/api/v2/blah"); + for (const [key, value] of Object.entries(input2)) { + expect(params).toContain(`${key}=${value}`); + } + + const oneWay3 = new OneWayWebSocket({ + apiRoute: dummyRoute, + websocketInit: createMockWebSocket, + searchParams: undefined, + location: { + protocol: "https:", + host: "www.blah.com", + }, + }); + [base, params] = oneWay3.url.split("?"); + expect(base).toBe("wss://www.blah.com/api/v2/blah"); + expect(params).toBe(undefined); + }); +}); diff --git a/site/src/utils/OneWayWebSocket.ts b/site/src/utils/OneWayWebSocket.ts index 572433d2c8ea5..7bbfb04a66774 100644 --- a/site/src/utils/OneWayWebSocket.ts +++ b/site/src/utils/OneWayWebSocket.ts @@ -18,9 +18,9 @@ // Not bothering with trying to borrow methods from the base WebSocket type // 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"; +export type WebSocketEventType = "close" | "error" | "message" | "open"; -type OneWayMessageEvent = Readonly< +export type OneWayMessageEvent = Readonly< | { sourceEvent: MessageEvent; parsedMessage: TData; @@ -46,13 +46,9 @@ type OneWayEventCallback = ( payload: OneWayEventPayloadMap[TEvent], ) => void; -type OneWayWebSocketInit = Readonly<{ - apiRoute: string; - location?: Location; - protocols?: string | string[]; -}>; - interface OneWayWebSocketApi { + get url(): string; + addEventListener: ( eventType: TEvent, callback: OneWayEventCallback, @@ -66,6 +62,22 @@ interface OneWayWebSocketApi { close: (closeCode?: number, reason?: string) => void; } +type OneWayWebSocketInit = Readonly<{ + apiRoute: string; + serverProtocols?: string | string[]; + searchParams?: Record | URLSearchParams; + binaryType?: BinaryType; + websocketInit?: (url: string, protocols?: string | string[]) => WebSocket; + location?: Readonly<{ + protocol: string; + host: string; + }>; +}>; + +function defaultInit(url: string, protocols?: string | string[]): WebSocket { + return new WebSocket(url, protocols); +} + export class OneWayWebSocket implements OneWayWebSocketApi { @@ -76,14 +88,34 @@ export class OneWayWebSocket >(); 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 { + apiRoute, + searchParams, + serverProtocols, + binaryType = "blob", + location = window.location, + websocketInit = defaultInit, + } = init; + + if (!apiRoute.startsWith("/api/v2/")) { + throw new Error(`API route '${apiRoute}' does not begin with a slash`); } - const protocol = location.protocol === "https:" ? "wss:" : "ws:"; - const url = `${protocol}//${location.host}${apiRoute}`; - this.#socket = new WebSocket(url, protocols); + const formattedParams = + searchParams instanceof URLSearchParams + ? searchParams + : new URLSearchParams(searchParams); + const paramsString = formattedParams.toString(); + const paramsSuffix = paramsString ? `?${paramsString}` : ""; + const wsProtocol = location.protocol === "https:" ? "wss:" : "ws:"; + const url = `${wsProtocol}//${location.host}${apiRoute}${paramsSuffix}`; + + this.#socket = websocketInit(url, serverProtocols); + this.#socket.binaryType = binaryType; + } + + get url(): string { + return this.#socket.url; } addEventListener( @@ -141,13 +173,13 @@ export class OneWayWebSocket WebSocketEventType >; - if (!this.#messageCallbackWrappers.has(looseCallback)) { - return; - } if (event !== "message") { this.#socket.removeEventListener(event, looseCallback); return; } + if (!this.#messageCallbackWrappers.has(looseCallback)) { + return; + } const wrapper = this.#messageCallbackWrappers.get(looseCallback); if (wrapper === undefined) { From 423910fa65c706e56474b5441128ab30e640302a Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Tue, 18 Mar 2025 23:20:13 +0000 Subject: [PATCH 13/19] fix: update notifications code to use OWWS --- site/src/api/api.ts | 39 +++++------------- .../NotificationsInbox/NotificationsInbox.tsx | 40 ++++++++++++------- 2 files changed, 34 insertions(+), 45 deletions(-) diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 4c97a547474a3..256ee19588f26 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -102,7 +102,6 @@ const getMissingParameters = ( }; /** - * * @param agentId * @returns {OneWayWebSocket} A OneWayWebSocket that emits Server-Sent Events. */ @@ -125,38 +124,18 @@ export const watchWorkspace = ( }); }; -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 { + return new OneWayWebSocket({ + apiRoute: "/api/v2/notifications/inbox/watch", + searchParams: params, }); - - return socket; -}; +} export const getURLWithSearchParams = ( basePath: string, @@ -2513,7 +2492,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 index bf8d3622e35f1..a4c9b776e72e0 100644 --- a/site/src/modules/notifications/NotificationsInbox/NotificationsInbox.tsx +++ b/site/src/modules/notifications/NotificationsInbox/NotificationsInbox.tsx @@ -1,4 +1,4 @@ -import { API, watchInboxNotifications } from "api/api"; +import { watchInboxNotifications } from "api/api"; import { getErrorDetail, getErrorMessage } from "api/errors"; import type { ListInboxNotificationsResponse, @@ -6,7 +6,7 @@ import type { } from "api/typesGenerated"; import { displayError } from "components/GlobalSnackbar/utils"; import { useEffectEvent } from "hooks/hookPolyfills"; -import { type FC, useEffect, useRef } from "react"; +import { type FC, useEffect } from "react"; import { useMutation, useQuery, useQueryClient } from "react-query"; import { InboxPopover } from "./InboxPopover"; @@ -58,21 +58,31 @@ export const NotificationsInbox: FC = ({ ); 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((prev) => { + return { + unread_count: msg.unread_count, + notifications: [msg.notification, ...prev.notifications], + }; + }); + }); + + socket.addEventListener("error", () => { + displayError( + "Unable to retrieve latest inbox notifications. Please try refreshing the browser.", + ); socket.close(); - }; + }); + + return () => socket.close(); }, [updateNotificationsCache]); const markAllAsReadMutation = useMutation({ From 6e3e0d8406d31b56474d7d898fae50df469a094e Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Tue, 18 Mar 2025 23:25:22 +0000 Subject: [PATCH 14/19] fix: remove comment about what to test --- site/src/utils/OneWayWebSocket.test.ts | 9 --------- 1 file changed, 9 deletions(-) diff --git a/site/src/utils/OneWayWebSocket.test.ts b/site/src/utils/OneWayWebSocket.test.ts index 753f8451708d3..9b6110f9b9439 100644 --- a/site/src/utils/OneWayWebSocket.test.ts +++ b/site/src/utils/OneWayWebSocket.test.ts @@ -131,15 +131,6 @@ function createMockWebSocket( }; } -/** - * Qualities to test: - * 1. Lets a consumer to add an event listener of each type - * 2. Lets a consumer to add multiple event listeners of each type - * 3. Registering a callback that is already registered will do nothing - * 4. Surfaces a parse error if a message is not formatted as JSON - * 4. Lets a consumer remove event listeners of each type - * 5. Closing the socket renders it inert - */ describe(OneWayWebSocket.name, () => { const dummyRoute = "/api/v2/blah"; From c42237921f49ecd430da8eb0422fcd894e5565db Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Wed, 19 Mar 2025 17:32:26 +0000 Subject: [PATCH 15/19] fix: make class fields readonly --- site/src/utils/OneWayWebSocket.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/site/src/utils/OneWayWebSocket.ts b/site/src/utils/OneWayWebSocket.ts index 7bbfb04a66774..94ed1f1efc868 100644 --- a/site/src/utils/OneWayWebSocket.ts +++ b/site/src/utils/OneWayWebSocket.ts @@ -81,8 +81,8 @@ function defaultInit(url: string, protocols?: string | string[]): WebSocket { export class OneWayWebSocket implements OneWayWebSocketApi { - #socket: WebSocket; - #messageCallbackWrappers = new Map< + readonly #socket: WebSocket; + readonly #messageCallbackWrappers = new Map< OneWayEventCallback, WebSocketMessageCallback >(); From 0824dd43d1a3b96b7a5a3c6ce436abf1ab674ab3 Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Wed, 19 Mar 2025 18:28:43 +0000 Subject: [PATCH 16/19] fix: sort imports --- site/src/utils/OneWayWebSocket.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/site/src/utils/OneWayWebSocket.test.ts b/site/src/utils/OneWayWebSocket.test.ts index 9b6110f9b9439..61d4ede9110a2 100644 --- a/site/src/utils/OneWayWebSocket.test.ts +++ b/site/src/utils/OneWayWebSocket.test.ts @@ -9,8 +9,8 @@ import { type OneWayMessageEvent, - type WebSocketEventType, OneWayWebSocket, + type WebSocketEventType, } from "./OneWayWebSocket"; type MockSocket = WebSocket & { From c1cee57f0e8803768ea19f4f5493ecf1da6e5578 Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Tue, 25 Mar 2025 19:14:13 +0000 Subject: [PATCH 17/19] refactor: make tests more maintainable --- .../NotificationsInbox/NotificationsInbox.tsx | 4 +- site/src/utils/OneWayWebSocket.test.ts | 116 +++++++++++------- 2 files changed, 73 insertions(+), 47 deletions(-) diff --git a/site/src/modules/notifications/NotificationsInbox/NotificationsInbox.tsx b/site/src/modules/notifications/NotificationsInbox/NotificationsInbox.tsx index acaa9b0803d0a..cdbf0941b7fdb 100644 --- a/site/src/modules/notifications/NotificationsInbox/NotificationsInbox.tsx +++ b/site/src/modules/notifications/NotificationsInbox/NotificationsInbox.tsx @@ -70,10 +70,10 @@ export const NotificationsInbox: FC = ({ } const msg = e.parsedMessage; - updateNotificationsCache((prev) => { + updateNotificationsCache((current) => { return { unread_count: msg.unread_count, - notifications: [msg.notification, ...prev.notifications], + notifications: [msg.notification, ...current.notifications], }; }); }); diff --git a/site/src/utils/OneWayWebSocket.test.ts b/site/src/utils/OneWayWebSocket.test.ts index 61d4ede9110a2..3061e4a15e103 100644 --- a/site/src/utils/OneWayWebSocket.test.ts +++ b/site/src/utils/OneWayWebSocket.test.ts @@ -13,17 +13,21 @@ import { type WebSocketEventType, } from "./OneWayWebSocket"; -type MockSocket = WebSocket & { +type MockPublisher = Readonly<{ publishMessage: (event: MessageEvent) => void; publishError: (event: ErrorEvent) => void; publishClose: (event: CloseEvent) => void; publishOpen: (event: Event) => void; +}>; + +type MockSocket = WebSocket & { + publisher: MockPublisher; }; function createMockWebSocket( url: string, protocols?: string | string[], -): MockSocket { +): readonly [WebSocket, MockPublisher] { type EventMap = { message: MessageEvent; error: ErrorEvent; @@ -51,7 +55,7 @@ function createMockWebSocket( open: [], }; - return { + const mockSocket: WebSocket = { CONNECTING: 0, OPEN: 1, CLOSING: 2, @@ -104,7 +108,9 @@ function createMockWebSocket( close: () => { closed = true; }, + }; + const publisher: MockPublisher = { publishOpen: (event) => { for (const sub of store.open) { sub(event); @@ -129,6 +135,8 @@ function createMockWebSocket( } }, }; + + return [mockSocket, publisher] as const; } describe(OneWayWebSocket.name, () => { @@ -141,19 +149,22 @@ describe(OneWayWebSocket.name, () => { expect(() => { new OneWayWebSocket({ apiRoute: r, - websocketInit: createMockWebSocket, + websocketInit: (url, protocols) => { + const [socket] = createMockWebSocket(url, protocols); + return socket; + }, }); }).toThrow(Error); } }); it("Lets a consumer add an event listener of each type", () => { - let mock!: MockSocket; + let publisher!: MockPublisher; const oneWay = new OneWayWebSocket({ apiRoute: dummyRoute, websocketInit: (url, protocols) => { - const socket = createMockWebSocket(url, protocols); - mock = socket; + const [socket, pub] = createMockWebSocket(url, protocols); + publisher = pub; return socket; }, }); @@ -168,14 +179,14 @@ describe(OneWayWebSocket.name, () => { oneWay.addEventListener("error", onError); oneWay.addEventListener("message", onMessage); - mock.publishOpen(new Event("open")); - mock.publishClose(new CloseEvent("close")); - mock.publishError( + publisher.publishOpen(new Event("open")); + publisher.publishClose(new CloseEvent("close")); + publisher.publishError( new ErrorEvent("error", { error: new Error("Whoops - connection broke"), }), ); - mock.publishMessage( + publisher.publishMessage( new MessageEvent("message", { data: "null", }), @@ -188,12 +199,12 @@ describe(OneWayWebSocket.name, () => { }); it("Lets a consumer remove an event listener of each type", () => { - let mock!: MockSocket; + let publisher!: MockPublisher; const oneWay = new OneWayWebSocket({ apiRoute: dummyRoute, websocketInit: (url, protocols) => { - const socket = createMockWebSocket(url, protocols); - mock = socket; + const [socket, pub] = createMockWebSocket(url, protocols); + publisher = pub; return socket; }, }); @@ -213,14 +224,14 @@ describe(OneWayWebSocket.name, () => { oneWay.removeEventListener("error", onError); oneWay.removeEventListener("message", onMessage); - mock.publishOpen(new Event("open")); - mock.publishClose(new CloseEvent("close")); - mock.publishError( + publisher.publishOpen(new Event("open")); + publisher.publishClose(new CloseEvent("close")); + publisher.publishError( new ErrorEvent("error", { error: new Error("Whoops - connection broke"), }), ); - mock.publishMessage( + publisher.publishMessage( new MessageEvent("message", { data: "null", }), @@ -233,12 +244,12 @@ describe(OneWayWebSocket.name, () => { }); it("Only calls each callback once if callback is added multiple times", () => { - let mock!: MockSocket; + let publisher!: MockPublisher; const oneWay = new OneWayWebSocket({ apiRoute: dummyRoute, websocketInit: (url, protocols) => { - const socket = createMockWebSocket(url, protocols); - mock = socket; + const [socket, pub] = createMockWebSocket(url, protocols); + publisher = pub; return socket; }, }); @@ -255,14 +266,14 @@ describe(OneWayWebSocket.name, () => { oneWay.addEventListener("message", onMessage); } - mock.publishOpen(new Event("open")); - mock.publishClose(new CloseEvent("close")); - mock.publishError( + publisher.publishOpen(new Event("open")); + publisher.publishClose(new CloseEvent("close")); + publisher.publishError( new ErrorEvent("error", { error: new Error("Whoops - connection broke"), }), ); - mock.publishMessage( + publisher.publishMessage( new MessageEvent("message", { data: "null", }), @@ -275,12 +286,12 @@ describe(OneWayWebSocket.name, () => { }); it("Lets consumers register multiple callbacks for each event type", () => { - let mock!: MockSocket; + let publisher!: MockPublisher; const oneWay = new OneWayWebSocket({ apiRoute: dummyRoute, websocketInit: (url, protocols) => { - const socket = createMockWebSocket(url, protocols); - mock = socket; + const [socket, pub] = createMockWebSocket(url, protocols); + publisher = pub; return socket; }, }); @@ -303,14 +314,14 @@ describe(OneWayWebSocket.name, () => { oneWay.addEventListener("error", onError2); oneWay.addEventListener("message", onMessage2); - mock.publishOpen(new Event("open")); - mock.publishClose(new CloseEvent("close")); - mock.publishError( + publisher.publishOpen(new Event("open")); + publisher.publishClose(new CloseEvent("close")); + publisher.publishError( new ErrorEvent("error", { error: new Error("Whoops - connection broke"), }), ); - mock.publishMessage( + publisher.publishMessage( new MessageEvent("message", { data: "null", }), @@ -330,7 +341,10 @@ describe(OneWayWebSocket.name, () => { it("Computes the socket protocol based on the browser location protocol", () => { const oneWay1 = new OneWayWebSocket({ apiRoute: dummyRoute, - websocketInit: createMockWebSocket, + websocketInit: (url, protocols) => { + const [socket] = createMockWebSocket(url, protocols); + return socket; + }, location: { protocol: "https:", host: "www.cool.com", @@ -338,7 +352,10 @@ describe(OneWayWebSocket.name, () => { }); const oneWay2 = new OneWayWebSocket({ apiRoute: dummyRoute, - websocketInit: createMockWebSocket, + websocketInit: (url, protocols) => { + const [socket] = createMockWebSocket(url, protocols); + return socket; + }, location: { protocol: "http:", host: "www.cool.com", @@ -350,12 +367,12 @@ describe(OneWayWebSocket.name, () => { }); it("Gives consumers pre-parsed versions of message events", () => { - let mock!: MockSocket; + let publisher!: MockPublisher; const oneWay = new OneWayWebSocket({ apiRoute: dummyRoute, websocketInit: (url, protocols) => { - const socket = createMockWebSocket(url, protocols); - mock = socket; + const [socket, pub] = createMockWebSocket(url, protocols); + publisher = pub; return socket; }, }); @@ -371,7 +388,7 @@ describe(OneWayWebSocket.name, () => { data: JSON.stringify(payload), }); - mock.publishMessage(event); + publisher.publishMessage(event); expect(onMessage).toHaveBeenCalledWith({ sourceEvent: event, parsedMessage: payload, @@ -380,12 +397,12 @@ describe(OneWayWebSocket.name, () => { }); it("Exposes parsing error if message payload could not be parsed as JSON", () => { - let mock!: MockSocket; + let publisher!: MockPublisher; const oneWay = new OneWayWebSocket({ apiRoute: dummyRoute, websocketInit: (url, protocols) => { - const socket = createMockWebSocket(url, protocols); - mock = socket; + const [socket, pub] = createMockWebSocket(url, protocols); + publisher = pub; return socket; }, }); @@ -397,7 +414,7 @@ describe(OneWayWebSocket.name, () => { const event = new MessageEvent("message", { data: payload, }); - mock.publishMessage(event); + publisher.publishMessage(event); const arg: OneWayMessageEvent = onMessage.mock.lastCall[0]; expect(arg.sourceEvent).toEqual(event); @@ -413,7 +430,10 @@ describe(OneWayWebSocket.name, () => { }; const oneWay1 = new OneWayWebSocket({ apiRoute: dummyRoute, - websocketInit: createMockWebSocket, + websocketInit: (url, protocols) => { + const [socket] = createMockWebSocket(url, protocols); + return socket; + }, searchParams: input1, location: { protocol: "https:", @@ -429,7 +449,10 @@ describe(OneWayWebSocket.name, () => { const input2 = new URLSearchParams(input1); const oneWay2 = new OneWayWebSocket({ apiRoute: dummyRoute, - websocketInit: createMockWebSocket, + websocketInit: (url, protocols) => { + const [socket] = createMockWebSocket(url, protocols); + return socket; + }, searchParams: input2, location: { protocol: "https:", @@ -444,7 +467,10 @@ describe(OneWayWebSocket.name, () => { const oneWay3 = new OneWayWebSocket({ apiRoute: dummyRoute, - websocketInit: createMockWebSocket, + websocketInit: (url, protocols) => { + const [socket] = createMockWebSocket(url, protocols); + return socket; + }, searchParams: undefined, location: { protocol: "https:", From 70b74e203036472241b238e6dc87b21a8de6688d Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Tue, 25 Mar 2025 19:15:43 +0000 Subject: [PATCH 18/19] fix: remove unused type alias --- site/src/utils/OneWayWebSocket.test.ts | 4 ---- 1 file changed, 4 deletions(-) diff --git a/site/src/utils/OneWayWebSocket.test.ts b/site/src/utils/OneWayWebSocket.test.ts index 3061e4a15e103..f3ad586690f16 100644 --- a/site/src/utils/OneWayWebSocket.test.ts +++ b/site/src/utils/OneWayWebSocket.test.ts @@ -20,10 +20,6 @@ type MockPublisher = Readonly<{ publishOpen: (event: Event) => void; }>; -type MockSocket = WebSocket & { - publisher: MockPublisher; -}; - function createMockWebSocket( url: string, protocols?: string | string[], From 8e34e915bef8553ebe44abf840702900193ae242 Mon Sep 17 00:00:00 2001 From: Michael Smith Date: Tue, 25 Mar 2025 19:19:43 +0000 Subject: [PATCH 19/19] fix: make mock publisher more robust --- site/src/utils/OneWayWebSocket.test.ts | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/site/src/utils/OneWayWebSocket.test.ts b/site/src/utils/OneWayWebSocket.test.ts index f3ad586690f16..c6b00b593111f 100644 --- a/site/src/utils/OneWayWebSocket.test.ts +++ b/site/src/utils/OneWayWebSocket.test.ts @@ -108,24 +108,36 @@ function createMockWebSocket( const publisher: MockPublisher = { publishOpen: (event) => { + if (closed) { + return; + } for (const sub of store.open) { sub(event); } }, publishError: (event) => { + if (closed) { + return; + } for (const sub of store.error) { sub(event); } }, publishMessage: (event) => { + if (closed) { + return; + } for (const sub of store.message) { sub(event); } }, publishClose: (event) => { + if (closed) { + return; + } for (const sub of store.close) { sub(event); }