diff --git a/litellm/proxy/auth/customer_hierarchy_auth.py b/litellm/proxy/auth/customer_hierarchy_auth.py index cf566bc6300..6c6bac924ce 100644 --- a/litellm/proxy/auth/customer_hierarchy_auth.py +++ b/litellm/proxy/auth/customer_hierarchy_auth.py @@ -29,6 +29,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_ from litellm.proxy.common_utils.customer_usage import CustomerIdentifier from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl from litellm.proxy.management_helpers.customer_hierarchy import ( + ErrorCode, Failure, PrismaRecord, customer_row_reset_at, @@ -552,6 +553,13 @@ async def _register_organization( _check_key_scope(customer_org_alias, key_metadata) registered: Final = await register(customer_org_alias, prisma_client) if isinstance(registered, Failure): + if registered.code is ErrorCode.FREE_TIER_LIMIT: + verbose_proxy_logger.warning( + "CUSTOMER TRACKING OFF for organization %s because the free customer portal limit was reached. " + "Upgrade to Enterprise to remove the limits: https://www.litellm.ai/enterprise", + customer_org_alias, + ) + return None verbose_proxy_logger.warning( "Could not register customer organization %s: %s", customer_org_alias, registered.code.value ) @@ -612,6 +620,13 @@ async def _register_member( return None registered: Final = await register(organization.customer_org_id, customer_member_id, prisma_client) if isinstance(registered, Failure): + if registered.code is ErrorCode.FREE_TIER_LIMIT: + verbose_proxy_logger.warning( + "CUSTOMER TRACKING OFF for member %s because the free customer portal limit was reached. " + "Upgrade to Enterprise to remove the limits: https://www.litellm.ai/enterprise", + customer_member_id, + ) + return None verbose_proxy_logger.warning( "Could not register customer member %s in organization %s: %s", customer_member_id, @@ -681,7 +696,7 @@ async def load_customer_hierarchy( f"aliased {identity.customer_org_alias}", CUSTOMER_ORG_ALIAS_HEADER, ) - if identity.customer_team_id is None: + if identity.customer_team_id is None or hierarchy.customer_member_id is None: return hierarchy selected: Final = select_customer_team(hierarchy, identity.customer_team_id) if selected is None: diff --git a/litellm/proxy/auth/customer_member_auto_create.py b/litellm/proxy/auth/customer_member_auto_create.py index 37353d990a1..ee176bb0cf7 100644 --- a/litellm/proxy/auth/customer_member_auto_create.py +++ b/litellm/proxy/auth/customer_member_auto_create.py @@ -51,7 +51,9 @@ async def ensure_customer_member( ) -> CustomerMember | Failure: """The member row for ``customer_member_id`` in ``customer_org_id``'s default team, created when absent. Two workers registering the same member at once both end up with the one row.""" - hierarchy: Final = CustomerHierarchy(prisma_client) + from litellm.proxy.proxy_server import premium_user + + hierarchy: Final = CustomerHierarchy(prisma_client, premium_user=premium_user) team: Final = await _default_team(hierarchy, customer_org_id) if isinstance(team, Failure): return team @@ -78,7 +80,9 @@ async def ensure_customer_organization( """The organization row carrying ``customer_org_alias``, created with the alias as its id when absent, so a customer's first request is charged to it. Two workers registering the same alias at once both end up with the one row.""" - hierarchy: Final = CustomerHierarchy(prisma_client) + from litellm.proxy.proxy_server import premium_user + + hierarchy: Final = CustomerHierarchy(prisma_client, premium_user=premium_user) created: Final = await hierarchy.create_organization( OrganizationCreate(customer_org_id=customer_org_alias, alias=customer_org_alias), AUTO_CREATE_ORG_ACTOR ) diff --git a/litellm/proxy/management_endpoints/management_v1/customers.py b/litellm/proxy/management_endpoints/management_v1/customers.py index 3d978e2925c..2e323dcd3da 100644 --- a/litellm/proxy/management_endpoints/management_v1/customers.py +++ b/litellm/proxy/management_endpoints/management_v1/customers.py @@ -16,7 +16,13 @@ from litellm.proxy.list_api.common import ( reject_unknown_query_params, ) from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX -from litellm.proxy.management_helpers.customer_hierarchy import CustomerHierarchy, ErrorCode, Failure, Page +from litellm.proxy.management_helpers.customer_hierarchy import ( + CustomerHierarchy, + CustomerPortalLimits, + ErrorCode, + Failure, + Page, +) from litellm.proxy.management_helpers.customer_upsert import CustomerUpsert from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, # pyright: ignore[reportUnknownVariableType] # legacy decorator is untyped @@ -45,6 +51,13 @@ T: Final = TypeVar("T") _ERROR_DETAILS: Final[Mapping[ErrorCode, tuple[int, str, str]]] = MappingProxyType( { + ErrorCode.FREE_TIER_LIMIT: ( + 403, + "Free customer portal limit reached", + "Free customer portal limit reached: up to 5 customer organizations, 5 teams per organization and " + "5 members per team without a LiteLLM Enterprise license. Upgrade to Enterprise to remove the limits: " + "https://www.litellm.ai/enterprise", + ), ErrorCode.NOT_FOUND: (404, "Not found", "Customer entity not found."), ErrorCode.PARENT_NOT_FOUND: (404, "Parent not found", "Customer parent not found."), ErrorCode.ALREADY_EXISTS: (409, "Already exists", "Customer entity ID already exists."), @@ -131,11 +144,11 @@ def customer_actor(auth: Annotated[UserAPIKeyAuth, Depends(require_customer_admi def get_customer_hierarchy() -> CustomerHierarchy: - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import premium_user, prisma_client if prisma_client is None: raise_public(Failure(ErrorCode.DATABASE_UNAVAILABLE)) - return CustomerHierarchy(prisma_client) + return CustomerHierarchy(prisma_client, premium_user=premium_user) def resolve(result: T | Failure) -> T: @@ -186,6 +199,18 @@ def get_user_api_key_cache() -> UserApiKeyCache: AuthCache: TypeAlias = Annotated[UserApiKeyCache, Depends(get_user_api_key_cache)] +@router.get( + "/customer_portal/limits", + response_model=CustomerPortalLimits, + description="Get customer portal free-tier limits and current capacity. Proxy admins only.", +) +@management_endpoint_wrapper +async def get_customer_portal_limits( + hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)], +) -> CustomerPortalLimits: + return resolve(await hierarchy.portal_limits()) + + @router.post( "/customer_organizations", status_code=201, diff --git a/litellm/proxy/management_helpers/customer_hierarchy.py b/litellm/proxy/management_helpers/customer_hierarchy.py index d3ef6ffb6ac..631f338c32f 100644 --- a/litellm/proxy/management_helpers/customer_hierarchy.py +++ b/litellm/proxy/management_helpers/customer_hierarchy.py @@ -41,6 +41,7 @@ if TYPE_CHECKING: class ErrorCode(str, Enum): + FREE_TIER_LIMIT = "free_tier_limit" NOT_FOUND = "not_found" PARENT_NOT_FOUND = "parent_not_found" ALREADY_EXISTS = "already_exists" @@ -66,6 +67,30 @@ class Failure: code: ErrorCode +FREE_ORGANIZATION_LIMIT: Final = 5 +FREE_TEAMS_PER_ORGANIZATION_LIMIT: Final = 5 +FREE_MEMBERS_PER_TEAM_LIMIT: Final = 5 +FREE_CAPACITY_ADVISORY_LOCK: Final = 714023581 + + +class CustomerPortalLimits(BaseModel): + model_config = ConfigDict(frozen=True) + + limits_apply: bool + organization_limit: int = FREE_ORGANIZATION_LIMIT + teams_per_organization_limit: int = FREE_TEAMS_PER_ORGANIZATION_LIMIT + members_per_team_limit: int = FREE_MEMBERS_PER_TEAM_LIMIT + organization_count: int + any_limit_reached: bool + + +class _CapacityCounts(BaseModel): + model_config = ConfigDict(frozen=True) + + full_organizations: int + full_teams: int + + _LOCK_ORGANIZATION_SQL: Final = ( 'SELECT 1 AS locked FROM "LiteLLM_CustomerOrganizationTable" WHERE customer_org_id = $1 FOR UPDATE' ) @@ -436,8 +461,11 @@ def _member_ids(rows: Sequence[PrismaRecord]) -> tuple[str, ...]: class CustomerHierarchyStore: - def __init__(self, db: "Prisma", now: datetime, settings: BudgetResetSettings | None = None) -> None: + def __init__( + self, db: "Prisma", now: datetime, settings: BudgetResetSettings | None = None, *, premium_user: bool + ) -> None: self.db = db + self.premium_user = premium_user self.now = now self.settings = settings if settings is not None else get_budget_reset_settings() self.organizations: TableActions[PrismaRecord] = db.litellm_customerorganizationtable # pyright: ignore[reportAttributeAccessIssue] # protocol accepts immutable query mappings @@ -445,6 +473,41 @@ class CustomerHierarchyStore: self.members: TableActions[PrismaRecord] = db.litellm_customermembertable # pyright: ignore[reportAttributeAccessIssue] # protocol accepts immutable query mappings self.budgets: TableActions[PrismaRecord] = db.litellm_budgettable # pyright: ignore[reportAttributeAccessIssue] # protocol accepts immutable query mappings + async def lock_free_capacity(self) -> None: + await self.db.query_raw("SELECT 1 AS locked FROM pg_advisory_xact_lock($1)", FREE_CAPACITY_ADVISORY_LOCK) + + async def member_capacity_reached(self, team_ids: Sequence[str]) -> bool: + if self.premium_user or not team_ids: + return False + await self.lock_free_capacity() + for team_id in team_ids: + if await self.members.count(where={"parent_team_ids": {"has": team_id}}) >= FREE_MEMBERS_PER_TEAM_LIMIT: + return True + return False + + async def portal_limits(self) -> CustomerPortalLimits: + organization_count: Final = await self.organizations.count() + counts: Final = TypeAdapter(tuple[_CapacityCounts, ...]).validate_python( + await self.db.query_raw( + 'SELECT (SELECT COUNT(*) FROM (SELECT parent_org_id FROM "LiteLLM_CustomerTeamTable" ' + "GROUP BY parent_org_id HAVING COUNT(*) >= $1) AS full_orgs) AS full_organizations, " + '(SELECT COUNT(*) FROM (SELECT team_id FROM "LiteLLM_CustomerMemberTable", ' + "unnest(parent_team_ids) AS team_id GROUP BY team_id HAVING COUNT(*) >= $2) AS full_members) AS full_teams", + FREE_TEAMS_PER_ORGANIZATION_LIMIT, + FREE_MEMBERS_PER_TEAM_LIMIT, + ) + ) + return CustomerPortalLimits( + limits_apply=not self.premium_user, + organization_count=organization_count, + any_limit_reached=not self.premium_user + and ( + organization_count >= FREE_ORGANIZATION_LIMIT + or counts[0].full_organizations > 0 + or counts[0].full_teams > 0 + ), + ) + async def organization_responses(self, rows: Sequence[PrismaRecord]) -> tuple[Organization, ...]: records: Final = tuple( OrganizationRecord.model_validate( @@ -582,6 +645,13 @@ class CustomerHierarchyStore: alias_failure: Final = await self.claim_organization_alias(payload.alias, None) if alias_failure is not None: return alias_failure + if not self.premium_user: + await self.lock_free_capacity() + if await self.organizations.count() >= FREE_ORGANIZATION_LIMIT and ( + payload.customer_org_id is None + or await self.organizations.find_unique(where={"customer_org_id": payload.customer_org_id}) is None + ): + return Failure(ErrorCode.FREE_TIER_LIMIT) default_index: Final = default_template_index(payload.default_budgets, payload.default_budget_template) if isinstance(default_index, Failure): return default_index @@ -717,6 +787,15 @@ class CustomerHierarchyStore: parent: Final = await self.organizations.find_unique(where={"customer_org_id": payload.parent_org_id}) if parent is None: return Failure(ErrorCode.PARENT_NOT_FOUND) + if not self.premium_user: + await self.lock_free_capacity() + if await self.teams.count( + where={"parent_org_id": payload.parent_org_id} + ) >= FREE_TEAMS_PER_ORGANIZATION_LIMIT and ( + payload.customer_team_id is None + or await self.teams.find_unique(where={"customer_team_id": payload.customer_team_id}) is None + ): + return Failure(ErrorCode.FREE_TIER_LIMIT) budget_id: Final = await self.create_budget(payload.budget, actor) if payload.budget is not None else None override_id: Final = ( await self.create_budget(payload.budget_override, actor) if payload.budget_override is not None else None @@ -773,6 +852,11 @@ class CustomerHierarchyStore: organization_id: Final = await self.member_organization(payload) if isinstance(organization_id, Failure): return organization_id + if await self.member_capacity_reached(payload.parent_team_ids) and ( + payload.customer_member_id is None + or await self.members.find_unique(where={"customer_member_id": payload.customer_member_id}) is None + ): + return Failure(ErrorCode.FREE_TIER_LIMIT) template_id: Final = await self.resolve_member_template(organization_id, payload.budget_id) if isinstance(template_id, Failure): return template_id @@ -996,9 +1080,13 @@ class CustomerHierarchyStore: class CustomerHierarchy: - def __init__(self, db: TransactionSource, clock: Callable[[], datetime] = utc_now) -> None: + def __init__(self, db: TransactionSource, *, premium_user: bool, clock: Callable[[], datetime] = utc_now) -> None: self.db = db self.clock = clock + self.premium_user = premium_user + + async def portal_limits(self) -> CustomerPortalLimits | Failure: + return await self.transaction(lambda store: store.portal_limits()) async def transaction( self, @@ -1010,7 +1098,7 @@ class CustomerHierarchy: try: async with self.db.tx() as tx: - return await operation(CustomerHierarchyStore(tx, self.clock())) + return await operation(CustomerHierarchyStore(tx, self.clock(), premium_user=self.premium_user)) except UniqueViolationError: return Failure(ErrorCode.ALREADY_EXISTS) except RecordNotFoundError: diff --git a/litellm/proxy/management_helpers/customer_upsert.py b/litellm/proxy/management_helpers/customer_upsert.py index 495232195be..a89049e34f8 100644 --- a/litellm/proxy/management_helpers/customer_upsert.py +++ b/litellm/proxy/management_helpers/customer_upsert.py @@ -5,6 +5,7 @@ from pydantic import BaseModel from litellm.proxy.auth.customer_member_auto_create import DEFAULT_TEAM_ALIAS, default_team_id from litellm.proxy.management_helpers.customer_hierarchy import ( + FREE_TEAMS_PER_ORGANIZATION_LIMIT, CustomerHierarchy, CustomerHierarchyStore, ErrorCode, @@ -157,6 +158,13 @@ async def customer_team( if explicit_team.parent_org_id == organization_id else Failure(ErrorCode.CUSTOMER_ORG_MISMATCH) ) + if not store.premium_user: + await store.lock_free_capacity() + if ( + await store.teams.find_unique(where={"customer_team_id": default_team_id(organization_id)}) is None + and await store.teams.count(where={"parent_org_id": organization_id}) >= FREE_TEAMS_PER_ORGANIZATION_LIMIT + ): + return Failure(ErrorCode.FREE_TIER_LIMIT) await store.teams.upsert( where={"customer_team_id": default_team_id(organization_id)}, data={ @@ -257,6 +265,8 @@ async def upsert_member( and expiry_matches ): return MemberUpsertResponse.model_validate({**existing.model_dump(), "created": False}) + if await store.member_capacity_reached(tuple(tid for tid in team_ids if tid not in existing.parent_team_ids)): + return Failure(ErrorCode.FREE_TIER_LIMIT) await store.members.update( where={"customer_member_id": payload.customer_member_id}, data={"parent_team_ids": team_ids} ) 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 345079c286c..477ecde42e2 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 @@ -1,13 +1,17 @@ import { useState } from "react"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import userEvent from "@testing-library/user-event"; +import { toast } from "sonner"; import { fireEvent, renderWithProviders, screen, testQueryClient, waitFor } from "../../../../../tests/test-utils"; import CreateCustomerDialog from "./CreateCustomerDialog"; const creations: unknown[] = []; let rejectCreation = false; +let freeTierReached = false; const onCreated = vi.fn(); +vi.mock("sonner", () => ({ toast: { success: vi.fn(), error: vi.fn() } })); + function Harness() { const [open, setOpen] = useState(true); return ( @@ -21,6 +25,8 @@ function Harness() { beforeEach(() => { creations.length = 0; rejectCreation = false; + freeTierReached = false; + vi.mocked(toast.error).mockClear(); onCreated.mockClear(); testQueryClient.clear(); vi.stubGlobal( @@ -32,6 +38,11 @@ beforeEach(() => { } if (url.pathname === "/management/v1/customer_organizations" && request.method === "POST") { creations.push(await request.json()); + if (freeTierReached) + return Response.json( + { detail: "Free customer portal limit reached. Upgrade to Enterprise." }, + { status: 403 }, + ); return rejectCreation ? Response.json({ detail: "Customer entity ID already exists." }, { status: 409 }) : Response.json({ customer_org_id: "created-org", alias: "Example customer" }, { status: 201 }); @@ -47,6 +58,17 @@ afterEach(() => { }); describe("CreateCustomerDialog", () => { + it("surfaces the free-tier 403 detail", async () => { + freeTierReached = true; + renderWithProviders(); + fireEvent.change(screen.getByLabelText("Customer name"), { target: { value: "Sixth customer" } }); + fireEvent.click(screen.getByRole("button", { name: "Create customer" })); + await waitFor(() => + expect(toast.error).toHaveBeenCalledWith( + expect.stringContaining("Free customer portal limit reached. Upgrade to Enterprise."), + ), + ); + }); it("adds a custom group and selects a remaining default when that group is removed", async () => { const user = userEvent.setup(); renderWithProviders(); 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 ff9f823d491..74b67942f9e 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 @@ -1,6 +1,7 @@ "use client"; import { useState } from "react"; +import { useQueryClient } from "@tanstack/react-query"; import { ChevronRight, Plus, Trash2 } from "lucide-react"; import { useTranslation } from "react-i18next"; import { toast } from "sonner"; @@ -70,10 +71,12 @@ const groupBudgetFields = [ ] as const; export default function CreateCustomerDialog({ open, onOpenChange, onCreated }: Props) { + const queryClient = useQueryClient(); const { t } = useTranslation(); const [fields, setFields] = useState(initialCustomerOrganizationFields); const [error, setError] = useState(null); const create = $api.useMutation("post", "/management/v1/customer_organizations", { + onSettled: () => queryClient.invalidateQueries({ queryKey: ["get", "/management/v1/customer_portal/limits"] }), onSuccess: (customer) => { toast.success(t("customerPortal.toasts.created", { name: customer.alias || fields.name.trim() })); onCreated(customer.customer_org_id); 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 b6935aff2b0..aac5807f490 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 @@ -1,7 +1,7 @@ import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { registerBaseUrlGetter } from "@/lib/http/runtime"; -import { act, fireEvent, render, screen, waitFor } from "../../../../../tests/test-utils"; +import { act, fireEvent, renderWithProviders as render, screen, waitFor } from "../../../../../tests/test-utils"; import CreateCustomerHierarchyDialog from "./CreateCustomerHierarchyDialog"; import type { Organization, Team } from "./customerHierarchyForms"; @@ -161,4 +161,25 @@ describe("CreateCustomerHierarchyDialog", () => { expect(onOpenChange).not.toHaveBeenCalled(); expect(onCreated).not.toHaveBeenCalled(); }); + + it.each(["team", "member"] as const)("surfaces the free-tier 403 detail when creating a %s", async (kind) => { + fetchMock.mockResolvedValue(json({ detail: "Free customer portal limit reached. Upgrade to Enterprise." }, 403)); + render( + , + ); + fireEvent.change(screen.getByRole("textbox", { name: "Display name" }), { target: { value: "Example" } }); + fireEvent.click(screen.getByRole("button", { name: "Create" })); + expect(await screen.findByRole("alert")).toHaveTextContent( + "Free customer portal limit reached. Upgrade to Enterprise.", + ); + expect(onOpenChange).not.toHaveBeenCalled(); + expect(onCreated).not.toHaveBeenCalled(); + }); }); diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerHierarchyDialog.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerHierarchyDialog.tsx index 479d21296f8..d9b6b9d693d 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerHierarchyDialog.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerHierarchyDialog.tsx @@ -1,6 +1,7 @@ "use client"; import { useState, type FormEvent } from "react"; +import { useQueryClient } from "@tanstack/react-query"; import { useTranslation } from "react-i18next"; import { toast } from "sonner"; import { fetchClient } from "@/lib/http/api"; @@ -45,6 +46,7 @@ const createEntity = async (create: CustomerHierarchyCreate): Promise void }) { + const queryClient = useQueryClient(); const { t } = useTranslation(); const [fields, setFields] = useState(() => newCustomerHierarchyFields(props)); const [identifier, setIdentifier] = useState(""); @@ -87,6 +89,7 @@ function CreateForm(props: Props & { pending: boolean; onPendingChange: (pending } catch (cause) { setError(extractProxyErrorMessage(cause)); } finally { + void queryClient.invalidateQueries({ queryKey: ["get", "/management/v1/customer_portal/limits"] }); setPending(false); } }; 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 174e7815669..5472b1570e4 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 @@ -1,7 +1,17 @@ import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { registerBaseUrlGetter } from "@/lib/http/runtime"; -import { act, chooseSelectOption, fireEvent, render, screen, waitFor, within } from "../../../../../tests/test-utils"; +import { + act, + chooseSelectOption, + fireEvent, + renderWithProviders as render, + screen, + testQueryClient, + waitFor, + within, +} from "../../../../../tests/test-utils"; +import CustomersLayout from "../layout"; import CustomerHierarchySettings from "./CustomerHierarchySettings"; import type { Organization, Member, CustomerHierarchyEntity } from "./customerHierarchyForms"; @@ -64,12 +74,16 @@ const lastRequest = (): Request => { describe("CustomerHierarchySettings", () => { beforeEach(() => { registerBaseUrlGetter(() => "http://localhost"); + testQueryClient.clear(); fetchMock.mockReset(); onSaved.mockReset(); onDeleted.mockReset(); vi.stubGlobal("fetch", fetchMock); }); - afterEach(() => vi.unstubAllGlobals()); + afterEach(() => { + testQueryClient.clear(); + vi.unstubAllGlobals(); + }); it("should render the existing organization settings and read-only identifier", () => { render( @@ -342,6 +356,43 @@ describe("CustomerHierarchySettings", () => { }, ); + it("clears the free-tier warning when a successful deletion frees capacity", async () => { + let atCapacity = true; + fetchMock.mockImplementation(async (input) => { + const request = input as Request; + if (request.method === "DELETE") { + atCapacity = false; + return new Response(null, { status: 204 }); + } + const limits = { + limits_apply: true, + organization_limit: 5, + teams_per_organization_limit: 5, + members_per_team_limit: 5, + organization_count: 1, + any_limit_reached: atCapacity, + }; + return json(limits); + }); + render( + + + , + ); + expect(await screen.findByRole("alert")).toHaveTextContent("Free customer portal limit reached"); + fireEvent.click(screen.getByRole("button", { name: "Delete customer member" })); + const dialog = await screen.findByRole("dialog", { name: "Delete customer member" }); + fireEvent.change(within(dialog).getByPlaceholderText("member-1"), { target: { value: "member-1" } }); + fireEvent.click(within(dialog).getByRole("button", { name: "Delete" })); + await waitFor(() => expect(screen.queryByRole("alert")).not.toBeInTheDocument()); + expect(screen.getByText(/Free tier: up to 5 customers/)).toBeInTheDocument(); + }); + it("should keep the confirmation pending until a delete returns204", async () => { const user = userEvent.setup(); let settle: ((response: Response) => void) | undefined; diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchySettings.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchySettings.tsx index 32368c989aa..1bdc8880c59 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchySettings.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerHierarchySettings.tsx @@ -1,6 +1,7 @@ "use client"; import { useState, type FormEvent } from "react"; +import { useQueryClient } from "@tanstack/react-query"; import { useTranslation } from "react-i18next"; import { toast } from "sonner"; import { Trash2 } from "lucide-react"; @@ -85,6 +86,7 @@ const DELETE_MESSAGES = { } as const; function SettingsForm({ entity, organization, onSaved, onDeleted }: Props) { + const queryClient = useQueryClient(); const { t } = useTranslation(); const initial = customerHierarchyFields(entity); const [fields, setFields] = useState(initial); @@ -128,6 +130,7 @@ function SettingsForm({ entity, organization, onSaved, onDeleted }: Props) { setDeleteError(null); try { await deleteEntity(entity); + void queryClient.invalidateQueries({ queryKey: ["get", "/management/v1/customer_portal/limits"] }); setDeleteOpen(false); onDeleted(); } catch (cause) { diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/layout.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/layout.tsx new file mode 100644 index 00000000000..adac885663a --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/layout.tsx @@ -0,0 +1,60 @@ +"use client"; + +import type { ReactNode } from "react"; +import { Trans, useTranslation } from "react-i18next"; +import { AlertTriangle } from "lucide-react"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert"; +import { $api } from "@/lib/http/api"; + +const enterpriseUrl = "https://www.litellm.ai/enterprise"; + +export default function CustomersLayout({ children }: { children: ReactNode }) { + const { t } = useTranslation(); + const { accessToken } = useAuthorized(); + const { data: limits } = $api.useQuery( + "get", + "/management/v1/customer_portal/limits", + {}, + { + enabled: Boolean(accessToken), + }, + ); + + return ( + <> + {limits?.limits_apply && ( +
+

+ + ), + }} + /> +

+ {limits.any_limit_reached && ( + + + {t("customerPortal.limits.reachedTitle")} + +

{t("customerPortal.limits.reachedDescription")}

+ + {t("customerPortal.limits.upgrade")} + + + + )} +
+ )} + {children} + + ); +} diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/limits.integration.test.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/limits.integration.test.tsx new file mode 100644 index 00000000000..7b8f6cdf7ce --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/limits.integration.test.tsx @@ -0,0 +1,85 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { act, renderWithProviders, screen, testQueryClient } from "../../../../tests/test-utils"; +import CustomersLayout from "./layout"; + +beforeEach(() => testQueryClient.clear()); +afterEach(() => { + testQueryClient.clear(); + vi.unstubAllGlobals(); +}); + +describe("Customer portal limits", () => { + it.each([ + { limitsApply: true, reached: false }, + { limitsApply: true, reached: true }, + { limitsApply: false, reached: true }, + { limitsApply: false, reached: false }, + ])("shows the notices for limits_apply=$limitsApply and reached=$reached", async ({ limitsApply, reached }) => { + vi.stubGlobal( + "fetch", + vi.fn(async (request: Request) => { + expect(new URL(request.url).pathname).toBe("/management/v1/customer_portal/limits"); + const limits = { + limits_apply: limitsApply, + organization_limit: 5, + teams_per_organization_limit: 5, + members_per_team_limit: 5, + organization_count: reached ? 5 : 1, + any_limit_reached: reached, + }; + return Response.json(limits); + }), + ); + renderWithProviders( + +

Customers

+ , + ); + await act(async () => { + await testQueryClient.ensureQueryData({ queryKey: ["get", "/management/v1/customer_portal/limits", {}] }); + }); + if (limitsApply) { + expect(await screen.findByText(/Free tier: up to 5 customers/)).toHaveTextContent( + "Free tier: up to 5 customers, 5 teams per customer, 5 members per team. LiteLLM Enterprise removes these limits.", + ); + expect(screen.getByRole("link", { name: "Enterprise" })).toHaveAttribute( + "href", + "https://www.litellm.ai/enterprise", + ); + } else { + expect(screen.queryByText(/Free tier:/)).not.toBeInTheDocument(); + } + if (limitsApply && reached) { + expect(await screen.findByRole("alert")).toHaveTextContent("Free customer portal limit reached"); + expect(screen.getByRole("alert")).toHaveTextContent( + "Requests from new customers or members still go through but are not tracked", + ); + expect(screen.getByRole("link", { name: "Upgrade to Enterprise" })).toHaveAttribute( + "href", + "https://www.litellm.ai/enterprise", + ); + } else { + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + } + }); + + it.each([204, 503])("keeps the page available without notices when limits return %s", async (status) => { + vi.stubGlobal( + "fetch", + vi.fn(async () => new Response(null, { status })), + ); + renderWithProviders( + +

Customers

+ , + ); + await act(async () => { + await testQueryClient + .ensureQueryData({ queryKey: ["get", "/management/v1/customer_portal/limits", {}] }) + .catch(() => undefined); + }); + expect(screen.getByRole("heading", { name: "Customers" })).toBeInTheDocument(); + expect(screen.queryByText(/Free tier:/)).not.toBeInTheDocument(); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + }); +}); diff --git a/litellm/ui/litellm-dashboard/src/i18n/locales/en.json b/litellm/ui/litellm-dashboard/src/i18n/locales/en.json index fe60af7f715..e9d0071a49e 100644 --- a/litellm/ui/litellm-dashboard/src/i18n/locales/en.json +++ b/litellm/ui/litellm-dashboard/src/i18n/locales/en.json @@ -20630,6 +20630,12 @@ "loadFailed": "Could not load this model group's deployments. Saved weights are kept." }, "customerPortal": { + "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", + "reachedDescription": "New customers, teams and members can't be created. Requests from new customers or members still go through but are not tracked in the customer portal.", + "upgrade": "Upgrade to Enterprise" + }, "hierarchy": { "customerNavigation": "Customer navigation", "usage": "Usage", 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 0ef4354a1a8..d53543fb937 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9921,6 +9921,26 @@ export interface paths { patch: operations["update_customer_budget_template_management_v1_customer_organizations__customer_org_id__default_budgets__budget_id__patch"]; trace?: never; }; + "/management/v1/customer_portal/limits": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Customer Portal Limits + * @description Get customer portal free-tier limits and current capacity. Proxy admins only. + */ + get: operations["get_customer_portal_limits_management_v1_customer_portal_limits_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/management/v1/customer_teams": { parameters: { query?: never; @@ -33420,6 +33440,30 @@ export interface components { /** Updated By */ updated_by: string; }; + /** CustomerPortalLimits */ + CustomerPortalLimits: { + /** Any Limit Reached */ + any_limit_reached: boolean; + /** Limits Apply */ + limits_apply: boolean; + /** + * Members Per Team Limit + * @default 5 + */ + members_per_team_limit: number; + /** Organization Count */ + organization_count: number; + /** + * Organization Limit + * @default 5 + */ + organization_limit: number; + /** + * Teams Per Organization Limit + * @default 5 + */ + teams_per_organization_limit: number; + }; /** * CustomerResponse * @description Customer object returned by the /customer read+write endpoints. @@ -69338,6 +69382,26 @@ export interface operations { }; }; }; + get_customer_portal_limits_management_v1_customer_portal_limits_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["CustomerPortalLimits"]; + }; + }; + }; + }; list_customer_teams_management_v1_customer_teams_get: { parameters: { query?: { diff --git a/tests/unit/proxy/auth/test_customer_hierarchy_auth.py b/tests/unit/proxy/auth/test_customer_hierarchy_auth.py index 90ba3e66e1b..034ffb3143a 100644 --- a/tests/unit/proxy/auth/test_customer_hierarchy_auth.py +++ b/tests/unit/proxy/auth/test_customer_hierarchy_auth.py @@ -661,6 +661,26 @@ async def test_an_alias_next_to_a_member_must_name_that_members_organization() - assert (denied.value.code, denied.value.param) == ("403", CUSTOMER_ORG_ALIAS_HEADER) +@pytest.mark.asyncio +@pytest.mark.parametrize("member_id", (None, "new-member")) +async def test_a_scoped_key_cannot_create_another_organization_before_being_rejected(member_id: str | None) -> None: + organizations: Final = _OrganizationTable((_organization("existing"),)) + registrar: Final = _OrganizationRegistrar(organizations) + prisma: Final = _Prisma(_MemberTable({}), organizations) + + with pytest.raises(ProxyException) as denied: + await _resolve( + prisma, + CustomerRequestIdentity(customer_org_alias="new-organization", customer_member_id=member_id), + key_metadata={"customer_org_id": "org-1"}, + register_organization=registrar, + ) + + assert (denied.value.code, denied.value.param) == ("403", CUSTOMER_MEMBER_ID_HEADER) + assert registrar.calls == () + assert tuple(row["customer_org_id"] for row in organizations.rows) == ("org-1",) + + @pytest.mark.asyncio async def test_an_unknown_alias_becomes_a_new_organization_charged_for_the_request() -> None: organizations: Final = _OrganizationTable(()) @@ -687,8 +707,11 @@ async def test_an_unknown_alias_becomes_a_new_organization_charged_for_the_reque @pytest.mark.asyncio -async def test_an_alias_that_cannot_be_registered_runs_untracked_and_a_shared_alias_is_rejected_as_ambiguous() -> None: - failing: Final = _OrganizationRegistrar(_OrganizationTable(()), Failure(ErrorCode.ALREADY_EXISTS)) +@pytest.mark.parametrize("error", (ErrorCode.ALREADY_EXISTS, ErrorCode.FREE_TIER_LIMIT)) +async def test_an_alias_that_cannot_be_registered_runs_untracked_and_a_shared_alias_is_rejected_as_ambiguous( + error: ErrorCode, +) -> None: + failing: Final = _OrganizationRegistrar(_OrganizationTable(()), Failure(error)) untracked: Final = await _resolve( _Prisma(_MemberTable({})), CustomerRequestIdentity(customer_org_alias="ghost"), register_organization=failing ) @@ -754,9 +777,12 @@ async def test_a_member_the_aliased_organization_has_never_seen_is_registered_th @pytest.mark.asyncio -async def test_an_unknown_member_runs_untracked_without_an_alias_and_on_the_organization_when_registration_fails() -> ( - None -): +@pytest.mark.parametrize("error", (ErrorCode.BUDGET_REQUIRED, ErrorCode.FREE_TIER_LIMIT)) +@pytest.mark.parametrize("team_id", (None, "team-1")) +async def test_an_unknown_member_runs_untracked_without_an_alias_and_on_the_organization_when_registration_fails( + error: ErrorCode, + team_id: str | None, +) -> None: members: Final = _MemberTable({}) registrar: Final = _Registrar(members) no_alias: Final = await _resolve( @@ -764,10 +790,12 @@ async def test_an_unknown_member_runs_untracked_without_an_alias_and_on_the_orga ) assert (no_alias, registrar.calls) == (None, ()) - failing: Final = _Registrar(members, Failure(ErrorCode.BUDGET_REQUIRED)) + failing: Final = _Registrar(members, Failure(error)) prisma: Final = _Prisma(members, _OrganizationTable((_organization("acme", max_budget=5.0),))) organization_only: Final = await _resolve( - prisma, CustomerRequestIdentity(customer_member_id="user-7", customer_org_alias="acme"), register=failing + prisma, + CustomerRequestIdentity(customer_member_id="user-7", customer_org_alias="acme", customer_team_id=team_id), + register=failing, ) assert failing.calls == (("org-1", "user-7"),) assert organization_only is not None diff --git a/tests/unit/proxy/auth/test_customer_member_auto_create.py b/tests/unit/proxy/auth/test_customer_member_auto_create.py index 88bfa74f916..6701953b584 100644 --- a/tests/unit/proxy/auth/test_customer_member_auto_create.py +++ b/tests/unit/proxy/auth/test_customer_member_auto_create.py @@ -29,9 +29,9 @@ def _prisma(database: Database) -> "PrismaClient": async def _organization(database: Database, org_id: str, default_budget: bool = True) -> None: templates: Final = (BudgetTemplate(name="std", max_budget=10.0),) if default_budget else () - created: Final = await CustomerHierarchy(_prisma(database), clock=lambda: NOW).create_organization( - OrganizationCreate(customer_org_id=org_id, default_budgets=templates), "admin" - ) + created: Final = await CustomerHierarchy( + _prisma(database), premium_user=False, clock=lambda: NOW + ).create_organization(OrganizationCreate(customer_org_id=org_id, default_budgets=templates), "admin") assert not isinstance(created, Failure) @@ -53,7 +53,7 @@ async def test_a_new_member_joins_the_organizations_default_team_on_the_organiza async def test_registering_the_same_member_twice_keeps_one_row_and_reuses_an_existing_default_team() -> None: database: Final = Database() await _organization(database, "org-1") - hierarchy: Final = CustomerHierarchy(_prisma(database), clock=lambda: NOW) + hierarchy: Final = CustomerHierarchy(_prisma(database), premium_user=False, clock=lambda: NOW) existing_team: Final = await hierarchy.create_team( TeamCreate(customer_team_id=default_team_id("org-1"), alias="pre-existing", parent_org_id="org-1"), "admin" ) @@ -84,7 +84,7 @@ async def test_a_new_member_of_an_organization_without_a_default_budget_is_regis async def test_an_admin_created_member_still_needs_a_budget_and_an_uncapped_member_stays_editable() -> None: database: Final = Database() await _organization(database, "org-1", default_budget=False) - hierarchy: Final = CustomerHierarchy(_prisma(database), clock=lambda: NOW) + hierarchy: Final = CustomerHierarchy(_prisma(database), premium_user=True, clock=lambda: NOW) registered: Final = await ensure_customer_member("org-1", "user-7", _prisma(database)) assert isinstance(registered, CustomerMember), registered @@ -102,6 +102,28 @@ async def test_an_unknown_organization_cannot_register_members() -> None: assert await ensure_customer_member("ghost", "user-7", _prisma(Database())) == Failure(ErrorCode.PARENT_NOT_FOUND) +@pytest.mark.asyncio +async def test_auto_registration_stops_at_free_capacity_but_reuses_existing_ids( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "premium_user", False) + database: Final = Database() + await _organization(database, "org-0") + for index in range(1, 5): + assert not isinstance(await ensure_customer_organization(f"org-{index}", _prisma(database)), Failure) + assert await ensure_customer_organization("sixth", _prisma(database)) == Failure(ErrorCode.FREE_TIER_LIMIT) + assert not isinstance(await ensure_customer_organization("org-0", _prisma(database)), Failure) + for index in range(5): + assert isinstance(await ensure_customer_member("org-0", f"member-{index}", _prisma(database)), CustomerMember) + assert await ensure_customer_member("org-0", "sixth", _prisma(database)) == Failure(ErrorCode.FREE_TIER_LIMIT) + assert isinstance(await ensure_customer_member("org-0", "member-0", _prisma(database)), CustomerMember) + monkeypatch.setattr(proxy_server, "premium_user", True) + assert not isinstance(await ensure_customer_organization("sixth", _prisma(database)), Failure) + assert isinstance(await ensure_customer_member("org-0", "sixth", _prisma(database)), CustomerMember) + + @pytest.mark.asyncio async def test_an_unknown_alias_becomes_an_organization_once_and_a_known_one_is_reused() -> None: database: Final = Database() @@ -111,9 +133,9 @@ async def test_an_unknown_alias_becomes_an_organization_once_and_a_known_one_is_ assert not isinstance(created, Failure) and not isinstance(again, Failure) assert (created.customer_org_id, created.alias, again.customer_org_id) == ("acme", "acme", "acme") - listed: Final = await CustomerHierarchy(_prisma(database), clock=lambda: NOW).list_organizations( - page=1, page_size=10 - ) + listed: Final = await CustomerHierarchy( + _prisma(database), premium_user=False, clock=lambda: NOW + ).list_organizations(page=1, page_size=10) assert not isinstance(listed, Failure) and len(listed.data) == 1 @@ -121,7 +143,7 @@ async def test_an_unknown_alias_becomes_an_organization_once_and_a_known_one_is_ async def test_renaming_a_member_keeps_a_budget_assigned_after_the_rename_read_the_member() -> None: database: Final = Database() await _organization(database, "org-1", default_budget=False) - hierarchy: Final = CustomerHierarchy(_prisma(database), clock=lambda: NOW) + hierarchy: Final = CustomerHierarchy(_prisma(database), premium_user=True, clock=lambda: NOW) registered: Final = await ensure_customer_member("org-1", "user-7", _prisma(database)) assert isinstance(registered, CustomerMember), registered capped: Final = await hierarchy.update_member("user-7", MemberPatch(budget_override=Budget(max_budget=3.0)), "a") 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 84aadd7c018..e17656231bf 100644 --- a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py +++ b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py @@ -31,6 +31,7 @@ def client( role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN, cache: UserApiKeyCache | None = None, database: Database | None = None, + premium_user: bool = True, actor: str = "admin", ) -> TestClient: app: Final = FastAPI() @@ -42,7 +43,7 @@ def client( auth: Final = UserAPIKeyAuth(user_id=actor, user_role=role) hierarchy: Final = CustomerHierarchy( - cast("Prisma", database if database is not None else Database()), clock=lambda: NOW + cast("Prisma", database if database is not None else Database()), premium_user=premium_user, clock=lambda: NOW ) # cast-ok: in-memory Prisma edge app.dependency_overrides[user_api_key_auth] = lambda: auth app.dependency_overrides[get_customer_hierarchy] = lambda: hierarchy @@ -50,6 +51,95 @@ def client( return TestClient(app) +@pytest.mark.parametrize("premium", (False, True)) +def test_customer_portal_limits_and_free_tier_problem(premium: bool) -> None: + http: Final = client(premium_user=premium) + empty: Final = http.get("/management/v1/customer_portal/limits") + assert empty.status_code == 200, empty.text + assert empty.json() == { + "limits_apply": not premium, + "organization_limit": 5, + "teams_per_organization_limit": 5, + "members_per_team_limit": 5, + "organization_count": 0, + "any_limit_reached": False, + } + for index in range(5): + response: Final = http.post("/management/v1/customer_organizations", json={"customer_org_id": f"org-{index}"}) + assert response.status_code == 201, response.text + full: Final = http.get("/management/v1/customer_portal/limits") + assert full.status_code == 200, full.text + assert full.json() == {**empty.json(), "organization_count": 5, "any_limit_reached": not premium} + sixth: Final = http.post("/management/v1/customer_organizations", json={"customer_org_id": "sixth"}) + assert sixth.status_code == (201 if premium else 403), sixth.text + if not premium: + assert sixth.json()["detail"] == ( + "Free customer portal limit reached: up to 5 customer organizations, 5 teams per organization and " + "5 members per team without a LiteLLM Enterprise license. Upgrade to Enterprise to remove the limits: " + "https://www.litellm.ai/enterprise" + ) + + +@pytest.mark.parametrize("level", ("teams", "members")) +def test_customer_portal_limits_detect_full_children_before_organization_cap(level: str) -> None: + http: Final = client(premium_user=False) + family(http) + for index in range(5): + if level == "teams": + response: Final = http.post( + "/management/v1/customer_teams", json={"parent_org_id": "org", "customer_team_id": f"extra-{index}"} + ) + assert response.status_code == (201 if index < 4 else 403), response.text + else: + member: Final = http.post( + "/management/v1/customer_members", + json={"customer_member_id": f"member-{index}", "parent_team_ids": ["team"]}, + ) + assert member.status_code == 201, member.text + limits: Final = http.get("/management/v1/customer_portal/limits") + assert limits.status_code == 200, limits.text + assert limits.json()["organization_count"] == 1 + assert limits.json()["any_limit_reached"] is True + + +@pytest.mark.parametrize("premium", (False, True)) +def test_member_upsert_respects_capacity_for_default_team_creation_and_added_membership(premium: bool) -> None: + http: Final = client(premium_user=premium) + family(http) + for index in range(4): + created: Final = http.post( + "/management/v1/customer_teams", json={"parent_org_id": "org", "customer_team_id": f"extra-{index}"} + ) + assert created.status_code == 201, created.text + default_team: Final = http.post( + "/management/v1/customer_members/upsert", json={"customer_org_id": "org", "customer_member_id": "defaulted"} + ) + assert default_team.status_code == (200 if premium else 403), default_team.text + for index in range(5): + member: Final = http.post( + "/management/v1/customer_members", + json={"customer_member_id": f"full-{index}", "parent_team_ids": ["team"]}, + ) + assert member.status_code == 201, member.text + original: Final = http.post( + "/management/v1/customer_members", json={"customer_member_id": "moving", "parent_team_ids": ["extra-0"]} + ) + assert original.status_code == 201, original.text + moved: Final = http.post( + "/management/v1/customer_members/upsert", + json={"customer_org_id": "org", "customer_member_id": "moving", "parent_team_id": "team"}, + ) + assert moved.status_code == (200 if premium else 403), moved.text + if not premium: + assert http.get("/management/v1/customer_members/moving").json()["parent_team_ids"] == ["extra-0"] + assert http.get("/management/v1/customer_teams/org:default").status_code == 404 + + +def test_customer_portal_limits_require_admin() -> None: + response: Final = client(role=LitellmUserRoles.INTERNAL_USER).get("/management/v1/customer_portal/limits") + assert response.status_code == 403, response.text + + def family(http: TestClient) -> dict[str, object]: org: Final = http.post( "/management/v1/customer_organizations", diff --git a/tests/unit/proxy/management_helpers/test_customer_hierarchy.py b/tests/unit/proxy/management_helpers/test_customer_hierarchy.py index 91c21271cd3..8d9be312ac9 100644 --- a/tests/unit/proxy/management_helpers/test_customer_hierarchy.py +++ b/tests/unit/proxy/management_helpers/test_customer_hierarchy.py @@ -4,7 +4,8 @@ from collections.abc import AsyncIterator, Mapping, Sequence from contextlib import asynccontextmanager from datetime import datetime, timezone from functools import reduce -from typing import Final +from itertools import chain +from typing import TYPE_CHECKING, Final, cast import pytest from prisma import Json @@ -13,11 +14,22 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.models.budget import BudgetLimit, LiteLLM_BudgetTable from litellm.proxy._types import BudgetWriteRequest from litellm.proxy.management_helpers.customer_hierarchy import ( + CustomerHierarchy, + ErrorCode, + Failure, budget_row_values, effective_member_budget, related_budget, ) -from litellm.types.proxy.management_endpoints.customer_hierarchy import Budget +from litellm.types.proxy.management_endpoints.customer_hierarchy import ( + Budget, + MemberCreate, + OrganizationCreate, + TeamCreate, +) + +if TYPE_CHECKING: + from prisma import Prisma NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) @@ -265,6 +277,7 @@ class Database: connection.executescript( 'CREATE TABLE "LiteLLM_CustomerTeamTable" (customer_team_id TEXT, parent_org_id TEXT);' 'CREATE TABLE "LiteLLM_CustomerMemberTable" (customer_member_id TEXT, parent_org_id TEXT);' + "CREATE TABLE member_teams (team_id TEXT);" ) connection.executemany( 'INSERT INTO "LiteLLM_CustomerTeamTable" VALUES (?, ?)', @@ -273,6 +286,16 @@ class Database: for row in self.litellm_customerteamtable.rows.values() ), ) + connection.executemany( + "INSERT INTO member_teams VALUES (?)", + tuple( + (team_id,) + for team_id in chain.from_iterable( + TypeAdapter(tuple[str, ...]).validate_python(row.model_dump()["parent_team_ids"]) + for row in self.litellm_customermembertable.rows.values() + ) + ), + ) connection.executemany( 'INSERT INTO "LiteLLM_CustomerMemberTable" VALUES (?, ?)', tuple( @@ -283,8 +306,13 @@ class Database: return tuple( dict(row) for row in connection.execute( - query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each($1))"), - tuple(json.dumps(arg) for arg in args), + query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each($1))").replace( + '"LiteLLM_CustomerMemberTable", unnest(parent_team_ids) AS team_id', "member_teams" + ), + { + str(index): json.dumps(arg) if not isinstance(arg, int) else arg + for index, arg in enumerate(args, 1) + }, ) ) @@ -307,6 +335,53 @@ class Database: raise +@pytest.mark.asyncio +@pytest.mark.parametrize("premium", (False, True)) +async def test_customer_portal_capacity_caps_every_parent(premium: bool) -> None: + database: Final = Database() + hierarchy: Final = CustomerHierarchy( + cast("Prisma", database), premium_user=premium, clock=lambda: NOW + ) # cast-ok: in-memory Prisma edge + for index in range(5): + assert not isinstance( + await hierarchy.create_organization(OrganizationCreate(customer_org_id=f"org-{index}"), "admin"), Failure + ) + sixth_org: Final = await hierarchy.create_organization(OrganizationCreate(customer_org_id="sixth"), "admin") + assert isinstance(sixth_org, Failure) is not premium + if not premium: + assert sixth_org == Failure(ErrorCode.FREE_TIER_LIMIT) + for index in range(5): + assert not isinstance( + await hierarchy.create_team(TeamCreate(customer_team_id=f"team-{index}", parent_org_id="org-0"), "admin"), + Failure, + ) + sixth_team: Final = await hierarchy.create_team( + TeamCreate(customer_team_id="sixth", parent_org_id="org-0"), "admin" + ) + assert isinstance(sixth_team, Failure) is not premium + if not premium: + assert sixth_team == Failure(ErrorCode.FREE_TIER_LIMIT) + for index in range(5): + assert not isinstance( + await hierarchy.create_member( + MemberCreate( + customer_member_id=f"member-{index}", + parent_team_ids=("team-0", "team-1"), + budget_override=Budget(max_budget=10), + ), + "admin", + ), + Failure, + ) + sixth_member: Final = await hierarchy.create_member( + MemberCreate(customer_member_id="sixth", parent_team_ids=("team-1",), budget_override=Budget(max_budget=10)), + "admin", + ) + assert isinstance(sixth_member, Failure) is not premium + if not premium: + assert sixth_member == Failure(ErrorCode.FREE_TIER_LIMIT) + + @pytest.mark.parametrize("expires_at", (None, datetime(2027, 1, 1, tzinfo=timezone.utc))) def test_effective_budget_merges_non_null_fields_and_empty_models_inherit(expires_at: datetime | None) -> None: template: Final = Budget(max_budget=100, rpm_limit=10, allowed_models=("model-a",))