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

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>
 
 

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"
@@ -70,4 +71,14 @@ export function getLicenseInfo() {
             resolve(err)
         })
     })
+}
+
+export function getClientId() {
+    let clientId = Storage.get('client_id')
+    if (clientId) {
+        return clientId
+    }
+    clientId = randString(42)
+    Storage.set('client_id', clientId)
+    return clientId
 }

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)
+            }
         }
     }
 });

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"
@@ -71,4 +70,4 @@ export function setRoute(path) {
 
 export function getRoute() {
     return Storage.get(process.env.VUE_APP_KEY_PREFIX + 'ROUTE_')
-}
+}

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

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();
+  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)
+    })
+  })
 });
-
-onUnmounted(() => {
-  if (socket.value !== null) {
-    socket.value = null
-  }
-})
-
 // 初始化数据
 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('')

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>

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("")
-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)
-  // };
+const data = ref('abc')
+
+onMounted(() => {
+  
 })
+
 </script>