fix websocket state bug

src/renderer/src/context/websocket-context.tsx

@@ -1,5 +1,7 @@
-import React, { useContext } from 'react';
+/* eslint-disable react/jsx-no-constructed-context-values */
+import React, { useContext, useCallback } from 'react';
 import { wsService } from '@/services/websocket-service';
+import { useLocalStorage } from '@/hooks/utils/use-local-storage';
 
 const DEFAULT_WS_URL = 'ws://127.0.0.1:12393/client-ws';
 const DEFAULT_BASE_URL = 'http://127.0.0.1:12393';
@@ -44,3 +46,28 @@ export function useWebSocket() {
 
 export const defaultWsUrl = DEFAULT_WS_URL;
 export const defaultBaseUrl = DEFAULT_BASE_URL;
+
+export function WebSocketProvider({ children }: { children: React.ReactNode }) {
+  const [wsUrl, setWsUrl] = useLocalStorage('wsUrl', DEFAULT_WS_URL);
+  const [baseUrl, setBaseUrl] = useLocalStorage('baseUrl', DEFAULT_BASE_URL);
+  const handleSetWsUrl = useCallback((url: string) => {
+    setWsUrl(url);
+    wsService.connect(url);
+  }, [setWsUrl]);
+
+  const value = {
+    sendMessage: wsService.sendMessage.bind(wsService),
+    wsState: 'CLOSED',
+    reconnect: () => wsService.connect(wsUrl),
+    wsUrl,
+    setWsUrl: handleSetWsUrl,
+    baseUrl,
+    setBaseUrl,
+  };
+
+  return (
+    <WebSocketContext.Provider value={value}>
+      {children}
+    </WebSocketContext.Provider>
+  );
+}

src/renderer/src/hooks/sidebar/setting/use-general-settings.ts

@@ -120,6 +120,13 @@ export const useGeneralSettings = ({
     value: GeneralSettings[keyof GeneralSettings],
   ): void => {
     setSettings((prev) => ({ ...prev, [key]: value }));
+
+    if (key === 'wsUrl') {
+      onWsUrlChange(value as string);
+    }
+    if (key === 'baseUrl') {
+      onBaseUrlChange(value as string);
+    }
   };
 
   const handleSave = (): void => {

src/renderer/src/hooks/utils/use-local-storage.ts

@@ -4,13 +4,18 @@ export function useLocalStorage<T>(
   key: string,
   initialValue: T,
   options?: {
-    filter?: (value: T) => T
+    filter?: (value: T) => T;
+    validate?: (value: T) => boolean;
   },
 ) {
   const [storedValue, setStoredValue] = useState<T>(() => {
     try {
       const item = window.localStorage.getItem(key);
-      return item ? JSON.parse(item) : initialValue;
+      const parsedValue = item ? JSON.parse(item) : initialValue;
+      if (options?.validate && !options.validate(parsedValue)) {
+        return initialValue;
+      }
+      return parsedValue;
     } catch (error) {
       console.error(`Error reading localStorage key "${key}":`, error);
       return initialValue;
@@ -21,6 +26,10 @@ export function useLocalStorage<T>(
     try {
       const valueToStore = value instanceof Function ? value(storedValue) : value;
       const filteredValue = options?.filter ? options.filter(valueToStore) : valueToStore;
+      if (options?.validate && !options.validate(filteredValue)) {
+        console.warn(`Invalid value for localStorage key "${key}":`, filteredValue);
+        return;
+      }
       setStoredValue(valueToStore);
       window.localStorage.setItem(key, JSON.stringify(filteredValue));
     } catch (error) {

src/renderer/src/services/websocket-handler.tsx

@@ -16,11 +16,30 @@ import { useChatHistory } from '@/context/chat-history-context';
 import { toaster } from '@/components/ui/toaster';
 import { useVAD } from '@/context/vad-context';
 import { AiState, useAiState } from "@/context/ai-state-context";
+import { useLocalStorage } from '@/hooks/utils/use-local-storage';
 
 function WebSocketHandler({ children }: { children: React.ReactNode }) {
   const [wsState, setWsState] = useState<string>('CLOSED');
-  const [wsUrl, setWsUrl] = useState(defaultWsUrl);
-  const [baseUrl, setBaseUrl] = useState(defaultBaseUrl);
+  const [wsUrl, setWsUrl] = useLocalStorage<string>('wsUrl', defaultWsUrl, {
+    validate: (url: string) => {
+      try {
+        const urlObj = new URL(url);
+        return ['ws:', 'wss:'].includes(urlObj.protocol);
+      } catch (e) {
+        return false;
+      }
+    },
+  });
+  const [baseUrl, setBaseUrl] = useLocalStorage<string>('baseUrl', defaultBaseUrl, {
+    validate: (url: string) => {
+      try {
+        const urlObj = new URL(url);
+        return ['http:', 'https:'].includes(urlObj.protocol);
+      } catch (e) {
+        return false;
+      }
+    },
+  });
   const { aiState, setAiState } = useAiState();
   const { setModelInfo } = useLive2DConfig();
   const { setSubtitleText } = useSubtitle();

src/renderer/src/services/websocket-service.tsx

@@ -46,6 +46,8 @@ class WebSocketService {
 
   private stateSubject = new Subject<'CONNECTING' | 'OPEN' | 'CLOSING' | 'CLOSED'>();
 
+  private currentState: 'CONNECTING' | 'OPEN' | 'CLOSING' | 'CLOSED' = 'CLOSED';
+
   static getInstance() {
     if (!WebSocketService.instance) {
       WebSocketService.instance = new WebSocketService();
@@ -68,45 +70,77 @@ class WebSocketService {
     });
   }
 
+  private validateUrl(url: string): boolean {
+    try {
+      const urlObj = new URL(url);
+      return ['ws:', 'wss:'].includes(urlObj.protocol);
+    } catch (e) {
+      return false;
+    }
+  }
+
   connect(url: string) {
-    if (this.ws?.readyState === WebSocket.CONNECTING ||
-        this.ws?.readyState === WebSocket.OPEN) {
+    if (!this.validateUrl(url)) {
+      console.error('Invalid WebSocket URL:', url);
+      this.currentState = 'CLOSED';
+      this.stateSubject.next('CLOSED');
+      toaster.create({
+        title: `Invalid WebSocket URL: ${url}`,
+        type: 'error',
+        duration: 2000,
+      });
       return;
     }
 
-    if (this.ws) {
-      this.ws.close();
+    if (this.ws?.readyState === WebSocket.CONNECTING ||
+        this.ws?.readyState === WebSocket.OPEN) {
+      this.disconnect();
     }
 
-    this.ws = new WebSocket(url);
-    this.stateSubject.next('CONNECTING');
-
-    this.ws.onopen = () => {
-      this.stateSubject.next('OPEN');
-      this.initializeConnection();
-    };
-
-    this.ws.onmessage = (event) => {
-      try {
-        const message = JSON.parse(event.data);
-        this.messageSubject.next(message);
-      } catch (error) {
-        console.error('Failed to parse WebSocket message:', error);
-        toaster.create({
-          title: `Failed to parse WebSocket message: ${error}`,
-          type: "error",
-          duration: 2000,
-        });
-      }
-    };
-
-    this.ws.onclose = () => {
-      this.stateSubject.next('CLOSED');
-    };
-
-    this.ws.onerror = () => {
+    try {
+      this.ws = new WebSocket(url);
+      this.currentState = 'CONNECTING';
+      this.stateSubject.next('CONNECTING');
+
+      this.ws.onopen = () => {
+        this.currentState = 'OPEN';
+        this.stateSubject.next('OPEN');
+        this.initializeConnection();
+      };
+
+      this.ws.onmessage = (event) => {
+        try {
+          const message = JSON.parse(event.data);
+          this.messageSubject.next(message);
+        } catch (error) {
+          console.error('Failed to parse WebSocket message:', error);
+          toaster.create({
+            title: `Failed to parse WebSocket message: ${error}`,
+            type: "error",
+            duration: 2000,
+          });
+        }
+      };
+
+      this.ws.onclose = () => {
+        this.currentState = 'CLOSED';
+        this.stateSubject.next('CLOSED');
+      };
+
+      this.ws.onerror = () => {
+        this.currentState = 'CLOSED';
+        this.stateSubject.next('CLOSED');
+      };
+    } catch (error) {
+      console.error('Failed to connect to WebSocket:', error);
+      this.currentState = 'CLOSED';
       this.stateSubject.next('CLOSED');
-    };
+      toaster.create({
+        title: `Failed to connect to WebSocket: ${error}`,
+        type: 'error',
+        duration: 2000,
+      });
+    }
   }
 
   sendMessage(message: object) {
@@ -134,6 +168,10 @@ class WebSocketService {
     this.ws?.close();
     this.ws = null;
   }
+
+  getCurrentState() {
+    return this.currentState;
+  }
 }
 
 export const wsService = WebSocketService.getInstance();