diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts index 198058803eb..0df809bc582 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts @@ -101,13 +101,14 @@ export const useKeys = ( page: number, pageSize: number, options: KeyListCallOptions = {}, + enabled: boolean = true, ): UseQueryResult => { const { accessToken } = useAuthorized(); return useQuery({ queryKey: keyKeys.list({ page, limit: pageSize, ...options }), queryFn: async () => await keyListCall(accessToken!, page, pageSize, options), - enabled: Boolean(accessToken), + enabled: Boolean(accessToken) && enabled, staleTime: 30000, // 30 seconds placeholderData: keepPreviousData, }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/page.test.tsx new file mode 100644 index 00000000000..6fb64333089 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/page.test.tsx @@ -0,0 +1,111 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { render, screen } from "@testing-library/react"; +import CreateKeyPage from "./page"; + +interface KeyRow { + token: string; +} + +const { mockReplace, mockUseKeys, mockMigratedHref, state } = vi.hoisted(() => { + const state = { + login: "success" as string | null, + userRole: "Internal User", + keys: [] as KeyRow[], + keysLoading: false, + returnUrl: null as string | null, + }; + return { + state, + mockReplace: vi.fn(), + mockMigratedHref: vi.fn((segment: string) => `/mocked-ui/${segment}`), + mockUseKeys: vi.fn((_page: number, _size: number, _opts: unknown, _enabled: boolean) => ({ + data: state.keysLoading ? undefined : { keys: state.keys, total_count: state.keys.length }, + isLoading: state.keysLoading, + })), + }; +}); + +vi.mock("next/navigation", () => ({ + useRouter: () => ({ replace: mockReplace }), + useSearchParams: () => ({ get: (key: string) => (key === "login" ? state.login : null) }), +})); +vi.mock("@/contexts/AuthContext", () => ({ + useAuth: () => ({ + authLoading: false, + token: "tok", + userRole: state.userRole, + userID: "user-1", + }), +})); +vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ useKeys: mockUseKeys })); +vi.mock("@/app/(dashboard)/api-keys/ApiKeysDashboard", () => ({ + default: () =>
, +})); +vi.mock("@/components/common_components/LoadingScreen", () => ({ + default: () =>
, +})); +vi.mock("@/components/networking", () => ({ proxyBaseUrl: "" })); +vi.mock("@/utils/migratedPages", () => ({ MIGRATED_PAGES: {}, migratedHref: mockMigratedHref })); +vi.mock("@/utils/returnUrlUtils", () => ({ + buildLoginUrlWithReturn: (u: string) => u, + consumeReturnUrl: () => state.returnUrl, + getLoginUrl: () => "/login", + isValidReturnUrl: () => true, + normalizeUrlForCompare: (u: string) => u, + storeReturnUrl: () => undefined, +})); + +describe("dashboard landing keyless redirect", () => { + afterEach(() => { + state.login = "success"; + state.userRole = "Internal User"; + state.keys = []; + state.keysLoading = false; + state.returnUrl = null; + mockReplace.mockClear(); + mockUseKeys.mockClear(); + mockMigratedHref.mockClear(); + }); + + it("sends a keyless non-admin to the connect page after login", () => { + render(); + expect(mockReplace).toHaveBeenCalledWith("/mocked-ui/connect"); + expect(screen.queryByTestId("api-keys-dashboard")).not.toBeInTheDocument(); + }); + + it("leaves an admin with no keys on the dashboard", () => { + state.userRole = "Admin"; + render(); + expect(mockReplace).not.toHaveBeenCalled(); + expect(screen.getByTestId("api-keys-dashboard")).toBeInTheDocument(); + }); + + it("leaves a user who already has a key on the dashboard", () => { + state.keys = [{ token: "sk-abc" }]; + render(); + expect(mockReplace).not.toHaveBeenCalled(); + expect(screen.getByTestId("api-keys-dashboard")).toBeInTheDocument(); + }); + + it("does not redirect outside the post-login landing, and skips the key lookup entirely", () => { + state.login = null; + render(); + expect(mockReplace).not.toHaveBeenCalled(); + expect(screen.getByTestId("api-keys-dashboard")).toBeInTheDocument(); + expect(mockUseKeys.mock.calls[0][3]).toBe(false); + }); + + it("holds the loading screen while the key lookup is in flight", () => { + state.keysLoading = true; + render(); + expect(screen.getByTestId("loading-screen")).toBeInTheDocument(); + expect(screen.queryByTestId("api-keys-dashboard")).not.toBeInTheDocument(); + expect(mockReplace).not.toHaveBeenCalled(); + }); + + it("yields to an explicit return URL instead of the connect redirect", () => { + state.returnUrl = "/ui/models-and-endpoints"; + render(); + expect(mockReplace).not.toHaveBeenCalledWith("/mocked-ui/connect"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx index 02b2ccf5357..883aff8b36b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx @@ -3,6 +3,8 @@ import ApiKeysDashboard from "@/app/(dashboard)/api-keys/ApiKeysDashboard"; import LoadingScreen from "@/components/common_components/LoadingScreen"; import { proxyBaseUrl } from "@/components/networking"; +import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { isAdminRole } from "@/utils/roles"; import { useAuth } from "@/contexts/AuthContext"; import { buildLoginUrlWithReturn, @@ -17,7 +19,7 @@ import { useRouter, useSearchParams } from "next/navigation"; import { Suspense, useEffect, useRef } from "react"; function CreateKeyPageContent() { - const { authLoading, token } = useAuth(); + const { authLoading, token, userRole, userID } = useAuth(); const router = useRouter(); const searchParams = useSearchParams()!; @@ -26,6 +28,7 @@ function CreateKeyPageContent() { // Track if we've already attempted a return URL redirect to prevent race conditions const hasAttemptedReturnRedirectRef = useRef(false); + const didReturnRedirectRef = useRef(false); const redirectToLogin = authLoading === false && token === null; @@ -75,6 +78,7 @@ function CreateKeyPageContent() { // Only redirect if the return URL is different from the current URL // This prevents infinite redirect loops if (normalizedReturnUrl !== normalizedCurrentUrl) { + didReturnRedirectRef.current = true; window.location.replace(safeUrl.href); } } @@ -83,10 +87,26 @@ function CreateKeyPageContent() { useEffect(() => { if (!token) { hasAttemptedReturnRedirectRef.current = false; + didReturnRedirectRef.current = false; } }, [token]); - if (authLoading || redirectToLogin || isLegacyRedirect) { + const isPostLoginLanding = searchParams.get("login") === "success"; + const isSignedIn = !authLoading && Boolean(token); + const shouldCheckForKeys = isPostLoginLanding && isSignedIn && !isAdminRole(userRole); + const { data: keysData, isLoading: keysLoading } = useKeys(1, 1, { userID }, shouldCheckForKeys); + const isKeylessLanding = shouldCheckForKeys && !keysLoading && keysData?.keys?.length === 0; + const isResolvingKeylessLanding = (shouldCheckForKeys && keysLoading) || isKeylessLanding; + + useEffect(() => { + if (isKeylessLanding && !didReturnRedirectRef.current) { + router.replace(migratedHref("connect")); + } + }, [isKeylessLanding, router]); + + const isRedirecting = redirectToLogin || isLegacyRedirect || isResolvingKeylessLanding; + + if (authLoading || isRedirecting) { return ; }