From 48592e3c49857032553896eacaad1488f78a8cbd Mon Sep 17 00:00:00 2001 From: Tanishq <74068132+tanishqbaweja@users.noreply.github.com> Date: Sun, 25 Oct 2026 20:59:31 +0530 Subject: [PATCH] fix(ui): honor custom auth headers during onboarding --- .../OnboardingForm.integration.test.tsx | 91 +++++++++++++++++++ .../app/onboarding/OnboardingForm.test.tsx | 1 + .../src/app/onboarding/OnboardingForm.tsx | 3 +- 3 files changed, 94 insertions(+), 1 deletion(-) create mode 100644 ui/litellm-dashboard/src/app/onboarding/OnboardingForm.integration.test.tsx diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.integration.test.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.integration.test.tsx new file mode 100644 index 00000000000..30d6d265ef3 --- /dev/null +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.integration.test.tsx @@ -0,0 +1,91 @@ +import { setGlobalLitellmHeaderName } from "@/components/networking"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { OnboardingForm } from "./OnboardingForm"; + +vi.mock("next/navigation", () => ({ + useSearchParams: () => new URLSearchParams("invitation_id=inv-123"), +})); + +afterEach(() => { + setGlobalLitellmHeaderName("Authorization"); + sessionStorage.clear(); + vi.unstubAllGlobals(); +}); + +describe.each(["signup", "reset_password"] as const)("OnboardingForm %s authentication", (variant) => { + it.each([ + { name: "custom header", authHeader: "x-litellm-api-key", previousHeader: "Authorization" }, + { name: "default header", authHeader: "Authorization", previousHeader: "x-stale-auth" }, + { name: "missing header claim", authHeader: undefined, previousHeader: "x-stale-auth" }, + ])("submits with the token's $name instead of prior session state", async ({ authHeader, previousHeader }) => { + const expectedHeader = authHeader ?? "Authorization"; + const onboardingClaims = { + user_email: "alice@example.test", + user_id: "user-123", + key: "onboarding-key", + auth_header_name: authHeader, + }; + const token = `e30.${btoa(JSON.stringify(onboardingClaims))}.signature`; + const claims: Request[] = []; + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }); + setGlobalLitellmHeaderName(previousHeader); + vi.stubGlobal("location", { + ...window.location, + href: "http://localhost/ui/onboarding/?invitation_id=inv-123", + pathname: "/ui/onboarding/", + }); + vi.stubGlobal( + "fetch", + vi.fn(async (input, init) => { + const url = new URL(input instanceof Request ? input.url : String(input), "http://localhost"); + switch (url.pathname) { + case "/ui/litellm-runtime-config.json": + return Response.json({}); + case "/litellm/.well-known/litellm-ui-config": + return Response.json({ server_root_path: "", proxy_base_url: null }); + case "/onboarding/get_token": + return Response.json({ token }); + case "/onboarding/claim_token": { + const claim = new Request(url, init); + claims.push(claim); + return claim.headers.get(expectedHeader) === "Bearer onboarding-key" + ? Response.json({ token: "new-session-token" }) + : Response.json( + { detail: { error: "Missing onboarding session for invitation link." } }, + { status: 401 }, + ); + } + default: + throw new Error(`Unexpected request ${url.pathname}`); + } + }), + ); + + try { + render( + + + , + ); + await waitFor(() => expect(screen.getByLabelText("Email Address")).toHaveValue("alice@example.test")); + fireEvent.change(screen.getByLabelText("Password"), { target: { value: "New-password-for-test" } }); + fireEvent.click(screen.getByRole("button", { name: variant === "signup" ? "Sign Up" : "Reset Password" })); + + await waitFor(() => expect(claims).toHaveLength(1)); + expect(claims[0].headers.get(expectedHeader)).toBe("Bearer onboarding-key"); + expect(claims[0].headers.has(previousHeader)).toBe(false); + expect(await claims[0].json()).toEqual({ + invitation_link: "inv-123", + user_id: "user-123", + password: "New-password-for-test", + }); + await waitFor(() => expect(sessionStorage.getItem("token")).toBe("new-session-token")); + } finally { + queryClient.clear(); + } + }); +}); diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx index 0891de52736..05b6a2965e2 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx @@ -26,6 +26,7 @@ vi.mock("@/app/(dashboard)/hooks/onboarding/useOnboarding", () => ({ vi.mock("@/components/networking", () => ({ getProxyBaseUrl: vi.fn(() => ""), getUiPageUrl: (url: string) => url, + setGlobalLitellmHeaderName: vi.fn(), })); vi.mock("./OnboardingLoadingView", () => ({ diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx index 304ee161a9a..2193020827d 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx @@ -3,7 +3,7 @@ import React from "react"; import { useSearchParams } from "next/navigation"; import { jwtDecode } from "jwt-decode"; import { useOnboardingCredentials, useClaimOnboardingToken } from "@/app/(dashboard)/hooks/onboarding/useOnboarding"; -import { getProxyBaseUrl, getUiPageUrl } from "@/components/networking"; +import { getProxyBaseUrl, getUiPageUrl, setGlobalLitellmHeaderName } from "@/components/networking"; import { clearTokenCookies, storeLoginToken } from "@/utils/cookieUtils"; import { OnboardingLoadingView } from "./OnboardingLoadingView"; import { OnboardingErrorView } from "./OnboardingErrorView"; @@ -35,6 +35,7 @@ export function OnboardingForm({ variant }: OnboardingFormProps) { if (!accessToken || !userId || !inviteId) return; setClaimError(null); + setGlobalLitellmHeaderName(decoded?.auth_header_name ?? "Authorization"); claimToken( { accessToken, inviteId, userId, password: formValues.password },