From edc19b483063ad0a70ae697c132627b1b482e0d0 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 28 Feb 2026 12:30:47 -0800 Subject: [PATCH] [Fix] UI - Virtual Keys: Fix beforeunload listener memory leak in UserDashboard Move beforeunload event listener from the component render body into a useEffect with cleanup. Previously, every re-render added a new duplicate listener that was never removed. Co-Authored-By: Claude Opus 4.6 --- .../src/components/user_dashboard.test.tsx | 132 ++++++++++++++++++ .../src/components/user_dashboard.tsx | 18 +-- 2 files changed, 142 insertions(+), 8 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/user_dashboard.test.tsx diff --git a/ui/litellm-dashboard/src/components/user_dashboard.test.tsx b/ui/litellm-dashboard/src/components/user_dashboard.test.tsx new file mode 100644 index 0000000000..189c794224 --- /dev/null +++ b/ui/litellm-dashboard/src/components/user_dashboard.test.tsx @@ -0,0 +1,132 @@ +import { vi, describe, it, expect, beforeEach, afterEach } from "vitest"; +import { cleanup } from "@testing-library/react"; +import React from "react"; +import { renderWithProviders } from "../../tests/test-utils"; + +// Track addEventListener/removeEventListener calls for "beforeunload" +const addEventListenerSpy = vi.spyOn(window, "addEventListener"); +const removeEventListenerSpy = vi.spyOn(window, "removeEventListener"); + +// Mock next/navigation +vi.mock("next/navigation", () => ({ + useSearchParams: () => new URLSearchParams(), +})); + +// Mock networking with importOriginal so all exports are available +vi.mock("./networking", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + getProxyBaseUrl: vi.fn().mockReturnValue("http://localhost:4000"), + getProxyUISettings: vi.fn().mockResolvedValue({}), + keyInfoCall: vi.fn().mockResolvedValue({}), + modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), + userInfoCall: vi.fn().mockResolvedValue({ user_info: {}, keys: [], teams: [] }), + }; +}); + +// Mock jwt-decode to return a valid token structure +vi.mock("jwt-decode", () => ({ + jwtDecode: vi.fn().mockReturnValue({ + key: "test-access-token", + user_role: "proxy_admin", + user_email: "test@example.com", + exp: Math.floor(Date.now() / 1000) + 3600, + }), +})); + +// Mock cookie utility +vi.mock("@/utils/cookieUtils", () => ({ + clearTokenCookies: vi.fn(), +})); + +// Mock fetchTeams +vi.mock("./common_components/fetch_teams", () => ({ + fetchTeams: vi.fn(), +})); + +// Mock heavy child components to isolate UserDashboard behavior +vi.mock("./organisms/create_key_button", () => ({ + default: () =>
, +})); + +vi.mock("./VirtualKeysPage/VirtualKeysTable", () => ({ + VirtualKeysTable: () =>
, +})); + +vi.mock("../app/onboarding/page", () => ({ + default: () =>
, +})); + +// Provide a token cookie so the component doesn't redirect to login +Object.defineProperty(document, "cookie", { + writable: true, + value: "token=fake-jwt-token", +}); + +import UserDashboard from "./user_dashboard"; + +const defaultProps = { + userID: "user-1", + userRole: "Admin", + userEmail: "test@example.com", + teams: [] as any[], + keys: [] as any[], + setUserRole: vi.fn(), + setUserEmail: vi.fn(), + setTeams: vi.fn(), + setKeys: vi.fn(), + premiumUser: false, + organizations: [] as any[], + addKey: vi.fn(), + createClicked: false, +}; + +function renderDashboard(props = {}) { + return renderWithProviders(); +} + +describe("UserDashboard beforeunload listener", () => { + beforeEach(() => { + addEventListenerSpy.mockClear(); + removeEventListenerSpy.mockClear(); + }); + + afterEach(() => { + cleanup(); + }); + + it("registers exactly one beforeunload listener on mount", () => { + renderDashboard(); + + const beforeUnloadCalls = addEventListenerSpy.mock.calls.filter( + ([event]) => event === "beforeunload", + ); + expect(beforeUnloadCalls).toHaveLength(1); + }); + + it("does not add duplicate listeners on re-render", () => { + const { rerender } = renderWithProviders(); + + addEventListenerSpy.mockClear(); + + // Re-render with different props to trigger a render cycle + rerender(); + + const beforeUnloadCalls = addEventListenerSpy.mock.calls.filter( + ([event]) => event === "beforeunload", + ); + expect(beforeUnloadCalls).toHaveLength(0); + }); + + it("removes the beforeunload listener on unmount", () => { + const { unmount } = renderDashboard(); + + unmount(); + + const removeCalls = removeEventListenerSpy.mock.calls.filter( + ([event]) => event === "beforeunload", + ); + expect(removeCalls).toHaveLength(1); + }); +}); diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index ec6de82fdf..f0efe6bbaa 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -93,15 +93,17 @@ const UserDashboard: React.FC = ({ const [userModels, setUserModels] = useState([]); const [proxySettings, setProxySettings] = useState(null); const [selectedTeam, setSelectedTeam] = useState(null); - // check if window is not undefined - if (typeof window !== "undefined") { - window.addEventListener("beforeunload", function () { - // Clear session storage + + // Clear session storage on page unload so next load fetches fresh data. + // Note: MCP auth tokens are persistent and should not be cleared on page refresh + // They are only cleared on logout + useEffect(() => { + const handleBeforeUnload = () => { sessionStorage.clear(); - // Note: MCP auth tokens are persistent and should not be cleared on page refresh - // They are only cleared on logout - }); - } + }; + window.addEventListener("beforeunload", handleBeforeUnload); + return () => window.removeEventListener("beforeunload", handleBeforeUnload); + }, []); function formatUserRole(userRole: string) { if (!userRole) {