feat(proxy): register an unknown customer organization alias as a new organization charged for the request (#51637)

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)

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

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

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