scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 175

feat(proxy): register unknown customer members in the aliased organization's default team (#51546)

Browse source 
* 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>
This commit is contained in:
ishaan-j-berri 2026-11-13 13:15:45 +08:00 • committed by GitHub
parent e2dfd3dad1
commit 27e9a04d62
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 650 additions and 61 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

245
litellm/proxy/auth/customer_hierarchy_auth.py
View file

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

61
litellm/proxy/auth/customer_member_auto_create.py Normal file
View file

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

8
litellm/proxy/auth/user_api_key_auth.py
View file

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

4
litellm/proxy/db/db_spend_update_writer.py
View file

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

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

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

50
litellm/types/proxy/auth/customer_hierarchy.py
View file

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

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

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

232
tests/unit/proxy/auth/test_customer_hierarchy_auth.py
View file

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

77
tests/unit/proxy/auth/test_customer_member_auto_create.py Normal file
View file

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

18
tests/unit/types/proxy/auth/test_customer_hierarchy.py
View file

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