shuzhiren-comfyui/backend/src/websocket-client/index.js

/**
 * websocket-client模块 - 与ComfyUI实例的WebSocket通信
 */

import WebSocket from 'ws';
import logger from '../logger/index.js';
import EventEmitter from 'events';
import comfyUIMonitor from '../comfyui-monitor/index.js';

class WebSocketClient extends EventEmitter {
  constructor() {
    super();
    this.connections = new Map();
  }

  /**
   * 连接到指定实例
   * @param {string} instanceId - 实例ID
   * @param {string} wsUrl - WebSocket地址
   * @returns {Promise<WebSocket>} WebSocket连接
   */
  connect(instanceId, wsUrl) {
    return new Promise((resolve, reject) => {
      if (this.connections.has(instanceId)) {
        const conn = this.connections.get(instanceId);
        if (conn.readyState === WebSocket.OPEN) {
          resolve(conn);
          return;
        }
      }

      logger.info(`正在连接到实例 ${instanceId}: ${wsUrl}`);
      
      const ws = new WebSocket(wsUrl);
      
      ws.on('open', () => {
        logger.info(`成功连接到实例 ${instanceId}`);
        this.connections.set(instanceId, ws);
        
        const stateChange = comfyUIMonitor.setInstanceState(instanceId, 'connected');
        if (stateChange) {
          const config = { wsUrl };
          comfyUIMonitor.logConnectionStateChange(
            instanceId,
            stateChange.oldState,
            'connected',
            'WebSocket连接成功',
            config
          );
        }
        
        resolve(ws);
      });

      ws.on('message', (data) => {
        try {
          const message = JSON.parse(data.toString());
          this.handleMessage(instanceId, message);
        } catch (error) {
          logger.error(`解析消息失败 (${instanceId}):`, error);
        }
      });

      ws.on('error', (error) => {
        logger.error(`WebSocket连接错误 (${instanceId}):`, error);
        
        const config = { wsUrl };
        comfyUIMonitor.logConnectionError(instanceId, error, config);
        
        reject(error);
      });

      ws.on('close', (code, reason) => {
        logger.warn(`与实例 ${instanceId} 的连接已关闭`);
        
        const stateChange = comfyUIMonitor.setInstanceState(instanceId, 'disconnected');
        if (stateChange) {
          const disconnectReason = reason ? reason.toString() : `关闭代码: ${code}`;
          const config = { wsUrl, closeCode: code };
          comfyUIMonitor.logConnectionStateChange(
            instanceId,
            stateChange.oldState,
            'disconnected',
            disconnectReason,
            config
          );
        }
        
        this.connections.delete(instanceId);
        this.emit('disconnected', { instanceId });
      });
    });
  }

  /**
   * 处理收到的消息
   * @param {string} instanceId - 实例ID
   * @param {object} message - 消息对象
   */
  handleMessage(instanceId, message) {
    this.emit('message', { instanceId, message });
    
    switch (message.type) {
      case 'status':
        this.emit('status', { instanceId, status: message.data });
        break;
      case 'progress':
        this.emit('progress', { instanceId, data: message.data });
        break;
      case 'execution_start':
        this.emit('execution_start', { instanceId, promptId: message.data.prompt_id });
        break;
      case 'execution_cached':
        this.emit('execution_cached', { instanceId, data: message.data });
        break;
      case 'executed':
        this.emit('executed', { instanceId, data: message.data });
        break;
      case 'execution_error':
        this.emit('execution_error', { instanceId, data: message.data });
        break;
    }
  }

  /**
   * 发送消息到指定实例
   * @param {string} instanceId - 实例ID
   * @param {object} message - 消息对象
   */
  send(instanceId, message) {
    const ws = this.connections.get(instanceId);
    if (!ws || ws.readyState !== WebSocket.OPEN) {
      throw new Error(`实例 ${instanceId} 未连接`);
    }
    ws.send(JSON.stringify(message));
  }

  /**
   * 断开指定实例的连接
   * @param {string} instanceId - 实例ID
   */
  disconnect(instanceId) {
    const ws = this.connections.get(instanceId);
    if (ws) {
      ws.close();
      this.connections.delete(instanceId);
    }
  }

  /**
   * 断开所有连接
   */
  disconnectAll() {
    for (const [instanceId, ws] of this.connections) {
      ws.close();
    }
    this.connections.clear();
  }

  /**
   * 检查实例是否已连接
   * @param {string} instanceId - 实例ID
   * @returns {boolean} 连接状态
   */
  isConnected(instanceId) {
    const ws = this.connections.get(instanceId);
    return ws && ws.readyState === WebSocket.OPEN;
  }
}

export default new WebSocketClient();