Line data Source code
1 : import type { LoggerLike } from '../platform/logger/logger.js';
2 : import {
3 : AGENT_FEATURES,
4 : AGENT_PROTOCOL_VERSION,
5 : type AgentClientSession,
6 : type AgentCommandParams,
7 : type AgentCommandResult,
8 : type AgentErrorCode,
9 : type AgentNotification,
10 : type AgentRpcError,
11 : type AgentRpcMessage,
12 : type AgentServerMethods,
13 : type ClientSessionsParams,
14 : } from './agent_protocol.js';
15 :
16 1 : const LOG_CTX = 'agent_client';
17 :
18 : /// The connection to the local agent server
19 1 : export enum AgentConnectionStatus {
20 : /// No session has agent access and nothing holds the connection, the client does not connect
21 1 : Off = 'off',
22 : /// Connecting, or waiting to reconnect
23 1 : Connecting = 'connecting',
24 : /// Connected, the server took the sessions
25 1 : Connected = 'connected',
26 : /// The server turned the client away, it speaks another protocol version. The client retries slowly.
27 1 : Incompatible = 'incompatible',
28 1 : }
29 :
30 : /// The part of the WebSocket API that the client uses
31 : export interface AgentSocket {
32 : readonly readyState: number;
33 : onopen: ((event: unknown) => void) | null;
34 : onclose: ((event: unknown) => void) | null;
35 : onerror: ((event: unknown) => void) | null;
36 : onmessage: ((event: { data: unknown }) => void) | null;
37 : send(data: string): void;
38 : close(code?: number, reason?: string): void;
39 : }
40 :
41 1 : const SOCKET_OPEN = 1;
42 : /// Messages are cut after this many characters in the log
43 1 : const LOG_TEXT_LIMIT = 64 * 1024;
44 : /// Values up to this many characters on one line are logged as they are, longer ones as their size
45 1 : const LOG_VALUE_LIMIT = 120;
46 :
47 : /// A message for the log, cut after LOG_TEXT_LIMIT characters
48 1 : function clip(text: string): string {
49 1 : if (text.length <= LOG_TEXT_LIMIT) return text;
50 0 : return `${text.slice(0, LOG_TEXT_LIMIT)}... (${text.length - LOG_TEXT_LIMIT} more characters)`;
51 1 : }
52 :
53 : /// Is a value too long for a log line?
54 1 : function isLongText(text: string): boolean {
55 1 : return text.length > LOG_VALUE_LIMIT || text.includes('\n');
56 1 : }
57 :
58 : /// A value for a log line: short values as they are, long ones as their size
59 1 : function logValue(value: unknown): string {
60 1 : if (typeof value === 'string') {
61 1 : if (!isLongText(value)) return value;
62 1 : const lines = value.split('\n').length;
63 1 : return `${value.length} characters, ${lines} line${lines === 1 ? '' : 's'}`;
64 1 : }
65 1 : if (value == null || typeof value !== 'object') return String(value);
66 0 : const json = JSON.stringify(value);
67 0 : return json.length <= LOG_VALUE_LIMIT ? json : `${Array.isArray(value) ? 'array' : 'object'} of ${json.length} characters`;
68 1 : }
69 :
70 : /// The fields of an object for a log line
71 1 : function logFields(object: unknown): Record<string, string> {
72 1 : if (object == null || typeof object !== 'object' || Array.isArray(object)) return {};
73 1 : return Object.fromEntries(Object.entries(object).map(([key, value]) => [key, logValue(value)]));
74 1 : }
75 :
76 : /// The first line of a text for a log line
77 1 : function firstLine(text: string): string {
78 1 : const line = text.split('\n', 1)[0];
79 1 : return line.length <= LOG_VALUE_LIMIT ? line : `${line.slice(0, LOG_VALUE_LIMIT)}...`;
80 1 : }
81 : /// The close code of a server that turns a client away, e.g. for another protocol version
82 1 : const CLOSE_POLICY_VIOLATION = 1008;
83 :
84 : /// An error that a command reports to the agent
85 1 : export class AgentCommandError extends Error {
86 1 : constructor(readonly code: AgentErrorCode, message: string) {
87 1 : super(message);
88 1 : }
89 1 : }
90 :
91 : export interface AgentClientOptions {
92 : /// The WebSocket endpoint, e.g. ws://127.0.0.1:47100/v1/ws
93 : url: string;
94 : /// The kind of client, "web" or "desktop"
95 : clientKind: string;
96 : logger: LoggerLike;
97 : /// Run a command of the command line tool
98 : handleCommand: (params: AgentCommandParams) => Promise<AgentCommandResult>;
99 : /// Receive a notification, e.g. agent/status
100 : onNotification: (notification: AgentNotification) => void;
101 : /// Follow the connection status
102 : onStatus: (status: AgentConnectionStatus) => void;
103 : /// The connection is gone, the notifications of the turns that the server started for it are lost
104 : onDisconnect?: (reason: string) => void;
105 : /// Create a socket, replaceable in tests
106 : createSocket?: (url: string) => AgentSocket;
107 : /// The first reconnect delay, doubled after every failure
108 : initialBackoffMs?: number;
109 : /// The longest reconnect delay
110 : maxBackoffMs?: number;
111 : }
112 :
113 : interface PendingRequest {
114 : resolve: (result: unknown) => void;
115 : reject: (error: Error) => void;
116 : }
117 :
118 : /// The WebSocket client of the agent bridge.
119 : ///
120 : /// Connects while at least one session has agent access and reconnects with backoff.
121 : /// The server queues calls while the page is gone, so a reload loses nothing.
122 1 : export class AgentClient {
123 : readonly #options: AgentClientOptions;
124 1 : #socket: AgentSocket | null = null;
125 1 : #sessions: AgentClientSession[] = [];
126 1 : #focused: string | null = null;
127 : /// Stay connected without sessions, e.g. while the settings show the agent server
128 1 : #held = false;
129 1 : #status = AgentConnectionStatus.Off;
130 : #backoffMs: number;
131 1 : #reconnectTimer: ReturnType<typeof setTimeout> | null = null;
132 1 : #nextRequestId = 1;
133 1 : #pending = new Map<number, PendingRequest>();
134 :
135 1 : constructor(options: AgentClientOptions) {
136 1 : this.#options = options;
137 1 : this.#backoffMs = options.initialBackoffMs ?? 1000;
138 1 : }
139 :
140 1 : get status(): AgentConnectionStatus {
141 1 : return this.#status;
142 1 : }
143 :
144 : /// Announce the sessions with agent access, connects or disconnects as needed
145 1 : setSessions(sessions: AgentClientSession[], focused: string | null): void {
146 1 : this.#sessions = sessions;
147 1 : this.#focused = focused != null && sessions.some(s => s.id === focused) ? focused : null;
148 1 : this.#update();
149 1 : }
150 :
151 : /// Stay connected while no session has agent access
152 1 : setHeld(held: boolean): void {
153 1 : this.#held = held;
154 1 : this.#update();
155 1 : }
156 :
157 1 : #wanted(): boolean {
158 1 : return this.#sessions.length > 0 || this.#held;
159 1 : }
160 :
161 1 : #update(): void {
162 1 : if (!this.#wanted()) {
163 1 : this.#disconnect();
164 1 : return;
165 1 : }
166 1 : if (this.#socket == null && this.#reconnectTimer == null) {
167 1 : this.#connect();
168 : } else {
169 1 : void this.#announce();
170 : }
171 1 : }
172 :
173 : /// Send a request to the server
174 1 : request<M extends keyof AgentServerMethods>(method: M, params: AgentServerMethods[M]['params']): Promise<AgentServerMethods[M]['result']> {
175 1 : const socket = this.#socket;
176 1 : if (socket == null || socket.readyState !== SOCKET_OPEN) {
177 1 : this.#options.logger.warn('Cannot send a request, the agent server is not connected', { method }, LOG_CTX);
178 1 : return Promise.reject(new Error('The agent server is not connected'));
179 1 : }
180 1 : const id = this.#nextRequestId++;
181 1 : this.#options.logger.info('Sent a request to the agent server', { id: String(id), method }, LOG_CTX);
182 1 : return new Promise<AgentServerMethods[M]['result']>((resolve, reject) => {
183 1 : this.#pending.set(id, { resolve: resolve as (result: unknown) => void, reject });
184 1 : this.#send(socket, JSON.stringify({ id, method, params } satisfies AgentRpcMessage));
185 1 : });
186 1 : }
187 :
188 : /// Disconnect for good
189 1 : dispose(): void {
190 1 : this.#sessions = [];
191 1 : this.#held = false;
192 1 : this.#disconnect();
193 1 : }
194 :
195 : /// Send a message, the raw message is logged at debug level
196 1 : #send(socket: AgentSocket, text: string): void {
197 1 : this.#options.logger.debug('Sent to the agent server', { message: clip(text) }, LOG_CTX);
198 1 : socket.send(text);
199 1 : }
200 :
201 1 : #setStatus(status: AgentConnectionStatus): void {
202 1 : if (this.#status === status) return;
203 1 : this.#options.logger.info('Changed the connection status', { from: this.#status, to: status }, LOG_CTX);
204 1 : this.#status = status;
205 1 : this.#options.onStatus(status);
206 1 : }
207 :
208 : /// Announce the sessions, resolves once the server took them
209 1 : async #announce(): Promise<boolean> {
210 1 : const socket = this.#socket;
211 1 : if (socket == null || socket.readyState !== SOCKET_OPEN) return false;
212 1 : const params: ClientSessionsParams = {
213 1 : protocol: AGENT_PROTOCOL_VERSION,
214 1 : client: this.#options.clientKind,
215 1 : sessions: this.#sessions,
216 1 : focused: this.#focused,
217 1 : features: AGENT_FEATURES,
218 1 : };
219 : // A server that turns the client away closes the socket instead of answering
220 1 : return this.request('client/sessions', params).then(() => true, () => false);
221 1 : }
222 :
223 1 : #connect(): void {
224 1 : this.#reconnectTimer = null;
225 : // An incompatible server stays incompatible until it takes the sessions
226 1 : if (this.#status !== AgentConnectionStatus.Incompatible) this.#setStatus(AgentConnectionStatus.Connecting);
227 1 : this.#options.logger.info('Connecting to the agent server', { url: this.#options.url }, LOG_CTX);
228 1 : let socket: AgentSocket;
229 1 : try {
230 1 : socket = (this.#options.createSocket ?? (url => new WebSocket(url) as unknown as AgentSocket))(this.#options.url);
231 1 : } catch (e: unknown) {
232 0 : this.#options.logger.warn('Cannot connect to the agent server', { url: this.#options.url, error: String(e) }, LOG_CTX);
233 0 : this.#scheduleReconnect();
234 0 : return;
235 0 : }
236 1 : this.#socket = socket;
237 1 : socket.onopen = () => {
238 1 : this.#options.logger.info('Opened the socket to the agent server, announcing the sessions', { url: this.#options.url }, LOG_CTX);
239 1 : void this.#announce().then(accepted => {
240 1 : if (!accepted || this.#socket !== socket) return;
241 1 : this.#options.logger.info('Connected to the agent server', { url: this.#options.url }, LOG_CTX);
242 1 : this.#backoffMs = this.#options.initialBackoffMs ?? 1000;
243 1 : this.#setStatus(AgentConnectionStatus.Connected);
244 1 : });
245 1 : };
246 1 : socket.onmessage = event => {
247 1 : if (typeof event.data === 'string') this.#receive(socket, event.data);
248 1 : };
249 1 : socket.onerror = () => {
250 : // Browsers do not tell why a WebSocket failed, the close event follows
251 0 : this.#options.logger.warn('The socket to the agent server failed', { url: this.#options.url }, LOG_CTX);
252 0 : };
253 1 : socket.onclose = event => {
254 1 : if (this.#socket !== socket) return;
255 1 : this.#socket = null;
256 1 : const { code, reason, wasClean } = (event ?? {}) as { code?: number; reason?: string; wasClean?: boolean };
257 1 : this.#options.logger.info('The socket to the agent server closed', {
258 1 : code: String(code ?? ''),
259 1 : reason: reason ?? '',
260 1 : wasClean: String(wasClean ?? ''),
261 1 : pending: String(this.#pending.size),
262 1 : }, LOG_CTX);
263 1 : if (code === CLOSE_POLICY_VIOLATION) {
264 1 : this.#options.logger.warn('The agent server turned the client away', { reason: reason ?? '' }, LOG_CTX);
265 1 : this.#lose(`The agent server turned DashQL away: ${reason || 'unknown reason'}`);
266 1 : if (this.#wanted()) {
267 : // Reloading the page or updating the server fixes this, retry slowly until then
268 1 : this.#setStatus(AgentConnectionStatus.Incompatible);
269 1 : this.#backoffMs = this.#options.maxBackoffMs ?? 30_000;
270 1 : this.#scheduleReconnect();
271 1 : return;
272 1 : }
273 1 : } else {
274 1 : this.#lose('The agent server disconnected');
275 : }
276 1 : if (this.#wanted()) {
277 1 : this.#scheduleReconnect();
278 : } else {
279 0 : this.#setStatus(AgentConnectionStatus.Off);
280 : }
281 1 : };
282 1 : }
283 :
284 1 : #scheduleReconnect(): void {
285 1 : if (this.#status !== AgentConnectionStatus.Incompatible) this.#setStatus(AgentConnectionStatus.Connecting);
286 1 : const delay = this.#backoffMs;
287 1 : this.#backoffMs = Math.min(this.#backoffMs * 2, this.#options.maxBackoffMs ?? 30_000);
288 1 : this.#options.logger.info('Reconnecting to the agent server later', { delayMs: String(delay) }, LOG_CTX);
289 1 : this.#reconnectTimer = setTimeout(() => this.#connect(), delay);
290 1 : }
291 :
292 1 : #disconnect(): void {
293 1 : if (this.#reconnectTimer != null) {
294 1 : clearTimeout(this.#reconnectTimer);
295 1 : this.#reconnectTimer = null;
296 1 : }
297 1 : const socket = this.#socket;
298 1 : this.#socket = null;
299 1 : if (socket != null) {
300 1 : socket.onclose = null;
301 1 : socket.close(1000, 'agent access turned off');
302 1 : this.#options.logger.info('Disconnected from the agent server', {}, LOG_CTX);
303 1 : this.#lose('Agent access was turned off');
304 1 : }
305 1 : this.#setStatus(AgentConnectionStatus.Off);
306 1 : }
307 :
308 : /// The connection is gone: fail the pending requests and tell the owner
309 1 : #lose(reason: string): void {
310 1 : if (this.#pending.size > 0) {
311 1 : this.#options.logger.warn('Failing the pending requests', { count: String(this.#pending.size), reason }, LOG_CTX);
312 : }
313 1 : for (const pending of this.#pending.values()) pending.reject(new Error(reason));
314 1 : this.#pending.clear();
315 1 : this.#options.onDisconnect?.(reason);
316 1 : }
317 :
318 1 : #receive(socket: AgentSocket, text: string): void {
319 1 : this.#options.logger.debug('Received from the agent server', { message: clip(text) }, LOG_CTX);
320 1 : let message: AgentRpcMessage;
321 1 : try {
322 1 : message = JSON.parse(text) as AgentRpcMessage;
323 1 : } catch (e: unknown) {
324 0 : this.#options.logger.warn('Received invalid JSON from the agent server', { error: String(e) }, LOG_CTX);
325 0 : return;
326 0 : }
327 1 : if (message.method == null) {
328 1 : const pending = this.#pending.get(Number(message.id));
329 1 : if (pending == null) {
330 0 : this.#options.logger.warn('Received a response that no request waits for', { id: String(message.id) }, LOG_CTX);
331 0 : return;
332 0 : }
333 1 : this.#pending.delete(Number(message.id));
334 1 : this.#options.logger.debug('Received a response from the agent server', { id: String(message.id), error: message.error?.message }, LOG_CTX);
335 1 : if (message.error != null) {
336 1 : pending.reject(new Error(message.error.message));
337 : } else {
338 1 : pending.resolve(message.result);
339 : }
340 1 : return;
341 1 : }
342 1 : if (message.id == null) {
343 1 : this.#options.logger.info('Received a notification from the agent server', { method: message.method, ...logFields(message.params) }, LOG_CTX);
344 1 : this.#options.onNotification({ method: message.method, params: message.params } as AgentNotification);
345 1 : return;
346 1 : }
347 1 : const id = message.id;
348 1 : const reply = (body: { result: AgentCommandResult } | { error: AgentRpcError }) => {
349 1 : if (socket.readyState === SOCKET_OPEN) {
350 1 : this.#send(socket, JSON.stringify({ id, ...body }));
351 : } else {
352 0 : this.#options.logger.warn('Cannot answer the agent server, the socket is closed', { id: String(id) }, LOG_CTX);
353 : }
354 1 : };
355 1 : if (message.method !== 'command') {
356 0 : reply({ error: { code: 'usage', message: `unknown method: ${message.method}` } });
357 0 : return;
358 0 : }
359 1 : const params = message.params as AgentCommandParams;
360 1 : const command = { id: String(id), command: String(params.command) };
361 1 : const logger = this.#options.logger;
362 1 : logger.info('Running a command of the agent', { ...logFields(params.args), ...command, session: String(params.session) }, LOG_CTX);
363 : // The texts that the agent passes, e.g. the script of edit or the spec of vis edit
364 1 : for (const [name, value] of Object.entries(params.args ?? {})) {
365 1 : if (typeof value === 'string' && isLongText(value)) logger.debug(`The ${name} of the command`, { ...command, [name]: value }, LOG_CTX);
366 : }
367 1 : const started = performance.now();
368 1 : const duration = () => String(Math.round(performance.now() - started));
369 1 : this.#options.handleCommand(params).then(
370 1 : result => {
371 1 : logger.info('Finished the command of the agent', { ...command, durationMs: duration(), result: firstLine(result.text) }, LOG_CTX);
372 1 : if (isLongText(result.text)) logger.debug('The result of the command', { ...command, text: clip(result.text) }, LOG_CTX);
373 1 : reply({ result });
374 1 : },
375 1 : (e: unknown) => {
376 1 : const code = e instanceof AgentCommandError ? e.code : 'failed';
377 1 : const error = e instanceof Error ? e.message : String(e);
378 1 : logger.warn('The command of the agent failed', { ...command, durationMs: duration(), code, error }, LOG_CTX);
379 1 : reply({ error: { code, message: error } });
380 1 : },
381 1 : );
382 1 : }
383 1 : }
|
