refactor: websocket connect

src/context/websocket-context.ts

@@ -1,11 +1,11 @@
 import { createContext } from "react";
 
 export interface WebSocketContextType {
-  sendMessage: (message: string) => void;
-  lastMessage: MessageEvent | null;
-  readyState: number;
+  lastMessage: { data: string } | null;
+  connected: boolean;
 }
 
-export const WebSocketContext = createContext<WebSocketContextType | undefined>(
-  undefined,
-);
+export const WebSocketContext = createContext<WebSocketContextType>({
+  lastMessage: null,
+  connected: false,
+});

src/context/websocket-provider.tsx

@@ -1,9 +1,5 @@
-import React from "react";
-import {
-  WebSocketContext,
-  type WebSocketContextType,
-} from "./websocket-context";
-import useWebSocket from "react-use-websocket";
+import React, { useEffect, useRef, useState } from "react";
+import { WebSocketContext, WebSocketContextType } from "./websocket-context";
 
 interface WebSocketProviderProps {
   url: string;
@@ -14,16 +10,67 @@ export const WebSocketProvider: React.FC<WebSocketProviderProps> = ({
   url,
   children,
 }) => {
-  const { sendMessage, lastMessage, readyState } = useWebSocket(url, {
-    reconnectAttempts: 10,
-    reconnectInterval: 3000,
-    shouldReconnect: () => true,
-  });
+  const [lastMessage, setLastMessage] = useState<{ data: string } | null>(null);
+  const [connected, setConnected] = useState(false);
+  const ws = useRef<WebSocket | null>(null);
+  const reconnectTimeout = useRef<NodeJS.Timeout>();
+  const maxReconnectAttempts = 30;
+  const reconnectAttempts = useRef(0);
+
+  const connect = () => {
+    try {
+      const wsUrl = new URL(url, window.location.origin);
+      wsUrl.protocol = wsUrl.protocol.replace("http", "ws");
+
+      ws.current = new WebSocket(wsUrl.toString());
+
+      ws.current.onopen = () => {
+        console.log("WebSocket connected");
+        setConnected(true);
+        reconnectAttempts.current = 0;
+      };
+
+      ws.current.onclose = () => {
+        console.log("WebSocket disconnected");
+        setConnected(false);
+
+        // 重连逻辑
+        if (reconnectAttempts.current < maxReconnectAttempts) {
+          reconnectTimeout.current = setTimeout(() => {
+            reconnectAttempts.current++;
+            connect();
+          }, 3000);
+        }
+      };
+
+      ws.current.onmessage = (event) => {
+        setLastMessage({ data: event.data });
+      };
+
+      ws.current.onerror = (error) => {
+        console.error("WebSocket error:", error);
+      };
+    } catch (error) {
+      console.error("WebSocket connection error:", error);
+    }
+  };
+
+  useEffect(() => {
+    connect();
+
+    return () => {
+      if (ws.current) {
+        ws.current.close();
+      }
+      if (reconnectTimeout.current) {
+        clearTimeout(reconnectTimeout.current);
+      }
+    };
+  }, [url]);
 
   const contextValue: WebSocketContextType = {
-    sendMessage,
     lastMessage,
-    readyState,
+    connected,
   };
 
   return (