Skip to content

Commit 80d2958

Browse files
EhabYEhabY
committed
Add logging interceptor for axios + Refactor WS (WIP)
1 parent a70f4d9 commit 80d2958

File tree

9 files changed

+369
-119
lines changed

9 files changed

+369
-119
lines changed

‎src/agentMetadataHelper.ts

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
import { Api } from "coder/site/src/api/api";
22
import { WorkspaceAgent } from "coder/site/src/api/typesGenerated";
3-
import { EventSource } from "eventsource";
3+
import { ProxyAgent } from "proxy-agent";
44
import * as vscode from "vscode";
5-
import { createStreamingFetchAdapter } from "./api";
65
import {
76
AgentMetadataEvent,
87
AgentMetadataEventSchemaArray,
98
errToStr,
109
} from "./api-helper";
10+
import { watchAgentMetadata } from "./websocket/ws-helper";
1111

1212
export type AgentMetadataWatcher = {
1313
onChange: vscode.EventEmitter<null>["event"];
@@ -17,38 +17,33 @@ export type AgentMetadataWatcher = {
1717
};
1818

1919
/**
20-
* Opens an SSE connection to watch metadata for a given workspace agent.
20+
* Opens a websocket connection to watch metadata for a given workspace agent.
2121
* Emits onChange when metadata updates or an error occurs.
2222
*/
2323
export function createAgentMetadataWatcher(
2424
agentId: WorkspaceAgent["id"],
2525
restClient: Api,
26+
httpAgent: ProxyAgent,
2627
): AgentMetadataWatcher {
27-
// TODO: Is there a better way to grab the url and token?
28-
const url = restClient.getAxiosInstance().defaults.baseURL;
29-
const metadataUrl = new URL(
30-
`${url}/api/v2/workspaceagents/${agentId}/watch-metadata`,
31-
);
32-
const eventSource = new EventSource(metadataUrl.toString(), {
33-
fetch: createStreamingFetchAdapter(restClient.getAxiosInstance()),
34-
});
28+
const socket = watchAgentMetadata(restClient, httpAgent, agentId);
3529

3630
let disposed = false;
3731
const onChange = new vscode.EventEmitter<null>();
3832
const watcher: AgentMetadataWatcher = {
3933
onChange: onChange.event,
4034
dispose: () => {
4135
if (!disposed) {
42-
eventSource.close();
36+
socket.close();
4337
disposed = true;
4438
}
4539
},
4640
};
4741

48-
eventSource.addEventListener("data", (event) => {
42+
socket.addEventListener("message", (event) => {
4943
try {
50-
const dataEvent = JSON.parse(event.data);
51-
const metadata = AgentMetadataEventSchemaArray.parse(dataEvent);
44+
const metadata = AgentMetadataEventSchemaArray.parse(
45+
event.parsedMessage?.data,
46+
);
5247

5348
// Overwrite metadata if it changed.
5449
if (JSON.stringify(watcher.metadata) !== JSON.stringify(metadata)) {
@@ -61,6 +56,11 @@ export function createAgentMetadataWatcher(
6156
}
6257
});
6358

59+
socket.addEventListener("error", (error) => {
60+
watcher.error = error;
61+
onChange.fire(null);
62+
});
63+
6464
return watcher;
6565
}
6666

‎src/api.ts

Lines changed: 70 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
1-
import { AxiosInstance } from "axios";
1+
import { AxiosInstance, InternalAxiosRequestConfig, isAxiosError } from "axios";
22
import { spawn } from "child_process";
33
import { Api } from "coder/site/src/api/api";
44
import {
55
ProvisionerJobLog,
66
Workspace,
77
} from "coder/site/src/api/typesGenerated";
8-
import { FetchLikeInit } from "eventsource";
98
import fs from "fs/promises";
109
import { ProxyAgent } from "proxy-agent";
1110
import * as vscode from "vscode";
@@ -83,6 +82,9 @@ export function makeCoderSdk(
8382
restClient.setSessionToken(token);
8483
}
8584

85+
// Logging interceptor
86+
addLoggingInterceptors(restClient.getAxiosInstance(), storage.output);
87+
8688
restClient.getAxiosInstance().interceptors.request.use(async (config) => {
8789
// Add headers from the header command.
8890
Object.entries(await storage.getHeaders(baseUrl)).forEach(
@@ -113,57 +115,75 @@ export function makeCoderSdk(
113115
return restClient;
114116
}
115117

116-
/**
117-
* Creates a fetch adapter using an Axios instance that returns streaming responses.
118-
* This can be used with APIs that accept fetch-like interfaces.
119-
*/
120-
export function createStreamingFetchAdapter(axiosInstance: AxiosInstance) {
121-
return async (url: string | URL, init?: FetchLikeInit) => {
122-
const urlStr = url.toString();
123-
124-
const response = await axiosInstance.request({
125-
url: urlStr,
126-
signal: init?.signal,
127-
headers: init?.headers as Record<string, string>,
128-
responseType: "stream",
129-
validateStatus: () => true, // Don't throw on any status code
130-
});
131-
const stream = new ReadableStream({
132-
start(controller) {
133-
response.data.on("data", (chunk: Buffer) => {
134-
controller.enqueue(chunk);
135-
});
118+
interface RequestConfigWithMetadata extends InternalAxiosRequestConfig {
119+
metadata?: {
120+
requestId: string;
121+
startedAt: number;
122+
};
123+
}
136124

137-
response.data.on("end", () => {
138-
controller.close();
139-
});
125+
function addLoggingInterceptors(
126+
client: AxiosInstance,
127+
logger: vscode.LogOutputChannel,
128+
) {
129+
client.interceptors.request.use(
130+
(config) => {
131+
const requestId = crypto.randomUUID();
132+
(config as RequestConfigWithMetadata).metadata = {
133+
requestId,
134+
startedAt: Date.now(),
135+
};
140136

141-
response.data.on("error", (err: Error) => {
142-
controller.error(err);
143-
});
144-
},
137+
logger.trace(
138+
`Request ${requestId}: ${config.method?.toUpperCase()} ${config.url}`,
139+
config.data ?? "",
140+
);
145141

146-
cancel() {
147-
response.data.destroy();
148-
return Promise.resolve();
149-
},
150-
});
142+
return config;
143+
},
144+
(error: unknown) => {
145+
let message: string = "Request error";
146+
if (isAxiosError(error)) {
147+
const meta = (error.config as RequestConfigWithMetadata)?.metadata;
148+
const requestId = meta?.requestId ?? "n/a";
149+
message = `Request ${requestId} error`;
150+
}
151+
logger.warn(message, error);
151152

152-
return {
153-
body: {
154-
getReader: () => stream.getReader(),
155-
},
156-
url: urlStr,
157-
status: response.status,
158-
redirected: response.request.res.responseUrl !== urlStr,
159-
headers: {
160-
get: (name: string) => {
161-
const value = response.headers[name.toLowerCase()];
162-
return value === undefined ? null : String(value);
163-
},
164-
},
165-
};
166-
};
153+
return Promise.reject(error);
154+
},
155+
);
156+
157+
client.interceptors.response.use(
158+
(response) => {
159+
const { requestId, startedAt } =
160+
(response.config as RequestConfigWithMetadata).metadata ?? {};
161+
const ms = startedAt ? Date.now() - startedAt : undefined;
162+
163+
logger.trace(
164+
`Response ${requestId ?? "n/a"}: ${response.status}${
165+
ms !== undefined ? ` in ${ms}ms` : ""
166+
}`,
167+
// { responseBody: response.data }, // TODO too noisy
168+
);
169+
return response;
170+
},
171+
(error: unknown) => {
172+
let message = "Response error";
173+
if (isAxiosError(error)) {
174+
const { metadata } = (error.config as RequestConfigWithMetadata) ?? {};
175+
const requestId = metadata?.requestId ?? "n/a";
176+
const startedAt = metadata?.startedAt;
177+
const ms = startedAt ? Date.now() - startedAt : undefined;
178+
179+
const status = error.response?.status ?? "unknown";
180+
message = `Response ${requestId}: ${status}${ms !== undefined ? ` in ${ms}ms` : ""}`;
181+
}
182+
logger.warn(message, error);
183+
184+
return Promise.reject(error);
185+
},
186+
);
167187
}
168188

169189
/**
@@ -269,6 +289,7 @@ export async function waitForBuild(
269289
const agent = await createHttpAgent();
270290
await new Promise<void>((resolve, reject) => {
271291
try {
292+
// TODO move to `ws-helper`
272293
const baseUrl = new URL(baseUrlRaw);
273294
const proto = baseUrl.protocol === "https:" ? "wss:" : "ws:";
274295
const socketUrlRaw = `${proto}//${baseUrl.host}${path}`;

‎src/extension.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import axios, { isAxiosError } from "axios";
33
import { getErrorMessage } from "coder/site/src/api/errors";
44
import * as module from "module";
55
import * as vscode from "vscode";
6-
import { makeCoderSdk, needToken } from "./api";
6+
import { createHttpAgent, makeCoderSdk, needToken } from "./api";
77
import { errToStr } from "./api-helper";
88
import { Commands } from "./commands";
99
import { CertificateError, getErrorDetail } from "./error";
@@ -67,16 +67,20 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
6767
storage,
6868
);
6969

70+
// TODO this won't get updated when users change their settings; Listen to changes and update this
71+
const httpAgent = await createHttpAgent();
7072
const myWorkspacesProvider = new WorkspaceProvider(
7173
WorkspaceQuery.Mine,
7274
restClient,
7375
storage,
76+
httpAgent,
7477
5,
7578
);
7679
const allWorkspacesProvider = new WorkspaceProvider(
7780
WorkspaceQuery.All,
7881
restClient,
7982
storage,
83+
httpAgent,
8084
);
8185

8286
// createTreeView, unlike registerTreeDataProvider, gives us the tree view API

‎src/inbox.ts

Lines changed: 13 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,10 @@ import {
55
} from "coder/site/src/api/typesGenerated";
66
import { ProxyAgent } from "proxy-agent";
77
import * as vscode from "vscode";
8-
import { WebSocket } from "ws";
9-
import { coderSessionTokenHeader } from "./api";
108
import { errToStr } from "./api-helper";
119
import { type Storage } from "./storage";
10+
import { OneWayCodeWebSocket } from "./websocket/OneWayCodeWebSocket";
11+
import { watchInboxNotifications } from "./websocket/ws-helper";
1212

1313
// These are the template IDs of our notifications.
1414
// Maybe in the future we should avoid hardcoding
@@ -19,7 +19,7 @@ const TEMPLATE_WORKSPACE_OUT_OF_DISK = "f047f6a3-5713-40f7-85aa-0394cce9fa3a";
1919
export class Inbox implements vscode.Disposable {
2020
readonly #storage: Storage;
2121
#disposed = false;
22-
#socket: WebSocket;
22+
#socket: OneWayCodeWebSocket<GetInboxNotificationResponse>;
2323

2424
constructor(
2525
workspace: Workspace,
@@ -29,54 +29,32 @@ export class Inbox implements vscode.Disposable {
2929
) {
3030
this.#storage = storage;
3131

32-
const baseUrlRaw = restClient.getAxiosInstance().defaults.baseURL;
33-
if (!baseUrlRaw) {
34-
throw new Error("No base URL set on REST client");
35-
}
36-
3732
const watchTemplates = [
3833
TEMPLATE_WORKSPACE_OUT_OF_DISK,
3934
TEMPLATE_WORKSPACE_OUT_OF_MEMORY,
4035
];
41-
const watchTemplatesParam = encodeURIComponent(watchTemplates.join(","));
4236

4337
const watchTargets = [workspace.id];
44-
const watchTargetsParam = encodeURIComponent(watchTargets.join(","));
45-
46-
// We shouldn't need to worry about this throwing. Whilst `baseURL` could
47-
// be an invalid URL, that would've caused issues before we got to here.
48-
const baseUrl = new URL(baseUrlRaw);
49-
const socketProto = baseUrl.protocol === "https:" ? "wss:" : "ws:";
50-
const socketUrl = `${socketProto}//${baseUrl.host}/api/v2/notifications/inbox/watch?format=plaintext&templates=${watchTemplatesParam}&targets=${watchTargetsParam}`;
5138

52-
const token = restClient.getAxiosInstance().defaults.headers.common[
53-
coderSessionTokenHeader
54-
] as string | undefined;
55-
this.#socket = new WebSocket(new URL(socketUrl), {
56-
agent: httpAgent,
57-
followRedirects: true,
58-
headers: token
59-
? {
60-
[coderSessionTokenHeader]: token,
61-
}
62-
: undefined,
63-
});
39+
this.#socket = watchInboxNotifications(
40+
restClient,
41+
httpAgent,
42+
watchTemplates,
43+
watchTargets,
44+
);
6445

65-
this.#socket.on("open", () => {
46+
this.#socket.addEventListener("open", () => {
6647
this.#storage.output.info("Listening to Coder Inbox");
6748
});
6849

69-
this.#socket.on("error", (error) => {
50+
this.#socket.addEventListener("error", (error) => {
7051
this.notifyError(error);
7152
this.dispose();
7253
});
7354

74-
this.#socket.on("message", (data) => {
55+
this.#socket.addEventListener("message", (data) => {
7556
try {
76-
const inboxMessage = JSON.parse(
77-
data.toString(),
78-
) as GetInboxNotificationResponse;
79-
57+
const inboxMessage = data.parsedMessage!;
8058
vscode.window.showInformationMessage(inboxMessage.notification.title);
8159
} catch (error) {
8260
this.notifyError(error);

0 commit comments

Comments
 (0)