kandev/apps/web/lib/ws/client.ts

import type { BackendMessageMap, BackendMessageType } from "@/lib/types/backend";
import type { ConnectionStatus } from "@/lib/types/connection";
import { generateUUID } from "@/lib/utils";
import { createDebugLogger, isDebug } from "@/lib/debug/log";
import { dispatchToPluginWsHandlers } from "@/lib/ws/plugin-bridge";

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.message.deleted",
  "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>();
  private systemMetricsSubscriptionCount = 0;

  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 (isDebug()) {
      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);
  }

  /**
   * Re-send a `session.focus` frame for an already-focused session WITHOUT
   * touching the focus ref-count. The backend's focus handler pushes a fresh
   * live git-status snapshot as a side effect (see `handleSessionFocus` →
   * `sendSessionData` → `appendLiveGitStatusMessage`, which forces a
   * `GetGitStatusMultiFresh` query that bypasses the workspace poll cache).
   *
   * This is the deterministic recovery for the focus-gated polling race: the
   * agent edits a file, but if the workspace is still in slow poll mode (the
   * focus→fast push lost a race with agentctl startup) the next
   * `session.git.event` may be up to 30s away. Components that become visible
   * and need fresh diff data (the diff panel becoming the active tab) call
   * this to pull the latest git status immediately.
   *
   * No-op when the session isn't currently focused (the regular focus frame,
   * sent on 0→1, already pushed fresh data).
   */
  refreshSessionData(sessionId: string) {
    if (this.status !== "connected") return;
    if (!this.sessionFocusCounts.get(sessionId)) return;
    this.send({
      id: generateUUID(),
      type: "request",
      action: "session.focus",
      payload: { session_id: 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);
  }

  subscribeSystemMetrics() {
    this.systemMetricsSubscriptionCount += 1;
    if (this.status === "connected" && this.systemMetricsSubscriptionCount === 1) {
      this.send({
        id: generateUUID(),
        type: "request",
        action: "system.metrics.subscribe",
        payload: {},
      });
    }
    return () => this.unsubscribeSystemMetrics();
  }

  unsubscribeSystemMetrics() {
    this.systemMetricsSubscriptionCount = Math.max(0, this.systemMetricsSubscriptionCount - 1);
    if (this.status === "connected" && this.systemMetricsSubscriptionCount === 0) {
      this.send({
        id: generateUUID(),
        type: "request",
        action: "system.metrics.unsubscribe",
        payload: {},
      });
    }
  }

  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 (!isDebug() || 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 (isDebug()) debugDispatch("response", { id: msgWithId.id });
      this.resolvePendingRequest(msgWithId.id, message.payload);
      return;
    }
    if (msgWithId.type === "error" && msgWithId.id) {
      if (isDebug()) 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));
    }
    // Plugin bridge: forward the same notification to any handlers a loaded
    // plugin registered for this action (registry.registerWsHandler).
    dispatchToPluginWsHandlers(action, message.payload);
  }

  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: {},
      });
    }
    if (this.systemMetricsSubscriptionCount > 0) {
      this.send({
        id: generateUUID(),
        type: "request",
        action: "system.metrics.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);
  }
}