From 3ece94bb7ad6088ce1bd85fdbac510feef2e0797 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 9 Jun 2026 12:09:48 -0700 Subject: [PATCH] refactor(ui): single auth engine; useAuthorized becomes a policy layer over AuthContext The dashboard had two parallel auth systems: AuthContext (provided at the root, used by the legacy ?page= shell) and useAuthorized (used by ~85 files in the migrated tree), each reading the cookie, decoding the JWT, and deciding when to redirect on its own. Only AuthContext applied the decoded auth_header_name via setGlobalLitellmHeaderName, so the two trees could disagree on custom auth headers, and each ran its own login-redirect effect AuthContext stays the engine: it still resolves uiConfig before clearing authLoading (so proxy-rooted URLs are correct) and performs the single decode. useAuthorized now consumes the context and only layers on policy: the admin_ui_disabled check and the redirect-to-login side effect. Its return shape is unchanged, so none of the ~85 consumers are touched Context semantics are aligned to what the hook's consumers already expected: userRole defaults to formatUserRole("") instead of "", and showSSOBanner is computed strictly from login_method (default false) instead of defaulting to true. The context's unused public setters (setToken, setUserID, setAccessToken, setPremiumUser, setShowSSOBanner) are dropped from the exported surface; only setUserRole/setUserEmail have consumers. Decoded-JWT fields keep their legacy `any` typing at the hook boundary since ~25 call sites rely on it; tightening is a follow-up The hook's tests now exercise the real provider + hook together with real JWTs instead of mocking the decode away, which also covers the provider's expiry and undecodable-token paths --- .../(dashboard)/hooks/useAuthorized.test.ts | 161 ++++++------------ .../app/(dashboard)/hooks/useAuthorized.ts | 55 +++--- .../src/contexts/AuthContext.tsx | 18 +- ui/litellm-dashboard/tests/setupTests.ts | 2 + 4 files changed, 89 insertions(+), 147 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts index 94f9d9173f0..e435d9ac242 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts @@ -3,28 +3,20 @@ import React from "react"; import { renderHook, waitFor } from "@testing-library/react"; import { afterEach, describe, expect, it, vi } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { AuthProvider } from "@/contexts/AuthContext"; import useAuthorized from "./useAuthorized"; // Unmock useAuthorized to test the actual implementation vi.unmock("@/app/(dashboard)/hooks/useAuthorized"); -const { - replaceMock, - clearTokenCookiesMock, - getProxyBaseUrlMock, - getUiConfigMock, - decodeTokenMock, - checkTokenValidityMock, - buildLoginUrlWithReturnMock, -} = vi.hoisted(() => ({ - replaceMock: vi.fn(), - clearTokenCookiesMock: vi.fn(), - getProxyBaseUrlMock: vi.fn(() => "http://proxy.example"), - getUiConfigMock: vi.fn(), - decodeTokenMock: vi.fn(), - checkTokenValidityMock: vi.fn(), - buildLoginUrlWithReturnMock: vi.fn((baseUrl: string) => baseUrl), -})); +const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock, buildLoginUrlWithReturnMock } = + vi.hoisted(() => ({ + replaceMock: vi.fn(), + clearTokenCookiesMock: vi.fn(), + getProxyBaseUrlMock: vi.fn(() => "http://proxy.example"), + getUiConfigMock: vi.fn(), + buildLoginUrlWithReturnMock: vi.fn((baseUrl: string) => baseUrl), + })); vi.mock("next/navigation", () => ({ useRouter: () => ({ @@ -49,15 +41,6 @@ vi.mock("@/utils/cookieUtils", async (importOriginal) => { }; }); -vi.mock("@/utils/jwtUtils", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - decodeToken: decodeTokenMock, - checkTokenValidity: checkTokenValidityMock, - }; -}); - vi.mock("@/utils/returnUrlUtils", async (importOriginal) => { const actual = await importOriginal(); return { @@ -66,6 +49,7 @@ vi.mock("@/utils/returnUrlUtils", async (importOriginal) => { storeReturnUrl: vi.fn(), }; }); + const createQueryClient = () => new QueryClient({ defaultOptions: { @@ -78,7 +62,11 @@ const createQueryClient = () => const wrapper = ({ children }: { children: React.ReactNode }) => { const queryClient = createQueryClient(); - return React.createElement(QueryClientProvider, { client: queryClient }, children); + return React.createElement( + QueryClientProvider, + { client: queryClient }, + React.createElement(AuthProvider, null, children), + ); }; const createJwt = (payload: Record) => { @@ -90,39 +78,37 @@ const clearCookie = () => { document.cookie = "token=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=/;"; }; +const uiConfig = (overrides: Record = {}) => ({ + server_root_path: "/", + proxy_base_url: null, + auto_redirect_to_sso: false, + admin_ui_disabled: false, + sso_configured: false, + ...overrides, +}); + +const decodedPayload = { + key: "api-key-123", + user_id: "user-1", + user_email: "user@example.com", + user_role: "app_admin", + premium_user: true, + disabled_non_admin_personal_key_creation: false, + login_method: "username_password", +}; + describe("useAuthorized", () => { afterEach(() => { replaceMock.mockReset(); clearTokenCookiesMock.mockReset(); getProxyBaseUrlMock.mockClear(); getUiConfigMock.mockReset(); - decodeTokenMock.mockReset(); - checkTokenValidityMock.mockReset(); buildLoginUrlWithReturnMock.mockClear(); clearCookie(); }); it("should decode the token and expose user details", async () => { - getUiConfigMock.mockResolvedValue({ - server_root_path: "/", - proxy_base_url: null, - auto_redirect_to_sso: false, - admin_ui_disabled: false, - sso_configured: false, - }); - - const decodedPayload = { - key: "api-key-123", - user_id: "user-1", - user_email: "user@example.com", - user_role: "app_admin", - premium_user: true, - disabled_non_admin_personal_key_creation: false, - login_method: "username_password", - }; - - decodeTokenMock.mockReturnValue(decodedPayload); - checkTokenValidityMock.mockReturnValue(true); + getUiConfigMock.mockResolvedValue(uiConfig()); const token = createJwt(decodedPayload); document.cookie = `token=${token}; path=/;`; @@ -133,6 +119,7 @@ describe("useAuthorized", () => { expect(result.current.token).toBe(token); }); + expect(result.current.isAuthorized).toBe(true); expect(result.current.accessToken).toBe("api-key-123"); expect(result.current.userId).toBe("user-1"); expect(result.current.userEmail).toBe("user@example.com"); @@ -144,52 +131,25 @@ describe("useAuthorized", () => { expect(clearTokenCookiesMock).not.toHaveBeenCalled(); }); - it("should clear cookies and redirect on an invalid token", async () => { - getUiConfigMock.mockResolvedValue({ - server_root_path: "/", - proxy_base_url: null, - auto_redirect_to_sso: false, - admin_ui_disabled: false, - sso_configured: false, - }); - - decodeTokenMock.mockReturnValue(null); - checkTokenValidityMock.mockReturnValue(false); + it("should clear cookies and redirect on an undecodable token", async () => { + getUiConfigMock.mockResolvedValue(uiConfig()); document.cookie = "token=invalid-token; path=/;"; const { result } = renderHook(() => useAuthorized(), { wrapper }); await waitFor(() => { - expect(clearTokenCookiesMock).toHaveBeenCalled(); + expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login"); }); - expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login"); + expect(clearTokenCookiesMock).toHaveBeenCalled(); + expect(result.current.token).toBeNull(); expect(result.current.accessToken).toBeNull(); expect(result.current.userRole).toBe("Undefined Role"); }); it("should redirect even with valid token if admin_ui_disabled is true", async () => { - getUiConfigMock.mockResolvedValue({ - server_root_path: "/", - proxy_base_url: null, - auto_redirect_to_sso: false, - admin_ui_disabled: true, - sso_configured: false, - }); - - const decodedPayload = { - key: "api-key-123", - user_id: "user-1", - user_email: "user@example.com", - user_role: "app_admin", - premium_user: true, - disabled_non_admin_personal_key_creation: false, - login_method: "username_password", - }; - - decodeTokenMock.mockReturnValue(decodedPayload); - checkTokenValidityMock.mockReturnValue(true); + getUiConfigMock.mockResolvedValue(uiConfig({ admin_ui_disabled: true })); const token = createJwt(decodedPayload); document.cookie = `token=${token}; path=/;`; @@ -200,22 +160,16 @@ describe("useAuthorized", () => { expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login"); }); + expect(clearTokenCookiesMock).toHaveBeenCalled(); + expect(result.current.isAuthorized).toBe(false); + expect(result.current.token).toBeNull(); expect(result.current.accessToken).toBe("api-key-123"); expect(result.current.userId).toBe("user-1"); expect(result.current.userEmail).toBe("user@example.com"); }); it("should redirect when token is missing", async () => { - getUiConfigMock.mockResolvedValue({ - server_root_path: "/", - proxy_base_url: null, - auto_redirect_to_sso: false, - admin_ui_disabled: false, - sso_configured: false, - }); - - decodeTokenMock.mockReturnValue(null); - checkTokenValidityMock.mockReturnValue(false); + getUiConfigMock.mockResolvedValue(uiConfig()); // No token cookie set const { result } = renderHook(() => useAuthorized(), { wrapper }); @@ -229,25 +183,9 @@ describe("useAuthorized", () => { }); it("should clear cookies and redirect when token is expired", async () => { - getUiConfigMock.mockResolvedValue({ - server_root_path: "/", - proxy_base_url: null, - auto_redirect_to_sso: false, - admin_ui_disabled: false, - sso_configured: false, - }); + getUiConfigMock.mockResolvedValue(uiConfig()); - const decodedPayload = { - key: "api-key-123", - user_id: "user-1", - user_email: "user@example.com", - user_role: "app_admin", - }; - - decodeTokenMock.mockReturnValue(decodedPayload); - checkTokenValidityMock.mockReturnValue(false); - - const token = createJwt(decodedPayload); + const token = createJwt({ ...decodedPayload, exp: Math.floor(Date.now() / 1000) - 60 }); document.cookie = `token=${token}; path=/;`; const { result } = renderHook(() => useAuthorized(), { wrapper }); @@ -257,6 +195,7 @@ describe("useAuthorized", () => { }); expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login"); - expect(checkTokenValidityMock).toHaveBeenCalledWith(token); + expect(result.current.token).toBeNull(); + expect(result.current.accessToken).toBeNull(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts index 8f8c403a4e9..622f96c2ce0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts @@ -1,34 +1,47 @@ "use client"; import { getProxyBaseUrl } from "@/components/networking"; -import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; -import { checkTokenValidity, decodeToken } from "@/utils/jwtUtils"; +import { clearTokenCookies } from "@/utils/cookieUtils"; import { buildLoginUrlWithReturn, storeReturnUrl } from "@/utils/returnUrlUtils"; import { useRouter } from "next/navigation"; -import { useCallback, useEffect, useMemo } from "react"; -import { formatUserRole } from "@/utils/roles"; +import { useCallback, useEffect } from "react"; +import { useAuth } from "@/contexts/AuthContext"; import { useUIConfig } from "./uiConfig/useUIConfig"; +// Decoded-JWT fields keep the pre-consolidation `any` typing: the old hook returned +// `any` here and ~25 call sites pass these where `string` is expected. Tightening to +// the context's `string | null` is a follow-up that has to fix those call sites. +type LegacyDecodedField = any; + +/** + * Policy hook for pages that require an authenticated user. Auth state itself + * lives in AuthContext (single decode at the root); this hook layers on the + * admin_ui_disabled check and the redirect-to-login side effect. + */ const useAuthorized = () => { const router = useRouter(); const { data: uiConfig, isLoading: isUIConfigLoading } = useUIConfig(); + const { + authLoading, + token, + userID, + userRole, + userEmail, + accessToken, + premiumUser, + disabledPersonalKeyCreation, + showSSOBanner, + } = useAuth(); - const token = typeof document !== "undefined" ? getCookie("token") : null; + const isLoading = authLoading || isUIConfigLoading; + const isAuthorized = token !== null && !uiConfig?.admin_ui_disabled; - const decoded = useMemo(() => decodeToken(token), [token]); - const isTokenValid = useMemo(() => checkTokenValidity(token), [token]); - const isLoading = isUIConfigLoading; - const isAuthorized = isTokenValid && !uiConfig?.admin_ui_disabled; - - // Helper function to redirect to login while preserving the current URL const redirectToLogin = useCallback(() => { storeReturnUrl(); const baseLoginUrl = `${getProxyBaseUrl()}/ui/login`; - const loginUrlWithReturn = buildLoginUrlWithReturn(baseLoginUrl); - router.replace(loginUrlWithReturn); + router.replace(buildLoginUrlWithReturn(baseLoginUrl)); }, [router]); - // Single useEffect for all redirect logic useEffect(() => { if (isLoading) return; @@ -44,13 +57,13 @@ const useAuthorized = () => { isLoading, isAuthorized, token: isAuthorized ? token : null, - accessToken: decoded?.key ?? null, - userId: decoded?.user_id ?? null, - userEmail: decoded?.user_email ?? null, - userRole: formatUserRole(decoded?.user_role), - premiumUser: decoded?.premium_user ?? null, - disabledPersonalKeyCreation: decoded?.disabled_non_admin_personal_key_creation ?? null, - showSSOBanner: decoded?.login_method === "username_password", + accessToken: accessToken as LegacyDecodedField, + userId: userID as LegacyDecodedField, + userEmail: userEmail as LegacyDecodedField, + userRole, + premiumUser: premiumUser as LegacyDecodedField, + disabledPersonalKeyCreation: disabledPersonalKeyCreation as LegacyDecodedField, + showSSOBanner, }; }; diff --git a/ui/litellm-dashboard/src/contexts/AuthContext.tsx b/ui/litellm-dashboard/src/contexts/AuthContext.tsx index 3693d858952..591f6ad1a21 100644 --- a/ui/litellm-dashboard/src/contexts/AuthContext.tsx +++ b/ui/litellm-dashboard/src/contexts/AuthContext.tsx @@ -25,13 +25,8 @@ type AuthContextValue = { disabledPersonalKeyCreation: boolean; showSSOBanner: boolean; - setToken: React.Dispatch>; - setUserID: React.Dispatch>; setUserRole: React.Dispatch>; setUserEmail: React.Dispatch>; - setAccessToken: React.Dispatch>; - setPremiumUser: React.Dispatch>; - setShowSSOBanner: React.Dispatch>; }; const AuthContext = createContext(null); @@ -40,12 +35,12 @@ export function AuthProvider({ children }: { children: React.ReactNode }) { const [authLoading, setAuthLoading] = useState(true); const [token, setToken] = useState(null); const [userID, setUserID] = useState(null); - const [userRole, setUserRole] = useState(""); + const [userRole, setUserRole] = useState(formatUserRole("")); const [userEmail, setUserEmail] = useState(null); const [accessToken, setAccessToken] = useState(null); const [premiumUser, setPremiumUser] = useState(false); const [disabledPersonalKeyCreation, setDisabledPersonalKeyCreation] = useState(false); - const [showSSOBanner, setShowSSOBanner] = useState(true); + const [showSSOBanner, setShowSSOBanner] = useState(false); // Load runtime UI config (populates proxyBaseUrl etc.) before clearing // authLoading, so any consumer that builds proxy-rooted URLs from authLoading=false @@ -105,6 +100,7 @@ export function AuthProvider({ children }: { children: React.ReactNode }) { setAccessToken(decoded.key); setDisabledPersonalKeyCreation(decoded.disabled_non_admin_personal_key_creation); + setShowSSOBanner(decoded.login_method === "username_password"); if (decoded.user_role) { setUserRole(formatUserRole(decoded.user_role)); @@ -112,9 +108,6 @@ export function AuthProvider({ children }: { children: React.ReactNode }) { 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); } @@ -136,13 +129,8 @@ export function AuthProvider({ children }: { children: React.ReactNode }) { premiumUser, disabledPersonalKeyCreation, showSSOBanner, - setToken, - setUserID, setUserRole, setUserEmail, - setAccessToken, - setPremiumUser, - setShowSSOBanner, }; return {children}; diff --git a/ui/litellm-dashboard/tests/setupTests.ts b/ui/litellm-dashboard/tests/setupTests.ts index 6bc2e7775a1..c38ab71a4eb 100644 --- a/ui/litellm-dashboard/tests/setupTests.ts +++ b/ui/litellm-dashboard/tests/setupTests.ts @@ -138,6 +138,8 @@ vi.mock("@tremor/react", async (importOriginal) => { // Global mock for useAuthorized hook to avoid repeating the same mock in every test file vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => ({ + isLoading: false, + isAuthorized: true, token: "123", accessToken: "123", userId: "user-1",