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..89975f231aa --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/page.test.tsx @@ -0,0 +1,127 @@ +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.each(["Internal User", "Internal Viewer"])("sends a keyless %s to the connect page after login", (role) => { + state.userRole = role; + render(); + expect(mockReplace).toHaveBeenCalledWith("/mocked-ui/connect"); + expect(screen.queryByTestId("api-keys-dashboard")).not.toBeInTheDocument(); + }); + + it.each(["Admin", "Admin Viewer", "Org Admin"])("leaves a keyless %s on the dashboard", (role) => { + state.userRole = role; + 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 on the landing until the role hydrates, instead of flashing the dashboard", () => { + state.userRole = ""; + render(); + expect(screen.getByTestId("loading-screen")).toBeInTheDocument(); + expect(screen.queryByTestId("api-keys-dashboard")).not.toBeInTheDocument(); + expect(mockReplace).not.toHaveBeenCalled(); + }); + + it("does not hold the dashboard for an unhydrated role outside the post-login landing", () => { + state.login = null; + state.userRole = ""; + render(); + expect(screen.getByTestId("api-keys-dashboard")).toBeInTheDocument(); + }); + + 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..c43ba12985d 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 { internalUserRoles } 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,28 @@ 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 isAwaitingRole = isPostLoginLanding && isSignedIn && userRole === ""; + const shouldCheckForKeys = isPostLoginLanding && isSignedIn && internalUserRoles.includes(userRole); + const { data: keysData, isLoading: keysLoading } = useKeys(1, 1, { userID }, shouldCheckForKeys); + const isKeylessLanding = shouldCheckForKeys && !keysLoading && keysData?.keys?.length === 0; + const isResolvingKeylessLanding = (shouldCheckForKeys && keysLoading) || isKeylessLanding; + const isResolvingLanding = isAwaitingRole || isResolvingKeylessLanding; + + useEffect(() => { + if (isKeylessLanding && !didReturnRedirectRef.current) { + router.replace(migratedHref("connect")); + } + }, [isKeylessLanding, router]); + + const isRedirecting = redirectToLogin || isLegacyRedirect || isResolvingLanding; + + if (authLoading || isRedirecting) { return ; } diff --git a/ui/litellm-dashboard/src/app/connect/layout.test.tsx b/ui/litellm-dashboard/src/app/connect/layout.test.tsx new file mode 100644 index 00000000000..795a79d77e3 --- /dev/null +++ b/ui/litellm-dashboard/src/app/connect/layout.test.tsx @@ -0,0 +1,65 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { render, screen } from "@testing-library/react"; +import ConnectLayout from "./layout"; + +const { mockUseAuthorized, state } = vi.hoisted(() => { + const state = { + accessToken: "token-123" as string | null, + isAuthorized: true, + isLoading: false, + }; + return { + state, + mockUseAuthorized: vi.fn(() => ({ + accessToken: state.accessToken, + isAuthorized: state.isAuthorized, + isLoading: state.isLoading, + })), + }; +}); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: mockUseAuthorized })); +vi.mock("@/components/navbar", () => ({ default: () =>
})); +vi.mock("@/contexts/ThemeContext", () => ({ + ThemeProvider: ({ children }: { children: React.ReactNode }) => <>{children}, +})); + +describe("ConnectLayout", () => { + afterEach(() => { + state.accessToken = "token-123"; + state.isAuthorized = true; + state.isLoading = false; + }); + + it("renders the connect surface for an authorized user without any chat-ui flag", () => { + render( + +
+ , + ); + expect(screen.getByTestId("navbar")).toBeInTheDocument(); + expect(screen.getByTestId("page-content")).toBeInTheDocument(); + }); + + it("renders nothing when the user is not authorized", () => { + state.isAuthorized = false; + render( + +
+ , + ); + expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); + expect(screen.queryByTestId("navbar")).not.toBeInTheDocument(); + }); + + it("renders nothing while authorization is still loading", () => { + state.isLoading = true; + render( + +
+ , + ); + expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); + expect(screen.queryByTestId("navbar")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/connect/layout.tsx b/ui/litellm-dashboard/src/app/connect/layout.tsx new file mode 100644 index 00000000000..63b1c484094 --- /dev/null +++ b/ui/litellm-dashboard/src/app/connect/layout.tsx @@ -0,0 +1,20 @@ +"use client"; + +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import Navbar from "@/components/navbar"; +import { ThemeProvider } from "@/contexts/ThemeContext"; + +export default function ConnectLayout({ children }: { children: React.ReactNode }) { + const { accessToken, isAuthorized, isLoading } = useAuthorized(); + + if (isLoading || !isAuthorized) return null; + + return ( + +
+ +
{children}
+
+
+ ); +} diff --git a/ui/litellm-dashboard/src/app/connect/page.test.tsx b/ui/litellm-dashboard/src/app/connect/page.test.tsx new file mode 100644 index 00000000000..07b0e7a305a --- /dev/null +++ b/ui/litellm-dashboard/src/app/connect/page.test.tsx @@ -0,0 +1,55 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { render, screen } from "@testing-library/react"; +import ConnectPage from "./page"; + +interface PanelProps { + accessToken: string; + selectedServers: string[]; + onChange: (servers: string[]) => void; +} + +const { mockReplace, mockPanel, state } = vi.hoisted(() => { + const state = { + oauthReturn: null as string | null, + }; + return { + state, + mockReplace: vi.fn(), + mockPanel: vi.fn((_props: PanelProps) =>
), + }; +}); + +vi.mock("next/navigation", () => ({ + useRouter: () => ({ replace: mockReplace }), + useSearchParams: () => ({ get: (key: string) => (key === "mcpOauthReturn" ? state.oauthReturn : null) }), +})); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "token-123" }), +})); +vi.mock("@/components/chat/MCPAppsPanel", () => ({ default: mockPanel })); + +describe("ConnectPage", () => { + afterEach(() => { + state.oauthReturn = null; + mockReplace.mockClear(); + mockPanel.mockClear(); + }); + + it("renders the MCP connect panel with the user's access token", () => { + render(); + expect(screen.getByTestId("mcp-apps-panel")).toBeInTheDocument(); + expect(mockPanel.mock.calls[0][0]).toMatchObject({ accessToken: "token-123", selectedServers: [] }); + }); + + it("strips the mcpOauthReturn param from the URL after an OAuth return", () => { + state.oauthReturn = "apps"; + window.history.replaceState({}, "", "/connect?mcpOauthReturn=apps"); + render(); + expect(mockReplace).toHaveBeenCalledWith("/connect"); + }); + + it("does not rewrite the URL when there is no OAuth return param", () => { + render(); + expect(mockReplace).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/connect/page.tsx b/ui/litellm-dashboard/src/app/connect/page.tsx new file mode 100644 index 00000000000..84770915e46 --- /dev/null +++ b/ui/litellm-dashboard/src/app/connect/page.tsx @@ -0,0 +1,36 @@ +"use client"; + +import { Suspense, useEffect, useState } from "react"; +import { useRouter, useSearchParams } from "next/navigation"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import MCPAppsPanel from "@/components/chat/MCPAppsPanel"; + +function ConnectPageContent() { + const { accessToken } = useAuthorized(); + const [selectedServers, setSelectedServers] = useState([]); + const router = useRouter(); + const searchParams = useSearchParams(); + const oauthReturn = searchParams.get("mcpOauthReturn"); + + useEffect(() => { + if (oauthReturn) { + const url = new URL(window.location.href); + url.searchParams.delete("mcpOauthReturn"); + router.replace(url.pathname + url.search); + } + }, [oauthReturn, router]); + + return ( +
+ +
+ ); +} + +export default function ConnectPage() { + return ( + + + + ); +}