start

websocket.js

@@ -0,0 +1,131 @@
+import WebSocket, { WebSocketServer } from 'ws';
+import http from 'http';
+import { DBModel } from "./models/index.js";
+import { chatTask } from "./agent/task.js";
+
+export default class WebSocketServerManager {
+  constructor(port = 8080) {
+    this.port = port;
+    this.wss = null;
+    this.server = null;
+  }
+
+  start() {
+    this.server = http.createServer((req, res) => {
+      res.writeHead(426, { 'Content-Type': 'text/plain' });
+      res.end('Upgrade Required: WebSocket expected');
+    });
+
+    this.wss = new WebSocketServer({
+      server: this.server,
+      perMessageDeflate: {
+        zlibDeflateOptions: {
+          chunkSize: 1024,
+          memLevel: 7,
+          level: 3
+        },
+        zlibInflateOptions: {
+          chunkSize: 10 * 1024
+        },
+        clientNoContextTakeover: true,
+        serverNoContextTakeover: true,
+        serverMaxWindowBits: 10,
+        concurrencyLimit: 10,
+        threshold: 1024
+      }
+    });
+
+    this.wss.on('connection', (ws, req) => {
+      const path = req.url;
+      console.log(`WebSocket connected: ${req.socket.remoteAddress}, path: ${path}`);
+
+      if (path === '/chat') {
+        this._handleChatConnection(ws);
+      } else {
+        ws.send(JSON.stringify({ type: 'system', content: 'Unknown path' }));
+        ws.close(1000, 'Unknown path');
+      }
+    });
+
+    this.server.listen(this.port, () => {
+      console.log(`WebSocket server started on ws://localhost:${this.port}`);
+    });
+  }
+
+  _handleChatConnection(ws) {
+    ws.send(JSON.stringify({ type: 'system', content: 'Connected to chat server' }));
+
+    ws.on('message', (data) => this._handleMessage(ws, data));
+
+    ws.on('close', () => {
+      console.log('Chat client disconnected');
+    });
+
+    ws.on('error', (err) => {
+      console.error('WebSocket error:', err);
+    });
+  }
+
+  async _handleMessage(ws, data) {
+    try {
+      const raw = data.toString();
+      console.log('Received:', raw);
+
+      const msg = JSON.parse(raw);
+
+      if (msg.type === 'ping') {
+        ws.send(JSON.stringify({ type: 'pong' }));
+        return;
+      }
+
+      if (msg.type === 'chat' || msg.type === 'clear') {
+        const userInfo = await this.getUserInfo(msg.userId);
+        chatTask.streamChat(userInfo, msg, (source, type, content, id) => {
+          ws.send(JSON.stringify({ source, type, content, id }));
+        });
+        return;
+      }
+
+      if (msg.type === 'system') {
+        console.log('[System]', msg.content);
+        return;
+      }
+
+      console.log('[Unknown type]', msg);
+    } catch (e) {
+      console.error('Message parse error:', e);
+      ws.send(JSON.stringify({ type: 'error', content: 'Invalid JSON' }));
+    }
+  }
+
+  broadcast(message) {
+    if (!this.wss) return;
+    this.wss.clients.forEach((client) => {
+      if (client.readyState === WebSocket.OPEN) {
+        client.send(message);
+      }
+    });
+  }
+
+  stop() {
+    if (this.wss) {
+      this.wss.close();
+      this.wss = null;
+    }
+    if (this.server) {
+      this.server.close();
+      this.server = null;
+    }
+    console.log('WebSocket server stopped');
+  }
+
+  getClientsCount() {
+    if (!this.wss) return 0;
+    return this.wss.clients.size;
+  }
+
+  async getUserInfo(userId) {
+    if (!DBModel.User || userId.length === 0) return null;
+    return await DBModel.User.findById(userId).lean().exec();
+  }
+}