diff --git a/litellm/proxy/auth/customer_hierarchy_auth.py b/litellm/proxy/auth/customer_hierarchy_auth.py index 51c6344465b..b9d12b00c1a 100644 --- a/litellm/proxy/auth/customer_hierarchy_auth.py +++ b/litellm/proxy/auth/customer_hierarchy_auth.py @@ -2,39 +2,52 @@ A request names its customer member with the ``x-litellm-customer-member-id`` header. The member row carries its team and that team's organization, so one lookup yields every level whose budget and rate -limits apply to the request. +limits apply to the request. A request may instead, or additionally, name the organization by its alias +with ``x-litellm-customer-org-alias``: alone it charges the organization level only, next to the member +header it must agree with the member's organization. """ -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, TypeAlias from fastapi import status from pydantic import BaseModel, ConfigDict, TypeAdapter import litellm from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.auth.auth_checks import can_object_call_model from litellm.proxy.auth.auth_object_prefetch import AUTH_OBJECTS_TARGET +from litellm.proxy.auth.customer_member_auto_create import ensure_customer_member from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl from litellm.proxy.management_helpers.customer_hierarchy import ( + Failure, PrismaRecord, effective_member_budget, related_budget, utc_now, ) from litellm.repositories.prisma_protocols import TableActions -from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth, CustomerScopeKey -from litellm.types.proxy.management_endpoints.customer_hierarchy import Budget +from litellm.types.proxy.auth.customer_hierarchy import ( + CustomerHierarchyAuth, + CustomerRequestIdentity, + CustomerScopeAuth, + CustomerScopeKey, +) +from litellm.types.proxy.management_endpoints.customer_hierarchy import Budget, CustomerMember if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient from litellm.router_api import RegistryReader CUSTOMER_MEMBER_ID_HEADER: Final = "x-litellm-customer-member-id" +CUSTOMER_ORG_ALIAS_HEADER: Final = "x-litellm-customer-org-alias" +_ORGANIZATION_INCLUDE: Final = {"budget": True} _MEMBER_INCLUDE: Final = { "budget": True, "budget_override": True, @@ -71,21 +84,51 @@ class _MemberRow(BaseModel): parent_team: _TeamRow +@dataclass(frozen=True, slots=True) +class AmbiguousOrganizationAlias: + """More than one customer organization carries the alias, so the request cannot be charged to one of them.""" + + alias: str + + def customer_member_cache_key(customer_member_id: str) -> str: return f"customer_member_id:{customer_member_id}" -def customer_member_id_from_headers(headers: Mapping[str, object] | None) -> str | None: +def customer_org_alias_cache_key(customer_org_alias: str) -> str: + return f"customer_org_alias:{customer_org_alias}" + + +def _header_value(headers: Mapping[str, object] | None, header: str) -> str | None: if headers is None: return None values: Final = tuple( value.strip() for name, value in headers.items() - if name.lower() == CUSTOMER_MEMBER_ID_HEADER and isinstance(value, str) and value.strip() + if name.lower() == header and isinstance(value, str) and value.strip() ) return values[0] if values else None +def customer_member_id_from_headers(headers: Mapping[str, object] | None) -> str | None: + return _header_value(headers, CUSTOMER_MEMBER_ID_HEADER) + + +def customer_org_alias_from_headers(headers: Mapping[str, object] | None) -> str | None: + return _header_value(headers, CUSTOMER_ORG_ALIAS_HEADER) + + +def customer_identity_from_headers(headers: Mapping[str, object] | None) -> CustomerRequestIdentity | None: + """The customer the request named, or None when it carried neither customer header.""" + identity: Final = CustomerRequestIdentity( + customer_member_id=customer_member_id_from_headers(headers), + customer_org_alias=customer_org_alias_from_headers(headers), + ) + if identity.customer_member_id is None and identity.customer_org_alias is None: + return None + return identity + + def _scope(key: CustomerScopeKey, entity_id: str, spend: float, budget: Budget | None) -> CustomerScopeAuth: if budget is None: return CustomerScopeAuth(key=key, entity_id=entity_id, spend=spend) @@ -132,10 +175,30 @@ def customer_hierarchy_from_member_row(row: Mapping[str, object], now: datetime) ) +def customer_hierarchy_from_organization_row(row: Mapping[str, object], now: datetime) -> CustomerHierarchyAuth: + """The organization scope alone, for a request that named an organization alias and no member.""" + organization: Final = _OrganizationRow.model_validate(row) + return CustomerHierarchyAuth( + customer_org_id=organization.customer_org_id, + scopes=( + _scope( + "customer_org", + organization.customer_org_id, + organization.spend, + related_budget(organization.budget, now), + ), + ), + ) + + def _member_table(prisma_client: "PrismaClient") -> TableActions[PrismaRecord]: return prisma_client.db.litellm_customermembertable # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportAny] # prisma tables are generated at build time +def _organization_table(prisma_client: "PrismaClient") -> TableActions[PrismaRecord]: + return prisma_client.db.litellm_customerorganizationtable # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportAny] # prisma tables are generated at build time + + @with_service_target(AUTH_OBJECTS_TARGET) async def get_customer_hierarchy( customer_member_id: str, @@ -163,6 +226,34 @@ async def get_customer_hierarchy( return hierarchy +@with_service_target(AUTH_OBJECTS_TARGET) +async def get_customer_organization_hierarchy( + customer_org_alias: str, + prisma_client: "PrismaClient", + user_api_key_cache: UserApiKeyCache, + now: datetime | None = None, +) -> CustomerHierarchyAuth | AmbiguousOrganizationAlias | None: + """The organization-only hierarchy behind an alias, None when no organization carries it.""" + cache_key: Final = customer_org_alias_cache_key(customer_org_alias) + cached: Final = await user_api_key_cache.async_get_cache(key=cache_key, model_type=CustomerHierarchyAuth) + if cached is not None: + return cached + rows: Final = await _organization_table(prisma_client).find_many( + where={"alias": customer_org_alias}, include=_ORGANIZATION_INCLUDE, take=2 + ) + if not rows: + return None + if len(rows) > 1: + return AmbiguousOrganizationAlias(customer_org_alias) + hierarchy: Final = customer_hierarchy_from_organization_row( + rows[0].model_dump(), now if now is not None else utc_now() + ) + await user_api_key_cache.async_set_cache( + key=cache_key, value=hierarchy, ttl=get_management_object_ttl(user_api_key_cache) + ) + return hierarchy + + def _cache_ttl(row: Mapping[str, object], now: datetime, default_ttl: float) -> float: """A cached hierarchy may not outlive the member's budget override, or an expired override would keep granting its limits until the management-object TTL ran out.""" @@ -199,30 +290,128 @@ def check_customer_hierarchy_model_access( ) -async def resolve_customer_hierarchy( +def _not_found(message: str, param: str) -> ProxyException: + return ProxyException( + message=f"Authentication Error - {message}", + type=ProxyErrorTypes.auth_error, + param=param, + code=status.HTTP_401_UNAUTHORIZED, + ) + + +async def _organization_named_by_alias( + customer_org_alias: str, prisma_client: "PrismaClient | None", user_api_key_cache: UserApiKeyCache +) -> CustomerHierarchyAuth: + organization: Final = ( + await get_customer_organization_hierarchy(customer_org_alias, prisma_client, user_api_key_cache) + if prisma_client is not None + else None + ) + if organization is None: + raise _not_found(f"customer organization alias not found: {customer_org_alias}", CUSTOMER_ORG_ALIAS_HEADER) + if isinstance(organization, AmbiguousOrganizationAlias): + raise ProxyException( + message=f"Customer organization alias is not unique: {customer_org_alias}", + type=ProxyErrorTypes.bad_request_error, + param=CUSTOMER_ORG_ALIAS_HEADER, + code=status.HTTP_400_BAD_REQUEST, + ) + return organization + + +MemberRegistrar: TypeAlias = Callable[[str, str, "PrismaClient"], Awaitable[CustomerMember | Failure]] + + +async def _register_member( customer_member_id: str, + organization: CustomerHierarchyAuth | None, + prisma_client: "PrismaClient", + user_api_key_cache: UserApiKeyCache, + register: MemberRegistrar, +) -> CustomerHierarchyAuth | None: + """A member the request named for the first time joins the aliased organization's default team. Without an + alias there is no organization to place it in, and the member stays unknown.""" + if organization is None: + return None + registered: Final = await register(organization.customer_org_id, customer_member_id, prisma_client) + if isinstance(registered, Failure): + verbose_proxy_logger.warning( + "Could not register customer member %s in organization %s: %s", + customer_member_id, + organization.customer_org_id, + registered.code.value, + ) + return None + return await get_customer_hierarchy(customer_member_id, prisma_client, user_api_key_cache) + + +async def _hierarchy_named_by_member( + customer_member_id: str, + organization: CustomerHierarchyAuth | None, + prisma_client: "PrismaClient | None", + user_api_key_cache: UserApiKeyCache, + register: MemberRegistrar, +) -> CustomerHierarchyAuth: + if prisma_client is None: + raise _not_found(f"customer member not found: {customer_member_id}", CUSTOMER_MEMBER_ID_HEADER) + known: Final = await get_customer_hierarchy(customer_member_id, prisma_client, user_api_key_cache) + hierarchy: Final = ( + known + if known is not None + else await _register_member(customer_member_id, organization, prisma_client, user_api_key_cache, register) + ) + if hierarchy is None: + raise _not_found(f"customer member not found: {customer_member_id}", CUSTOMER_MEMBER_ID_HEADER) + return hierarchy + + +async def load_customer_hierarchy( + identity: CustomerRequestIdentity, + prisma_client: "PrismaClient | None", + user_api_key_cache: UserApiKeyCache, + register: MemberRegistrar = ensure_customer_member, +) -> CustomerHierarchyAuth: + """The hierarchy the headers name. An unknown alias is refused rather than treated as unrestricted: the caller + asked for a budgeted identity that does not exist. A member header next to an alias header must name a member + of that organization; a member the organization has never seen is registered in its default team first.""" + organization: Final = ( + await _organization_named_by_alias(identity.customer_org_alias, prisma_client, user_api_key_cache) + if identity.customer_org_alias is not None + else None + ) + if identity.customer_member_id is None: + if organization is None: + raise _not_found("request named no customer member or organization", CUSTOMER_MEMBER_ID_HEADER) + return organization + hierarchy: Final = await _hierarchy_named_by_member( + identity.customer_member_id, organization, prisma_client, user_api_key_cache, register + ) + if organization is not None and organization.customer_org_id != hierarchy.customer_org_id: + raise ProxyException( + message=( + f"Customer member {identity.customer_member_id} does not belong to the organization " + f"aliased {identity.customer_org_alias}" + ), + type=ProxyErrorTypes.auth_error, + param=CUSTOMER_ORG_ALIAS_HEADER, + code=status.HTTP_403_FORBIDDEN, + ) + return hierarchy + + +async def resolve_customer_hierarchy( + identity: CustomerRequestIdentity, prisma_client: "PrismaClient | None", user_api_key_cache: UserApiKeyCache, llm_router: "RegistryReader | None", request_model: str | Sequence[str] | None, route: str, key_metadata: Mapping[str, object] | None = None, + register: MemberRegistrar = ensure_customer_member, ) -> CustomerHierarchyAuth: - """Load the hierarchy the header names and admit the request against it. An unknown member is refused - rather than treated as unrestricted: the caller asked for a budgeted identity that does not exist. - ``request_model`` is the model the route will actually serve, which may differ from the body.""" - hierarchy: Final = ( - await get_customer_hierarchy(customer_member_id, prisma_client, user_api_key_cache) - if prisma_client is not None - else None - ) - if hierarchy is None: - raise ProxyException( - message=f"Authentication Error - customer member not found: {customer_member_id}", - type=ProxyErrorTypes.auth_error, - param=CUSTOMER_MEMBER_ID_HEADER, - code=status.HTTP_401_UNAUTHORIZED, - ) + """Load the hierarchy the headers name and admit the request against it. ``request_model`` is the model the + route will actually serve, which may differ from the body.""" + hierarchy: Final = await load_customer_hierarchy(identity, prisma_client, user_api_key_cache, register) check_customer_hierarchy_key_scope(hierarchy, key_metadata or {}) if not RouteChecks.is_info_route(route): check_customer_hierarchy_budget(hierarchy) @@ -250,3 +439,13 @@ def check_customer_hierarchy_key_scope(hierarchy: CustomerHierarchyAuth, key_met async def forget_customer_hierarchies(customer_member_ids: Sequence[str], user_api_key_cache: UserApiKeyCache) -> None: """Drop the cached hierarchies of members whose own, team, org or template budget changed so auth reloads them.""" await evict_and_broadcast(tuple(customer_member_cache_key(m) for m in customer_member_ids), user_api_key_cache) + + +@with_service_target(AUTH_OBJECTS_TARGET) +async def forget_customer_organization_aliases( + aliases: Sequence[str | None], user_api_key_cache: UserApiKeyCache +) -> None: + """Drop the cached organization hierarchies behind aliases whose organization changed so request auth reloads them.""" + keys: Final = tuple(customer_org_alias_cache_key(alias) for alias in dict.fromkeys(aliases) if alias) + if keys: + await evict_and_broadcast(keys, user_api_key_cache) diff --git a/litellm/proxy/auth/customer_member_auto_create.py b/litellm/proxy/auth/customer_member_auto_create.py new file mode 100644 index 00000000000..6bafcaba785 --- /dev/null +++ b/litellm/proxy/auth/customer_member_auto_create.py @@ -0,0 +1,61 @@ +"""Registration of a customer member the request named but the proxy has never seen. + +Customers send a shared service key with ``x-litellm-customer-org-alias`` and ``x-litellm-customer-member-id`` +on every call, and their users appear for the first time on their first request. The member is created under +the organization's ``default`` team, so it inherits the organization's default budget. An organization without +a default budget cannot host an unregistered member, because a member must have some budget to be charged to. +""" + +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.management_helpers.customer_hierarchy import CustomerHierarchy, ErrorCode, Failure +from litellm.types.proxy.management_endpoints.customer_hierarchy import CustomerMember, MemberCreate, Team, TeamCreate + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +AUTO_CREATE_ACTOR: Final = "x-litellm-customer-member-id" +DEFAULT_TEAM_ALIAS: Final = "default" + + +def default_team_id(customer_org_id: str) -> str: + return f"{customer_org_id}:{DEFAULT_TEAM_ALIAS}" + + +async def _default_team(hierarchy: CustomerHierarchy, customer_org_id: str) -> Team | Failure: + team_id: Final = default_team_id(customer_org_id) + existing: Final = await hierarchy.get_team(team_id) + if not isinstance(existing, Failure): + return existing + created: Final = await hierarchy.create_team( + TeamCreate(customer_team_id=team_id, alias=DEFAULT_TEAM_ALIAS, parent_org_id=customer_org_id), + AUTO_CREATE_ACTOR, + ) + if isinstance(created, Failure) and created.code is ErrorCode.ALREADY_EXISTS: + return await hierarchy.get_team(team_id) + return created + + +async def ensure_customer_member( + customer_org_id: str, customer_member_id: str, prisma_client: "PrismaClient" +) -> 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) + team: Final = await _default_team(hierarchy, customer_org_id) + if isinstance(team, Failure): + return team + created: Final = await hierarchy.create_member( + MemberCreate(customer_member_id=customer_member_id, parent_team_id=team.customer_team_id), AUTO_CREATE_ACTOR + ) + if isinstance(created, Failure) and created.code is ErrorCode.ALREADY_EXISTS: + return await hierarchy.get_member(customer_member_id) + if not isinstance(created, Failure): + verbose_proxy_logger.info( + "Registered customer member %s in team %s of organization %s", + customer_member_id, + team.customer_team_id, + customer_org_id, + ) + return created diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 28d4d5efcd5..f9b642a99cf 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -131,7 +131,7 @@ from litellm.proxy.auth.auth_utils import ( route_in_additonal_public_routes, ) from litellm.proxy.auth.customer_hierarchy_auth import ( - customer_member_id_from_headers, + customer_identity_from_headers, resolve_customer_hierarchy, ) from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler @@ -4544,14 +4544,14 @@ async def _authorize_authenticated_request( user_model, ) - customer_member_id: Final = customer_member_id_from_headers( + customer_identity: Final = customer_identity_from_headers( safe_get_request_headers(request) # pyright: ignore[reportUnknownArgumentType] # request headers are untyped ) user_api_key_auth_obj.customer_hierarchy = None # rebind-ok: request state, never inherited from a cached key - if customer_member_id is not None: + if customer_identity is not None: with tracer.trace("litellm.proxy.auth.resolve_customer_hierarchy"): # pyright: ignore[reportUnknownMemberType] # tracer is the untyped proxy tracer user_api_key_auth_obj.customer_hierarchy = await resolve_customer_hierarchy( # rebind-ok: the gate updates the request-owned auth object in place - customer_member_id=customer_member_id, + identity=customer_identity, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, llm_router=llm_router, diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 98817dac3f1..87a94e121a6 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -3579,8 +3579,8 @@ class DBSpendUpdateWriter: ) -> None: if prisma_client is None or customer_hierarchy is None: return - member_id: Final = customer_hierarchy.customer_member_id - team_id: Final = customer_hierarchy.customer_team_id + member_id: Final = customer_hierarchy.customer_member_id or "" + team_id: Final = customer_hierarchy.customer_team_id or "" org_id: Final = customer_hierarchy.customer_org_id base: Final = await self._common_add_spend_log_transaction_to_daily_transaction( # pyright: ignore[reportUnknownMemberType] # existing helper accepts legacy untyped payloads {**payload, "customer_member_id": member_id, "customer_team_id": team_id, "customer_org_id": org_id}, diff --git a/litellm/proxy/management_endpoints/management_v1/customers.py b/litellm/proxy/management_endpoints/management_v1/customers.py index eb7e727fe87..fcf4a0992b4 100644 --- a/litellm/proxy/management_endpoints/management_v1/customers.py +++ b/litellm/proxy/management_endpoints/management_v1/customers.py @@ -5,7 +5,7 @@ from fastapi import APIRouter, Depends, Query, Request, Response from typing_extensions import assert_never from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth -from litellm.proxy.auth.customer_hierarchy_auth import forget_customer_hierarchies +from litellm.proxy.auth.customer_hierarchy_auth import forget_customer_hierarchies, forget_customer_organization_aliases 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 ( @@ -192,8 +192,10 @@ async def update_customer_organization( hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)], cache: AuthCache, ) -> Organization: + previous: Final = resolve(await hierarchy.get_organization(customer_org_id)) organization: Final = resolve(await hierarchy.update_organization(customer_org_id, data, actor)) await forget_customer_hierarchies(resolve(await hierarchy.member_ids_in_organization(customer_org_id)), cache) + await forget_customer_organization_aliases((previous.alias, organization.alias), cache) return organization @@ -206,8 +208,11 @@ async def update_customer_organization( async def delete_customer_organization( customer_org_id: str, hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)], + cache: AuthCache, ) -> Response: + previous: Final = resolve(await hierarchy.get_organization(customer_org_id)) resolve(await hierarchy.delete_organization(customer_org_id)) + await forget_customer_organization_aliases((previous.alias,), cache) return Response(status_code=204) diff --git a/litellm/types/proxy/auth/customer_hierarchy.py b/litellm/types/proxy/auth/customer_hierarchy.py index c131df439f5..a1141c9d113 100644 --- a/litellm/types/proxy/auth/customer_hierarchy.py +++ b/litellm/types/proxy/auth/customer_hierarchy.py @@ -1,5 +1,6 @@ from collections.abc import Mapping from dataclasses import dataclass +from types import MappingProxyType from typing import Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError @@ -9,6 +10,14 @@ from litellm.types.integrations.slack_alerting import Litellm_EntityType CustomerScopeKey: TypeAlias = Literal["customer_org", "customer_team", "customer_member"] +@dataclass(frozen=True, slots=True) +class CustomerRequestIdentity: + """What a request's headers said about the customer it runs for: the member id, the organization alias, or both.""" + + customer_member_id: str | None = None + customer_org_alias: str | None = None + + class CustomerScopeAuth(BaseModel): """One level of the customer hierarchy a request runs under, with the limits its budget row grants.""" @@ -26,23 +35,26 @@ class CustomerScopeAuth(BaseModel): class CustomerHierarchyAuth(BaseModel): + """The levels a request is charged to. A request that named only an organization alias has no team or member.""" + model_config = ConfigDict(frozen=True) - customer_member_id: str - customer_team_id: str customer_org_id: str + customer_team_id: str | None = None + customer_member_id: str | None = None scopes: tuple[CustomerScopeAuth, ...] @dataclass(frozen=True, slots=True) class CustomerHierarchyIds: - """The three rows a request's spend is charged to, as the request metadata recorded them.""" + """The rows a request's spend is charged to, as the request metadata recorded them. Team and member are absent + when the request named only an organization.""" - customer_member_id: str - customer_team_id: str customer_org_id: str + customer_team_id: str | None = None + customer_member_id: str | None = None - def entities(self) -> Mapping[str, str]: + def entities(self) -> Mapping[str, str | None]: return { "customer_member_id": self.customer_member_id, "customer_team_id": self.customer_team_id, @@ -50,34 +62,44 @@ class CustomerHierarchyIds: } def spend_entities(self) -> tuple[tuple[Litellm_EntityType, str], ...]: - return ( + levels: Final = ( (Litellm_EntityType.CUSTOMER_ORG, self.customer_org_id), (Litellm_EntityType.CUSTOMER_TEAM, self.customer_team_id), (Litellm_EntityType.CUSTOMER_MEMBER, self.customer_member_id), ) + return tuple((entity_type, entity_id) for entity_type, entity_id in levels if entity_id) _FIELDS: Final = TypeAdapter(Mapping[str, object]) +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) def _mapping(value: object) -> Mapping[str, object]: try: return _FIELDS.validate_python(value) except ValidationError: - return {} + return _NO_FIELDS def _stamped_metadata(litellm_params: Mapping[str, object]) -> Mapping[str, object]: """The proxy stamps auth ids into ``litellm_metadata`` on routes whose body owns ``metadata``, else into ``metadata``.""" buckets: Final = tuple(_mapping(litellm_params.get(name)) for name in ("litellm_metadata", "metadata")) - return next((bucket for bucket in buckets if "user_api_key_customer_member_id" in bucket), {}) + return next((bucket for bucket in buckets if "user_api_key_customer_org_id" in bucket), _NO_FIELDS) + + +def _stamped_id(metadata: Mapping[str, object], level: str) -> str | None: + value: Final = metadata.get(f"user_api_key_customer_{level}_id") + return value if isinstance(value, str) and value else None def customer_hierarchy_ids_from_request(kwargs: Mapping[str, object]) -> CustomerHierarchyIds | None: - """The ids auth stamped into the request's proxy metadata, or None when the request named no customer member.""" + """The ids auth stamped into the request's proxy metadata, or None when the request named no customer.""" metadata: Final = _stamped_metadata(_mapping(kwargs.get("litellm_params"))) - ids: Final = tuple(metadata.get(f"user_api_key_customer_{level}_id") for level in ("member", "team", "org")) - if not all(isinstance(value, str) and value for value in ids): + org_id: Final = _stamped_id(metadata, "org") + if org_id is None: return None - member_id, team_id, org_id = ids - return CustomerHierarchyIds(str(member_id), str(team_id), str(org_id)) + return CustomerHierarchyIds( + customer_org_id=org_id, + customer_team_id=_stamped_id(metadata, "team"), + customer_member_id=_stamped_id(metadata, "member"), + ) 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 d0ec4789a5a..c6a636d8c50 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -33144,14 +33144,17 @@ export interface components { */ type: "custom"; }; - /** CustomerHierarchyAuth */ + /** + * CustomerHierarchyAuth + * @description The levels a request is charged to. A request that named only an organization alias has no team or member. + */ CustomerHierarchyAuth: { /** Customer Member Id */ - customer_member_id: string; + customer_member_id?: string | null; /** Customer Org Id */ customer_org_id: string; /** Customer Team Id */ - customer_team_id: string; + customer_team_id?: string | null; /** Scopes */ scopes: components["schemas"]["CustomerScopeAuth"][]; }; diff --git a/tests/unit/proxy/auth/test_customer_hierarchy_auth.py b/tests/unit/proxy/auth/test_customer_hierarchy_auth.py index 84f0b86f219..9efb9167ecd 100644 --- a/tests/unit/proxy/auth/test_customer_hierarchy_auth.py +++ b/tests/unit/proxy/auth/test_customer_hierarchy_auth.py @@ -1,24 +1,34 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import Final import pytest import litellm -from litellm.caching.caching import DualCache from litellm.proxy._types import ProxyException from litellm.proxy.auth.customer_hierarchy_auth import ( CUSTOMER_MEMBER_ID_HEADER, + CUSTOMER_ORG_ALIAS_HEADER, + AmbiguousOrganizationAlias, check_customer_hierarchy_budget, check_customer_hierarchy_key_scope, check_customer_hierarchy_model_access, customer_hierarchy_from_member_row, + customer_identity_from_headers, customer_member_id_from_headers, + forget_customer_organization_aliases, get_customer_hierarchy, + get_customer_organization_hierarchy, resolve_customer_hierarchy, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth +from litellm.proxy.management_helpers.customer_hierarchy import ErrorCode, Failure +from litellm.types.proxy.auth.customer_hierarchy import ( + CustomerHierarchyAuth, + CustomerRequestIdentity, + CustomerScopeAuth, +) +from litellm.types.proxy.management_endpoints.customer_hierarchy import CustomerMember NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) @@ -178,20 +188,87 @@ class _Record: return self.values +class _OrganizationTable: + def __init__(self, rows: Sequence[Mapping[str, object]] = ()) -> None: + self.rows = rows + self.lookups = 0 + + async def find_many( + self, where: Mapping[str, object], include: Mapping[str, object] | None = None, take: int | None = None + ): + self.lookups += 1 + matching: Final = [_Record(row) for row in self.rows if row.get("alias") == where["alias"]] + return matching[:take] if take is not None else matching + + class _Db: - def __init__(self, table: _MemberTable) -> None: + def __init__(self, table: _MemberTable, organizations: _OrganizationTable) -> None: self.litellm_customermembertable = table + self.litellm_customerorganizationtable = organizations class _Prisma: - def __init__(self, table: _MemberTable) -> None: - self.db = _Db(table) + def __init__(self, table: _MemberTable, organizations: _OrganizationTable | None = None) -> None: + self.db = _Db(table, organizations if organizations is not None else _OrganizationTable()) + + +def _organization( + alias: str, org_id: str = "org-1", spend: float = 0.0, max_budget: float | None = None +) -> Mapping[str, object]: + budget: Final = {"max_budget": max_budget} if max_budget is not None else None + return {"customer_org_id": org_id, "alias": alias, "spend": spend, "budget": budget} + + +class _Registrar: + """Stands in for ensure_customer_member: records the ask and adds the member row, or fails.""" + + def __init__(self, members: _MemberTable, failure: Failure | None = None) -> None: + self.members = members + self.failure = failure + self.calls: tuple[tuple[str, str], ...] = () + + async def __call__(self, org_id: str, member_id: str, prisma_client: object) -> CustomerMember | Failure: + self.calls = (*self.calls, (org_id, member_id)) + if self.failure is not None: + return self.failure + self.members.rows = { + **self.members.rows, + member_id: {**_row(template={"max_budget": 10.0}), "customer_member_id": member_id}, + } + return CustomerMember.model_validate( + { + "customer_member_id": member_id, + "parent_team_id": "team-1", + "effective_budget": {"max_budget": 10.0}, + "created_at": NOW, + "created_by": "x", + "updated_at": NOW, + "updated_by": "x", + } + ) + + +def _resolve( + prisma: _Prisma, + identity: CustomerRequestIdentity, + route: str = "/chat/completions", + register: _Registrar | None = None, +): + return resolve_customer_hierarchy( + identity=identity, + prisma_client=prisma, # pyright: ignore[reportArgumentType] # test double + user_api_key_cache=UserApiKeyCache(), + llm_router=None, + request_model="gpt-4o", + route=route, + register=register if register is not None else _Registrar(_MemberTable({}), Failure(ErrorCode.NOT_FOUND)), + ) @pytest.mark.asyncio async def test_hierarchy_is_loaded_once_and_then_served_from_the_cache() -> None: table: Final = _MemberTable({"member-1": _row(template={"max_budget": 10.0})}) - cache: Final = UserApiKeyCache(DualCache()) + cache: Final = UserApiKeyCache() prisma: Final = _Prisma(table) first: Final = await get_customer_hierarchy("member-1", prisma, cache, now=NOW) # pyright: ignore[reportArgumentType] # test double @@ -204,10 +281,10 @@ async def test_hierarchy_is_loaded_once_and_then_served_from_the_cache() -> None @pytest.mark.asyncio async def test_resolution_refuses_an_unknown_member_and_a_proxy_without_a_database() -> None: - cache: Final = UserApiKeyCache(DualCache()) + cache: Final = UserApiKeyCache() with pytest.raises(ProxyException) as missing: await resolve_customer_hierarchy( - customer_member_id="ghost", + identity=CustomerRequestIdentity(customer_member_id="ghost"), prisma_client=_Prisma(_MemberTable({})), # pyright: ignore[reportArgumentType] # test double user_api_key_cache=cache, llm_router=None, @@ -217,7 +294,7 @@ async def test_resolution_refuses_an_unknown_member_and_a_proxy_without_a_databa assert missing.value.code == "401" with pytest.raises(ProxyException): await resolve_customer_hierarchy( - customer_member_id="member-1", + identity=CustomerRequestIdentity(customer_member_id="member-1"), prisma_client=None, user_api_key_cache=cache, llm_router=None, @@ -228,11 +305,11 @@ async def test_resolution_refuses_an_unknown_member_and_a_proxy_without_a_databa @pytest.mark.asyncio async def test_resolution_enforces_budget_on_model_routes_but_not_on_info_routes() -> None: - cache: Final = UserApiKeyCache(DualCache()) + cache: Final = UserApiKeyCache() prisma: Final = _Prisma(_MemberTable({"member-1": _row(template={"max_budget": 1.0}, spend=(0.0, 0.0, 1.0))})) with pytest.raises(litellm.BudgetExceededError): await resolve_customer_hierarchy( - customer_member_id="member-1", + identity=CustomerRequestIdentity(customer_member_id="member-1"), prisma_client=prisma, # pyright: ignore[reportArgumentType] # test double user_api_key_cache=cache, llm_router=None, @@ -240,7 +317,7 @@ async def test_resolution_enforces_budget_on_model_routes_but_not_on_info_routes route="/chat/completions", ) hierarchy: Final = await resolve_customer_hierarchy( - customer_member_id="member-1", + identity=CustomerRequestIdentity(customer_member_id="member-1"), prisma_client=prisma, # pyright: ignore[reportArgumentType] # test double user_api_key_cache=cache, llm_router=None, @@ -257,3 +334,132 @@ def test_a_key_pinned_to_another_organization_cannot_act_for_the_member() -> Non assert denied.value.code == "403" check_customer_hierarchy_key_scope(hierarchy, {"customer_org_id": "org-1"}) check_customer_hierarchy_key_scope(hierarchy, {}) + + +@pytest.mark.parametrize( + "headers, expected", + [ + ({"X-LiteLLM-Customer-Org-Alias": " acme "}, CustomerRequestIdentity(customer_org_alias="acme")), + ( + {CUSTOMER_ORG_ALIAS_HEADER: "acme", CUSTOMER_MEMBER_ID_HEADER: "member-1"}, + CustomerRequestIdentity(customer_member_id="member-1", customer_org_alias="acme"), + ), + ({CUSTOMER_MEMBER_ID_HEADER: "member-1"}, CustomerRequestIdentity(customer_member_id="member-1")), + ({CUSTOMER_ORG_ALIAS_HEADER: " ", "x-litellm-end-user-id": "end-user"}, None), + (None, None), + ], +) +def test_the_request_identity_is_read_from_both_customer_headers( + headers: Mapping[str, object] | None, expected: CustomerRequestIdentity | None +) -> None: + assert customer_identity_from_headers(headers) == expected + + +@pytest.mark.asyncio +async def test_an_alias_alone_charges_only_the_organization_and_enforces_its_budget() -> None: + prisma: Final = _Prisma(_MemberTable({}), _OrganizationTable((_organization("acme", max_budget=5.0, spend=1.0),))) + hierarchy: Final = await _resolve(prisma, CustomerRequestIdentity(customer_org_alias="acme")) + + assert (hierarchy.customer_org_id, hierarchy.customer_team_id, hierarchy.customer_member_id) == ( + "org-1", + None, + None, + ) + assert tuple((scope.key, scope.entity_id, scope.max_budget, scope.spend) for scope in hierarchy.scopes) == ( + ("customer_org", "org-1", 5.0, 1.0), + ) + + exhausted: Final = _Prisma( + _MemberTable({}), _OrganizationTable((_organization("acme", max_budget=5.0, spend=5.0),)) + ) + with pytest.raises(litellm.BudgetExceededError) as exc: + await _resolve(exhausted, CustomerRequestIdentity(customer_org_alias="acme")) + assert (exc.value.entity_type, exc.value.entity_id) == ("customer_org", "org-1") + + +@pytest.mark.asyncio +async def test_an_alias_next_to_a_member_must_name_that_members_organization() -> None: + members: Final = _MemberTable({"member-1": _row(template={"max_budget": 10.0})}) + organizations: Final = _OrganizationTable((_organization("acme"), _organization("other", org_id="org-2"))) + prisma: Final = _Prisma(members, organizations) + + matching: Final = await _resolve( + prisma, CustomerRequestIdentity(customer_member_id="member-1", customer_org_alias="acme") + ) + assert (matching.customer_org_id, matching.customer_member_id) == ("org-1", "member-1") + assert tuple(scope.key for scope in matching.scopes) == ("customer_org", "customer_team", "customer_member") + + with pytest.raises(ProxyException) as denied: + await _resolve(prisma, CustomerRequestIdentity(customer_member_id="member-1", customer_org_alias="other")) + assert (denied.value.code, denied.value.param) == ("403", CUSTOMER_ORG_ALIAS_HEADER) + + +@pytest.mark.asyncio +async def test_an_unknown_alias_is_refused_and_a_shared_alias_is_rejected_as_ambiguous() -> None: + with pytest.raises(ProxyException) as missing: + await _resolve(_Prisma(_MemberTable({})), CustomerRequestIdentity(customer_org_alias="ghost")) + assert (missing.value.code, missing.value.param) == ("401", CUSTOMER_ORG_ALIAS_HEADER) + + shared: Final = _Prisma( + _MemberTable({}), _OrganizationTable((_organization("acme"), _organization("acme", org_id="org-2"))) + ) + with pytest.raises(ProxyException) as ambiguous: + await _resolve(shared, CustomerRequestIdentity(customer_org_alias="acme")) + assert ambiguous.value.code == "400" + + +@pytest.mark.asyncio +async def test_an_organization_alias_is_loaded_once_then_served_from_the_cache_until_forgotten() -> None: + organizations: Final = _OrganizationTable((_organization("acme"),)) + prisma: Final = _Prisma(_MemberTable({}), organizations) + cache: Final = UserApiKeyCache() + + first: Final = await get_customer_organization_hierarchy("acme", prisma, cache, now=NOW) # pyright: ignore[reportArgumentType] # test double + second: Final = await get_customer_organization_hierarchy("acme", prisma, cache, now=NOW) # pyright: ignore[reportArgumentType] # test double + assert isinstance(first, CustomerHierarchyAuth) and second == first + assert organizations.lookups == 1 + + await forget_customer_organization_aliases(("acme", None), cache) + await get_customer_organization_hierarchy("acme", prisma, cache, now=NOW) # pyright: ignore[reportArgumentType] # test double + assert organizations.lookups == 2 + assert await get_customer_organization_hierarchy("ghost", prisma, cache, now=NOW) is None # pyright: ignore[reportArgumentType] # test double + shared: Final = _Prisma( + _MemberTable({}), _OrganizationTable((_organization("acme"), _organization("acme", org_id="org-2"))) + ) + assert isinstance( + await get_customer_organization_hierarchy("acme", shared, UserApiKeyCache(), now=NOW), # pyright: ignore[reportArgumentType] # test double + AmbiguousOrganizationAlias, + ) + + +@pytest.mark.asyncio +async def test_a_member_the_aliased_organization_has_never_seen_is_registered_there_and_served() -> None: + members: Final = _MemberTable({}) + registrar: Final = _Registrar(members) + prisma: Final = _Prisma(members, _OrganizationTable((_organization("acme"),))) + + hierarchy: Final = await _resolve( + prisma, CustomerRequestIdentity(customer_member_id="user-7", customer_org_alias="acme"), register=registrar + ) + + assert registrar.calls == (("org-1", "user-7"),) + assert (hierarchy.customer_org_id, hierarchy.customer_member_id) == ("org-1", "user-7") + assert _scopes(hierarchy)["customer_member"].max_budget == 10.0 + + +@pytest.mark.asyncio +async def test_an_unknown_member_stays_refused_without_an_alias_or_when_registration_fails() -> None: + members: Final = _MemberTable({}) + registrar: Final = _Registrar(members) + with pytest.raises(ProxyException) as no_alias: + await _resolve(_Prisma(members), CustomerRequestIdentity(customer_member_id="user-7"), register=registrar) + assert (no_alias.value.code, registrar.calls) == ("401", ()) + + failing: Final = _Registrar(members, Failure(ErrorCode.BUDGET_REQUIRED)) + prisma: Final = _Prisma(members, _OrganizationTable((_organization("acme"),))) + with pytest.raises(ProxyException) as failed: + await _resolve( + prisma, CustomerRequestIdentity(customer_member_id="user-7", customer_org_alias="acme"), register=failing + ) + assert (failed.value.code, failing.calls) == ("401", (("org-1", "user-7"),)) + assert "user-7" not in members.rows diff --git a/tests/unit/proxy/auth/test_customer_member_auto_create.py b/tests/unit/proxy/auth/test_customer_member_auto_create.py new file mode 100644 index 00000000000..7933e05e4de --- /dev/null +++ b/tests/unit/proxy/auth/test_customer_member_auto_create.py @@ -0,0 +1,77 @@ +from typing import TYPE_CHECKING, Final, cast + +import pytest + +from litellm.proxy.auth.customer_member_auto_create import default_team_id, ensure_customer_member +from litellm.proxy.management_helpers.customer_hierarchy import CustomerHierarchy, ErrorCode, Failure +from litellm.types.proxy.management_endpoints.customer_hierarchy import ( + BudgetTemplate, + CustomerMember, + OrganizationCreate, + TeamCreate, +) +from tests.unit.proxy.management_helpers.test_customer_hierarchy import NOW, Database + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +def _prisma(database: Database) -> "PrismaClient": + return cast("PrismaClient", database) # cast-ok: in-memory Prisma edge + + +async def _organization(database: Database, org_id: str, default_budget: bool = True) -> None: + templates: Final = (BudgetTemplate(name="std", max_budget=10.0, is_default=True),) 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" + ) + assert not isinstance(created, Failure) + + +@pytest.mark.asyncio +async def test_a_new_member_joins_the_organizations_default_team_on_the_organizations_default_budget() -> None: + database: Final = Database() + await _organization(database, "org-1") + + member: Final = await ensure_customer_member("org-1", "user-7", _prisma(database)) + + assert isinstance(member, CustomerMember) + assert (member.customer_member_id, member.parent_team_id) == ("user-7", default_team_id("org-1")) + assert member.effective_budget.max_budget == 10.0 + team: Final = database.litellm_customerteamtable.rows[default_team_id("org-1")].model_dump() + assert (team["parent_org_id"], team["alias"]) == ("org-1", "default") + + +@pytest.mark.asyncio +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) + 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" + ) + assert not isinstance(existing_team, Failure) + + first: Final = await ensure_customer_member("org-1", "user-7", _prisma(database)) + second: Final = await ensure_customer_member("org-1", "user-7", _prisma(database)) + + assert isinstance(first, CustomerMember) and isinstance(second, CustomerMember) + assert first.parent_team_id == second.parent_team_id == default_team_id("org-1") + assert tuple(database.litellm_customermembertable.rows) == ("user-7",) + assert tuple(database.litellm_customerteamtable.rows) == (default_team_id("org-1"),) + + +@pytest.mark.asyncio +async def test_an_organization_without_a_default_budget_cannot_register_members() -> None: + database: Final = Database() + await _organization(database, "org-1", default_budget=False) + + result: Final = await ensure_customer_member("org-1", "user-7", _prisma(database)) + + assert result == Failure(ErrorCode.BUDGET_REQUIRED) + assert "user-7" not in database.litellm_customermembertable.rows + + +@pytest.mark.asyncio +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) diff --git a/tests/unit/types/proxy/auth/test_customer_hierarchy.py b/tests/unit/types/proxy/auth/test_customer_hierarchy.py index a6d72cdc3b1..6b9fa4046b3 100644 --- a/tests/unit/types/proxy/auth/test_customer_hierarchy.py +++ b/tests/unit/types/proxy/auth/test_customer_hierarchy.py @@ -10,10 +10,26 @@ _IDS: Final = { def test_ids_are_read_from_whichever_metadata_bucket_the_proxy_stamped() -> None: - expected: Final = CustomerHierarchyIds("m1", "t1", "o1") + expected: Final = CustomerHierarchyIds(customer_org_id="o1", customer_team_id="t1", customer_member_id="m1") assert customer_hierarchy_ids_from_request({"litellm_params": {"metadata": _IDS}}) == expected assert ( customer_hierarchy_ids_from_request({"litellm_params": {"metadata": {"note": "x"}, "litellm_metadata": _IDS}}) == expected ) assert customer_hierarchy_ids_from_request({"litellm_params": {"metadata": {"note": "x"}}}) is None + + +def test_a_request_that_named_only_an_organization_charges_the_organization_alone() -> None: + ids: Final = customer_hierarchy_ids_from_request( + { + "litellm_params": { + "metadata": {"user_api_key_customer_org_id": "o1", "user_api_key_customer_member_id": None} + } + } + ) + assert ids == CustomerHierarchyIds(customer_org_id="o1") + assert ids is not None + assert tuple(entity_id for _, entity_id in ids.spend_entities()) == ("o1",) + assert ids.entities() == {"customer_member_id": None, "customer_team_id": None, "customer_org_id": "o1"} + member_only: Final = {"litellm_params": {"metadata": {"user_api_key_customer_member_id": "m1"}}} + assert customer_hierarchy_ids_from_request(member_only) is None