From c374126f69e36e7b7ab8e0b1e1f7cbe45fff5c72 Mon Sep 17 00:00:00 2001
From: RockYang <yangjian102621@gmail.com>
Date: Tue, 8 Oct 2024 16:41:19 +0800
Subject: [PATCH] fixed bug for websocket message handler rebind

---
 web/src/App.vue                      | 17 +--------
 web/src/store/sharedata.js           | 56 ++++++++++++++--------------
 web/src/views/ChatPlus.vue           |  6 +--
 web/src/views/mobile/ChatSession.vue |  6 +--
 4 files changed, 35 insertions(+), 50 deletions(-)

diff --git a/web/src/App.vue b/web/src/App.vue
index 0ea27ad1..ea30e4bf 100644
--- a/web/src/App.vue
+++ b/web/src/App.vue
@@ -54,7 +54,7 @@ onMounted(() => {
 })
 
 watch(() => store.isLogin, (val) => {
-  if (val && store.socket.readyState !== WebSocket.OPEN) {
+  if (val) {
     connect()
   }
 })
@@ -79,26 +79,11 @@ const connect = () => {
         _socket.send(JSON.stringify({"type":"ping"}))
       }
     },5000)
-
-    // 绑定事件监听
-    for (const key in store.messageHandlers) {
-      console.log(store.messageHandlers[key])
-      store.setMessageHandler(store.messageHandlers[key])
-    }
   })
-
   _socket.addEventListener('close', () => {
-    // 移除事件监听
-    for (const key in store.messageHandlers) {
-      if (store.socket) {
-        store.socket.removeEventListener('message', store.messageHandlers[key])
-      }
-    }
-    store.setSocket(null)
     clearInterval(handler.value)
     connect()
   });
-
   store.setSocket(_socket)
 }
 
diff --git a/web/src/store/sharedata.js b/web/src/store/sharedata.js
index d3a3993f..66317697 100644
--- a/web/src/store/sharedata.js
+++ b/web/src/store/sharedata.js
@@ -6,8 +6,7 @@ export const useSharedStore = defineStore('shared', {
         showLoginDialog: false,
         chatListStyle: Storage.get("chat_list_style","chat"),
         chatStream: Storage.get("chat_stream",true),
-        socket: WebSocket,
-        messageHandlers:{},
+        socket: {conn:null, handlers:{}},
         mobileTheme: Storage.get("mobile_theme", "light"),
         adminTheme: Storage.get("admin_theme", "light"),
         isLogin: false
@@ -26,41 +25,42 @@ export const useSharedStore = defineStore('shared', {
             Storage.set("chat_stream", value);
         },
         setSocket(value) {
-            this.socket = value;
+            for (const key in this.socket.handlers) {
+                this.setMessageHandler(value, this.socket.handlers[key])
+            }
+            this.socket.conn = value
         },
         addMessageHandler(key, callback) {
-            if (!this.messageHandlers[key]) {
-                this.setMessageHandler(callback)
+            if (!this.socket.handlers[key]) {
+                this.socket.handlers[key] = callback;
             }
-            this.messageHandlers[key] = callback;
+            this.setMessageHandler(this.socket.conn, callback)
         },
-        setMessageHandler(callback) {
-            if (this.socket instanceof WebSocket && this.socket.readyState === WebSocket.OPEN) {
-                console.log(callback)
-                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)))
-                            }
+
+        setMessageHandler(conn, callback) {
+            if (!conn) {
+                return
+            }
+            conn.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)
-            }
+                } catch (e) {
+                    console.warn(e)
+                }
+            })
         },
+
         removeMessageHandler(key) {
-            if (this.socket.readyState === WebSocket.OPEN) {
-                this.socket.removeEventListener('message', this.messageHandlers[key])
+            if (this.socket.conn && this.socket.conn.readyState === WebSocket.OPEN) {
+                this.socket.conn.removeEventListener('message', this.socket.handlers[key])
             }
-            delete this.messageHandlers[key]
+            delete this.socket.handlers[key]
         },
         setMobileTheme(theme) {
             this.mobileTheme = theme
diff --git a/web/src/views/ChatPlus.vue b/web/src/views/ChatPlus.vue
index 1352f9b3..19cd964c 100644
--- a/web/src/views/ChatPlus.vue
+++ b/web/src/views/ChatPlus.vue
@@ -689,7 +689,7 @@ const sendMessage = function () {
     return;
   }
 
-  if (store.socket.readyState !== WebSocket.OPEN) {
+  if (store.socket.conn.readyState !== WebSocket.OPEN) {
     ElMessage.warning("连接断开，正在重连...");
     return
   }
@@ -727,7 +727,7 @@ const sendMessage = function () {
 
   showHello.value = false
   disableInput(false)
-  store.socket.send(JSON.stringify({
+  store.socket.conn.send(JSON.stringify({
     channel: 'chat',
     type:'text',
     body:{
@@ -825,7 +825,7 @@ const reGenerate = function (prompt) {
     icon: loginUser.value.avatar,
     content: text
   });
-  store.socket.send(JSON.stringify({
+  store.socket.conn.send(JSON.stringify({
     channel: 'chat',
     type:'text',
     body:{
diff --git a/web/src/views/mobile/ChatSession.vue b/web/src/views/mobile/ChatSession.vue
index ff2e7fa4..5fdcd338 100644
--- a/web/src/views/mobile/ChatSession.vue
+++ b/web/src/views/mobile/ChatSession.vue
@@ -512,7 +512,7 @@ const sendMessage = () => {
     return
   }
 
-  if (store.socket.readyState !== WebSocket.OPEN) {
+  if (store.socket.conn.readyState !== WebSocket.OPEN) {
     showToast("连接断开，正在重连...");
     return
   }
@@ -536,7 +536,7 @@ const sendMessage = () => {
   })
 
   disableInput(false)
-  store.socket.send(JSON.stringify({
+  store.socket.conn.send(JSON.stringify({
     channel: 'chat',
     type:'text',
     body:{
@@ -569,7 +569,7 @@ const reGenerate = () => {
     icon: loginUser.value.avatar,
     content: renderInputText(text)
   });
-  store.socket.send(JSON.stringify({
+  store.socket.conn.send(JSON.stringify({
     channel: 'chat',
     type:'text',
     body:{