# All files / services websocket.service.ts

0% Statements 0/75
0% Branches 0/1
0% Functions 0/1
0% Lines 0/75

Press n or j to go to the next uncovered block, b, p or k for the previous block.

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85
    86
    87
    88
    89
    90
    91
    92
    93
    94
    95
    96
    97
    98
    99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136

     /**
     * WebSocket 服务
     * 用于推送音频/视频生成完成事件
     */

    import { Server as HttpServer } from 'http';
    import WebSocket, { WebSocketServer } from 'ws';

    const wss = new WebSocketServer({ noServer: true });

    // 客户端连接管理
    const clients = new Map<string, WebSocket>();

    // ============ 客户端管理 ============

    /**
     * 注册客户端连接
     */
    export function addClient(clientId: string, ws: WebSocket) {
      clients.set(clientId, ws);
      console.log(`[WS] 客户端连接: ${clientId}, 当前在线: ${clients.size}`);
    }

    /**
     * 移除客户端连接
     */
    export function removeClient(clientId: string) {
      clients.delete(clientId);
      console.log(`[WS] 客户端断开: ${clientId}, 当前在线: ${clients.size}`);
    }

    /**
     * 通过 clientId 发送消息
     */
    export function sendToClient(clientId: string, event: string, data: any): boolean {
      const ws = clients.get(clientId);
      if (!ws || ws.readyState !== WebSocket.OPEN) {
        return false;
      }
      try {
        ws.send(JSON.stringify({ event, data }));
        return true;
      } catch (error) {
        console.error(`[WS] 发送消息失败: ${clientId}`, error);
        return false;
      }
    }

    /**
     * 广播消息到所有客户端
     */
    export function broadcast(event: string, data: any) {
      const message = JSON.stringify({ event, data });
      clients.forEach((ws, clientId) => {
        if (ws.readyState === WebSocket.OPEN) {
          try {
            ws.send(message);
          } catch (error) {
            console.error(`[WS] 广播失败: ${clientId}`, error);
          }
        }
      });
    }

    // ============ 事件推送 ============

    /**
     * 推送音频生成完成事件
     */
    export function pushAudioGenerationComplete(bookId: string, chapterId: number, status: 'completed' | 'failed') {
      const event = 'audio_generation_complete';
      const data = { bookId, chapterId, status };
      console.log(`[WS] 推送 ${event}:`, data);
      // 广播给所有客户端（前端可根据 bookId 过滤）
      broadcast(event, data);
    }

    /**
     * 推送视频生成完成事件
     */
    export function pushVideoGenerationComplete(bookId: string, chapterId: number, status: 'completed' | 'failed') {
      const event = 'video_generation_complete';
      const data = { bookId, chapterId, status };
      console.log(`[WS] 推送 ${event}:`, data);
      broadcast(event, data);
    }

    /**
     * 推送批量生成进度
     */
    export function pushBatchGenerationProgress(taskId: string, step: string, progress: number) {
      const event = 'batch_generation_progress';
      const data = { taskId, step, progress };
      broadcast(event, data);
    }

    // ============ 初始化 ============

    /**
     * 初始化 WebSocket 服务
     */
    export function initWebSocket(server: HttpServer) {
      // 处理 HTTP upgrade 请求
      server.on('upgrade', (request, socket, head) => {
        const url = new URL(request.url || '', `http://${request.headers.host}`);

        // 只处理 /ws 路径的连接
        if (url.pathname !== '/ws') {
          socket.destroy();
          return;
        }

        wss.handleUpgrade(request, socket, head, (ws) => {
          // 从 query 参数获取 clientId
          const clientId = url.searchParams.get('clientId') || 'anonymous';
          addClient(clientId, ws);

          ws.on('close', () => {
            removeClient(clientId);
          });

          ws.on('error', (error) => {
            console.error(`[WS] 客户端错误: ${clientId}`, error);
            removeClient(clientId);
          });

          // 发送连接成功消息
          ws.send(JSON.stringify({ event: 'connected', data: { clientId } }));
        });
      });

      console.log('[WS] WebSocket 服务已初始化，挂载于 /ws 路径');
    }

    // 导出 wss 实例供外部使用
    export { wss };