From 0f059701414651ab857d3c20b731cee2984fde55 Mon Sep 17 00:00:00 2001
From: RockYang <yangjian102621@gmail.com>
Date: Mon, 22 Jan 2024 18:42:51 +0800
Subject: [PATCH] opt: add heartbeat message for websocket connects

---
 CHANGELOG.md                         |  2 ++
 api/handler/chatimpl/chat_handler.go | 22 ++++++++++++++++------
 web/src/views/ChatPlus.vue           | 11 ++++++++++-
 web/src/views/mobile/ChatSession.vue | 11 ++++++++++-
 4 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index defc683e..3f382664 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,8 @@
 * 功能新增：新增阿里通义千问大模型支持
 * Bug修复：修复 MJ 放大任务失败时候 img_call 会增加的 Bug
 * 功能优化：新增虎皮椒和PayJS订单状态校验功能，增加安全性
+* Bug修复：修复微信转账交易 ID 提取失败 Bug
+* 功能优化：给所有的 websocket 连接加上心跳，解决 "close 1006 (abnormal closure): unexpected EOF" Bug
 
 
 ## v3.2.5
diff --git a/api/handler/chatimpl/chat_handler.go b/api/handler/chatimpl/chat_handler.go
index 0b19b238..99f4f64f 100644
--- a/api/handler/chatimpl/chat_handler.go
+++ b/api/handler/chatimpl/chat_handler.go
@@ -134,7 +134,6 @@ func (h *ChatHandler) ChatHandle(c *gin.Context) {
 		for {
 			_, msg, err := client.Receive()
 			if err != nil {
-				logger.Error(err)
 				client.Close()
 				h.App.ChatClients.Delete(sessionId)
 				cancelFunc := h.App.ReqCancelFunc.Get(sessionId)
@@ -145,19 +144,30 @@ func (h *ChatHandler) ChatHandle(c *gin.Context) {
 				return
 			}
 
-			message := string(msg)
-			logger.Info("Receive a message: ", message)
-			//utils.ReplyMessage(client, "这是一条测试消息！")
+			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)
+
 			ctx, cancel := context.WithCancel(context.Background())
 			h.App.ReqCancelFunc.Put(sessionId, cancel)
 			// 回复消息
-			err = h.sendMessage(ctx, session, chatRole, message, client)
+			err = h.sendMessage(ctx, session, chatRole, utils.InterfaceToString(message.Content), client)
 			if err != nil {
 				logger.Error(err)
 				utils.ReplyChunkMessage(client, types.WsMessage{Type: types.WsEnd})
 			} else {
 				utils.ReplyChunkMessage(client, types.WsMessage{Type: types.WsEnd})
-				logger.Info("回答完毕: " + string(message))
+				logger.Infof("回答完毕: %v", message.Content)
 			}
 
 		}
diff --git a/web/src/views/ChatPlus.vue b/web/src/views/ChatPlus.vue
index 456a8ac4..cd7f8b4a 100644
--- a/web/src/views/ChatPlus.vue
+++ b/web/src/views/ChatPlus.vue
@@ -543,6 +543,7 @@ const lineBuffer = ref(''); // 输出缓冲行
 const socket = ref(null);
 const activelyClose = ref(false); // 主动关闭
 const canSend = ref(true);
+const heartbeatHandle = ref(null)
 const connect = function (chat_id, role_id) {
   let isNewChat = false;
   if (!chat_id) {
@@ -589,6 +590,14 @@ const connect = function (chat_id, role_id) {
       loadChatHistory(chat_id);
     }
 
+    // 发送心跳消息
+    clearInterval(heartbeatHandle.value)
+    heartbeatHandle.value = setInterval(() => {
+      if (socket.value !== null) {
+        socket.value.send(JSON.stringify({type: "heartbeat", content: "ping"}))
+      }
+    }, 5000);
+
   });
 
   _socket.addEventListener('message', event => {
@@ -718,7 +727,7 @@ const sendMessage = function () {
 
   showHello.value = false
   disableInput(false)
-  socket.value.send(prompt.value);
+  socket.value.send(JSON.stringify({type: "chat", content: prompt.value}));
   previousText.value = prompt.value;
   prompt.value = '';
   return true;
diff --git a/web/src/views/mobile/ChatSession.vue b/web/src/views/mobile/ChatSession.vue
index 0d708fbe..973365a3 100644
--- a/web/src/views/mobile/ChatSession.vue
+++ b/web/src/views/mobile/ChatSession.vue
@@ -200,6 +200,7 @@ const lineBuffer = ref(''); // 输出缓冲行
 const socket = ref(null);
 const activelyClose = ref(false); // 主动关闭
 const canSend = ref(true);
+const heartbeatHandle = ref(null)
 const connect = function (chat_id, role_id) {
   let isNewChat = false;
   if (!chat_id) {
@@ -238,6 +239,14 @@ const connect = function (chat_id, role_id) {
         orgContent: role.helloMsg,
       })
     }
+
+    // 发送心跳消息
+    clearInterval(heartbeatHandle.value)
+    heartbeatHandle.value = setInterval(() => {
+      if (socket.value !== null) {
+        socket.value.send(JSON.stringify({type: "heartbeat", content: "ping"}))
+      }
+    }, 5000);
   });
 
   _socket.addEventListener('message', event => {
@@ -370,7 +379,7 @@ const sendMessage = () => {
   })
 
   disableInput(false)
-  socket.value.send(prompt.value);
+  socket.value.send(JSON.stringify({type: "chat", content: prompt.value}));
   previousText.value = prompt.value;
   prompt.value = '';
   return true;