fix: persist auth across refreshes with onFinishHydration and safety timeout

frontend/src/App.tsx

@@ -1,8 +1,8 @@
+import { useEffect } from 'react'
 import { Routes, Route, Navigate } from 'react-router-dom'
-import { CssBaseline, ThemeProvider } from '@mui/material'
+import { CssBaseline, ThemeProvider, CircularProgress, Box } from '@mui/material'
 import { darkTheme } from './theme/theme'
 import { useAuthStore } from './store/authStore'
-import { CircularProgress, Box } from '@mui/material'
 import AppLayout from './components/AppLayout'
 import LoginPage from './pages/LoginPage'
 import MfaSetupPage from './pages/MfaSetupPage'
@@ -33,10 +33,58 @@ function RequireAuth({ children }: { children: React.ReactNode }) {
   return isAuthenticated ? <>{children}</> : <Navigate to="/login" replace />
 }
 
+/**
+ * Waits for Zustand persist to finish rehydrating from localStorage,
+ * then calls restoreSession() so it can see the persisted refreshToken.
+ * Includes a safety timeout in case anything hangs.
+ */
+function AuthRestorer({ children }: { children: React.ReactNode }) {
+  const restoreSession = useAuthStore((s) => s.restoreSession)
+
+  useEffect(() => {
+    let cancelled = false
+
+    // Safety timeout: force isRestoring=false if restoration doesn't complete in 15s
+    const timeout = setTimeout(() => {
+      if (!cancelled) {
+        console.warn('[auth] Restoration timeout — forcing isRestoring=false')
+        useAuthStore.setState({ isRestoring: false })
+      }
+    }, 15_000)
+
+    const doRestore = () => {
+      if (!cancelled) restoreSession()
+    }
+
+    let unsub: (() => void) | undefined
+
+    // Only call restoreSession AFTER Zustand has rehydrated the persisted state
+    if (useAuthStore.persist.hasHydrated()) {
+      console.log('[auth] Store already hydrated, restoring session')
+      doRestore()
+    } else {
+      console.log('[auth] Waiting for Zustand hydration...')
+      unsub = useAuthStore.persist.onFinishHydration(() => {
+        console.log('[auth] Hydration complete, restoring session')
+        doRestore()
+      })
+    }
+
+    return () => {
+      cancelled = true
+      clearTimeout(timeout)
+      unsub?.()
+    }
+  }, [restoreSession])
+
+  return <>{children}</>
+}
+
 function App() {
   return (
     <ThemeProvider theme={darkTheme}>
       <CssBaseline />
+      <AuthRestorer>
         <Routes>
           {/* Public */}
           <Route path="/login" element={<LoginPage />} />
@@ -60,6 +108,7 @@ function App() {
 
           <Route path="*" element={<Navigate to="/dashboard" replace />} />
         </Routes>
+      </AuthRestorer>
     </ThemeProvider>
   )
 }

frontend/src/store/authStore.ts

@@ -12,11 +12,12 @@ interface AuthState {
   setTokens: (access: string, refresh: string) => void
   setUser: (user: User) => void
   logout: () => void
+  restoreSession: () => Promise<void>
 }
 
 export const useAuthStore = create<AuthState>()(
   persist(
-    (set) => ({
+    (set, get) => ({
       accessToken: null,
       refreshToken: null,
       user: null,
@@ -30,40 +31,47 @@ export const useAuthStore = create<AuthState>()(
 
       logout: () =>
         set({ accessToken: null, refreshToken: null, user: null, isAuthenticated: false, isRestoring: false }),
-    }),
-    {
-      name: 'pm-auth',
-      // Only persist refresh token; access token regenerated on load
-      partialize: (state) => ({ refreshToken: state.refreshToken, user: state.user }),
-      onRehydrateStorage: () => {
-        return (state) => {
-          if (state?.refreshToken) {
-            // Proactively refresh the access token using the persisted refresh token
-            axios.post('/api/v1/auth/refresh', { refresh_token: state.refreshToken })
-              .then(({ data }) => {
-                useAuthStore.setState({
+
+      restoreSession: async () => {
+        const { refreshToken } = get()
+        if (!refreshToken) {
+          console.log('[auth] No refresh token found, skipping restoration')
+          set({ isRestoring: false })
+          return
+        }
+
+        try {
+          const { data } = await axios.post(
+            '/api/v1/auth/refresh',
+            { refresh_token: refreshToken },
+            { timeout: 10000 }
+          )
+          console.log('[auth] Token refresh successful')
+          set({
             accessToken: data.access_token,
             refreshToken: data.refresh_token,
+            user: data.user ?? get().user,
             isAuthenticated: true,
             isRestoring: false,
           })
-              })
-              .catch(() => {
-                // Refresh token expired or invalid — clear all auth state
-                useAuthStore.setState({
+        } catch (err: unknown) {
+          const status = (err as { response?: { status?: number } })?.response?.status
+          const message = (err as Error)?.message
+          console.warn('[auth] Token refresh failed:', status, message)
+          set({
             accessToken: null,
             refreshToken: null,
             user: null,
             isAuthenticated: false,
             isRestoring: false,
           })
-              })
-          } else {
-            // No refresh token — not logged in, skip restoration
-            useAuthStore.setState({ isRestoring: false })
-          }
         }
       },
+    }),
+    {
+      name: 'pm-auth',
+      // Only persist refresh token; access token regenerated on load
+      partialize: (state) => ({ refreshToken: state.refreshToken, user: state.user }),
     }
   )
 )