From 0e97fec98017da1055e7fe3fab1119329a4311a0 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 18 Sep 2026 00:27:33 +0000 Subject: [PATCH 1/2] feat(ui): add per-model budgets to the create user form Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../users/_components/view_users.tsx | 9 +- .../src/app/(dashboard)/users/page.tsx | 3 +- .../src/components/CreateUserButton.test.tsx | 85 ++++++++++++++++++- .../src/components/CreateUserButton.tsx | 25 +++++- .../organisms/create_key_button.tsx | 1 + 5 files changed, 119 insertions(+), 4 deletions(-) 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 20ce22b6444..658e71d12a3 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..9555dd28e1e 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,85 @@ 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"); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.tsx index 83a72e50f5c..8f2ccc1a2b6 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,7 @@ export const CreateUserButton: React.FC = ({ const [isModalVisible, setIsModalVisible] = useState(false); const [apiuser, setApiuser] = useState(false); const [userModels, setUserModels] = useState([]); + const [modelMaxBudget, setModelMaxBudget] = useState({}); const [isPersonalKeyOpen, setIsPersonalKeyOpen] = useState(false); const [isInvitationLinkModalVisible, setIsInvitationLinkModalVisible] = useState(false); const [invitationLinkData, setInvitationLinkData] = useState(null); @@ -197,6 +202,7 @@ export const CreateUserButton: React.FC = ({ setIsModalVisible(false); setApiuser(false); form.reset(defaultValues); + setModelMaxBudget({}); }; const handleCreate = async (formValues: CreateUserFormValues) => { @@ -205,7 +211,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 +223,7 @@ export const CreateUserButton: React.FC = ({ if (onUserCreated && isEmbedded) { onUserCreated(user_id); form.reset(defaultValues); + setModelMaxBudget({}); return; } @@ -230,6 +240,7 @@ export const CreateUserButton: React.FC = ({ toast.success("API user Created"); form.reset(defaultValues); + setModelMaxBudget({}); localStorage.removeItem("userData" + userID); } catch (error: any) { const errorMessage = error.response?.data?.detail || error?.message || "Error creating the user"; @@ -277,6 +288,16 @@ export const CreateUserButton: React.FC = ({ ); + const modelMaxBudgetField = ( + + ); + const roleField = (label: React.ReactNode) => ( {({ id, value, onChange }) => ( @@ -312,6 +333,7 @@ export const CreateUserButton: React.FC = ({ {roleField("User Role")} {teamField} {metadataField} + {modelMaxBudgetField} {sendInviteEmailField}
@@ -386,6 +408,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 b8ea8de7f59..f34b70190c5 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1820,6 +1820,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp possibleUIRoles={possibleUIRoles} onUserCreated={handleUserCreated} isEmbedded={true} + premiumUser={premiumUser === true} /> From 4e90a80275cd040cb54c9f5b14b81ffc983aee67 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 18 Sep 2026 21:50:11 +0000 Subject: [PATCH 2/2] fix(ui): remount the per-model budget editor when the create user form resets Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/components/CreateUserButton.test.tsx | 65 +++++++++++++++++++ .../src/components/CreateUserButton.tsx | 13 +++- 2 files changed, 75 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx index 9555dd28e1e..6ab857e6b61 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx @@ -833,5 +833,70 @@ describe("CreateUserButton", () => { }); 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 8f2ccc1a2b6..88b96e8b540 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.tsx @@ -166,6 +166,7 @@ export const CreateUserButton: React.FC = ({ 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); @@ -198,11 +199,16 @@ export const CreateUserButton: React.FC = ({ fetchData(); }, []); + const resetModelMaxBudget = () => { + setModelMaxBudget({}); + setModelMaxBudgetKey((prev) => prev + 1); + }; + const handleCancel = () => { setIsModalVisible(false); setApiuser(false); form.reset(defaultValues); - setModelMaxBudget({}); + resetModelMaxBudget(); }; const handleCreate = async (formValues: CreateUserFormValues) => { @@ -223,7 +229,7 @@ export const CreateUserButton: React.FC = ({ if (onUserCreated && isEmbedded) { onUserCreated(user_id); form.reset(defaultValues); - setModelMaxBudget({}); + resetModelMaxBudget(); return; } @@ -240,7 +246,7 @@ export const CreateUserButton: React.FC = ({ toast.success("API user Created"); form.reset(defaultValues); - setModelMaxBudget({}); + resetModelMaxBudget(); localStorage.removeItem("userData" + userID); } catch (error: any) { const errorMessage = error.response?.data?.detail || error?.message || "Error creating the user"; @@ -290,6 +296,7 @@ export const CreateUserButton: React.FC = ({ const modelMaxBudgetField = (