From 2a0c2c0e4162163783ff05ba58a4e803a095f294 Mon Sep 17 00:00:00 2001
From: James Magahern <james@magahern.com>
Date: Sun, 23 Feb 2025 12:34:00 -0800
Subject: [PATCH] frontend: better websocket handling

---
 frontend/src/components/App.tsx         | 43 +++++++--------
 frontend/src/hooks/useEventWebsocket.ts | 70 +++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 24 deletions(-)
 create mode 100644 frontend/src/hooks/useEventWebsocket.ts

diff --git a/frontend/src/components/App.tsx b/frontend/src/components/App.tsx
index c951088..8b3777b 100644
--- a/frontend/src/components/App.tsx
+++ b/frontend/src/components/App.tsx
@@ -3,6 +3,7 @@ import SongTable from './SongTable';
 import NowPlaying from './NowPlaying';
 import AddSongPanel from './AddSongPanel';
 import { API, getDisplayTitle, PlaylistItem } from '../api/player';
+import { useEventWebSocket } from '../hooks/useEventWebsocket';
 
 const App: React.FC = () => {
   const [isPlaying, setIsPlaying] = useState(false);
@@ -11,7 +12,6 @@ const App: React.FC = () => {
   const [volume, setVolume] = useState(100);
   const [volumeSettingIsLocked, setVolumeSettingIsLocked] = useState(false);
   const [songs, setSongs] = useState<PlaylistItem[]>([]);
-  const [ws, setWs] = useState<WebSocket | null>(null);
 
   const fetchPlaylist = useCallback(async () => {
     const playlist = await API.getPlaylist();
@@ -98,35 +98,30 @@ const App: React.FC = () => {
     }
   }, [fetchPlaylist, fetchNowPlaying]);
 
+  // Use the hook
+  useEventWebSocket(handleWebSocketEvent);
+
+  // Handle visibility changes
+  useEffect(() => {
+    const handleVisibilityChange = () => {
+      if (document.visibilityState === 'visible') {
+        fetchPlaylist();
+        fetchNowPlaying();
+      }
+    };
+
+    document.addEventListener('visibilitychange', handleVisibilityChange);
+    return () => {
+      document.removeEventListener('visibilitychange', handleVisibilityChange);
+    };
+  }, [fetchPlaylist, fetchNowPlaying]);
+
   // Initial data fetch
   useEffect(() => {
     fetchPlaylist();
     fetchNowPlaying();
   }, [fetchPlaylist, fetchNowPlaying]);
 
-  // Update WebSocket connection
-  useEffect(() => {
-    const websocket = API.subscribeToEvents(handleWebSocketEvent);
-    setWs(websocket);
-
-    // Handle page visibility changes, so if the user navigates back to this tab, we reconnect the WebSocket
-    const handleVisibilityChange = () => {
-      if (document.visibilityState === 'visible' && (!ws || ws.readyState === WebSocket.CLOSED)) {
-        const newWs = API.subscribeToEvents(handleWebSocketEvent);
-        setWs(newWs);
-      }
-    };
-
-    document.addEventListener('visibilitychange', handleVisibilityChange);
-
-    return () => {
-      document.removeEventListener('visibilitychange', handleVisibilityChange);
-      if (websocket) {
-        websocket.close();
-      }
-    };
-  }, [handleWebSocketEvent]);
-
   return (
     <div className="flex items-center justify-center h-screen w-screen bg-black md:py-10">
       <div className="bg-violet-900 w-full md:max-w-2xl h-full md:max-h-xl md:border md:rounded-2xl flex flex-col">
diff --git a/frontend/src/hooks/useEventWebsocket.ts b/frontend/src/hooks/useEventWebsocket.ts
new file mode 100644
index 0000000..5c9e226
--- /dev/null
+++ b/frontend/src/hooks/useEventWebsocket.ts
@@ -0,0 +1,70 @@
+import { useState, useEffect, useCallback } from 'react';
+import { API } from '../api/player';
+
+export const useEventWebSocket = (onEvent: (event: any) => void) => {
+  const [ws, setWs] = useState<WebSocket | null>(null);
+  const [wsReconnectTimeout, setWsReconnectTimeout] = useState<number | null>(null);
+
+  const connectWebSocket = useCallback(() => {
+    // Clear any existing reconnection timeout
+    if (wsReconnectTimeout) {
+      window.clearTimeout(wsReconnectTimeout);
+      setWsReconnectTimeout(null);
+    }
+
+    // Close existing websocket if it exists
+    if (ws) {
+      ws.close();
+    }
+
+    const websocket = API.subscribeToEvents(onEvent);
+    
+    websocket.addEventListener('close', () => {
+      console.log('WebSocket closed. Attempting to reconnect...');
+
+      // Attempt to reconnect after 2 seconds
+      const timeout = window.setTimeout(() => {
+        connectWebSocket();
+      }, 2000);
+
+      setWsReconnectTimeout(timeout);
+    });
+
+    websocket.addEventListener('error', (error) => {
+      console.error('WebSocket error:', error);
+      websocket.close();
+    });
+
+    setWs(websocket);
+  }, [wsReconnectTimeout, onEvent]);
+
+  // Check WebSocket health periodically
+  useEffect(() => {
+    const healthCheck = setInterval(() => {
+      if (!ws || ws.readyState === WebSocket.CLOSED || ws.readyState === WebSocket.CLOSING) {
+        console.log('WebSocket unhealthy, reconnecting...');
+        connectWebSocket();
+      }
+    }, 5000);
+
+    return () => {
+      clearInterval(healthCheck);
+    };
+  }, [ws, connectWebSocket]);
+
+  // Initial WebSocket connection
+  useEffect(() => {
+    connectWebSocket();
+
+    return () => {
+      if (ws) {
+        ws.close();
+      }
+      if (wsReconnectTimeout) {
+        window.clearTimeout(wsReconnectTimeout);
+      }
+    };
+  }, [connectWebSocket]);
+
+  return ws;
+}; 
\ No newline at end of file