refactor websocket message protocol, keep the only connection for all clients

2025-11-04 16:23:42 +08:00 · 2024-09-27 17:50:54 +08:00
parent 478bc32ddd
commit 2debe7e927
29 changed files with 407 additions and 567 deletions
--- a/api/core/app_server.go
+++ b/api/core/app_server.go
@@ -51,9 +51,9 @@ func NewServer(appConfig *types.AppConfig) *AppServer {
 func (s *AppServer) Init(debug bool, client *redis.Client) {
 	if debug { // 调试模式允许跨域请求 API
 		s.Debug = debug
+		s.Engine.Use(corsMiddleware())
 		logger.Info("Enabled debug mode")
 	}
-	s.Engine.Use(corsMiddleware())
 	s.Engine.Use(staticResourceMiddleware())
 	s.Engine.Use(authorizeMiddleware(s, client))
 	s.Engine.Use(parameterHandlerMiddleware())
@@ -101,9 +101,9 @@ func corsMiddleware() gin.HandlerFunc {
 			c.Header("Access-Control-Allow-Origin", origin)
 			c.Header("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, UPDATE")
 			//允许跨域设置可以返回其他子段，可以自定义字段
-			c.Header("Access-Control-Allow-Headers", "Authorization, Content-Length, Content-Type, Chat-Token, Admin-Authorization")
+			c.Header("Access-Control-Allow-Headers", "Authorization, Body-Length, Body-Type, Admin-Authorization,content-type")
 			// 允许浏览器（客户端）可以解析的头部 （重要）
-			c.Header("Access-Control-Expose-Headers", "Content-Length, Access-Control-Allow-Origin, Access-Control-Allow-Headers")
+			c.Header("Access-Control-Expose-Headers", "Body-Length, Access-Control-Allow-Origin, Access-Control-Allow-Headers")
 			//设置缓存时间
 			c.Header("Access-Control-Max-Age", "172800")
 			//允许客户端传递校验信息比如 cookie (重要)
@@ -131,7 +131,7 @@ func authorizeMiddleware(s *AppServer, client *redis.Client) gin.HandlerFunc {
 		isAdminApi := strings.Contains(c.Request.URL.Path, "/api/admin/")
 		if isAdminApi { // 后台管理 API
 			tokenString = c.GetHeader(types.AdminAuthHeader)
-		} else if c.Request.URL.Path == "/api/chat/new" {
+		} else if c.Request.URL.Path == "/api/ws" { // Websocket 连接
 			tokenString = c.Query("token")
 		} else {
 			tokenString = c.GetHeader(types.UserAuthHeader)
@@ -209,23 +209,18 @@ func needLogin(c *gin.Context) bool {
 		c.Request.URL.Path == "/api/app/list/user" ||
 		c.Request.URL.Path == "/api/model/list" ||
 		c.Request.URL.Path == "/api/mj/imgWall" ||
-		c.Request.URL.Path == "/api/mj/client" ||
 		c.Request.URL.Path == "/api/mj/notify" ||
 		c.Request.URL.Path == "/api/invite/hits" ||
 		c.Request.URL.Path == "/api/sd/imgWall" ||
-		c.Request.URL.Path == "/api/sd/client" ||
 		c.Request.URL.Path == "/api/dall/imgWall" ||
-		c.Request.URL.Path == "/api/dall/client" ||
 		c.Request.URL.Path == "/api/product/list" ||
 		c.Request.URL.Path == "/api/menu/list" ||
 		c.Request.URL.Path == "/api/markMap/client" ||
 		c.Request.URL.Path == "/api/payment/doPay" ||
 		c.Request.URL.Path == "/api/payment/payWays" ||
-		c.Request.URL.Path == "/api/suno/client" ||
 		c.Request.URL.Path == "/api/suno/detail" ||
 		c.Request.URL.Path == "/api/suno/play" ||
 		c.Request.URL.Path == "/api/download" ||
-		c.Request.URL.Path == "/api/video/client" ||
 		strings.HasPrefix(c.Request.URL.Path, "/api/test") ||
 		strings.HasPrefix(c.Request.URL.Path, "/api/payment/notify/") ||
 		strings.HasPrefix(c.Request.URL.Path, "/api/user/clogin") ||
--- a/api/core/types/chat.go
+++ b/api/core/types/chat.go
@@ -52,7 +52,6 @@ type Delta struct {

 // ChatSession 聊天会话对象
 type ChatSession struct {
-	SessionId string    `json:"session_id"`
 	UserId   uint      `json:"user_id"`
 	ClientIP string    `json:"client_ip"` // 客户端 IP
 	ChatId   string    `json:"chat_id"`   // 客户端聊天会话 ID, 多会话模式专用字段
--- a/api/core/types/client.go
+++ b/api/core/types/client.go
@@ -17,15 +17,17 @@ var ErrConClosed = errors.New("connection Closed")

 // WsClient websocket client
 type WsClient struct {
+	Id     string
 	Conn   *websocket.Conn
 	lock   sync.Mutex
 	mt     int
 	Closed bool
 }

-func NewWsClient(conn *websocket.Conn) *WsClient {
+func NewWsClient(conn *websocket.Conn, id string) *WsClient {
 	return &WsClient{
 		Conn:   conn,
+		Id:     id,
 		lock:   sync.Mutex{},
 		mt:     2, // fixed bug for 'Invalid UTF-8 in text frame'
 		Closed: false,
--- a/api/core/types/web.go
+++ b/api/core/types/web.go
@@ -20,34 +20,43 @@ type BizVo struct {
 // ReplyMessage 对话回复消息结构
 type ReplyMessage struct {
 	Channel  WsChannel   `json:"channel"`  // 消息频道，目前只有 chat
+	ClientId string      `json:"clientId"` // 客户端ID
 	Type     WsMsgType   `json:"type"`     // 消息类别
-	Content interface{} `json:"content"`
+	Body     interface{} `json:"body"`
 }

 type WsMsgType string
 type WsChannel string

 const (
-	WsMsgTypeContent = WsMsgType("content") // 输出内容
-	WsMsgTypeEnd     = WsMsgType("end")
-	WsMsgTypeErr     = WsMsgType("error")
-	WsMsgTypePing    = WsMsgType("ping") // 心跳消息
+	MsgTypeText = WsMsgType("text") // 输出内容
+	MsgTypeEnd  = WsMsgType("end")
+	MsgTypeErr  = WsMsgType("error")
+	MsgTypePing = WsMsgType("ping") // 心跳消息

-	WsChat = WsChannel("chat")
-	WsMj   = WsChannel("mj")
-	WsSd   = WsChannel("sd")
-	WsDall = WsChannel("dall")
-	WsSuno = WsChannel("suno")
-	WsLuma = WsChannel("luma")
+	ChPing = WsChannel("ping")
+	ChChat = WsChannel("chat")
+	ChMj   = WsChannel("mj")
+	ChSd   = WsChannel("sd")
+	ChDall = WsChannel("dall")
+	ChSuno = WsChannel("suno")
+	ChLuma = WsChannel("luma")
 )

 // InputMessage 对话输入消息结构
 type InputMessage struct {
 	Channel WsChannel   `json:"channel"` // 消息频道
 	Type    WsMsgType   `json:"type"`    // 消息类别
+	Body    interface{} `json:"body"`
+}
+
+type ChatMessage struct {
+	Tools   []int  `json:"tools,omitempty"`  // 允许调用工具列表
+	Stream  bool   `json:"stream,omitempty"` // 是否采用流式输出
+	RoleId  int    `json:"role_id"`
+	ModelId int    `json:"model_id"`
+	ChatId  string `json:"chat_id"`
 	Content string `json:"content"`
-	Tools   []int     `json:"tools"`  // 允许调用工具列表
-	Stream  bool      `json:"stream"` // 是否采用流式输出
 }

 type BizCode int
--- a/api/handler/chatimpl/chat_handler.go
+++ b/api/handler/chatimpl/chat_handler.go
@@ -1,4 +1,4 @@
-package chatimpl
+package handler

 // * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 // * Copyright 2023 The Geek-AI Authors. All rights reserved.
@@ -15,8 +15,6 @@ import (
 	"fmt"
 	"geekai/core"
 	"geekai/core/types"
-	"geekai/handler"
-	logger2 "geekai/logger"
 	"geekai/service"
 	"geekai/service/oss"
 	"geekai/store/model"
@@ -33,14 +31,11 @@ import (

 	"github.com/gin-gonic/gin"
 	"github.com/go-redis/redis/v8"
-	"github.com/gorilla/websocket"
 	"gorm.io/gorm"
 )

-var logger = logger2.GetLogger()
-
 type ChatHandler struct {
-	handler.BaseHandler
+	BaseHandler
 	redis          *redis.Client
 	uploadManager  *oss.UploaderManager
 	licenseService *service.LicenseService
@@ -51,7 +46,7 @@ type ChatHandler struct {

 func NewChatHandler(app *core.AppServer, db *gorm.DB, redis *redis.Client, manager *oss.UploaderManager, licenseService *service.LicenseService, userService *service.UserService) *ChatHandler {
 	return &ChatHandler{
-		BaseHandler:    handler.BaseHandler{App: app, DB: db},
+		BaseHandler:    BaseHandler{App: app, DB: db},
 		redis:          redis,
 		uploadManager:  manager,
 		licenseService: licenseService,
@@ -61,106 +56,6 @@ func NewChatHandler(app *core.AppServer, db *gorm.DB, redis *redis.Client, manag
 	}
 }

-// ChatHandle 处理聊天 WebSocket 请求
-func (h *ChatHandler) ChatHandle(c *gin.Context) {
-	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
-	if err != nil {
-		logger.Error(err)
-		return
-	}
-
-	sessionId := c.Query("session_id")
-	roleId := h.GetInt(c, "role_id", 0)
-	chatId := c.Query("chat_id")
-	modelId := h.GetInt(c, "model_id", 0)
-
-	client := types.NewWsClient(ws)
-	var chatRole model.ChatRole
-	res := h.DB.First(&chatRole, roleId)
-	if res.Error != nil || !chatRole.Enable {
-		utils.ReplyErrorMessage(client, "当前聊天角色不存在或者未启用，对话已关闭！！！")
-		c.Abort()
-		return
-	}
-	// if the role bind a model_id, use role's bind model_id
-	if chatRole.ModelId > 0 {
-		modelId = chatRole.ModelId
-	}
-	// get model info
-	var chatModel model.ChatModel
-	res = h.DB.First(&chatModel, modelId)
-	if res.Error != nil || chatModel.Enabled == false {
-		utils.ReplyErrorMessage(client, "当前AI模型暂未启用，对话已关闭！！！")
-		c.Abort()
-		return
-	}
-
-	session := &types.ChatSession{
-		SessionId: sessionId,
-		ClientIP:  c.ClientIP(),
-		UserId:    h.GetLoginUserId(c),
-	}
-
-	// use old chat data override the chat model and role ID
-	var chat model.ChatItem
-	res = h.DB.Where("chat_id = ?", chatId).First(&chat)
-	if res.Error == nil {
-		chatModel.Id = chat.ModelId
-		roleId = int(chat.RoleId)
-	}
-
-	session.ChatId = chatId
-	session.Model = types.ChatModel{
-		Id:          chatModel.Id,
-		Name:        chatModel.Name,
-		Value:       chatModel.Value,
-		Power:       chatModel.Power,
-		MaxTokens:   chatModel.MaxTokens,
-		MaxContext:  chatModel.MaxContext,
-		Temperature: chatModel.Temperature,
-		KeyId:       chatModel.KeyId}
-	logger.Infof("New websocket connected, IP: %s", c.ClientIP())
-
-	go func() {
-		for {
-			_, msg, err := client.Receive()
-			if err != nil {
-				logger.Debugf("close connection: %s", client.Conn.RemoteAddr())
-				client.Close()
-				cancelFunc := h.ReqCancelFunc.Get(sessionId)
-				if cancelFunc != nil {
-					cancelFunc()
-					h.ReqCancelFunc.Delete(sessionId)
-				}
-				return
-			}
-
-			var message types.InputMessage
-			err = utils.JsonDecode(string(msg), &message)
-			if err != nil {
-				continue
-			}
-
-			logger.Infof("Receive a message:%+v", message)
-
-			session.Tools = message.Tools
-			session.Stream = message.Stream
-			ctx, cancel := context.WithCancel(context.Background())
-			h.ReqCancelFunc.Put(sessionId, cancel)
-			// 回复消息
-			err = h.sendMessage(ctx, session, chatRole, utils.InterfaceToString(message.Content), client)
-			if err != nil {
-				logger.Error(err)
-				utils.SendMessage(client, err.Error())
-			} else {
-				utils.SendChunkMessage(client, types.ReplyMessage{Type: types.WsMsgTypeEnd})
-				logger.Infof("回答完毕: %v", message.Content)
-			}
-
-		}
-	}()
-}
-
 func (h *ChatHandler) sendMessage(ctx context.Context, session *types.ChatSession, role model.ChatRole, prompt string, ws *types.WsClient) error {
 	if !h.App.Debug {
 		defer func() {
@@ -206,7 +101,7 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session *types.ChatSessio
 	}
 	// 兼容 GPT-O1 模型
 	if strings.HasPrefix(session.Model.Value, "o1-") {
-		utils.ReplyContent(ws, "AI 正在思考...\n")
+		utils.SendChunkMsg(ws, "AI 正在思考...\n")
 		req.Stream = false
 		session.Start = time.Now().Unix()
 	} else {
--- a/api/handler/chatimpl/chat_item_handler.go
+++ b/api/handler/chatimpl/chat_item_handler.go
@@ -1,4 +1,4 @@
-package chatimpl
+package handler

 // * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 // * Copyright 2023 The Geek-AI Authors. All rights reserved.
@@ -28,32 +28,41 @@ func (h *ChatHandler) List(c *gin.Context) {
 	userId := h.GetLoginUserId(c)
 	var items = make([]vo.ChatItem, 0)
 	var chats []model.ChatItem
-	res := h.DB.Where("user_id = ?", userId).Order("id DESC").Find(&chats)
-	if res.Error == nil {
+	h.DB.Where("user_id", userId).Order("id DESC").Find(&chats)
+	if len(chats) == 0 {
+		resp.SUCCESS(c, items)
+		return
+	}
+
 	var roleIds = make([]uint, 0)
+	var modelValues = make([]string, 0)
 	for _, chat := range chats {
 		roleIds = append(roleIds, chat.RoleId)
+		modelValues = append(modelValues, chat.Model)
 	}
+
 	var roles []model.ChatRole
-		res = h.DB.Find(&roles, roleIds)
-		if res.Error == nil {
+	var models []model.ChatModel
 	roleMap := make(map[uint]model.ChatRole)
+	modelMap := make(map[string]model.ChatModel)
+	h.DB.Where("id IN ?", roleIds).Find(&roles)
+	h.DB.Where("value IN ?", modelValues).Find(&models)
 	for _, role := range roles {
 		roleMap[role.Id] = role
 	}
-
+	for _, m := range models {
+		modelMap[m.Value] = m
+	}
 	for _, chat := range chats {
 		var item vo.ChatItem
 		err := utils.CopyObject(chat, &item)
 		if err == nil {
 			item.Id = chat.Id
 			item.Icon = roleMap[chat.RoleId].Icon
+			item.ModelId = modelMap[chat.Model].Id
 			items = append(items, item)
 		}
 	}
-		}
-
-	}
 	resp.SUCCESS(c, items)
 }

--- a/api/handler/dalle_handler.go
+++ b/api/handler/dalle_handler.go
@@ -20,9 +20,7 @@ import (
 	"geekai/utils/resp"
 	"github.com/gin-gonic/gin"
 	"github.com/go-redis/redis/v8"
-	"github.com/gorilla/websocket"
 	"gorm.io/gorm"
-	"net/http"
 )

 type DallJobHandler struct {
@@ -45,49 +43,6 @@ func NewDallJobHandler(app *core.AppServer, db *gorm.DB, service *dalle.Service,
 	}
 }

-// Client WebSocket 客户端，用于通知任务状态变更
-func (h *DallJobHandler) Client(c *gin.Context) {
-	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
-	if err != nil {
-		logger.Error(err)
-		c.Abort()
-		return
-	}
-
-	userId := h.GetInt(c, "user_id", 0)
-	if userId == 0 {
-		logger.Info("Invalid user ID")
-		c.Abort()
-		return
-	}
-
-	client := types.NewWsClient(ws)
-	h.dallService.Clients.Put(uint(userId), client)
-	logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
-	go func() {
-		for {
-			_, msg, err := client.Receive()
-			if err != nil {
-				client.Close()
-				h.dallService.Clients.Delete(uint(userId))
-				return
-			}
-
-			var message types.ReplyMessage
-			err = utils.JsonDecode(string(msg), &message)
-			if err != nil {
-				continue
-			}
-
-			// 心跳消息
-			if message.Type == "heartbeat" {
-				logger.Debug("收到 DallE 心跳消息：", message.Content)
-				continue
-			}
-		}
-	}()
-}
-
 func (h *DallJobHandler) preCheck(c *gin.Context) bool {
 	user, err := h.GetLoginUser(c)
 	if err != nil {
--- a/api/handler/markmap_handler.go
+++ b/api/handler/markmap_handler.go
@@ -19,7 +19,6 @@ import (
 	"geekai/store/model"
 	"geekai/utils"
 	"github.com/gin-gonic/gin"
-	"github.com/gorilla/websocket"
 	"gorm.io/gorm"
 	"io"
 	"net/http"
@@ -43,55 +42,9 @@ func NewMarkMapHandler(app *core.AppServer, db *gorm.DB, userService *service.Us
 	}
 }

-func (h *MarkMapHandler) Client(c *gin.Context) {
-	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
-	if err != nil {
-		logger.Error(err)
-		return
-	}
+// Generate 生成思维导图
+func (h *MarkMapHandler) Generate(c *gin.Context) {

-	modelId := h.GetInt(c, "model_id", 0)
-	userId := h.GetInt(c, "user_id", 0)
-
-	client := types.NewWsClient(ws)
-	h.clients.Put(userId, client)
-	go func() {
-		for {
-			_, msg, err := client.Receive()
-			if err != nil {
-				client.Close()
-				h.clients.Delete(userId)
-				return
-			}
-
-			var message types.ReplyMessage
-			err = utils.JsonDecode(string(msg), &message)
-			if err != nil {
-				continue
-			}
-
-			// 心跳消息
-			if message.Type == "heartbeat" {
-				logger.Debug("收到 MarkMap 心跳消息：", message.Content)
-				continue
-			}
-			// change model
-			if message.Type == "model_id" {
-				modelId = utils.IntValue(utils.InterfaceToString(message.Content), 0)
-				continue
-			}
-
-			logger.Info("Receive a message: ", message.Content)
-			err = h.sendMessage(client, utils.InterfaceToString(message.Content), modelId, userId)
-			if err != nil {
-				logger.Error(err)
-				utils.ReplyErrorMessage(client, err.Error())
-			} else {
-				utils.SendMessage(client, types.ReplyMessage{Type: types.WsMsgTypeEnd})
-			}
-
-		}
-	}()
 }

 func (h *MarkMapHandler) sendMessage(client *types.WsClient, prompt string, modelId int, userId int) error {
@@ -170,13 +123,13 @@ func (h *MarkMapHandler) sendMessage(client *types.WsClient, prompt string, mode
 				break
 			}

-			utils.SendChunkMessage(client, types.ReplyMessage{
-				Type:    types.WsMsgTypeContent,
-				Content: utils.InterfaceToString(responseBody.Choices[0].Delta.Content),
+			utils.SendMsg(client, types.ReplyMessage{
+				Type: types.MsgTypeText,
+				Body: utils.InterfaceToString(responseBody.Choices[0].Delta.Content),
 			})
 		} // end for

-		utils.SendChunkMessage(client, types.ReplyMessage{Type: types.WsMsgTypeEnd})
+		utils.SendMsg(client, types.ReplyMessage{Type: types.MsgTypeEnd})

 	} else {
 		body, _ := io.ReadAll(response.Body)
--- a/api/handler/mj_handler.go
+++ b/api/handler/mj_handler.go
@@ -19,12 +19,10 @@ import (
 	"geekai/store/vo"
 	"geekai/utils"
 	"geekai/utils/resp"
-	"net/http"
 	"strings"
 	"time"

 	"github.com/gin-gonic/gin"
-	"github.com/gorilla/websocket"
 	"gorm.io/gorm"
 )

@@ -65,27 +63,6 @@ func (h *MidJourneyHandler) preCheck(c *gin.Context) bool {

 }

-// Client WebSocket 客户端，用于通知任务状态变更
-func (h *MidJourneyHandler) Client(c *gin.Context) {
-	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
-	if err != nil {
-		logger.Error(err)
-		c.Abort()
-		return
-	}
-
-	userId := h.GetInt(c, "user_id", 0)
-	if userId == 0 {
-		logger.Info("Invalid user ID")
-		c.Abort()
-		return
-	}
-
-	client := types.NewWsClient(ws)
-	h.mjService.Clients.Put(uint(userId), client)
-	logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
-}
-
 // Image 创建一个绘画任务
 func (h *MidJourneyHandler) Image(c *gin.Context) {
 	var data struct {
--- a/api/handler/chatimpl/openai_handler.go
+++ b/api/handler/chatimpl/openai_handler.go
@@ -1,4 +1,4 @@
-package chatimpl
+package handler

 // * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 // * Copyright 2023 The Geek-AI Authors. All rights reserved.
@@ -108,7 +108,7 @@ func (h *ChatHandler) sendOpenAiMessage(
 			}

 			if responseBody.Choices[0].FinishReason == "stop" && len(contents) == 0 {
-				utils.SendMessage(ws, "抱歉😔😔😔，AI助手由于未知原因已经停止输出内容。")
+				utils.SendChunkMsg(ws, "抱歉😔😔😔，AI助手由于未知原因已经停止输出内容。")
 				break
 			}

@@ -136,7 +136,7 @@ func (h *ChatHandler) sendOpenAiMessage(
 				if res.Error == nil {
 					toolCall = true
 					callMsg := fmt.Sprintf("正在调用工具 `%s` 作答 ...\n\n", function.Label)
-					utils.SendChunkMessage(ws, types.ReplyMessage{Type: types.WsMsgTypeContent, Content: callMsg})
+					utils.SendChunkMsg(ws, callMsg)
 					contents = append(contents, callMsg)
 				}
 				continue
@@ -153,10 +153,7 @@ func (h *ChatHandler) sendOpenAiMessage(
 			} else {
 				content := responseBody.Choices[0].Delta.Content
 				contents = append(contents, utils.InterfaceToString(content))
-				utils.SendChunkMessage(ws, types.ReplyMessage{
-					Type:    types.WsMsgTypeContent,
-					Content: utils.InterfaceToString(responseBody.Choices[0].Delta.Content),
-				})
+				utils.SendChunkMsg(ws, responseBody.Choices[0].Delta.Content)
 			}
 		} // end for

@@ -174,7 +171,7 @@ func (h *ChatHandler) sendOpenAiMessage(
 			logger.Debugf("函数名称: %s, 函数参数：%s", function.Name, params)
 			params["user_id"] = userVo.Id
 			var apiRes types.BizVo
-			r, err := req2.C().R().SetHeader("Content-Type", "application/json").
+			r, err := req2.C().R().SetHeader("Body-Type", "application/json").
 				SetHeader("Authorization", function.Token).
 				SetBody(params).
 				SetSuccessResult(&apiRes).Post(function.Action)
@@ -185,19 +182,13 @@ func (h *ChatHandler) sendOpenAiMessage(
 				errMsg = r.Status
 			}
 			if errMsg != "" || apiRes.Code != types.Success {
-				msg := "调用函数工具出错：" + apiRes.Message + errMsg
-				utils.SendChunkMessage(ws, types.ReplyMessage{
-					Type:    types.WsMsgTypeContent,
-					Content: msg,
-				})
-				contents = append(contents, msg)
+				errMsg = "调用函数工具出错：" + apiRes.Message + errMsg
+				contents = append(contents, errMsg)
 			} else {
-				utils.SendChunkMessage(ws, types.ReplyMessage{
-					Type:    types.WsMsgTypeContent,
-					Content: apiRes.Data,
-				})
-				contents = append(contents, utils.InterfaceToString(apiRes.Data))
+				errMsg = utils.InterfaceToString(apiRes.Data)
+				contents = append(contents, errMsg)
 			}
+			utils.SendChunkMsg(ws, errMsg)
 		}

 		// 消息发送成功
@@ -226,7 +217,7 @@ func (h *ChatHandler) sendOpenAiMessage(
 		if strings.HasPrefix(req.Model, "o1-") {
 			content = fmt.Sprintf("AI思考结束，耗时：%d 秒。\n%s", time.Now().Unix()-session.Start, respVo.Choices[0].Message.Content)
 		}
-		utils.SendMessage(ws, content)
+		utils.SendChunkMsg(ws, content)
 		respVo.Usage.Prompt = prompt
 		respVo.Usage.Content = content
 		h.saveChatHistory(req, respVo.Usage, respVo.Choices[0].Message, chatCtx, session, role, userVo, promptCreatedAt, time.Now())
--- a/api/handler/sd_handler.go
+++ b/api/handler/sd_handler.go
@@ -19,11 +19,8 @@ import (
 	"geekai/store/vo"
 	"geekai/utils"
 	"geekai/utils/resp"
-	"net/http"
 	"time"

-	"github.com/gorilla/websocket"
-
 	"github.com/gin-gonic/gin"
 	"github.com/go-redis/redis/v8"
 	"gorm.io/gorm"
@@ -59,27 +56,6 @@ func NewSdJobHandler(app *core.AppServer,
 	}
 }

-// Client WebSocket 客户端，用于通知任务状态变更
-func (h *SdJobHandler) Client(c *gin.Context) {
-	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
-	if err != nil {
-		logger.Error(err)
-		c.Abort()
-		return
-	}
-
-	userId := h.GetInt(c, "user_id", 0)
-	if userId == 0 {
-		logger.Info("Invalid user ID")
-		c.Abort()
-		return
-	}
-
-	client := types.NewWsClient(ws)
-	h.sdService.Clients.Put(uint(userId), client)
-	logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
-}
-
 func (h *SdJobHandler) preCheck(c *gin.Context) bool {
 	user, err := h.GetLoginUser(c)
 	if err != nil {
--- a/api/handler/suno_handler.go
+++ b/api/handler/suno_handler.go
@@ -19,9 +19,7 @@ import (
 	"geekai/utils"
 	"geekai/utils/resp"
 	"github.com/gin-gonic/gin"
-	"github.com/gorilla/websocket"
 	"gorm.io/gorm"
-	"net/http"
 	"time"
 )

@@ -44,27 +42,6 @@ func NewSunoHandler(app *core.AppServer, db *gorm.DB, service *suno.Service, upl
 	}
 }

-// Client WebSocket 客户端，用于通知任务状态变更
-func (h *SunoHandler) Client(c *gin.Context) {
-	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
-	if err != nil {
-		logger.Error(err)
-		c.Abort()
-		return
-	}
-
-	userId := h.GetInt(c, "user_id", 0)
-	if userId == 0 {
-		logger.Info("Invalid user ID")
-		c.Abort()
-		return
-	}
-
-	client := types.NewWsClient(ws)
-	h.sunoService.Clients.Put(uint(userId), client)
-	logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
-}
-
 func (h *SunoHandler) Create(c *gin.Context) {

 	var data struct {
--- a/api/handler/test_handler.go
+++ b/api/handler/test_handler.go
@@ -19,7 +19,7 @@ func NewTestHandler(db *gorm.DB, snowflake *service.Snowflake, js *payment.GeekP
 }

 func (h *TestHandler) SseTest(c *gin.Context) {
-	//c.Header("Content-Type", "text/event-stream")
+	//c.Header("Body-Type", "text/event-stream")
 	//c.Header("Cache-Control", "no-cache")
 	//c.Header("Connection", "keep-alive")
 	//
--- a/api/handler/video_handler.go
+++ b/api/handler/video_handler.go
@@ -19,9 +19,7 @@ import (
 	"geekai/utils"
 	"geekai/utils/resp"
 	"github.com/gin-gonic/gin"
-	"github.com/gorilla/websocket"
 	"gorm.io/gorm"
-	"net/http"
 	"time"
 )

@@ -44,27 +42,6 @@ func NewVideoHandler(app *core.AppServer, db *gorm.DB, service *video.Service, u
 	}
 }

-// Client WebSocket 客户端，用于通知任务状态变更
-func (h *VideoHandler) Client(c *gin.Context) {
-	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
-	if err != nil {
-		logger.Error(err)
-		c.Abort()
-		return
-	}
-
-	userId := h.GetInt(c, "user_id", 0)
-	if userId == 0 {
-		logger.Info("Invalid user ID")
-		c.Abort()
-		return
-	}
-
-	client := types.NewWsClient(ws)
-	h.videoService.Clients.Put(uint(userId), client)
-	logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
-}
-
 func (h *VideoHandler) LumaCreate(c *gin.Context) {

 	var data struct {
--- a/api/handler/ws_handler.go
+++ b/api/handler/ws_handler.go
@@ -8,6 +8,7 @@ package handler
 // * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

 import (
+	"context"
 	"geekai/core"
 	"geekai/core/types"
 	"geekai/service"
@@ -15,6 +16,7 @@ import (
 	"geekai/utils"
 	"github.com/gin-gonic/gin"
 	"github.com/gorilla/websocket"
+	"gorm.io/gorm"
 	"net/http"
 )

@@ -23,11 +25,13 @@ import (
 type WebsocketHandler struct {
 	BaseHandler
 	wsService   *service.WebsocketService
+	chatHandler *ChatHandler
 }

-func NewWebsocketHandler(app *core.AppServer, s *service.WebsocketService) *WebsocketHandler {
+func NewWebsocketHandler(app *core.AppServer, s *service.WebsocketService, db *gorm.DB, chatHandler *ChatHandler) *WebsocketHandler {
 	return &WebsocketHandler{
-		BaseHandler: BaseHandler{App: app},
+		BaseHandler: BaseHandler{App: app, DB: db},
+		chatHandler: chatHandler,
 		wsService:   s,
 	}
 }
@@ -40,9 +44,9 @@ func (h *WebsocketHandler) Client(c *gin.Context) {
 		return
 	}

-	userId := h.GetInt(c, "user_id", 0)
-	clientId := c.Query("client")
-	client := types.NewWsClient(ws)
+	clientId := c.Query("client_id")
+	client := types.NewWsClient(ws, clientId)
+	userId := h.GetLoginUserId(c)
 	if userId == 0 {
 		_ = client.Send([]byte("Invalid user_id"))
 		c.Abort()
@@ -63,6 +67,8 @@ func (h *WebsocketHandler) Client(c *gin.Context) {
 			if err != nil {
 				logger.Debugf("close connection: %s", client.Conn.RemoteAddr())
 				client.Close()
+				h.wsService.Clients.Delete(clientId)
+				break
 			}

 			var message types.InputMessage
@@ -72,12 +78,66 @@ func (h *WebsocketHandler) Client(c *gin.Context) {
 			}

 			logger.Infof("Receive a message:%+v", message)
-			if message.Type == types.WsMsgTypePing {
-				_ = client.Send([]byte(`{"type":"pong"}`))
+			if message.Type == types.MsgTypePing {
+				utils.SendChannelMsg(client, types.ChPing, "pong")
+				continue
 			}

-			switch message.Channel {
+			// 当前只处理聊天消息，其他消息全部丢弃
+			var chatMessage types.ChatMessage
+			err = utils.JsonDecode(utils.JsonEncode(message.Body), &chatMessage)
+			if err != nil || message.Channel != types.ChChat {
+				logger.Warnf("invalid message body:%+v", message.Body)
+				continue
+			}
+			var chatRole model.ChatRole
+			err = h.DB.First(&chatRole, chatMessage.RoleId).Error
+			if err != nil || !chatRole.Enable {
+				utils.SendAndFlush(client, "当前聊天角色不存在或者未启用，请更换角色之后再发起对话！！！")
+				continue
+			}
+			// if the role bind a model_id, use role's bind model_id
+			if chatRole.ModelId > 0 {
+				chatMessage.RoleId = chatRole.ModelId
+			}
+			// get model info
+			var chatModel model.ChatModel
+			err = h.DB.Where("id", chatMessage.ModelId).First(&chatModel).Error
+			if err != nil || chatModel.Enabled == false {
+				utils.SendAndFlush(client, "当前AI模型暂未启用，请更换模型后再发起对话！！！")
+				continue
+			}

+			session := &types.ChatSession{
+				ClientIP: c.ClientIP(),
+				UserId:   userId,
+			}
+
+			// use old chat data override the chat model and role ID
+			var chat model.ChatItem
+			h.DB.Where("chat_id", chatMessage.ChatId).First(&chat)
+			if chat.Id > 0 {
+				chatModel.Id = chat.ModelId
+				chatMessage.RoleId = int(chat.RoleId)
+			}
+
+			session.ChatId = chatMessage.ChatId
+			session.Tools = chatMessage.Tools
+			session.Stream = chatMessage.Stream
+			// 复制模型数据
+			err = utils.CopyObject(chatModel, &session.Model)
+			if err != nil {
+				logger.Error(err, chatModel)
+			}
+			ctx, cancel := context.WithCancel(context.Background())
+			h.chatHandler.ReqCancelFunc.Put(clientId, cancel)
+			err = h.chatHandler.sendMessage(ctx, session, chatRole, chatMessage.Content, client)
+			if err != nil {
+				logger.Error(err)
+				utils.SendAndFlush(client, err.Error())
+			} else {
+				utils.SendMsg(client, types.ReplyMessage{Channel: types.ChChat, Type: types.MsgTypeEnd})
+				logger.Infof("回答完毕: %v", message.Body)
 			}

 		}
--- a/api/main.go
+++ b/api/main.go
@@ -14,7 +14,6 @@ import (
 	"geekai/core/types"
 	"geekai/handler"
 	"geekai/handler/admin"
-	"geekai/handler/chatimpl"
 	logger2 "geekai/logger"
 	"geekai/service"
 	"geekai/service/dalle"
@@ -128,7 +127,7 @@ func main() {
 		// 创建控制器
 		fx.Provide(handler.NewChatRoleHandler),
 		fx.Provide(handler.NewUserHandler),
-		fx.Provide(chatimpl.NewChatHandler),
+		fx.Provide(handler.NewChatHandler),
 		fx.Provide(handler.NewNetHandler),
 		fx.Provide(handler.NewSmsHandler),
 		fx.Provide(handler.NewRedeemHandler),
@@ -246,9 +245,8 @@ func main() {
 			group.GET("clogin", h.CLogin)
 			group.GET("clogin/callback", h.CLoginCallback)
 		}),
-		fx.Invoke(func(s *core.AppServer, h *chatimpl.ChatHandler) {
+		fx.Invoke(func(s *core.AppServer, h *handler.ChatHandler) {
 			group := s.Engine.Group("/api/chat/")
-			group.Any("new", h.ChatHandle)
 			group.GET("list", h.List)
 			group.GET("detail", h.Detail)
 			group.POST("update", h.Update)
@@ -281,7 +279,6 @@ func main() {
 		}),
 		fx.Invoke(func(s *core.AppServer, h *handler.MidJourneyHandler) {
 			group := s.Engine.Group("/api/mj/")
-			group.Any("client", h.Client)
 			group.POST("image", h.Image)
 			group.POST("upscale", h.Upscale)
 			group.POST("variation", h.Variation)
@@ -292,7 +289,6 @@ func main() {
 		}),
 		fx.Invoke(func(s *core.AppServer, h *handler.SdJobHandler) {
 			group := s.Engine.Group("/api/sd")
-			group.Any("client", h.Client)
 			group.POST("image", h.Image)
 			group.GET("jobs", h.JobList)
 			group.GET("imgWall", h.ImgWall)
@@ -467,13 +463,11 @@ func main() {
 		}),
 		fx.Provide(handler.NewMarkMapHandler),
 		fx.Invoke(func(s *core.AppServer, h *handler.MarkMapHandler) {
-			group := s.Engine.Group("/api/markMap/")
-			group.Any("client", h.Client)
+			s.Engine.POST("/api/markMap/gen", h.Generate)
 		}),
 		fx.Provide(handler.NewDallJobHandler),
 		fx.Invoke(func(s *core.AppServer, h *handler.DallJobHandler) {
 			group := s.Engine.Group("/api/dall")
-			group.Any("client", h.Client)
 			group.POST("image", h.Image)
 			group.GET("jobs", h.JobList)
 			group.GET("imgWall", h.ImgWall)
@@ -483,7 +477,6 @@ func main() {
 		fx.Provide(handler.NewSunoHandler),
 		fx.Invoke(func(s *core.AppServer, h *handler.SunoHandler) {
 			group := s.Engine.Group("/api/suno")
-			group.Any("client", h.Client)
 			group.POST("create", h.Create)
 			group.GET("list", h.List)
 			group.GET("remove", h.Remove)
@@ -496,7 +489,6 @@ func main() {
 		fx.Provide(handler.NewVideoHandler),
 		fx.Invoke(func(s *core.AppServer, h *handler.VideoHandler) {
 			group := s.Engine.Group("/api/video")
-			group.Any("client", h.Client)
 			group.POST("luma/create", h.LumaCreate)
 			group.GET("list", h.List)
 			group.GET("remove", h.Remove)
@@ -521,6 +513,11 @@ func main() {
 			group := s.Engine.Group("/api/test")
 			group.Any("sse", h.PostTest, h.SseTest)
 		}),
+		fx.Provide(service.NewWebsocketService),
+		fx.Provide(handler.NewWebsocketHandler),
+		fx.Invoke(func(s *core.AppServer, h *handler.WebsocketHandler) {
+			s.Engine.Any("/api/ws", h.Client)
+		}),
 		fx.Invoke(func(s *core.AppServer, db *gorm.DB) {
 			go func() {
 				err := s.Run(db)
--- a/api/service/dalle/service.go
+++ b/api/service/dalle/service.go
@@ -158,7 +158,7 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
 		Quality: task.Quality,
 	}
 	logger.Infof("Channel:%s, API KEY:%s, BODY: %+v", apiURL, apiKey.Value, reqBody)
-	r, err := s.httpClient.R().SetHeader("Content-Type", "application/json").
+	r, err := s.httpClient.R().SetHeader("Body-Type", "application/json").
 		SetHeader("Authorization", "Bearer "+apiKey.Value).
 		SetBody(reqBody).
 		SetErrorResult(&errRes).
--- a/api/service/oss/minio_oss.go
+++ b/api/service/oss/minio_oss.go
@@ -89,7 +89,7 @@ func (s MiniOss) PutFile(ctx *gin.Context, name string) (File, error) {
 	fileExt := utils.GetImgExt(file.Filename)
 	filename := fmt.Sprintf("%s/%d%s", s.config.SubDir, time.Now().UnixMicro(), fileExt)
 	info, err := s.client.PutObject(ctx, s.config.Bucket, filename, fileReader, file.Size, minio.PutObjectOptions{
-		ContentType: file.Header.Get("Content-Type"),
+		ContentType: file.Header.Get("Body-Type"),
 	})
 	if err != nil {
 		return File{}, fmt.Errorf("error uploading to MinIO: %v", err)
--- a/api/service/ws_service.go
+++ b/api/service/ws_service.go
@@ -5,3 +5,9 @@ import "geekai/core/types"
 type WebsocketService struct {
 	Clients *types.LMap[string, *types.WsClient] // clientId => Client
 }
+
+func NewWebsocketService() *WebsocketService {
+	return &WebsocketService{
+		Clients: types.NewLMap[string, *types.WsClient](),
+	}
+}
--- a/api/utils/net.go
+++ b/api/utils/net.go
@@ -19,8 +19,9 @@ import (

 var logger = logger2.GetLogger()

-// SendChunkMessage 回复客户片段端消息
-func SendChunkMessage(client *types.WsClient, message interface{}) {
+// SendMsg 回复客户片段端消息
+func SendMsg(client *types.WsClient, message types.ReplyMessage) {
+	message.ClientId = client.Id
 	msg, err := json.Marshal(message)
 	if err != nil {
 		logger.Errorf("Error for decoding json data: %v", err.Error())
@@ -32,19 +33,23 @@ func SendChunkMessage(client *types.WsClient, message interface{}) {
 	}
 }

-// SendMessage 回复客户端一条完整的消息
-func SendMessage(ws *types.WsClient, message interface{}) {
-	SendChunkMessage(ws, types.ReplyMessage{Type: types.WsMsgTypeContent, Content: message})
-	SendChunkMessage(ws, types.ReplyMessage{Type: types.WsMsgTypeEnd})
+// SendAndFlush 回复客户端一条完整的消息
+func SendAndFlush(ws *types.WsClient, message interface{}) {
+	SendMsg(ws, types.ReplyMessage{Channel: types.ChChat, Type: types.MsgTypeText, Body: message})
+	SendMsg(ws, types.ReplyMessage{Channel: types.ChChat, Type: types.MsgTypeEnd})
 }

-func ReplyContent(ws *types.WsClient, message interface{}) {
-	SendChunkMessage(ws, types.ReplyMessage{Type: types.WsMsgTypeContent, Content: message})
+func SendChunkMsg(ws *types.WsClient, message interface{}) {
+	SendMsg(ws, types.ReplyMessage{Channel: types.ChChat, Type: types.MsgTypeText, Body: message})
 }

-// ReplyErrorMessage 向客户端发送错误消息
-func ReplyErrorMessage(ws *types.WsClient, message interface{}) {
-	SendChunkMessage(ws, types.ReplyMessage{Type: types.WsMsgTypeErr, Content: message})
+// SendErrMsg 向客户端发送错误消息
+func SendErrMsg(ws *types.WsClient, message interface{}) {
+	SendMsg(ws, types.ReplyMessage{Channel: types.ChChat, Type: types.MsgTypeErr, Body: message})
+}
+
+func SendChannelMsg(ws *types.WsClient, channel types.WsChannel, message interface{}) {
+	SendMsg(ws, types.ReplyMessage{Channel: channel, Type: types.MsgTypeText, Body: message})
 }

 func DownloadImage(imageURL string, proxy string) ([]byte, error) {
@@ -68,7 +73,9 @@ func DownloadImage(imageURL string, proxy string) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	defer func(Body io.ReadCloser) {
+		_ = Body.Close()
+	}(resp.Body)

 	imageBytes, err := io.ReadAll(resp.Body)
 	if err != nil {
--- a/api/utils/openai.go
+++ b/api/utils/openai.go
@@ -65,7 +65,7 @@ func OpenAIRequest(db *gorm.DB, prompt string, modelName string) (string, error)
 	}
 	apiURL := fmt.Sprintf("%s/v1/chat/completions", apiKey.ApiURL)
 	logger.Debugf("Sending %s request, API KEY:%s, PROXY: %s, Model: %s", apiKey.ApiURL, apiURL, apiKey.ProxyURL, modelName)
-	r, err := client.R().SetHeader("Content-Type", "application/json").
+	r, err := client.R().SetHeader("Body-Type", "application/json").
 		SetHeader("Authorization", "Bearer "+apiKey.Value).
 		SetBody(types.ApiRequest{
 			Model:       modelName,
--- a/web/src/App.vue
+++ b/web/src/App.vue
@@ -6,10 +6,13 @@

 <script setup>
 import {ElConfigProvider} from 'element-plus';
-import {onMounted} from "vue";
-import {getSystemInfo} from "@/store/cache";
+import {onMounted, ref} from "vue";
+import {checkSession, getClientId, getSystemInfo} from "@/store/cache";
 import {isChrome, isMobile} from "@/utils/libs";
 import {showMessageInfo} from "@/utils/dialog";
+import {useSharedStore} from "@/store/sharedata";
+import {getUserToken} from "@/store/session";
+import {clear} from "core-js/internals/task";

 const debounce = (fn, delay) => {
  let timer
@@ -32,6 +35,7 @@ window.ResizeObserver = class ResizeObserver extends _ResizeObserver {
 }

 onMounted(() => {
+  // 获取系统参数
  getSystemInfo().then((res) => {
    const link = document.createElement('link')
    link.rel = 'shortcut icon'
@@ -39,9 +43,50 @@ onMounted(() => {
    document.head.appendChild(link)
  })
  if (!isChrome() && !isMobile()) {
-    showMessageInfo("检测到您使用的浏览器不是 Chrome，可能会导致部分功能无法正常使用，建议使用 Chrome 浏览器。")
+    showMessageInfo("建议使用 Chrome 浏览器以获得最佳体验。")
  }
+
+  checkSession().then(() => {
+    connect()
+  }).catch(()=>{})
 })
+
+const store = useSharedStore()
+const handler = ref(0)
+
+// 初始化 websocket 连接
+const connect = () => {
+  let host = process.env.VUE_APP_WS_HOST
+  if (host === '') {
+    if (location.protocol === 'https:') {
+      host = 'wss://' + location.host;
+    } else {
+      host = 'ws://' + location.host;
+    }
+  }
+  const clientId = getClientId()
+  const _socket = new WebSocket(host + `/api/ws?client_id=${clientId}&token=${getUserToken()}`);
+  _socket.addEventListener('open', () => {
+    console.log('WebSocket 已连接')
+    handler.value = setInterval(() => {
+      _socket.send(JSON.stringify({"type":"ping"}))
+    },5000)
+
+    for (const key in store.messageHandlers) {
+      console.log(key, store.messageHandlers[key])
+      store.setMessageHandler(store.messageHandlers[key])
+    }
+  });
+
+  _socket.addEventListener('close', () => {
+    store.setSocket(null)
+    clearInterval(handler.value)
+    connect()
+  });
+
+  store.setSocket(_socket)
+}
+
 </script>


--- a/web/src/store/cache.js
+++ b/web/src/store/cache.js
@@ -1,5 +1,6 @@
 import {httpGet} from "@/utils/http";
 import Storage from "good-storage";
+import {randString} from "@/utils/libs";

 const userDataKey = "USER_INFO_CACHE_KEY"
 const adminDataKey = "ADMIN_INFO_CACHE_KEY"
@@ -71,3 +72,13 @@ export function getLicenseInfo() {
        })
    })
 }
+
+export function getClientId() {
+    let clientId = Storage.get('client_id')
+    if (clientId) {
+        return clientId
+    }
+    clientId = randString(42)
+    Storage.set('client_id', clientId)
+    return clientId
+}
--- a/web/src/store/sharedata.js
+++ b/web/src/store/sharedata.js
@@ -6,6 +6,8 @@ export const useSharedStore = defineStore('shared', {
        showLoginDialog: false,
        chatListStyle: Storage.get("chat_list_style","chat"),
        chatStream: Storage.get("chat_stream",true),
+        socket: WebSocket,
+        messageHandlers:{},
    }),
    getters: {},
    actions: {
@@ -19,6 +21,36 @@ export const useSharedStore = defineStore('shared', {
        setChatStream(value) {
            this.chatStream = value;
            Storage.set("chat_stream", value);
+        },
+        setSocket(value) {
+            this.socket = value;
+        },
+        addMessageHandler(key, callback) {
+            if (!this.messageHandlers[key]) {
+                this.messageHandlers[key] = callback;
+                this.setMessageHandler(callback)
+            }
+        },
+        setMessageHandler(callback) {
+            if (this.socket instanceof WebSocket && this.socket.readyState === WebSocket.OPEN) {
+                this.socket.addEventListener('message', (event) => {
+                    try {
+                        if (event.data instanceof Blob) {
+                            const reader = new FileReader();
+                            reader.readAsText(event.data, "UTF-8");
+                            reader.onload = () => {
+                                callback(JSON.parse(String(reader.result)))
+                            }
+                        }
+                    } catch (e) {
+                        console.warn(e)
+                    }
+                })
+            } else {
+                setTimeout(() => {
+                    this.setMessageHandler(callback)
+                }, 1000)
+            }
        }
    }
 });
--- a/web/src/store/system.js
+++ b/web/src/store/system.js
@@ -6,7 +6,6 @@
 // * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

 import Storage from "good-storage";
-import {useRouter} from "vue-router";

 const MOBILE_THEME = process.env.VUE_APP_KEY_PREFIX + "MOBILE_THEME"
 const ADMIN_THEME = process.env.VUE_APP_KEY_PREFIX + "ADMIN_THEME"
--- a/web/src/utils/http.js
+++ b/web/src/utils/http.js
@@ -17,7 +17,6 @@ axios.defaults.headers.post['Content-Type'] = 'application/json'
 axios.interceptors.request.use(
    config => {
        // set token
-        config.headers['Chat-Token'] = getSessionId();
        config.headers['Authorization'] = getUserToken();
        config.headers['Admin-Authorization'] = getAdminToken();
        return config
--- a/web/src/views/ChatPlus.vue
+++ b/web/src/views/ChatPlus.vue
@@ -213,7 +213,7 @@

 </template>
 <script setup>
-import {nextTick, onMounted, onUnmounted, ref, watch} from 'vue'
+import {nextTick, onMounted, ref, watch} from 'vue'
 import ChatPrompt from "@/components/ChatPrompt.vue";
 import ChatReply from "@/components/ChatReply.vue";
 import {Delete, Edit, InfoFilled, More, Plus, Promotion, Search, Share, VideoPause} from '@element-plus/icons-vue'
@@ -225,11 +225,10 @@ import {
  UUID
 } from "@/utils/libs";
 import {ElMessage, ElMessageBox} from "element-plus";
-import {getSessionId, getUserToken} from "@/store/session";
 import {httpGet, httpPost} from "@/utils/http";
 import {useRouter} from "vue-router";
 import Clipboard from "clipboard";
-import {checkSession, getSystemInfo} from "@/store/cache";
+import {checkSession, getClientId, getSystemInfo} from "@/store/cache";
 import Welcome from "@/components/Welcome.vue";
 import {useSharedStore} from "@/store/sharedata";
 import FileSelect from "@/components/FileSelect.vue";
@@ -270,7 +269,6 @@ watch(() => store.chatListStyle, (newValue) => {
 });
 const tools = ref([])
 const toolSelected = ref([])
-const loadHistory = ref(false)
 const stream = ref(store.chatStream)

 watch(() => store.chatStream, (newValue) => {
@@ -337,6 +335,13 @@ httpGet("/api/function/list").then(res => {
  showMessageError("获取工具函数失败：" + e.message)
 })

+// 创建 socket 连接
+const prompt = ref('');
+const showStopGenerate = ref(false); // 停止生成
+const lineBuffer = ref(''); // 输出缓冲行
+const canSend = ref(true);
+const isNewMsg = ref(true)
+
 onMounted(() => {
  resizeElement();
  initData()
@@ -351,14 +356,73 @@ onMounted(() => {
  })

  window.onresize = () => resizeElement();
-});
-
-onUnmounted(() => {
-  if (socket.value !== null) {
-    socket.value = null
+  store.addMessageHandler("chat", (data) => {
+    console.log(data)
+    // 丢去非本频道和本客户端的消息
+    if (data.channel !== 'chat' || data.clientId !== getClientId()) {
+      return
    }
-})

+    if (data.type === 'error') {
+      ElMessage.error(data.body)
+      return
+    }
+
+    const chatRole = getRoleById(roleId.value)
+    if (isNewMsg.value && data.type !== 'end') {
+      const prePrompt = chatData.value[chatData.value.length-1]?.content
+      chatData.value.push({
+        type: "reply",
+        id: randString(32),
+        icon: chatRole['icon'],
+        prompt:prePrompt,
+        content: data.body,
+      });
+      isNewMsg.value = false
+      lineBuffer.value = data.body;
+    } else if (data.type === 'end') { // 消息接收完毕
+      // 追加当前会话到会话列表
+      if (newChatItem.value !== null) {
+        newChatItem.value['title'] = tmpChatTitle.value;
+        newChatItem.value['chat_id'] = chatId.value;
+        chatList.value.unshift(newChatItem.value);
+        newChatItem.value = null; // 只追加一次
+      }
+
+      enableInput()
+      lineBuffer.value = ''; // 清空缓冲
+
+      // 获取 token
+      const reply = chatData.value[chatData.value.length - 1]
+      httpPost("/api/chat/tokens", {
+        text: "",
+        model: getModelValue(modelID.value),
+        chat_id: chatId.value,
+      }).then(res => {
+        reply['created_at'] = new Date().getTime();
+        reply['tokens'] = res.data;
+        // 将聊天框的滚动条滑动到最底部
+        nextTick(() => {
+          document.getElementById('chat-box').scrollTo(0, document.getElementById('chat-box').scrollHeight)
+        })
+      }).catch(() => {
+      })
+      isNewMsg.value = true
+
+    } else if (data.type === 'text') {
+      lineBuffer.value += data.body;
+      const reply = chatData.value[chatData.value.length - 1]
+      if (reply) {
+        reply['content'] = lineBuffer.value;
+      }
+    }
+    // 将聊天框的滚动条滑动到最底部
+    nextTick(() => {
+      document.getElementById('chat-box').scrollTo(0, document.getElementById('chat-box').scrollHeight)
+      localStorage.setItem("chat_id", chatId.value)
+    })
+  })
+});
 // 初始化数据
 const initData = () => {

@@ -492,9 +556,8 @@ const newChat = () => {
    removing: false,
  };
  showStopGenerate.value = false;
+  loadChatHistory(chatId.value)
  router.push(`/chat/${chatId.value}`)
-  loadHistory.value = true
-  connect()
 }

 // 切换会话
@@ -507,14 +570,12 @@ const loadChat = function (chat) {
  if (chatId.value === chat.chat_id) {
    return;
  }
-
  newChatItem.value = null;
  roleId.value = chat.role_id;
  modelID.value = chat.model_id;
  chatId.value = chat.chat_id;
  showStopGenerate.value = false;
-  loadHistory.value = true
-  connect()
+  loadChatHistory(chatId.value)
  router.replace(`/chat/${chatId.value}`)
 }

@@ -587,118 +648,6 @@ const removeChat = function (chat) {

 }

-// 创建 socket 连接
-const prompt = ref('');
-const showStopGenerate = ref(false); // 停止生成
-const lineBuffer = ref(''); // 输出缓冲行
-const socket = ref(null);
-const canSend = ref(true);
-const sessionId = ref("")
-const isNewMsg = ref(true)
-const connect = function () {
-  const chatRole = getRoleById(roleId.value);
-  // 初始化 WebSocket 对象
-  sessionId.value = getSessionId();
-  let host = process.env.VUE_APP_WS_HOST
-  if (host === '') {
-    if (location.protocol === 'https:') {
-      host = 'wss://' + location.host;
-    } else {
-      host = 'ws://' + location.host;
-    }
-  }
-
-  loading.value = true
-  const _socket = new WebSocket(host + `/api/chat/new?session_id=${sessionId.value}&role_id=${roleId.value}&chat_id=${chatId.value}&model_id=${modelID.value}&token=${getUserToken()}`);
-  _socket.addEventListener('open', () => {
-    enableInput()
-    if (loadHistory.value) {
-      loadChatHistory(chatId.value)
-    }
-    loading.value = false
-  });
-
-  _socket.addEventListener('message', event => {
-    try {
-      if (event.data instanceof Blob) {
-        const reader = new FileReader();
-        reader.readAsText(event.data, "UTF-8");
-        reader.onload = () => {
-          const data = JSON.parse(String(reader.result));
-          if (data.type === 'error') {
-            ElMessage.error(data.message)
-            return
-          }
-
-          if (isNewMsg.value && data.type !== 'end') {
-            const prePrompt = chatData.value[chatData.value.length-1]?.content
-            chatData.value.push({
-              type: "reply",
-              id: randString(32),
-              icon: chatRole['icon'],
-              prompt:prePrompt,
-              content: data.content,
-            });
-            isNewMsg.value = false
-            lineBuffer.value = data.content;
-          } else if (data.type === 'end') { // 消息接收完毕
-            // 追加当前会话到会话列表
-            if (newChatItem.value !== null) {
-              newChatItem.value['title'] = tmpChatTitle.value;
-              newChatItem.value['chat_id'] = chatId.value;
-              chatList.value.unshift(newChatItem.value);
-              newChatItem.value = null; // 只追加一次
-            }
-
-            enableInput()
-            lineBuffer.value = ''; // 清空缓冲
-
-            // 获取 token
-            const reply = chatData.value[chatData.value.length - 1]
-            httpPost("/api/chat/tokens", {
-              text: "",
-              model: getModelValue(modelID.value),
-              chat_id: chatId.value,
-            }).then(res => {
-              reply['created_at'] = new Date().getTime();
-              reply['tokens'] = res.data;
-              // 将聊天框的滚动条滑动到最底部
-              nextTick(() => {
-                document.getElementById('chat-box').scrollTo(0, document.getElementById('chat-box').scrollHeight)
-              })
-              isNewMsg.value = true
-            }).catch(() => {
-            })
-
-          } else {
-            lineBuffer.value += data.content;
-            const reply = chatData.value[chatData.value.length - 1]
-            if (reply) {
-              reply['content'] = lineBuffer.value;
-            }
-          }
-          // 将聊天框的滚动条滑动到最底部
-          nextTick(() => {
-            document.getElementById('chat-box').scrollTo(0, document.getElementById('chat-box').scrollHeight)
-            localStorage.setItem("chat_id", chatId.value)
-          })
-        };
-      }
-    } catch (e) {
-      console.warn(e)
-    }
-
-  });
-
-  _socket.addEventListener('close', () => {
-    disableInput(true)
-    loadHistory.value = false
-    connect()
-  });
-
-  socket.value = _socket;
-}
-
 const disableInput = (force) => {
  canSend.value = false;
  showStopGenerate.value = !force;
@@ -747,6 +696,11 @@ const sendMessage = function () {
    return;
  }

+  if (store.socket.readyState !== WebSocket.OPEN) {
+    ElMessage.warning("连接断开，正在重连...");
+    return
+  }
+
  if (canSend.value === false) {
    ElMessage.warning("AI 正在作答中，请稍后...");
    return
@@ -780,7 +734,18 @@ const sendMessage = function () {

  showHello.value = false
  disableInput(false)
-  socket.value.send(JSON.stringify({tools: toolSelected.value, content: content, stream: stream.value}));
+  store.socket.send(JSON.stringify({
+    channel: 'chat',
+    type:'text',
+    body:{
+      role_id: roleId.value,
+      model_id: modelID.value,
+      chat_id: chatId.value,
+      content: content,
+      tools:toolSelected.value,
+      stream: stream.value
+    }
+  }));
  tmpChatTitle.value = content
  prompt.value = ''
  files.value = []
@@ -849,7 +814,7 @@ const loadChatHistory = function (chatId) {

 const stopGenerate = function () {
  showStopGenerate.value = false;
-  httpGet("/api/chat/stop?session_id=" + sessionId.value).then(() => {
+  httpGet("/api/chat/stop?session_id=" + getClientId()).then(() => {
    enableInput()
  })
 }
@@ -865,7 +830,18 @@ const reGenerate = function (prompt) {
    icon: loginUser.value.avatar,
    content: text
  });
-  socket.value.send(JSON.stringify({tools: toolSelected.value, content: text, stream: stream.value}));
+  store.socket.send(JSON.stringify({
+    channel: 'chat',
+    type:'text',
+    body:{
+      role_id: roleId.value,
+      model_id: modelID.value,
+      chat_id: chatId.value,
+      content: text,
+      tools:toolSelected.value,
+      stream: stream.value
+    }
+  }));
 }

 const chatName = ref('')
--- a/web/src/views/Index.vue
+++ b/web/src/views/Index.vue
@@ -31,6 +31,7 @@
            <el-button :color="theme.btnBgColor" :style="{color: theme.btnTextColor}" @click="router.push('/login')" class="shadow" round>登录</el-button>
            <el-button :color="theme.btnBgColor" :style="{color: theme.btnTextColor}" @click="router.push('/register')" class="shadow" round>注册</el-button>
          </span>
+          <el-button :color="theme.btnBgColor" :style="{color: theme.btnTextColor}" @click="router.push('/test')" class="shadow" round>测试</el-button>
        </div>
      </el-menu>
    </div>
--- a/web/src/views/Test.vue
+++ b/web/src/views/Test.vue
@@ -5,20 +5,12 @@
 </template>

 <script setup>
-import {ref, onMounted, onUpdated} from 'vue';
-import {Markmap} from 'markmap-view';
-import {loadJS, loadCSS} from 'markmap-common';
-import {Transformer} from 'markmap-lib';
-import {httpPost} from "@/utils/http";
+import {onMounted, ref} from "vue";
+
+const data = ref('abc')
+
+onMounted(() => {
  
-const data=ref("")
-httpPost("/api/test/sse",{
-  "message":"你是什么模型",
-  "user_id":123
-}).then(res=>{
-  // const source = new EventSource("http://localhost:5678/api/test/sse");
-  // source.onmessage = function(event) {
-  //   console.log(event.data)
-  // };
 })
+
 </script>