优化 websocket 验证和断线重连

2025-11-04 16:23:42 +08:00 · 2023-04-04 10:32:52 +08:00
parent 65ad5fb632
commit a78fb83f30
3 changed files with 95 additions and 48 deletions
--- a/server/chat_handler.go
+++ b/server/chat_handler.go
@@ -27,7 +27,12 @@ func (s *Server) ChatHandle(c *gin.Context) {
 	}
 	sessionId := c.Query("sessionId")
 	roleKey := c.Query("role")
-	session := s.ChatSession[sessionId]
+	session, ok := s.ChatSession[sessionId]
+	if !ok { // 用户未登录
+		c.Abort()
+		return
+	}
+	
 	logger.Infof("New websocket connected, IP: %s, Username: %s", c.Request.RemoteAddr, session.Username)
 	client := NewWsClient(ws)
 	var roles = GetChatRoles()
@@ -64,7 +69,13 @@ func (s *Server) ChatHandle(c *gin.Context) {
 func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, prompt string, ws Client, resetContext bool) error {
 	user, err := GetUser(session.Username)
 	if err != nil {
-		replyMessage(ws, "当前 user 无效，请使用合法的 user 登录！", false)
+		replyMessage(ws, "当前 TOKEN 无效，请使用合法的 TOKEN 登录！", false)
+		return err
+	}
+
+	if user.Status == false {
+		replyMessage(ws, "当前 TOKEN 已经被禁用，如果疑问，请联系管理员！", false)
+		replyMessage(ws, "![](images/wx.png)", true)
 		return err
 	}

--- a/server/config_handler.go
+++ b/server/config_handler.go
@@ -16,7 +16,16 @@ func (s *Server) TestHandle(c *gin.Context) {
 		c.JSON(http.StatusBadRequest, nil)
 		return
 	}
-	c.JSON(http.StatusOK, types.BizVo{Code: types.Success, Data: data})
+	if v, ok := data["opt"]; ok && v == "init_user" {
+		users := GetUsers()
+		for _, user := range users {
+			user.Status = true
+			_ = PutUser(user)
+		}
+
+		c.JSON(http.StatusOK, types.BizVo{Code: types.Success, Data: GetUsers()})
+	}
+
 }

 // ConfigSetHandle set configs
@@ -102,7 +111,7 @@ func (s *Server) AddUserHandle(c *gin.Context) {
 		return
 	}

-	user := types.User{Name: data.Name, MaxCalls: data.MaxCalls, RemainingCalls: data.MaxCalls, EnableHistory: data.EnableHistory}
+	user := types.User{Name: data.Name, MaxCalls: data.MaxCalls, RemainingCalls: data.MaxCalls, EnableHistory: data.EnableHistory, Status: true}
 	err = PutUser(user)
 	if err != nil {
 		c.JSON(http.StatusOK, types.BizVo{Code: types.Failed, Message: "Failed to save configs"})
@@ -132,7 +141,7 @@ func (s *Server) BatchAddUserHandle(c *gin.Context) {
 		for err == nil {
 			name = utils.RandString(12)
 		}
-		err = PutUser(types.User{Name: name, MaxCalls: data.MaxCalls, RemainingCalls: data.MaxCalls, EnableHistory: data.EnableHistory})
+		err = PutUser(types.User{Name: name, MaxCalls: data.MaxCalls, RemainingCalls: data.MaxCalls, EnableHistory: data.EnableHistory, Status: true})
 		if err == nil {
 			users = append(users, name)
 		}
@@ -177,6 +186,12 @@ func (s *Server) SetUserHandle(c *gin.Context) {
 	if v, ok := data["enable_history"]; ok {
 		user.EnableHistory = v.(bool)
 	}
+	if v, ok := data["remaining_calls"]; ok {
+		user.RemainingCalls = v.(int)
+	}
+	if v, ok := data["api_key"]; ok {
+		user.ApiKey = v.(string)
+	}

 	err = PutUser(*user)
 	if err != nil {
--- a/web/src/views/Chat.vue
+++ b/web/src/views/Chat.vue
@@ -129,11 +129,12 @@ export default defineComponent({
      replyIcon: 'images/avatar/gpt.png', // 回复信息的头像

      lineBuffer: '', // 输出缓冲行
-      connectingMessageBox: null, // 保存重连的消息框对象load
+      connectingMessageBox: null, // 保存重连的消息框对象
+      errorMessage: null, // 错误信息提示框
      socket: null,
      toolBoxHeight: 61 + 42, // 工具框的高度
      inputBoxWidth: window.innerWidth - 20,
-      sending: false,
+      sending: true,
      loading: true
    }
  },
@@ -214,21 +215,24 @@ export default defineComponent({
      const socket = new WebSocket(process.env.VUE_APP_WS_HOST + `/api/chat?sessionId=${sessionId}&role=${this.role}`);
      socket.addEventListener('open', () => {
        // 获取聊天角色
-        httpGet("/api/config/chat-roles/get").then((res) => {
-          // ElMessage.success('创建会话成功！');
-          this.chatRoles = res.data;
+        if (this.chatRoles.length === 0) {
+          httpGet("/api/config/chat-roles/get").then((res) => {
+            // ElMessage.success('创建会话成功！');
+            this.chatRoles = res.data;
+            this.loading = false
+          }).catch(() => {
+            ElMessage.error("获取聊天角色失败");
+          })
+        } else {
          this.loading = false
-        }).catch(() => {
-          ElMessage.error("获取聊天角色失败");
-        })
+        }

+        this.sending = false; // 允许用户发送消息
+        if (this.errorMessage !== null) {
+          this.errorMessage.close(); // 关闭错误提示信息
+        }
        // 加载聊天记录
        this.fetchChatHistory();
-
-        if (this.connectingMessageBox && typeof this.connectingMessageBox.close === 'function') {
-          this.connectingMessageBox.close();
-          this.connectingMessageBox = null;
-        }
      });

      socket.addEventListener('message', event => {
@@ -270,42 +274,59 @@ export default defineComponent({
        }

      });
-      socket.addEventListener('close', () => {
-        // 检查会话
-        httpGet("/api/session/get").then(() => {
-          if (this.connectingMessageBox === null) {
-            this.connectingMessageBox = ElMessageBox.confirm(
-                '^_^ 会话发生异常，您已经从服务器断开连接!',
-                '注意：',
-                {
-                  confirmButtonText: '重连会话',
-                  cancelButtonText: '不聊了',
-                  type: 'warning',
-                  showClose: false,
-                  closeOnClickModal: false
-                }
-            ).then(() => {
-              this.connect();
-            }).catch(() => {
-              ElMessage({
-                type: 'info',
-                message: '您关闭了会话',
-              })
-            })
-          }
-        }).catch((res) => {
-          if (res.code === 400) {
-            this.showLoginDialog = true;
-          } else {
-            ElMessage.error(res.message)
-          }
-        })

+      socket.addEventListener('close', () => {
+        // 停止送消息
+        this.sending = true;
+        if (this.errorMessage === null) {
+          this.errorMessage = ElMessage({
+            message: '当前无法连接服务器，可检查网络设置是否正常',
+            type: 'error',
+            duration: 0,
+            showClose: false
+          });
+        }
+        this.checkSession();
      });

      this.socket = socket;
    },

+    checkSession: function () {
+      // 检查会话
+      httpGet("/api/session/get").then(() => {
+        // 自动重新连接
+        this.connect();
+        // if (this.connectingMessageBox === null) {
+        //   this.connectingMessageBox = ElMessageBox.confirm(
+        //       '^_^ 会话发生异常，您已经从服务器断开连接!',
+        //       '注意：',
+        //       {
+        //         confirmButtonText: '重连会话',
+        //         cancelButtonText: '不聊了',
+        //         type: 'warning',
+        //         showClose: false,
+        //         closeOnClickModal: false
+        //       }
+        //   ).then(() => {
+        //     this.connect();
+        //   }).catch(() => {
+        //     ElMessage({
+        //       type: 'info',
+        //       message: '您关闭了会话',
+        //     })
+        //   })
+        // }
+      }).catch((res) => {
+        if (res.code === 400) {
+          this.showLoginDialog = true;
+        } else {
+          // 3 秒后继续重连
+          setTimeout(() => this.checkSession(), 3000)
+        }
+      })
+    },
+
    // 更换角色
    changeRole: function () {
      this.loading = true