Carlos Florêncio
522 lines
16 KiB
TypeScript
522 lines
16 KiB
TypeScript
import type { BackendMessageMap, BackendMessageType } from "@/lib/types/backend";
|
|
import type { ConnectionStatus } from "@/lib/types/connection";
|
|
import { generateUUID } from "@/lib/utils";
|
|
import { createDebugLogger, IS_DEBUG } from "@/lib/debug/log";
|
|
|
|
const debugDispatch = createDebugLogger("ws:dispatch");
|
|
|
|
// High-frequency notification types we skip in the dispatch log to avoid
|
|
// drowning the console during agent streams. Filter [ws:dispatch] to see
|
|
// everything else; if you need the streaming traffic, comment this out.
|
|
const DISPATCH_LOG_DENYLIST = new Set<string>([
|
|
"session.message.added",
|
|
"session.message.updated",
|
|
"session.shell.output",
|
|
"session.process.output",
|
|
]);
|
|
|
|
type MessageHandler<T extends BackendMessageType> = (message: BackendMessageMap[T]) => void;
|
|
|
|
// Internal alias kept for readability — the WS client emits the same vocabulary
|
|
// the UI consumes. `disconnected` covers both "never connected" (formerly
|
|
// `idle`) and "socket closed" (formerly `closed`); `connected` replaces `open`.
|
|
type WebSocketStatus = ConnectionStatus;
|
|
|
|
export interface ReconnectOptions {
|
|
enabled?: boolean;
|
|
maxAttempts?: number;
|
|
initialDelay?: number;
|
|
maxDelay?: number;
|
|
backoffMultiplier?: number;
|
|
}
|
|
|
|
const DEFAULT_RECONNECT_OPTIONS: Required<ReconnectOptions> = {
|
|
enabled: true,
|
|
maxAttempts: 10,
|
|
initialDelay: 1000,
|
|
maxDelay: 30000,
|
|
backoffMultiplier: 1.5,
|
|
};
|
|
|
|
export class WebSocketClient {
|
|
private socket: WebSocket | null = null;
|
|
private status: WebSocketStatus = "disconnected";
|
|
private handlers = new Map<BackendMessageType, Set<MessageHandler<BackendMessageType>>>();
|
|
private pendingRequests = new Map<
|
|
string,
|
|
{
|
|
resolve: (payload: unknown) => void;
|
|
reject: (error: Error) => void;
|
|
timeout: ReturnType<typeof setTimeout>;
|
|
}
|
|
>();
|
|
private pendingQueue: string[] = [];
|
|
private reconnectOptions: Required<ReconnectOptions>;
|
|
private reconnectAttempts = 0;
|
|
private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
|
|
private intentionalClose = false;
|
|
private subscriptions = new Map<string, number>();
|
|
private sessionSubscriptions = new Map<string, number>();
|
|
// Ref-counted focus signals: a session can be focused by both the task panel
|
|
// and the task details page if both are mounted. Backend wakes its workspace
|
|
// tracker into fast-poll mode while any client has focus, falling back to
|
|
// slow when the count reaches 0 (debounced server-side).
|
|
private sessionFocusCounts = new Map<string, number>();
|
|
private userSubscriptionCount = 0;
|
|
private runSubscriptions = new Map<string, number>();
|
|
|
|
constructor(
|
|
private url: string,
|
|
private onStatusChange?: (status: WebSocketStatus) => void,
|
|
reconnectOptions?: ReconnectOptions,
|
|
) {
|
|
this.reconnectOptions = { ...DEFAULT_RECONNECT_OPTIONS, ...reconnectOptions };
|
|
}
|
|
|
|
getStatus() {
|
|
return this.status;
|
|
}
|
|
|
|
connect() {
|
|
if (this.socket) return;
|
|
this.intentionalClose = false;
|
|
this.clearReconnectTimer();
|
|
this.setStatus("connecting");
|
|
this.socket = new WebSocket(this.url);
|
|
|
|
this.socket.onopen = () => {
|
|
this.reconnectAttempts = 0;
|
|
this.setStatus("connected");
|
|
this.flushQueue();
|
|
this.resubscribe();
|
|
};
|
|
|
|
this.socket.onmessage = (event) => {
|
|
const parts = (event.data as string).split("\n");
|
|
for (const part of parts) {
|
|
const trimmed = part.trim();
|
|
if (!trimmed) continue;
|
|
try {
|
|
const message = JSON.parse(trimmed) as BackendMessageMap[BackendMessageType];
|
|
this.handleParsedMessage(message);
|
|
} catch {
|
|
// Ignore parse errors for individual messages
|
|
}
|
|
}
|
|
};
|
|
|
|
this.socket.onerror = () => {
|
|
this.setStatus("error");
|
|
};
|
|
|
|
this.socket.onclose = (event) => {
|
|
this.socket = null;
|
|
this.handleDisconnect(event);
|
|
};
|
|
}
|
|
|
|
disconnect() {
|
|
this.intentionalClose = true;
|
|
this.clearReconnectTimer();
|
|
if (this.socket) {
|
|
this.socket.close();
|
|
this.socket = null;
|
|
}
|
|
this.setStatus("disconnected");
|
|
this.cleanupPendingRequests();
|
|
}
|
|
|
|
send(payload: unknown) {
|
|
const data = JSON.stringify(payload);
|
|
if (IS_DEBUG) {
|
|
const p = payload as { action?: string; id?: string; type?: string } | null;
|
|
const action = p?.action ?? "?";
|
|
if (!DISPATCH_LOG_DENYLIST.has(action)) {
|
|
const sessionId = (p as { payload?: { session_id?: string } } | null)?.payload?.session_id;
|
|
debugDispatch("send", {
|
|
action,
|
|
id: p?.id ?? null,
|
|
type: p?.type ?? null,
|
|
sessionId: sessionId ?? null,
|
|
queued: this.status !== "connected" || !this.socket,
|
|
});
|
|
}
|
|
}
|
|
if (this.status !== "connected" || !this.socket) {
|
|
this.pendingQueue.push(data);
|
|
return;
|
|
}
|
|
this.socket.send(data);
|
|
}
|
|
|
|
request<T>(action: string, payload: unknown, timeoutMs = 5000): Promise<T> {
|
|
const id = generateUUID();
|
|
return new Promise((resolve, reject) => {
|
|
const timeout = setTimeout(() => {
|
|
this.pendingRequests.delete(id);
|
|
reject(new Error(`WebSocket request timed out: ${action}`));
|
|
}, timeoutMs);
|
|
this.pendingRequests.set(id, {
|
|
resolve: resolve as (payload: unknown) => void,
|
|
reject,
|
|
timeout,
|
|
});
|
|
this.send({ id, type: "request", action, payload });
|
|
});
|
|
}
|
|
|
|
subscribe(taskId: string) {
|
|
const currentCount = this.subscriptions.get(taskId) ?? 0;
|
|
const nextCount = currentCount + 1;
|
|
this.subscriptions.set(taskId, nextCount);
|
|
if (this.status === "connected" && nextCount === 1) {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "task.subscribe",
|
|
payload: { task_id: taskId },
|
|
});
|
|
}
|
|
return () => this.unsubscribe(taskId);
|
|
}
|
|
|
|
subscribeSession(sessionId: string) {
|
|
const currentCount = this.sessionSubscriptions.get(sessionId) ?? 0;
|
|
const nextCount = currentCount + 1;
|
|
this.sessionSubscriptions.set(sessionId, nextCount);
|
|
if (this.status === "connected" && nextCount === 1) {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "session.subscribe",
|
|
payload: { session_id: sessionId },
|
|
});
|
|
}
|
|
return () => this.unsubscribeSession(sessionId);
|
|
}
|
|
|
|
focusSession(sessionId: string) {
|
|
const currentCount = this.sessionFocusCounts.get(sessionId) ?? 0;
|
|
const nextCount = currentCount + 1;
|
|
this.sessionFocusCounts.set(sessionId, nextCount);
|
|
if (this.status === "connected" && nextCount === 1) {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "session.focus",
|
|
payload: { session_id: sessionId },
|
|
});
|
|
}
|
|
return () => this.unfocusSession(sessionId);
|
|
}
|
|
|
|
unfocusSession(sessionId: string) {
|
|
const currentCount = this.sessionFocusCounts.get(sessionId);
|
|
if (!currentCount) return;
|
|
const nextCount = currentCount - 1;
|
|
if (nextCount <= 0) {
|
|
this.sessionFocusCounts.delete(sessionId);
|
|
if (this.status === "connected") {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "session.unfocus",
|
|
payload: { session_id: sessionId },
|
|
});
|
|
}
|
|
return;
|
|
}
|
|
this.sessionFocusCounts.set(sessionId, nextCount);
|
|
}
|
|
|
|
subscribeUser() {
|
|
this.userSubscriptionCount += 1;
|
|
if (this.status === "connected" && this.userSubscriptionCount === 1) {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "user.subscribe",
|
|
payload: {},
|
|
});
|
|
}
|
|
}
|
|
|
|
unsubscribe(taskId: string) {
|
|
const currentCount = this.subscriptions.get(taskId);
|
|
if (!currentCount) return;
|
|
const nextCount = currentCount - 1;
|
|
if (nextCount <= 0) {
|
|
this.subscriptions.delete(taskId);
|
|
if (this.status === "connected") {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "task.unsubscribe",
|
|
payload: { task_id: taskId },
|
|
});
|
|
}
|
|
return;
|
|
}
|
|
this.subscriptions.set(taskId, nextCount);
|
|
}
|
|
|
|
unsubscribeSession(sessionId: string) {
|
|
const currentCount = this.sessionSubscriptions.get(sessionId);
|
|
if (!currentCount) {
|
|
return;
|
|
}
|
|
const nextCount = currentCount - 1;
|
|
|
|
if (nextCount <= 0) {
|
|
this.sessionSubscriptions.delete(sessionId);
|
|
if (this.status === "connected") {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "session.unsubscribe",
|
|
payload: { session_id: sessionId },
|
|
});
|
|
}
|
|
return;
|
|
}
|
|
this.sessionSubscriptions.set(sessionId, nextCount);
|
|
}
|
|
|
|
/**
|
|
* Subscribe to office run-event notifications for the given run id.
|
|
* Ref-counted so multiple components can subscribe independently;
|
|
* only the first call sends `run.subscribe` over the wire, only the
|
|
* last unsubscribe sends `run.unsubscribe`. Returns an
|
|
* unsubscribe function.
|
|
*/
|
|
subscribeRun(runId: string) {
|
|
const currentCount = this.runSubscriptions.get(runId) ?? 0;
|
|
const nextCount = currentCount + 1;
|
|
this.runSubscriptions.set(runId, nextCount);
|
|
if (this.status === "connected" && nextCount === 1) {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "run.subscribe",
|
|
payload: { run_id: runId },
|
|
});
|
|
}
|
|
return () => this.unsubscribeRun(runId);
|
|
}
|
|
|
|
unsubscribeRun(runId: string) {
|
|
const currentCount = this.runSubscriptions.get(runId);
|
|
if (!currentCount) return;
|
|
const nextCount = currentCount - 1;
|
|
if (nextCount <= 0) {
|
|
this.runSubscriptions.delete(runId);
|
|
if (this.status === "connected") {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "run.unsubscribe",
|
|
payload: { run_id: runId },
|
|
});
|
|
}
|
|
return;
|
|
}
|
|
this.runSubscriptions.set(runId, nextCount);
|
|
}
|
|
|
|
unsubscribeUser() {
|
|
this.userSubscriptionCount = Math.max(0, this.userSubscriptionCount - 1);
|
|
if (this.status === "connected" && this.userSubscriptionCount === 0) {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "user.unsubscribe",
|
|
payload: {},
|
|
});
|
|
}
|
|
}
|
|
|
|
on<T extends BackendMessageType>(type: T, handler: MessageHandler<T>) {
|
|
const handlers = this.handlers.get(type) ?? new Set();
|
|
handlers.add(handler as MessageHandler<BackendMessageType>);
|
|
this.handlers.set(type, handlers);
|
|
return () => this.off(type, handler);
|
|
}
|
|
|
|
off<T extends BackendMessageType>(type: T, handler: MessageHandler<T>) {
|
|
const handlers = this.handlers.get(type);
|
|
if (!handlers) return;
|
|
handlers.delete(handler as MessageHandler<BackendMessageType>);
|
|
if (!handlers.size) {
|
|
this.handlers.delete(type);
|
|
}
|
|
}
|
|
|
|
private debugNotification(action: BackendMessageType, payload: unknown, handlerCount: number) {
|
|
if (!IS_DEBUG || DISPATCH_LOG_DENYLIST.has(action)) return;
|
|
const payloadSessionId = (payload as { session_id?: string } | undefined)?.session_id;
|
|
debugDispatch("notification", {
|
|
action,
|
|
sessionId: payloadSessionId ?? null,
|
|
handlers: handlerCount,
|
|
});
|
|
}
|
|
|
|
private handleParsedMessage(message: BackendMessageMap[BackendMessageType]) {
|
|
const msgWithId = message as { id?: string; type: string };
|
|
|
|
if (msgWithId.type === "response" && msgWithId.id) {
|
|
if (IS_DEBUG) debugDispatch("response", { id: msgWithId.id });
|
|
this.resolvePendingRequest(msgWithId.id, message.payload);
|
|
return;
|
|
}
|
|
if (msgWithId.type === "error" && msgWithId.id) {
|
|
if (IS_DEBUG) debugDispatch("error-response", { id: msgWithId.id });
|
|
this.rejectPendingRequest(msgWithId.id, message.payload);
|
|
return;
|
|
}
|
|
if (message.type !== "notification") return;
|
|
|
|
const action = (message as { action?: string })?.action as BackendMessageType | undefined;
|
|
if (!action) return;
|
|
const handlers = this.handlers.get(action);
|
|
this.debugNotification(action, message.payload, handlers?.size ?? 0);
|
|
if (handlers) {
|
|
handlers.forEach((handler) => handler(message));
|
|
}
|
|
}
|
|
|
|
private resolvePendingRequest(msgId: string, payload: unknown) {
|
|
const pending = this.pendingRequests.get(msgId);
|
|
if (!pending) return;
|
|
clearTimeout(pending.timeout);
|
|
this.pendingRequests.delete(msgId);
|
|
pending.resolve(payload);
|
|
}
|
|
|
|
private rejectPendingRequest(msgId: string, payload: unknown) {
|
|
const pending = this.pendingRequests.get(msgId);
|
|
if (!pending) return;
|
|
clearTimeout(pending.timeout);
|
|
this.pendingRequests.delete(msgId);
|
|
const errorMessage =
|
|
typeof payload === "object" && payload && "message" in payload
|
|
? String((payload as { message?: string }).message)
|
|
: "WebSocket request failed";
|
|
pending.reject(new Error(errorMessage));
|
|
}
|
|
|
|
private handleDisconnect(event: CloseEvent) {
|
|
this.setStatus("disconnected");
|
|
|
|
// Don't reconnect if this was an intentional close
|
|
if (this.intentionalClose) {
|
|
return;
|
|
}
|
|
|
|
// Don't reconnect if reconnect is disabled
|
|
if (!this.reconnectOptions.enabled) {
|
|
return;
|
|
}
|
|
|
|
// Don't reconnect if we've exceeded max attempts
|
|
if (this.reconnectAttempts >= this.reconnectOptions.maxAttempts) {
|
|
console.warn(
|
|
`WebSocket max reconnect attempts (${this.reconnectOptions.maxAttempts}) reached`,
|
|
);
|
|
this.setStatus("error");
|
|
this.cleanupPendingRequests();
|
|
return;
|
|
}
|
|
|
|
// Calculate delay with exponential backoff
|
|
const delay = Math.min(
|
|
this.reconnectOptions.initialDelay *
|
|
Math.pow(this.reconnectOptions.backoffMultiplier, this.reconnectAttempts),
|
|
this.reconnectOptions.maxDelay,
|
|
);
|
|
|
|
this.reconnectAttempts++;
|
|
this.setStatus("reconnecting");
|
|
|
|
console.log(
|
|
`WebSocket disconnected (code: ${event.code}, reason: ${event.reason || "none"}). ` +
|
|
`Reconnecting in ${delay}ms (attempt ${this.reconnectAttempts}/${this.reconnectOptions.maxAttempts})...`,
|
|
);
|
|
|
|
this.reconnectTimer = setTimeout(() => {
|
|
this.connect();
|
|
}, delay);
|
|
}
|
|
|
|
private clearReconnectTimer() {
|
|
if (this.reconnectTimer) {
|
|
clearTimeout(this.reconnectTimer);
|
|
this.reconnectTimer = null;
|
|
}
|
|
}
|
|
|
|
private cleanupPendingRequests() {
|
|
// Reject all pending requests
|
|
this.pendingRequests.forEach(({ reject, timeout }) => {
|
|
clearTimeout(timeout);
|
|
reject(new Error("WebSocket connection closed"));
|
|
});
|
|
this.pendingRequests.clear();
|
|
}
|
|
|
|
private resubscribe() {
|
|
// Re-subscribe to all tasks after reconnection
|
|
this.subscriptions.forEach((_count, taskId) => {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "task.subscribe",
|
|
payload: { task_id: taskId },
|
|
});
|
|
});
|
|
this.sessionSubscriptions.forEach((_count, sessionId) => {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "session.subscribe",
|
|
payload: { session_id: sessionId },
|
|
});
|
|
});
|
|
this.sessionFocusCounts.forEach((_count, sessionId) => {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "session.focus",
|
|
payload: { session_id: sessionId },
|
|
});
|
|
});
|
|
this.runSubscriptions.forEach((_count, runId) => {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "run.subscribe",
|
|
payload: { run_id: runId },
|
|
});
|
|
});
|
|
if (this.userSubscriptionCount > 0) {
|
|
this.send({
|
|
id: generateUUID(),
|
|
type: "request",
|
|
action: "user.subscribe",
|
|
payload: {},
|
|
});
|
|
}
|
|
}
|
|
|
|
private flushQueue() {
|
|
if (!this.socket || this.status !== "connected") return;
|
|
this.pendingQueue.forEach((data) => this.socket?.send(data));
|
|
this.pendingQueue = [];
|
|
}
|
|
|
|
private setStatus(status: WebSocketStatus) {
|
|
this.status = status;
|
|
this.onStatusChange?.(status);
|
|
}
|
|
}
|