feat(ui): add key creation deep-links with SSO return URL support
Enables deep-linking directly to the key creation modal with prefilled form data via URL parameters, including support for preserving these deep-links through SSO authentication flows. Key Creation Deep-links: - Auto-open key creation modal via ?create=true parameter - Prefill form fields from URL parameters (team_id, key_alias, models, etc.) - Role-based access control for auto-open (requires write access) - Race condition protection for redirect handling Example: /ui?create=true&team_id=abc&key_alias=my-key&models=gpt-4,claude-3 SSO Return URL Preservation: - Cookie-based return URL storage (works across ports for SSO flows) - URL validation to prevent open redirect attacks - Support for both dev and production environments Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Dibyo Mukherjee
parent
719b7fd013
commit
518cd3ef60
@ -8,13 +8,14 @@ import useAuthorized from "./useAuthorized";
|
||||
// Unmock useAuthorized to test the actual implementation
|
||||
vi.unmock("@/app/(dashboard)/hooks/useAuthorized");
|
||||
|
||||
const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock, decodeTokenMock, checkTokenValidityMock } = vi.hoisted(() => ({
|
||||
const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock, decodeTokenMock, checkTokenValidityMock, buildLoginUrlWithReturnMock } = vi.hoisted(() => ({
|
||||
replaceMock: vi.fn(),
|
||||
clearTokenCookiesMock: vi.fn(),
|
||||
getProxyBaseUrlMock: vi.fn(() => "http://proxy.example"),
|
||||
getUiConfigMock: vi.fn(),
|
||||
decodeTokenMock: vi.fn(),
|
||||
checkTokenValidityMock: vi.fn(),
|
||||
buildLoginUrlWithReturnMock: vi.fn((baseUrl: string) => baseUrl),
|
||||
}));
|
||||
|
||||
vi.mock("next/navigation", () => ({
|
||||
@ -49,6 +50,14 @@ vi.mock("@/utils/jwtUtils", async (importOriginal) => {
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock("@/utils/returnUrlUtils", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import("@/utils/returnUrlUtils")>();
|
||||
return {
|
||||
...actual,
|
||||
buildLoginUrlWithReturn: buildLoginUrlWithReturnMock,
|
||||
storeReturnUrl: vi.fn(),
|
||||
};
|
||||
});
|
||||
const createQueryClient = () =>
|
||||
new QueryClient({
|
||||
defaultOptions: {
|
||||
@ -81,6 +90,7 @@ describe("useAuthorized", () => {
|
||||
getUiConfigMock.mockReset();
|
||||
decodeTokenMock.mockReset();
|
||||
checkTokenValidityMock.mockReset();
|
||||
buildLoginUrlWithReturnMock.mockClear();
|
||||
clearCookie();
|
||||
});
|
||||
|
||||
|
||||
@ -3,39 +3,12 @@
|
||||
import { getProxyBaseUrl } from "@/components/networking";
|
||||
import { clearTokenCookies, getCookie } from "@/utils/cookieUtils";
|
||||
import { checkTokenValidity, decodeToken } from "@/utils/jwtUtils";
|
||||
import { buildLoginUrlWithReturn, storeReturnUrl } from "@/utils/returnUrlUtils";
|
||||
import { useRouter } from "next/navigation";
|
||||
import { useEffect, useMemo } from "react";
|
||||
import { useCallback, useEffect, useMemo } from "react";
|
||||
import { formatUserRole } from "@/utils/roles";
|
||||
import { useUIConfig } from "./uiConfig/useUIConfig";
|
||||
|
||||
function formatUserRole(userRole: string) {
|
||||
if (!userRole) {
|
||||
return "Undefined Role";
|
||||
}
|
||||
switch (userRole.toLowerCase()) {
|
||||
case "app_owner":
|
||||
return "App Owner";
|
||||
case "demo_app_owner":
|
||||
return "App Owner";
|
||||
case "app_admin":
|
||||
return "Admin";
|
||||
case "proxy_admin":
|
||||
return "Admin";
|
||||
case "proxy_admin_viewer":
|
||||
return "Admin Viewer";
|
||||
case "org_admin":
|
||||
return "Org Admin";
|
||||
case "internal_user":
|
||||
return "Internal User";
|
||||
case "internal_user_viewer":
|
||||
case "internal_viewer": // TODO:remove if deprecated
|
||||
return "Internal Viewer";
|
||||
case "app_user":
|
||||
return "App User";
|
||||
default:
|
||||
return "Unknown Role";
|
||||
}
|
||||
}
|
||||
|
||||
const useAuthorized = () => {
|
||||
const router = useRouter();
|
||||
const { data: uiConfig, isLoading: isUIConfigLoading } = useUIConfig();
|
||||
@ -47,6 +20,14 @@ const useAuthorized = () => {
|
||||
const isLoading = isUIConfigLoading;
|
||||
const isAuthorized = isTokenValid && !uiConfig?.admin_ui_disabled;
|
||||
|
||||
// Helper function to redirect to login while preserving the current URL
|
||||
const redirectToLogin = useCallback(() => {
|
||||
storeReturnUrl();
|
||||
const baseLoginUrl = `${getProxyBaseUrl()}/ui/login`;
|
||||
const loginUrlWithReturn = buildLoginUrlWithReturn(baseLoginUrl);
|
||||
router.replace(loginUrlWithReturn);
|
||||
}, [router]);
|
||||
|
||||
// Single useEffect for all redirect logic
|
||||
useEffect(() => {
|
||||
if (isLoading) return;
|
||||
@ -55,9 +36,9 @@ const useAuthorized = () => {
|
||||
if (token) {
|
||||
clearTokenCookies();
|
||||
}
|
||||
router.replace(`${getProxyBaseUrl()}/ui/login`);
|
||||
redirectToLogin();
|
||||
}
|
||||
}, [isLoading, isAuthorized, token, router]);
|
||||
}, [isLoading, isAuthorized, token, redirectToLogin]);
|
||||
|
||||
return {
|
||||
isLoading,
|
||||
|
||||
@ -6,6 +6,7 @@ import LoadingScreen from "@/components/common_components/LoadingScreen";
|
||||
import { getProxyBaseUrl } from "@/components/networking";
|
||||
import { getCookie } from "@/utils/cookieUtils";
|
||||
import { isJwtExpired } from "@/utils/jwtUtils";
|
||||
import { consumeReturnUrl, getReturnUrl, isValidReturnUrl } from "@/utils/returnUrlUtils";
|
||||
import { InfoCircleOutlined } from "@ant-design/icons";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { Alert, Button, Card, Form, Input, Popover, Space, Typography } from "antd";
|
||||
@ -33,12 +34,24 @@ function LoginPageContent() {
|
||||
|
||||
const rawToken = getCookie("token");
|
||||
if (rawToken && !isJwtExpired(rawToken)) {
|
||||
router.replace(`${getProxyBaseUrl()}/ui`);
|
||||
// User already logged in - redirect to return URL or default
|
||||
const returnUrl = consumeReturnUrl();
|
||||
if (returnUrl) {
|
||||
router.replace(returnUrl);
|
||||
} else {
|
||||
router.replace(`${getProxyBaseUrl()}/ui`);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (uiConfig && uiConfig.auto_redirect_to_sso) {
|
||||
router.push(`${getProxyBaseUrl()}/sso/key/generate`);
|
||||
// For SSO, pass the return URL to the SSO endpoint
|
||||
const returnUrl = getReturnUrl();
|
||||
let ssoUrl = `${getProxyBaseUrl()}/sso/key/generate`;
|
||||
if (returnUrl && isValidReturnUrl(returnUrl)) {
|
||||
ssoUrl += `?redirect_to=${encodeURIComponent(returnUrl)}`;
|
||||
}
|
||||
router.push(ssoUrl);
|
||||
return;
|
||||
}
|
||||
|
||||
@ -50,7 +63,13 @@ function LoginPageContent() {
|
||||
{ username, password },
|
||||
{
|
||||
onSuccess: (data) => {
|
||||
router.push(data.redirect_url);
|
||||
// Check if we have a return URL to use instead of the default redirect
|
||||
const returnUrl = consumeReturnUrl();
|
||||
if (returnUrl) {
|
||||
router.push(returnUrl);
|
||||
} else {
|
||||
router.push(data.redirect_url);
|
||||
}
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
@ -23,7 +23,7 @@ import Navbar from "@/components/navbar";
|
||||
import { getUiConfig, Organization, proxyBaseUrl, setGlobalLitellmHeaderName, getInProductNudgesCall } from "@/components/networking";
|
||||
import NewUsagePage from "@/components/UsagePage/components/UsagePageView";
|
||||
import OldTeams from "@/components/OldTeams";
|
||||
import { fetchUserModels } from "@/components/organisms/create_key_button";
|
||||
import { fetchUserModels, CreateKeyPrefillData } from "@/components/organisms/create_key_button";
|
||||
import Organizations, { fetchOrganizations } from "@/components/organizations";
|
||||
import PassThroughSettings from "@/components/pass_through_settings";
|
||||
import PromptsPanel from "@/components/prompts";
|
||||
@ -43,11 +43,12 @@ import SpendLogsTable from "@/components/view_logs";
|
||||
import ViewUserDashboard from "@/components/view_users";
|
||||
import { ThemeProvider } from "@/contexts/ThemeContext";
|
||||
import { isJwtExpired } from "@/utils/jwtUtils";
|
||||
import { isAdminRole } from "@/utils/roles";
|
||||
import { buildLoginUrlWithReturn, consumeReturnUrl, normalizeUrlForCompare, storeReturnUrl } from "@/utils/returnUrlUtils";
|
||||
import { formatUserRole, isAdminRole } from "@/utils/roles";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { jwtDecode } from "jwt-decode";
|
||||
import { useSearchParams } from "next/navigation";
|
||||
import { Suspense, useEffect, useState } from "react";
|
||||
import { Suspense, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { ConfigProvider, theme } from "antd";
|
||||
|
||||
function getCookie(name: string) {
|
||||
@ -67,35 +68,6 @@ function deleteCookie(name: string, path = "/") {
|
||||
document.cookie = `${name}=; Max-Age=0; Path=${path}`;
|
||||
}
|
||||
|
||||
function formatUserRole(userRole: string) {
|
||||
if (!userRole) {
|
||||
return "Undefined Role";
|
||||
}
|
||||
switch (userRole.toLowerCase()) {
|
||||
case "app_owner":
|
||||
return "App Owner";
|
||||
case "demo_app_owner":
|
||||
return "App Owner";
|
||||
case "app_admin":
|
||||
return "Admin";
|
||||
case "proxy_admin":
|
||||
return "Admin";
|
||||
case "proxy_admin_viewer":
|
||||
return "Admin Viewer";
|
||||
case "org_admin":
|
||||
return "Org Admin";
|
||||
case "internal_user":
|
||||
return "Internal User";
|
||||
case "internal_user_viewer":
|
||||
case "internal_viewer": // TODO:remove if deprecated
|
||||
return "Internal Viewer";
|
||||
case "app_user":
|
||||
return "App User";
|
||||
default:
|
||||
return "Unknown Role";
|
||||
}
|
||||
}
|
||||
|
||||
interface ProxySettings {
|
||||
PROXY_BASE_URL: string;
|
||||
PROXY_LOGOUT_URL: string;
|
||||
@ -143,6 +115,58 @@ function CreateKeyPageContent() {
|
||||
|
||||
const invitation_id = searchParams.get("invitation_id");
|
||||
|
||||
// Parse URL query parameters for pre-filling the create key form
|
||||
// Includes validation to prevent injection and DoS attacks
|
||||
const autoOpenCreate = searchParams.get("create") === "true";
|
||||
const prefillData: CreateKeyPrefillData | undefined = useMemo(() => {
|
||||
if (!autoOpenCreate) return undefined;
|
||||
|
||||
const ownedBy = searchParams.get("owned_by");
|
||||
const teamId = searchParams.get("team_id");
|
||||
const keyAlias = searchParams.get("key_alias");
|
||||
const modelsParam = searchParams.get("models");
|
||||
const keyType = searchParams.get("key_type");
|
||||
|
||||
// Only return prefill data if at least one field is provided
|
||||
if (!ownedBy && !teamId && !keyAlias && !modelsParam && !keyType) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
// Validate owned_by against allowed values
|
||||
const validOwnedByValues = ["you", "service_account", "another_user"];
|
||||
const validatedOwnedBy = ownedBy && validOwnedByValues.includes(ownedBy)
|
||||
? (ownedBy as CreateKeyPrefillData["owned_by"])
|
||||
: undefined;
|
||||
|
||||
// Validate key_type against allowed values
|
||||
const validKeyTypes = ["default", "llm_api", "management"];
|
||||
const validatedKeyType = keyType && validKeyTypes.includes(keyType)
|
||||
? (keyType as CreateKeyPrefillData["key_type"])
|
||||
: undefined;
|
||||
|
||||
// Sanitize key_alias (limit length, trim whitespace)
|
||||
const sanitizedKeyAlias = keyAlias
|
||||
? keyAlias.trim().slice(0, 256) // Reasonable max length
|
||||
: undefined;
|
||||
|
||||
// Sanitize models (limit array size and individual model name length)
|
||||
const sanitizedModels = modelsParam
|
||||
? modelsParam
|
||||
.split(",")
|
||||
.slice(0, 100) // Limit number of models to prevent DoS
|
||||
.map(m => m.trim().slice(0, 256)) // Limit individual model name length
|
||||
.filter(m => m.length > 0) // Remove empty strings
|
||||
: undefined;
|
||||
|
||||
return {
|
||||
owned_by: validatedOwnedBy,
|
||||
team_id: teamId?.trim() || undefined,
|
||||
key_alias: sanitizedKeyAlias,
|
||||
models: sanitizedModels && sanitizedModels.length > 0 ? sanitizedModels : undefined,
|
||||
key_type: validatedKeyType,
|
||||
};
|
||||
}, [searchParams, autoOpenCreate]);
|
||||
|
||||
// Get page from URL, default to 'api-keys' if not present
|
||||
const [page, setPage] = useState(() => {
|
||||
return searchParams.get("page") || "api-keys";
|
||||
@ -163,6 +187,9 @@ function CreateKeyPageContent() {
|
||||
const [accessToken, setAccessToken] = useState<string | null>(null);
|
||||
const [sidebarCollapsed, setSidebarCollapsed] = useState(false);
|
||||
|
||||
// Track if we've already attempted a return URL redirect to prevent race conditions
|
||||
const hasAttemptedReturnRedirectRef = useRef(false);
|
||||
|
||||
const toggleSidebar = () => {
|
||||
setSidebarCollapsed(!sidebarCollapsed);
|
||||
};
|
||||
@ -207,12 +234,48 @@ function CreateKeyPageContent() {
|
||||
|
||||
useEffect(() => {
|
||||
if (redirectToLogin) {
|
||||
// Store the current URL so we can redirect back after login
|
||||
storeReturnUrl();
|
||||
// Build login URL with return URL parameter
|
||||
const baseLoginUrl = (proxyBaseUrl || "") + "/ui/login";
|
||||
const dest = buildLoginUrlWithReturn(baseLoginUrl);
|
||||
// Replace instead of assigning to avoid back-button loops
|
||||
const dest = (proxyBaseUrl || "") + "/ui/login";
|
||||
window.location.replace(dest);
|
||||
}
|
||||
}, [redirectToLogin]);
|
||||
|
||||
// Check for a stored return URL after successful authentication
|
||||
// This handles the case where user comes back from SSO and we need to redirect to the original URL
|
||||
useEffect(() => {
|
||||
// Skip if still loading, no token, or we've already attempted a redirect
|
||||
if (authLoading || !token || hasAttemptedReturnRedirectRef.current) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Mark that we've attempted the redirect to prevent race conditions
|
||||
// This prevents duplicate redirects if token changes (e.g., refresh)
|
||||
hasAttemptedReturnRedirectRef.current = true;
|
||||
|
||||
// Check for a stored return URL
|
||||
const returnUrl = consumeReturnUrl();
|
||||
if (returnUrl) {
|
||||
const currentUrl = window.location.href;
|
||||
const normalizedReturnUrl = normalizeUrlForCompare(returnUrl);
|
||||
const normalizedCurrentUrl = normalizeUrlForCompare(currentUrl);
|
||||
// Only redirect if the return URL is different from the current URL
|
||||
// This prevents infinite redirect loops
|
||||
if (normalizedReturnUrl !== normalizedCurrentUrl) {
|
||||
window.location.replace(returnUrl);
|
||||
}
|
||||
}
|
||||
}, [authLoading, token]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!token) {
|
||||
hasAttemptedReturnRedirectRef.current = false;
|
||||
}
|
||||
}, [token]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!token) {
|
||||
return;
|
||||
@ -410,9 +473,8 @@ function CreateKeyPageContent() {
|
||||
/>
|
||||
<div className="flex flex-1">
|
||||
<div className="mt-2">
|
||||
<SidebarProvider setPage={updatePage} defaultSelectedKey={page} sidebarCollapsed={sidebarCollapsed} />
|
||||
</div>
|
||||
|
||||
<SidebarProvider setPage={updatePage} defaultSelectedKey={page} sidebarCollapsed={sidebarCollapsed} />
|
||||
</div>
|
||||
{page == "api-keys" ? (
|
||||
<UserDashboard
|
||||
userID={userID}
|
||||
@ -428,6 +490,8 @@ function CreateKeyPageContent() {
|
||||
organizations={organizations}
|
||||
addKey={addKey}
|
||||
createClicked={createClicked}
|
||||
autoOpenCreate={autoOpenCreate}
|
||||
prefillData={prefillData}
|
||||
/>
|
||||
) : page == "models" ? (
|
||||
<OldModelDashboard
|
||||
|
||||
@ -1,20 +1,183 @@
|
||||
import { act, fireEvent, waitFor } from "@testing-library/react";
|
||||
import { act, fireEvent } from "@testing-library/react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { renderWithProviders, screen } from "../../../tests/test-utils";
|
||||
import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils";
|
||||
import CreateKey from "./create_key_button";
|
||||
|
||||
const { mockKeyCreateCall } = vi.hoisted(() => {
|
||||
const fn = vi.fn().mockResolvedValue({
|
||||
const { formMock, setFieldsValueMock, radioGroupValueRef, formStateRef, mockKeyCreateCall } = vi.hoisted(() => {
|
||||
const formStateRef = { current: {} as Record<string, any> };
|
||||
const mockKeyCreateCall = vi.fn().mockResolvedValue({
|
||||
key: "test-api-key",
|
||||
soft_budget: null,
|
||||
});
|
||||
return { mockKeyCreateCall: fn };
|
||||
const formMock = {
|
||||
setFieldsValue: vi.fn((values: Record<string, any>) => {
|
||||
Object.assign(formStateRef.current, values);
|
||||
}),
|
||||
setFieldValue: vi.fn((name: string, value: any) => {
|
||||
formStateRef.current[name] = value;
|
||||
}),
|
||||
getFieldValue: vi.fn((name: string) => formStateRef.current[name]),
|
||||
resetFields: vi.fn(() => {
|
||||
formStateRef.current = {};
|
||||
}),
|
||||
};
|
||||
const radioGroupValueRef = { current: null as string | null };
|
||||
return {
|
||||
formMock,
|
||||
setFieldsValueMock: formMock.setFieldsValue,
|
||||
radioGroupValueRef,
|
||||
formStateRef,
|
||||
mockKeyCreateCall,
|
||||
};
|
||||
});
|
||||
|
||||
const defaultAuthorizedState = {
|
||||
accessToken: "test-token",
|
||||
userId: "test-user-id",
|
||||
userRole: "Admin",
|
||||
premiumUser: false,
|
||||
};
|
||||
|
||||
let authorizedState = { ...defaultAuthorizedState };
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
|
||||
default: () => authorizedState,
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({
|
||||
keyKeys: {
|
||||
lists: () => ["keys"],
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock("@ant-design/icons", () => ({
|
||||
InfoCircleOutlined: () => null,
|
||||
}));
|
||||
|
||||
vi.mock("react-copy-to-clipboard", () => ({
|
||||
CopyToClipboard: ({ children }: { children: any }) => children,
|
||||
}));
|
||||
|
||||
vi.mock("@tremor/react", () => {
|
||||
const React = require("react");
|
||||
const Stub = ({ children }: { children?: any }) => React.createElement("div", null, children);
|
||||
const Button = ({ children, ...props }: { children?: any }) =>
|
||||
React.createElement("button", props, children);
|
||||
const TextInput = (props: any) => React.createElement("input", props);
|
||||
|
||||
return {
|
||||
Accordion: Stub,
|
||||
AccordionBody: Stub,
|
||||
AccordionHeader: Stub,
|
||||
Button,
|
||||
Col: Stub,
|
||||
Grid: Stub,
|
||||
Text: Stub,
|
||||
TextInput,
|
||||
Title: Stub,
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock("antd", () => {
|
||||
const React = require("react");
|
||||
|
||||
const getValueFromEvent = (event: any) => {
|
||||
if (event?.target) {
|
||||
if (event.target.type === "checkbox") {
|
||||
return event.target.checked;
|
||||
}
|
||||
return event.target.value;
|
||||
}
|
||||
return event;
|
||||
};
|
||||
|
||||
const Form = ({ children, onFinish, ...props }: { children?: any; onFinish?: (values: Record<string, any>) => void }) =>
|
||||
React.createElement(
|
||||
"form",
|
||||
{
|
||||
...props,
|
||||
onSubmit: (event: Event) => {
|
||||
event.preventDefault();
|
||||
onFinish?.({ ...formStateRef.current });
|
||||
},
|
||||
},
|
||||
children,
|
||||
);
|
||||
|
||||
Form.Item = ({ children, name }: { children?: any; name?: string }) => {
|
||||
if (!name || !React.isValidElement(children)) {
|
||||
return React.createElement(React.Fragment, null, children);
|
||||
}
|
||||
|
||||
return React.cloneElement(children, {
|
||||
value: formStateRef.current[name],
|
||||
onChange: (event: any) => {
|
||||
formStateRef.current[name] = getValueFromEvent(event);
|
||||
},
|
||||
});
|
||||
};
|
||||
|
||||
Form.useForm = () => [formMock];
|
||||
|
||||
const Select = ({ children, onChange, ...props }: { children?: any; onChange?: (value: string) => void }) =>
|
||||
React.createElement(
|
||||
"select",
|
||||
{
|
||||
...props,
|
||||
onChange: (event: any) => onChange?.(event.target.value),
|
||||
},
|
||||
children,
|
||||
);
|
||||
|
||||
Select.Option = ({ children, ...props }: { children?: any }) =>
|
||||
React.createElement("option", props, children);
|
||||
|
||||
const Input = (props: any) => React.createElement("input", props);
|
||||
Input.Password = (props: any) => React.createElement("input", { ...props, type: "password" });
|
||||
Input.TextArea = (props: any) => React.createElement("textarea", props);
|
||||
|
||||
const Modal = ({ children, open }: { children?: any; open?: boolean }) =>
|
||||
open ? React.createElement("div", null, children) : null;
|
||||
|
||||
const Radio = ({ children, ...props }: { children?: any }) =>
|
||||
React.createElement("div", props, children);
|
||||
|
||||
Radio.Group = ({ children, value }: { children?: any; value?: string }) => {
|
||||
radioGroupValueRef.current = value ?? null;
|
||||
return React.createElement("div", null, children);
|
||||
};
|
||||
|
||||
const Switch = (props: any) => React.createElement("input", { ...props, type: "checkbox" });
|
||||
const Tag = ({ children }: { children?: any }) => React.createElement("span", null, children);
|
||||
const Tooltip = ({ children }: { children?: any }) => React.createElement(React.Fragment, null, children);
|
||||
|
||||
const Button = ({ children, htmlType, ...props }: { children?: any; htmlType?: string }) =>
|
||||
React.createElement("button", { ...props, type: htmlType ?? props.type }, children);
|
||||
|
||||
return {
|
||||
Button,
|
||||
Form,
|
||||
Input,
|
||||
message: {
|
||||
success: vi.fn(),
|
||||
error: vi.fn(),
|
||||
warning: vi.fn(),
|
||||
info: vi.fn(),
|
||||
},
|
||||
Modal,
|
||||
Radio,
|
||||
Select,
|
||||
Switch,
|
||||
Tag,
|
||||
Tooltip,
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock("../networking", () => ({
|
||||
keyCreateCall: mockKeyCreateCall,
|
||||
modelAvailableCall: vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }, { id: "gpt-3.5-turbo" }] }),
|
||||
modelAvailableCall: vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }] }),
|
||||
getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }),
|
||||
getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }),
|
||||
getPromptsList: vi.fn().mockResolvedValue({ prompts: [] }),
|
||||
proxyBaseUrl: "http://localhost:4000",
|
||||
getPossibleUserRoles: vi.fn().mockResolvedValue({
|
||||
@ -41,12 +204,31 @@ vi.mock("../molecules/notifications_manager", () => ({
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock("../agent_management/AgentSelector", () => ({ default: () => null }));
|
||||
vi.mock("../common_components/budget_duration_dropdown", () => ({ default: () => null }));
|
||||
vi.mock("../common_components/check_openapi_schema", () => ({ default: () => null }));
|
||||
vi.mock("../common_components/KeyLifecycleSettings", () => ({ default: () => null }));
|
||||
vi.mock("../common_components/ModelAliasManager", () => ({ default: () => null }));
|
||||
vi.mock("../common_components/PassThroughRoutesSelector", () => ({ default: () => null }));
|
||||
vi.mock("../common_components/PremiumLoggingSettings", () => ({ default: () => null }));
|
||||
vi.mock("../common_components/RateLimitTypeFormItem", () => ({ default: () => null }));
|
||||
vi.mock("../common_components/RouterSettingsAccordion", () => ({ default: () => null }));
|
||||
vi.mock("../common_components/team_dropdown", () => ({ default: () => null }));
|
||||
vi.mock("../CreateUserButton", () => ({ CreateUserButton: () => null }));
|
||||
vi.mock("../mcp_server_management/MCPServerSelector", () => ({ default: () => null }));
|
||||
vi.mock("../mcp_server_management/MCPToolPermissions", () => ({ default: () => null }));
|
||||
vi.mock("../shared/numerical_input", () => ({ default: () => null }));
|
||||
vi.mock("../vector_store_management/VectorStoreSelector", () => ({ default: () => null }));
|
||||
vi.mock("../key_team_helpers/fetch_available_models_team_key", () => ({
|
||||
getModelDisplayName: (model: string) => model,
|
||||
}));
|
||||
|
||||
vi.mock("../common_components/AccessGroupSelector", () => ({
|
||||
default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => (
|
||||
<input
|
||||
data-testid="access-group-selector"
|
||||
value={Array.isArray(value) ? value.join(",") : ""}
|
||||
onChange={(e) => onChange?.(e.target.value ? e.target.value.split(",").map((s) => s.trim()) : [])}
|
||||
onChange={(event) => onChange?.(event.target.value ? event.target.value.split(",").map((v) => v.trim()) : [])}
|
||||
/>
|
||||
),
|
||||
}));
|
||||
@ -54,14 +236,19 @@ vi.mock("../common_components/AccessGroupSelector", () => ({
|
||||
describe("CreateKey", () => {
|
||||
const defaultProps = {
|
||||
team: null,
|
||||
data: [],
|
||||
teams: [],
|
||||
data: [],
|
||||
addKey: vi.fn(),
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
localStorage.clear();
|
||||
if (typeof window !== "undefined" && window.localStorage && typeof window.localStorage.clear === "function") {
|
||||
window.localStorage.clear();
|
||||
}
|
||||
authorizedState = { ...defaultAuthorizedState };
|
||||
radioGroupValueRef.current = null;
|
||||
formStateRef.current = {};
|
||||
mockKeyCreateCall.mockResolvedValue({
|
||||
key: "test-api-key",
|
||||
soft_budget: null,
|
||||
@ -81,26 +268,8 @@ describe("CreateKey", () => {
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Key Type")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// Open the Key Type dropdown
|
||||
const keyTypeSection = screen.getByText("Key Type").closest(".ant-form-item")!;
|
||||
const selectElement = keyTypeSection.querySelector(".ant-select-selector")!;
|
||||
act(() => {
|
||||
fireEvent.mouseDown(selectElement);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
// Verify "AI APIs" appears as an option
|
||||
const options = document.querySelectorAll(".ant-select-item-option");
|
||||
const optionTexts = Array.from(options).map((el) => el.textContent);
|
||||
const hasAIAPIs = optionTexts.some((text) => text?.includes("AI APIs"));
|
||||
expect(hasAIAPIs).toBe(true);
|
||||
|
||||
// Verify old "LLM API" label does NOT appear
|
||||
const hasLLMAPI = optionTexts.some((text) => text?.includes("LLM API"));
|
||||
expect(hasLLMAPI).toBe(false);
|
||||
expect(screen.getByText("AI APIs")).toBeInTheDocument();
|
||||
expect(screen.queryByText("LLM API")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
@ -111,46 +280,118 @@ describe("CreateKey", () => {
|
||||
fireEvent.click(screen.getByRole("button", { name: /create new key/i }));
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByLabelText(/key name/i)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
fireEvent.change(screen.getByLabelText(/key name/i), { target: { value: "Test Key" } });
|
||||
|
||||
const optionalSettingsAccordion = screen.getByText("Optional Settings");
|
||||
act(() => {
|
||||
fireEvent.click(optionalSettingsAccordion);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByTestId("access-group-selector")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
fireEvent.change(screen.getByTestId("access-group-selector"), { target: { value: "ag-1,ag-2" } });
|
||||
act(() => {
|
||||
fireEvent.change(screen.getByTestId("access-group-selector"), { target: { value: "ag-1,ag-2" } });
|
||||
formMock.setFieldValue("key_alias", "Test Key");
|
||||
});
|
||||
|
||||
const modelsCombobox = screen.getAllByRole("combobox").find((el) => el.closest('[class*="ant-form-item"]')?.textContent?.includes("Models")) ||
|
||||
screen.getAllByRole("combobox")[1];
|
||||
if (modelsCombobox) {
|
||||
act(() => fireEvent.mouseDown(modelsCombobox));
|
||||
await waitFor(() => {
|
||||
const allTeamModels = [...document.body.querySelectorAll(".ant-select-item")].find(
|
||||
(el) => el.textContent?.includes("All Team Models"),
|
||||
);
|
||||
if (allTeamModels) fireEvent.click(allTeamModels);
|
||||
});
|
||||
}
|
||||
act(() => {
|
||||
fireEvent.click(screen.getByRole("button", { name: /create key/i }));
|
||||
});
|
||||
|
||||
const createButton = screen.getByRole("button", { name: /create key/i });
|
||||
act(() => fireEvent.click(createButton));
|
||||
await waitFor(() => {
|
||||
expect(mockKeyCreateCall).toHaveBeenCalled();
|
||||
const formValues = mockKeyCreateCall.mock.calls[0][2];
|
||||
expect(formValues).toHaveProperty("access_group_ids");
|
||||
expect(formValues.access_group_ids).toEqual(["ag-1", "ag-2"]);
|
||||
});
|
||||
});
|
||||
|
||||
await waitFor(
|
||||
() => {
|
||||
expect(mockKeyCreateCall).toHaveBeenCalled();
|
||||
const formValues = mockKeyCreateCall.mock.calls[0][2];
|
||||
expect(formValues).toHaveProperty("access_group_ids");
|
||||
expect(formValues.access_group_ids).toEqual(["ag-1", "ag-2"]);
|
||||
},
|
||||
{ timeout: 15000 },
|
||||
it("should prefill models when provided without team_id", async () => {
|
||||
renderWithProviders(
|
||||
<CreateKey
|
||||
{...defaultProps}
|
||||
autoOpenCreate={true}
|
||||
prefillData={{
|
||||
models: ["gpt-4"],
|
||||
}}
|
||||
/>,
|
||||
);
|
||||
}, { timeout: 30000 });
|
||||
|
||||
await waitFor(() => {
|
||||
expect(setFieldsValueMock).toHaveBeenCalledWith({ models: ["gpt-4"] });
|
||||
});
|
||||
});
|
||||
|
||||
it("should prefill team_id when it exists in teams", async () => {
|
||||
renderWithProviders(
|
||||
<CreateKey
|
||||
{...defaultProps}
|
||||
teams={[{ team_id: "team-1", models: [] } as any]}
|
||||
autoOpenCreate={true}
|
||||
prefillData={{ team_id: "team-1" }}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(setFieldsValueMock).toHaveBeenCalledWith({ team_id: "team-1" });
|
||||
});
|
||||
});
|
||||
|
||||
it("should ignore team_id when it does not exist in teams", async () => {
|
||||
renderWithProviders(
|
||||
<CreateKey
|
||||
{...defaultProps}
|
||||
teams={[{ team_id: "team-1", models: [] } as any]}
|
||||
autoOpenCreate={true}
|
||||
prefillData={{ team_id: "team-404", key_alias: "example-key" }}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(setFieldsValueMock).toHaveBeenCalledWith({ key_alias: "example-key" });
|
||||
});
|
||||
|
||||
expect(setFieldsValueMock).not.toHaveBeenCalledWith({ team_id: "team-404" });
|
||||
});
|
||||
|
||||
it('should fall back to "you" when owned_by is another_user for non-admin', async () => {
|
||||
authorizedState = { ...defaultAuthorizedState, userRole: "Internal User" };
|
||||
|
||||
renderWithProviders(
|
||||
<CreateKey
|
||||
{...defaultProps}
|
||||
autoOpenCreate={true}
|
||||
prefillData={{ owned_by: "another_user", key_alias: "example-key" }}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(setFieldsValueMock).toHaveBeenCalledWith({ key_alias: "example-key" });
|
||||
});
|
||||
|
||||
expect(radioGroupValueRef.current).toBe("you");
|
||||
});
|
||||
|
||||
it("should apply owned_by another_user for admin", async () => {
|
||||
renderWithProviders(
|
||||
<CreateKey
|
||||
{...defaultProps}
|
||||
autoOpenCreate={true}
|
||||
prefillData={{ owned_by: "another_user" }}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(radioGroupValueRef.current).toBe("another_user");
|
||||
});
|
||||
});
|
||||
|
||||
it("should prefill key_type when provided", async () => {
|
||||
renderWithProviders(
|
||||
<CreateKey
|
||||
{...defaultProps}
|
||||
autoOpenCreate={true}
|
||||
prefillData={{ key_type: "management" }}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(setFieldsValueMock).toHaveBeenCalledWith({ key_type: "management" });
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@ -46,11 +46,24 @@ import { simplifyKeyGenerateError } from "./utils";
|
||||
|
||||
const { Option } = Select;
|
||||
|
||||
/**
|
||||
* Interface for pre-filling the create key form from URL parameters
|
||||
*/
|
||||
export interface CreateKeyPrefillData {
|
||||
owned_by?: "you" | "service_account" | "another_user";
|
||||
team_id?: string;
|
||||
key_alias?: string;
|
||||
models?: string[];
|
||||
key_type?: "default" | "llm_api" | "management";
|
||||
}
|
||||
|
||||
interface CreateKeyProps {
|
||||
team: Team | null;
|
||||
data: any[] | null;
|
||||
teams: Team[] | null;
|
||||
addKey: (data: any) => void;
|
||||
autoOpenCreate?: boolean;
|
||||
prefillData?: CreateKeyPrefillData;
|
||||
}
|
||||
|
||||
interface User {
|
||||
@ -141,7 +154,7 @@ export const fetchUserModels = async (
|
||||
* Please contribute to the new refactor.
|
||||
* ─────────────────────────────────────────────────────────────────────────
|
||||
*/
|
||||
const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
|
||||
const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOpenCreate, prefillData }) => {
|
||||
const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized();
|
||||
const queryClient = useQueryClient();
|
||||
const [form] = Form.useForm();
|
||||
@ -152,6 +165,8 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
|
||||
const [modelsToPick, setModelsToPick] = useState<string[]>([]);
|
||||
const [keyOwner, setKeyOwner] = useState("you");
|
||||
const [predefinedTags, setPredefinedTags] = useState(getPredefinedTags(data));
|
||||
const [hasPrefilled, setHasPrefilled] = useState(false);
|
||||
const [pendingPrefillModels, setPendingPrefillModels] = useState<string[] | null>(null);
|
||||
const [guardrailsList, setGuardrailsList] = useState<string[]>([]);
|
||||
const [policiesList, setPoliciesList] = useState<string[]>([]);
|
||||
const [promptsList, setPromptsList] = useState<string[]>([]);
|
||||
@ -274,6 +289,55 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
|
||||
fetchPossibleRoles();
|
||||
}, [accessToken]);
|
||||
|
||||
// Auto-open modal and prefill form from URL params (deep link).
|
||||
// Guarded by write access so we don't open for read-only users.
|
||||
useEffect(() => {
|
||||
if (autoOpenCreate && !hasPrefilled && teams && userRole && rolesWithWriteAccess.includes(userRole)) {
|
||||
// Open the modal
|
||||
setIsModalVisible(true);
|
||||
setHasPrefilled(true);
|
||||
|
||||
// Apply prefill data if provided
|
||||
if (prefillData) {
|
||||
// Set key owner (owned_by) - validate that "another_user" is only allowed for Admin
|
||||
if (prefillData.owned_by) {
|
||||
if (prefillData.owned_by === "another_user" && userRole !== "Admin") {
|
||||
// Ignore invalid owned_by for non-admin users, fall back to default
|
||||
setKeyOwner("you");
|
||||
} else {
|
||||
setKeyOwner(prefillData.owned_by);
|
||||
}
|
||||
}
|
||||
|
||||
// Set team - find the team by ID and set it (only if team exists in user's teams)
|
||||
if (prefillData.team_id) {
|
||||
const selectedTeam = teams?.find((t) => t.team_id === prefillData.team_id) || null;
|
||||
if (selectedTeam) {
|
||||
setSelectedCreateKeyTeam(selectedTeam);
|
||||
form.setFieldsValue({ team_id: prefillData.team_id });
|
||||
}
|
||||
// Silently ignore invalid team_id - don't prefill with a team user doesn't have access to
|
||||
}
|
||||
|
||||
// Set key alias
|
||||
if (prefillData.key_alias) {
|
||||
form.setFieldsValue({ key_alias: prefillData.key_alias });
|
||||
}
|
||||
|
||||
// Defer model selection until we load the allowed model list.
|
||||
if (prefillData.models && prefillData.models.length > 0) {
|
||||
setPendingPrefillModels(prefillData.models);
|
||||
}
|
||||
|
||||
// Set key type
|
||||
if (prefillData.key_type) {
|
||||
setKeyType(prefillData.key_type);
|
||||
form.setFieldsValue({ key_type: prefillData.key_type });
|
||||
}
|
||||
}
|
||||
}
|
||||
}, [autoOpenCreate, prefillData, teams, hasPrefilled, form, userRole]);
|
||||
|
||||
// Check if team selection is required
|
||||
const isTeamSelectionRequired = modelsToPick.includes("no-default-models");
|
||||
const isFormDisabled = isTeamSelectionRequired && !selectedCreateKeyTeam;
|
||||
@ -467,6 +531,9 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
|
||||
NotificationsManager.success("Virtual Key copied to clipboard");
|
||||
};
|
||||
|
||||
// Fetch available models when team or auth changes.
|
||||
// Note: Model prefill from URL params is handled by the useEffect below, which
|
||||
// watches for pendingPrefillModels + modelsToPick to both be populated.
|
||||
useEffect(() => {
|
||||
if (userID && userRole && accessToken) {
|
||||
fetchTeamModels(userID, userRole, accessToken, selectedCreateKeyTeam?.team_id ?? null).then((models) => {
|
||||
@ -474,8 +541,28 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
|
||||
setModelsToPick(allModels);
|
||||
});
|
||||
}
|
||||
form.setFieldValue("models", []);
|
||||
}, [selectedCreateKeyTeam, accessToken, userID, userRole]);
|
||||
// Only clear models if we don't have pending prefill models
|
||||
if (!pendingPrefillModels) {
|
||||
form.setFieldValue("models", []);
|
||||
}
|
||||
}, [selectedCreateKeyTeam, accessToken, userID, userRole, form]);
|
||||
|
||||
// Apply deferred model prefill once the available model list arrives.
|
||||
// This handles timing where prefill data arrives before or after models are fetched.
|
||||
useEffect(() => {
|
||||
if (!pendingPrefillModels || pendingPrefillModels.length === 0) {
|
||||
return;
|
||||
}
|
||||
if (!modelsToPick || modelsToPick.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const validModels = pendingPrefillModels.filter((model) => modelsToPick.includes(model));
|
||||
if (validModels.length > 0) {
|
||||
form.setFieldsValue({ models: validModels });
|
||||
}
|
||||
setPendingPrefillModels(null);
|
||||
}, [pendingPrefillModels, modelsToPick, form]);
|
||||
|
||||
// Add a callback function to handle user creation
|
||||
const handleUserCreated = (userId: string) => {
|
||||
|
||||
@ -16,7 +16,7 @@ import {
|
||||
Organization,
|
||||
userInfoCall,
|
||||
} from "./networking";
|
||||
import CreateKey from "./organisms/create_key_button";
|
||||
import CreateKey, { CreateKeyPrefillData } from "./organisms/create_key_button";
|
||||
import { VirtualKeysTable } from "./VirtualKeysPage/VirtualKeysTable";
|
||||
|
||||
export interface ProxySettings {
|
||||
@ -55,6 +55,8 @@ interface UserDashboardProps {
|
||||
organizations: Organization[] | null;
|
||||
addKey: (data: any) => void;
|
||||
createClicked: boolean;
|
||||
autoOpenCreate?: boolean;
|
||||
prefillData?: CreateKeyPrefillData;
|
||||
}
|
||||
|
||||
type TeamInterface = {
|
||||
@ -77,6 +79,8 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
|
||||
organizations,
|
||||
addKey,
|
||||
createClicked,
|
||||
autoOpenCreate,
|
||||
prefillData,
|
||||
}) => {
|
||||
const [userSpendData, setUserSpendData] = useState<UserInfo | null>(null);
|
||||
const [currentOrg, setCurrentOrg] = useState<Organization | null>(null);
|
||||
@ -350,6 +354,8 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
|
||||
teams={teams as Team[]}
|
||||
data={keys}
|
||||
addKey={addKey}
|
||||
autoOpenCreate={autoOpenCreate}
|
||||
prefillData={prefillData}
|
||||
/>
|
||||
<VirtualKeysTable teams={teams} organizations={organizations} />
|
||||
</Col>
|
||||
|
||||
383
ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts
Normal file
383
ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts
Normal file
@ -0,0 +1,383 @@
|
||||
import {
|
||||
buildLoginUrlWithReturn,
|
||||
clearStoredReturnUrl,
|
||||
consumeReturnUrl,
|
||||
getCurrentUrl,
|
||||
getReturnUrl,
|
||||
getReturnUrlFromParams,
|
||||
getStoredReturnUrl,
|
||||
isValidReturnUrl,
|
||||
storeReturnUrl,
|
||||
} from "./returnUrlUtils";
|
||||
|
||||
describe("returnUrlUtils", () => {
|
||||
const originalLocation = window.location;
|
||||
|
||||
beforeEach(() => {
|
||||
// Clear cookies before each test
|
||||
document.cookie.split(";").forEach((c) => {
|
||||
document.cookie = c
|
||||
.replace(/^ +/, "")
|
||||
.replace(/=.*/, "=;expires=" + new Date().toUTCString() + ";path=/");
|
||||
});
|
||||
|
||||
// Reset location mock
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
href: "http://localhost:3000/ui?page=api-keys",
|
||||
origin: "http://localhost:3000",
|
||||
hostname: "localhost",
|
||||
pathname: "/ui",
|
||||
search: "?page=api-keys",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
// Restore original location
|
||||
Object.defineProperty(window, "location", {
|
||||
value: originalLocation,
|
||||
writable: true,
|
||||
});
|
||||
});
|
||||
|
||||
describe("getCurrentUrl", () => {
|
||||
it("should return the current URL", () => {
|
||||
const url = getCurrentUrl();
|
||||
expect(url).toBe("http://localhost:3000/ui?page=api-keys");
|
||||
});
|
||||
});
|
||||
|
||||
describe("storeReturnUrl and getStoredReturnUrl", () => {
|
||||
it("should store and retrieve the return URL from cookie", () => {
|
||||
storeReturnUrl();
|
||||
const storedUrl = getStoredReturnUrl();
|
||||
expect(storedUrl).toBe("http://localhost:3000/ui?page=api-keys");
|
||||
});
|
||||
|
||||
it("should return null if no URL is stored", () => {
|
||||
const storedUrl = getStoredReturnUrl();
|
||||
expect(storedUrl).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("clearStoredReturnUrl", () => {
|
||||
it("should clear the stored return URL", () => {
|
||||
storeReturnUrl();
|
||||
expect(getStoredReturnUrl()).not.toBeNull();
|
||||
|
||||
clearStoredReturnUrl();
|
||||
expect(getStoredReturnUrl()).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("getReturnUrlFromParams", () => {
|
||||
it("should return the redirect_to parameter from URL", () => {
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
...window.location,
|
||||
search: "?redirect_to=http%3A%2F%2Flocalhost%3A3000%2Fui%3Fcreate%3Dtrue",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
const returnUrl = getReturnUrlFromParams();
|
||||
expect(returnUrl).toBe("http://localhost:3000/ui?create=true");
|
||||
});
|
||||
|
||||
it("should return null if redirect_to parameter is not present", () => {
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
...window.location,
|
||||
search: "?page=api-keys",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
const returnUrl = getReturnUrlFromParams();
|
||||
expect(returnUrl).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("buildLoginUrlWithReturn", () => {
|
||||
it("should build login URL with return URL parameter", () => {
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
...window.location,
|
||||
href: "http://localhost:3000/ui?create=true&team_id=123",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
const loginUrl = buildLoginUrlWithReturn("/ui/login");
|
||||
expect(loginUrl).toBe(
|
||||
"/ui/login?redirect_to=http%3A%2F%2Flocalhost%3A3000%2Fui%3Fcreate%3Dtrue%26team_id%3D123"
|
||||
);
|
||||
});
|
||||
|
||||
it("should not add return URL if already on login page", () => {
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
...window.location,
|
||||
href: "http://localhost:3000/ui/login",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
const loginUrl = buildLoginUrlWithReturn("/ui/login");
|
||||
expect(loginUrl).toBe("/ui/login");
|
||||
});
|
||||
|
||||
it("should handle login URL with existing query parameters", () => {
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
...window.location,
|
||||
href: "http://localhost:3000/ui?page=api-keys",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
const loginUrl = buildLoginUrlWithReturn("/ui/login?foo=bar");
|
||||
expect(loginUrl).toContain("&redirect_to=");
|
||||
});
|
||||
});
|
||||
|
||||
describe("getReturnUrl", () => {
|
||||
it("should prefer URL params over cookie", () => {
|
||||
// Store a URL in cookie
|
||||
storeReturnUrl();
|
||||
|
||||
// Set a different URL in the params
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
...window.location,
|
||||
search: "?redirect_to=http%3A%2F%2Flocalhost%3A3000%2Fui%3Fpage%3Dteams",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
const returnUrl = getReturnUrl();
|
||||
expect(returnUrl).toBe("http://localhost:3000/ui?page=teams");
|
||||
});
|
||||
|
||||
it("should fall back to cookie if no URL param", () => {
|
||||
// Store a URL in cookie first
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
href: "http://localhost:3000/ui?create=true",
|
||||
origin: "http://localhost:3000",
|
||||
hostname: "localhost",
|
||||
pathname: "/ui",
|
||||
search: "?create=true",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
storeReturnUrl();
|
||||
|
||||
// Clear the URL params
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
...window.location,
|
||||
search: "",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
const returnUrl = getReturnUrl();
|
||||
expect(returnUrl).toBe("http://localhost:3000/ui?create=true");
|
||||
});
|
||||
|
||||
it("should return null if no return URL found", () => {
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
...window.location,
|
||||
search: "",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
const returnUrl = getReturnUrl();
|
||||
expect(returnUrl).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("isValidReturnUrl", () => {
|
||||
it("should validate relative URLs starting with /", () => {
|
||||
expect(isValidReturnUrl("/ui?page=api-keys")).toBe(true);
|
||||
expect(isValidReturnUrl("/ui/teams")).toBe(true);
|
||||
});
|
||||
|
||||
it("should reject protocol-relative URLs", () => {
|
||||
expect(isValidReturnUrl("//evil.com")).toBe(false);
|
||||
});
|
||||
|
||||
it("should validate same-hostname URLs (even with different ports) in dev", () => {
|
||||
// Same hostname, same port
|
||||
expect(isValidReturnUrl("http://localhost:3000/ui?page=teams")).toBe(true);
|
||||
// Same hostname, different port (important for dev environments)
|
||||
expect(isValidReturnUrl("http://localhost:4000/ui?page=teams")).toBe(true);
|
||||
});
|
||||
|
||||
it("should reject different-hostname URLs", () => {
|
||||
expect(isValidReturnUrl("http://evil.com/ui")).toBe(false);
|
||||
expect(isValidReturnUrl("https://google.com")).toBe(false);
|
||||
});
|
||||
|
||||
it("should reject empty URLs", () => {
|
||||
expect(isValidReturnUrl("")).toBe(false);
|
||||
});
|
||||
|
||||
it("should reject invalid URLs", () => {
|
||||
expect(isValidReturnUrl("not-a-url")).toBe(false);
|
||||
});
|
||||
|
||||
it("should reject XSS attempts with javascript: protocol", () => {
|
||||
expect(isValidReturnUrl('javascript:alert("xss")')).toBe(false);
|
||||
expect(isValidReturnUrl("javascript:void(0)")).toBe(false);
|
||||
});
|
||||
|
||||
it("should reject data: URLs", () => {
|
||||
expect(isValidReturnUrl("data:text/html,<script>alert(1)</script>")).toBe(false);
|
||||
});
|
||||
|
||||
it("should allow 127.x.x.x addresses in dev environment", () => {
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
href: "http://127.0.0.1:3000/ui",
|
||||
origin: "http://127.0.0.1:3000",
|
||||
hostname: "127.0.0.1",
|
||||
protocol: "http:",
|
||||
pathname: "/ui",
|
||||
search: "",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
expect(isValidReturnUrl("http://127.0.0.1:4000/ui")).toBe(true);
|
||||
});
|
||||
|
||||
it("should allow .local domains in dev environment", () => {
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
href: "http://myapp.local:3000/ui",
|
||||
origin: "http://myapp.local:3000",
|
||||
hostname: "myapp.local",
|
||||
protocol: "http:",
|
||||
pathname: "/ui",
|
||||
search: "",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
// Same hostname with different port should be allowed in dev
|
||||
expect(isValidReturnUrl("http://myapp.local:4000/ui")).toBe(true);
|
||||
});
|
||||
|
||||
it("should reject cross-port redirects in production environment", () => {
|
||||
// Simulate production environment
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
href: "https://app.example.com/ui",
|
||||
origin: "https://app.example.com",
|
||||
hostname: "app.example.com",
|
||||
protocol: "https:",
|
||||
pathname: "/ui",
|
||||
search: "",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
// Same origin should work
|
||||
expect(isValidReturnUrl("https://app.example.com/ui?page=teams")).toBe(true);
|
||||
// Different port should be rejected in production
|
||||
expect(isValidReturnUrl("https://app.example.com:8080/ui")).toBe(false);
|
||||
// Different hostname should be rejected
|
||||
expect(isValidReturnUrl("https://evil.com/ui")).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("consumeReturnUrl", () => {
|
||||
it("should return and clear the stored return URL", () => {
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
href: "http://localhost:3000/ui?create=true",
|
||||
origin: "http://localhost:3000",
|
||||
hostname: "localhost",
|
||||
pathname: "/ui",
|
||||
search: "?create=true",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
storeReturnUrl();
|
||||
|
||||
// Clear the URL params for the consume call
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
href: "http://localhost:3000/ui/login",
|
||||
origin: "http://localhost:3000",
|
||||
hostname: "localhost",
|
||||
pathname: "/ui/login",
|
||||
search: "",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
const returnUrl = consumeReturnUrl();
|
||||
expect(returnUrl).toBe("http://localhost:3000/ui?create=true");
|
||||
expect(getStoredReturnUrl()).toBeNull();
|
||||
});
|
||||
|
||||
it("should return null for invalid return URLs (different hostname)", () => {
|
||||
// Manually set an invalid URL in cookie
|
||||
document.cookie = "litellm_return_url=" + encodeURIComponent("http://evil.com/phishing") + "; path=/";
|
||||
|
||||
const returnUrl = consumeReturnUrl();
|
||||
expect(returnUrl).toBeNull();
|
||||
});
|
||||
|
||||
it("should allow URLs with different ports on same hostname", () => {
|
||||
// Store URL with port 3000
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
href: "http://localhost:3000/ui?create=true",
|
||||
origin: "http://localhost:3000",
|
||||
hostname: "localhost",
|
||||
pathname: "/ui",
|
||||
search: "?create=true",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
storeReturnUrl();
|
||||
|
||||
// Now we're on port 4000
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
href: "http://localhost:4000/ui",
|
||||
origin: "http://localhost:4000",
|
||||
hostname: "localhost",
|
||||
pathname: "/ui",
|
||||
search: "",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
const returnUrl = consumeReturnUrl();
|
||||
// Should be valid because same hostname (localhost)
|
||||
expect(returnUrl).toBe("http://localhost:3000/ui?create=true");
|
||||
});
|
||||
|
||||
it("should return null if no return URL found", () => {
|
||||
Object.defineProperty(window, "location", {
|
||||
value: {
|
||||
...window.location,
|
||||
search: "",
|
||||
},
|
||||
writable: true,
|
||||
});
|
||||
|
||||
const returnUrl = consumeReturnUrl();
|
||||
expect(returnUrl).toBeNull();
|
||||
});
|
||||
});
|
||||
});
|
||||
304
ui/litellm-dashboard/src/utils/returnUrlUtils.ts
Normal file
304
ui/litellm-dashboard/src/utils/returnUrlUtils.ts
Normal file
@ -0,0 +1,304 @@
|
||||
/**
|
||||
* Utility functions for managing return URLs during authentication flows.
|
||||
*
|
||||
* When a user is redirected to login, we store the original URL so they can be
|
||||
* redirected back after successful authentication.
|
||||
*
|
||||
* NOTE: We use cookies instead of sessionStorage because the SSO flow may cross
|
||||
* different ports (e.g., localhost:3000 -> localhost:4000), and sessionStorage
|
||||
* is not shared across different origins. Cookies on the same hostname are shared
|
||||
* across different ports.
|
||||
*/
|
||||
|
||||
const RETURN_URL_COOKIE_NAME = "litellm_return_url";
|
||||
const RETURN_URL_PARAM = "redirect_to";
|
||||
|
||||
/**
|
||||
* Gets the current URL with all query parameters.
|
||||
* Returns null if running on server-side.
|
||||
*/
|
||||
export function getCurrentUrl(): string | null {
|
||||
if (typeof window === "undefined") {
|
||||
return null;
|
||||
}
|
||||
return window.location.href;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets a cookie with the given name and value.
|
||||
* Automatically adds Secure flag when running over HTTPS.
|
||||
*/
|
||||
function setCookie(name: string, value: string, maxAgeSeconds: number = 300): void {
|
||||
if (typeof document === "undefined") {
|
||||
return;
|
||||
}
|
||||
// Set cookie with path=/ so it's available across all paths
|
||||
// Use SameSite=Lax to allow the cookie to be sent on navigation from external sites (SSO redirect)
|
||||
// Add Secure flag when running over HTTPS to prevent cookie from being sent over unencrypted connections
|
||||
const isSecure = typeof window !== "undefined" && window.location.protocol === "https:";
|
||||
const secureFlag = isSecure ? "; Secure" : "";
|
||||
document.cookie = `${name}=${encodeURIComponent(value)}; path=/; max-age=${maxAgeSeconds}; SameSite=Lax${secureFlag}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets a cookie value by name.
|
||||
*/
|
||||
function getCookie(name: string): string | null {
|
||||
if (typeof document === "undefined") {
|
||||
return null;
|
||||
}
|
||||
const match = document.cookie.match(new RegExp(`(^| )${name}=([^;]+)`));
|
||||
if (match) {
|
||||
try {
|
||||
return decodeURIComponent(match[2]);
|
||||
} catch {
|
||||
return match[2];
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Deletes a cookie by name.
|
||||
*/
|
||||
function deleteCookie(name: string): void {
|
||||
if (typeof document === "undefined") {
|
||||
return;
|
||||
}
|
||||
document.cookie = `${name}=; path=/; max-age=0`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Stores the current URL in a cookie before redirecting to login.
|
||||
* This allows us to redirect the user back to their original destination after login.
|
||||
* Cookie expires in 5 minutes (300 seconds).
|
||||
*/
|
||||
export function storeReturnUrl(): void {
|
||||
if (typeof window === "undefined") {
|
||||
return;
|
||||
}
|
||||
|
||||
const currentUrl = getCurrentUrl();
|
||||
if (currentUrl) {
|
||||
setCookie(RETURN_URL_COOKIE_NAME, currentUrl, 300);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieves the stored return URL from the cookie.
|
||||
* Returns null if no return URL is stored or if running on server-side.
|
||||
*/
|
||||
export function getStoredReturnUrl(): string | null {
|
||||
if (typeof window === "undefined") {
|
||||
return null;
|
||||
}
|
||||
return getCookie(RETURN_URL_COOKIE_NAME);
|
||||
}
|
||||
|
||||
/**
|
||||
* Clears the stored return URL from the cookie.
|
||||
* Should be called after redirecting to the return URL.
|
||||
*/
|
||||
export function clearStoredReturnUrl(): void {
|
||||
if (typeof window === "undefined") {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
deleteCookie(RETURN_URL_COOKIE_NAME);
|
||||
} catch (error) {
|
||||
console.error("Failed to clear return URL cookie:", error);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets the return URL from URL query parameters.
|
||||
* Used when the return URL is passed via query string to the login page.
|
||||
*/
|
||||
export function getReturnUrlFromParams(): string | null {
|
||||
if (typeof window === "undefined") {
|
||||
return null;
|
||||
}
|
||||
|
||||
const searchParams = new URLSearchParams(window.location.search);
|
||||
return searchParams.get(RETURN_URL_PARAM);
|
||||
}
|
||||
|
||||
/**
|
||||
* Builds a login URL with the return URL as a query parameter.
|
||||
*
|
||||
* @param baseLoginUrl - The base login URL (e.g., "/ui/login")
|
||||
* @param returnUrl - The URL to redirect to after login (defaults to current URL)
|
||||
*/
|
||||
export function buildLoginUrlWithReturn(baseLoginUrl: string, returnUrl?: string): string {
|
||||
const url = returnUrl || getCurrentUrl();
|
||||
|
||||
if (!url) {
|
||||
return baseLoginUrl;
|
||||
}
|
||||
|
||||
// Don't add return URL if we're already on the login page
|
||||
if (url.includes("/login")) {
|
||||
return baseLoginUrl;
|
||||
}
|
||||
|
||||
const separator = baseLoginUrl.includes("?") ? "&" : "?";
|
||||
return `${baseLoginUrl}${separator}${RETURN_URL_PARAM}=${encodeURIComponent(url)}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets the best return URL to use after login.
|
||||
* Priority:
|
||||
* 1. URL query parameter (redirect_to)
|
||||
* 2. Cookie
|
||||
* 3. null (caller should use default)
|
||||
*/
|
||||
export function getReturnUrl(): string | null {
|
||||
// First check URL params
|
||||
const paramUrl = getReturnUrlFromParams();
|
||||
if (paramUrl) {
|
||||
return paramUrl;
|
||||
}
|
||||
|
||||
// Then check cookie
|
||||
const storedUrl = getStoredReturnUrl();
|
||||
if (storedUrl) {
|
||||
return storedUrl;
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Checks if we're running in a development environment.
|
||||
* Returns true for localhost, 127.0.0.1, IPv6 localhost, or .local domains.
|
||||
* This determines whether cross-port redirects are allowed (dev only).
|
||||
*/
|
||||
function isDevEnvironment(): boolean {
|
||||
if (typeof window === "undefined") {
|
||||
return false;
|
||||
}
|
||||
const hostname = window.location.hostname;
|
||||
return (
|
||||
hostname === "localhost" ||
|
||||
hostname === "127.0.0.1" ||
|
||||
hostname === "::1" ||
|
||||
hostname.startsWith("127.") || // Full IPv4 loopback range (127.0.0.0/8)
|
||||
hostname.endsWith(".local") // Common dev domain suffix
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates a return URL to prevent open redirect attacks.
|
||||
* - Always allows relative URLs (starting with / but not //)
|
||||
* - In dev (localhost): allows same hostname with any port
|
||||
* - In production: requires exact origin match (protocol + hostname + port)
|
||||
*
|
||||
* @param url - The URL to validate
|
||||
* @returns true if the URL is safe to redirect to
|
||||
*/
|
||||
export function isValidReturnUrl(url: string): boolean {
|
||||
if (!url) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Allow relative URLs
|
||||
if (url.startsWith("/") && !url.startsWith("//")) {
|
||||
return true;
|
||||
}
|
||||
|
||||
// For absolute URLs, validate against current origin
|
||||
if (typeof window === "undefined") {
|
||||
return false;
|
||||
}
|
||||
|
||||
try {
|
||||
const returnUrlObj = new URL(url);
|
||||
const currentHostname = window.location.hostname;
|
||||
|
||||
// Hostname must always match
|
||||
if (returnUrlObj.hostname !== currentHostname) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// In dev environments (localhost), allow any port on the same hostname
|
||||
// This supports SSO flows that cross ports (e.g., localhost:3000 -> localhost:4000)
|
||||
if (isDevEnvironment()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
// In production, require exact origin match (protocol + hostname + port)
|
||||
return returnUrlObj.origin === window.location.origin;
|
||||
} catch {
|
||||
// Invalid URL
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
export function normalizeUrlForCompare(url: string): string {
|
||||
if (typeof window === "undefined") {
|
||||
return url;
|
||||
}
|
||||
|
||||
try {
|
||||
const parsed = new URL(url, window.location.origin);
|
||||
let pathname = parsed.pathname;
|
||||
if (pathname.length > 1 && pathname.endsWith("/")) {
|
||||
pathname = pathname.slice(0, -1);
|
||||
}
|
||||
|
||||
const params = new URLSearchParams(parsed.search);
|
||||
const sortedParams = new URLSearchParams();
|
||||
Array.from(params.entries())
|
||||
.sort(([a], [b]) => a.localeCompare(b))
|
||||
.forEach(([key, value]) => {
|
||||
sortedParams.append(key, value);
|
||||
});
|
||||
|
||||
const search = sortedParams.toString();
|
||||
const hash = parsed.hash || "";
|
||||
return `${parsed.origin}${pathname}${search ? `?${search}` : ""}${hash}`;
|
||||
} catch {
|
||||
return url;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets and clears the return URL in one operation.
|
||||
* Returns the validated return URL or null if invalid/not found.
|
||||
*
|
||||
* Priority:
|
||||
* 1. If redirect_to param is valid, use it and clear cookie
|
||||
* 2. If redirect_to param is invalid/missing, check cookie
|
||||
* 3. Only clear cookie when we have a valid URL to return
|
||||
*/
|
||||
export function consumeReturnUrl(): string | null {
|
||||
// Check URL param first
|
||||
const paramUrl = getReturnUrlFromParams();
|
||||
if (paramUrl) {
|
||||
if (isValidReturnUrl(paramUrl)) {
|
||||
clearStoredReturnUrl();
|
||||
return paramUrl;
|
||||
}
|
||||
// Log rejected URLs in development for debugging
|
||||
if (isDevEnvironment()) {
|
||||
console.warn("[returnUrlUtils] Invalid return URL in params rejected:", paramUrl);
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to cookie
|
||||
const storedUrl = getStoredReturnUrl();
|
||||
if (storedUrl) {
|
||||
if (isValidReturnUrl(storedUrl)) {
|
||||
clearStoredReturnUrl();
|
||||
return storedUrl;
|
||||
}
|
||||
// Log rejected URLs in development for debugging
|
||||
if (isDevEnvironment()) {
|
||||
console.warn("[returnUrlUtils] Invalid return URL in cookie rejected:", storedUrl);
|
||||
}
|
||||
}
|
||||
|
||||
// No valid URL found - don't clear cookie (nothing to clear or already invalid)
|
||||
return null;
|
||||
}
|
||||
@ -31,3 +31,32 @@ export const isUserTeamAdminForSingleTeam = (teamMemberWithRoles: Member[] | nul
|
||||
}
|
||||
return teamMemberWithRoles.some((member) => member.user_id === userID && member.role === "admin");
|
||||
};
|
||||
|
||||
export const formatUserRole = (userRole: string): string => {
|
||||
if (!userRole) {
|
||||
return "Undefined Role";
|
||||
}
|
||||
switch (userRole.toLowerCase()) {
|
||||
case "app_owner":
|
||||
return "App Owner";
|
||||
case "demo_app_owner":
|
||||
return "App Owner";
|
||||
case "app_admin":
|
||||
return "Admin";
|
||||
case "proxy_admin":
|
||||
return "Admin";
|
||||
case "proxy_admin_viewer":
|
||||
return "Admin Viewer";
|
||||
case "org_admin":
|
||||
return "Org Admin";
|
||||
case "internal_user":
|
||||
return "Internal User";
|
||||
case "internal_user_viewer":
|
||||
case "internal_viewer": // TODO:remove if deprecated
|
||||
return "Internal Viewer";
|
||||
case "app_user":
|
||||
return "App User";
|
||||
default:
|
||||
return "Unknown Role";
|
||||
}
|
||||
};
|
||||
|
||||
@ -5,12 +5,13 @@ import { vi, describe, it, beforeEach, afterEach, expect } from "vitest";
|
||||
/** ----------------------------
|
||||
* Hoisted helpers for mocks (required by Vitest)
|
||||
* --------------------------- */
|
||||
const { stub, jwtDecodeMock } = vi.hoisted(() => {
|
||||
const { stub, jwtDecodeMock, consumeReturnUrlMock } = vi.hoisted(() => {
|
||||
const React = require("react");
|
||||
const stub = (name: string) => () => React.createElement("div", { "data-testid": name });
|
||||
return {
|
||||
stub,
|
||||
jwtDecodeMock: vi.fn(),
|
||||
consumeReturnUrlMock: vi.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
@ -84,6 +85,14 @@ vi.mock("jwt-decode", () => ({
|
||||
jwtDecode: (token: string) => jwtDecodeMock(token),
|
||||
}));
|
||||
|
||||
vi.mock("@/utils/returnUrlUtils", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import("@/utils/returnUrlUtils")>();
|
||||
return {
|
||||
...actual,
|
||||
consumeReturnUrl: consumeReturnUrlMock,
|
||||
};
|
||||
});
|
||||
|
||||
// Super-light stubs for all heavy components so rendering doesn't explode
|
||||
vi.mock("@/components/navbar", () => ({ default: stub("navbar") }));
|
||||
vi.mock("@/components/user_dashboard", () => ({ default: stub("user-dashboard") }));
|
||||
@ -152,6 +161,7 @@ beforeEach(() => {
|
||||
// Fresh module state & DOM
|
||||
vi.clearAllMocks();
|
||||
clearAllCookies();
|
||||
consumeReturnUrlMock.mockReturnValue(null);
|
||||
|
||||
// Make location.replace spy-able to validate redirect
|
||||
delete (window as any).location;
|
||||
@ -191,9 +201,11 @@ describe("CreateKeyPage auth behavior", () => {
|
||||
// Act
|
||||
render(<CreateKeyPage />);
|
||||
|
||||
// Assert: we eventually redirect to SSO login (single replace, not assign/href)
|
||||
// Assert: we eventually redirect to SSO login with return URL (single replace, not assign/href)
|
||||
await waitFor(() => {
|
||||
expect(window.location.replace).toHaveBeenCalledWith("https://example.com/ui/login");
|
||||
expect(window.location.replace).toHaveBeenCalledWith(
|
||||
expect.stringContaining("https://example.com/ui/login?redirect_to=")
|
||||
);
|
||||
});
|
||||
|
||||
// And we attempted to clear the cookie (defensive deletion)
|
||||
@ -235,4 +247,41 @@ describe("CreateKeyPage auth behavior", () => {
|
||||
expect(screen.getByTestId("navbar")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
it("should not redirect when return URL only differs by query order", async () => {
|
||||
setCookie("token=validtoken");
|
||||
|
||||
jwtDecodeMock.mockImplementation((tok: string) => {
|
||||
expect(tok).toBe("validtoken");
|
||||
return {
|
||||
exp: Math.floor(Date.now() / 1000) + 60 * 60,
|
||||
key: "accessKey-123",
|
||||
user_role: "app_user",
|
||||
user_email: "user@example.com",
|
||||
login_method: "username_password",
|
||||
premium_user: false,
|
||||
auth_header_name: "x-litellm-auth",
|
||||
user_id: "u_123",
|
||||
};
|
||||
});
|
||||
|
||||
// Current URL has params in a different order
|
||||
delete (window as any).location;
|
||||
(window as any).location = {
|
||||
...originalLocation,
|
||||
href: "http://localhost/ui?b=2&a=1",
|
||||
origin: "http://localhost",
|
||||
assign: vi.fn(),
|
||||
replace: vi.fn(),
|
||||
};
|
||||
|
||||
// Return URL has the same params in a different order
|
||||
consumeReturnUrlMock.mockReturnValue("http://localhost/ui?a=1&b=2");
|
||||
|
||||
render(<CreateKeyPage />);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(window.location.replace).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
Loading…
Reference in New Issue
Block a user