Merge pull request #34334 from BerriAI/litellm_connect_page_standalone

feat(ui): standalone /connect route for MCP OAuth, decoupled from Chat UI flag

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<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: Boolean(accessToken) && enabled,
     staleTime: 30000, // 30 seconds
     placeholderData: keepPreviousData,
   });

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: () => <div data-testid="api-keys-dashboard" />,
+}));
+vi.mock("@/components/common_components/LoadingScreen", () => ({
+  default: () => <div data-testid="loading-screen" />,
+}));
+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(<CreateKeyPage />);
+    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(<CreateKeyPage />);
+    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(<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 = "";
+    render(<CreateKeyPage />);
+    expect(screen.getByTestId("api-keys-dashboard")).toBeInTheDocument();
+  });
+
+  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", () => {
+    state.returnUrl = "/ui/models-and-endpoints";
+    render(<CreateKeyPage />);
+    expect(mockReplace).not.toHaveBeenCalledWith("/mocked-ui/connect");
+  });
+});

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

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: () => <div data-testid="navbar" /> }));
+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(
+      <ConnectLayout>
+        <div data-testid="page-content" />
+      </ConnectLayout>,
+    );
+    expect(screen.getByTestId("navbar")).toBeInTheDocument();
+    expect(screen.getByTestId("page-content")).toBeInTheDocument();
+  });
+
+  it("renders nothing when the user is not authorized", () => {
+    state.isAuthorized = false;
+    render(
+      <ConnectLayout>
+        <div data-testid="page-content" />
+      </ConnectLayout>,
+    );
+    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(
+      <ConnectLayout>
+        <div data-testid="page-content" />
+      </ConnectLayout>,
+    );
+    expect(screen.queryByTestId("page-content")).not.toBeInTheDocument();
+    expect(screen.queryByTestId("navbar")).not.toBeInTheDocument();
+  });
+});

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 (
+    <ThemeProvider accessToken={accessToken}>
+      <div className="flex h-screen flex-col">
+        <Navbar accessToken={accessToken} isPublicPage={false} />
+        <div className="min-h-0 flex-1 overflow-auto">{children}</div>
+      </div>
+    </ThemeProvider>
+  );
+}

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) => <div data-testid="mcp-apps-panel" />),
+  };
+});
+
+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(<ConnectPage />);
+    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(<ConnectPage />);
+    expect(mockReplace).toHaveBeenCalledWith("/connect");
+  });
+
+  it("does not rewrite the URL when there is no OAuth return param", () => {
+    render(<ConnectPage />);
+    expect(mockReplace).not.toHaveBeenCalled();
+  });
+});

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<string[]>([]);
+  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 (
+    <div className="mx-auto w-full max-w-5xl px-8 py-8">
+      <MCPAppsPanel accessToken={accessToken ?? ""} selectedServers={selectedServers} onChange={setSelectedServers} />
+    </div>
+  );
+}
+
+export default function ConnectPage() {
+  return (
+    <Suspense>
+      <ConnectPageContent />
+    </Suspense>
+  );
+}