diff --git a/litellm/proxy/auth/customer_member_auto_create.py b/litellm/proxy/auth/customer_member_auto_create.py index 34f1b91cfc4..59c69a7e63e 100644 --- a/litellm/proxy/auth/customer_member_auto_create.py +++ b/litellm/proxy/auth/customer_member_auto_create.py @@ -79,7 +79,7 @@ async def ensure_customer_organization( created: Final = await hierarchy.create_organization( OrganizationCreate(customer_org_id=customer_org_alias, alias=customer_org_alias), AUTO_CREATE_ORG_ACTOR ) - if isinstance(created, Failure) and created.code is ErrorCode.ALREADY_EXISTS: + if isinstance(created, Failure) and created.code in (ErrorCode.ALREADY_EXISTS, ErrorCode.ALIAS_ALREADY_EXISTS): return await hierarchy.get_organization(customer_org_alias) if not isinstance(created, Failure): verbose_proxy_logger.info("Registered customer organization %s", customer_org_alias) diff --git a/litellm/proxy/management_endpoints/management_v1/customers.py b/litellm/proxy/management_endpoints/management_v1/customers.py index fac17c02ca0..3d978e2925c 100644 --- a/litellm/proxy/management_endpoints/management_v1/customers.py +++ b/litellm/proxy/management_endpoints/management_v1/customers.py @@ -17,6 +17,7 @@ from litellm.proxy.list_api.common import ( ) 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_upsert import CustomerUpsert from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, # pyright: ignore[reportUnknownVariableType] # legacy decorator is untyped ) @@ -25,9 +26,13 @@ from litellm.types.proxy.management_endpoints.customer_hierarchy import ( CustomerMember, MemberCreate, MemberPatch, + MemberUpsert, + MemberUpsertResponse, Organization, OrganizationCreate, OrganizationPatch, + OrganizationUpsert, + OrganizationUpsertResponse, Team, TeamCreate, TeamPatch, @@ -84,6 +89,16 @@ _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.ALIAS_ALREADY_EXISTS: ( + 409, + "Customer org alias already exists", + "Customer org alias is already used by another customer org.", + ), + ErrorCode.CUSTOMER_ORG_MISMATCH: ( + 409, + "Customer org mismatch", + "Customer member and customer team must belong to the specified customer org.", + ), } ) assert _ERROR_DETAILS.keys() == frozenset(ErrorCode), "every ErrorCode needs a public problem" @@ -186,6 +201,32 @@ async def create_customer_organization( return resolve(await hierarchy.create_organization(data, actor)) +@router.post( + "/customer_organizations/upsert", + response_model=OrganizationUpsertResponse, + description=( + "Create or update a customer org by customer_org_id, or by alias when no ID is supplied. " + "Returns 200 with created=true for creation and created=false for updates. Repeated identical calls are no-ops. " + "Supplied default_budgets are merged by template name, preserving existing template IDs and omitted templates. " + "A single supplied template becomes the default; multiple templates require default_budget_template to name one. " + "Omitted fields retain their values. Conflicting or ambiguous aliases return 409. Proxy admins only." + ), +) +@management_endpoint_wrapper +async def upsert_customer_organization( + data: OrganizationUpsert, + actor: Annotated[str, Depends(customer_actor)], + hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)], + cache: AuthCache, +) -> OrganizationUpsertResponse: + outcome: Final = resolve(await CustomerUpsert(hierarchy).organization(data, actor)) + await forget_customer_hierarchies( + resolve(await hierarchy.member_ids_in_organization(outcome.data.customer_org_id)), cache + ) + await forget_customer_organization_aliases((outcome.previous_alias, outcome.data.alias), cache) + return outcome.data + + @router.get( "/customer_organizations/{customer_org_id}", response_model=Organization, @@ -351,6 +392,32 @@ async def create_customer_member( return resolve(await hierarchy.create_member(data, actor)) +@router.post( + "/customer_members/upsert", + response_model=MemberUpsertResponse, + description=( + "Create or update a customer member by customer_member_id within a customer org identified by " + "customer_org_id or customer_org_alias. If both are supplied they must match. " + "A new customer member with no parent_team_id joins the customer org's default customer team, created if necessary. An existing customer member keeps its customer teams, and a supplied parent_team_id becomes its first customer team. " + "Alias, budget_id, budget_override and budget_override_expires_at follow the existing customer member models. " + "Omitted update fields retain their values. Moving between customer teams in the same customer org is allowed. " + "Moving between customer orgs or selecting another customer org's customer team returns 409. " + "Returns 200 with created=true for creation and created=false for updates. Repeated identical calls are no-ops. " + "Proxy admins only." + ), +) +@management_endpoint_wrapper +async def upsert_customer_member( + data: MemberUpsert, + actor: Annotated[str, Depends(customer_actor)], + hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)], + cache: AuthCache, +) -> MemberUpsertResponse: + member: Final = resolve(await CustomerUpsert(hierarchy).member(data, actor)) + await forget_customer_hierarchies((member.customer_member_id,), cache) + return member + + @router.get( "/customer_members/{customer_member_id}", response_model=CustomerMember, diff --git a/litellm/proxy/management_helpers/customer_hierarchy.py b/litellm/proxy/management_helpers/customer_hierarchy.py index 30ba85eaa6e..f8f40971fbd 100644 --- a/litellm/proxy/management_helpers/customer_hierarchy.py +++ b/litellm/proxy/management_helpers/customer_hierarchy.py @@ -52,6 +52,8 @@ class ErrorCode(str, Enum): INVALID_TEMPLATE_NAMES = "invalid_template_names" TEMPLATE_IN_USE = "template_in_use" BUDGET_REQUIRED = "budget_required" + ALIAS_ALREADY_EXISTS = "alias_already_exists" + CUSTOMER_ORG_MISMATCH = "customer_org_mismatch" INCREASE_NOT_ALLOWED = "increase_not_allowed" REQUEST_NOT_PENDING = "request_not_pending" REQUEST_ALREADY_PENDING = "request_already_pending" @@ -68,6 +70,10 @@ _LOCK_ORGANIZATION_SQL: Final = ( 'SELECT 1 AS locked FROM "LiteLLM_CustomerOrganizationTable" WHERE customer_org_id = $1 FOR UPDATE' ) +_LOCK_CUSTOMER_IDENTITIES_SQL: Final = ( + "SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtextextended('litellm:customer-identities', 0))" +) + class TransactionOperation(str, Enum): READ_WRITE = "read_write" @@ -555,7 +561,25 @@ class CustomerHierarchyStore: async def lock_organization(self, organization_id: str) -> None: await self.db.query_raw(_LOCK_ORGANIZATION_SQL, organization_id) + async def lock_customer_identities(self) -> None: + await self.db.query_raw(_LOCK_CUSTOMER_IDENTITIES_SQL) + + async def claim_organization_alias(self, alias: str | None, organization_id: str | None) -> Failure | None: + await self.lock_customer_identities() + if alias is None: + return None + duplicates: Final = await self.organizations.find_many( + where={"alias": alias, "customer_org_id": {"not": organization_id}} + if organization_id + else {"alias": alias}, + take=1, + ) + return Failure(ErrorCode.ALIAS_ALREADY_EXISTS) if duplicates else None + async def create_organization(self, payload: OrganizationCreate, actor: str) -> Organization | Failure: + alias_failure: Final = await self.claim_organization_alias(payload.alias, None) + if alias_failure is not None: + return alias_failure default_index: Final = default_template_index(payload.default_budgets, payload.default_budget_template) if isinstance(default_index, Failure): return default_index @@ -605,6 +629,9 @@ class CustomerHierarchyStore: async def update_organization( self, organization_id: str, payload: OrganizationPatch, actor: str ) -> Organization | Failure: + alias_failure: Final = await self.claim_organization_alias(payload.alias, organization_id) + if alias_failure is not None: + return alias_failure organization: Final = await self.get_organization(organization_id) if isinstance(organization, Failure): return organization diff --git a/litellm/proxy/management_helpers/customer_upsert.py b/litellm/proxy/management_helpers/customer_upsert.py new file mode 100644 index 00000000000..495232195be --- /dev/null +++ b/litellm/proxy/management_helpers/customer_upsert.py @@ -0,0 +1,277 @@ +from dataclasses import dataclass +from typing import Final + +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 ( + CustomerHierarchy, + CustomerHierarchyStore, + ErrorCode, + Failure, + default_template_index, +) +from litellm.types.proxy.management_endpoints.customer_hierarchy import ( + Budget, + BudgetTemplate, + MemberCreate, + MemberPatch, + MemberUpsert, + MemberUpsertResponse, + Organization, + OrganizationPatch, + OrganizationUpsert, + OrganizationUpsertResponse, + Team, +) + + +@dataclass(frozen=True, slots=True) +class OrganizationOutcome: + data: OrganizationUpsertResponse + previous_alias: str | None + + +def matches_patch(current: BaseModel | None, patch: BaseModel | None) -> bool: + if current is None or patch is None: + return current is patch + values: Final = current.model_dump() + return all(values.get(key) == value for key, value in patch.model_dump(exclude_unset=True).items()) + + +async def organization_by_identity( + store: CustomerHierarchyStore, customer_org_id: str | None, alias: str | None +) -> Organization | Failure: + if customer_org_id is not None: + await store.lock_organization(customer_org_id) + return await store.get_organization(customer_org_id) + rows: Final = await store.organizations.find_many(where={"alias": alias}, take=2) + if len(rows) > 1: + return Failure(ErrorCode.ALIAS_ALREADY_EXISTS) + if not rows: + return Failure(ErrorCode.NOT_FOUND) + organization: Final = Organization.model_validate(rows[0].model_dump()) + await store.lock_organization(organization.customer_org_id) + return await store.get_organization(organization.customer_org_id) + + +def template_budget(template: BudgetTemplate) -> Budget: + return Budget.model_validate(template.model_dump(exclude_unset=True, exclude={"name"})) + + +def templates_match(organization: Organization, payload: OrganizationUpsert, default_index: int | None) -> bool: + if not payload.default_budgets: + return True + existing: Final = {template.name: template for template in organization.default_budgets} + return all( + template.name in existing and matches_patch(existing[template.name], template_budget(template)) + for template in payload.default_budgets + ) and ( + default_index is not None + and payload.default_budgets[default_index].name in existing + and existing[payload.default_budgets[default_index].name].budget_id == organization.default_budget_template_id + ) + + +async def write_templates( + store: CustomerHierarchyStore, + organization: Organization, + payload: OrganizationUpsert, + default_index: int, + actor: str, +) -> None: + existing: Final = {template.name: template.budget_id for template in organization.default_budgets} + written: Final = tuple( + [ + await store.update_budget(existing[template.name], template_budget(template), actor) + if template.name in existing + else await store.create_budget(template_budget(template), actor, template.name) + for template in payload.default_budgets + ] + ) + current_ids: Final = tuple(template.budget_id for template in organization.default_budgets) + await store.write_template_metadata( + organization.customer_org_id, + (*current_ids, *(budget_id for budget_id in written if budget_id not in current_ids)), + written[default_index], + actor, + ) + + +async def upsert_organization( + store: CustomerHierarchyStore, payload: OrganizationUpsert, actor: str +) -> OrganizationOutcome | Failure: + await store.lock_customer_identities() + default_index: Final = default_template_index(payload.default_budgets, payload.default_budget_template) + if isinstance(default_index, Failure): + return default_index + names: Final = tuple(template.name for template in payload.default_budgets) + if len(set(names)) != len(names): + return Failure(ErrorCode.INVALID_TEMPLATE_NAMES) + existing: Final = await organization_by_identity(store, payload.customer_org_id, payload.alias) + if isinstance(existing, Failure): + if existing.code is not ErrorCode.NOT_FOUND: + return existing + created: Final = await store.create_organization(payload, actor) + if isinstance(created, Failure): + return created + return OrganizationOutcome( + OrganizationUpsertResponse.model_validate({**created.model_dump(), "created": True}), None + ) + patch: Final = OrganizationPatch.model_validate( + payload.model_dump( + exclude_unset=True, exclude={"customer_org_id", "default_budgets", "default_budget_template"} + ) + ) + alias_failure: Final = await store.claim_organization_alias(payload.alias, existing.customer_org_id) + if alias_failure is not None: + return alias_failure + alias_unchanged: Final = "alias" not in patch.model_fields_set or existing.alias == patch.alias + budget_unchanged: Final = "budget" not in patch.model_fields_set or matches_patch(existing.budget, patch.budget) + if alias_unchanged and budget_unchanged and templates_match(existing, payload, default_index): + return OrganizationOutcome( + OrganizationUpsertResponse.model_validate({**existing.model_dump(), "created": False}), existing.alias + ) + updated: Final = await store.update_organization(existing.customer_org_id, patch, actor) + if isinstance(updated, Failure): + return updated + if default_index is not None: + await write_templates(store, updated, payload, default_index, actor) + result: Final = await store.get_organization(existing.customer_org_id) + if isinstance(result, Failure): + return result + return OrganizationOutcome( + OrganizationUpsertResponse.model_validate({**result.model_dump(), "created": False}), existing.alias + ) + + +async def customer_team( + store: CustomerHierarchyStore, organization_id: str, team_id: str | None, actor: str +) -> Team | Failure: + if team_id is not None: + explicit_team: Final = await store.get_team(team_id) + if isinstance(explicit_team, Failure): + return Failure(ErrorCode.PARENT_NOT_FOUND) + return ( + explicit_team + if explicit_team.parent_org_id == organization_id + else Failure(ErrorCode.CUSTOMER_ORG_MISMATCH) + ) + await store.teams.upsert( + where={"customer_team_id": default_team_id(organization_id)}, + data={ + "create": { + "customer_team_id": default_team_id(organization_id), + "alias": DEFAULT_TEAM_ALIAS, + "parent_org_id": organization_id, + "created_by": actor, + "updated_by": actor, + }, + "update": {}, + }, + ) + team: Final = await store.get_team(default_team_id(organization_id)) + if isinstance(team, Failure): + return team + return team if team.parent_org_id == organization_id else Failure(ErrorCode.CUSTOMER_ORG_MISMATCH) + + +async def upsert_member( + store: CustomerHierarchyStore, payload: MemberUpsert, actor: str +) -> MemberUpsertResponse | Failure: + await store.lock_customer_identities() + organization: Final = await organization_by_identity(store, payload.customer_org_id, payload.customer_org_alias) + if isinstance(organization, Failure): + return Failure(ErrorCode.PARENT_NOT_FOUND) if organization.code is ErrorCode.NOT_FOUND else organization + if payload.customer_org_alias is not None and organization.alias != payload.customer_org_alias: + return Failure(ErrorCode.CUSTOMER_ORG_MISMATCH) + existing: Final = await store.get_member(payload.customer_member_id) + if not isinstance(existing, Failure): + if existing.parent_org_id != organization.customer_org_id: + return Failure(ErrorCode.CUSTOMER_ORG_MISMATCH) + elif existing.code is not ErrorCode.NOT_FOUND: + return existing + patch: Final = MemberPatch.model_validate( + payload.model_dump( + exclude_unset=True, + exclude={"customer_member_id", "customer_org_id", "customer_org_alias", "parent_team_id"}, + ) + ) + requested_budget: Final = ( + patch.budget_id + if "budget_id" in patch.model_fields_set + else organization.default_budget_template_id + if isinstance(existing, Failure) + else existing.budget_id + ) + selected: Final = ( + requested_budget or organization.default_budget_template_id + if isinstance(existing, Failure) + else requested_budget + ) + if selected is not None and not any(template.budget_id == selected for template in organization.default_budgets): + return Failure(ErrorCode.TEMPLATE_NOT_IN_ORG) + override: Final = ( + patch.budget_override + if "budget_override" in patch.model_fields_set or isinstance(existing, Failure) + else existing.budget_override + ) + if selected is None and override is None: + return Failure(ErrorCode.BUDGET_REQUIRED) + team: Final = await customer_team(store, organization.customer_org_id, payload.parent_team_id, actor) + if isinstance(team, Failure): + return team + team_ids: Final = ( + (team.customer_team_id,) + if isinstance(existing, Failure) + else existing.parent_team_ids + if payload.parent_team_id is None + else (team.customer_team_id, *(tid for tid in existing.parent_team_ids if tid != team.customer_team_id)) + ) + if isinstance(existing, Failure): + member: Final = await store.create_member( + MemberCreate.model_validate( + { + **patch.model_dump(exclude_unset=True), + "customer_member_id": payload.customer_member_id, + "parent_team_ids": (team.customer_team_id,), + } + ), + actor, + ) + if isinstance(member, Failure): + return member + return MemberUpsertResponse.model_validate({**member.model_dump(), "created": True}) + alias_matches: Final = "alias" not in patch.model_fields_set or existing.alias == patch.alias + override_matches: Final = "budget_override" not in patch.model_fields_set or matches_patch( + existing.budget_override, patch.budget_override + ) + expiry_matches: Final = "budget_override_expires_at" not in patch.model_fields_set or ( + existing.budget_override_expires_at == patch.budget_override_expires_at + ) + if ( + existing.parent_team_ids == team_ids + and selected == existing.budget_id + and alias_matches + and override_matches + and expiry_matches + ): + return MemberUpsertResponse.model_validate({**existing.model_dump(), "created": False}) + await store.members.update( + where={"customer_member_id": payload.customer_member_id}, data={"parent_team_ids": team_ids} + ) + updated: Final = await store.update_member(payload.customer_member_id, patch, actor) + if isinstance(updated, Failure): + return updated + return MemberUpsertResponse.model_validate({**updated.model_dump(), "created": False}) + + +class CustomerUpsert: + def __init__(self, hierarchy: CustomerHierarchy) -> None: + self.hierarchy = hierarchy + + async def organization(self, payload: OrganizationUpsert, actor: str) -> OrganizationOutcome | Failure: + return await self.hierarchy.transaction(lambda store: upsert_organization(store, payload, actor)) + + async def member(self, payload: MemberUpsert, actor: str) -> MemberUpsertResponse | Failure: + return await self.hierarchy.transaction(lambda store: upsert_member(store, payload, actor)) diff --git a/litellm/types/proxy/management_endpoints/customer_hierarchy.py b/litellm/types/proxy/management_endpoints/customer_hierarchy.py index 029729e069c..aa977174bec 100644 --- a/litellm/types/proxy/management_endpoints/customer_hierarchy.py +++ b/litellm/types/proxy/management_endpoints/customer_hierarchy.py @@ -2,7 +2,7 @@ from collections.abc import Mapping from datetime import datetime from typing import Final, Literal, TypeAlias -from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, NonNegativeInt, field_validator +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, NonNegativeInt, field_validator, model_validator from litellm.models.budget import BudgetLimits from litellm.proxy.common_utils.timezone_utils import budget_duration_error @@ -67,6 +67,16 @@ class OrganizationCreate(OrganizationFields): default_budget_template: str | None = Field(default=None, min_length=1) +class OrganizationUpsert(OrganizationCreate): + customer_org_id: str | None = Field(default=None, min_length=1) + + @model_validator(mode="after") + def require_identity(self) -> "OrganizationUpsert": + if not self.customer_org_id and not self.alias: + raise ValueError("A customer org ID or alias is required for upsert") + return self + + class TeamPatch(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") @@ -96,6 +106,19 @@ class MemberCreate(MemberPatch): parent_team_ids: tuple[str, ...] = Field(min_length=1, description="Teams the member belongs to, first is default") +class MemberUpsert(MemberPatch): + customer_member_id: str = Field(min_length=1) + customer_org_id: str | None = None + customer_org_alias: str | None = None + parent_team_id: str | None = None + + @model_validator(mode="after") + def require_customer_org(self) -> "MemberUpsert": + if not self.customer_org_id and not self.customer_org_alias: + raise ValueError("A customer org ID or alias is required for customer member upsert") + return self + + class AuditFields(BaseModel): model_config = ConfigDict(frozen=True) @@ -139,6 +162,14 @@ class CustomerMember(AuditFields): effective_budget: Budget +class OrganizationUpsertResponse(Organization): + created: bool + + +class MemberUpsertResponse(CustomerMember): + created: bool + + class BudgetIncreaseRequestCreate(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") 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 9ebba5b46d0..250e7fb3ed8 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9717,6 +9717,26 @@ export interface paths { patch?: never; trace?: never; }; + "/management/v1/customer_members/upsert": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Upsert Customer Member + * @description Create or update a customer member by customer_member_id within a customer org identified by customer_org_id or customer_org_alias. If both are supplied they must match. A new customer member with no parent_team_id joins the customer org's default customer team, created if necessary. An existing customer member keeps its customer teams, and a supplied parent_team_id becomes its first customer team. Alias, budget_id, budget_override and budget_override_expires_at follow the existing customer member models. Omitted update fields retain their values. Moving between customer teams in the same customer org is allowed. Moving between customer orgs or selecting another customer org's customer team returns 409. Returns 200 with created=true for creation and created=false for updates. Repeated identical calls are no-ops. Proxy admins only. + */ + post: operations["upsert_customer_member_management_v1_customer_members_upsert_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/management/v1/customer_members/{customer_member_id}": { parameters: { query?: never; @@ -9809,6 +9829,26 @@ export interface paths { patch?: never; trace?: never; }; + "/management/v1/customer_organizations/upsert": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Upsert Customer Organization + * @description Create or update a customer org by customer_org_id, or by alias when no ID is supplied. Returns 200 with created=true for creation and created=false for updates. Repeated identical calls are no-ops. Supplied default_budgets are merged by template name, preserving existing template IDs and omitted templates. A single supplied template becomes the default; multiple templates require default_budget_template to name one. Omitted fields retain their values. Conflicting or ambiguous aliases return 409. Proxy admins only. + */ + post: operations["upsert_customer_organization_management_v1_customer_organizations_upsert_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/management/v1/customer_organizations/{customer_org_id}": { parameters: { query?: never; @@ -40699,6 +40739,66 @@ export interface components { /** Budget Override Expires At */ budget_override_expires_at?: string | null; }; + /** MemberUpsert */ + MemberUpsert: { + /** Alias */ + alias?: string | null; + /** Budget Id */ + budget_id?: string | null; + budget_override?: components["schemas"]["Budget"] | null; + /** Budget Override Expires At */ + budget_override_expires_at?: string | null; + /** Customer Member Id */ + customer_member_id: string; + /** Customer Org Alias */ + customer_org_alias?: string | null; + /** Customer Org Id */ + customer_org_id?: string | null; + /** Parent Team Id */ + parent_team_id?: string | null; + }; + /** MemberUpsertResponse */ + MemberUpsertResponse: { + /** Alias */ + alias?: string | null; + /** Budget Id */ + budget_id?: string | null; + budget_override?: components["schemas"]["Budget"] | null; + /** Budget Override Expires At */ + budget_override_expires_at?: string | null; + /** Budget Override Id */ + budget_override_id?: string | null; + /** Budget Reset At */ + budget_reset_at?: string | null; + /** Created */ + created: boolean; + /** + * Created At + * Format: date-time + */ + created_at: string; + /** Created By */ + created_by: string; + /** Customer Member Id */ + customer_member_id: string; + effective_budget: components["schemas"]["Budget"]; + /** Parent Org Id */ + parent_org_id: string; + /** Parent Team Ids */ + parent_team_ids: string[]; + /** + * Spend + * @default 0 + */ + spend: number; + /** + * Updated At + * Format: date-time + */ + updated_at: string; + /** Updated By */ + updated_by: string; + }; /** MemoryCreateRequest */ MemoryCreateRequest: { /** @@ -42741,6 +42841,67 @@ export interface components { /** Tpm Limit */ tpm_limit?: number | null; }; + /** OrganizationUpsert */ + OrganizationUpsert: { + /** Alias */ + alias?: string | null; + budget?: components["schemas"]["Budget"] | null; + /** Customer Org Id */ + customer_org_id?: string | null; + /** Default Budget Template */ + default_budget_template?: string | null; + /** + * Default Budgets + * @default [] + */ + default_budgets: components["schemas"]["BudgetTemplate"][]; + }; + /** OrganizationUpsertResponse */ + OrganizationUpsertResponse: { + /** Alias */ + alias?: string | null; + budget?: components["schemas"]["Budget"] | null; + /** Budget Id */ + budget_id?: string | null; + /** Budget Reset At */ + budget_reset_at?: string | null; + /** Created */ + created: boolean; + /** + * Created At + * Format: date-time + */ + created_at: string; + /** Created By */ + created_by: string; + /** Customer Org Id */ + customer_org_id: string; + /** Default Budget Template Id */ + default_budget_template_id?: string | null; + /** + * Default Budgets + * @default [] + */ + default_budgets: components["schemas"]["OrganizationBudget"][]; + /** + * Member Count + * @description Current registered members across all customer teams + * @default 0 + */ + member_count: number; + /** + * Spend + * @default 0 + */ + spend: number; + /** + * Updated At + * Format: date-time + */ + updated_at: string; + /** Updated By */ + updated_by: string; + }; /** * OutcomeExit * @description Indicates that the shell commands finished and returned an exit code. @@ -68677,6 +68838,39 @@ export interface operations { }; }; }; + upsert_customer_member_management_v1_customer_members_upsert_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["MemberUpsert"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["MemberUpsertResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_customer_member_management_v1_customer_members__customer_member_id__get: { parameters: { query?: never; @@ -68915,6 +69109,39 @@ export interface operations { }; }; }; + upsert_customer_organization_management_v1_customer_organizations_upsert_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["OrganizationUpsert"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["OrganizationUpsertResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_customer_organization_management_v1_customer_organizations__customer_org_id__get: { parameters: { query?: never; diff --git a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py index 78146687515..f7213be2eca 100644 --- a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py +++ b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py @@ -7,7 +7,7 @@ from fastapi import FastAPI, Request from fastapi.testclient import TestClient from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth -from litellm.proxy.auth.customer_hierarchy_auth import customer_member_cache_key +from litellm.proxy.auth.customer_hierarchy_auth import customer_member_cache_key, customer_org_alias_cache_key from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.list_api.common import ManagementProblem, problem_response @@ -31,6 +31,7 @@ def client( role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN, cache: UserApiKeyCache | None = None, database: Database | None = None, + actor: str = "admin", ) -> TestClient: app: Final = FastAPI() app.include_router(router) @@ -39,7 +40,7 @@ def client( async def problem(request: Request, exc: ManagementProblem): return problem_response(exc.problem) - auth: Final = UserAPIKeyAuth(user_id="admin", user_role=role) + 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-ok: in-memory Prisma edge @@ -173,6 +174,7 @@ def test_member_window_override_round_trip_and_clear(clear_limits: object) -> No "method,path,body", ( ("POST", "/customer_organizations", {}), + ("POST", "/customer_organizations/upsert", {"alias": "customer"}), ("GET", "/customer_organizations", None), ("GET", "/customer_organizations/id", None), ("PATCH", "/customer_organizations/id", {}), @@ -185,6 +187,7 @@ def test_member_window_override_round_trip_and_clear(clear_limits: object) -> No ("DELETE", "/customer_teams/id", None), ("DELETE", "/customer_teams/id?cascade=true", None), ("POST", "/customer_members", {"parent_team_ids": ["id"]}), + ("POST", "/customer_members/upsert", {"customer_member_id": "m", "customer_org_id": "id"}), ("GET", "/customer_members", None), ("GET", "/customer_members/id", None), ("PATCH", "/customer_members/id", {}), @@ -1002,6 +1005,271 @@ def family_default_budget_id(http: TestClient) -> str: return str(organization["default_budget_template_id"]) +@pytest.mark.parametrize("use_id", (True, False)) +def test_customer_org_upsert_repeats_without_new_budgets_and_updates_by_identity(use_id: bool) -> None: + http: Final = client() + identity: Final = {"customer_org_id": "org"} if use_id else {} + payload: Final = { + **identity, + "alias": "customer", + "budget": {"max_budget": 10}, + "default_budgets": [{"name": "standard", "max_budget": 2}], + } + first: Final = http.post("/management/v1/customer_organizations/upsert", json=payload) + assert first.status_code == 200, first.text + assert first.json()["created"] is True + second: Final = http.post("/management/v1/customer_organizations/upsert", json=payload) + assert second.status_code == 200, second.text + assert second.json() == {**first.json(), "created": False} + changed: Final = http.post( + "/management/v1/customer_organizations/upsert", + json={**payload, "budget": {"max_budget": 20}, "default_budgets": [{"name": "standard", "max_budget": 4}]}, + ) + assert changed.status_code == 200, changed.text + assert changed.json()["created"] is False + assert changed.json()["customer_org_id"] == first.json()["customer_org_id"] + assert changed.json()["budget_id"] == first.json()["budget_id"] + assert changed.json()["default_budget_template_id"] == first.json()["default_budget_template_id"] + assert changed.json()["budget"]["max_budget"] == 20 + assert changed.json()["default_budgets"][0]["max_budget"] == 4 + assert http.get("/management/v1/customer_organizations").json()["meta"]["total_count"] == 1 + + +@pytest.mark.parametrize("method", ("POST", "PATCH", "UPSERT")) +def test_duplicate_alias_is_rejected_without_changing_the_customer_org(method: str) -> None: + http: Final = client() + first: Final = http.post( + "/management/v1/customer_organizations", json={"customer_org_id": "first", "alias": "taken"} + ) + assert first.status_code == 201, first.text + other: Final = http.post( + "/management/v1/customer_organizations", json={"customer_org_id": "other", "alias": "other"} + ) + assert other.status_code == 201, other.text + requests: Final = { + "POST": ("POST", "/customer_organizations", {"customer_org_id": "new", "alias": "taken"}), + "PATCH": ("PATCH", "/customer_organizations/other", {"alias": "taken"}), + "UPSERT": ("POST", "/customer_organizations/upsert", {"customer_org_id": "other", "alias": "taken"}), + } + request_method, path, body = requests[method] + response: Final = http.request(request_method, "/management/v1" + path, json=body) + assert response.status_code == 409, response.text + assert response.json()["detail"] == "Customer org alias is already used by another customer org." + assert http.get("/management/v1/customer_organizations/other").json() == other.json() + assert http.get("/management/v1/customer_organizations").json()["meta"]["total_count"] == 2 + + +def test_customer_member_upsert_inherits_default_and_preserves_budget_ids_on_update() -> None: + http: Final = client() + family(http) + payload: Final = {"customer_member_id": "m", "customer_org_id": "org", "alias": "customer member"} + first: Final = http.post("/management/v1/customer_members/upsert", json=payload) + assert first.status_code == 200, first.text + assert first.json()["created"] is True + assert first.json()["parent_team_ids"] == ["org:default"] + assert first.json()["budget_id"] == family_default_budget_id(http) + repeated: Final = http.post("/management/v1/customer_members/upsert", json=payload) + assert repeated.status_code == 200, repeated.text + assert repeated.json() == {**first.json(), "created": False} + overridden: Final = http.post( + "/management/v1/customer_members/upsert", json={**payload, "budget_override": {"max_budget": 3}} + ) + assert overridden.status_code == 200, overridden.text + assert overridden.json()["effective_budget"]["max_budget"] == 3 + updated: Final = http.post( + "/management/v1/customer_members/upsert", json={**payload, "budget_override": {"max_budget": 5}} + ) + assert updated.status_code == 200, updated.text + assert updated.json()["budget_override_id"] == overridden.json()["budget_override_id"] + assert updated.json()["effective_budget"]["max_budget"] == 5 + cleared: Final = http.post("/management/v1/customer_members/upsert", json={**payload, "budget_override": None}) + assert cleared.status_code == 200, cleared.text + assert cleared.json()["budget_override_id"] is None + assert cleared.json()["effective_budget"]["max_budget"] == 100 + assert http.get("/management/v1/customer_members").json()["meta"]["total_count"] == 1 + + +def test_customer_member_upsert_accepts_customer_org_alias_and_explicit_customer_team() -> None: + http: Final = client() + family(http) + assert http.patch("/management/v1/customer_organizations/org", json={"alias": "customer"}).status_code == 200 + response: Final = http.post( + "/management/v1/customer_members/upsert", + json={"customer_member_id": "m", "customer_org_alias": "customer", "parent_team_id": "team"}, + ) + assert response.status_code == 200, response.text + assert response.json()["created"] is True + assert response.json()["parent_team_ids"] == ["team"] + repeated: Final = http.post( + "/management/v1/customer_members/upsert", json={"customer_member_id": "m", "customer_org_id": "org"} + ) + assert repeated.status_code == 200, repeated.text + assert repeated.json()["parent_team_ids"] == ["team"] + assert repeated.json()["created"] is False + + +@pytest.mark.parametrize("mismatch", ("customer_org", "customer_team", "customer_org_alias")) +def test_customer_member_upsert_rejects_cross_customer_org_identity(mismatch: str) -> None: + http: Final = client() + family(http) + assert http.post("/management/v1/customer_organizations", json={"customer_org_id": "other"}).status_code == 201 + assert ( + http.post( + "/management/v1/customer_teams", json={"customer_team_id": "foreign", "parent_org_id": "other"} + ).status_code + == 201 + ) + original: Final = http.post( + "/management/v1/customer_members", json={"customer_member_id": "m", "parent_team_ids": ["team"]} + ) + assert original.status_code == 201, original.text + edits: Final = { + "customer_org": {"customer_org_id": "other"}, + "customer_team": {"parent_team_id": "foreign"}, + "customer_org_alias": {"customer_org_alias": "wrong"}, + } + response: Final = http.post( + "/management/v1/customer_members/upsert", + json={"customer_member_id": "m", "customer_org_id": "org", **edits[mismatch]}, + ) + assert response.status_code == 409, response.text + assert response.json()["title"] == "Customer org mismatch" + assert http.get("/management/v1/customer_members/m").json() == original.json() + + +@pytest.mark.parametrize( + "resource,payload,status", + ( + ("customer_organizations", {}, 422), + ("customer_organizations", {"customer_org_id": "", "alias": "new"}, 422), + ("customer_members", {"customer_member_id": "m"}, 422), + ("customer_members", {"customer_org_id": "org"}, 422), + ("customer_members", {"customer_member_id": "m", "customer_org_id": "missing"}, 404), + ("customer_members", {"customer_member_id": "m", "customer_org_id": "org", "parent_team_id": "missing"}, 404), + ("customer_members", {"customer_member_id": "m", "customer_org_id": "org", "budget_id": "missing"}, 422), + ), +) +def test_customer_upsert_missing_identity_or_parent_is_rejected( + resource: str, payload: dict[str, object], status: int +) -> None: + http: Final = client() + family(http) + response: Final = http.post(f"/management/v1/{resource}/upsert", json=payload) + assert response.status_code == status, response.text + assert http.get("/management/v1/customer_members").json()["meta"]["total_count"] == 0 + + +def test_customer_upsert_rejects_missing_budget_without_creating_a_default_customer_team() -> None: + http: Final = client() + assert http.post("/management/v1/customer_organizations", json={"customer_org_id": "org"}).status_code == 201 + response: Final = http.post( + "/management/v1/customer_members/upsert", json={"customer_member_id": "m", "customer_org_id": "org"} + ) + assert response.status_code == 422, response.text + assert http.get("/management/v1/customer_teams").json()["meta"]["total_count"] == 0 + + +def test_customer_upsert_invalid_templates_do_not_partially_update_customer_org() -> None: + http: Final = client() + original: Final = family(http) + invalid: Final = http.post( + "/management/v1/customer_organizations/upsert", + json={ + "customer_org_id": "org", + "alias": "changed", + "budget": {"max_budget": 1}, + "default_budgets": [{"name": "duplicate"}, {"name": "duplicate"}], + "default_budget_template": "duplicate", + }, + ) + assert invalid.status_code == 422, invalid.text + assert http.get("/management/v1/customer_organizations/org").json() == original + + +def test_customer_upserts_invalidate_customer_member_and_customer_org_alias_caches() -> None: + cache: Final = UserApiKeyCache() + http: Final = client(cache=cache) + family(http) + assert http.patch("/management/v1/customer_organizations/org", json={"alias": "old"}).status_code == 200 + assert ( + http.post( + "/management/v1/customer_members", json={"customer_member_id": "m", "parent_team_ids": ["team"]} + ).status_code + == 201 + ) + member_key: Final = customer_member_cache_key("m") + for key in (member_key, customer_org_alias_cache_key("old"), customer_org_alias_cache_key("new")): + asyncio.run(cache.async_set_cache(key=key, value="stale customer hierarchy")) + updated: Final = http.post( + "/management/v1/customer_organizations/upsert", json={"customer_org_id": "org", "alias": "new"} + ) + assert updated.status_code == 200, updated.text + for key in (member_key, customer_org_alias_cache_key("old"), customer_org_alias_cache_key("new")): + assert cache.get_cache(key) is None + asyncio.run(cache.async_set_cache(key=member_key, value="stale customer hierarchy")) + member: Final = http.post( + "/management/v1/customer_members/upsert", + json={"customer_org_alias": "new", "customer_member_id": "m", "budget_override": {"max_budget": 1}}, + ) + assert member.status_code == 200, member.text + assert cache.get_cache(member_key) is None + + +def test_customer_org_upsert_rejects_ambiguous_legacy_alias() -> None: + database: Final = Database() + for identifier in ("first", "second"): + asyncio.run( + database.litellm_customerorganizationtable.create( + {"customer_org_id": identifier, "alias": "legacy", "created_by": "admin", "updated_by": "admin"} + ) + ) + http: Final = client(database=database) + response: Final = http.post("/management/v1/customer_organizations/upsert", json={"alias": "legacy"}) + assert response.status_code == 409, response.text + assert http.get("/management/v1/customer_organizations").json()["meta"]["total_count"] == 2 + + +def test_customer_org_upsert_retry_does_not_rewrite_audit_fields() -> None: + database: Final = Database() + payload: Final = {"alias": "customer", "budget": {"max_budget": 10}} + original: Final = client(database=database, actor="creator").post( + "/management/v1/customer_organizations/upsert", json=payload + ) + assert original.status_code == 200, original.text + repeated: Final = client(database=database, actor="retry").post( + "/management/v1/customer_organizations/upsert", json=payload + ) + assert repeated.status_code == 200, repeated.text + assert repeated.json() == {**original.json(), "created": False} + + +def test_customer_org_upsert_merges_templates_without_replacing_customer_member_budgets() -> None: + http: Final = client() + original: 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 + payload: Final = {"customer_org_id": "org", "default_budgets": [{"name": "premium", "max_budget": 200}]} + updated: Final = http.post("/management/v1/customer_organizations/upsert", json=payload) + assert updated.status_code == 200, updated.text + original_ids: Final = {template["name"]: template["budget_id"] for template in original["default_budgets"]} + updated_ids: Final = {template["name"]: template["budget_id"] for template in updated.json()["default_budgets"]} + assert set(updated_ids) == {"standard", "small", "premium"} + assert {name: updated_ids[name] for name in original_ids} == original_ids + assert updated.json()["default_budget_template_id"] == updated_ids["premium"] + assert http.get("/management/v1/customer_members/m").json() == member.json() + repeated: Final = http.post("/management/v1/customer_organizations/upsert", json=payload) + assert repeated.status_code == 200, repeated.text + assert repeated.json() == updated.json() + inherited: Final = http.post( + "/management/v1/customer_members/upsert", json={"customer_member_id": "new", "customer_org_id": "org"} + ) + assert inherited.status_code == 200, inherited.text + assert inherited.json()["budget_id"] == updated_ids["premium"] + assert inherited.json()["effective_budget"]["max_budget"] == 200 + + def test_resubmitting_a_templates_duration_keeps_its_reset_clock() -> None: """An edit that repeats the duration a template already runs must not push its window out: the reset job owes that window a clearing, and a re-stamped clock would skip it. A changed duration starts a new window.""" diff --git a/tests/unit/proxy/management_helpers/test_customer_hierarchy.py b/tests/unit/proxy/management_helpers/test_customer_hierarchy.py index 1c8e2fab826..91c21271cd3 100644 --- a/tests/unit/proxy/management_helpers/test_customer_hierarchy.py +++ b/tests/unit/proxy/management_helpers/test_customer_hierarchy.py @@ -122,6 +122,15 @@ class Table: self.rows = {**self.rows, row_id: row} return row + async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> Record: + existing: Final = await self.find_unique(where) + values: Final = TypeAdapter(Mapping[str, object]).validate_python(data["update" if existing else "create"]) + if existing is None: + return await self.create(values) + updated: Final = await self.update(where, values) + assert updated is not None + return updated + def matches(self, row: Record, where: Mapping[str, object]) -> bool: values: Final = row.model_dump() for key, value in where.items(): @@ -244,6 +253,8 @@ class Database: return any(row.model_dump().get(field) == budget_id for row in table.rows.values()) async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + if "pg_advisory_xact_lock" in query: + return () if "SELECT EXISTS" in query: return ({"referenced": False},) if "FOR UPDATE" in query: