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

/**
 * websocket-client模块 - 与ComfyUI实例的WebSocket通信
 * 
 * 设计说明：
 * - clientId 使用实例 ID（固定不变），实现连接复用
 * - 同一实例的所有任务共享同一个 WebSocket 连接
 * - 通过 prompt_id 区分不同任务的消息
 */

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

  connect(instanceId, wsUrl) {
    return new Promise((resolve, reject) => {
      const existingConn = this.connections.get(instanceId);
      
      if (existingConn && existingConn.ws && existingConn.ws.readyState === WebSocket.OPEN) {
        logger.info(`[WebSocketClient] 实例 ${instanceId} 已有连接，直接复用`);
        resolve(existingConn.ws);
        return;
      }
      
      if (existingConn && existingConn.ws) {
        logger.info(`[WebSocketClient] 关闭旧连接，重新连接`);
        existingConn.ws.close();
        this.connections.delete(instanceId);
      }

      logger.info(`正在连接到实例 ${instanceId}: ${wsUrl}`);
      logger.info(`[WebSocketClient] 连接详情: instanceId=${instanceId}, wsUrl=${wsUrl}`);
      
      const ws = new WebSocket(wsUrl);
      
      ws.on('open', () => {
        logger.info(`成功连接到实例 ${instanceId}`);
        this.connections.set(instanceId, { ws, wsUrl });
        
        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 });
      });
    });
  }

  handleMessage(instanceId, message) {
    if (message.type !== 'progress_state' && message.type !== 'progress') {
      logger.info(`[WebSocketClient] 收到消息 from ${instanceId}: type=${message.type}, data=${JSON.stringify(message.data || {}).substring(0, 200)}`);
    }
    this.emit('message', { instanceId, message });
    
    switch (message.type) {
      case 'status':
        logger.info(`[WebSocketClient] status 消息: ${JSON.stringify(message.data)}`);
        this.emit('status', { instanceId, status: message.data });
        break;
      case 'progress':
        break;
      case 'execution_start':
        logger.info(`[WebSocketClient] execution_start 消息: prompt_id=${message.data?.prompt_id}`);
        this.emit('execution_start', { instanceId, promptId: message.data.prompt_id });
        break;
      case 'execution_cached':
        logger.info(`[WebSocketClient] execution_cached 消息: ${JSON.stringify(message.data)}`);
        this.emit('execution_cached', { instanceId, data: message.data });
        break;
      case 'executing':
        logger.info(`[WebSocketClient] executing 消息: node=${message.data?.node}, prompt_id=${message.data?.prompt_id}`);
        this.emit('executing', { instanceId, data: message.data });
        break;
      case 'executed':
        logger.info(`[WebSocketClient] executed 消息: prompt_id=${message.data?.prompt_id}`);
        this.emit('executed', { instanceId, data: message.data });
        break;
      case 'execution_success':
        logger.info(`[WebSocketClient] execution_success 消息: prompt_id=${message.data?.prompt_id}`);
        this.emit('execution_success', { instanceId, data: message.data });
        break;
      case 'execution_error':
        logger.error(`[WebSocketClient] execution_error 消息: ${JSON.stringify(message.data)}`);
        this.emit('execution_error', { instanceId, data: message.data });
        break;
      case 'progress_state':
        break;
      default:
        logger.info(`[WebSocketClient] 未处理的消息类型: ${message.type}`);
    }
  }

  send(instanceId, message) {
    const conn = this.connections.get(instanceId);
    const ws = conn?.ws;
    if (!ws || ws.readyState !== WebSocket.OPEN) {
      logger.error(`[WebSocketClient] 实例 ${instanceId} 未连接，无法发送消息`);
      throw new Error(`实例 ${instanceId} 未连接`);
    }
    const messageStr = JSON.stringify(message);
    logger.info(`[WebSocketClient] 发送消息到实例 ${instanceId}: ${messageStr.substring(0, 500)}${messageStr.length > 500 ? '...' : ''}`);
    ws.send(messageStr);
  }

  disconnect(instanceId) {
    const conn = this.connections.get(instanceId);
    if (conn && conn.ws) {
      conn.ws.close();
      this.connections.delete(instanceId);
    }
  }

  disconnectAll() {
    for (const [instanceId, conn] of this.connections) {
      if (conn.ws) {
        conn.ws.close();
      }
    }
    this.connections.clear();
  }

  isConnected(instanceId) {
    const conn = this.connections.get(instanceId);
    return conn?.ws && conn.ws.readyState === WebSocket.OPEN;
  }
}

export default new WebSocketClient();