refresh token on ws error

apps/web/src/components/realtime-switch.tsx

@@ -2,37 +2,27 @@
 import { Switch } from "@/components/ui/switch";
 import { useWebSocket } from "@/lib/hooks/useWebsocket";
 import { useRealtimeCountActions } from "@/lib/store";
-import { useAuth } from "@clerk/nextjs";
 import { useParams, useRouter } from "next/navigation";
-import React, { useEffect, useState } from "react";
+import React, { useState } from "react";
 
 export function RealtimeSwitch() {
   const router = useRouter();
   const params = useParams<{ waitlist: string }>();
   const { incrementRealtimeCount, resetRealtimeCount } = useRealtimeCountActions();
 
   const [enabled, setEnabled] = useState(false);
-  const [token, setToken] = useState<string | null>(null);
-  const { getToken } = useAuth();
 
   const { ws, reconnect } = useWebSocket(`${process.env.NEXT_PUBLIC_REALTIME_URL}/${params.waitlist}/ws`, {
     onMessage: () => {
       incrementRealtimeCount();
     },
-    onClose: () => setEnabled(false),
+    onClose: () => {
+      setEnabled(false);
+    },
     enabled,
-    authToken: token ?? undefined,
+    authToken: true,
   });
 
-  useEffect(() => {
-    async function fetchToken() {
-      const token = await getToken();
-      setToken(token);
-    }
-
-    fetchToken();
-  }, []);
-
   return (
     <div className="flex items-center gap-2">
       <p className="items-center font-medium text-lg tracking-tight">Enable Realtime</p>

apps/web/src/lib/hooks/useWebsocket.ts

@@ -1,3 +1,4 @@
+import { useAuth } from "@clerk/nextjs";
 import { useCallback, useEffect, useState } from "react";
 
 export const useWebSocket = (
@@ -6,23 +7,36 @@ export const useWebSocket = (
     onMessage: (e: MessageEvent) => void;
     onClose?: (e: CloseEvent) => void;
     enabled?: boolean;
-    authToken?: string;
+    authToken?: boolean;
   },
 ) => {
   const [ws, setWs] = useState<WebSocket | null>(null);
   const [key, setKey] = useState(0);
+  const [token, setToken] = useState<string | null>(null);
+  const { getToken } = useAuth();
 
   const link = new URL(url);
 
-  if (opts.authToken) link.searchParams.append("token", opts.authToken);
+  async function fetchToken() {
+    setToken(await getToken());
+  }
+
+  if (opts.authToken) {
+    fetchToken();
+  }
 
   const reconnect = useCallback(() => {
     setKey((prev) => prev + 1);
   }, []);
 
   useEffect(() => {
     if (opts.enabled) {
+      if (token) {
+        link.searchParams.set("token", token);
+      }
+
       const ws = new WebSocket(link);
+
       ws.onopen = () => {
         console.log("[Connected to realtime]");
       };
@@ -33,6 +47,8 @@ export const useWebSocket = (
       ws.onmessage = opts.onMessage;
       ws.onerror = (e) => {
         console.error("[Realtime error]", e);
+        console.log("refreshing token");
+        fetchToken();
       };
 
       setWs(ws);