From 6d15d68bc84e9068ac8243fa085026b8ebfc42e1 Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Sat, 14 Nov 2026 01:12:45 +0800 Subject: [PATCH] feat(proxy): register an unknown customer organization alias as a new organization charged for the request (#51637) --- litellm/proxy/auth/customer_hierarchy_auth.py | 72 +++++++++++++++--- .../proxy/auth/customer_member_auto_create.py | 27 ++++++- .../auth/test_customer_hierarchy_auth.py | 76 ++++++++++++++++++- .../auth/test_customer_member_auto_create.py | 21 ++++- 4 files changed, 179 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/auth/customer_hierarchy_auth.py b/litellm/proxy/auth/customer_hierarchy_auth.py index 0ebcf588bbc..530410117c7 100644 --- a/litellm/proxy/auth/customer_hierarchy_auth.py +++ b/litellm/proxy/auth/customer_hierarchy_auth.py @@ -22,7 +22,7 @@ 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.customer_member_auto_create import ensure_customer_member, ensure_customer_organization from litellm.proxy.auth.customer_spend import current_customer_spend from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast @@ -45,7 +45,12 @@ from litellm.types.proxy.auth.customer_hierarchy import ( CustomerScopeAuth, CustomerScopeKey, ) -from litellm.types.proxy.management_endpoints.customer_hierarchy import INCREASE_REQUEST_ROUTE, Budget, CustomerMember +from litellm.types.proxy.management_endpoints.customer_hierarchy import ( + INCREASE_REQUEST_ROUTE, + Budget, + CustomerMember, + Organization, +) if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -491,14 +496,12 @@ def _untracked(message: str, header: str) -> None: verbose_proxy_logger.warning("Customer %s not tracked, %s", header, message) -async def _organization_named_by_alias( - customer_org_alias: str, prisma_client: "PrismaClient | None", user_api_key_cache: UserApiKeyCache +OrganizationRegistrar: TypeAlias = Callable[[str, "PrismaClient"], Awaitable[Organization | Failure]] + + +def _organization_from_lookup( + customer_org_alias: str, organization: CustomerHierarchyAuth | AmbiguousOrganizationAlias | None ) -> CustomerHierarchyAuth | None: - 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: return _untracked(f"organization alias not found: {customer_org_alias}", CUSTOMER_ORG_ALIAS_HEADER) if isinstance(organization, AmbiguousOrganizationAlias): @@ -511,6 +514,45 @@ async def _organization_named_by_alias( return organization +async def _register_organization( + customer_org_alias: str, + prisma_client: "PrismaClient", + user_api_key_cache: UserApiKeyCache, + register: OrganizationRegistrar, +) -> CustomerHierarchyAuth | None: + """An alias the proxy has never seen becomes an organization with that id and alias, so a customer's first + request is charged to it instead of running untracked.""" + if not _registrable(customer_org_alias): + verbose_proxy_logger.warning( + "Could not register customer organization: alias is not a valid customer identifier" + ) + return _untracked(f"organization alias not found: {customer_org_alias}", CUSTOMER_ORG_ALIAS_HEADER) + registered: Final = await register(customer_org_alias, prisma_client) + if isinstance(registered, Failure): + verbose_proxy_logger.warning( + "Could not register customer organization %s: %s", customer_org_alias, registered.code.value + ) + return _untracked(f"organization alias not found: {customer_org_alias}", CUSTOMER_ORG_ALIAS_HEADER) + return _organization_from_lookup( + customer_org_alias, + await get_customer_organization_hierarchy(customer_org_alias, prisma_client, user_api_key_cache), + ) + + +async def _organization_named_by_alias( + customer_org_alias: str, + prisma_client: "PrismaClient | None", + user_api_key_cache: UserApiKeyCache, + register: OrganizationRegistrar, +) -> CustomerHierarchyAuth | None: + if prisma_client is None: + return _untracked(f"organization alias not found: {customer_org_alias}", CUSTOMER_ORG_ALIAS_HEADER) + found: Final = await get_customer_organization_hierarchy(customer_org_alias, prisma_client, user_api_key_cache) + if found is None: + return await _register_organization(customer_org_alias, prisma_client, user_api_key_cache, register) + return _organization_from_lookup(customer_org_alias, found) + + MemberRegistrar: TypeAlias = Callable[[str, str, "PrismaClient"], Awaitable[CustomerMember | Failure]] _MEMBER_ID: Final[TypeAdapter[str]] = TypeAdapter(CustomerIdentifier) @@ -583,13 +625,16 @@ async def load_customer_hierarchy( user_api_key_cache: UserApiKeyCache, register: MemberRegistrar = ensure_customer_member, key_metadata: Mapping[str, object] = MappingProxyType({}), + register_organization: OrganizationRegistrar = ensure_customer_organization, ) -> CustomerHierarchyAuth | None: """The hierarchy the headers name, or None when the proxy knows none of the ids, in which case the request runs - untracked. A member header next to an alias header must name a member of that organization; a member the + untracked. An alias the proxy has never seen is registered as a new organization first. 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, and when that fails the organization alone is charged. A key pinned to another organization is refused before any row is created.""" organization: Final = ( - await _organization_named_by_alias(identity.customer_org_alias, prisma_client, user_api_key_cache) + await _organization_named_by_alias( + identity.customer_org_alias, prisma_client, user_api_key_cache, register_organization + ) if identity.customer_org_alias is not None else None ) @@ -634,13 +679,16 @@ async def resolve_customer_hierarchy( route: str, key_metadata: Mapping[str, object] | None = None, register: MemberRegistrar = ensure_customer_member, + register_organization: OrganizationRegistrar = ensure_customer_organization, ) -> CustomerHierarchyAuth | None: """Load the hierarchy the headers name and admit the request against it, or None when no id is known. ``request_model`` is the model the route will actually serve, which may differ from the body.""" from litellm.proxy.proxy_server import get_current_spend metadata: Final = key_metadata or {} - cached: Final = await load_customer_hierarchy(identity, prisma_client, user_api_key_cache, register, metadata) + cached: Final = await load_customer_hierarchy( + identity, prisma_client, user_api_key_cache, register, metadata, register_organization + ) if cached is None: return None hierarchy: Final = await current_customer_spend(cached, get_current_spend) diff --git a/litellm/proxy/auth/customer_member_auto_create.py b/litellm/proxy/auth/customer_member_auto_create.py index 06936d26043..34f1b91cfc4 100644 --- a/litellm/proxy/auth/customer_member_auto_create.py +++ b/litellm/proxy/auth/customer_member_auto_create.py @@ -10,12 +10,20 @@ 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 +from litellm.types.proxy.management_endpoints.customer_hierarchy import ( + CustomerMember, + MemberCreate, + Organization, + OrganizationCreate, + Team, + TeamCreate, +) if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient AUTO_CREATE_ACTOR: Final = "x-litellm-customer-member-id" +AUTO_CREATE_ORG_ACTOR: Final = "x-litellm-customer-org-alias" DEFAULT_TEAM_ALIAS: Final = "default" @@ -59,3 +67,20 @@ async def ensure_customer_member( customer_org_id, ) return created + + +async def ensure_customer_organization( + customer_org_alias: str, prisma_client: "PrismaClient" +) -> Organization | Failure: + """The organization row carrying ``customer_org_alias``, created with the alias as its id when absent, so a + customer's first request is charged to it. Two workers registering the same alias at once both end up with + the one row.""" + hierarchy: Final = CustomerHierarchy(prisma_client) + created: Final = await hierarchy.create_organization( + OrganizationCreate(customer_org_id=customer_org_alias, alias=customer_org_alias), AUTO_CREATE_ORG_ACTOR + ) + if isinstance(created, Failure) and created.code is ErrorCode.ALREADY_EXISTS: + return await hierarchy.get_organization(customer_org_alias) + if not isinstance(created, Failure): + verbose_proxy_logger.info("Registered customer organization %s", customer_org_alias) + return created diff --git a/tests/unit/proxy/auth/test_customer_hierarchy_auth.py b/tests/unit/proxy/auth/test_customer_hierarchy_auth.py index e8d10a6ef04..f313050cd40 100644 --- a/tests/unit/proxy/auth/test_customer_hierarchy_auth.py +++ b/tests/unit/proxy/auth/test_customer_hierarchy_auth.py @@ -40,7 +40,7 @@ from litellm.types.proxy.auth.customer_hierarchy import ( CustomerRequestIdentity, CustomerScopeAuth, ) -from litellm.types.proxy.management_endpoints.customer_hierarchy import CustomerMember +from litellm.types.proxy.management_endpoints.customer_hierarchy import CustomerMember, Organization from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) @@ -464,12 +464,39 @@ class _Registrar: ) +class _OrganizationRegistrar: + """Stands in for ensure_customer_organization: records the alias and adds the organization row, or fails.""" + + def __init__(self, organizations: _OrganizationTable, failure: Failure | None = None) -> None: + self.organizations = organizations + self.failure = failure + self.calls: tuple[str, ...] = () + + async def __call__(self, alias: str, prisma_client: object) -> Organization | Failure: + self.calls = (*self.calls, alias) + if self.failure is not None: + return self.failure + self.organizations.rows = (*self.organizations.rows, _organization(alias, org_id=alias)) + return Organization.model_validate( + { + "customer_org_id": alias, + "alias": alias, + "member_count": 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, key_metadata: Mapping[str, object] | None = None, + register_organization: _OrganizationRegistrar | None = None, ): return resolve_customer_hierarchy( identity=identity, @@ -480,6 +507,9 @@ def _resolve( route=route, key_metadata=key_metadata, register=register if register is not None else _Registrar(_MemberTable({}), Failure(ErrorCode.NOT_FOUND)), + register_organization=register_organization + if register_organization is not None + else _OrganizationRegistrar(_OrganizationTable(), Failure(ErrorCode.NOT_FOUND)), ) @@ -614,8 +644,48 @@ async def test_an_alias_next_to_a_member_must_name_that_members_organization() - @pytest.mark.asyncio -async def test_an_unknown_alias_runs_untracked_and_a_shared_alias_is_rejected_as_ambiguous() -> None: - assert await _resolve(_Prisma(_MemberTable({})), CustomerRequestIdentity(customer_org_alias="ghost")) is None +async def test_an_unknown_alias_becomes_a_new_organization_charged_for_the_request() -> None: + organizations: Final = _OrganizationTable(()) + registrar: Final = _OrganizationRegistrar(organizations) + prisma: Final = _Prisma(_MemberTable({}), organizations) + + hierarchy: Final = await _resolve( + prisma, CustomerRequestIdentity(customer_org_alias="ghost"), register_organization=registrar + ) + + assert registrar.calls == ("ghost",) + assert hierarchy is not None + assert (hierarchy.customer_org_id, hierarchy.customer_team_id, hierarchy.customer_member_id) == ( + "ghost", + None, + None, + ) + assert tuple((scope.key, scope.entity_id) for scope in hierarchy.scopes) == (("customer_org", "ghost"),) + + again: Final = await _resolve( + prisma, CustomerRequestIdentity(customer_org_alias="ghost"), register_organization=registrar + ) + assert again is not None and registrar.calls == ("ghost",) + + +@pytest.mark.asyncio +async def test_an_alias_that_cannot_be_registered_runs_untracked_and_a_shared_alias_is_rejected_as_ambiguous() -> None: + failing: Final = _OrganizationRegistrar(_OrganizationTable(()), Failure(ErrorCode.ALREADY_EXISTS)) + untracked: Final = await _resolve( + _Prisma(_MemberTable({})), CustomerRequestIdentity(customer_org_alias="ghost"), register_organization=failing + ) + assert (untracked, failing.calls) == (None, ("ghost",)) + + oversized: Final = _OrganizationRegistrar(_OrganizationTable(())) + assert ( + await _resolve( + _Prisma(_MemberTable({})), + CustomerRequestIdentity(customer_org_alias="g" * 257), + register_organization=oversized, + ) + is None + ) + assert oversized.calls == () shared: Final = _Prisma( _MemberTable({}), _OrganizationTable((_organization("acme"), _organization("acme", org_id="org-2"))) diff --git a/tests/unit/proxy/auth/test_customer_member_auto_create.py b/tests/unit/proxy/auth/test_customer_member_auto_create.py index 696e1e0b4c2..2320b05d3cf 100644 --- a/tests/unit/proxy/auth/test_customer_member_auto_create.py +++ b/tests/unit/proxy/auth/test_customer_member_auto_create.py @@ -2,7 +2,11 @@ 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.auth.customer_member_auto_create import ( + default_team_id, + ensure_customer_member, + ensure_customer_organization, +) from litellm.proxy.management_helpers.customer_hierarchy import CustomerHierarchy, ErrorCode, Failure from litellm.types.proxy.management_endpoints.customer_hierarchy import ( BudgetTemplate, @@ -75,3 +79,18 @@ async def test_an_organization_without_a_default_budget_cannot_register_members( @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) + + +@pytest.mark.asyncio +async def test_an_unknown_alias_becomes_an_organization_once_and_a_known_one_is_reused() -> None: + database: Final = Database() + + created: Final = await ensure_customer_organization("acme", _prisma(database)) + again: Final = await ensure_customer_organization("acme", _prisma(database)) + + assert not isinstance(created, Failure) and not isinstance(again, Failure) + assert (created.customer_org_id, created.alias, again.customer_org_id) == ("acme", "acme", "acme") + listed: Final = await CustomerHierarchy(_prisma(database), clock=lambda: NOW).list_organizations( + page=1, page_size=10 + ) + assert not isinstance(listed, Failure) and len(listed.data) == 1