Skip to content

Commit 3a9e159

Browse files
EhabYEhabY
committed
Handle CoderApi review comments
1 parent e5bdad3 commit 3a9e159

File tree

3 files changed

+88
-115
lines changed

3 files changed

+88
-115
lines changed

‎src/api/coderApi.ts‎

Lines changed: 67 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ const coderSessionTokenHeader = "Coder-Session-Token";
5757
*/
5858
export class CoderApi extends Api implements vscode.Disposable {
5959
private readonly reconnectingSockets = new Set<
60-
ReconnectingWebSocket<unknown>
60+
ReconnectingWebSocket<never>
6161
>();
6262

6363
private constructor(private readonly output: Logger) {
@@ -80,6 +80,16 @@ export class CoderApi extends Api implements vscode.Disposable {
8080
return client;
8181
}
8282

83+
getHost(): string | undefined {
84+
return this.getAxiosInstance().defaults.baseURL;
85+
}
86+
87+
getSessionToken(): string | undefined {
88+
return this.getAxiosInstance().defaults.headers.common[
89+
coderSessionTokenHeader
90+
] as string | undefined;
91+
}
92+
8393
/**
8494
* Set both host and token together. Useful for login/logout/switch to
8595
* avoid triggering multiple reconnection events.
@@ -88,14 +98,15 @@ export class CoderApi extends Api implements vscode.Disposable {
8898
host: string | undefined,
8999
token: string | undefined,
90100
): void => {
91-
const defaults = this.getAxiosInstance().defaults;
92-
const currentHost = defaults.baseURL;
93-
const currentToken = defaults.headers.common[coderSessionTokenHeader];
101+
const currentHost = this.getHost();
102+
const currentToken = this.getSessionToken();
94103

104+
// We cannot use the super.setHost/setSessionToken methods because they are shadowed here
105+
const defaults = this.getAxiosInstance().defaults;
95106
defaults.baseURL = host;
96107
defaults.headers.common[coderSessionTokenHeader] = token;
97108

98-
const hostChanged = currentHost !== host;
109+
const hostChanged = (currentHost || "") !== (host || "");
99110
const tokenChanged = currentToken !== token;
100111

101112
if (hostChanged || tokenChanged) {
@@ -109,16 +120,12 @@ export class CoderApi extends Api implements vscode.Disposable {
109120
}
110121
};
111122

112-
setSessionToken = (token: string): void => {
113-
const currentHost = this.getAxiosInstance().defaults.baseURL;
114-
this.setCredentials(currentHost, token);
123+
override setSessionToken = (token: string): void => {
124+
this.setCredentials(this.getHost(), token);
115125
};
116126

117-
setHost = (host: string | undefined): void => {
118-
const currentToken = this.getAxiosInstance().defaults.headers.common[
119-
coderSessionTokenHeader
120-
] as string | undefined;
121-
this.setCredentials(host, currentToken);
127+
override setHost = (host: string | undefined): void => {
128+
this.setCredentials(host, this.getSessionToken());
122129
};
123130

124131
/**
@@ -137,37 +144,40 @@ export class CoderApi extends Api implements vscode.Disposable {
137144
watchTargets: string[],
138145
options?: ClientOptions,
139146
) => {
140-
return this.createWebSocket<GetInboxNotificationResponse>({
141-
apiRoute: "/api/v2/notifications/inbox/watch",
142-
searchParams: {
143-
format: "plaintext",
144-
templates: watchTemplates.join(","),
145-
targets: watchTargets.join(","),
146-
},
147-
options,
148-
enableRetry: true,
149-
});
147+
return this.createReconnectingSocket(() =>
148+
this.createOneWayWebSocket<GetInboxNotificationResponse>({
149+
apiRoute: "/api/v2/notifications/inbox/watch",
150+
searchParams: {
151+
format: "plaintext",
152+
templates: watchTemplates.join(","),
153+
targets: watchTargets.join(","),
154+
},
155+
options,
156+
}),
157+
);
150158
};
151159

152160
watchWorkspace = async (workspace: Workspace, options?: ClientOptions) => {
153-
return this.createWebSocketWithFallback({
154-
apiRoute: `/api/v2/workspaces/${workspace.id}/watch-ws`,
155-
fallbackApiRoute: `/api/v2/workspaces/${workspace.id}/watch`,
156-
options,
157-
enableRetry: true,
158-
});
161+
return this.createReconnectingSocket(() =>
162+
this.createStreamWithSseFallback({
163+
apiRoute: `/api/v2/workspaces/${workspace.id}/watch-ws`,
164+
fallbackApiRoute: `/api/v2/workspaces/${workspace.id}/watch`,
165+
options,
166+
}),
167+
);
159168
};
160169

161170
watchAgentMetadata = async (
162171
agentId: WorkspaceAgent["id"],
163172
options?: ClientOptions,
164173
) => {
165-
return this.createWebSocketWithFallback({
166-
apiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata-ws`,
167-
fallbackApiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata`,
168-
options,
169-
enableRetry: true,
170-
});
174+
return this.createReconnectingSocket(() =>
175+
this.createStreamWithSseFallback({
176+
apiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata-ws`,
177+
fallbackApiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata`,
178+
options,
179+
}),
180+
);
171181
};
172182

173183
watchBuildLogsByBuildId = async (
@@ -205,33 +215,13 @@ export class CoderApi extends Api implements vscode.Disposable {
205215
searchParams.append("after", lastLog.id.toString());
206216
}
207217

208-
return this.createWebSocket<TData>({
218+
return this.createOneWayWebSocket<TData>({
209219
apiRoute,
210220
searchParams,
211221
options,
212222
});
213223
}
214224

215-
private async createWebSocket<TData = unknown>(
216-
configs: Omit<OneWayWebSocketInit, "location"> & { enableRetry?: boolean },
217-
): Promise<UnidirectionalStream<TData>> {
218-
const { enableRetry, ...socketConfigs } = configs;
219-
220-
const socketFactory: SocketFactory<TData> = async () => {
221-
const baseUrlRaw = this.getAxiosInstance().defaults.baseURL;
222-
if (!baseUrlRaw) {
223-
throw new Error("No base URL set on REST client");
224-
}
225-
226-
return this.createOneWayWebSocket<TData>(socketConfigs);
227-
};
228-
229-
if (enableRetry) {
230-
return this.createReconnectingSocket(socketFactory, configs.apiRoute);
231-
}
232-
return socketFactory();
233-
}
234-
235225
private async createOneWayWebSocket<TData>(
236226
configs: Omit<OneWayWebSocketInit, "location">,
237227
): Promise<OneWayWebSocket<TData>> {
@@ -307,46 +297,34 @@ export class CoderApi extends Api implements vscode.Disposable {
307297
/**
308298
* Create a WebSocket connection with SSE fallback on 404.
309299
*
310-
* The factory tries WS first, falls back to SSE on 404. Since the factory
311-
* is called on every reconnect.
300+
* Tries WS first, falls back to SSE on 404.
312301
*
313302
* Note: The fallback on SSE ignores all passed client options except the headers.
314303
*/
315-
private async createWebSocketWithFallback(
304+
private async createStreamWithSseFallback(
316305
configs: Omit<OneWayWebSocketInit, "location"> & {
317306
fallbackApiRoute: string;
318-
enableRetry?: boolean;
319307
},
320308
): Promise<UnidirectionalStream<ServerSentEvent>> {
321-
const { fallbackApiRoute, enableRetry, ...socketConfigs } = configs;
322-
const socketFactory: SocketFactory<ServerSentEvent> = async () => {
323-
try {
324-
const ws =
325-
await this.createOneWayWebSocket<ServerSentEvent>(socketConfigs);
326-
return await this.waitForOpen(ws);
327-
} catch (error) {
328-
if (this.is404Error(error)) {
329-
this.output.warn(
330-
`WebSocket failed, using SSE fallback: ${socketConfigs.apiRoute}`,
331-
);
332-
const sse = this.createSseConnection(
333-
fallbackApiRoute,
334-
socketConfigs.searchParams,
335-
socketConfigs.options?.headers,
336-
);
337-
return await this.waitForOpen(sse);
338-
}
339-
throw error;
309+
const { fallbackApiRoute, ...socketConfigs } = configs;
310+
try {
311+
const ws =
312+
await this.createOneWayWebSocket<ServerSentEvent>(socketConfigs);
313+
return await this.waitForOpen(ws);
314+
} catch (error) {
315+
if (this.is404Error(error)) {
316+
this.output.warn(
317+
`WebSocket failed, using SSE fallback: ${socketConfigs.apiRoute}`,
318+
);
319+
const sse = this.createSseConnection(
320+
fallbackApiRoute,
321+
socketConfigs.searchParams,
322+
socketConfigs.options?.headers,
323+
);
324+
return await this.waitForOpen(sse);
340325
}
341-
};
342-
343-
if (enableRetry) {
344-
return this.createReconnectingSocket(
345-
socketFactory,
346-
socketConfigs.apiRoute,
347-
);
326+
throw error;
348327
}
349-
return socketFactory();
350328
}
351329

352330
/**
@@ -416,22 +394,15 @@ export class CoderApi extends Api implements vscode.Disposable {
416394
*/
417395
private async createReconnectingSocket<TData>(
418396
socketFactory: SocketFactory<TData>,
419-
apiRoute: string,
420397
): Promise<ReconnectingWebSocket<TData>> {
421398
const reconnectingSocket = await ReconnectingWebSocket.create<TData>(
422399
socketFactory,
423400
this.output,
424-
apiRoute,
425401
undefined,
426-
() =>
427-
this.reconnectingSockets.delete(
428-
reconnectingSocket as ReconnectingWebSocket<unknown>,
429-
),
402+
() => this.reconnectingSockets.delete(reconnectingSocket),
430403
);
431404

432-
this.reconnectingSockets.add(
433-
reconnectingSocket as ReconnectingWebSocket<unknown>,
434-
);
405+
this.reconnectingSockets.add(reconnectingSocket);
435406

436407
return reconnectingSocket;
437408
}

‎src/websocket/reconnectingWebSocket.ts‎

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@ export class ReconnectingWebSocket<TData = unknown>
2727
{
2828
readonly #socketFactory: SocketFactory<TData>;
2929
readonly #logger: Logger;
30-
readonly #apiRoute: string;
3130
readonly #options: Required<ReconnectingWebSocketOptions>;
3231
readonly #eventHandlers: {
3332
[K in WebSocketEventType]: Set<EventHandler<TData, K>>;
@@ -39,6 +38,7 @@ export class ReconnectingWebSocket<TData = unknown>
3938
};
4039

4140
#currentSocket: UnidirectionalStream<TData> | null = null;
41+
#lastRoute = "unknown"; // Cached route for logging when socket is closed
4242
#backoffMs: number;
4343
#reconnectTimeoutId: NodeJS.Timeout | null = null;
4444
#isSuspended = false; // Temporary pause, can be resumed via reconnect()
@@ -50,13 +50,11 @@ export class ReconnectingWebSocket<TData = unknown>
5050
private constructor(
5151
socketFactory: SocketFactory<TData>,
5252
logger: Logger,
53-
apiRoute: string,
5453
options: ReconnectingWebSocketOptions = {},
5554
onDispose?: () => void,
5655
) {
5756
this.#socketFactory = socketFactory;
5857
this.#logger = logger;
59-
this.#apiRoute = apiRoute;
6058
this.#options = {
6159
initialBackoffMs: options.initialBackoffMs ?? 250,
6260
maxBackoffMs: options.maxBackoffMs ?? 30000,
@@ -69,14 +67,12 @@ export class ReconnectingWebSocket<TData = unknown>
6967
static async create<TData>(
7068
socketFactory: SocketFactory<TData>,
7169
logger: Logger,
72-
apiRoute: string,
7370
options: ReconnectingWebSocketOptions = {},
7471
onDispose?: () => void,
7572
): Promise<ReconnectingWebSocket<TData>> {
7673
const instance = new ReconnectingWebSocket<TData>(
7774
socketFactory,
7875
logger,
79-
apiRoute,
8076
options,
8177
onDispose,
8278
);
@@ -88,6 +84,19 @@ export class ReconnectingWebSocket<TData = unknown>
8884
return this.#currentSocket?.url ?? "";
8985
}
9086

87+
/**
88+
* Extract the route (pathname + search) from the current socket URL for logging.
89+
* Falls back to the last known route when the socket is closed.
90+
*/
91+
get #route(): string {
92+
const socketUrl = this.#currentSocket?.url;
93+
if (!socketUrl) {
94+
return this.#lastRoute;
95+
}
96+
const url = new URL(socketUrl);
97+
return url.pathname + url.search;
98+
}
99+
91100
addEventListener<TEvent extends WebSocketEventType>(
92101
event: TEvent,
93102
callback: EventHandler<TData, TEvent>,
@@ -174,6 +183,7 @@ export class ReconnectingWebSocket<TData = unknown>
174183

175184
const socket = await this.#socketFactory();
176185
this.#currentSocket = socket;
186+
this.#lastRoute = this.#route;
177187

178188
socket.addEventListener("open", (event) => {
179189
this.#backoffMs = this.#options.initialBackoffMs;
@@ -193,7 +203,7 @@ export class ReconnectingWebSocket<TData = unknown>
193203
const errorMessage = event.error?.message ?? event.message ?? "";
194204
if (this.isUnrecoverableHttpError(errorMessage)) {
195205
this.#logger.error(
196-
`Unrecoverable HTTP error for ${this.#apiRoute}: ${errorMessage}`,
206+
`Unrecoverable HTTP error for ${this.#route}: ${errorMessage}`,
197207
);
198208
this.suspend();
199209
}
@@ -247,7 +257,7 @@ export class ReconnectingWebSocket<TData = unknown>
247257
const delayMs = Math.max(0, this.#backoffMs + jitter);
248258

249259
this.#logger.debug(
250-
`Reconnecting WebSocket in ${Math.round(delayMs)}ms for ${this.#apiRoute}`,
260+
`Reconnecting WebSocket in ${Math.round(delayMs)}ms for ${this.#route}`,
251261
);
252262

253263
this.#reconnectTimeoutId = setTimeout(() => {
@@ -267,7 +277,7 @@ export class ReconnectingWebSocket<TData = unknown>
267277
handler(eventData);
268278
} catch (error) {
269279
this.#logger.error(
270-
`Error in ${event} handler for ${this.#apiRoute}`,
280+
`Error in ${event} handler for ${this.#route}`,
271281
error,
272282
);
273283
}
@@ -285,17 +295,14 @@ export class ReconnectingWebSocket<TData = unknown>
285295

286296
if (this.isUnrecoverableHttpError(error)) {
287297
this.#logger.error(
288-
`Unrecoverable HTTP error during connection for ${this.#apiRoute}`,
298+
`Unrecoverable HTTP error during connection for ${this.#route}`,
289299
error,
290300
);
291301
this.suspend();
292302
return;
293303
}
294304

295-
this.#logger.warn(
296-
`WebSocket connection failed for ${this.#apiRoute}`,
297-
error,
298-
);
305+
this.#logger.warn(`WebSocket connection failed for ${this.#route}`, error);
299306
this.scheduleReconnect();
300307
}
301308

‎test/unit/websocket/reconnectingWebSocket.test.ts‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -91,11 +91,7 @@ describe("ReconnectingWebSocket", () => {
9191
});
9292

9393
await expect(
94-
ReconnectingWebSocket.create(
95-
factory,
96-
createMockLogger(),
97-
"/api/test",
98-
),
94+
ReconnectingWebSocket.create(factory, createMockLogger()),
9995
).rejects.toThrow(`Unexpected server response: ${statusCode}`);
10096

10197
// Should not retry after unrecoverable HTTP error
@@ -598,7 +594,6 @@ async function fromFactory<T>(
598594
return await ReconnectingWebSocket.create(
599595
factory,
600596
createMockLogger(),
601-
"/random/api",
602597
undefined,
603598
onDispose,
604599
);

0 commit comments

Comments
 (0)