Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 24 commits into from
Mar 28, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
b8cfe76
feat: add support for one-way websockets to UI
Parkreiner Mar 7, 2025
367906d
fix: apply formatting
Parkreiner Mar 7, 2025
09f5e95
docs: remove outdated comment
Parkreiner Mar 7, 2025
04a3846
fix: add missing clear call
Parkreiner Mar 7, 2025
ca8e94f
fix: streamline biome fixes
Parkreiner Mar 7, 2025
81a723a
fix: resolve Storybook metadata setup bug
Parkreiner Mar 12, 2025
4364a3d
docs: make warning more obvious
Parkreiner Mar 12, 2025
ecb2940
fix: beef up socket retry logic
Parkreiner Mar 12, 2025
bfe4d9f
fix: make it harder to initialize OWWS
Parkreiner Mar 12, 2025
6cdfc21
fix: apply feedback
Parkreiner Mar 12, 2025
682e2f4
fix: update JSDoc
Parkreiner Mar 12, 2025
20ad778
Merge branch 'main' into mes/one-way-ws-02
Parkreiner Mar 18, 2025
9b19ceb
chore: add missing socket unit tests
Parkreiner Mar 18, 2025
423910f
fix: update notifications code to use OWWS
Parkreiner Mar 18, 2025
247dbb6
Merge branch 'mes/one-way-ws-01' into mes/one-way-ws-02
Parkreiner Mar 18, 2025
6e3e0d8
fix: remove comment about what to test
Parkreiner Mar 18, 2025
c422379
fix: make class fields readonly
Parkreiner Mar 19, 2025
0824dd4
fix: sort imports
Parkreiner Mar 19, 2025
db448d7
Merge branch 'mes/one-way-ws-01' into mes/one-way-ws-02
Parkreiner Mar 19, 2025
60bf505
Merge branch 'mes/one-way-ws-01' into mes/one-way-ws-02
Parkreiner Mar 25, 2025
c1cee57
refactor: make tests more maintainable
Parkreiner Mar 25, 2025
70b74e2
fix: remove unused type alias
Parkreiner Mar 25, 2025
8e34e91
fix: make mock publisher more robust
Parkreiner Mar 25, 2025
8db068a
Merge branch 'mes/one-way-ws-01' into mes/one-way-ws-02
Parkreiner Mar 28, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion site/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
8 changes: 0 additions & 8 deletions site/pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 0 additions & 1 deletion site/src/@types/eventsourcemock.d.ts

This file was deleted.

76 changes: 28 additions & 48 deletions site/src/api/api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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[],
Expand Down Expand Up @@ -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> => {
Copy link
Member Author

@Parkreiner Parkreiner Mar 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sadly, this is the best we can do as far as giving ourselves type-safety (without updating GUTS). It'll be less of an issue for WebSocket connections that don't send SSE-formatted events, but the ServerSentEvent type is structured like this:

// From codersdk/serversentevents.go
export interface ServerSentEvent {
	readonly type: ServerSentEventType;
	// empty interface{} type, falling back to unknown
	readonly data: unknown;
}

There's no way to pass a type parameter to the type and make data more specific

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,
Expand Down Expand Up @@ -1125,7 +1105,7 @@ class ApiMethods {
};

getWorkspaceByOwnerAndName = async (
username = "me",
username: string,
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Biome was complaining about having default parameters not be at the end of the function signature, but also, we were never calling these functions with an explicit value of undefined

workspaceName: string,
params?: TypesGen.WorkspaceOptions,
): Promise<TypesGen.Workspace> => {
Expand All @@ -1138,7 +1118,7 @@ class ApiMethods {
};

getWorkspaceBuildByNumber = async (
username = "me",
username: string,
workspaceName: string,
buildNumber: number,
): Promise<TypesGen.WorkspaceBuild> => {
Expand Down Expand Up @@ -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>(
Expand Down Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
93 changes: 64 additions & 29 deletions site/src/modules/resources/AgentMetadata.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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";

Expand All @@ -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);
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was removed because it introduces unnecessary renders – better to just use storybookMetadata as the initial value for the state and avoid a bunch of state syncs

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 = () => {
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The previous create function was super weird because it was responsible for three things

  • Creating a new WebSocket connection from the outside effect
  • Being able to create new connections by calling itself recursively
  • Also acting as a useEffect cleanup function

Figured it'd be better to split up those responsibilities

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 = () => {
Expand Down
Loading
Loading