refactor(ui): extract auth state into AuthContext (#28910)

* refactor(ui): extract auth state into AuthContext

Move auth state (token, userID, userRole, accessToken, premiumUser, userEmail,
disabledPersonalKeyCreation, showSSOBanner) out of src/app/page.tsx into a
new AuthProvider at src/contexts/AuthContext.tsx. Wrapped at the root layout
so login/onboarding/dashboard routes all have access via useAuth().

Day 1 foundation for the App Router migration: migrated (dashboard)/X/page.tsx
route entry points won't have a parent passing props, so shared auth state
must live in a context they can read from.

Sub-components are unchanged — they still receive accessToken/userID/userRole
as props from page.tsx (which now reads them from useAuth()). Only the
page.tsx → top-level-page-component handoff is de-drilled; deeper prop
drilling is left for the per-page migration to address.

Net change: -86 lines from page.tsx (state + two effects moved), +5 in
layout.tsx (provider wrap), new AuthContext.tsx (~140 lines), test update
to wrap CreateKeyPage in AuthProvider.

Fixes LIT-3366
Part of LIT-3128

* fix(ui): await getUiConfig before clearing authLoading

The AuthContext refactor flipped authLoading to false synchronously on mount
while letting getUiConfig() run fire-and-forget. On SERVER_ROOT_PATH deployments
this races the unauthenticated login-redirect effect: the redirect fires with
proxyBaseUrl still at its module-init value, sending users to /ui/login instead
of {SERVER_ROOT_PATH}/ui/login.

Restores the original sequencing inside AuthProvider's mount effect and adds a
Playwright spec wired into the existing SERVER_ROOT_PATH workflow matrix. The
spec delays the config endpoint via page.route() to make the race deterministic
across CI runners.

.github/workflows/test_server_root_path.yml

@@ -101,6 +101,31 @@ jobs:
           docker logs litellm-test
           exit 1
 
+      - name: Setup Node for Playwright
+        uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
+        with:
+          node-version: "20"
+
+      - name: Install UI deps and Chromium
+        working-directory: ui/litellm-dashboard
+        run: |
+          npm ci
+          npx playwright install --with-deps chromium
+
+      - name: Run SERVER_ROOT_PATH redirect e2e
+        working-directory: ui/litellm-dashboard
+        env:
+          SERVER_ROOT_PATH: ${{ matrix.root_path }}
+        run: npx playwright test --config=e2e_tests/serverRootPath.config.ts
+
+      - name: Upload Playwright artifacts on failure
+        if: failure()
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
+        with:
+          name: playwright-trace-${{ strategy.job-index }}
+          path: ui/litellm-dashboard/test-results/
+          retention-days: 7
+
       - name: Cleanup
         if: always()
         run: |

ui/litellm-dashboard/e2e_tests/serverRootPath.config.ts

@@ -0,0 +1,32 @@
+import { defineConfig, devices } from "@playwright/test";
+
+// Minimal config for the SERVER_ROOT_PATH redirect spec. Deliberately does NOT
+// reuse the main e2e config because:
+//   - globalSetup logs in via http://localhost:4000/ui/login, which 404s when
+//     the proxy is mounted under a non-root path.
+//   - The redirect spec must run against a clean, unauthenticated session, so
+//     no storage state should be loaded.
+export default defineConfig({
+  testDir: "./tests/login",
+  testMatch: ["serverRootPathRedirect.spec.ts"],
+  fullyParallel: false,
+  forbidOnly: !!process.env.CI,
+  retries: process.env.CI ? 2 : 0,
+  workers: 1,
+  reporter: "list",
+  use: {
+    trace: "on-first-retry",
+    actionTimeout: 15 * 1000,
+    navigationTimeout: 30 * 1000,
+  },
+  projects: [
+    {
+      name: "chromium",
+      use: { ...devices["Desktop Chrome"] },
+    },
+  ],
+  timeout: 60 * 1000,
+  expect: {
+    timeout: 10 * 1000,
+  },
+});

ui/litellm-dashboard/e2e_tests/tests/login/serverRootPathRedirect.spec.ts

@@ -0,0 +1,32 @@
+import { expect, test } from "@playwright/test";
+
+// Driven by the SERVER_ROOT_PATH env var injected by the workflow; the container
+// is booted with the same value, so the asset paths and the runtime config it
+// serves at /litellm/.well-known/litellm-ui-config will both reflect it.
+const ROOT_PATH = process.env.SERVER_ROOT_PATH ?? "";
+
+test.skip(!ROOT_PATH, "Requires SERVER_ROOT_PATH env var");
+
+// Contract: an unauthenticated visit must redirect to a login URL that preserves
+// the SERVER_ROOT_PATH prefix. The redirect URL is built client-side from
+// `proxyBaseUrl`, which is populated by an async fetch of the runtime UI config.
+// If the redirect fires before that fetch resolves, the URL is missing the
+// prefix and the user lands on a 404. To make the race deterministic across
+// runners, the config endpoint is intentionally delayed.
+test("unauth redirect preserves SERVER_ROOT_PATH prefix", async ({ page }) => {
+  // Matches both `/litellm/.well-known/litellm-ui-config` and
+  // `${SERVER_ROOT_PATH}/.well-known/litellm-ui-config` (the proxy rewrites the
+  // bundle at boot when a root path is set).
+  await page.route("**/.well-known/litellm-ui-config", async (route) => {
+    await new Promise((resolve) => setTimeout(resolve, 500));
+    await route.continue();
+  });
+
+  await page.context().clearCookies();
+
+  await page.goto(`http://localhost:4000${ROOT_PATH}/ui/?page=virtual-keys`);
+
+  await page.waitForURL((url) => url.pathname.endsWith("/ui/login"), { timeout: 15_000 });
+
+  expect(page.url()).toContain(`${ROOT_PATH}/ui/login`);
+});

ui/litellm-dashboard/src/app/layout.tsx

@@ -3,6 +3,7 @@ import { Inter } from "next/font/google";
 import "./globals.css";
 
 import AntdGlobalProvider from "@/contexts/AntdGlobalProvider";
+import { AuthProvider } from "@/contexts/AuthContext";
 import ReactQueryProvider from "@/contexts/ReactQueryProvider";
 
 const inter = Inter({ subsets: ["latin"] });
@@ -22,7 +23,9 @@ export default function RootLayout({
     <html lang="en">
       <body className={inter.className}>
         <ReactQueryProvider>
-          <AntdGlobalProvider>{children}</AntdGlobalProvider>
+          <AntdGlobalProvider>
+            <AuthProvider>{children}</AuthProvider>
+          </AntdGlobalProvider>
         </ReactQueryProvider>
       </body>
     </html>

ui/litellm-dashboard/src/app/page.tsx

@@ -20,7 +20,7 @@ import { Team } from "@/components/key_team_helpers/key_list";
 import { MCPServers } from "@/components/mcp_tools";
 import ModelHubTable from "@/components/AIHub/ModelHubTable";
 import Navbar from "@/components/navbar";
-import { getUiConfig, Organization, proxyBaseUrl, setGlobalLitellmHeaderName, getInProductNudgesCall } from "@/components/networking";
+import { Organization, proxyBaseUrl, getInProductNudgesCall } from "@/components/networking";
 import NewUsagePage from "@/components/UsagePage/components/UsagePageView";
 import OldTeams from "@/components/OldTeams";
 import { fetchUserModels, CreateKeyPrefillData } from "@/components/organisms/create_key_button";
@@ -45,24 +45,14 @@ import WorkflowRuns from "@/components/workflow_runs";
 import SpendLogsTable from "@/components/view_logs";
 import ViewUserDashboard from "@/components/view_users";
 import { ThemeProvider } from "@/contexts/ThemeContext";
-import { clearTokenCookies, getCookie } from "@/utils/cookieUtils";
-import { isJwtExpired } from "@/utils/jwtUtils";
+import { useAuth } from "@/contexts/AuthContext";
 import { buildLoginUrlWithReturn, consumeReturnUrl, isValidReturnUrl, normalizeUrlForCompare, storeReturnUrl } from "@/utils/returnUrlUtils";
-import { formatUserRole, isAdminRole } from "@/utils/roles";
+import { isAdminRole } from "@/utils/roles";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
-import { jwtDecode } from "jwt-decode";
 import { useRouter, useSearchParams } from "next/navigation";
 import { Suspense, useEffect, useMemo, useRef, useState } from "react";
 import { ConfigProvider, theme } from "antd";
 
-function deleteCookie(name: string, path = "/") {
-  // Best-effort client-side clear (works for non-HttpOnly cookies without Domain)
-  document.cookie = `${name}=; Max-Age=0; Path=${path}`;
-  if (name === "token") {
-    clearTokenCookies();
-  }
-}
-
 interface ProxySettings {
   PROXY_BASE_URL: string;
   PROXY_LOGOUT_URL: string;
@@ -77,10 +67,18 @@ interface ProxySettings {
 const LEGACY_REDIRECTS: Record<string, string> = {};
 
 function CreateKeyPageContent() {
-  const [userRole, setUserRole] = useState("");
-  const [premiumUser, setPremiumUser] = useState(false);
-  const [disabledPersonalKeyCreation, setDisabledPersonalKeyCreation] = useState(false);
-  const [userEmail, setUserEmail] = useState<null | string>(null);
+  const {
+    authLoading,
+    token,
+    userID,
+    userRole,
+    userEmail,
+    accessToken,
+    premiumUser,
+    setUserRole,
+    setUserEmail,
+  } = useAuth();
+
   const [teams, setTeams] = useState<Team[] | null>(null);
   const [keys, setKeys] = useState<null | any[]>([]);
   const [organizations, setOrganizations] = useState<Organization[]>([]);
@@ -90,14 +88,10 @@ function CreateKeyPageContent() {
     PROXY_LOGOUT_URL: "",
   });
 
-  const [showSSOBanner, setShowSSOBanner] = useState<boolean>(true);
   const router = useRouter();
   const searchParams = useSearchParams()!;
   const [modelData, setModelData] = useState<any>({ data: [] });
-  const [token, setToken] = useState<string | null>(null);
   const [createClicked, setCreateClicked] = useState<boolean>(false);
-  const [authLoading, setAuthLoading] = useState(true);
-  const [userID, setUserID] = useState<string | null>(null);
 
   // Survey state - always show by default
   const [showSurveyPrompt, setShowSurveyPrompt] = useState(true);
@@ -185,7 +179,6 @@ function CreateKeyPageContent() {
     setPage(newPage);
   };
 
-  const [accessToken, setAccessToken] = useState<string | null>(null);
   const [sidebarCollapsed, setSidebarCollapsed] = useState(false);
 
   // Track if we've already attempted a return URL redirect to prevent race conditions
@@ -201,38 +194,6 @@ function CreateKeyPageContent() {
   };
   const redirectToLogin = authLoading === false && token === null && invitation_id === null;
 
-  useEffect(() => {
-    let cancelled = false;
-
-    (async () => {
-      try {
-        await getUiConfig(); // ensures proxyBaseUrl etc. are ready
-      } catch {
-        // proceed regardless; we still need to decide auth state
-      }
-
-      if (cancelled) return;
-
-      const raw = getCookie("token");
-      const valid = raw && !isJwtExpired(raw) ? raw : null;
-
-      // If token exists but is invalid/expired, clear it so downstream code
-      // doesn't keep trying to use it and cause redirect spasms.
-      if (raw && !valid) {
-        deleteCookie("token", "/");
-      }
-
-      if (!cancelled) {
-        setToken(valid);
-        setAuthLoading(false);
-      }
-    })();
-
-    return () => {
-      cancelled = true;
-    };
-  }, []);
-
   useEffect(() => {
     if (redirectToLogin) {
       // Store the current URL so we can redirect back after login
@@ -291,62 +252,6 @@ function CreateKeyPageContent() {
     }
   }, [token]);
 
-  useEffect(() => {
-    if (!token) {
-      return;
-    }
-
-    // Defensive: re-check expiry in case cookie changed after mount
-    if (isJwtExpired(token)) {
-      deleteCookie("token", "/");
-      setToken(null);
-      return;
-    }
-
-    let decoded: any = null;
-    try {
-      decoded = jwtDecode(token);
-    } catch {
-      // Malformed token → treat as unauthenticated
-      deleteCookie("token", "/");
-      setToken(null);
-      return;
-    }
-
-    if (decoded) {
-      // set accessToken
-      setAccessToken(decoded.key);
-
-      setDisabledPersonalKeyCreation(decoded.disabled_non_admin_personal_key_creation);
-
-      // check if userRole is defined
-      if (decoded.user_role) {
-        const formattedUserRole = formatUserRole(decoded.user_role);
-        setUserRole(formattedUserRole);
-      }
-
-      if (decoded.user_email) {
-        setUserEmail(decoded.user_email);
-      }
-
-      if (decoded.login_method) {
-        setShowSSOBanner(decoded.login_method == "username_password" ? true : false);
-      }
-
-      if (decoded.premium_user) {
-        setPremiumUser(decoded.premium_user);
-      }
-
-      if (decoded.auth_header_name) {
-        setGlobalLitellmHeaderName(decoded.auth_header_name);
-      }
-
-      if (decoded.user_id) {
-        setUserID(decoded.user_id);
-      }
-    }
-  }, [token]);
-
   useEffect(() => {
     if (accessToken && userID && userRole) {
       fetchUserModels(userID, userRole, accessToken, setUserModels);

ui/litellm-dashboard/src/contexts/AuthContext.tsx

@@ -0,0 +1,157 @@
+"use client";
+
+import React, { createContext, useContext, useEffect, useState } from "react";
+import { jwtDecode } from "jwt-decode";
+import { clearTokenCookies, getCookie } from "@/utils/cookieUtils";
+import { isJwtExpired } from "@/utils/jwtUtils";
+import { formatUserRole } from "@/utils/roles";
+import { getUiConfig, setGlobalLitellmHeaderName } from "@/components/networking";
+
+function deleteCookie(name: string, path = "/") {
+  document.cookie = `${name}=; Max-Age=0; Path=${path}`;
+  if (name === "token") {
+    clearTokenCookies();
+  }
+}
+
+type AuthContextValue = {
+  authLoading: boolean;
+  token: string | null;
+  userID: string | null;
+  userRole: string;
+  userEmail: string | null;
+  accessToken: string | null;
+  premiumUser: boolean;
+  disabledPersonalKeyCreation: boolean;
+  showSSOBanner: boolean;
+
+  setToken: React.Dispatch<React.SetStateAction<string | null>>;
+  setUserID: React.Dispatch<React.SetStateAction<string | null>>;
+  setUserRole: React.Dispatch<React.SetStateAction<string>>;
+  setUserEmail: React.Dispatch<React.SetStateAction<string | null>>;
+  setAccessToken: React.Dispatch<React.SetStateAction<string | null>>;
+  setPremiumUser: React.Dispatch<React.SetStateAction<boolean>>;
+  setShowSSOBanner: React.Dispatch<React.SetStateAction<boolean>>;
+};
+
+const AuthContext = createContext<AuthContextValue | null>(null);
+
+export function AuthProvider({ children }: { children: React.ReactNode }) {
+  const [authLoading, setAuthLoading] = useState(true);
+  const [token, setToken] = useState<string | null>(null);
+  const [userID, setUserID] = useState<string | null>(null);
+  const [userRole, setUserRole] = useState("");
+  const [userEmail, setUserEmail] = useState<string | null>(null);
+  const [accessToken, setAccessToken] = useState<string | null>(null);
+  const [premiumUser, setPremiumUser] = useState(false);
+  const [disabledPersonalKeyCreation, setDisabledPersonalKeyCreation] = useState(false);
+  const [showSSOBanner, setShowSSOBanner] = useState(true);
+
+  // Load runtime UI config (populates proxyBaseUrl etc.) before clearing
+  // authLoading, so any consumer that builds proxy-rooted URLs from authLoading=false
+  // (e.g. the unauthenticated login redirect) sees the resolved value rather than
+  // the module-init default. Then read the cookie and validate JWT expiry.
+  useEffect(() => {
+    let cancelled = false;
+
+    (async () => {
+      try {
+        await getUiConfig();
+      } catch {
+        // proceed regardless; auth state must still be resolved
+      }
+
+      if (cancelled) return;
+
+      const raw = getCookie("token");
+      const valid = raw && !isJwtExpired(raw) ? raw : null;
+
+      // Clear expired/invalid token so downstream code doesn't keep trying to use it.
+      if (raw && !valid) {
+        deleteCookie("token", "/");
+      }
+
+      setToken(valid);
+      setAuthLoading(false);
+    })();
+
+    return () => {
+      cancelled = true;
+    };
+  }, []);
+
+  // Decode JWT and populate derived auth state whenever the token changes.
+  useEffect(() => {
+    if (!token) {
+      return;
+    }
+
+    if (isJwtExpired(token)) {
+      deleteCookie("token", "/");
+      setToken(null);
+      return;
+    }
+
+    let decoded: { [k: string]: any } | null = null;
+    try {
+      decoded = jwtDecode(token);
+    } catch {
+      deleteCookie("token", "/");
+      setToken(null);
+      return;
+    }
+
+    if (!decoded) return;
+
+    setAccessToken(decoded.key);
+    setDisabledPersonalKeyCreation(decoded.disabled_non_admin_personal_key_creation);
+
+    if (decoded.user_role) {
+      setUserRole(formatUserRole(decoded.user_role));
+    }
+    if (decoded.user_email) {
+      setUserEmail(decoded.user_email);
+    }
+    if (decoded.login_method) {
+      setShowSSOBanner(decoded.login_method === "username_password");
+    }
+    if (decoded.premium_user) {
+      setPremiumUser(decoded.premium_user);
+    }
+    if (decoded.auth_header_name) {
+      setGlobalLitellmHeaderName(decoded.auth_header_name);
+    }
+    if (decoded.user_id) {
+      setUserID(decoded.user_id);
+    }
+  }, [token]);
+
+  const value: AuthContextValue = {
+    authLoading,
+    token,
+    userID,
+    userRole,
+    userEmail,
+    accessToken,
+    premiumUser,
+    disabledPersonalKeyCreation,
+    showSSOBanner,
+    setToken,
+    setUserID,
+    setUserRole,
+    setUserEmail,
+    setAccessToken,
+    setPremiumUser,
+    setShowSSOBanner,
+  };
+
+  return <AuthContext.Provider value={value}>{children}</AuthContext.Provider>;
+}
+
+export function useAuth(): AuthContextValue {
+  const ctx = useContext(AuthContext);
+  if (!ctx) {
+    throw new Error("useAuth must be used within an AuthProvider");
+  }
+  return ctx;
+}

ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx

@@ -147,6 +147,18 @@ vi.mock("@/lib/cva.config", () => ({
 }));
 
 import CreateKeyPage from "@/app/page";
+import { AuthProvider } from "@/contexts/AuthContext";
+
+// The page consumes auth state via useAuth(). Wrap it so the hook resolves
+// against a real provider — the provider's effects (cookie read, JWT decode,
+// redirect-on-expired) are what these tests exercise.
+function PageUnderTest() {
+  return (
+    <AuthProvider>
+      <CreateKeyPage />
+    </AuthProvider>
+  );
+}
 
 /** ----------------------------
  * Helpers
@@ -207,7 +219,7 @@ describe("CreateKeyPage auth behavior", () => {
     const cookieSetSpy = vi.spyOn(document, "cookie", "set");
 
     // Act
-    render(<CreateKeyPage />);
+    render(<PageUnderTest />);
 
     // Assert: we eventually redirect to SSO login with return URL (single replace, not assign/href)
     await waitFor(() => {
@@ -243,7 +255,7 @@ describe("CreateKeyPage auth behavior", () => {
     });
 
     // Act
-    render(<CreateKeyPage />);
+    render(<PageUnderTest />);
 
     // Assert: no redirect
     await waitFor(() => {
@@ -286,7 +298,7 @@ describe("CreateKeyPage auth behavior", () => {
     // Return URL has the same params in a different order
     consumeReturnUrlMock.mockReturnValue("http://localhost/ui?a=1&b=2");
 
-    render(<CreateKeyPage />);
+    render(<PageUnderTest />);
 
     await waitFor(() => {
       expect(window.location.replace).not.toHaveBeenCalled();