scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

feat(proxy): run requests with unknown customer ids untracked instead of refusing them (#51548)

Browse source
This commit is contained in:
ishaan-j-berri 2026-11-14 00:56:21 +08:00 • committed by GitHub
parent ca4c6c0a6d
commit 616471d400
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 186 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

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

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

48
tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py
View file

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