opt: close the old connection for mj and sd clients

api/handler/markmap_handler.go

@@ -3,10 +3,9 @@ package handler
 import (
 	"chatplus/core"
 	"chatplus/core/types"
-	"chatplus/store/model"
-	"chatplus/store/vo"
 	"chatplus/utils"
-	"chatplus/utils/resp"
+	"github.com/gorilla/websocket"
+	"net/http"
 
 	"github.com/gin-gonic/gin"
 	"gorm.io/gorm"
@@ -15,34 +14,53 @@ import (
 // MarkMapHandler 生成思维导图
 type MarkMapHandler struct {
 	BaseHandler
+	clients *types.LMap[uint, *types.WsClient]
 }
 
 func NewMarkMapHandler(app *core.AppServer, db *gorm.DB) *MarkMapHandler {
-	return &MarkMapHandler{BaseHandler: BaseHandler{App: app, DB: db}}
+	return &MarkMapHandler{
+		BaseHandler: BaseHandler{App: app, DB: db},
+		clients:     types.NewLMap[uint, *types.WsClient](),
+	}
 }
 
-// GetModel get the chat model for generating Markdown text
-func (h *MarkMapHandler) GetModel(c *gin.Context) {
-	modelId := h.App.SysConfig.XMindModelId
-	session := h.DB.Session(&gorm.Session{}).Where("enabled", true)
-	if modelId > 0 {
-		session = session.Where("id", modelId)
-	} else {
-		session = session.Where("platform", types.OpenAI)
-	}
-	var chatModel model.ChatModel
-	res := session.First(&chatModel)
-	if res.Error != nil {
-		resp.ERROR(c, "No available AI model")
-		return
-	}
-
-	var modelVo vo.ChatModel
-	err := utils.CopyObject(chatModel, &modelVo)
+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 {
-		resp.ERROR(c, "error with copy object: "+err.Error())
+		logger.Error(err)
 		return
 	}
 
-	resp.SUCCESS(c, modelVo)
+	modelId := h.GetInt(c, "model_id", 0)
+	userId := h.GetLoginUserId(c)
+	logger.Info(modelId)
+	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.WsMessage
+			err = utils.JsonDecode(string(msg), &message)
+			if err != nil {
+				continue
+			}
+
+			// 心跳消息
+			if message.Type == "heartbeat" {
+				logger.Debug("收到 Chat 心跳消息：", message.Content)
+				continue
+			}
+
+			logger.Info("Receive a message: ", message.Content)
+
+		}
+	}()
 }

api/handler/mj_handler.go

@@ -78,6 +78,10 @@ func (h *MidJourneyHandler) Client(c *gin.Context) {
 	}
 
 	client := types.NewWsClient(ws)
+	// close the existed connections
+	if cli := h.pool.Clients.Get(uint(userId)); cli != nil {
+		cli.Close()
+	}
 	h.pool.Clients.Put(uint(userId), client)
 	logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
 }

api/handler/sd_handler.go

@@ -61,6 +61,10 @@ func (h *SdJobHandler) Client(c *gin.Context) {
 	}
 
 	client := types.NewWsClient(ws)
+	// close the existed connections
+	if cli := h.pool.Clients.Get(uint(userId)); cli != nil {
+		cli.Close()
+	}
 	h.pool.Clients.Put(uint(userId), client)
 	logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
 }

web/src/views/ImageMj.vue

@@ -797,9 +797,23 @@ const connect = () => {
   });
 
   _socket.addEventListener('close', () => {
-    if (socket.value !== null) {
-      connect()
+    ElMessageBox.confirm(
+        '检测到您已经在其他客户端创建了新的连接，当前连接将被关闭！',
+        '提示',
+        {
+          dangerouslyUseHTMLString: true,
+          confirmButtonText: '重新连接',
+          cancelButtonText: '关闭',
+          type: 'warning',
         }
+    ).then(() => {
+      connect()
+    }).catch(() => {
+      ElMessage({
+        type: 'info',
+        message: '连接已关闭',
+      })
+    })
   });
 }
 

web/src/views/ImageSd.vue

@@ -576,9 +576,23 @@ const connect = () => {
   });
 
   _socket.addEventListener('close', () => {
-    if (socket.value !== null) {
-      connect()
+    ElMessageBox.confirm(
+        '检测到您已经在其他客户端创建了新的连接，当前连接将被关闭！',
+        '提示',
+        {
+          dangerouslyUseHTMLString: true,
+          confirmButtonText: '重新连接',
+          cancelButtonText: '关闭',
+          type: 'warning',
         }
+    ).then(() => {
+      connect()
+    }).catch(() => {
+      ElMessage({
+        type: 'info',
+        message: '连接已关闭',
+      })
+    })
   });
 }