From 69f5fe35d8ac7970df175fda3a76cc4f853a59e5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 1 Aug 2026 17:17:18 -0700 Subject: [PATCH] Merge pull request #35523 from BerriAI/litellm_ui_login_no_mcp_landing fix(ui): land general login on the keys dashboard, send MCP consent to /ui/connect (cherry picked from commit ceaf556b2e7eab737517fc4e494816e4b4359c71) --- .../mcp_server/gateway_dcr_flow.py | 2 +- .../internal-user/internalUserNoTeam.spec.ts | 5 +- .../mcp_server/test_gateway_dcr_flow.py | 2 +- .../src/app/(dashboard)/hooks/keys/useKeys.ts | 3 +- .../src/app/(dashboard)/page.test.tsx | 88 ++++++++----------- .../src/app/(dashboard)/page.tsx | 24 +---- .../src/app/connect/page.test.tsx | 51 ++++++++++- ui/litellm-dashboard/src/app/connect/page.tsx | 11 ++- 8 files changed, 103 insertions(+), 83 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 58233c4c9e5..20094ce3e21 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -361,7 +361,7 @@ def aggregate_authorize( exp=int(now.timestamp()) + CONNECT_FLOW_TTL_SECONDS, ) connect_url = _append_query_params( - f"{base_url}/ui/chat/integrations", + f"{base_url}/ui/connect", {"connect_flow": handle, "connect_client": _origin_only(redirect_uri)}, ) response = RedirectResponse(connect_url, status_code=303) diff --git a/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts index 3affa4d898d..1b048198456 100644 --- a/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts @@ -17,9 +17,8 @@ test.describe("Internal User with no team memberships", () => { await page.getByPlaceholder("Enter your username").fill("noteam@test.local"); await page.getByPlaceholder("Enter your password").fill("test"); await page.getByRole("button", { name: "Login", exact: true }).click(); - // A non-admin with no keys lands on /ui/connect, so the keys dashboard has - // to be asked for explicitly once that redirect settles. - await page.waitForURL(/\/ui\/connect/, { timeout: 30_000 }); + await expect(page.getByRole("complementary").getByText("Virtual Keys")).toBeVisible({ timeout: 30_000 }); + expect(new URL(page.url()).pathname).not.toMatch(/\/connect$/); await navigateToPage(page, Page.ApiKeys); // Scope to the sidebar; the top-bar breadcrumb also shows "Virtual Keys". await expect(page.getByRole("complementary").getByText("Virtual Keys")).toBeVisible({ timeout: 15_000 }); diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 375ec022115..b7f2a93f8fb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -208,7 +208,7 @@ async def test_authorize_with_session_hands_browser_to_connect_page_with_flow_co response = _authorize(client_id, session_user_id="u1") assert response.status_code == 303 location = urlparse(response.headers["location"]) - assert location.path == "/ui/chat/integrations" + assert location.path == "/ui/connect" params = parse_qs(location.query) handle = params["connect_flow"][0] assert params["connect_client"] == ["https://claude.ai"] 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 0df809bc582..198058803eb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts @@ -101,14 +101,13 @@ 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, + enabled: Boolean(accessToken), 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 index 89975f231aa..5abb219f019 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/page.test.tsx @@ -1,4 +1,4 @@ -import { afterEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { render, screen } from "@testing-library/react"; import CreateKeyPage from "./page"; @@ -11,16 +11,15 @@ const { mockReplace, mockUseKeys, mockMigratedHref, state } = vi.hoisted(() => { 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, + mockUseKeys: vi.fn(() => ({ + data: { keys: state.keys, total_count: state.keys.length }, + isLoading: false, })), }; }); @@ -55,73 +54,60 @@ vi.mock("@/utils/returnUrlUtils", () => ({ storeReturnUrl: () => undefined, })); -describe("dashboard landing keyless redirect", () => { +const realLocation = window.location; +const mockLocationReplace = vi.fn(); + +describe("dashboard landing", () => { + beforeEach(() => { + Object.defineProperty(window, "location", { + configurable: true, + value: { + origin: "http://localhost:3000", + href: "http://localhost:3000/ui/?login=success", + replace: mockLocationReplace, + }, + }); + }); + afterEach(() => { + Object.defineProperty(window, "location", { configurable: true, value: realLocation }); state.login = "success"; state.userRole = "Internal User"; state.keys = []; - state.keysLoading = false; state.returnUrl = null; mockReplace.mockClear(); mockUseKeys.mockClear(); mockMigratedHref.mockClear(); + mockLocationReplace.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(["Internal User", "Internal Viewer", "Admin", "Admin Viewer", "Org Admin", ""])( + "lands a keyless %s on the keys dashboard, never on the MCP connect page", + (role) => { + state.userRole = role; + render(); + expect(screen.getByTestId("api-keys-dashboard")).toBeInTheDocument(); + expect(screen.queryByTestId("loading-screen")).not.toBeInTheDocument(); + expect(mockReplace).not.toHaveBeenCalled(); + expect(mockMigratedHref).not.toHaveBeenCalledWith("connect"); + }, + ); - 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", () => { + it("lands a user who already owns a key on the keys 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 = ""; + it("never looks a user's keys up to decide where the landing goes", () => { render(); - expect(screen.getByTestId("api-keys-dashboard")).toBeInTheDocument(); + expect(mockUseKeys).not.toHaveBeenCalled(); }); - 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", () => { + it("still sends the user to an explicit stored return URL", () => { state.returnUrl = "/ui/models-and-endpoints"; render(); - expect(mockReplace).not.toHaveBeenCalledWith("/mocked-ui/connect"); + expect(mockLocationReplace).toHaveBeenCalledWith("http://localhost:3000/ui/models-and-endpoints"); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx index c43ba12985d..9e82d33dc2b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx @@ -3,8 +3,6 @@ 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, @@ -19,7 +17,7 @@ import { useRouter, useSearchParams } from "next/navigation"; import { Suspense, useEffect, useRef } from "react"; function CreateKeyPageContent() { - const { authLoading, token, userRole, userID } = useAuth(); + const { authLoading, token } = useAuth(); const router = useRouter(); const searchParams = useSearchParams()!; @@ -28,7 +26,6 @@ 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; @@ -78,7 +75,6 @@ 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); } } @@ -87,26 +83,10 @@ function CreateKeyPageContent() { useEffect(() => { if (!token) { hasAttemptedReturnRedirectRef.current = false; - didReturnRedirectRef.current = false; } }, [token]); - 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; + const isRedirecting = redirectToLogin || isLegacyRedirect; if (authLoading || isRedirecting) { return ; diff --git a/ui/litellm-dashboard/src/app/connect/page.test.tsx b/ui/litellm-dashboard/src/app/connect/page.test.tsx index 07b0e7a305a..7d49a8b6a4c 100644 --- a/ui/litellm-dashboard/src/app/connect/page.test.tsx +++ b/ui/litellm-dashboard/src/app/connect/page.test.tsx @@ -6,33 +6,53 @@ interface PanelProps { accessToken: string; selectedServers: string[]; onChange: (servers: string[]) => void; + connectMode?: boolean; } -const { mockReplace, mockPanel, state } = vi.hoisted(() => { +interface BannerProps { + flowHandle: string; + clientOrigin: string | null; +} + +const { mockReplace, mockPanel, mockBanner, state } = vi.hoisted(() => { const state = { oauthReturn: null as string | null, + connectFlow: null as string | null, + connectClient: null as string | null, }; return { state, mockReplace: vi.fn(), mockPanel: vi.fn((_props: PanelProps) =>
), + mockBanner: vi.fn((_props: BannerProps) =>
), }; }); vi.mock("next/navigation", () => ({ useRouter: () => ({ replace: mockReplace }), - useSearchParams: () => ({ get: (key: string) => (key === "mcpOauthReturn" ? state.oauthReturn : null) }), + useSearchParams: () => ({ + get: (key: string) => { + if (key === "mcpOauthReturn") return state.oauthReturn; + if (key === "connect_flow") return state.connectFlow; + if (key === "connect_client") return state.connectClient; + return null; + }, + }), })); vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => ({ accessToken: "token-123" }), })); vi.mock("@/components/chat/MCPAppsPanel", () => ({ default: mockPanel })); +vi.mock("@/components/chat/ConnectFlowBanner", () => ({ default: mockBanner })); describe("ConnectPage", () => { afterEach(() => { state.oauthReturn = null; + state.connectFlow = null; + state.connectClient = null; mockReplace.mockClear(); mockPanel.mockClear(); + mockBanner.mockClear(); }); it("renders the MCP connect panel with the user's access token", () => { @@ -52,4 +72,31 @@ describe("ConnectPage", () => { render(); expect(mockReplace).not.toHaveBeenCalled(); }); + + it("mounts the gateway connect banner and puts the panel in connect mode for a DCR flow", () => { + state.connectFlow = "flow-handle-123"; + state.connectClient = "https://claude.ai"; + render(); + expect(screen.getByTestId("connect-flow-banner")).toBeInTheDocument(); + expect(mockBanner.mock.calls[0][0]).toMatchObject({ + flowHandle: "flow-handle-123", + clientOrigin: "https://claude.ai", + }); + expect(mockPanel.mock.calls[0][0].connectMode).toBe(true); + }); + + it("shows no connect banner and leaves connect mode off for a plain visit", () => { + render(); + expect(screen.queryByTestId("connect-flow-banner")).not.toBeInTheDocument(); + expect(mockBanner).not.toHaveBeenCalled(); + expect(mockPanel.mock.calls[0][0].connectMode).toBe(false); + }); + + it("keeps the connect flow handle in the URL while stripping the OAuth return param", () => { + state.oauthReturn = "apps"; + state.connectFlow = "flow-handle-123"; + window.history.replaceState({}, "", "/connect?connect_flow=flow-handle-123&mcpOauthReturn=apps"); + render(); + expect(mockReplace).toHaveBeenCalledWith("/connect?connect_flow=flow-handle-123"); + }); }); diff --git a/ui/litellm-dashboard/src/app/connect/page.tsx b/ui/litellm-dashboard/src/app/connect/page.tsx index 84770915e46..3f0c269e86b 100644 --- a/ui/litellm-dashboard/src/app/connect/page.tsx +++ b/ui/litellm-dashboard/src/app/connect/page.tsx @@ -4,6 +4,7 @@ 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"; +import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner"; function ConnectPageContent() { const { accessToken } = useAuthorized(); @@ -11,6 +12,8 @@ function ConnectPageContent() { const router = useRouter(); const searchParams = useSearchParams(); const oauthReturn = searchParams.get("mcpOauthReturn"); + const connectFlow = searchParams.get("connect_flow"); + const connectClient = searchParams.get("connect_client"); useEffect(() => { if (oauthReturn) { @@ -22,7 +25,13 @@ function ConnectPageContent() { return (
- + {connectFlow && } +
); }