From 27e9a04d62f5d9b92a69704c9195754c30bf6ecf Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Fri, 13 Nov 2026 13:15:45 +0800 Subject: [PATCH] feat(proxy): register unknown customer members in the aliased organization's default team (#51546) * feat(proxy): name the customer organization by alias with x-litellm-customer-org-alias A request may carry x-litellm-customer-org-alias next to, or instead of, x-litellm-customer-member-id. Alone, the alias charges and enforces the organization level only. Next to the member header it must name the member's organization, else 403. Unknown alias is 401, an alias shared by several organizations is 400. Organization alias changes evict the cached lookup. Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * chore(ui): regenerate schema.d.ts for optional customer team and member ids Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * feat(proxy): register unknown customer members in the aliased organization's default team A request carrying x-litellm-customer-org-alias and an x-litellm-customer-member-id the proxy has never seen creates the member under the organization's default team, on the organization's default budget, and serves the request. Without an alias, or when the organization has no default budget, the member stays unknown and the request is still refused. Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(proxy): declare the member registrar callable as a TypeAlias Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm/proxy/auth/customer_hierarchy_auth.py | 245 ++++++++++++++++-- .../proxy/auth/customer_member_auto_create.py | 61 +++++ litellm/proxy/auth/user_api_key_auth.py | 8 +- litellm/proxy/db/db_spend_update_writer.py | 4 +- .../management_v1/customers.py | 7 +- .../types/proxy/auth/customer_hierarchy.py | 50 +++- .../src/lib/http/schema.d.ts | 9 +- .../auth/test_customer_hierarchy_auth.py | 232 ++++++++++++++++- .../auth/test_customer_member_auto_create.py | 77 ++++++ .../proxy/auth/test_customer_hierarchy.py | 18 +- 10 files changed, 650 insertions(+), 61 deletions(-) create mode 100644 litellm/proxy/auth/customer_member_auto_create.py create mode 100644 tests/unit/proxy/auth/test_customer_member_auto_create.py 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