xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-12-26 01:55:58 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

优化 websocket 验证和断线重连

Browse Source
This commit is contained in:
RockYang
2023-04-04 10:32:52 +08:00
parent 65ad5fb632
commit a78fb83f30
3 changed files with 95 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
server/chat_handler.go
View File

@@ -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
}

21
server/config_handler.go
View File

@@ -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 {