From 1c3684faaf8a1028f1050189ad419887e5f06d3c Mon Sep 17 00:00:00 2001
From: hamster1963 <1410514192@qq.com>
Date: Wed, 25 Dec 2024 14:28:35 +0800
Subject: [PATCH] feat: sync login status with dashboard

---
 src/components/Header.tsx          | 22 +++++++++++--
 src/context/websocket-context.ts   |  2 ++
 src/context/websocket-provider.tsx | 52 ++++++++++++++++++++++++++----
 3 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/src/components/Header.tsx b/src/components/Header.tsx
index 8ab54e4..404fea8 100644
--- a/src/components/Header.tsx
+++ b/src/components/Header.tsx
@@ -140,14 +140,32 @@ function Links() {
 
 function DashboardLink() {
   const { t } = useTranslation()
+  const { reconnect } = useWebSocketContext()
+  const initRef = useRef(false)
   const { data: userData } = useQuery({
     queryKey: ["login-user"],
     queryFn: () => fetchLoginUser(),
     refetchOnMount: true,
     refetchOnWindowFocus: true,
+    refetchIntervalInBackground: true,
     refetchInterval: 1000 * 60 * 1,
   })
 
+  let isLogin = !!userData?.data?.id
+
+  if (!document.cookie) {
+    isLogin = false
+  }
+
+  useEffect(() => {
+    // 当登录状态变化时重新连接 WebSocket
+    if (initRef.current) {
+      reconnect()
+    } else {
+      initRef.current = true
+    }
+  }, [isLogin])
+
   return (
     <div className="flex items-center gap-2">
       <a
@@ -156,8 +174,8 @@ function DashboardLink() {
         rel="noopener noreferrer"
         className="flex items-center text-nowrap gap-1 text-sm font-medium opacity-50 transition-opacity hover:opacity-100"
       >
-        {!userData?.data?.id && t("login")}
-        {userData?.data?.id && t("dashboard")}
+        {!isLogin && t("login")}
+        {isLogin && t("dashboard")}
       </a>
     </div>
   )
diff --git a/src/context/websocket-context.ts b/src/context/websocket-context.ts
index 2bf955c..682a50f 100644
--- a/src/context/websocket-context.ts
+++ b/src/context/websocket-context.ts
@@ -4,10 +4,12 @@ export interface WebSocketContextType {
   lastMessage: { data: string } | null
   connected: boolean
   messageHistory: { data: string }[]
+  reconnect: () => void
 }
 
 export const WebSocketContext = createContext<WebSocketContextType>({
   lastMessage: null,
   connected: false,
   messageHistory: [],
+  reconnect: () => {},
 })
diff --git a/src/context/websocket-provider.tsx b/src/context/websocket-provider.tsx
index b41aed7..7e818ba 100644
--- a/src/context/websocket-provider.tsx
+++ b/src/context/websocket-provider.tsx
@@ -15,8 +15,37 @@ export const WebSocketProvider: React.FC<WebSocketProviderProps> = ({ url, child
   const reconnectTimeout = useRef<NodeJS.Timeout>(null)
   const maxReconnectAttempts = 30
   const reconnectAttempts = useRef(0)
+  const isConnecting = useRef(false)
+
+  const cleanup = () => {
+    if (ws.current) {
+      // 移除所有事件监听器
+      ws.current.onopen = null
+      ws.current.onclose = null
+      ws.current.onmessage = null
+      ws.current.onerror = null
+
+      if (ws.current.readyState === WebSocket.OPEN || ws.current.readyState === WebSocket.CONNECTING) {
+        ws.current.close()
+      }
+      ws.current = null
+    }
+    if (reconnectTimeout.current) {
+      clearTimeout(reconnectTimeout.current)
+      reconnectTimeout.current = null
+    }
+    setConnected(false)
+  }
 
   const connect = () => {
+    if (isConnecting.current) {
+      console.log("Connection already in progress")
+      return
+    }
+
+    cleanup()
+    isConnecting.current = true
+
     try {
       const wsUrl = new URL(url, window.location.origin)
       wsUrl.protocol = wsUrl.protocol.replace("http", "ws")
@@ -27,13 +56,15 @@ export const WebSocketProvider: React.FC<WebSocketProviderProps> = ({ url, child
         console.log("WebSocket connected")
         setConnected(true)
         reconnectAttempts.current = 0
+        isConnecting.current = false
       }
 
       ws.current.onclose = () => {
         console.log("WebSocket disconnected")
         setConnected(false)
+        ws.current = null
+        isConnecting.current = false
 
-        // 重连逻辑
         if (reconnectAttempts.current < maxReconnectAttempts) {
           reconnectTimeout.current = setTimeout(() => {
             reconnectAttempts.current++
@@ -54,22 +85,28 @@ export const WebSocketProvider: React.FC<WebSocketProviderProps> = ({ url, child
 
       ws.current.onerror = (error) => {
         console.error("WebSocket error:", error)
+        isConnecting.current = false
       }
     } catch (error) {
       console.error("WebSocket connection error:", error)
+      isConnecting.current = false
     }
   }
 
+  const reconnect = () => {
+    reconnectAttempts.current = 0
+    // 等待一个小延时确保清理完成
+    cleanup()
+    setTimeout(() => {
+      connect()
+    }, 100)
+  }
+
   useEffect(() => {
     connect()
 
     return () => {
-      if (ws.current) {
-        ws.current.close()
-      }
-      if (reconnectTimeout.current) {
-        clearTimeout(reconnectTimeout.current)
-      }
+      cleanup()
     }
   }, [url])
 
@@ -77,6 +114,7 @@ export const WebSocketProvider: React.FC<WebSocketProviderProps> = ({ url, child
     lastMessage,
     connected,
     messageHistory, // 添加到 context value
+    reconnect, // 添加到 context value
   }
 
   return <WebSocketContext.Provider value={contextValue}>{children}</WebSocketContext.Provider>