Merge pull request #19687 from BerriAI/litellm_ui_refresh_mcp

[Fix] UI - Redirect to ui/login on expired JWT

ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts

@@ -8,11 +8,12 @@ import useAuthorized from "./useAuthorized";
 // Unmock useAuthorized to test the actual implementation
 vi.unmock("@/app/(dashboard)/hooks/useAuthorized");
 
-const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock } = vi.hoisted(() => ({
+const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock, isJwtExpiredMock } = vi.hoisted(() => ({
   replaceMock: vi.fn(),
   clearTokenCookiesMock: vi.fn(),
   getProxyBaseUrlMock: vi.fn(() => "http://proxy.example"),
   getUiConfigMock: vi.fn(),
+  isJwtExpiredMock: vi.fn(),
 }));
 
 vi.mock("next/navigation", () => ({
@@ -38,6 +39,14 @@ vi.mock("@/utils/cookieUtils", async (importOriginal) => {
   };
 });
 
+vi.mock("@/utils/jwtUtils", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("@/utils/jwtUtils")>();
+  return {
+    ...actual,
+    isJwtExpired: isJwtExpiredMock,
+  };
+});
+
 const createQueryClient = () =>
   new QueryClient({
     defaultOptions: {
@@ -68,6 +77,7 @@ describe("useAuthorized", () => {
     clearTokenCookiesMock.mockReset();
     getProxyBaseUrlMock.mockClear();
     getUiConfigMock.mockReset();
+    isJwtExpiredMock.mockReset();
     clearCookie();
   });
 
@@ -78,6 +88,7 @@ describe("useAuthorized", () => {
       auto_redirect_to_sso: false,
       admin_ui_disabled: false,
     });
+    isJwtExpiredMock.mockReturnValue(false);
 
     const token = createJwt({
       key: "api-key-123",
@@ -104,6 +115,7 @@ describe("useAuthorized", () => {
     expect(result.current.disabledPersonalKeyCreation).toBe(false);
     expect(result.current.showSSOBanner).toBe(true);
     expect(replaceMock).not.toHaveBeenCalled();
+    expect(clearTokenCookiesMock).not.toHaveBeenCalled();
   });
 
   it("should clear cookies and redirect on an invalid token", async () => {
@@ -134,6 +146,7 @@ describe("useAuthorized", () => {
       auto_redirect_to_sso: false,
       admin_ui_disabled: true,
     });
+    isJwtExpiredMock.mockReturnValue(false);
 
     const token = createJwt({
       key: "api-key-123",
@@ -156,4 +169,50 @@ describe("useAuthorized", () => {
     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,
+    });
+
+    // No token cookie set
+    const { result } = renderHook(() => useAuthorized(), { wrapper });
+
+    await waitFor(() => {
+      expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login");
+    });
+
+    expect(clearTokenCookiesMock).not.toHaveBeenCalled();
+    expect(result.current.token).toBeNull();
+  });
+
+  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,
+    });
+    isJwtExpiredMock.mockReturnValue(true);
+
+    const token = createJwt({
+      key: "api-key-123",
+      user_id: "user-1",
+      user_email: "user@example.com",
+      user_role: "app_admin",
+    });
+    document.cookie = `token=${token}; path=/;`;
+
+    const { result } = renderHook(() => useAuthorized(), { wrapper });
+
+    await waitFor(() => {
+      expect(clearTokenCookiesMock).toHaveBeenCalled();
+    });
+
+    expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login");
+    expect(isJwtExpiredMock).toHaveBeenCalledWith(token);
+  });
 });

ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts

@@ -2,6 +2,7 @@
 
 import { getProxyBaseUrl } from "@/components/networking";
 import { clearTokenCookies, getCookie } from "@/utils/cookieUtils";
+import { isJwtExpired } from "@/utils/jwtUtils";
 import { jwtDecode } from "jwt-decode";
 import { useRouter } from "next/navigation";
 import { useEffect, useMemo } from "react";
@@ -42,15 +43,24 @@ const useAuthorized = () => {
 
   const token = typeof document !== "undefined" ? getCookie("token") : null;
 
-  // Redirect after mount if missing/invalid token
+  // Step 1: Check for missing token or expired JWT - kick out immediately (even if UI Config is loading)
+  useEffect(() => {
+    if (!token || (token && isJwtExpired(token))) {
+      if (token) {
+        clearTokenCookies();
+      }
+      router.replace(`${getProxyBaseUrl()}/ui/login`);
+    }
+  }, [token, router]);
+
   useEffect(() => {
     if (isUIConfigLoading) {
       return;
     }
-    if (!token || uiConfig?.admin_ui_disabled) {
+    if (uiConfig?.admin_ui_disabled) {
       router.replace(`${getProxyBaseUrl()}/ui/login`);
     }
-  }, [token, router, isUIConfigLoading, uiConfig]);
+  }, [router, isUIConfigLoading, uiConfig]);
 
   // Decode safely
   const decoded = useMemo(() => {