diff --git a/litellm/proxy/auth/customer_hierarchy_auth.py b/litellm/proxy/auth/customer_hierarchy_auth.py index 61c225f0afa..0ebcf588bbc 100644 --- a/litellm/proxy/auth/customer_hierarchy_auth.py +++ b/litellm/proxy/auth/customer_hierarchy_auth.py @@ -10,10 +10,11 @@ header it must agree with the member's organization. from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Final, TypeAlias from fastapi import HTTPException, status -from pydantic import BaseModel, ConfigDict, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm._internal_context import with_service_target @@ -25,6 +26,7 @@ from litellm.proxy.auth.customer_member_auto_create import ensure_customer_membe 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 +from litellm.proxy.common_utils.customer_usage import CustomerIdentifier from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl from litellm.proxy.management_helpers.customer_hierarchy import ( Failure, @@ -483,27 +485,22 @@ def check_customer_hierarchy_model_access( ) -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, - ) +def _untracked(message: str, header: str) -> None: + """An id the proxy does not know is logged and the request runs without that customer level, never refused: + the headers are attribution, and a customer whose rows are not set up yet must keep working.""" + 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, -) -> CustomerHierarchyAuth: + customer_org_alias: str, prisma_client: "PrismaClient | None", user_api_key_cache: UserApiKeyCache +) -> 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: - raise _not_found(f"customer organization alias not found: {customer_org_alias}", CUSTOMER_ORG_ALIAS_HEADER) + return _untracked(f"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}", @@ -515,6 +512,17 @@ async def _organization_named_by_alias( MemberRegistrar: TypeAlias = Callable[[str, str, "PrismaClient"], Awaitable[CustomerMember | Failure]] +_MEMBER_ID: Final[TypeAdapter[str]] = TypeAdapter(CustomerIdentifier) + + +def _registrable(customer_member_id: str) -> bool: + """A header value becomes a member row only within the identifier bound the customer tables share, so the spend + ledger can index every row the proxy creates.""" + try: + _MEMBER_ID.validate_python(customer_member_id) + except ValidationError: + return False + return True async def _register_member( @@ -528,6 +536,12 @@ async def _register_member( alias there is no organization to place it in, and the member stays unknown.""" if organization is None: return None + if not _registrable(customer_member_id): + verbose_proxy_logger.warning( + "Could not register customer member in organization %s: id is not a valid customer identifier", + organization.customer_org_id, + ) + return None registered: Final = await register(organization.customer_org_id, customer_member_id, prisma_client) if isinstance(registered, Failure): verbose_proxy_logger.warning( @@ -546,9 +560,11 @@ async def _hierarchy_named_by_member( prisma_client: "PrismaClient | None", user_api_key_cache: UserApiKeyCache, register: MemberRegistrar, -) -> CustomerHierarchyAuth: +) -> CustomerHierarchyAuth | None: + """The member's hierarchy, or the aliased organization alone when the member is unknown and cannot be + registered, or None when neither is known.""" if prisma_client is None: - raise _not_found(f"customer member not found: {customer_member_id}", CUSTOMER_MEMBER_ID_HEADER) + return _untracked(f"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 @@ -556,7 +572,8 @@ async def _hierarchy_named_by_member( 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) + _untracked(f"member not found: {customer_member_id}", CUSTOMER_MEMBER_ID_HEADER) + return organization return hierarchy @@ -565,24 +582,28 @@ async def load_customer_hierarchy( 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.""" + key_metadata: Mapping[str, object] = MappingProxyType({}), +) -> 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 + 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) if identity.customer_org_alias is not None else None ) + if organization is not None: + check_customer_hierarchy_key_scope(organization, key_metadata) 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) if identity.customer_team_id is not None: raise _forbidden("a customer team header needs a customer member header", CUSTOMER_TEAM_ID_HEADER) return organization hierarchy: Final = await _hierarchy_named_by_member( identity.customer_member_id, organization, prisma_client, user_api_key_cache, register ) + if hierarchy is None: + return None if organization is not None and organization.customer_org_id != hierarchy.customer_org_id: raise _forbidden( f"Customer member {identity.customer_member_id} does not belong to the organization " @@ -613,14 +634,17 @@ async def resolve_customer_hierarchy( route: str, key_metadata: Mapping[str, object] | None = None, register: MemberRegistrar = ensure_customer_member, -) -> CustomerHierarchyAuth: - """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.""" +) -> 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 - cached: Final = await load_customer_hierarchy(identity, prisma_client, user_api_key_cache, register) + metadata: Final = key_metadata or {} + cached: Final = await load_customer_hierarchy(identity, prisma_client, user_api_key_cache, register, metadata) + if cached is None: + return None hierarchy: Final = await current_customer_spend(cached, get_current_spend) - check_customer_hierarchy_key_scope(hierarchy, key_metadata or {}) + check_customer_hierarchy_key_scope(hierarchy, metadata) if not RouteChecks.is_info_route(route): check_customer_hierarchy_budget(hierarchy) if any(scope.budget_limits for scope in hierarchy.scopes): diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index f9b642a99cf..d1f78c24477 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -250,6 +250,13 @@ def _route_requires_auth_despite_public(route: str, general_settings: dict | Non return False +def _is_unauthenticated_public_route(route: str) -> bool: + """A route the builder serves without credentials, so nothing on the request may create or charge anything.""" + return not _route_requires_auth_despite_public(route=route, general_settings=None) and ( + route in LiteLLMRoutes.public_routes.value or route_in_additonal_public_routes(current_route=route) + ) + + custom_litellm_key_header: Final = APIKeyHeader( name=SpecialHeaders.custom_litellm_api_key.value, auto_error=False, @@ -2612,9 +2619,7 @@ async def user_api_key_auth_builder( """ ######## Route Checks Before Reading DB / Cache for "token" ################ - if not _route_requires_auth_despite_public(route=route, general_settings=general_settings) and ( - route in LiteLLMRoutes.public_routes.value or route_in_additonal_public_routes(current_route=route) - ): + if _is_unauthenticated_public_route(route): # check if public endpoint return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) @@ -4548,7 +4553,7 @@ async def _authorize_authenticated_request( 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_identity is not None: + if customer_identity is not None and not _is_unauthenticated_public_route(route): 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 identity=customer_identity, diff --git a/tests/unit/proxy/auth/test_customer_hierarchy_auth.py b/tests/unit/proxy/auth/test_customer_hierarchy_auth.py index 4c5de74a774..e8d10a6ef04 100644 --- a/tests/unit/proxy/auth/test_customer_hierarchy_auth.py +++ b/tests/unit/proxy/auth/test_customer_hierarchy_auth.py @@ -469,6 +469,7 @@ def _resolve( identity: CustomerRequestIdentity, route: str = "/chat/completions", register: _Registrar | None = None, + key_metadata: Mapping[str, object] | None = None, ): return resolve_customer_hierarchy( identity=identity, @@ -477,6 +478,7 @@ def _resolve( llm_router=None, request_model="gpt-4o", route=route, + key_metadata=key_metadata, register=register if register is not None else _Registrar(_MemberTable({}), Failure(ErrorCode.NOT_FOUND)), ) @@ -496,27 +498,26 @@ 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: +async def test_an_unknown_member_or_a_proxy_without_a_database_runs_the_request_untracked() -> None: cache: Final = UserApiKeyCache() - with pytest.raises(ProxyException) as missing: - await resolve_customer_hierarchy( - identity=CustomerRequestIdentity(customer_member_id="ghost"), - prisma_client=_Prisma(_MemberTable({})), # pyright: ignore[reportArgumentType] # test double - user_api_key_cache=cache, - llm_router=None, - request_model="gpt-4o", - route="/chat/completions", - ) - assert missing.value.code == "401" - with pytest.raises(ProxyException): - await resolve_customer_hierarchy( - identity=CustomerRequestIdentity(customer_member_id="member-1"), - prisma_client=None, - user_api_key_cache=cache, - llm_router=None, - request_model="gpt-4o", - route="/chat/completions", - ) + untracked: Final = await resolve_customer_hierarchy( + identity=CustomerRequestIdentity(customer_member_id="ghost"), + prisma_client=_Prisma(_MemberTable({})), # pyright: ignore[reportArgumentType] # test double + user_api_key_cache=cache, + llm_router=None, + request_model="gpt-4o", + route="/chat/completions", + ) + assert untracked is None + without_database: Final = await resolve_customer_hierarchy( + identity=CustomerRequestIdentity(customer_member_id="member-1"), + prisma_client=None, + user_api_key_cache=cache, + llm_router=None, + request_model="gpt-4o", + route="/chat/completions", + ) + assert without_database is None @pytest.mark.asyncio @@ -540,7 +541,7 @@ async def test_resolution_enforces_budget_on_model_routes_but_not_on_info_routes request_model=None, route="/key/info", ) - assert hierarchy.customer_member_id == "member-1" + assert hierarchy is not None and hierarchy.customer_member_id == "member-1" def test_a_key_pinned_to_another_organization_cannot_act_for_the_member() -> None: @@ -576,6 +577,7 @@ async def test_an_alias_alone_charges_only_the_organization_and_enforces_its_bud 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 is not None assert (hierarchy.customer_org_id, hierarchy.customer_team_id, hierarchy.customer_member_id) == ( "org-1", None, @@ -602,6 +604,7 @@ async def test_an_alias_next_to_a_member_must_name_that_members_organization() - matching: Final = await _resolve( prisma, CustomerRequestIdentity(customer_member_id="member-1", customer_org_alias="acme") ) + assert matching is not None 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") @@ -611,10 +614,8 @@ async def test_an_alias_next_to_a_member_must_name_that_members_organization() - @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) +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 shared: Final = _Prisma( _MemberTable({}), _OrganizationTable((_organization("acme"), _organization("acme", org_id="org-2"))) @@ -659,28 +660,70 @@ async def test_a_member_the_aliased_organization_has_never_seen_is_registered_th ) assert registrar.calls == (("org-1", "user-7"),) + assert hierarchy is not None 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: +async def test_an_unknown_member_runs_untracked_without_an_alias_and_on_the_organization_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", ()) + no_alias: Final = await _resolve( + _Prisma(members), CustomerRequestIdentity(customer_member_id="user-7"), register=registrar + ) + assert (no_alias, registrar.calls) == (None, ()) 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"),)) + prisma: Final = _Prisma(members, _OrganizationTable((_organization("acme", max_budget=5.0),))) + organization_only: Final = await _resolve( + prisma, CustomerRequestIdentity(customer_member_id="user-7", customer_org_alias="acme"), register=failing + ) + assert failing.calls == (("org-1", "user-7"),) + assert organization_only is not None + assert (organization_only.customer_org_id, organization_only.customer_member_id) == ("org-1", None) + assert tuple(scope.key for scope in organization_only.scopes) == ("customer_org",) assert "user-7" not in members.rows +@pytest.mark.asyncio +async def test_a_key_pinned_to_another_organization_is_refused_before_any_member_is_registered() -> None: + members: Final = _MemberTable({}) + registrar: Final = _Registrar(members) + prisma: Final = _Prisma(members, _OrganizationTable((_organization("acme"),))) + + with pytest.raises(ProxyException) as denied: + await _resolve( + prisma, + CustomerRequestIdentity(customer_member_id="user-7", customer_org_alias="acme"), + register=registrar, + key_metadata={"customer_org_id": "some-other-org"}, + ) + + assert (denied.value.code, registrar.calls, members.rows) == ("403", (), {}) + + +@pytest.mark.asyncio +async def test_a_member_id_past_the_identifier_bound_is_not_registered_and_only_the_organization_is_charged() -> None: + members: Final = _MemberTable({}) + registrar: Final = _Registrar(members) + prisma: Final = _Prisma(members, _OrganizationTable((_organization("acme"),))) + + oversized: Final = await _resolve( + prisma, CustomerRequestIdentity(customer_member_id="u" * 257, customer_org_alias="acme"), register=registrar + ) + assert oversized is not None + assert (oversized.customer_org_id, oversized.customer_member_id, registrar.calls) == ("org-1", None, ()) + + longest: Final = await _resolve( + prisma, CustomerRequestIdentity(customer_member_id="u" * 256, customer_org_alias="acme"), register=registrar + ) + assert longest is not None + assert (longest.customer_member_id, registrar.calls) == ("u" * 256, (("org-1", "u" * 256),)) + + @pytest.mark.asyncio @pytest.mark.parametrize("exhausted_duration", ("1mo", "12mo")) async def test_either_customer_cap_blocks_while_the_other_has_capacity(exhausted_duration: str, monkeypatch) -> None: diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index e76b2e85afe..fe8b79f5815 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -7340,6 +7340,54 @@ async def test_centralized_common_checks_skips_public_routes(): setattr(_proxy_server_mod, k, v) +class _CustomerTablesOffLimits: + """Prisma stand-in whose customer tables fail the test the moment anything reads them.""" + + class _Db: + def __getattr__(self, name: str) -> MagicMock: + assert not name.startswith("litellm_customer"), f"public route read {name}" + return MagicMock() + + db = _Db() + + +@pytest.mark.asyncio +async def test_authorizing_a_public_route_ignores_the_customer_headers(): + """A credential-free GET / carrying both customer headers must not look up, register or charge a customer: + the builder served it without a key, so nothing on it is trusted.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.auth.user_api_key_auth import _authorize_authenticated_request + + token = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) + request = Request( + scope={ + "type": "http", + "method": "GET", + "path": "/", + "query_string": b"", + "headers": [(b"x-litellm-customer-org-alias", b"acme"), (b"x-litellm-customer-member-id", b"user-7")], + } + ) + attrs = { + **_proxy_attrs_for_centralized_checks(), + "prisma_client": _CustomerTablesOffLimits(), + "llm_model_list": None, + "user_model": None, + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + recovered = await _authorize_authenticated_request( + user_api_key_auth_obj=token, request=request, request_data={}, route="/", api_key="" + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert (recovered, token.customer_hierarchy) == (None, None) + + @pytest.mark.asyncio async def test_centralized_common_checks_skips_passthrough_endpoint_with_auth_false(): """Regression: user-configured pass-through endpoints with