fix(ws): only allow go to main app when ws is established

app/[locale]/layout.client.tsx

@@ -17,16 +17,22 @@ import 'react-photo-view/dist/react-photo-view.css'
 import { trpc } from '@/utils/trpc'
 import { EDeviceStatus } from '@/entities/enum'
 import { useActionDebounce } from '@/hooks/useAction'
+import { useConnectionStore } from '@/states/connection'
 
 export const ClientLayout: IComponent = ({ children }) => {
   const pathname = usePathname()
   const session = useSession()
   const router = useRouter()
   const isDarkMode = useDarkMode()
+  const { wsConnected } = useConnectionStore()
   const debounce = useActionDebounce(500, true)
   const keepAlive = trpc.userClient.ping.useMutation()
   const userStatusUpdater = trpc.userClient.updateStatus.useMutation()
 
+  const isMain = pathname.includes('/main')
+
+  const isLoading = session.status === 'loading' || (isMain && !wsConnected)
+
   const updateStatus = useCallback(
     (status: EDeviceStatus) => {
       debounce(() => {
@@ -48,14 +54,14 @@ export const ClientLayout: IComponent = ({ children }) => {
   }, [isDarkMode])
 
   useEffect(() => {
-    if (session.status === 'authenticated' && !pathname.includes('/main')) {
+    if (session.status === 'authenticated' && !isMain) {
       router.replace('/main')
     }
     if (session.status === 'unauthenticated' && !pathname.includes('/auth')) {
       router.replace('/auth/basic')
     }
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [session])
+  }, [session, isMain])
 
   useEffect(() => {
     // Update status to online when component is mounted
@@ -121,7 +127,7 @@ export const ClientLayout: IComponent = ({ children }) => {
       onIdle={() => updateStatus(EDeviceStatus.IDLE)}
       onActive={() => updateStatus(EDeviceStatus.ONLINE)}
     >
-      {session.status === 'loading' && (
+      {isLoading && (
         <div className='top-0 left-0 fixed w-screen h-[100dvh] md:h-screen z-10 bg-popover/50 flex justify-end items-end p-8'>
           <Card className='p-4 flex gap-4 items-center bg-background'>
             <LoadingSVG width={32} height={32} />
@@ -136,7 +142,7 @@ export const ClientLayout: IComponent = ({ children }) => {
             </div>
           }
         >
-          {children}
+          {((wsConnected && isMain) || (!wsConnected && !isMain)) && children}
         </PhotoProvider>
       </ReactFlowProvider>
       <Toaster />

app/[locale]/layout.trpc.tsx

@@ -1,5 +1,6 @@
 'use client'
 
+import { useConnectionStore } from '@/states/connection'
 import { setAuthToken, trpc } from '@/utils/trpc'
 import { useSession } from 'next-auth/react'
 
@@ -9,10 +10,17 @@ import { useEffect, type PropsWithChildren } from 'react'
 
 const TRPCLayout: React.FC<PropsWithChildren> = ({ children }) => {
   const { data: session } = useSession()
+  const { setWsConnected } = useConnectionStore()
 
   useEffect(() => {
     setAuthToken(session?.accessToken.token ?? '', session?.accessToken.wsToken ?? '')
-  }, [session])
+      .then((connected) => {
+        setWsConnected(connected)
+      })
+      .catch(() => {
+        alert('Websocket connection failed, please check your configuration')
+      })
+  }, [session, setWsConnected])
 
   return <>{children}</>
 }

server/run.ts

@@ -72,11 +72,6 @@ Bun.serve({
     const pathName = url.pathname
     const clientOrigin = req.headers.get('origin') || 'http://localhost:3000'
     if (pathName.startsWith('/ws')) {
-      // Queries have auth
-      if (!url.searchParams.get('auth')) {
-        // auth required
-        return new Response('UNAUTHORIZED', { status: 401 })
-      }
       if (server.upgrade(req, { data: { req: req } })) {
         return
       }

states/connection.ts

@@ -0,0 +1,10 @@
+import { create } from 'zustand'
+
+interface IState {
+  wsConnected: boolean
+  setWsConnected: (wsConnected: boolean) => void
+}
+export const useConnectionStore = create<IState>((set, get) => ({
+  wsConnected: false,
+  setWsConnected: (wsConnected) => set({ wsConnected })
+}))

utils/trpc.ts

@@ -18,16 +18,35 @@ let AuthToken = ''
 let wsAuthToken = ''
 
 export function setAuthToken(newToken: string, wsToken: string) {
-  /**
-   * You can also save the token to cookies, and initialize from
-   * cookies above.
-   */
-  AuthToken = newToken
-  wsAuthToken = wsToken
-  if (wsToken) {
-    wsClient.close()
-    wsClient.reconnect(null)
-  }
+  return new Promise<boolean>((resolve, rj) => {
+    /**
+     * You can also save the token to cookies, and initialize from
+     * cookies above.
+     */
+    AuthToken = newToken
+    if (!wsAuthToken && wsToken) {
+      wsAuthToken = wsToken
+      let checkThreshold = 0
+      if (wsClient.connection?.state === 'open') {
+        wsClient.close()
+      }
+      wsClient.reconnect(null)
+      const checkConnection = setInterval(() => {
+        if (checkThreshold === 100) {
+          // Around 20 seconds
+          clearInterval(checkConnection)
+          rj('Connection timeout')
+        }
+        if (wsClient.connection?.state === 'open') {
+          clearInterval(checkConnection)
+          resolve(true)
+        }
+        checkThreshold++
+      }, 200)
+    } else {
+      resolve(!!wsAuthToken && wsClient.connection?.state === 'open')
+    }
+  })
 }
 
 export function getBaseUrl() {