ai review

websocket.js

@@ -1,9 +1,9 @@
 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"
+import { getUserInfo } from "./utils/user_service.js";
 
 export default class WebSocketServerManager {
   constructor(port = 8080) {
@@ -81,16 +81,33 @@ export default class WebSocketServerManager {
       }
 
       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) => {
+        // 客户端断开后不再处理消息
+        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 }));
-          });
-        } 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 }));
-          });
+          }
+        };
+
+        try {
+          const userInfo = await getUserInfo(msg.token, msg.userId);
+          if (!userInfo) {
+            ws.send(JSON.stringify({ type: 'error', content: '登录已失效' }));
+            return;
+          }
+          if (msg.agent === 'escort-admin') {
+            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;
       }
@@ -132,25 +149,4 @@ export default class WebSocketServerManager {
     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;
-    }
-  }
 }