xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-09-18 01:06:39 +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
Show Stats 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 {

83
web/src/views/Chat.vue
View File

@ -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,6 +215,7 @@ export default defineComponent({
const socket = new WebSocket(process.env.VUE_APP_WS_HOST + `/api/chat?sessionId=${sessionId}&role=${this.role}`);
socket.addEventListener('open', () => {
// 获取聊天角色
if (this.chatRoles.length === 0) {
httpGet("/api/config/chat-roles/get").then((res) => {
// ElMessage.success('创建会话成功！');
this.chatRoles = res.data;
@ -221,14 +223,16 @@ export default defineComponent({
}).catch(() => {
ElMessage.error("获取聊天角色失败");
})
} else {
this.loading = false
}
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,40 +274,57 @@ export default defineComponent({
}
});
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(() => {
if (this.connectingMessageBox === null) {
this.connectingMessageBox = ElMessageBox.confirm(
'^_^ 会话发生异常，您已经从服务器断开连接!',
'注意：',
{
confirmButtonText: '重连会话',
cancelButtonText: '不聊了',
type: 'warning',
showClose: false,
closeOnClickModal: false
}
).then(() => {
// 自动重新连接
this.connect();
}).catch(() => {
ElMessage({
type: 'info',
message: '您关闭了会话',
})
})
}
// 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)
// 3 秒后继续重连
setTimeout(() => this.checkSession(), 3000)
}
})
});
this.socket = socket;
},
// 更换角色