api_health/websocket.js

import WebSocket, { WebSocketServer } from 'ws';
import http from 'http';
import { DBModel } from "./models/index.js";
import { chatTask } from "./agent/bots/bot-escort/task.js";
import { adminAgent } from "./agent/bots/bot-escort-admin/agent.js"
import { getUserInfo } from "./utils/user_service.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') {
        // 客户端断开后不再处理消息
        if (ws.readyState !== WebSocket.OPEN) return;

        const send = (source, type, content, id) => {
          if (ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({ source, type, content, id }));
          }
        };

        try {
          const userInfo = await getUserInfo(msg.token, msg.userId);
          if (!userInfo) {
            ws.send(JSON.stringify({ type: 'error', content: '登录已失效' }));
            return;
          }
          if (msg.agent === 'escort-admin') {
            // 管理端 agent 含写库/环境变量等高危工具，必须校验管理员角色，防止普通用户越权
            if (!userInfo.app || !("wxapp-escort-admin" in userInfo.app)) {
              ws.send(JSON.stringify({ type: 'error', content: '无管理员权限' }));
              return;
            }
            await adminAgent.streamChat(userInfo, [msg], (source, type, content, id) => {
              send(source, type, content, id);
            });
          } else {
            await chatTask.streamChat(userInfo, msg, (source, type, content, id) => {
              send(source, type, content, id);
            });
          }
        } catch (err) {
          console.error('streamChat error:', err);
          send('system', 'error', '处理消息失败');
        }
        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;
  }
}