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 ceaf556b2e)

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)

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 });

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"]

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<KeysResponse> => {
   const { accessToken } = useAuthorized();
 
   return useQuery<KeysResponse>({
     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,
   });

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(<CreateKeyPage />);
-    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(<CreateKeyPage />);
+      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(<CreateKeyPage />);
-    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(<CreateKeyPage />);
-    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(<CreateKeyPage />);
-    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(<CreateKeyPage />);
-    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(<CreateKeyPage />);
-    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(<CreateKeyPage />);
-    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(<CreateKeyPage />);
-    expect(mockReplace).not.toHaveBeenCalledWith("/mocked-ui/connect");
+    expect(mockLocationReplace).toHaveBeenCalledWith("http://localhost:3000/ui/models-and-endpoints");
   });
 });

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 <LoadingScreen />;

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) => <div data-testid="mcp-apps-panel" />),
+    mockBanner: vi.fn((_props: BannerProps) => <div data-testid="connect-flow-banner" />),
   };
 });
 
 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(<ConnectPage />);
     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(<ConnectPage />);
+    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(<ConnectPage />);
+    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(<ConnectPage />);
+    expect(mockReplace).toHaveBeenCalledWith("/connect?connect_flow=flow-handle-123");
+  });
 });

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 (
     <div className="mx-auto w-full max-w-5xl px-8 py-8">
-      <MCPAppsPanel accessToken={accessToken ?? ""} selectedServers={selectedServers} onChange={setSelectedServers} />
+      {connectFlow && <ConnectFlowBanner flowHandle={connectFlow} clientOrigin={connectClient} />}
+      <MCPAppsPanel
+        accessToken={accessToken ?? ""}
+        selectedServers={selectedServers}
+        onChange={setSelectedServers}
+        connectMode={!!connectFlow}
+      />
     </div>
   );
 }