Merge pull request #15309 from BerriAI/litellm_cookie_spasm_mitigation

potentially fixes a UI spasm issue with an expired cookie

Author: krrishdholakia

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

@@ -41,8 +41,33 @@ import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
 import { cx } from "@/lib/cva.config";
 
 function getCookie(name: string) {
-  const cookieValue = document.cookie.split("; ").find((row) => row.startsWith(name + "="));
-  return cookieValue ? cookieValue.split("=")[1] : null;
+  // Safer cookie read + decoding; handles '=' inside values
+  const match = document.cookie.split("; ").find((row) => row.startsWith(name + "="));
+  if (!match) return null;
+  const value = match.slice(name.length + 1);
+  try {
+    return decodeURIComponent(value);
+  } catch {
+    return value;
+  }
+}
+
+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}`;
+}
+
+function isJwtExpired(token: string): boolean {
+  try {
+    const decoded: any = jwtDecode(token);
+    if (decoded && typeof decoded.exp === "number") {
+      return decoded.exp * 1000 <= Date.now();
+    }
+    return false;
+  } catch {
+    // If we can't decode, treat as invalid/expired
+    return true;
+  }
 }
 
 function formatUserRole(userRole: string) {
@@ -149,17 +174,42 @@ export default function CreateKeyPage() {
   const redirectToLogin = authLoading === false && token === null && invitation_id === null;
 
   useEffect(() => {
-    const token = getCookie("token");
-    getUiConfig().then((data) => {
-      // get the information for constructing the proxy base url, and then set the token and auth loading
-      setToken(token);
-      setAuthLoading(false);
-    });
+    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) {
-      window.location.href = (proxyBaseUrl || "") + "/sso/key/generate";
+      // Replace instead of assigning to avoid back-button loops
+      const dest = (proxyBaseUrl || "") + "/sso/key/generate";
+      window.location.replace(dest);
     }
   }, [redirectToLogin]);
 
@@ -168,7 +218,23 @@ export default function CreateKeyPage() {
       return;
     }
 
-    const decoded = jwtDecode(token) as { [key: string]: any };
+    // 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);

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

@@ -0,0 +1,195 @@
+import React from "react";
+import { render, screen, waitFor } from "@testing-library/react";
+import { vi, describe, it, beforeEach, afterEach, expect } from "vitest";
+
+/** ----------------------------
+ * Hoisted helpers for mocks (required by Vitest)
+ * --------------------------- */
+const { stub, jwtDecodeMock } = vi.hoisted(() => {
+  const React = require("react");
+  const stub = (name: string) => () => React.createElement("div", { "data-testid": name });
+  return {
+    stub,
+    jwtDecodeMock: vi.fn(),
+  };
+});
+
+/** ----------------------------
+ * Mocks
+ * --------------------------- */
+
+// next/navigation: just return empty URLSearchParams (no invitation/page)
+vi.mock("next/navigation", () => ({
+  useSearchParams: () => new URLSearchParams(""),
+}));
+
+// Networking layer
+vi.mock("@/components/networking", () => {
+  return {
+    // Called on mount; we don't care about its contents, only that it resolves
+    getUiConfig: vi.fn().mockResolvedValue({}),
+    // Used to build the redirect URL
+    proxyBaseUrl: "https://example.com",
+    // Called when decoding a valid token
+    setGlobalLitellmHeaderName: vi.fn(),
+    Organization: {},
+  };
+});
+
+// jwt-decode: we’ll swap implementation per test via mockImplementation
+vi.mock("jwt-decode", () => ({
+  jwtDecode: (token: string) => jwtDecodeMock(token),
+}));
+
+// Super-light stubs for all heavy components so rendering doesn't explode
+vi.mock("@/components/navbar", () => ({ default: stub("navbar") }));
+vi.mock("@/components/user_dashboard", () => ({ default: stub("user-dashboard") }));
+vi.mock("@/components/templates/model_dashboard", () => ({ default: stub("model-dashboard") }));
+vi.mock("@/components/view_users", () => ({ default: stub("view-users") }));
+vi.mock("@/components/teams", () => ({ default: stub("teams") }));
+vi.mock("@/components/organizations", () => ({
+  default: stub("organizations"),
+  fetchOrganizations: vi.fn(), // consumed in effects
+}));
+vi.mock("@/components/admins", () => ({ default: stub("admin-panel") }));
+vi.mock("@/components/settings", () => ({ default: stub("settings") }));
+vi.mock("@/components/general_settings", () => ({ default: stub("general-settings") }));
+vi.mock("@/components/pass_through_settings", () => ({ default: stub("pass-through-settings") }));
+vi.mock("@/components/budgets/budget_panel", () => ({ default: stub("budget-panel") }));
+vi.mock("@/components/view_logs", () => ({ default: stub("spend-logs") }));
+vi.mock("@/components/model_hub_table", () => ({ default: stub("model-hub-table") }));
+vi.mock("@/components/new_usage", () => ({ default: stub("new-usage") }));
+vi.mock("@/components/api_ref", () => ({ default: stub("api-ref") }));
+vi.mock("@/components/chat_ui/ChatUI", () => ({ default: stub("chat-ui") }));
+vi.mock("@/components/leftnav", () => ({ default: stub("sidebar") }));
+vi.mock("@/components/usage", () => ({ default: stub("usage") }));
+vi.mock("@/components/cache_dashboard", () => ({ default: stub("cache-dashboard") }));
+vi.mock("@/components/guardrails", () => ({ default: stub("guardrails") }));
+vi.mock("@/components/prompts", () => ({ default: stub("prompts") }));
+vi.mock("@/components/transform_request", () => ({ default: stub("transform-request") }));
+vi.mock("@/components/mcp_tools", () => ({ MCPServers: stub("mcp-servers") }));
+vi.mock("@/components/tag_management", () => ({ default: stub("tag-management") }));
+vi.mock("@/components/vector_store_management", () => ({ default: stub("vector-stores") }));
+vi.mock("@/components/ui_theme_settings", () => ({ default: stub("ui-theme-settings") }));
+vi.mock("@/components/organisms/create_key_button", () => ({ fetchUserModels: vi.fn() }));
+vi.mock("@/components/common_components/fetch_teams", () => ({ fetchTeams: vi.fn() }));
+vi.mock("@/components/ui/ui-loading-spinner", () => ({
+  UiLoadingSpinner: stub("spinner"),
+}));
+vi.mock("@/contexts/ThemeContext", () => {
+  const React = require("react");
+  return {
+    ThemeProvider: ({ children }: any) => React.createElement(React.Fragment, null, children),
+  };
+});
+vi.mock("@/lib/cva.config", () => ({
+  cx: (...args: string[]) => args.join(" "),
+}));
+
+import CreateKeyPage from "@/app/page";
+
+/** ----------------------------
+ * Helpers
+ * --------------------------- */
+
+function setCookie(raw: string) {
+  // JSDOM allows simple string assignment to document.cookie
+  document.cookie = raw;
+}
+
+function clearAllCookies() {
+  // JSDOM doesn't give an API to clear; overwrite with empty string
+  // plus ensure we wipe known names used by this app.
+  document.cookie = "token=; Max-Age=0; Path=/";
+}
+
+const originalLocation = window.location;
+
+beforeEach(() => {
+  // Fresh module state & DOM
+  vi.clearAllMocks();
+  clearAllCookies();
+
+  // Make location.replace spy-able to validate redirect
+  delete (window as any).location;
+  // minimal location object with replace and assign stubs
+  (window as any).location = {
+    ...originalLocation,
+    href: "http://localhost/",
+    assign: vi.fn(),
+    replace: vi.fn(),
+  };
+});
+
+afterEach(() => {
+  // Restore location to avoid leaking across test envs
+  delete (window as any).location;
+  (window as any).location = originalLocation;
+});
+
+/** ----------------------------
+ * Tests
+ * --------------------------- */
+
+describe("CreateKeyPage auth behavior", () => {
+  it("redirects to SSO when cookie token is expired and clears it (no spasms)", async () => {
+    // Arrange: expired token in cookie
+    setCookie("token=expiredtoken");
+
+    // jwtDecode returns past exp → expired
+    jwtDecodeMock.mockImplementation((tok: string) => {
+      expect(tok).toBe("expiredtoken");
+      return { exp: Math.floor(Date.now() / 1000) - 60 }; // expired 60s ago
+    });
+
+    // Spy on cookie writes to ensure we clear with Max-Age=0
+    const cookieSetSpy = vi.spyOn(document, "cookie", "set");
+
+    // Act
+    render(<CreateKeyPage />);
+
+    // Assert: we eventually redirect to SSO login (single replace, not assign/href)
+    await waitFor(() => {
+      expect(window.location.replace).toHaveBeenCalledWith("https://example.com/sso/key/generate");
+    });
+
+    // And we attempted to clear the cookie (defensive deletion)
+    const wroteDeletion = cookieSetSpy.mock.calls.some(
+      (args) => typeof args[0] === "string" && args[0].includes("Max-Age=0") && args[0].startsWith("token="),
+    );
+    expect(wroteDeletion).toBe(true);
+  });
+
+  it("does NOT redirect when token is valid and renders the app chrome", async () => {
+    // Arrange: valid token in cookie
+    setCookie("token=validtoken");
+
+    // jwtDecode returns future exp and expected shape
+    jwtDecodeMock.mockImplementation((tok: string) => {
+      expect(tok).toBe("validtoken");
+      return {
+        exp: Math.floor(Date.now() / 1000) + 60 * 60, // 1h in the future
+        key: "accessKey-123",
+        user_role: "app_user",
+        user_email: "[email protected]",
+        login_method: "username_password",
+        premium_user: false,
+        auth_header_name: "x-litellm-auth",
+        user_id: "u_123",
+      };
+    });
+
+    // Act
+    render(<CreateKeyPage />);
+
+    // Assert: no redirect
+    await waitFor(() => {
+      expect(window.location.replace).not.toHaveBeenCalled();
+    });
+
+    // And some top-level UI appears (Navbar stub)
+    await waitFor(() => {
+      expect(screen.getByTestId("navbar")).toBeInTheDocument();
+    });
+  });
+});