From 42392ed9c7c42ec34d28e55292271e599e40afd4 Mon Sep 17 00:00:00 2001
From: Echo <echo@moon-dragon.us>
Date: Thu, 7 May 2026 13:39:14 +0000
Subject: [PATCH] fix: persist auth across refreshes with onFinishHydration and
 safety timeout

---
 frontend/src/App.tsx            | 95 +++++++++++++++++++++++++--------
 frontend/src/store/authStore.ts | 68 ++++++++++++-----------
 2 files changed, 110 insertions(+), 53 deletions(-)

diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index fc8a58b..85d38e9 100644
--- a/frontend/src/App.tsx
+++ b/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,33 +33,82 @@ 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 />
-      <Routes>
-        {/* Public */}
-        <Route path="/login" element={<LoginPage />} />
+      <AuthRestorer>
+        <Routes>
+          {/* Public */}
+          <Route path="/login" element={<LoginPage />} />
 
-        {/* Protected — wrapped in AppLayout with sidebar navigation */}
-        <Route element={<RequireAuth><AppLayout /></RequireAuth>}>
-          <Route path="/" element={<Navigate to="/dashboard" replace />} />
-          <Route path="/mfa/setup" element={<MfaSetupPage />} />
-          <Route path="/dashboard" element={<DashboardPage />} />
-          <Route path="/hosts" element={<HostsPage />} />
-          <Route path="/hosts/:id" element={<HostDetailPage />} />
-          <Route path="/groups" element={<GroupsPage />} />
-          <Route path="/users" element={<UsersPage />} />
-          <Route path="/jobs" element={<JobsPage />} />
-          <Route path="/deployment" element={<PatchDeploymentPage />} />
-          <Route path="/maintenance" element={<MaintenanceWindowsPage />} />
-          <Route path="/reports" element={<ReportsPage />} />
-          <Route path="/certificates" element={<CertificatesPage />} />
-          <Route path="/settings" element={<SettingsPage />} />
-        </Route>
+          {/* Protected — wrapped in AppLayout with sidebar navigation */}
+          <Route element={<RequireAuth><AppLayout /></RequireAuth>}>
+            <Route path="/" element={<Navigate to="/dashboard" replace />} />
+            <Route path="/mfa/setup" element={<MfaSetupPage />} />
+            <Route path="/dashboard" element={<DashboardPage />} />
+            <Route path="/hosts" element={<HostsPage />} />
+            <Route path="/hosts/:id" element={<HostDetailPage />} />
+            <Route path="/groups" element={<GroupsPage />} />
+            <Route path="/users" element={<UsersPage />} />
+            <Route path="/jobs" element={<JobsPage />} />
+            <Route path="/deployment" element={<PatchDeploymentPage />} />
+            <Route path="/maintenance" element={<MaintenanceWindowsPage />} />
+            <Route path="/reports" element={<ReportsPage />} />
+            <Route path="/certificates" element={<CertificatesPage />} />
+            <Route path="/settings" element={<SettingsPage />} />
+          </Route>
 
-        <Route path="*" element={<Navigate to="/dashboard" replace />} />
-      </Routes>
+          <Route path="*" element={<Navigate to="/dashboard" replace />} />
+        </Routes>
+      </AuthRestorer>
     </ThemeProvider>
   )
 }
diff --git a/frontend/src/store/authStore.ts b/frontend/src/store/authStore.ts
index d0812a3..3786222 100644
--- a/frontend/src/store/authStore.ts
+++ b/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 }),
+
+      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 (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,
+          })
+        }
+      },
     }),
     {
       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({
-                  accessToken: data.access_token,
-                  refreshToken: data.refresh_token,
-                  isAuthenticated: true,
-                  isRestoring: false,
-                })
-              })
-              .catch(() => {
-                // Refresh token expired or invalid — clear all auth state
-                useAuthStore.setState({
-                  accessToken: null,
-                  refreshToken: null,
-                  user: null,
-                  isAuthenticated: false,
-                  isRestoring: false,
-                })
-              })
-          } else {
-            // No refresh token — not logged in, skip restoration
-            useAuthStore.setState({ isRestoring: false })
-          }
-        }
-      },
     }
   )
 )