From b798cf30ab25ca180fd504a5cc2a281691a99f16 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 13 Nov 2026 12:29:36 -0800 Subject: [PATCH] feat(customer-portal): configure budget groups atomically (#51661) * feat(ui): add customer organization budget groups tab Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Add cap policy controls to customer budget groups Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Clarify per-group cap policy help text Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Fix customer budget groups formatting and lint budgets Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(customer-portal): save budget group configuration atomically * refactor(customer-portal): name unchanged budget group state * fix(customer-portal): chunk removed budget reference checks --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_v1/customers.py | 26 +++ .../management_helpers/customer_hierarchy.py | 81 +++++++- .../customer_hierarchy.py | 39 +++- .../CreateCustomerDialog.integration.test.tsx | 8 +- .../_components/CreateCustomerDialog.tsx | 165 ++------------- .../CreateCustomerHierarchyDialog.test.tsx | 1 + .../CustomerBudgetGroups.integration.test.tsx | 136 ++++++++++++ .../_components/CustomerBudgetGroups.tsx | 107 ++++++++++ .../CustomerBudgetGroupsEditor.tsx | 194 ++++++++++++++++++ .../_components/CustomerHierarchyManager.tsx | 13 +- .../CustomerHierarchyNavigation.tsx | 43 ++-- ...omerHierarchySettings.integration.test.tsx | 1 + .../customerBudgetGroupsPlan.test.ts | 125 +++++++++++ .../_components/customerBudgetGroupsPlan.ts | 58 ++++++ .../customerHierarchyForms.test.ts | 1 + .../customerOrganizationForm.test.ts | 21 +- .../_components/customerOrganizationForm.ts | 54 +++-- .../_components/useCustomerHierarchy.ts | 31 ++- .../manage/page.integration.test.tsx | 88 ++++++++ .../customers/page.integration.test.tsx | 3 + .../src/i18n/locales/en.json | 10 +- .../src/lib/http/schema.d.ts | 97 ++++++++- .../management_v1/test_customers.py | 194 ++++++++++++++++++ 23 files changed, 1297 insertions(+), 199 deletions(-) create mode 100644 litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroups.integration.test.tsx create mode 100644 litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroups.tsx create mode 100644 litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroupsEditor.tsx create mode 100644 litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerBudgetGroupsPlan.test.ts create mode 100644 litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerBudgetGroupsPlan.ts diff --git a/litellm/proxy/management_endpoints/management_v1/customers.py b/litellm/proxy/management_endpoints/management_v1/customers.py index 2e323dcd3da..be538603d6f 100644 --- a/litellm/proxy/management_endpoints/management_v1/customers.py +++ b/litellm/proxy/management_endpoints/management_v1/customers.py @@ -28,6 +28,7 @@ from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, # pyright: ignore[reportUnknownVariableType] # legacy decorator is untyped ) from litellm.types.proxy.management_endpoints.customer_hierarchy import ( + BudgetGroupsUpdate, BudgetTemplate, CustomerMember, MemberCreate, @@ -102,6 +103,11 @@ _ERROR_DETAILS: Final[Mapping[ErrorCode, tuple[int, str, str]]] = MappingProxyTy ), ErrorCode.FORBIDDEN: (403, "Forbidden", "Only proxy admins can manage customer hierarchies."), ErrorCode.DATABASE_UNAVAILABLE: (503, "Database not connected", "Database not connected."), + ErrorCode.CONFIGURATION_CONFLICT: ( + 409, + "Budget groups changed", + "Budget groups changed since you started editing. Reload before saving.", + ), ErrorCode.ALIAS_ALREADY_EXISTS: ( 409, "Customer org alias already exists", @@ -526,6 +532,26 @@ async def add_customer_budget_template( return resolve(await hierarchy.add_template(customer_org_id, data, actor)) +@router.patch( + "/customer_organizations/{customer_org_id}/default_budgets", + response_model=Organization, + description="Atomically configure budget groups and their default. Requires the current budget_groups_revision. " + "Existing groups retain omitted settings and window reset dates. Stale revisions and referenced deletions " + "return 409 without applying changes. Proxy admins only.", +) +@management_endpoint_wrapper +async def configure_customer_budget_groups( + customer_org_id: str, + data: BudgetGroupsUpdate, + actor: Annotated[str, Depends(customer_actor)], + hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)], + cache: AuthCache, +) -> Organization: + organization: Final = resolve(await hierarchy.configure_templates(customer_org_id, data, actor)) + await forget_customer_hierarchies(resolve(await hierarchy.member_ids_in_organization(customer_org_id)), cache) + return organization + + @router.patch( "/customer_organizations/{customer_org_id}/default_budgets/{budget_id}", response_model=Organization, diff --git a/litellm/proxy/management_helpers/customer_hierarchy.py b/litellm/proxy/management_helpers/customer_hierarchy.py index a65e3505715..63b51267b6d 100644 --- a/litellm/proxy/management_helpers/customer_hierarchy.py +++ b/litellm/proxy/management_helpers/customer_hierarchy.py @@ -19,10 +19,12 @@ from litellm.proxy.db.db_span import db_span from litellm.proxy.list_api.common import escape_like from litellm.proxy.management_helpers.customer_budget_usage import customer_budget_usage from litellm.repositories.budget_repository import create_budget_row, update_budget_row -from litellm.repositories.chunked_in import delete_many_in, find_many_in +from litellm.repositories.chunked_in import count_in, delete_many_in, find_many_in from litellm.repositories.prisma_protocols import TableActions from litellm.types.proxy.management_endpoints.customer_hierarchy import ( Budget, + BudgetGroupChange, + BudgetGroupsUpdate, BudgetTemplate, CustomerMember, MemberCreate, @@ -62,6 +64,7 @@ class ErrorCode(str, Enum): REQUEST_ALREADY_PENDING = "request_already_pending" EXPIRY_IN_PAST = "expiry_in_past" KEY_NOT_IN_ORG = "key_not_in_org" + CONFIGURATION_CONFLICT = "configuration_conflict" @dataclass(frozen=True, slots=True) @@ -96,6 +99,11 @@ class _CapacityCounts(BaseModel): _LOCK_ORGANIZATION_SQL: Final = ( 'SELECT 1 AS locked FROM "LiteLLM_CustomerOrganizationTable" WHERE customer_org_id = $1 FOR UPDATE' ) +_LOCK_TEMPLATES_SQL: Final = ( + 'SELECT budget_id FROM "LiteLLM_BudgetTable" WHERE budget_id IN ' + '(SELECT unnest(budget_template_ids) FROM "LiteLLM_CustomerOrganizationTable" WHERE customer_org_id = $1) ' + "ORDER BY budget_id FOR UPDATE" +) _LOCK_CUSTOMER_IDENTITIES_SQL: Final = ( "SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtextextended('litellm:customer-identities', 0))" @@ -184,6 +192,23 @@ def budget_values(budget: Budget) -> Mapping[str, object]: return _VALUES.validate_python(budget.model_dump(include=set(Budget.model_fields))) +def group_budget_patch(group: BudgetGroupChange, current: OrganizationBudget | None) -> Budget: + patch: Final = Budget.model_validate(group.model_dump(exclude_unset=True, exclude={"budget_id", "name"})) + if patch.budget_limits is None: + return patch + resets: Final = ( + {limit.budget_duration: limit.reset_at for limit in (current.budget_limits or ())} if current else {} + ) + return patch.model_copy( + update={ + "budget_limits": tuple( + limit.model_copy(update={"reset_at": resets.get(limit.budget_duration, limit.reset_at)}) + for limit in patch.budget_limits + ) + } + ) + + _JSON_BUDGET_COLUMNS: Final = frozenset({"model_tpm_limit", "model_rpm_limit"}) @@ -718,6 +743,7 @@ class CustomerHierarchyStore: alias_failure: Final = await self.claim_organization_alias(payload.alias, organization_id) if alias_failure is not None: return alias_failure + await self.lock_organization(organization_id) organization: Final = await self.get_organization(organization_id) if isinstance(organization, Failure): return organization @@ -787,6 +813,7 @@ class CustomerHierarchyStore: return None async def resolve_member_template(self, organization_id: str, template_id: str | None) -> str | Failure | None: + await self.lock_organization(organization_id) organization: Final = await self.get_organization(organization_id) if isinstance(organization, Failure): return Failure(ErrorCode.PARENT_NOT_FOUND) @@ -797,6 +824,53 @@ class CustomerHierarchyStore: return selected return Failure(ErrorCode.TEMPLATE_NOT_IN_ORG) + async def configure_templates( + self, organization_id: str, payload: BudgetGroupsUpdate, actor: str + ) -> Organization | Failure: + await self.lock_organization(organization_id) + await self.db.query_raw(_LOCK_TEMPLATES_SQL, organization_id) + organization: Final = await self.get_organization(organization_id) + if isinstance(organization, Failure): + return organization + if payload.expected_revision != organization.budget_groups_revision: + return Failure(ErrorCode.CONFIGURATION_CONFLICT) + originals: Final = {group.budget_id: group for group in organization.default_budgets} + retained: Final = tuple(group.budget_id for group in payload.default_budgets if group.budget_id is not None) + if len(set(retained)) != len(retained) or not set(retained).issubset(originals): + return Failure(ErrorCode.TEMPLATE_NOT_IN_ORG) + names: Final = tuple(group.name for group in payload.default_budgets) + if len(set(names)) != len(names): + return Failure(ErrorCode.INVALID_TEMPLATE_NAMES) + if payload.default_budget_template is not None and payload.default_budget_template not in names: + return Failure(ErrorCode.INVALID_DEFAULT_BUDGET) + removed: Final = tuple(budget_id for budget_id in originals if budget_id not in retained) + if await count_in(self.members, "budget_id", removed): + return Failure(ErrorCode.TEMPLATE_IN_USE) + ids: Final = tuple( + [ + await self.write_configured_template(group, originals.get(group.budget_id or ""), actor) + for group in payload.default_budgets + ] + ) + default_id: Final = ( + ids[names.index(payload.default_budget_template)] if payload.default_budget_template else None + ) + await self.write_template_metadata(organization_id, ids, default_id, actor) + await self._delete_unused_budgets(removed) + return await self.get_organization(organization_id) + + async def write_configured_template( + self, group: BudgetGroupChange, current: OrganizationBudget | None, actor: str + ) -> str: + patch: Final = group_budget_patch(group, current) + if current is None: + return await self.create_budget(patch, actor, group.name) + if patch.model_fields_set: + await self.update_budget(current.budget_id, patch, actor) + if group.name != current.name: + await update_budget_row(self.budgets, current.budget_id, {"name": group.name}, actor) + return current.budget_id + async def create_team(self, payload: TeamCreate, actor: str) -> Team | Failure: parent: Final = await self.organizations.find_unique(where={"customer_org_id": payload.parent_org_id}) if parent is None: @@ -1210,6 +1284,11 @@ class CustomerHierarchy: async def add_template(self, organization_id: str, payload: BudgetTemplate, actor: str) -> Organization | Failure: return await self.transaction(lambda store: store.add_template(organization_id, payload, actor)) + async def configure_templates( + self, organization_id: str, payload: BudgetGroupsUpdate, actor: str + ) -> Organization | Failure: + return await self.transaction(lambda store: store.configure_templates(organization_id, payload, actor)) + async def update_template( self, organization_id: str, budget_id: str, payload: TemplatePatch, actor: str ) -> Organization | Failure: diff --git a/litellm/types/proxy/management_endpoints/customer_hierarchy.py b/litellm/types/proxy/management_endpoints/customer_hierarchy.py index 06300affb76..5550e2c1f81 100644 --- a/litellm/types/proxy/management_endpoints/customer_hierarchy.py +++ b/litellm/types/proxy/management_endpoints/customer_hierarchy.py @@ -1,8 +1,19 @@ from collections.abc import Mapping from datetime import datetime +from hashlib import sha256 +from json import dumps from typing import Final, Literal, TypeAlias -from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, NonNegativeInt, field_validator, model_validator +from pydantic import ( + AwareDatetime, + BaseModel, + ConfigDict, + Field, + NonNegativeInt, + computed_field, + field_validator, + model_validator, +) from litellm.models.budget import BudgetLimits from litellm.proxy.common_utils.timezone_utils import budget_duration_error @@ -50,6 +61,20 @@ class OrganizationBudget(BudgetTemplate): budget_id: str +class BudgetGroupChange(BudgetTemplate): + budget_id: str | None = Field(default=None, description="Existing template ID. Omit to create a template.") + + +class BudgetGroupsUpdate(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + expected_revision: str = Field(min_length=64, max_length=64) + default_budgets: tuple[BudgetGroupChange, ...] = Field( + description="Complete retained group list. Existing groups patch only supplied fields. Omitted groups are deleted." + ) + default_budget_template: str | None = Field(description="Final name of the default group, or null for no default.") + + class OrganizationFields(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") @@ -155,6 +180,18 @@ class Organization(AuditFields): default_budget_template_id: str | None = None default_budgets: tuple[OrganizationBudget, ...] = () + @computed_field + @property + def budget_groups_revision(self) -> str: + configuration: Final = { + "default": self.default_budget_template_id, + "groups": tuple( + group.model_dump(mode="json", exclude={"budget_limits": {"__all__": {"reset_at"}}}) + for group in self.default_budgets + ), + } + return sha256(dumps(configuration, sort_keys=True).encode()).hexdigest() + class Team(AuditFields): customer_team_id: str diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.integration.test.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.integration.test.tsx index 477ecde42e2..ebc3ede926f 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.integration.test.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.integration.test.tsx @@ -73,6 +73,8 @@ describe("CreateCustomerDialog", () => { const user = userEvent.setup(); renderWithProviders(); fireEvent.change(screen.getByLabelText("Customer name"), { target: { value: "Example customer" } }); + await user.click(screen.getByRole("combobox", { name: "Group 2 at the cap" })); + await user.click(screen.getByRole("option", { name: "Pause, they can request more" })); await user.click(screen.getByRole("button", { name: "Add budget group" })); fireEvent.change(screen.getByLabelText("Group 4 name"), { target: { value: "custom_user" } }); fireEvent.change(screen.getByLabelText("Group 4 yearly budget (USD)"), { target: { value: "300" } }); @@ -83,7 +85,11 @@ describe("CreateCustomerDialog", () => { await waitFor(() => expect(onCreated).toHaveBeenCalledWith("created-org")); expect(creations).toHaveLength(1); expect(creations[0]).toMatchObject({ - default_budgets: [{ name: "power_user" }, { name: "default_user" }, { name: "low_user" }], + default_budgets: [ + { name: "power_user", cap_policy: "pause_until_reset" }, + { name: "default_user", cap_policy: "pause_allow_increase" }, + { name: "low_user", cap_policy: "pause_until_reset" }, + ], default_budget_template: "power_user", }); }); diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.tsx index 74b67942f9e..eba40f0df26 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.tsx @@ -2,7 +2,7 @@ import { useState } from "react"; import { useQueryClient } from "@tanstack/react-query"; -import { ChevronRight, Plus, Trash2 } from "lucide-react"; +import { ChevronRight } from "lucide-react"; import { useTranslation } from "react-i18next"; import { toast } from "sonner"; import { $api } from "@/lib/http/api"; @@ -17,59 +17,22 @@ import { DialogTitle, } from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; -import { InputGroup, InputGroupAddon, InputGroupInput, InputGroupText } from "@/components/ui/input-group"; import { Label } from "@/components/ui/label"; -import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import { buildCustomerOrganizationBody, initialCustomerOrganizationFields, - type CustomerBudgetGroupFields, type CustomerOrganizationFormError, } from "./customerOrganizationForm"; +import CustomerBudgetGroupsEditor, { BudgetInput } from "./CustomerBudgetGroupsEditor"; + interface Props { open: boolean; onOpenChange: (o: boolean) => void; onCreated: (customerOrgId: string) => void; } -function BudgetInput({ - id, - label, - value, - onChange, -}: { - id: string; - label: string; - value: string; - onChange: (value: string) => void; -}) { - const { t } = useTranslation(); - return ( - - - $ - - onChange(event.target.value)} - placeholder={t("customerPortal.createDialog.noLimit")} - /> - - ); -} - const rateFields = ["tpm", "rpm"] as const; -const groupBudgetFields = [ - ["monthlyBudget", "monthly"], - ["yearlyBudget", "yearly"], -] as const; - export default function CreateCustomerDialog({ open, onOpenChange, onCreated }: Props) { const queryClient = useQueryClient(); const { t } = useTranslation(); @@ -87,24 +50,6 @@ export default function CreateCustomerDialog({ open, onOpenChange, onCreated }: onError: (e: unknown) => toast.error(t("customerPortal.toasts.createFailed", { error: extractProxyErrorMessage(e) })), }); - const updateGroup = (id: string, patch: Partial>) => { - setFields((current) => ({ - ...current, - groups: current.groups.map((group) => (group.id === id ? { ...group, ...patch } : group)), - })); - setError(null); - }; - const removeGroup = (id: string) => { - setFields((current) => { - const groups = current.groups.filter((group) => group.id !== id); - return { - ...current, - groups, - defaultGroupId: current.defaultGroupId === id ? groups[0].id : current.defaultGroupId, - }; - }); - setError(null); - }; const submit = (event: React.FormEvent) => { event.preventDefault(); if (create.isPending) return; @@ -161,101 +106,15 @@ export default function CreateCustomerDialog({ open, onOpenChange, onCreated }: />

{t("customerPortal.createDialog.budgetHelp")}

-
-
-

- {t("customerPortal.createDialog.groups")} -

-

{t("customerPortal.createDialog.groupsHelp")}

-
- - setFields({ ...fields, defaultGroupId: String(defaultGroupId) })} - aria-label={t("customerPortal.createDialog.defaultGroup")} - disabled={create.isPending} - className="gap-1" - > - {fields.groups.map((group, index) => ( -
-
- - updateGroup(group.id, { name: event.target.value })} - className="font-mono text-xs" - /> -
- {groupBudgetFields.map(([field, period]) => ( -
- - updateGroup(group.id, { [field]: value })} - /> -
- ))} - - -
- ))} - - -

{t("customerPortal.createDialog.defaultHelp")}

-

{t("customerPortal.createDialog.resetHelp")}

-
+ { + setFields({ ...fields, groups, defaultGroupId }); + setError(null); + }} + />
diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerHierarchyDialog.test.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerHierarchyDialog.test.tsx index aac5807f490..5db40e0ce5a 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerHierarchyDialog.test.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerHierarchyDialog.test.tsx @@ -17,6 +17,7 @@ const audit = { spend: 0, }; const organization: Organization = { + budget_groups_revision: "a".repeat(64), ...audit, customer_org_id: "org-1", alias: "Example customer", diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroups.integration.test.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroups.integration.test.tsx new file mode 100644 index 00000000000..392d270f677 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroups.integration.test.tsx @@ -0,0 +1,136 @@ +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { registerBaseUrlGetter } from "@/lib/http/runtime"; +import { render, screen, fireEvent, waitFor } from "../../../../../tests/test-utils"; +import CustomerBudgetGroups from "./CustomerBudgetGroups"; +import type { Organization } from "./customerHierarchyForms"; + +const organization: Organization = { + customer_org_id: "org-1", + alias: "Example", + default_budget_template_id: "standard", + created_at: "", + updated_at: "", + created_by: "", + updated_by: "", + spend: 0, + budget_groups_revision: "a".repeat(64), + member_count: 0, + default_budgets: [ + { budget_id: "standard", name: "Standard", budget_limits: [{ budget_duration: "1mo", max_budget: 10 }] }, + { budget_id: "old", name: "Old", budget_limits: [{ budget_duration: "1mo", max_budget: 20 }] }, + ], +}; +const fetchMock = vi.fn(); +const onSaved = vi.fn(); +const json = (body: unknown, status = 200) => + new Response(JSON.stringify(body), { status, headers: { "Content-Type": "application/json" } }); + +describe("CustomerBudgetGroups", () => { + beforeEach(() => { + registerBaseUrlGetter(() => "http://localhost"); + fetchMock.mockReset(); + onSaved.mockReset(); + vi.stubGlobal("fetch", fetchMock); + }); + afterEach(() => vi.unstubAllGlobals()); + + it("should save additions, edited limits, the new default and deletions together", async () => { + const user = userEvent.setup(); + fetchMock.mockImplementation(async () => + json({ + ...organization, + default_budgets: [...organization.default_budgets, { budget_id: "created-id", name: "New" }], + }), + ); + render(); + expect(screen.getByRole("button", { name: "Save changes" })).toBeDisabled(); + fireEvent.change(screen.getByRole("spinbutton", { name: "Group 1 monthly budget (USD)" }), { + target: { value: "15" }, + }); + await user.click(screen.getByRole("combobox", { name: "Group 1 at the cap" })); + await user.click(screen.getByRole("option", { name: "Pause, they can request more" })); + await user.click(screen.getByRole("button", { name: "Add budget group" })); + fireEvent.change(screen.getByRole("textbox", { name: "Group 3 name" }), { target: { value: "New" } }); + fireEvent.change(screen.getByRole("spinbutton", { name: "Group 3 monthly budget (USD)" }), { + target: { value: "30" }, + }); + await user.click(screen.getByRole("radio", { name: "Use New as default" })); + await user.click(screen.getByRole("button", { name: "Remove Old" })); + await user.click(screen.getByRole("button", { name: "Save changes" })); + await waitFor(() => expect(onSaved).toHaveBeenCalledOnce()); + const requests = fetchMock.mock.calls.map(([request]) => { + if (!(request instanceof Request)) throw new Error("Expected shared HTTP client request"); + return request; + }); + expect(requests.map((request) => [request.method, new URL(request.url).pathname])).toEqual([ + ["PATCH", "/management/v1/customer_organizations/org-1/default_budgets"], + ]); + expect(await requests[0].clone().json()).toEqual({ + expected_revision: organization.budget_groups_revision, + default_budget_template: "New", + default_budgets: [ + { + budget_id: "standard", + name: "Standard", + cap_policy: "pause_allow_increase", + budget_limits: [{ budget_duration: "1mo", max_budget: 15 }], + }, + { name: "New", cap_policy: "pause_until_reset", budget_limits: [{ budget_duration: "1mo", max_budget: 30 }] }, + ], + }); + }); + + it("should keep the whole draft after a rejected save and allow retry", async () => { + const user = userEvent.setup(); + fetchMock.mockImplementation(async () => json({ detail: "A referenced template cannot be deleted." }, 409)); + render(); + fireEvent.change(screen.getByRole("spinbutton", { name: "Group 1 monthly budget (USD)" }), { + target: { value: "15" }, + }); + await user.click(screen.getByRole("button", { name: "Remove Old" })); + await user.click(screen.getByRole("button", { name: "Save changes" })); + expect(await screen.findByRole("alert")).toHaveTextContent("A referenced template cannot be deleted."); + expect(screen.getByRole("spinbutton", { name: "Group 1 monthly budget (USD)" })).toHaveValue(15); + expect(screen.queryByDisplayValue("Old")).not.toBeInTheDocument(); + expect(onSaved).not.toHaveBeenCalled(); + expect(fetchMock).toHaveBeenCalledOnce(); + fetchMock.mockImplementation(async () => json(organization)); + await user.click(screen.getByRole("button", { name: "Save changes" })); + await waitFor(() => expect(onSaved).toHaveBeenCalledOnce()); + }); + + it("should automatically make the first added group the default", async () => { + const user = userEvent.setup(); + fetchMock.mockImplementation(async () => json(organization)); + render( + , + ); + await user.click(screen.getByRole("button", { name: "Add budget group" })); + fireEvent.change(screen.getByRole("textbox", { name: "Group 1 name" }), { target: { value: "First" } }); + fireEvent.change(screen.getByRole("spinbutton", { name: "Group 1 monthly budget (USD)" }), { + target: { value: "10" }, + }); + expect(screen.getByRole("radio", { name: "Use First as default" })).toBeChecked(); + await user.click(screen.getByRole("button", { name: "Save changes" })); + await waitFor(() => expect(onSaved).toHaveBeenCalledOnce()); + const request = fetchMock.mock.calls[0][0]; + expect(request).toBeInstanceOf(Request); + expect(await (request as Request).clone().json()).toMatchObject({ default_budget_template: "First" }); + }); + + it("should start with an empty editor and no pending changes for an organization without groups", () => { + render( + , + ); + expect(screen.queryByRole("textbox")).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Add budget group" })).toBeEnabled(); + expect(screen.getByRole("button", { name: "Save changes" })).toBeDisabled(); + }); +}); diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroups.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroups.tsx new file mode 100644 index 00000000000..fa9a85716b4 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroups.tsx @@ -0,0 +1,107 @@ +"use client"; + +import { useState, type FormEvent } from "react"; +import { useTranslation } from "react-i18next"; +import { toast } from "sonner"; +import { fetchClient } from "@/lib/http/api"; +import { extractProxyErrorMessage } from "@/lib/http/client"; +import { Button } from "@/components/ui/button"; +import type { Organization } from "./customerHierarchyForms"; +import CustomerBudgetGroupsEditor from "./CustomerBudgetGroupsEditor"; +import { customerBudgetGroupRows, planCustomerBudgetGroups } from "./customerBudgetGroupsPlan"; +import type { CustomerBudgetGroupFields } from "./customerOrganizationForm"; + +interface Draft { + base: Organization; + groups: CustomerBudgetGroupFields[]; + defaultGroupId: string; +} + +export default function CustomerBudgetGroups({ + organization, + onSaved, +}: { + organization: Organization; + onSaved: (organization: Organization) => void | Promise; +}) { + const { t } = useTranslation(); + const [draft, setDraft] = useState(null); + const [pending, setPending] = useState(false); + const [error, setError] = useState(null); + const base = draft?.base ?? organization; + const groups = draft?.groups ?? customerBudgetGroupRows(base); + const defaultGroupId = draft?.defaultGroupId ?? base.default_budget_template_id ?? ""; + const plan = planCustomerBudgetGroups(base, groups, defaultGroupId); + const unchanged = plan.ok && !plan.changed; + const changeGroups = (groups: CustomerBudgetGroupFields[], defaultGroupId: string) => { + setDraft({ base, groups, defaultGroupId }); + setError(null); + }; + const save = async (event: FormEvent) => { + event.preventDefault(); + if (pending) return; + if (!plan.ok) { + setError(t(`customerPortal.createDialog.errors.${plan.error}`)); + return; + } + setPending(true); + setError(null); + try { + const response = await fetchClient.PATCH( + "/management/v1/customer_organizations/{customer_org_id}/default_budgets", + { params: { path: { customer_org_id: base.customer_org_id } }, body: plan.body }, + ); + if (!response.data) throw new Error(t("customerPortal.hierarchy.error")); + await onSaved(response.data); + setDraft(null); + toast.success(t("customerPortal.hierarchy.saved")); + } catch (failure) { + setError(extractProxyErrorMessage(failure)); + } finally { + setPending(false); + } + }; + const reload = async () => { + setPending(true); + try { + const response = await fetchClient.GET("/management/v1/customer_organizations/{customer_org_id}", { + params: { path: { customer_org_id: base.customer_org_id } }, + }); + if (!response.data) throw new Error(t("customerPortal.hierarchy.error")); + await onSaved(response.data); + setDraft(null); + setError(null); + } catch (failure) { + setError(extractProxyErrorMessage(failure)); + } finally { + setPending(false); + } + }; + return ( +
+
+ +
+ {error && ( +

+ {error} +

+ )} +
+ + {draft && ( + + )} +
+
+ ); +} diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroupsEditor.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroupsEditor.tsx new file mode 100644 index 00000000000..763db86f0a6 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroupsEditor.tsx @@ -0,0 +1,194 @@ +import { Plus, Trash2 } from "lucide-react"; +import { useTranslation } from "react-i18next"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { InputGroup, InputGroupAddon, InputGroupInput, InputGroupText } from "@/components/ui/input-group"; +import { Label } from "@/components/ui/label"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { defaultCustomerCapPolicy, type CustomerBudgetGroupFields } from "./customerOrganizationForm"; + +export function BudgetInput({ + id, + label, + value, + onChange, +}: { + id: string; + label: string; + value: string; + onChange: (value: string) => void; +}) { + const { t } = useTranslation(); + return ( + + + $ + + onChange(event.target.value)} + placeholder={t("customerPortal.createDialog.noLimit")} + /> + + ); +} + +const groupBudgetFields = [ + ["monthlyBudget", "monthly"], + ["yearlyBudget", "yearly"], +] as const; + +export default function CustomerBudgetGroupsEditor({ + groups, + defaultGroupId, + onChange, + disabled = false, +}: { + groups: CustomerBudgetGroupFields[]; + defaultGroupId: string; + onChange: (groups: CustomerBudgetGroupFields[], defaultGroupId: string) => void; + disabled?: boolean; +}) { + const { t } = useTranslation(); + const updateGroup = (id: string, patch: Partial>) => + onChange( + groups.map((group) => (group.id === id ? { ...group, ...patch } : group)), + defaultGroupId, + ); + const removeGroup = (id: string) => { + const remaining = groups.filter((group) => group.id !== id); + onChange(remaining, defaultGroupId === id ? remaining[0]?.id ?? "" : defaultGroupId); + }; + const addGroup = () => { + const group: CustomerBudgetGroupFields = { + id: crypto.randomUUID(), + name: "", + monthlyBudget: "", + yearlyBudget: "", + capPolicy: defaultCustomerCapPolicy, + }; + onChange([...groups, group], groups.length ? defaultGroupId : group.id); + }; + return ( +
+
+

+ {t("customerPortal.createDialog.groups")} +

+

{t("customerPortal.createDialog.groupsHelp")}

+
+ + onChange(groups, String(value))} + aria-label={t("customerPortal.createDialog.defaultGroup")} + disabled={disabled} + className="gap-1" + > + {groups.map((group, index) => ( +
+
+ + updateGroup(group.id, { name: event.target.value })} + className="font-mono text-xs" + /> +
+ {groupBudgetFields.map(([field, period]) => ( +
+ + updateGroup(group.id, { [field]: value })} + /> +
+ ))} +
+ + +
+ + +
+ ))} + + +

{t("customerPortal.createDialog.defaultHelp")}

+

{t("customerPortal.createDialog.resetHelp")}

+
+ ); +} diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchyManager.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchyManager.tsx index 0a72578a7c7..ad75a9b62f5 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchyManager.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchyManager.tsx @@ -12,6 +12,7 @@ import { extractProxyErrorMessage } from "@/lib/http/client"; import { uiHref } from "@/utils/uiHref"; import CustomerHierarchyNavigation, { customerHierarchyHref } from "./CustomerHierarchyNavigation"; import CustomerHierarchySettings from "./CustomerHierarchySettings"; +import CustomerBudgetGroups from "./CustomerBudgetGroups"; import CustomerBudgetSummary from "./CustomerBudgetSummary"; import type { CustomerUsageEntity } from "./customerUsageApi"; import CreateCustomerHierarchyDialog from "./CreateCustomerHierarchyDialog"; @@ -359,7 +360,14 @@ export default function CustomerHierarchyManager() {
- {tab === "settings" ? ( + {tab === "budget_groups" ? ( + + ) : null} + {tab === "settings" && ( - ) : ( + )} + {(tab === "teams" || tab === "members") && ( )}
diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchyNavigation.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchyNavigation.tsx index 3b32c36e1bd..44e164a0975 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchyNavigation.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchyNavigation.tsx @@ -3,11 +3,12 @@ import { cn } from "@/lib/cva.config"; import { uiHref } from "@/utils/uiHref"; import type { CustomerUsageScope } from "./customerUsageApi"; -type CustomerHierarchyView = "usage" | "teams" | "members" | "settings"; +type CustomerHierarchyView = "usage" | "teams" | "members" | "budget_groups" | "settings"; const LABELS = { usage: "customerPortal.hierarchy.usage", teams: "customerPortal.hierarchy.teams", members: "customerPortal.hierarchy.members", + budget_groups: "customerPortal.budgetGroups.title", settings: "customerPortal.hierarchy.settings", } as const; @@ -26,28 +27,32 @@ export default function CustomerHierarchyNavigation({ active: CustomerHierarchyView; }) { const { t } = useTranslation(); - let pages: readonly CustomerHierarchyView[] = ["usage", "teams", "members", "settings"]; + let pages: readonly CustomerHierarchyView[] = ["usage", "teams", "members", "budget_groups", "settings"]; if (scope.team) pages = ["usage", "members", "settings"]; if (scope.member) pages = ["usage", "settings"]; return ( ); } diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchySettings.integration.test.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchySettings.integration.test.tsx index 5472b1570e4..555b32f3450 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchySettings.integration.test.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchySettings.integration.test.tsx @@ -31,6 +31,7 @@ const audit = { spend: 5, }; const organization: Organization = { + budget_groups_revision: "a".repeat(64), ...audit, customer_org_id: "org-1", alias: "Example organization", diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerBudgetGroupsPlan.test.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerBudgetGroupsPlan.test.ts new file mode 100644 index 00000000000..2381c668b43 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerBudgetGroupsPlan.test.ts @@ -0,0 +1,125 @@ +import { describe, expect, it } from "vitest"; +import type { Organization } from "./customerHierarchyForms"; +import { customerBudgetGroupRows, planCustomerBudgetGroups } from "./customerBudgetGroupsPlan"; + +const organization: Organization = { + customer_org_id: "org-1", + budget_groups_revision: "a".repeat(64), + default_budget_template_id: "standard", + created_at: "", + updated_at: "", + created_by: "", + updated_by: "", + spend: 0, + member_count: 0, + default_budgets: [ + { + budget_id: "standard", + name: "Standard", + budget_limits: [ + { budget_duration: "1mo", max_budget: 10, reset_at: "2027-01-15T00:00:00Z" }, + { budget_duration: "12mo", max_budget: 100, reset_at: "2027-11-15T00:00:00Z" }, + ], + }, + { budget_id: "legacy", name: "Legacy", max_budget: 25, budget_duration: "7d", rpm_limit: 30 }, + { budget_id: "rate", name: "Rate only", tpm_limit: 500 }, + ], +}; + +describe("planCustomerBudgetGroups", () => { + it("should preserve an explicitly unset default while editing existing groups", () => { + const original = { ...organization, default_budget_template_id: null }; + const rows = customerBudgetGroupRows(original); + expect(planCustomerBudgetGroups(original, [{ ...rows[0], name: "Renamed" }, ...rows.slice(1)], "")).toMatchObject({ + ok: true, + body: { default_budget_template: null }, + }); + }); + + it("should omit unchanged rules, including scalar and rate-only templates", () => { + expect(planCustomerBudgetGroups(organization, customerBudgetGroupRows(organization), "standard")).toEqual({ + ok: true, + changed: false, + body: { + expected_revision: organization.budget_groups_revision, + default_budget_template: "Standard", + default_budgets: organization.default_budgets.map(({ budget_id, name }) => ({ budget_id, name })), + }, + }); + }); + + it("should group additions, deletions, name swaps, caps, policies and the default in one request", () => { + const rows = customerBudgetGroupRows(organization); + const result = planCustomerBudgetGroups( + organization, + [ + { ...rows[0], name: "Legacy", monthlyBudget: "0", yearlyBudget: "" }, + { ...rows[1], name: "Standard", capPolicy: "pause_allow_increase" }, + { id: "new", name: " New ", monthlyBudget: "30", yearlyBudget: "300", capPolicy: "pause_until_reset" }, + ], + "new", + ); + expect(result).toEqual({ + ok: true, + changed: true, + body: { + expected_revision: organization.budget_groups_revision, + default_budget_template: "New", + default_budgets: [ + { budget_id: "standard", name: "Legacy", budget_limits: [{ budget_duration: "1mo", max_budget: 0 }] }, + { budget_id: "legacy", name: "Standard", cap_policy: "pause_allow_increase" }, + { + name: "New", + cap_policy: "pause_until_reset", + budget_limits: [ + { budget_duration: "1mo", max_budget: 30 }, + { budget_duration: "12mo", max_budget: 300 }, + ], + }, + ], + }, + }); + }); + + it("should clear existing calendar caps without clearing other budget rules", () => { + const rows = customerBudgetGroupRows(organization); + expect( + planCustomerBudgetGroups( + organization, + [{ ...rows[0], monthlyBudget: "", yearlyBudget: "" }, ...rows.slice(1)], + "legacy", + ), + ).toMatchObject({ + ok: true, + changed: true, + body: { + default_budget_template: "Legacy", + default_budgets: [ + { budget_id: "standard", name: "Standard", budget_limits: [] }, + { budget_id: "legacy", name: "Legacy" }, + { budget_id: "rate", name: "Rate only" }, + ], + }, + }); + }); + + it.each([ + [{ name: "" }, "groupNameRequired"], + [{ name: "Legacy" }, "duplicateGroups"], + [{ monthlyBudget: "-1" }, "invalidBudget"], + [{ monthlyBudget: "Infinity" }, "invalidBudget"], + ])("should reject invalid edits before sending a request", (patch, error) => { + const rows = customerBudgetGroupRows(organization); + expect(planCustomerBudgetGroups(organization, [{ ...rows[0], ...patch }, ...rows.slice(1)], "standard")).toEqual({ + ok: false, + error, + }); + }); + + it("should still require a limit for a newly created group", () => { + const rows = customerBudgetGroupRows(organization); + expect( + planCustomerBudgetGroups(organization, [...rows, { ...rows[1], id: "new", name: "New" }], "standard"), + ).toEqual({ ok: false, error: "groupLimitRequired" }); + }); +}); diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerBudgetGroupsPlan.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerBudgetGroupsPlan.ts new file mode 100644 index 00000000000..3b099ea65a0 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerBudgetGroupsPlan.ts @@ -0,0 +1,58 @@ +import type { components } from "@/lib/http/schema"; +import type { Organization } from "./customerHierarchyForms"; +import { + defaultCustomerCapPolicy, + limit, + validateCustomerBudgetGroups, + type CustomerBudgetGroupFields, +} from "./customerOrganizationForm"; + +export const customerBudgetGroupRows = (organization: Organization): CustomerBudgetGroupFields[] => + organization.default_budgets.map((group) => ({ + id: group.budget_id, + name: group.name, + capPolicy: group.cap_policy ?? defaultCustomerCapPolicy, + monthlyBudget: String(group.budget_limits?.find((entry) => entry.budget_duration === "1mo")?.max_budget ?? ""), + yearlyBudget: String(group.budget_limits?.find((entry) => entry.budget_duration === "12mo")?.max_budget ?? ""), + })); + +export function planCustomerBudgetGroups( + organization: Organization, + groups: CustomerBudgetGroupFields[], + defaultGroupId: string, +) { + const originals = customerBudgetGroupRows(organization); + const validation = validateCustomerBudgetGroups( + { groups, defaultGroupId }, + originals.map((group) => group.id), + ); + if (!validation.ok) return validation; + const default_budgets = groups.map((group) => { + const original = originals.find((row) => row.id === group.id); + const limits = [...limit(group.monthlyBudget, "1mo"), ...limit(group.yearlyBudget, "12mo")]; + const previous = original + ? [...limit(original.monthlyBudget, "1mo"), ...limit(original.yearlyBudget, "12mo")] + : null; + return { + ...(original ? { budget_id: original.id } : {}), + name: group.name.trim(), + ...(group.capPolicy !== original?.capPolicy ? { cap_policy: group.capPolicy } : {}), + ...(JSON.stringify(limits) !== JSON.stringify(previous) ? { budget_limits: limits } : {}), + }; + }); + const body: components["schemas"]["BudgetGroupsUpdate"] = { + expected_revision: organization.budget_groups_revision, + default_budgets, + default_budget_template: groups.find((group) => group.id === defaultGroupId)?.name.trim() ?? null, + }; + const changed = + defaultGroupId !== (organization.default_budget_template_id ?? "") || + groups.length !== originals.length || + default_budgets.some( + (group) => + group.name !== originals.find((original) => original.id === group.budget_id)?.name || + group.cap_policy !== undefined || + group.budget_limits !== undefined, + ); + return { ok: true, body, changed } as const; +} diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerHierarchyForms.test.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerHierarchyForms.test.ts index f7f95ba2054..26d856d44f2 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerHierarchyForms.test.ts +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerHierarchyForms.test.ts @@ -18,6 +18,7 @@ const audit = { spend: 5, }; const organization: Organization = { + budget_groups_revision: "a".repeat(64), ...audit, customer_org_id: "org-1", alias: "Example", diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerOrganizationForm.test.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerOrganizationForm.test.ts index 9d72c0d5602..758aecf285b 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerOrganizationForm.test.ts +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerOrganizationForm.test.ts @@ -1,5 +1,9 @@ import { describe, expect, it } from "vitest"; -import { buildCustomerOrganizationBody, initialCustomerOrganizationFields } from "./customerOrganizationForm"; +import { + buildCustomerOrganizationBody, + defaultCustomerCapPolicy, + initialCustomerOrganizationFields, +} from "./customerOrganizationForm"; const fields = () => ({ ...initialCustomerOrganizationFields(), @@ -14,7 +18,15 @@ describe("buildCustomerOrganizationBody", () => { monthlyBudget: "1000.25", tpm: "3000", rpm: "25", - groups: [{ id: "custom", name: " power ", monthlyBudget: "200.50", yearlyBudget: "1500" }], + groups: [ + { + id: "custom", + name: " power ", + monthlyBudget: "200.50", + yearlyBudget: "1500", + capPolicy: defaultCustomerCapPolicy, + }, + ], defaultGroupId: "custom", }; const result = buildCustomerOrganizationBody(input); @@ -27,6 +39,7 @@ describe("buildCustomerOrganizationBody", () => { default_budgets: [ { name: "power", + cap_policy: "pause_until_reset", budget_limits: [ { budget_duration: "1mo", max_budget: 200.5 }, { budget_duration: "12mo", max_budget: 1500 }, @@ -52,8 +65,8 @@ describe("buildCustomerOrganizationBody", () => { ...fields(), monthlyBudget: "0", groups: [ - { id: "monthly", name: "monthly", monthlyBudget: "0", yearlyBudget: "" }, - { id: "yearly", name: "yearly", monthlyBudget: "", yearlyBudget: "120" }, + { id: "monthly", name: "monthly", monthlyBudget: "0", yearlyBudget: "", capPolicy: defaultCustomerCapPolicy }, + { id: "yearly", name: "yearly", monthlyBudget: "", yearlyBudget: "120", capPolicy: defaultCustomerCapPolicy }, ], defaultGroupId: "yearly", }; diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerOrganizationForm.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerOrganizationForm.ts index 99c67236d92..10b7953503e 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerOrganizationForm.ts +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerOrganizationForm.ts @@ -1,10 +1,14 @@ import type { components } from "@/lib/http/schema"; +export type CustomerCapPolicy = NonNullable; +export const defaultCustomerCapPolicy: CustomerCapPolicy = "pause_until_reset"; + export interface CustomerBudgetGroupFields { id: string; name: string; monthlyBudget: string; yearlyBudget: string; + capPolicy: CustomerCapPolicy; } export interface CustomerOrganizationFields { @@ -24,9 +28,21 @@ export const initialCustomerOrganizationFields = (): CustomerOrganizationFields tpm: "", rpm: "", groups: [ - { id: "power", name: "power_user", monthlyBudget: "200", yearlyBudget: "2000" }, - { id: "default", name: "default_user", monthlyBudget: "50", yearlyBudget: "500" }, - { id: "low", name: "low_user", monthlyBudget: "10", yearlyBudget: "100" }, + { + id: "power", + name: "power_user", + monthlyBudget: "200", + yearlyBudget: "2000", + capPolicy: defaultCustomerCapPolicy, + }, + { + id: "default", + name: "default_user", + monthlyBudget: "50", + yearlyBudget: "500", + capPolicy: defaultCustomerCapPolicy, + }, + { id: "low", name: "low_user", monthlyBudget: "10", yearlyBudget: "100", capPolicy: defaultCustomerCapPolicy }, ], defaultGroupId: "default", }); @@ -40,7 +56,7 @@ const validRate = (value: string): boolean => { const rate = optionalAmount(value); return rate === null || (Number.isSafeInteger(rate) && rate >= 0); }; -const limit = (value: string, duration: "1mo" | "12mo") => { +export const limit = (value: string, duration: "1mo" | "12mo") => { const amount = optionalAmount(value); return amount === null ? [] : [{ budget_duration: duration, max_budget: amount }]; }; @@ -54,23 +70,36 @@ export type CustomerOrganizationFormError = | "invalidRateLimit" | "defaultRequired"; -export function buildCustomerOrganizationBody(fields: CustomerOrganizationFields) { - if (!fields.name.trim()) return { ok: false, error: "nameRequired" } as const; - if (!fields.groups.length || !fields.groups.some((group) => group.id === fields.defaultGroupId)) { +export function validateCustomerBudgetGroups( + fields: Pick, + existingGroupIds: readonly string[] = [], +) { + const missingDefault = !fields.defaultGroupId && !existingGroupIds.length; + const invalidDefault = fields.defaultGroupId && !fields.groups.some((group) => group.id === fields.defaultGroupId); + if (missingDefault || invalidDefault) { return { ok: false, error: "defaultRequired" } as const; } const names = fields.groups.map((group) => group.name.trim()); if (names.some((name) => !name)) return { ok: false, error: "groupNameRequired" } as const; if (new Set(names).size !== names.length) return { ok: false, error: "duplicateGroups" } as const; - if (fields.groups.some((group) => !group.monthlyBudget.trim() && !group.yearlyBudget.trim())) { + if ( + fields.groups.some( + (group) => !existingGroupIds.includes(group.id) && !group.monthlyBudget.trim() && !group.yearlyBudget.trim(), + ) + ) { return { ok: false, error: "groupLimitRequired" } as const; } - if ( - !validAmount(fields.monthlyBudget) || - fields.groups.some((group) => !validAmount(group.monthlyBudget) || !validAmount(group.yearlyBudget)) - ) { + if (fields.groups.some((group) => !validAmount(group.monthlyBudget) || !validAmount(group.yearlyBudget))) { return { ok: false, error: "invalidBudget" } as const; } + return { ok: true } as const; +} + +export function buildCustomerOrganizationBody(fields: CustomerOrganizationFields) { + if (!fields.name.trim()) return { ok: false, error: "nameRequired" } as const; + const validation = validateCustomerBudgetGroups(fields); + if (!validation.ok) return validation; + if (!validAmount(fields.monthlyBudget)) return { ok: false, error: "invalidBudget" } as const; if (!validRate(fields.tpm) || !validRate(fields.rpm)) { return { ok: false, error: "invalidRateLimit" } as const; } @@ -91,6 +120,7 @@ export function buildCustomerOrganizationBody(fields: CustomerOrganizationFields : {}), default_budgets: fields.groups.map((group): components["schemas"]["BudgetTemplate"] => ({ name: group.name.trim(), + cap_policy: group.capPolicy, budget_limits: [...limit(group.monthlyBudget, "1mo"), ...limit(group.yearlyBudget, "12mo")], })), default_budget_template: fields.groups.find((group) => group.id === fields.defaultGroupId)?.name.trim(), diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useCustomerHierarchy.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useCustomerHierarchy.ts index b666dc83692..1155d3b52c8 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useCustomerHierarchy.ts +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useCustomerHierarchy.ts @@ -19,7 +19,7 @@ const MANAGE_PARSERS = { org: parseAsString, team: parseAsString, member: parseAsString, - tab: parseAsStringLiteral(["teams", "members", "settings"] as const), + tab: parseAsStringLiteral(["teams", "members", "budget_groups", "settings"] as const), hierarchyPage: parseAsInteger.withDefault(1), }; @@ -89,12 +89,20 @@ async function customerHierarchyChildren( function childView( selected: CustomerHierarchyContext | undefined, - requested: "teams" | "members" | "settings" | null, + requested: "teams" | "members" | "budget_groups" | "settings" | null, ): ChildList["kind"] { if (selected?.team) return "members"; return requested === "members" ? "members" : "teams"; } +function organizationView( + requested: "teams" | "members" | "budget_groups" | "settings" | null, + team: string | null, + children: ChildList["kind"], +) { + return requested === "budget_groups" && !team ? "budget_groups" : children; +} + export default function useCustomerHierarchy() { const { accessToken, userRole, isLoading } = useAuthorized(); const canManage = isProxyAdminRole(userRole); @@ -113,8 +121,10 @@ export default function useCustomerHierarchy() { const selected = context.data; const childTab = childView(selected, params.tab); const showSettings = params.member !== null || params.tab === "settings"; - const tab: "teams" | "members" | "settings" = showSettings ? "settings" : childTab; + const organizationTab = organizationView(params.tab, params.team, childTab); + const tab: "teams" | "members" | "budget_groups" | "settings" = showSettings ? "settings" : organizationTab; const childContext = selected && !selected.member ? selected : null; + const showChildren = tab === "teams" || tab === "members"; const page = Math.max(1, params.hierarchyPage); const children = useQuery({ queryKey: [ @@ -126,7 +136,7 @@ export default function useCustomerHierarchy() { page, ], queryFn: - enabled && childContext && tab !== "settings" + enabled && childContext && showChildren ? ({ signal }) => customerHierarchyChildren(childContext, childTab, page, signal) : skipToken, }); @@ -164,6 +174,18 @@ export default function useCustomerHierarchy() { { history: "push" }, ); const refresh = () => void queryClient.invalidateQueries({ queryKey: ["customer-hierarchy"] }); + const organizationSaved = async (organization: Organization) => { + const queryKey = ["customer-hierarchy", "manage", organization.customer_org_id, null, null]; + await queryClient.cancelQueries({ queryKey }); + const saved: CustomerHierarchyContext = { + organization, + team: null, + member: null, + entity: { kind: "organization", record: organization }, + }; + queryClient.setQueryData(queryKey, saved); + refresh(); + }; return { scope, context, @@ -179,6 +201,7 @@ export default function useCustomerHierarchy() { openMember, openCreatedScope, refresh, + organizationSaved, setPage: (hierarchyPage: number) => void setParams({ hierarchyPage }), }; } diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/manage/page.integration.test.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/manage/page.integration.test.tsx index e9c51fec78b..754578b43c0 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/manage/page.integration.test.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/manage/page.integration.test.tsx @@ -113,6 +113,94 @@ function BrowserNavigationProbe() { } describe("customer hierarchy management", () => { + it("should preserve a draft on refresh and conflict, then reload and save the current configuration", async () => { + const user = userEvent.setup(); + const standard = { + budget_id: "standard", + name: "Standard", + budget_limits: [{ budget_duration: "1mo", max_budget: 10 }], + }; + const alternate = { budget_id: "alternate", name: "Alternate", max_budget: 20, budget_duration: "7d" }; + const initial = { + ...organization, + budget_groups_revision: "a".repeat(64), + default_budgets: [standard, alternate], + default_budget_template_id: "standard", + }; + const external = { budget_id: "external", name: "Added elsewhere", max_budget: 30 }; + const refreshed = { + ...initial, + budget_groups_revision: "b".repeat(64), + alias: "Refreshed customer", + default_budgets: [...initial.default_budgets, external], + }; + const saved = { + ...refreshed, + budget_groups_revision: "c".repeat(64), + default_budget_template_id: "alternate", + default_budgets: [ + { ...standard, budget_limits: [{ budget_duration: "1mo", max_budget: 15 }] }, + alternate, + external, + ], + }; + fetchMock.mockImplementation(async () => json(initial)); + renderWithProviders(, { searchParams: "?org=org-1&tab=budget_groups" }); + const monthly = await screen.findByRole("spinbutton", { name: "Group 1 monthly budget (USD)" }); + fireEvent.change(monthly, { target: { value: "15" } }); + await user.click(screen.getByRole("radio", { name: "Use Alternate as default" })); + fetchMock.mockImplementation(async () => json(refreshed)); + await act(() => testQueryClient.invalidateQueries({ queryKey: ["customer-hierarchy"] })); + expect(await screen.findByRole("heading", { name: "Refreshed customer" })).toBeInTheDocument(); + expect(screen.getByRole("spinbutton", { name: "Group 1 monthly budget (USD)" })).toHaveValue(15); + expect(screen.getByRole("radio", { name: "Use Alternate as default" })).toBeChecked(); + fetchMock.mockClear(); + fetchMock.mockImplementation(async () => json({ detail: "Budget groups changed. Reload before saving." }, 409)); + await user.click(screen.getByRole("button", { name: "Save changes" })); + expect(await screen.findByRole("alert")).toHaveTextContent("Budget groups changed."); + expect(screen.getByRole("spinbutton", { name: "Group 1 monthly budget (USD)" })).toHaveValue(15); + const request = fetchMock.mock.calls[0][0]; + if (!(request instanceof Request)) throw new Error("Expected HTTP request"); + expect(await request.clone().json()).toMatchObject({ expected_revision: initial.budget_groups_revision }); + fetchMock.mockImplementation(async () => json(refreshed)); + await user.click(screen.getByRole("button", { name: "Discard edits and reload" })); + expect(await screen.findByDisplayValue("Added elsewhere")).toBeInTheDocument(); + expect(screen.getByRole("spinbutton", { name: "Group 1 monthly budget (USD)" })).toHaveValue(10); + fireEvent.change(screen.getByRole("spinbutton", { name: "Group 1 monthly budget (USD)" }), { + target: { value: "15" }, + }); + await user.click(screen.getByRole("radio", { name: "Use Alternate as default" })); + fetchMock.mockClear(); + fetchMock.mockImplementation(async () => json(saved)); + await user.click(screen.getByRole("button", { name: "Save changes" })); + await waitFor(() => expect(screen.getByRole("button", { name: "Save changes" })).toBeDisabled()); + expect(screen.getByRole("spinbutton", { name: "Group 1 monthly budget (USD)" })).toHaveValue(15); + expect(screen.getByRole("radio", { name: "Use Alternate as default" })).toBeChecked(); + const writes = fetchMock.mock.calls + .map(([input]) => input) + .filter((input): input is Request => input instanceof Request && input.method !== "GET"); + expect(writes).toHaveLength(1); + expect(await writes[0].clone().json()).toMatchObject({ + expected_revision: refreshed.budget_groups_revision, + default_budgets: [ + { budget_id: "standard", name: "Standard", budget_limits: [{ budget_duration: "1mo", max_budget: 15 }] }, + { budget_id: "alternate", name: "Alternate" }, + { budget_id: "external", name: "Added elsewhere" }, + ], + }); + }); + + it("should show budget groups only at organization scope without requesting child tables", async () => { + renderWithProviders(, { searchParams: "?org=org-1&tab=budget_groups" }); + expect(await screen.findByRole("link", { name: "Budget groups" })).toHaveAttribute("aria-current", "page"); + expect(screen.getByRole("link", { name: "Budget groups" })).toHaveAttribute( + "href", + `${uiHref("customers/manage")}?org=org-1&tab=budget_groups`, + ); + expect(screen.getByRole("button", { name: "Add budget group" })).toBeEnabled(); + expect(requestedUrls().map((url) => url.pathname)).toEqual(["/management/v1/customer_organizations/org-1"]); + }); + beforeEach(() => { auth.userRole = "Admin"; vi.spyOn(toast, "success").mockReturnValue("test-toast"); diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/page.integration.test.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/page.integration.test.tsx index 343c77cb93f..78dc38b92a7 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/page.integration.test.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/page.integration.test.tsx @@ -224,6 +224,7 @@ describe("customer organization usage", () => { default_budgets: [ { name: "power_user", + cap_policy: "pause_until_reset", budget_limits: [ { budget_duration: "1mo", max_budget: 200 }, { budget_duration: "12mo", max_budget: 1500 }, @@ -231,6 +232,7 @@ describe("customer organization usage", () => { }, { name: "default_user", + cap_policy: "pause_until_reset", budget_limits: [ { budget_duration: "1mo", max_budget: 50 }, { budget_duration: "12mo", max_budget: 500 }, @@ -238,6 +240,7 @@ describe("customer organization usage", () => { }, { name: "low_user", + cap_policy: "pause_until_reset", budget_limits: [ { budget_duration: "1mo", max_budget: 10 }, { budget_duration: "12mo", max_budget: 100 }, diff --git a/litellm/ui/litellm-dashboard/src/i18n/locales/en.json b/litellm/ui/litellm-dashboard/src/i18n/locales/en.json index e3d9a81d48e..3597b390209 100644 --- a/litellm/ui/litellm-dashboard/src/i18n/locales/en.json +++ b/litellm/ui/litellm-dashboard/src/i18n/locales/en.json @@ -20630,6 +20630,14 @@ "loadFailed": "Could not load this model group's deployments. Saved weights are kept." }, "customerPortal": { + "budgetGroups": { + "atCap": "At the cap", + "atCapLabel": "Group {{index}} at the cap", + "pauseUntilReset": "Pause until reset", + "pauseAllowIncrease": "Pause, they can request more", + "title": "Budget groups", + "discard": "Discard edits and reload" + }, "limits": { "byline": "Free tier: up to 5 customers, 5 teams per customer, 5 members per team. LiteLLM Enterprise removes these limits.", "reachedTitle": "Free customer portal limit reached", @@ -20846,7 +20854,7 @@ "removeGroup": "Remove {{name}}", "addGroup": "Add budget group", "defaultHelp": "The selected default group is assigned to new members automatically.", - "resetHelp": "Reaching either cap pauses usage until that budget resets.", + "resetHelp": "Each group's At the cap setting decides what happens when a member reaches either cap.", "advanced": "Advanced limits", "creating": "Creating customer...", "errors": { diff --git a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts index 4548eb989d7..26b57d62e7b 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9894,7 +9894,11 @@ export interface paths { delete?: never; options?: never; head?: never; - patch?: never; + /** + * Configure Customer Budget Groups + * @description Atomically configure budget groups and their default. Requires the current budget_groups_revision. Existing groups retain omitted settings and window reset dates. Stale revisions and referenced deletions return 409 without applying changes. Proxy admins only. + */ + patch: operations["configure_customer_budget_groups_management_v1_customer_organizations__customer_org_id__default_budgets_patch"]; trace?: never; }; "/management/v1/customer_organizations/{customer_org_id}/default_budgets/{budget_id}": { @@ -29635,6 +29639,58 @@ export interface components { /** Id */ id: string; }; + /** BudgetGroupChange */ + BudgetGroupChange: { + /** Allowed Models */ + allowed_models?: string[] | null; + /** Budget Duration */ + budget_duration?: string | null; + /** + * Budget Id + * @description Existing template ID. Omit to create a template. + */ + budget_id?: string | null; + /** Budget Limits */ + budget_limits?: components["schemas"]["BudgetLimit"][] | null; + /** + * Cap Policy + * @description What a paused member can do at the cap: wait for the reset (default) or request a temporary increase + */ + cap_policy?: ("pause_until_reset" | "pause_allow_increase") | null; + /** Max Budget */ + max_budget?: number | null; + /** Model Rpm Limit */ + model_rpm_limit?: { + [key: string]: number; + } | null; + /** Model Tpm Limit */ + model_tpm_limit?: { + [key: string]: number; + } | null; + /** Name */ + name: string; + /** Rpm Limit */ + rpm_limit?: number | null; + /** Soft Budget */ + soft_budget?: number | null; + /** Tpm Limit */ + tpm_limit?: number | null; + }; + /** BudgetGroupsUpdate */ + BudgetGroupsUpdate: { + /** + * Default Budget Template + * @description Final name of the default group, or null for no default. + */ + default_budget_template: string | null; + /** + * Default Budgets + * @description Complete retained group list. Existing groups patch only supplied fields. Omitted groups are deleted. + */ + default_budgets: components["schemas"]["BudgetGroupChange"][]; + /** Expected Revision */ + expected_revision: string; + }; /** BudgetIncreaseRejection */ BudgetIncreaseRejection: { /** Review Notes */ @@ -42738,6 +42794,8 @@ export interface components { /** Alias */ alias?: string | null; budget?: components["schemas"]["Budget"] | null; + /** Budget Groups Revision */ + readonly budget_groups_revision: string; /** Budget Id */ budget_id?: string | null; /** Budget Reset At */ @@ -42938,6 +42996,8 @@ export interface components { /** Alias */ alias?: string | null; budget?: components["schemas"]["Budget"] | null; + /** Budget Groups Revision */ + readonly budget_groups_revision: string; /** Budget Id */ budget_id?: string | null; /** Budget Reset At */ @@ -69359,6 +69419,41 @@ export interface operations { }; }; }; + configure_customer_budget_groups_management_v1_customer_organizations__customer_org_id__default_budgets_patch: { + parameters: { + query?: never; + header?: never; + path: { + customer_org_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["BudgetGroupsUpdate"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Organization"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; delete_customer_budget_template_management_v1_customer_organizations__customer_org_id__default_budgets__budget_id__delete: { parameters: { query?: never; diff --git a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py index ec8dcfd67b9..2b9eb2a8164 100644 --- a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py +++ b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py @@ -1,4 +1,5 @@ import asyncio +from collections.abc import Mapping from datetime import timedelta from typing import TYPE_CHECKING, Final, cast @@ -21,6 +22,7 @@ from tests.unit.proxy.management_helpers.test_customer_hierarchy import ( Database, DeleteConflictTable, Record, + Table, ) if TYPE_CHECKING: @@ -1624,6 +1626,198 @@ def test_member_in_two_teams_is_listed_under_each_and_teams_from_another_organiz assert fetched.json()["member_count"] == 1 +def test_grouped_configuration_preserves_settings_swaps_names_and_selects_a_new_default() -> None: + http: Final = client() + org: Final = family(http) + first, second = org["default_budgets"] + response: Final = http.patch( + "/management/v1/customer_organizations/org/default_budgets", + json={ + "expected_revision": org["budget_groups_revision"], + "default_budgets": [ + {"budget_id": first["budget_id"], "name": second["name"], "cap_policy": "pause_allow_increase"}, + {"budget_id": second["budget_id"], "name": first["name"]}, + {"name": "new", "budget_limits": [{"budget_duration": "1mo", "max_budget": 0}]}, + ], + "default_budget_template": "new", + }, + ) + assert response.status_code == 200, response.text + saved: Final = response.json() + assert saved["default_budgets"][0] == {**first, "name": second["name"], "cap_policy": "pause_allow_increase"} + assert saved["default_budgets"][1] == {**second, "name": first["name"]} + assert saved["default_budget_template_id"] == saved["default_budgets"][2]["budget_id"] + assert saved["budget_groups_revision"] != org["budget_groups_revision"] + assert http.get("/management/v1/customer_organizations/org").json() == saved + + +def test_grouped_configuration_rejects_stale_edits_without_losing_external_groups() -> None: + http: Final = client() + org: Final = family(http) + external: Final = http.post("/management/v1/customer_organizations/org/default_budgets", json={"name": "external"}) + assert external.status_code == 201, external.text + response: Final = http.patch( + "/management/v1/customer_organizations/org/default_budgets", + json={ + "expected_revision": org["budget_groups_revision"], + "default_budgets": [{"budget_id": row["budget_id"], "name": row["name"]} for row in org["default_budgets"]], + "default_budget_template": "standard", + }, + ) + assert response.status_code == 409, response.text + assert response.json()["type"].endswith("configuration-conflict") + assert http.get("/management/v1/customer_organizations/org").json() == external.json() + + +def test_grouped_configuration_referenced_deletion_rejects_the_whole_save() -> None: + http: Final = client() + org: Final = family(http) + member: Final = http.post( + "/management/v1/customer_members", json={"customer_member_id": "m", "parent_team_ids": ["team"]} + ) + assert member.status_code == 201, member.text + before: Final = http.get("/management/v1/customer_organizations/org").json() + response: Final = http.patch( + "/management/v1/customer_organizations/org/default_budgets", + json={ + "expected_revision": org["budget_groups_revision"], + "default_budgets": [ + {"budget_id": org["default_budgets"][1]["budget_id"], "name": "changed", "max_budget": 999}, + {"name": "new", "max_budget": 1}, + ], + "default_budget_template": "new", + }, + ) + assert response.status_code == 409, response.text + assert response.json()["type"].endswith("template-in-use") + assert http.get("/management/v1/customer_organizations/org").json() == before + + +def test_grouped_configuration_keeps_window_clocks_and_invalidates_member_cache() -> None: + cache: Final = UserApiKeyCache() + http: Final = client(cache=cache) + org: Final = family(http) + member: Final = http.post( + "/management/v1/customer_members", json={"customer_member_id": "m", "parent_team_ids": ["team"]} + ) + assert member.status_code == 201, member.text + first: Final = org["default_budgets"][0] + windows: Final = http.patch( + f"/management/v1/customer_organizations/org/default_budgets/{first['budget_id']}", + json={ + "budget_limits": [ + {"budget_duration": "1mo", "max_budget": 10, "reset_at": "2027-01-15T00:00:00Z"}, + {"budget_duration": "12mo", "max_budget": 100, "reset_at": "2027-11-15T00:00:00Z"}, + ] + }, + ) + assert windows.status_code == 200, windows.text + advanced: Final = http.patch( + f'/management/v1/customer_organizations/org/default_budgets/{first["budget_id"]}', + json={ + "budget_limits": [ + {**row, "reset_at": "2028-01-15T00:00:00Z"} + for row in windows.json()["default_budgets"][0]["budget_limits"] + ] + }, + ) + assert advanced.status_code == 200, advanced.text + assert advanced.json()["budget_groups_revision"] == windows.json()["budget_groups_revision"] + cache.set_cache(customer_member_cache_key("m"), "old-budget") + saved: Final = http.patch( + "/management/v1/customer_organizations/org/default_budgets", + json={ + "expected_revision": windows.json()["budget_groups_revision"], + "default_budgets": [ + { + "budget_id": first["budget_id"], + "name": "standard", + "budget_limits": [ + {"budget_duration": "1mo", "max_budget": 15}, + {"budget_duration": "12mo", "max_budget": 100}, + ], + } + ], + "default_budget_template": "standard", + }, + ) + assert saved.status_code == 200, saved.text + previous: Final = advanced.json()["default_budgets"][0] + assert saved.json()["default_budgets"] == [ + { + **previous, + "budget_limits": [ + {**previous["budget_limits"][0], "max_budget": 15}, + previous["budget_limits"][1], + ], + } + ] + assert cache.get_cache(customer_member_cache_key("m")) is None + assert ( + http.get("/management/v1/customer_members/m").json()["effective_budget"]["budget_limits"][0]["max_budget"] == 15 + ) + + +@pytest.mark.parametrize("case", ("duplicate_id", "foreign_id", "duplicate_name", "missing_default")) +def test_invalid_grouped_configuration_leaves_existing_groups_untouched(case: str) -> None: + http: Final = client() + org: Final = family(http) + first, second = org["default_budgets"] + groups: Final = [ + {"budget_id": first["budget_id"], "name": "changed", "max_budget": 999}, + { + "budget_id": first["budget_id"] + if case == "duplicate_id" + else "foreign" + if case == "foreign_id" + else second["budget_id"], + "name": "changed" if case == "duplicate_name" else "second", + }, + ] + response: Final = http.patch( + "/management/v1/customer_organizations/org/default_budgets", + json={ + "expected_revision": org["budget_groups_revision"], + "default_budgets": groups, + "default_budget_template": "missing" if case == "missing_default" else "changed", + }, + ) + assert response.status_code == 422, response.text + assert http.get("/management/v1/customer_organizations/org").json() == org + + +class FailingBudgetInsertTable(Table): + async def create(self, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> Record: + from prisma.errors import UniqueViolationError + + if data.get("name") == "failed-insert": + raise UniqueViolationError({}) + return await super().create(data, include) + + +def test_grouped_configuration_rolls_back_earlier_writes_when_a_later_database_write_fails() -> None: + database: Final = Database() + database.litellm_budgettable = FailingBudgetInsertTable("budget_id", database) + http: Final = client(database=database) + org: Final = family(http) + before_budgets: Final = dict(database.litellm_budgettable.rows) + response: Final = http.patch( + "/management/v1/customer_organizations/org/default_budgets", + json={ + "expected_revision": org["budget_groups_revision"], + "default_budgets": [ + {"budget_id": org["default_budget_template_id"], "name": "changed", "max_budget": 999}, + {"name": "created-first", "max_budget": 1}, + {"name": "failed-insert", "max_budget": 2}, + ], + "default_budget_template": "created-first", + }, + ) + assert response.status_code == 409, response.text + assert http.get("/management/v1/customer_organizations/org").json() == org + assert database.litellm_budgettable.rows == before_budgets + + def test_list_and_detail_expose_each_effective_budget_with_its_own_recorded_spend() -> None: database: Final = Database() http: Final = client(database=database)