gadget/gadget-code/frontend/src/lib/socket.ts

import { createContext } from "react";
import { io, Socket } from "socket.io-client";
import type { ChatSession } from "./api";

/**
 * Web Worker for heartbeat timing — avoids browser tab throttling.
 * Chrome throttles setInterval in background tabs to ~1/min, which
 * would cause heartbeat timeouts against the drone's 120s timer.
 */
let HeartbeatWorker: typeof Worker | null = null;
try {
  // Vite supports `new Worker` with `?worker` import syntax
  HeartbeatWorker = Worker;
} catch {
  // Web Workers not available (extremely rare in modern browsers)
}

const SOCKET_URL = "";

export interface ServerToClientEvents {
  thinking: (content: string) => void;
  response: (content: string) => void;
  toolCall: (
    callId: string,
    name: string,
    params: string,
    response: string,
  ) => void;
  workOrderComplete: (
    turnId: string,
    success: boolean,
    message?: string,
  ) => void;
  log: (
    timestamp: string,
    component: { name: string; slug: string },
    level: string,
    message: string,
    metadata?: unknown,
  ) => void;
}

export interface ClientToServerEvents {
  submitPrompt: (content: string) => void;
  abortWorkOrder: (cb: (success: boolean, message?: string) => void) => void;
  requestSessionLock: (
    registration: any,
    project: any,
    chatSession: any,
    cb: (success: boolean, chatSessionId: string) => void,
  ) => void;
  requestWorkspaceMode: (
    registration: any,
    project: any,
    chatSession: any,
    mode: string,
    cb: (success: boolean, mode: string) => void,
  ) => void;
  releaseSessionLock: (
    registration: any,
    project: any,
    chatSession: any,
    cb: (success: boolean) => void,
  ) => void;
  sessionHeartbeat: (cb: (ack: boolean) => void) => void;
  fileTreeRequest: (
    args: { path?: string; recursive?: boolean; maxDepth?: number },
    cb: (success: boolean, data: { entries?: any[]; error?: string }) => void,
  ) => void;
  fileReadRequest: (
    args: { path: string },
    cb: (success: boolean, data: { content?: string; error?: string }) => void,
  ) => void;
  fileWriteRequest: (
    args: { path: string; content: string },
    cb: (success: boolean, data: { success?: boolean; error?: string }) => void,
  ) => void;
}

export interface SocketEvents {
  thinking: (content: string) => void;
  response: (content: string) => void;
  toolCall: (
    callId: string,
    name: string,
    params: string,
    response: string,
  ) => void;
  workOrderComplete: (
    turnId: string,
    success: boolean,
    message?: string,
  ) => void;
  "agent:thinking": (data: { agentId: string; thinking: string }) => void;
  "agent:response": (data: { agentId: string; chunk: string }) => void;
  "agent:tool-call": (data: {
    agentId: string;
    tool: string;
    args: unknown;
  }) => void;
  "agent:tool-result": (data: {
    agentId: string;
    tool: string;
    result: unknown;
  }) => void;
  "agent:complete": (data: {
    agentId: string;
    response?: string;
    subagent?: Record<string, unknown>;
    stats?: {
      toolCallCount: number;
      inputTokens: number;
      thinkingTokenCount: number;
      responseTokens: number;
      durationMs: number;
      durationLabel: string;
    };
  }) => void;
  "log:entry": (data: {
    level: string;
    component: string;
    message: string;
    metadata?: unknown;
    timestamp: number;
  }) => void;
  "chat:message": (data: {
    agentId: string;
    message: string;
    role: "user" | "assistant" | "system";
  }) => void;
  workspaceModeChanged: (mode: string) => void;
  sessionUpdated: (updates: Partial<ChatSession>) => void;
  tabLockDenied: (data: { message: string }) => void;
  fileTreeResponse: (path: string, entries: any[], error?: string) => void;
  fileReadResponse: (path: string, content: string | null, error?: string) => void;
  fileWriteResponse: (path: string, success: boolean, error?: string) => void;
  status: (content: string) => void;
  "drone:log": (data: {
    timestamp: string;
    level: string;
    component: string;
    message: string;
    metadata?: unknown;
  }) => void;
  "drone:status": (data: {
    timestamp: string;
    message: string;
  }) => void;
  reconnect_attempt: (attempt: number) => void;
  reconnect_failed: () => void;
  reconnect: (attempt: number) => void;
  connect: () => void;
  disconnect: (reason: string) => void;
  error: (error: Error) => void;
}

class SocketClient {
  private _socket: Socket | null = null;
  private eventListeners: Map<string, Set<(...args: unknown[]) => void>> =
    new Map();
  private reconnectAttempts = 0;
  private maxReconnectAttempts = 5;
  private jwt: string | null = null;

  /**
   * Web Worker for heartbeat timing. Runs the 19-second interval in a
   * worker thread to avoid browser tab throttling, which would cause
   * heartbeat timeouts against the drone's 120-second timer.
   */
  private heartbeatWorker: Worker | null = null;

  /**
   * Fallback setInterval for environments where Web Workers are unavailable.
   * This is subject to browser tab throttling but is better than nothing.
   */
  private heartbeatFallbackInterval: ReturnType<typeof setInterval> | null = null;

  /**
   * Bound visibility handler so we can remove it on disconnect.
   */
  private boundVisibilityHandler: (() => void) | null = null;

  get connected(): boolean {
    return this._socket?.connected ?? false;
  }

  get socket(): Socket | null {
    return this._socket;
  }

  connect(token: string): void {
    if (this._socket?.connected) {
      return;
    }

    this.jwt = token;

    this._socket = io(SOCKET_URL, {
      auth: {
        token: this.jwt,
      },
      transports: ["websocket", "polling"],
      reconnection: true,
      reconnectionAttempts: this.maxReconnectAttempts,
      reconnectionDelay: 1000,
      reconnectionDelayMax: 5000,
    });

    if (!this.socket) {
      return;
    }

    // Forward server events to our event listeners
    this.socket.on("thinking", (content: string) => {
      this.emit("thinking", content);
    });

    this.socket.on("response", (content: string) => {
      this.emit("response", content);
    });

    this.socket.on(
      "toolCall",
      (callId: string, name: string, params: string, response: string) => {
        this.emit("toolCall", callId, name, params, response);
      },
    );

    this.socket.on(
      "workOrderComplete",
      (turnId: string, success: boolean, message?: string) => {
        this.emit("workOrderComplete", turnId, success, message);
      },
    );

    this.socket.on("workspaceModeChanged", (mode: string) => {
      this.emit("workspaceModeChanged", mode);
    });

    this.socket.on("sessionUpdated", (updates: unknown) => {
      this.emit("sessionUpdated", updates as Partial<ChatSession>);
    });

    this.socket.on("fileTreeResponse", (path: string, entries: unknown[], error?: string) => {
      this.emit("fileTreeResponse", path, entries as any[], error);
    });

    this.socket.on("fileReadResponse", (path: string, content: string | null, error?: string) => {
      this.emit("fileReadResponse", path, content, error);
    });

    this.socket.on("fileWriteResponse", (path: string, success: boolean, error?: string) => {
      this.emit("fileWriteResponse", path, success, error);
    });

    this.socket.on(
      "log",
      (
        timestamp: string,
        component: { name: string; slug: string },
        level: string,
        message: string,
        metadata?: unknown,
      ) => {
        this.emit("log:entry", {
          level,
          component: component.name,
          message,
          metadata,
          timestamp: new Date(timestamp).getTime(),
        });
      },
    );

    this.socket.on("agent:thinking", (data: unknown) => {
      this.emit("agent:thinking", data as { agentId: string; thinking: string });
    });

    this.socket.on("agent:response", (data: unknown) => {
      this.emit("agent:response", data as { agentId: string; chunk: string });
    });

    this.socket.on("agent:tool-call", (data: unknown) => {
      this.emit("agent:tool-call", data as { agentId: string; tool: string; args: unknown });
    });

    this.socket.on("agent:tool-result", (data: unknown) => {
      this.emit("agent:tool-result", data as { agentId: string; tool: string; result: unknown });
    });

    this.socket.on("agent:complete", (data: unknown) => {
      this.emit("agent:complete", data as SocketEvents["agent:complete"] extends (data: infer T) => void ? T : never);
    });

    this._socket.on("connect", () => {
      this.reconnectAttempts = 0;
      this.emit("connect");
    });

    this._socket.on("disconnect", (reason) => {
      this.emit("disconnect", reason);
    });

    this._socket.on("connect_error", (error) => {
      this.reconnectAttempts++;
      this.emit("error", error);
    });

    this._socket.on("tabLockDenied", (data: { message: string }) => {
      this.emit("tabLockDenied", data);
    });

    this.socket.on("status", (content: string) => {
      this.emit("status", content);
    });

    this.socket.on("drone:log", (data: unknown) => {
      this.emit("drone:log", data as SocketEvents["drone:log"] extends (data: infer T) => void ? T : never);
    });

    this.socket.on("drone:status", (data: unknown) => {
      this.emit("drone:status", data as SocketEvents["drone:status"] extends (data: infer T) => void ? T : never);
    });

    this._socket.on("reconnect_attempt", (attempt: number) => {
      this.emit("reconnect_attempt", attempt);
    });

    this._socket.on("reconnect_failed", () => {
      this.emit("reconnect_failed");
    });

    this._socket.on("reconnect", (attempt: number) => {
      this.emit("reconnect", attempt);
    });
  }

  disconnect(): void {
    this.stopSessionHeartbeat();
    this.removeVisibilityHandler();
    if (this._socket) {
      this._socket.disconnect();
      this._socket = null;
      this.jwt = null;
    }
  }

  on<K extends keyof SocketEvents>(event: K, callback: SocketEvents[K]): void {
    if (!this.eventListeners.has(event)) {
      this.eventListeners.set(event, new Set());
    }
    this.eventListeners
      .get(event)!
      .add(callback as (...args: unknown[]) => void);
  }

  off<K extends keyof SocketEvents>(event: K, callback: SocketEvents[K]): void {
    const listeners = this.eventListeners.get(event);
    if (listeners) {
      listeners.delete(callback as (...args: unknown[]) => void);
    }
  }

  emit<K extends keyof SocketEvents>(
    event: K,
    ...args: Parameters<SocketEvents[K]>
  ): void {
    const listeners = this.eventListeners.get(event);
    if (listeners) {
      listeners.forEach((callback) => callback(...args));
    }
  }

  send<K extends keyof SocketEvents>(
    event: K,
    ...args: Parameters<SocketEvents[K]>
  ): void {
    if (this._socket?.connected) {
      this._socket.emit(event, ...args);
    }
  }

  emitServer<K extends keyof ClientToServerEvents>(
    event: K,
    ...args: Parameters<ClientToServerEvents[K]>
  ): void {
    if (this._socket?.connected) {
      this._socket.emit(event, ...args);
    }
  }

  submitPrompt(
    prompt: string,
    cb: (success: boolean, data: { turnId?: string; message?: string }) => void,
  ): void {
    if (this._socket?.connected) {
      this._socket.emit('submitPrompt', prompt, cb);
    } else {
      cb(false, { message: 'Socket not connected' });
    }
  }

  abortWorkOrder(cb?: (success: boolean, message?: string) => void): void {
    if (this._socket?.connected) {
      this._socket.emit('abortWorkOrder', (success: boolean, message?: string) => {
        cb?.(success, message);
      });
    } else {
      cb?.(false, 'Socket not connected');
    }
  }

  requestSessionLock(
    registration: any,
    project: any,
    chatSession: any,
  ): Promise<boolean> {
    return new Promise((resolve) => {
      if (this._socket?.connected) {
        this._socket.emit(
          "requestSessionLock",
          registration,
          project,
          chatSession,
          resolve,
        );
      } else {
        resolve(false);
      }
    });
  }

  requestWorkspaceMode(
    registration: any,
    project: any,
    chatSession: any,
    mode: string,
  ): Promise<{ success: boolean; mode: string; reason?: string }> {
    return new Promise((resolve) => {
      if (this._socket?.connected) {
        this._socket.emit(
          "requestWorkspaceMode",
          registration,
          project,
          chatSession,
          mode,
          (success: boolean, mode: string, reason?: string) =>
            resolve({ success, mode, reason }),
        );
      } else {
        resolve({ success: false, mode: "" });
      }
    });
  }

  requestFileTree(
    args: { path?: string; recursive?: boolean; maxDepth?: number } = {},
  ): Promise<{ success: boolean; entries?: any[]; error?: string }> {
    return new Promise((resolve) => {
      if (this._socket?.connected) {
        this._socket.emit(
          "fileTreeRequest",
          args,
          (success: boolean, data: { entries?: any[]; error?: string }) =>
            resolve({ success, ...data }),
        );
      } else {
        resolve({ success: false, error: "Socket not connected" });
      }
    });
  }

  requestFileRead(
    args: { path: string },
  ): Promise<{ success: boolean; content?: string; error?: string }> {
    return new Promise((resolve) => {
      if (this._socket?.connected) {
        this._socket.emit(
          "fileReadRequest",
          args,
          (success: boolean, data: { content?: string; error?: string }) =>
            resolve({ success, ...data }),
        );
      } else {
        resolve({ success: false, error: "Socket not connected" });
      }
    });
  }

  requestFileWrite(
    args: { path: string; content: string },
  ): Promise<{ success: boolean; error?: string }> {
    return new Promise((resolve) => {
      if (this._socket?.connected) {
        this._socket.emit(
          "fileWriteRequest",
          args,
          (success: boolean, data: { success?: boolean; error?: string }) =>
            resolve({ success: success, ...data }),
        );
      } else {
        resolve({ success: false, error: "Socket not connected" });
      }
    });
  }

  releaseSessionLock(
    registration: any,
    project: any,
    chatSession: any,
  ): Promise<boolean> {
    return new Promise((resolve) => {
      if (this._socket?.connected) {
        this._socket.emit(
          "releaseSessionLock",
          registration,
          project,
          chatSession,
          resolve,
        );
      } else {
        resolve(false);
      }
    });
  }

  requestProcessStats(
    registrationId: string,
  ): Promise<{
    success: boolean;
    processes?: { pid: number; command: string; args: string[]; projectId: string; projectSlug: string; projectName: string; status: string; createdAt: string; updatedAt: string; stdoutFileName: string; stderrFileName: string }[];
    message?: string;
  }> {
    return new Promise((resolve) => {
      if (this._socket?.connected) {
        this._socket.emit(
          "requestProcessStats",
          registrationId,
          (success: boolean, data?: { processes?: any[]; message?: string }) => {
            resolve({ success, processes: data?.processes, message: data?.message });
          },
        );
      } else {
        resolve({ success: false, message: "Socket not connected" });
      }
    });
  }

  subscribeDrone(registrationId: string): Promise<boolean> {
    return new Promise((resolve) => {
      if (this._socket?.connected) {
        this._socket.emit("subscribeDrone", registrationId, resolve);
      } else {
        resolve(false);
      }
    });
  }

  unsubscribeDrone(registrationId: string): Promise<boolean> {
    return new Promise((resolve) => {
      if (this._socket?.connected) {
        this._socket.emit("unsubscribeDrone", registrationId, resolve);
      } else {
        resolve(false);
      }
    });
  }

  /**
   * Sends a single sessionHeartbeat event to the server (which relays
   * it to the drone). The drone resets its 120-second timeout timer
   * on each heartbeat.
   */
  private sendHeartbeat(): void {
    if (this._socket?.connected) {
      this._socket.emit("sessionHeartbeat", (ack: boolean) => {
        if (!ack) {
          console.warn("sessionHeartbeat: drone did not acknowledge");
        }
      });
    }
  }

  startSessionHeartbeat(): void {
    // Already running?
    if (this.heartbeatWorker || this.heartbeatFallbackInterval) return;

    if (HeartbeatWorker) {
      // Use a Web Worker for the heartbeat interval — not subject to
      // browser tab throttling, which can delay setInterval to 1/min
      // in background tabs and cause drone heartbeat timeouts.
      this.heartbeatWorker = new HeartbeatWorker(
        new URL("../workers/heartbeat.worker.ts", import.meta.url),
        { type: "module" },
      );
      this.heartbeatWorker.onmessage = (event: MessageEvent) => {
        if (event.data?.type === "tick") {
          this.sendHeartbeat();
        }
      };
      this.heartbeatWorker.onerror = (err) => {
        console.error("heartbeat worker error, falling back to setInterval", err);
        this.stopSessionHeartbeat();
        this.startFallbackHeartbeat();
      };
      this.heartbeatWorker.postMessage({ type: "start" });
    } else {
      // Fallback: use setInterval (subject to tab throttling)
      this.startFallbackHeartbeat();
    }

    // Install visibility handler: send an immediate heartbeat when
    // the tab becomes visible again, to reset the drone's timer
    // after any throttling that may have occurred.
    this.installVisibilityHandler();
  }

  private startFallbackHeartbeat(): void {
    if (this.heartbeatFallbackInterval) return;
    this.heartbeatFallbackInterval = setInterval(() => {
      this.sendHeartbeat();
    }, 19000);
  }

  stopSessionHeartbeat(): void {
    if (this.heartbeatWorker) {
      this.heartbeatWorker.postMessage({ type: "stop" });
      this.heartbeatWorker.terminate();
      this.heartbeatWorker = null;
    }
    if (this.heartbeatFallbackInterval) {
      clearInterval(this.heartbeatFallbackInterval);
      this.heartbeatFallbackInterval = null;
    }
  }

  private installVisibilityHandler(): void {
    if (this.boundVisibilityHandler) return; // already installed
    this.boundVisibilityHandler = () => {
      if (document.visibilityState === "visible") {
        // Tab is now visible — send an immediate heartbeat to reset
        // the drone's timeout timer, which may have been starved
        // while the tab was backgrounded and timers were throttled.
        this.sendHeartbeat();
      }
    };
    document.addEventListener("visibilitychange", this.boundVisibilityHandler);
  }

  private removeVisibilityHandler(): void {
    if (this.boundVisibilityHandler) {
      document.removeEventListener("visibilitychange", this.boundVisibilityHandler);
      this.boundVisibilityHandler = null;
    }
  }

  /**
   * Update the stored JWT used for socket authentication.
   * Called after a successful token renewal so that future socket
   * reconnections use the new token instead of the expired one.
   */
  updateJwt(newToken: string): void {
    this.jwt = newToken;
    if (this._socket) {
      this._socket.auth = { ...this._socket.auth, token: newToken };
    }
  }
}

export const socketClient = new SocketClient();

export const SocketContext = createContext<Socket | null>(null);