All files / services websocket.service.ts

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