diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx index 5b98ca82c16..9863a653adb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx @@ -41,6 +41,7 @@ interface ViewUserDashboardProps { userRole: string | null; userID: string | null; teams: any[] | null; + premiumUser?: boolean; orgAdminOrgIds?: Array<{ organization_id: string; organization_alias: string }> | null; } @@ -56,6 +57,7 @@ const ViewUserDashboard: React.FC = ({ userRole, userID, teams, + premiumUser = false, orgAdminOrgIds, }) => { const isProxyAdmin = userRole ? isProxyAdminRole(userRole) : false; @@ -329,7 +331,12 @@ const ViewUserDashboard: React.FC = ({ {!userListQuery.isLoading && userID && accessToken && ( <> {isProxyAdmin && ( - + )} {isProxyAdmin && ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/page.tsx index 75e5f27b07b..2c0e3d33f94 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/page.tsx @@ -5,7 +5,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; export default function UsersPage() { - const { accessToken, token, userRole, userId } = useAuthorized(); + const { accessToken, token, userRole, userId, premiumUser } = useAuthorized(); const { data: teams } = useTeams(); return ( ); } diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx index 68bfd94bd63..6ab857e6b61 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx @@ -1,4 +1,4 @@ -import { render, screen, waitFor, within } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { describe, it, expect, vi, beforeEach } from "vitest"; @@ -6,6 +6,8 @@ import { CreateUserButton } from "./CreateUserButton"; import * as networking from "./networking"; import { toast } from "@/lib/toast"; import { expectControlBesideLabel } from "../../tests/fieldOrientation"; +import { chooseSelectOption } from "../../tests/test-utils"; +import { MODEL_MAX_BUDGET_PREMIUM_HINT } from "./key_team_helpers/ModelMaxBudgetEditor"; vi.mock("./networking", () => ({ userCreateCall: vi.fn(), @@ -751,4 +753,150 @@ describe("CreateUserButton", () => { expect(submittedPayload().models).toEqual(["no-default-models"]); }); }); + + describe("per-model budgets", () => { + const openStandaloneModal = async (user: ReturnType) => { + expect(await screen.findByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + return screen.getByRole("dialog", { name: /invite user/i }); + }; + + it("locks the per-model budget editor when the proxy has no enterprise license", async () => { + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + renderWithProviders(); + const dialog = await openStandaloneModal(user); + + expect(within(dialog).getByRole("button", { name: /Add Model Budget/i })).toBeDisabled(); + expect(within(dialog).getByText(MODEL_MAX_BUDGET_PREMIUM_HINT)).toBeInTheDocument(); + }); + + it("sends model_max_budget with budget_limit and time_period when a budget is entered", async () => { + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + vi.mocked(networking.modelAvailableCall).mockResolvedValue({ data: [{ id: "gpt-4" }] } as any); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "budget-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-budget", + user_id: "budget-user", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + const dialog = await openStandaloneModal(user); + + await user.type(within(dialog).getByLabelText(/user email/i), "budget@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + await user.click(within(dialog).getByRole("button", { name: /Add Model Budget/i })); + await chooseSelectOption(user, within(dialog).getByPlaceholderText("Select model"), "gpt-4"); + fireEvent.change(within(dialog).getByPlaceholderText("Max spend ($)"), { target: { value: "3" } }); + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalled(); + }); + expect(mockUserCreateCall.mock.calls[0][2].model_max_budget).toStrictEqual({ + "gpt-4": { budget_limit: 3, time_period: "30d" }, + }); + }); + + it("omits model_max_budget from the payload when no budget is entered", async () => { + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "no-budget-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-nobudget", + user_id: "no-budget-user", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + const dialog = await openStandaloneModal(user); + + await user.type(within(dialog).getByLabelText(/user email/i), "nobudget@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalled(); + }); + expect(mockUserCreateCall.mock.calls[0][2]).not.toHaveProperty("model_max_budget"); + }); + + it("clears the per-model budget rows after a user is created in standalone mode", async () => { + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + vi.mocked(networking.modelAvailableCall).mockResolvedValue({ data: [{ id: "gpt-4" }] } as any); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "budget-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-budget", + user_id: "budget-user", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + const dialog = await openStandaloneModal(user); + + await user.type(within(dialog).getByLabelText(/user email/i), "budget@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + await user.click(within(dialog).getByRole("button", { name: /Add Model Budget/i })); + await chooseSelectOption(user, within(dialog).getByPlaceholderText("Select model"), "gpt-4"); + fireEvent.change(within(dialog).getByPlaceholderText("Max spend ($)"), { target: { value: "3" } }); + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalled(); + }); + expect(within(dialog).queryByPlaceholderText("Max spend ($)")).not.toBeInTheDocument(); + expect(within(dialog).getByRole("button", { name: /Add Model Budget/i, hidden: true })).toBeInTheDocument(); + }); + + it("clears the per-model budget rows after a user is created in embedded mode", async () => { + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + const onUserCreated = vi.fn(); + vi.mocked(networking.modelAvailableCall).mockResolvedValue({ data: [{ id: "gpt-4" }] } as any); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "embedded-budget-user" } }); + + renderWithProviders( + , + ); + + await user.type(screen.getByLabelText(/user email/i), "embedded-budget@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /Add Model Budget/i })); + await chooseSelectOption(user, screen.getByPlaceholderText("Select model"), "gpt-4"); + fireEvent.change(screen.getByPlaceholderText("Max spend ($)"), { target: { value: "3" } }); + await user.click(screen.getByRole("button", { name: /create user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalled(); + }); + expect(screen.queryByPlaceholderText("Max spend ($)")).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Add Model Budget/i })).toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.tsx index 83a72e50f5c..88b96e8b540 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.tsx @@ -17,6 +17,7 @@ import React, { useEffect, useState } from "react"; import { useForm } from "react-hook-form"; import TeamDropdown from "./common_components/team_dropdown"; import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key"; +import { ModelMaxBudget, ModelMaxBudgetField } from "./key_team_helpers/ModelMaxBudgetEditor"; import { toast } from "@/lib/toast"; import { getProxyBaseUrl, @@ -46,6 +47,7 @@ interface CreateuserProps { possibleUIRoles: null | Record>; onUserCreated?: (userId: string) => void; isEmbedded?: boolean; + premiumUser?: boolean; } // Define an interface for the UI settings @@ -64,6 +66,7 @@ interface CreateUserFormValues { metadata?: string; send_invite_email: boolean; models?: string[]; + model_max_budget?: ModelMaxBudget; } const EMBEDDED_DEFAULTS: CreateUserFormValues = { @@ -153,6 +156,7 @@ export const CreateUserButton: React.FC = ({ possibleUIRoles, onUserCreated, isEmbedded = false, + premiumUser = false, }) => { const queryClient = useQueryClient(); const [uiSettings, setUISettings] = useState(null); @@ -161,6 +165,8 @@ export const CreateUserButton: React.FC = ({ const [isModalVisible, setIsModalVisible] = useState(false); const [apiuser, setApiuser] = useState(false); const [userModels, setUserModels] = useState([]); + const [modelMaxBudget, setModelMaxBudget] = useState({}); + const [modelMaxBudgetKey, setModelMaxBudgetKey] = useState(0); const [isPersonalKeyOpen, setIsPersonalKeyOpen] = useState(false); const [isInvitationLinkModalVisible, setIsInvitationLinkModalVisible] = useState(false); const [invitationLinkData, setInvitationLinkData] = useState(null); @@ -193,10 +199,16 @@ export const CreateUserButton: React.FC = ({ fetchData(); }, []); + const resetModelMaxBudget = () => { + setModelMaxBudget({}); + setModelMaxBudgetKey((prev) => prev + 1); + }; + const handleCancel = () => { setIsModalVisible(false); setApiuser(false); form.reset(defaultValues); + resetModelMaxBudget(); }; const handleCreate = async (formValues: CreateUserFormValues) => { @@ -205,7 +217,10 @@ export const CreateUserButton: React.FC = ({ if (!isEmbedded) { setIsModalVisible(true); } - const payload = buildCreatePayload(withoutCollapsedModels(formValues, isPersonalKeyOpen)); + const payload = { + ...buildCreatePayload(withoutCollapsedModels(formValues, isPersonalKeyOpen)), + ...(Object.keys(modelMaxBudget).length > 0 && { model_max_budget: modelMaxBudget }), + }; const response = await userCreateCall(accessToken, null, payload); await queryClient.invalidateQueries({ queryKey: ["userList"] }); setApiuser(true); @@ -214,6 +229,7 @@ export const CreateUserButton: React.FC = ({ if (onUserCreated && isEmbedded) { onUserCreated(user_id); form.reset(defaultValues); + resetModelMaxBudget(); return; } @@ -230,6 +246,7 @@ export const CreateUserButton: React.FC = ({ toast.success("API user Created"); form.reset(defaultValues); + resetModelMaxBudget(); localStorage.removeItem("userData" + userID); } catch (error: any) { const errorMessage = error.response?.data?.detail || error?.message || "Error creating the user"; @@ -277,6 +294,17 @@ export const CreateUserButton: React.FC = ({ ); + const modelMaxBudgetField = ( + + ); + const roleField = (label: React.ReactNode) => ( {({ id, value, onChange }) => ( @@ -312,6 +340,7 @@ export const CreateUserButton: React.FC = ({ {roleField("User Role")} {teamField} {metadataField} + {modelMaxBudgetField} {sendInviteEmailField}
@@ -386,6 +415,7 @@ export const CreateUserButton: React.FC = ({ {metadataField} + {modelMaxBudgetField} {sendInviteEmailField} diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 45245ff95b3..96121afb7a8 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1847,6 +1847,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp possibleUIRoles={possibleUIRoles} onUserCreated={handleUserCreated} isEmbedded={true} + premiumUser={premiumUser === true} />