Skip to content

Commit 8f1eca0

Browse files
ParkreinerParkreiner
authored
chore: add support for one-way WebSockets to UI (#16855)
Closes #16777 ## Changes made - Added `OneWayWebSocket` utility class to help enforce one-way communication from the server to the client - Updated all client client code to use the new WebSocket-based endpoints made to replace the current SSE-based endpoints - Updated WebSocket event handlers to be aware of new protocols - Refactored existing `useEffect` calls and removed some synchronization bugs - Removed dependencies and types for dealing with SSEs - Addressed some minor Biome warnings
1 parent a567ff4 commit 8f1eca0

File tree

11 files changed

+843
-151
lines changed

11 files changed

+843
-151
lines changed

‎site/package.json

-1
Original file line numberDiff line numberDiff line change
@@ -166,7 +166,6 @@
166166
"@vitejs/plugin-react": "4.3.4",
167167
"autoprefixer": "10.4.20",
168168
"chromatic": "11.25.2",
169-
"eventsourcemock": "2.0.0",
170169
"express": "4.21.2",
171170
"jest": "29.7.0",
172171
"jest-canvas-mock": "2.5.2",

‎site/pnpm-lock.yaml

-8
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎site/src/@types/eventsourcemock.d.ts

-1
This file was deleted.

‎site/src/api/api.ts

+28-48
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,10 @@
2222
import globalAxios, { type AxiosInstance, isAxiosError } from "axios";
2323
import type dayjs from "dayjs";
2424
import userAgentParser from "ua-parser-js";
25+
import { OneWayWebSocket } from "utils/OneWayWebSocket";
2526
import { delay } from "../utils/delay";
26-
import * as TypesGen from "./typesGenerated";
2727
import type { PostWorkspaceUsageRequest } from "./typesGenerated";
28+
import * as TypesGen from "./typesGenerated";
2829

2930
const getMissingParameters = (
3031
oldBuildParameters: TypesGen.WorkspaceBuildParameter[],
@@ -101,61 +102,40 @@ const getMissingParameters = (
101102
};
102103

103104
/**
104-
*
105105
* @param agentId
106-
* @returns An EventSource that emits agent metadata event objects
107-
* (ServerSentEvent)
106+
* @returns {OneWayWebSocket} A OneWayWebSocket that emits Server-Sent Events.
108107
*/
109-
export const watchAgentMetadata = (agentId: string): EventSource => {
110-
return new EventSource(
111-
`${location.protocol}//${location.host}/api/v2/workspaceagents/${agentId}/watch-metadata`,
112-
{ withCredentials: true },
113-
);
108+
export const watchAgentMetadata = (
109+
agentId: string,
110+
): OneWayWebSocket<TypesGen.ServerSentEvent> => {
111+
return new OneWayWebSocket({
112+
apiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata-ws`,
113+
});
114114
};
115115

116116
/**
117-
* @returns {EventSource} An EventSource that emits workspace event objects
118-
* (ServerSentEvent)
117+
* @returns {OneWayWebSocket} A OneWayWebSocket that emits Server-Sent Events.
119118
*/
120-
export const watchWorkspace = (workspaceId: string): EventSource => {
121-
return new EventSource(
122-
`${location.protocol}//${location.host}/api/v2/workspaces/${workspaceId}/watch`,
123-
{ withCredentials: true },
124-
);
119+
export const watchWorkspace = (
120+
workspaceId: string,
121+
): OneWayWebSocket<TypesGen.ServerSentEvent> => {
122+
return new OneWayWebSocket({
123+
apiRoute: `/api/v2/workspaces/${workspaceId}/watch-ws`,
124+
});
125125
};
126126

127-
type WatchInboxNotificationsParams = {
127+
type WatchInboxNotificationsParams = Readonly<{
128128
read_status?: "read" | "unread" | "all";
129-
};
129+
}>;
130130

131-
export const watchInboxNotifications = (
132-
onNewNotification: (res: TypesGen.GetInboxNotificationResponse) => void,
131+
export function watchInboxNotifications(
133132
params?: WatchInboxNotificationsParams,
134-
) => {
135-
const searchParams = new URLSearchParams(params);
136-
const socket = createWebSocket(
137-
"/api/v2/notifications/inbox/watch",
138-
searchParams,
139-
);
140-
141-
socket.addEventListener("message", (event) => {
142-
try {
143-
const res = JSON.parse(
144-
event.data,
145-
) as TypesGen.GetInboxNotificationResponse;
146-
onNewNotification(res);
147-
} catch (error) {
148-
console.warn("Error parsing inbox notification: ", error);
149-
}
150-
});
151-
152-
socket.addEventListener("error", (event) => {
153-
console.warn("Watch inbox notifications error: ", event);
154-
socket.close();
133+
): OneWayWebSocket<TypesGen.GetInboxNotificationResponse> {
134+
return new OneWayWebSocket({
135+
apiRoute: "/api/v2/notifications/inbox/watch",
136+
searchParams: params,
155137
});
156-
157-
return socket;
158-
};
138+
}
159139

160140
export const getURLWithSearchParams = (
161141
basePath: string,
@@ -1125,7 +1105,7 @@ class ApiMethods {
11251105
};
11261106

11271107
getWorkspaceByOwnerAndName = async (
1128-
username = "me",
1108+
username: string,
11291109
workspaceName: string,
11301110
params?: TypesGen.WorkspaceOptions,
11311111
): Promise<TypesGen.Workspace> => {
@@ -1138,7 +1118,7 @@ class ApiMethods {
11381118
};
11391119

11401120
getWorkspaceBuildByNumber = async (
1141-
username = "me",
1121+
username: string,
11421122
workspaceName: string,
11431123
buildNumber: number,
11441124
): Promise<TypesGen.WorkspaceBuild> => {
@@ -1324,7 +1304,7 @@ class ApiMethods {
13241304
};
13251305

13261306
createWorkspace = async (
1327-
userId = "me",
1307+
userId: string,
13281308
workspace: TypesGen.CreateWorkspaceRequest,
13291309
): Promise<TypesGen.Workspace> => {
13301310
const response = await this.axios.post<TypesGen.Workspace>(
@@ -2542,7 +2522,7 @@ function createWebSocket(
25422522
) {
25432523
const protocol = location.protocol === "https:" ? "wss:" : "ws:";
25442524
const socket = new WebSocket(
2545-
`${protocol}//${location.host}${path}?${params.toString()}`,
2525+
`${protocol}//${location.host}${path}?${params}`,
25462526
);
25472527
socket.binaryType = "blob";
25482528
return socket;

‎site/src/modules/notifications/NotificationsInbox/NotificationsInbox.tsx

+23-13
Original file line numberDiff line numberDiff line change
@@ -61,21 +61,31 @@ export const NotificationsInbox: FC<NotificationsInboxProps> = ({
6161
);
6262

6363
useEffect(() => {
64-
const socket = watchInboxNotifications(
65-
(res) => {
66-
updateNotificationsCache((prev) => {
67-
return {
68-
unread_count: res.unread_count,
69-
notifications: [res.notification, ...prev.notifications],
70-
};
71-
});
72-
},
73-
{ read_status: "unread" },
74-
);
64+
const socket = watchInboxNotifications({ read_status: "unread" });
7565

76-
return () => {
66+
socket.addEventListener("message", (e) => {
67+
if (e.parseError) {
68+
console.warn("Error parsing inbox notification: ", e.parseError);
69+
return;
70+
}
71+
72+
const msg = e.parsedMessage;
73+
updateNotificationsCache((current) => {
74+
return {
75+
unread_count: msg.unread_count,
76+
notifications: [msg.notification, ...current.notifications],
77+
};
78+
});
79+
});
80+
81+
socket.addEventListener("error", () => {
82+
displayError(
83+
"Unable to retrieve latest inbox notifications. Please try refreshing the browser.",
84+
);
7785
socket.close();
78-
};
86+
});
87+
88+
return () => socket.close();
7989
}, [updateNotificationsCache]);
8090

8191
const {

‎site/src/modules/resources/AgentMetadata.tsx

+64-29
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,11 @@ import Skeleton from "@mui/material/Skeleton";
33
import Tooltip from "@mui/material/Tooltip";
44
import { watchAgentMetadata } from "api/api";
55
import type {
6+
ServerSentEvent,
67
WorkspaceAgent,
78
WorkspaceAgentMetadata,
89
} from "api/typesGenerated";
10+
import { displayError } from "components/GlobalSnackbar/utils";
911
import { Stack } from "components/Stack/Stack";
1012
import dayjs from "dayjs";
1113
import {
@@ -17,6 +19,7 @@ import {
1719
useState,
1820
} from "react";
1921
import { MONOSPACE_FONT_FAMILY } from "theme/constants";
22+
import type { OneWayWebSocket } from "utils/OneWayWebSocket";
2023

2124
type ItemStatus = "stale" | "valid" | "loading";
2225

@@ -42,58 +45,90 @@ interface AgentMetadataProps {
4245
storybookMetadata?: WorkspaceAgentMetadata[];
4346
}
4447

48+
const maxSocketErrorRetryCount = 3;
49+
4550
export const AgentMetadata: FC<AgentMetadataProps> = ({
4651
agent,
4752
storybookMetadata,
4853
}) => {
49-
const [metadata, setMetadata] = useState<
50-
WorkspaceAgentMetadata[] | undefined
51-
>(undefined);
52-
54+
const [activeMetadata, setActiveMetadata] = useState(storybookMetadata);
5355
useEffect(() => {
56+
// This is an unfortunate pitfall with this component's testing setup,
57+
// but even though we use the value of storybookMetadata as the initial
58+
// value of the activeMetadata, we cannot put activeMetadata itself into
59+
// the dependency array. If we did, we would destroy and rebuild each
60+
// connection every single time a new message comes in from the socket,
61+
// because the socket has to be wired up to the state setter
5462
if (storybookMetadata !== undefined) {
55-
setMetadata(storybookMetadata);
5663
return;
5764
}
5865

59-
let timeout: ReturnType<typeof setTimeout> | undefined = undefined;
60-
61-
const connect = (): (() => void) => {
62-
const source = watchAgentMetadata(agent.id);
66+
let timeoutId: number | undefined = undefined;
67+
let activeSocket: OneWayWebSocket<ServerSentEvent> | null = null;
68+
let retries = 0;
69+
70+
const createNewConnection = () => {
71+
const socket = watchAgentMetadata(agent.id);
72+
activeSocket = socket;
73+
74+
socket.addEventListener("error", () => {
75+
setActiveMetadata(undefined);
76+
window.clearTimeout(timeoutId);
77+
78+
// The error event is supposed to fire when an error happens
79+
// with the connection itself, which implies that the connection
80+
// would auto-close. Couldn't find a definitive answer on MDN,
81+
// though, so closing it manually just to be safe
82+
socket.close();
83+
activeSocket = null;
84+
85+
retries++;
86+
if (retries >= maxSocketErrorRetryCount) {
87+
displayError(
88+
"Unexpected disconnect while watching Metadata changes. Please try refreshing the page.",
89+
);
90+
return;
91+
}
6392

64-
source.onerror = (e) => {
65-
console.error("received error in watch stream", e);
66-
setMetadata(undefined);
67-
source.close();
93+
displayError(
94+
"Unexpected disconnect while watching Metadata changes. Creating new connection...",
95+
);
96+
timeoutId = window.setTimeout(() => {
97+
createNewConnection();
98+
}, 3_000);
99+
});
68100

69-
timeout = setTimeout(() => {
70-
connect();
71-
}, 3000);
72-
};
101+
socket.addEventListener("message", (e) => {
102+
if (e.parseError) {
103+
displayError(
104+
"Unable to process newest response from server. Please try refreshing the page.",
105+
);
106+
return;
107+
}
73108

74-
source.addEventListener("data", (e) => {
75-
const data = JSON.parse(e.data);
76-
setMetadata(data);
77-
});
78-
return () => {
79-
if (timeout !== undefined) {
80-
clearTimeout(timeout);
109+
const msg = e.parsedMessage;
110+
if (msg.type === "data") {
111+
setActiveMetadata(msg.data as WorkspaceAgentMetadata[]);
81112
}
82-
source.close();
83-
};
113+
});
114+
};
115+
116+
createNewConnection();
117+
return () => {
118+
window.clearTimeout(timeoutId);
119+
activeSocket?.close();
84120
};
85-
return connect();
86121
}, [agent.id, storybookMetadata]);
87122

88-
if (metadata === undefined) {
123+
if (activeMetadata === undefined) {
89124
return (
90125
<section css={styles.root}>
91126
<AgentMetadataSkeleton />
92127
</section>
93128
);
94129
}
95130

96-
return <AgentMetadataView metadata={metadata} />;
131+
return <AgentMetadataView metadata={activeMetadata} />;
97132
};
98133

99134
export const AgentMetadataSkeleton: FC = () => {

0 commit comments

Comments
 (0)