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

import { createContext } from "react";
import { io, Socket } from "socket.io-client";

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;
}

export interface ClientToServerEvents {
  submitPrompt: (content: string) => void;
  requestSessionLock: (
    registration: any,
    project: any,
    chatSession: any,
    cb: (success: boolean, chatSessionId: 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 }) => void;
  "log:entry": (data: {
    level: string;
    message: string;
    timestamp: number;
  }) => void;
  "chat:message": (data: {
    agentId: string;
    message: string;
    role: "user" | "assistant" | "system";
  }) => 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;

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

  disconnect(): void {
    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);
    }
  }

  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);
      }
    });
  }
}

export const socketClient = new SocketClient();

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