api_health/websocket.js

import WebSocket, { WebSocketServer } from 'ws';
import http from 'http';
import fetch from 'node-fetch';
import { DBModel } from "./models/index.js";
import { chatTask } from "./agent/escort/task.js";
import { adminAgent } from "./agent/escort-admin/agent.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 (msg.agent === 'escort-admin') {
          const userInfo = await this.getUserInfo(msg.token, msg.userId);
          adminAgent.streamChat(userInfo, [msg], (source, type, content, id) => {
            ws.send(JSON.stringify({ source, type, content, id }));
          });
        } else {
          const userInfo = await this.getUserInfo(msg.token, 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(token, userId) {
    try {
      if (!token && !userId) return null;

      const url = "http://127.0.0.1:9010/user/userInfo";
      const res = await fetch(url, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          token,
          userId
        })
      });
      const data = await res.json();
      return data.data.user;
    } catch (error) {
      console.error('Error fetching user info:', error);
      return null;
    }
  }
}