scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

feat(customer-portal): configure budget groups atomically (#51661)

Browse source 
* 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>
This commit is contained in:
tin-berri 2026-11-13 12:29:36 -08:00 • committed by GitHub
parent 534ef43526
commit b798cf30ab
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 1297 additions and 199 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
litellm/proxy/management_endpoints/management_v1/customers.py
View file

@ -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,

81
litellm/proxy/management_helpers/customer_hierarchy.py
View file

@ -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:

39
litellm/types/proxy/management_endpoints/customer_hierarchy.py
View file

@ -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

8
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.integration.test.tsx
View file

@ -73,6 +73,8 @@ describe("CreateCustomerDialog", () => {
const user = userEvent.setup();
renderWithProviders(<Harness />);
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",
});
});

165
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.tsx
View file

@ -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 (
<InputGroup>
<InputGroupAddon>
<InputGroupText>$</InputGroupText>
</InputGroupAddon>
<InputGroupInput
id={id}
aria-label={label}
type="number"
min={0}
step="any"
value={value}
onChange={(event) => onChange(event.target.value)}
placeholder={t("customerPortal.createDialog.noLimit")}
/>
</InputGroup>
);
}
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<Omit<CustomerBudgetGroupFields, "id">>) => {
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 }:
/>
<p className="text-xs text-muted-foreground">{t("customerPortal.createDialog.budgetHelp")}</p>
</div>
<section className="space-y-3 border-t pt-5" aria-labelledby="cc-groups-title">
<div>
<h3 id="cc-groups-title" className="font-medium">
{t("customerPortal.createDialog.groups")}
</h3>
<p className="mt-1 text-xs text-muted-foreground">{t("customerPortal.createDialog.groupsHelp")}</p>
</div>
<div
className="hidden grid-cols-[1.3fr_1fr_1fr_3rem_2rem] gap-3 px-2 text-xs text-muted-foreground sm:grid"
aria-hidden="true"
>
<span>{t("customerPortal.createDialog.groupName")}</span>
<span>{t("customerPortal.createDialog.monthly")}</span>
<span>{t("customerPortal.createDialog.yearly")}</span>
<span>{t("customerPortal.createDialog.default")}</span>
</div>
<RadioGroup
value={fields.defaultGroupId}
onValueChange={(defaultGroupId) => setFields({ ...fields, defaultGroupId: String(defaultGroupId) })}
aria-label={t("customerPortal.createDialog.defaultGroup")}
disabled={create.isPending}
className="gap-1"
>
{fields.groups.map((group, index) => (
<div
key={group.id}
className={`grid grid-cols-2 items-center gap-3 rounded-md p-2 sm:grid-cols-[1.3fr_1fr_1fr_3rem_2rem] ${group.id === fields.defaultGroupId ? "bg-muted/60" : ""}`}
>
<div className="col-span-2 min-w-0 space-y-1.5 sm:col-span-1">
<Label htmlFor={`cc-group-${group.id}`} className="sm:sr-only">
{t("customerPortal.createDialog.groupNameLabel", { index: index + 1 })}
</Label>
<Input
id={`cc-group-${group.id}`}
value={group.name}
onChange={(event) => updateGroup(group.id, { name: event.target.value })}
className="font-mono text-xs"
/>
</div>
{groupBudgetFields.map(([field, period]) => (
<div className="min-w-0 space-y-1.5" key={field}>
<Label htmlFor={`cc-${period}-${group.id}`} className="sm:sr-only">
{t(`customerPortal.createDialog.${period}`)}
</Label>
<BudgetInput
id={`cc-${period}-${group.id}`}
label={t(`customerPortal.createDialog.${period}Label`, { index: index + 1 })}
value={group[field]}
onChange={(value) => updateGroup(group.id, { [field]: value })}
/>
</div>
))}
<Label className="flex gap-2 sm:justify-center">
<RadioGroupItem
value={group.id}
aria-label={t("customerPortal.createDialog.useDefault", { name: group.name || index + 1 })}
/>
<span className="text-xs sm:sr-only" aria-hidden="true">
{t("customerPortal.createDialog.default")}
</span>
</Label>
<Button
type="button"
variant="ghost"
size="icon-sm"
className="justify-self-end"
disabled={fields.groups.length === 1}
onClick={() => removeGroup(group.id)}
aria-label={t("customerPortal.createDialog.removeGroup", { name: group.name || index + 1 })}
>
<Trash2 className="size-3.5" />
</Button>
</div>
))}
</RadioGroup>
<Button
type="button"
variant="ghost"
size="sm"
onClick={() =>
setFields({
...fields,
groups: [
...fields.groups,
{ id: crypto.randomUUID(), name: "", monthlyBudget: "", yearlyBudget: "" },
],
})
}
>
<Plus className="size-4" />
{t("customerPortal.createDialog.addGroup")}
</Button>
<p className="text-xs text-muted-foreground">{t("customerPortal.createDialog.defaultHelp")}</p>
<p className="text-xs text-muted-foreground">{t("customerPortal.createDialog.resetHelp")}</p>
</section>
<CustomerBudgetGroupsEditor
groups={fields.groups}
defaultGroupId={fields.defaultGroupId}
disabled={create.isPending}
onChange={(groups, defaultGroupId) => {
setFields({ ...fields, groups, defaultGroupId });
setError(null);
}}
/>
<details className="group border-t pt-4">
<summary className="flex cursor-pointer list-none items-center gap-1 text-sm font-medium">
<ChevronRight className="size-4 group-open:rotate-90" />

1
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerHierarchyDialog.test.tsx
View file

@ -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",

136
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroups.integration.test.tsx Normal file
View file

@ -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<typeof fetch>();
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(<CustomerBudgetGroups organization={organization} onSaved={onSaved} />);
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(<CustomerBudgetGroups organization={organization} onSaved={onSaved} />);
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(
<CustomerBudgetGroups
organization={{ ...organization, default_budgets: [], default_budget_template_id: null }}
onSaved={onSaved}
/>,
);
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(
<CustomerBudgetGroups
organization={{ ...organization, default_budgets: [], default_budget_template_id: null }}
onSaved={onSaved}
/>,
);
expect(screen.queryByRole("textbox")).not.toBeInTheDocument();
expect(screen.getByRole("button", { name: "Add budget group" })).toBeEnabled();
expect(screen.getByRole("button", { name: "Save changes" })).toBeDisabled();
});
});

107
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroups.tsx Normal file
View file

@ -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<void>;
}) {
const { t } = useTranslation();
const [draft, setDraft] = useState<Draft | null>(null);
const [pending, setPending] = useState(false);
const [error, setError] = useState<string | null>(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 (
<form onSubmit={save} className="max-w-3xl space-y-5">
<fieldset disabled={pending}>
<CustomerBudgetGroupsEditor
groups={groups}
defaultGroupId={defaultGroupId}
disabled={pending}
onChange={changeGroups}
/>
</fieldset>
{error && (
<p role="alert" className="text-sm text-destructive">
{error}
</p>
)}
<div className="flex gap-3">
<Button type="submit" disabled={pending || !draft || unchanged}>
{t("customerPortal.hierarchy.save")}
</Button>
{draft && (
<Button type="button" variant="outline" disabled={pending} onClick={() => void reload()}>
{t("customerPortal.budgetGroups.discard")}
</Button>
)}
</div>
</form>
);
}

194
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBudgetGroupsEditor.tsx Normal file
View file

@ -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 (
<InputGroup>
<InputGroupAddon>
<InputGroupText>$</InputGroupText>
</InputGroupAddon>
<InputGroupInput
id={id}
aria-label={label}
type="number"
min={0}
step="any"
value={value}
onChange={(event) => onChange(event.target.value)}
placeholder={t("customerPortal.createDialog.noLimit")}
/>
</InputGroup>
);
}
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<Omit<CustomerBudgetGroupFields, "id">>) =>
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 (
<section className="space-y-3 border-t pt-5" aria-labelledby="cc-groups-title">
<div>
<h3 id="cc-groups-title" className="font-medium">
{t("customerPortal.createDialog.groups")}
</h3>
<p className="mt-1 text-xs text-muted-foreground">{t("customerPortal.createDialog.groupsHelp")}</p>
</div>
<div
className="hidden grid-cols-[1.3fr_1fr_1fr_3rem_2rem] gap-3 px-2 text-xs text-muted-foreground sm:grid"
aria-hidden="true"
>
<span>{t("customerPortal.createDialog.groupName")}</span>
<span>{t("customerPortal.createDialog.monthly")}</span>
<span>{t("customerPortal.createDialog.yearly")}</span>
<span>{t("customerPortal.createDialog.default")}</span>
</div>
<RadioGroup
value={defaultGroupId}
onValueChange={(value) => onChange(groups, String(value))}
aria-label={t("customerPortal.createDialog.defaultGroup")}
disabled={disabled}
className="gap-1"
>
{groups.map((group, index) => (
<div
key={group.id}
className={`grid grid-cols-2 items-center gap-3 rounded-md p-2 sm:grid-cols-[1.3fr_1fr_1fr_3rem_2rem] ${group.id === defaultGroupId ? "bg-muted/60" : ""}`}
>
<div className="col-span-2 min-w-0 space-y-1.5 sm:col-span-1">
<Label htmlFor={`cc-group-${group.id}`} className="sm:sr-only">
{t("customerPortal.createDialog.groupNameLabel", { index: index + 1 })}
</Label>
<Input
id={`cc-group-${group.id}`}
value={group.name}
onChange={(event) => updateGroup(group.id, { name: event.target.value })}
className="font-mono text-xs"
/>
</div>
{groupBudgetFields.map(([field, period]) => (
<div className="min-w-0 space-y-1.5" key={field}>
<Label htmlFor={`cc-${period}-${group.id}`} className="sm:sr-only">
{t(`customerPortal.createDialog.${period}`)}
</Label>
<BudgetInput
id={`cc-${period}-${group.id}`}
label={t(`customerPortal.createDialog.${period}Label`, { index: index + 1 })}
value={group[field]}
onChange={(value) => updateGroup(group.id, { [field]: value })}
/>
</div>
))}
<div className="col-span-2 min-w-0 space-y-1.5 sm:col-span-3 sm:row-start-2">
<Label htmlFor={`cc-cap-${group.id}`}>{t("customerPortal.budgetGroups.atCap")}</Label>
<Select
value={group.capPolicy}
disabled={disabled}
onValueChange={(value) => {
if (value === "pause_until_reset" || value === "pause_allow_increase") {
updateGroup(group.id, { capPolicy: value });
}
}}
>
<SelectTrigger
id={`cc-cap-${group.id}`}
aria-label={t("customerPortal.budgetGroups.atCapLabel", { index: index + 1 })}
className="w-full min-w-0 text-xs"
>
<SelectValue>
{t(
group.capPolicy === "pause_allow_increase"
? "customerPortal.budgetGroups.pauseAllowIncrease"
: "customerPortal.budgetGroups.pauseUntilReset",
)}
</SelectValue>
</SelectTrigger>
<SelectContent>
<SelectItem value="pause_until_reset">{t("customerPortal.budgetGroups.pauseUntilReset")}</SelectItem>
<SelectItem value="pause_allow_increase">
{t("customerPortal.budgetGroups.pauseAllowIncrease")}
</SelectItem>
</SelectContent>
</Select>
</div>
<Label className="flex gap-2 sm:col-start-4 sm:row-start-1 sm:justify-center">
<RadioGroupItem
value={group.id}
aria-label={t("customerPortal.createDialog.useDefault", { name: group.name || index + 1 })}
/>
<span className="text-xs sm:sr-only" aria-hidden="true">
{t("customerPortal.createDialog.default")}
</span>
</Label>
<Button
type="button"
variant="ghost"
size="icon-sm"
className="justify-self-end sm:col-start-5 sm:row-start-1"
disabled={groups.length === 1}
onClick={() => removeGroup(group.id)}
aria-label={t("customerPortal.createDialog.removeGroup", { name: group.name || index + 1 })}
>
<Trash2 className="size-3.5" />
</Button>
</div>
))}
</RadioGroup>
<Button type="button" variant="ghost" size="sm" onClick={addGroup}>
<Plus className="size-4" />
{t("customerPortal.createDialog.addGroup")}
</Button>
<p className="text-xs text-muted-foreground">{t("customerPortal.createDialog.defaultHelp")}</p>
<p className="text-xs text-muted-foreground">{t("customerPortal.createDialog.resetHelp")}</p>
</section>
);
}

13
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchyManager.tsx
View file

@ -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() {
<HierarchyHeader hierarchy={hierarchy} selected={selected} />
<CustomerHierarchyNavigation scope={currentScope} active={tab} />
<section>
{tab === "settings" ? (
{tab === "budget_groups" ? (
<CustomerBudgetGroups
key={scope.org}
organization={selected.organization}
onSaved={hierarchy.organizationSaved}
/>
) : null}
{tab === "settings" && (
<CustomerHierarchySettings
key={customerEntityId(entity)}
organization={selected.organization}
@ -367,7 +375,8 @@ export default function CustomerHierarchyManager() {
onSaved={hierarchy.refresh}
onDeleted={onDeleted}
/>
) : (
)}
{(tab === "teams" || tab === "members") && (
<HierarchyChildTable hierarchy={hierarchy} selected={selected} scope={scope} />
)}
</section>

43
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchyNavigation.tsx
View file

@ -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 (
<nav aria-label={t("customerPortal.hierarchy.customerNavigation")} className="flex gap-6 border-b">
{pages.map((page) => (
<a
key={page}
href={
page === "usage"
? customerHierarchyHref("usage", scope)
: `${customerHierarchyHref("manage", scope)}${page === "settings" || (page === "members" && !scope.team) ? `&tab=${page}` : ""}`
}
aria-current={active === page ? "page" : undefined}
className={cn(
"border-b-2 border-transparent px-0.5 pb-3 text-sm font-medium text-muted-foreground hover:text-foreground focus-visible:outline-2 focus-visible:outline-ring",
active === page && "border-primary text-foreground",
)}
>
{t(LABELS[page])}
</a>
))}
{pages.map((page) => {
const organizationMembers = page === "members" && !scope.team;
const explicitTab = page === "settings" || page === "budget_groups" || organizationMembers;
return (
<a
key={page}
href={
page === "usage"
? customerHierarchyHref("usage", scope)
: `${customerHierarchyHref("manage", scope)}${explicitTab ? `&tab=${page}` : ""}`
}
aria-current={active === page ? "page" : undefined}
className={cn(
"border-b-2 border-transparent px-0.5 pb-3 text-sm font-medium text-muted-foreground hover:text-foreground focus-visible:outline-2 focus-visible:outline-ring",
active === page && "border-primary text-foreground",
)}
>
{t(LABELS[page])}
</a>
);
})}
</nav>
);
}

1
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchySettings.integration.test.tsx
View file

@ -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",

125
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerBudgetGroupsPlan.test.ts Normal file
View file

@ -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" });
});
});

58
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerBudgetGroupsPlan.ts Normal file
View file

@ -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;
}

1
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerHierarchyForms.test.ts
View file

@ -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",

21
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerOrganizationForm.test.ts
View file

@ -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",
};

54
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerOrganizationForm.ts
View file

@ -1,10 +1,14 @@
import type { components } from "@/lib/http/schema";
export type CustomerCapPolicy = NonNullable<components["schemas"]["BudgetTemplate"]["cap_policy"]>;
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<CustomerOrganizationFields, "groups" | "defaultGroupId">,
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(),

31
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useCustomerHierarchy.ts
View file

@ -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 }),
};
}

88
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/manage/page.integration.test.tsx
View file

@ -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(<CustomerHierarchyManager />, { 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(<CustomerHierarchyManager />, { 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");

3
litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/page.integration.test.tsx
View file

@ -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 },

10
litellm/ui/litellm-dashboard/src/i18n/locales/en.json
View file

@ -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>Enterprise</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": {

97
litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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;

194
tests/unit/proxy/management_endpoints/management_v1/test_customers.py
View file

@ -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)