ai-chatbot/lib/realtime/runware-ws.ts

type RunwareTask = {
  taskType: string;
  taskUUID?: string;
  [key: string]: unknown;
};

type RunwarePayload = {
  data?: Array<Record<string, unknown>>;
  errors?: Array<Record<string, unknown>>;
  [key: string]: unknown;
};

type PendingTask = {
  timeout: ReturnType<typeof setTimeout>;
  resolve: (value: RunwarePayload) => void;
  reject: (reason?: unknown) => void;
};

class RunwareWebSocketClient {
  private readonly endpoint = "wss://ws-api.runware.ai/v1";
  private socket: WebSocket | null = null;
  private connectionSessionUUID = "";
  private authenticated = false;
  private connectPromise: Promise<void> | null = null;
  private pingTimer: ReturnType<typeof setInterval> | null = null;
  private authWaiter:
    | {
        timeout: ReturnType<typeof setTimeout>;
        resolve: () => void;
        reject: (reason?: unknown) => void;
      }
    | null = null;
  private pending = new Map<string, PendingTask>();

  constructor(private readonly apiKey: string) {}

  private closeSocket() {
    if (!this.socket) return;
    try {
      this.socket.close();
    } catch {}
    this.socket = null;
  }

  private clearPingTimer() {
    if (!this.pingTimer) return;
    clearInterval(this.pingTimer);
    this.pingTimer = null;
  }

  private rejectPending(reason: unknown) {
    for (const [, task] of this.pending) {
      clearTimeout(task.timeout);
      task.reject(reason);
    }
    this.pending.clear();
  }

  private rejectAuthWaiter(reason: unknown) {
    if (!this.authWaiter) return;
    clearTimeout(this.authWaiter.timeout);
    this.authWaiter.reject(reason);
    this.authWaiter = null;
  }

  private startPing() {
    this.clearPingTimer();
    this.pingTimer = setInterval(() => {
      const socket = this.socket;
      if (!socket || socket.readyState !== WebSocket.OPEN) return;
      try {
        socket.send(
          JSON.stringify([
            {
              taskType: "ping",
              ping: true,
            },
          ])
        );
      } catch {}
    }, 55_000);
  }

  private sendAuth() {
    const socket = this.socket;
    if (!socket || socket.readyState !== WebSocket.OPEN) {
      throw new Error("Runware WebSocket 未连接");
    }
    const authTask: Record<string, unknown> = {
      taskType: "authentication",
      apiKey: this.apiKey,
    };
    if (this.connectionSessionUUID) {
      authTask.connectionSessionUUID = this.connectionSessionUUID;
    }
    socket.send(JSON.stringify([authTask]));
  }

  private parseMessage(raw: unknown) {
    if (typeof raw !== "string") return null;
    try {
      const parsed = JSON.parse(raw) as RunwarePayload;
      return parsed;
    } catch {
      return null;
    }
  }

  private handleMessage(raw: unknown) {
    const payload = this.parseMessage(raw);
    if (!payload) return;
    const dataList = Array.isArray(payload.data) ? payload.data : [];
    const errorList = Array.isArray(payload.errors) ? payload.errors : [];

    if (this.authWaiter) {
      const authError = errorList.find((item) => item?.taskType === "authentication");
      const authData = dataList.find((item) => item?.taskType === "authentication");
      if (authError) {
        const message = String(authError.message || "Runware WebSocket 鉴权失败");
        this.rejectAuthWaiter(new Error(message));
      } else if (authData) {
        const sessionId = String(authData.connectionSessionUUID || "").trim();
        if (sessionId) {
          this.connectionSessionUUID = sessionId;
        }
        this.authenticated = true;
        this.startPing();
        this.authWaiter.resolve();
        clearTimeout(this.authWaiter.timeout);
        this.authWaiter = null;
      }
    }

    for (const item of dataList) {
      const taskUUID = String(item?.taskUUID || "").trim();
      if (!taskUUID) continue;
      const pendingTask = this.pending.get(taskUUID);
      if (!pendingTask) continue;
      clearTimeout(pendingTask.timeout);
      this.pending.delete(taskUUID);
      pendingTask.resolve(payload);
    }

    for (const item of errorList) {
      const taskUUID = String(item?.taskUUID || "").trim();
      if (!taskUUID) continue;
      const pendingTask = this.pending.get(taskUUID);
      if (!pendingTask) continue;
      clearTimeout(pendingTask.timeout);
      this.pending.delete(taskUUID);
      const message = String(item?.message || item?.code || "Runware WebSocket 任务失败");
      pendingTask.reject(new Error(message));
    }
  }

  private handleSocketClose() {
    this.authenticated = false;
    this.clearPingTimer();
    this.rejectAuthWaiter(new Error("Runware WebSocket 连接已断开"));
    this.rejectPending(new Error("Runware WebSocket 连接已断开"));
    this.socket = null;
  }

  async ensureConnected(timeoutMs = 10_000) {
    if (this.socket?.readyState === WebSocket.OPEN && this.authenticated) {
      return;
    }
    if (this.connectPromise) {
      return this.connectPromise;
    }

    this.connectPromise = new Promise<void>((resolve, reject) => {
      this.closeSocket();
      this.authenticated = false;
      this.clearPingTimer();

      const socket = new WebSocket(this.endpoint);
      this.socket = socket;

      const cleanupAndReject = (error: unknown) => {
        this.handleSocketClose();
        reject(error instanceof Error ? error : new Error("Runware WebSocket 连接失败"));
      };

      socket.onmessage = (event) => {
        this.handleMessage(event.data);
      };

      socket.onerror = () => {
        if (!this.authenticated) {
          cleanupAndReject(new Error("Runware WebSocket 连接错误"));
        }
      };

      socket.onclose = () => {
        this.handleSocketClose();
      };

      socket.onopen = () => {
        this.authWaiter = {
          timeout: setTimeout(() => {
            this.rejectAuthWaiter(new Error("Runware WebSocket 鉴权超时"));
            this.closeSocket();
            reject(new Error("Runware WebSocket 鉴权超时"));
          }, timeoutMs),
          resolve,
          reject: (reason?: unknown) => {
            cleanupAndReject(reason);
          },
        };
        try {
          this.sendAuth();
        } catch (error) {
          cleanupAndReject(error);
        }
      };
    }).finally(() => {
      this.connectPromise = null;
    });

    return this.connectPromise;
  }

  async sendTask(task: RunwareTask, timeoutMs = 90_000) {
    const taskUUID = String(task.taskUUID || "").trim();
    if (!taskUUID) {
      throw new Error("Runware 任务缺少 taskUUID");
    }

    await this.ensureConnected();
    const socket = this.socket;
    if (!socket || socket.readyState !== WebSocket.OPEN || !this.authenticated) {
      throw new Error("Runware WebSocket 未就绪");
    }

    return await new Promise<RunwarePayload>((resolve, reject) => {
      const timeout = setTimeout(() => {
        this.pending.delete(taskUUID);
        reject(new Error("Runware WebSocket 推理超时"));
      }, timeoutMs);

      this.pending.set(taskUUID, {
        timeout,
        resolve,
        reject,
      });

      try {
        socket.send(JSON.stringify([task]));
      } catch (error) {
        clearTimeout(timeout);
        this.pending.delete(taskUUID);
        reject(error instanceof Error ? error : new Error("Runware WebSocket 发送失败"));
      }
    });
  }
}

declare global {
  // eslint-disable-next-line no-var
  var __runwareWebSocketPool: Map<string, RunwareWebSocketClient> | undefined;
}

function getRunwarePool() {
  if (!globalThis.__runwareWebSocketPool) {
    globalThis.__runwareWebSocketPool = new Map<string, RunwareWebSocketClient>();
  }
  return globalThis.__runwareWebSocketPool;
}

export function getRunwareWebSocketClient(apiKey: string) {
  const pool = getRunwarePool();
  let client = pool.get(apiKey);
  if (!client) {
    client = new RunwareWebSocketClient(apiKey);
    pool.set(apiKey, client);
  }
  return client;
}