type RunwareTask = { taskType: string; taskUUID?: string; [key: string]: unknown; }; type RunwarePayload = { data?: Array>; errors?: Array>; [key: string]: unknown; }; type PendingTask = { timeout: ReturnType; 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 | null = null; private pingTimer: ReturnType | null = null; private authWaiter: | { timeout: ReturnType; resolve: () => void; reject: (reason?: unknown) => void; } | null = null; private pending = new Map(); 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 = { 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((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((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 | undefined; } function getRunwarePool() { if (!globalThis.__runwareWebSocketPool) { globalThis.__runwareWebSocketPool = new Map(); } 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; }