From a78fb83f308d02a486c1db5adaad9c28d143ec4f Mon Sep 17 00:00:00 2001 From: RockYang Date: Tue, 4 Apr 2023 10:32:52 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=20websocket=20=E9=AA=8C?= =?UTF-8?q?=E8=AF=81=E5=92=8C=E6=96=AD=E7=BA=BF=E9=87=8D=E8=BF=9E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- server/chat_handler.go | 15 +++++- server/config_handler.go | 21 ++++++-- web/src/views/Chat.vue | 107 +++++++++++++++++++++++---------------- 3 files changed, 95 insertions(+), 48 deletions(-) diff --git a/server/chat_handler.go b/server/chat_handler.go index e8adf4ae..22cedfc9 100644 --- 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 } diff --git a/server/config_handler.go b/server/config_handler.go index b55fff89..8a841485 100644 --- 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 { diff --git a/web/src/views/Chat.vue b/web/src/views/Chat.vue index ec07db07..77d92aaa 100644 --- 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