From 7194cafbc106a5a1d5c367eb626ee0108f7fbc71 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 1 Aug 2026 16:09:30 -0700 Subject: [PATCH] fix(ui): land general login on the keys dashboard, send MCP consent to /ui/connect A keyless internal user signing in to the Admin UI was redirected off the post-login landing to /ui/connect, which renders nothing but the MCP apps panel, so a plain gateway sign-in ended on an MCP OAuth surface the user never asked for. The landing now renders the keys dashboard for every role. The key lookup that existed only to make that routing decision goes with it, along with the useKeys enabled flag it was the sole caller of and the role-hydration hold that guarded its one-frame dashboard flash The gateway DCR consent flow moves the other way. Its /authorize handed the browser to /ui/chat/integrations, whose layout hard-blocks when enable_chat_ui is off, which is the default, and client-side redirects to /ui/ without the query string; that destroys the connect_flow handle and strands the MCP client until the 600s flow cookie expires. It now lands on /ui/connect, which reads connect_flow and connect_client, mounts the consent banner and puts the apps panel in connect mode. /ui/chat/integrations keeps its connect-mode handling this release so flows sealed before the deploy still finish Resolves LIT-5104 Resolves LIT-4911 --- .../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 7177b798c5f..2835a1d75b9 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -369,7 +369,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 85a19331777..d8eab3ecb3b 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 @@ -212,7 +212,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 && } +
); }