From 1197a8cb06cc1f6a2160a25c8ceaefdc648b7da3 Mon Sep 17 00:00:00 2001
From: Arthur Belleville <arthur.belleville@datadoghq.com>
Date: Thu, 16 Oct 2025 17:45:35 +0200
Subject: [PATCH] Fix loading page hanging

---
 ui/src/components/ProtectedRoute.test.tsx | 54 +++++++++++------------
 ui/src/components/ProtectedRoute.tsx      | 22 +++++----
 ui/src/providers/UserStoreProvider.tsx    |  5 ++-
 3 files changed, 39 insertions(+), 42 deletions(-)

diff --git a/ui/src/components/ProtectedRoute.test.tsx b/ui/src/components/ProtectedRoute.test.tsx
index 7de09ef..8ef964e 100644
--- a/ui/src/components/ProtectedRoute.test.tsx
+++ b/ui/src/components/ProtectedRoute.test.tsx
@@ -3,6 +3,7 @@ import { ProtectedRoute } from "@ui/components/ProtectedRoute";
 import { SessionTestProvider } from "@ui/contexts/SessionContext";
 import { renderWithRouter } from "@ui/utils/testHelpers";
 import { Route, Routes } from "react-router-dom";
+import { TestUserStoreProvider } from "src/providers/UserStoreProvider";
 
 describe("ProtectedRoute", () => {
   beforeEach(() => {
@@ -30,14 +31,16 @@ describe("ProtectedRoute", () => {
 
   it("redirects to login when user is not authenticated", async () => {
     renderWithRouter(
-      <SessionTestProvider>
-        <Routes>
-          <Route element={<ProtectedRoute />}>
-            <Route path="/" element={<div>Protected Content</div>} />
-          </Route>
-          <Route path="/login" element={<div>Login Page</div>} />
-        </Routes>
-      </SessionTestProvider>,
+      <TestUserStoreProvider user={null}>
+        <SessionTestProvider>
+          <Routes>
+            <Route element={<ProtectedRoute />}>
+              <Route path="/" element={<div>Protected Content</div>} />
+            </Route>
+            <Route path="/login" element={<div>Login Page</div>} />
+          </Routes>
+        </SessionTestProvider>
+      </TestUserStoreProvider>,
       { route: "/" }
     );
 
@@ -71,31 +74,24 @@ describe("ProtectedRoute", () => {
 
   it("renders protected content when user is authenticated", async () => {
     renderWithRouter(
-      <SessionTestProvider
-        testUser={{
+      <TestUserStoreProvider
+        user={{
           id: "123",
-          app_metadata: {},
-          user_metadata: {
-            full_name: "Test User",
-            email: "test@example.com",
-            email_verified: true,
-            first_name: "Test",
-            last_name: "User",
-            business_name: "Test Business",
-          },
-          aud: "authenticated",
-          created_at: new Date().toISOString(),
+          name: "Test User",
           email: "test@example.com",
-          role: "authenticated",
-          updated_at: new Date().toISOString(),
+          avatar_url: "https://example.com/avatar.jpg",
+          streamToken: null,
+          short_user_id: "123",
         }}
       >
-        <Routes>
-          <Route element={<ProtectedRoute />}>
-            <Route path="/" element={<div>Protected Content</div>} />
-          </Route>
-        </Routes>
-      </SessionTestProvider>,
+        <SessionTestProvider>
+          <Routes>
+            <Route element={<ProtectedRoute />}>
+              <Route path="/" element={<div>Protected Content</div>} />
+            </Route>
+          </Routes>
+        </SessionTestProvider>
+      </TestUserStoreProvider>,
       { route: "/" }
     );
 
diff --git a/ui/src/components/ProtectedRoute.tsx b/ui/src/components/ProtectedRoute.tsx
index ebc6729..46e389c 100644
--- a/ui/src/components/ProtectedRoute.tsx
+++ b/ui/src/components/ProtectedRoute.tsx
@@ -3,6 +3,7 @@ import { useEffect, useState } from "react";
 import { Navigate, Outlet } from "react-router-dom";
 import { match } from "ts-pattern";
 import { LoadingSpinner } from "./LoadingSpinner";
+import { useMaybeUser } from "src/providers/UserStoreProvider";
 
 interface ProtectedRouteProps {
   fallback?: string;
@@ -22,13 +23,14 @@ export const ProtectedRoute = ({ fallback, shouldRedirectToCurrentPage }: Protec
 
   let status: "loading" | "should-land-user" | "should-redirect" | "should-pass" = "loading";
 
+  const user = useMaybeUser();
   const isFirstTimeUser = localStorage.getItem("xtablo-has-seen-landing-page") === null;
 
   if (isLoading) {
     status = "loading";
-  } else if (!session?.user && isFirstTimeUser) {
+  } else if (!user && isFirstTimeUser) {
     status = "should-land-user";
-  } else if (!session?.user) {
+  } else if (!user) {
     status = "should-redirect";
   } else {
     status = "should-pass";
@@ -40,14 +42,10 @@ export const ProtectedRoute = ({ fallback, shouldRedirectToCurrentPage }: Protec
       )}`
     : (fallback ?? "/login");
 
-  return (
-    <>
-      {match(status)
-        .with("loading", () => <LoadingSpinner />)
-        .with("should-land-user", () => <Navigate to="/landing" replace />)
-        .with("should-redirect", () => <Navigate to={redirectUrl} replace />)
-        .with("should-pass", () => <Outlet />)
-        .exhaustive()}
-    </>
-  );
+  return match(status)
+    .with("loading", () => <LoadingSpinner />)
+    .with("should-land-user", () => <Navigate to="/landing" replace />)
+    .with("should-redirect", () => <Navigate to={redirectUrl} replace />)
+    .with("should-pass", () => <Outlet />)
+    .exhaustive();
 };
diff --git a/ui/src/providers/UserStoreProvider.tsx b/ui/src/providers/UserStoreProvider.tsx
index c3aae1c..2d0d573 100644
--- a/ui/src/providers/UserStoreProvider.tsx
+++ b/ui/src/providers/UserStoreProvider.tsx
@@ -72,8 +72,11 @@ export const TestUserStoreProvider = ({
   user,
 }: {
   children: React.ReactNode;
-  user: User;
+  user: User | null;
 }) => {
+  if (!user) {
+    return children;
+  }
   const store = createStore<User>()(() => user);
 
   return (