From 95a2d5088a7aa62729062a9264dbdd4999bc2fc2 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:31:26 +0000 Subject: [PATCH] fix(proxy): keep custom-auth end-user caps under a key default budget Custom auth callables that already capped an end user keep their cap; the key default fills only unset limits. The proxy-wide default still reaches an uncapped custom-auth token, and the missing-budget log strips line breaks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 5 +- litellm/proxy/auth/user_api_key_auth.py | 26 +++++-- .../auth/test_custom_auth_end_user_budget.py | 76 ++++++++++++++++--- .../proxy/auth/test_user_api_key_auth.py | 48 +++++++++++- 4 files changed, 135 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3cb0f6255d9..32130f045c8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1407,7 +1407,10 @@ async def get_default_end_user_budget( ) if budget_record is None: - verbose_proxy_logger.warning("Default end user budget not found in database: %s", default_budget_id) + verbose_proxy_logger.warning( + "Default end user budget not found in database: %s", + default_budget_id.replace("\r", "").replace("\n", ""), + ) return None _budget_obj: Final = LiteLLM_BudgetTable.model_validate(budget_record.dict()) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 95fa5e01207..bf24ed19113 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -708,8 +708,6 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"] if end_user_params.get("end_user_tpd_limit") is not None: valid_token.end_user_tpd_limit = end_user_params["end_user_tpd_limit"] - if end_user_params.get("end_user_max_budget") is not None: - valid_token.end_user_max_budget = end_user_params["end_user_max_budget"] if end_user_params.get("allowed_model_region") is not None: valid_token.allowed_model_region = end_user_params["allowed_model_region"] if end_user_params.get("end_user_model_max_budget") is not None: @@ -2874,6 +2872,7 @@ async def _run_centralized_common_checks( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, + keep_token_limits=user_custom_auth is not None, ) skip_budget_checks: Final = _should_skip_budget_checks( @@ -2971,10 +2970,14 @@ async def _apply_key_end_user_default_budget_to_token( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, + keep_token_limits: bool, ) -> None: """The builder's end-user pass runs before the key is resolved, so only here can the key's ``end_user_budget_id`` win over the proxy-wide default on the token that reservation reads. - The budget replaces the proxy-wide one wholesale: a key budget with no cap also lifts the cap.""" + On the virtual-key path the token's end-user limits are the builder's proxy-wide defaults and + the key budget replaces them wholesale. With ``keep_token_limits`` (custom auth) the token's + limits are caps the custom auth callable set, so the key budget only fills the ones it left + unset.""" default_budget: Final = ( end_user_object.litellm_budget_table if end_user_object is not None @@ -2988,11 +2991,16 @@ async def _apply_key_end_user_default_budget_to_token( if default_budget is None: return - valid_token.end_user_max_budget = default_budget.max_budget - valid_token.end_user_tpm_limit = default_budget.tpm_limit - valid_token.end_user_rpm_limit = default_budget.rpm_limit - valid_token.end_user_tpd_limit = default_budget.tpd_limit - valid_token.end_user_model_max_budget = default_budget.model_max_budget + if not keep_token_limits or valid_token.end_user_max_budget is None: + valid_token.end_user_max_budget = default_budget.max_budget + if not keep_token_limits or valid_token.end_user_tpm_limit is None: + valid_token.end_user_tpm_limit = default_budget.tpm_limit + if not keep_token_limits or valid_token.end_user_rpm_limit is None: + valid_token.end_user_rpm_limit = default_budget.rpm_limit + if not keep_token_limits or valid_token.end_user_tpd_limit is None: + valid_token.end_user_tpd_limit = default_budget.tpd_limit + if not keep_token_limits or valid_token.end_user_model_max_budget is None: + valid_token.end_user_model_max_budget = default_budget.model_max_budget async def _reserve_budget_after_common_checks( @@ -3465,6 +3473,8 @@ async def _lookup_end_user_and_apply_budget( valid_token = update_valid_token_with_end_user_params( valid_token=valid_token, end_user_params=end_user_params ) + if valid_token.end_user_max_budget is None: + valid_token.end_user_max_budget = default_budget.max_budget except Exception as e: if isinstance(e, litellm.BudgetExceededError): raise e diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index bf425327d4a..e83c5cf8419 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -261,18 +261,76 @@ async def test_custom_auth_key_default_end_user_budget_reaches_the_token_for_a_n assert valid_token.end_user_max_budget == 0.5 -def test_end_user_budget_max_budget_reaches_the_token(): - from litellm.proxy.auth.user_api_key_auth import _apply_budget_limits_to_end_user_params +@pytest.mark.asyncio +async def test_custom_auth_cap_stays_below_the_key_default_end_user_budget(monkeypatch): + """A custom auth callable that already capped the end user tighter than the key's default + budget keeps its cap: the key default never loosens what custom auth set.""" + from unittest.mock import MagicMock - end_user_params = {"end_user_id": "user_1"} - _apply_budget_limits_to_end_user_params( - end_user_params=end_user_params, - budget_info=LiteLLM_BudgetTable(max_budget=20.0), - end_user_id="user_1", + from litellm.proxy.auth.user_api_key_auth import _lookup_end_user_and_apply_budget + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + monkeypatch.setattr(litellm, "max_end_user_budget_id", None) + + async def _find_budget(where): + row = MagicMock() + row.dict = lambda: {"budget_id": where["budget_id"], "max_budget": 0.5} + return row + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_budgettable.find_unique = AsyncMock(side_effect=_find_budget) + + valid_token, _ = await _lookup_end_user_and_apply_budget( + valid_token=UserAPIKeyAuth( + token="test_token", + end_user_id="customer-new", + end_user_max_budget=0.1, + metadata={"end_user_budget_id": "svc-a-budget"}, + ), + route="/v1/chat/completions", + parent_otel_span=None, + prisma_client=mock_prisma, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=MagicMock(), ) - result = update_valid_token_with_end_user_params(UserAPIKeyAuth(token="test_token"), end_user_params) - assert result.end_user_max_budget == 20.0 + assert valid_token.end_user_max_budget == 0.1 + + +@pytest.mark.asyncio +async def test_custom_auth_proxy_wide_default_end_user_budget_reaches_an_uncapped_token(monkeypatch): + """With no key default, a brand-new end user on a custom-auth token that set no cap gets the + proxy-wide default budget's cap, the same way the virtual-key path already applies it.""" + from unittest.mock import MagicMock + + from litellm.proxy.auth.user_api_key_auth import _lookup_end_user_and_apply_budget + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + monkeypatch.setattr(litellm, "max_end_user_budget_id", "global-eu-budget") + + async def _find_budget(where): + row = MagicMock() + row.dict = lambda: {"budget_id": where["budget_id"], "max_budget": 100.0} + return row + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_budgettable.find_unique = AsyncMock(side_effect=_find_budget) + + valid_token, end_user_object = await _lookup_end_user_and_apply_budget( + valid_token=UserAPIKeyAuth(token="test_token", end_user_id="customer-new"), + route="/v1/chat/completions", + parent_otel_span=None, + prisma_client=mock_prisma, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=MagicMock(), + ) + + assert end_user_object is None + assert valid_token.end_user_max_budget == 100.0 def test_update_valid_token_does_not_override_custom_auth_values_with_none(): diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 8052970684b..674eec9738c 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -4387,8 +4387,11 @@ async def _run_centralized_checks_with_key_end_user_budget( budgets: Mapping[str, float], request_user: str | None = None, user_api_key_cache: DualCache | None = None, + custom_auth: bool = False, ) -> UserAPIKeyAuth: - """Run the centralized checks with a fake DB and return the token handed to budget reservation.""" + """Run the centralized checks with a fake DB and return the token handed to budget reservation. + With ``custom_auth`` the token stands for one a custom auth callable returned and the checks + run under ``custom_auth_run_common_checks``.""" from fastapi import Request from starlette.datastructures import URL @@ -4408,7 +4411,9 @@ async def _run_centralized_checks_with_key_end_user_budget( request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") attrs = { - **_proxy_attrs_for_centralized_checks(user_custom_auth=None), + **_proxy_attrs_for_centralized_checks( + user_custom_auth=AsyncMock() if custom_auth else None, flag=custom_auth + ), "prisma_client": prisma_client, "user_api_key_cache": user_api_key_cache if user_api_key_cache is not None else DualCache(), "proxy_logging_obj": proxy_logging_obj, @@ -4511,6 +4516,45 @@ async def test_centralized_common_checks_keeps_an_end_users_own_budget_over_the_ assert reserved_token.end_user_max_budget == 500.0 +@pytest.mark.asyncio +async def test_centralized_common_checks_keeps_a_stricter_custom_auth_cap_over_the_key_default(monkeypatch): + """A custom auth callable that caps the end user tighter than the key's default budget keeps + its cap and its rate limit. The key default only fills the limits the callable left unset.""" + monkeypatch.setattr(litellm, "max_end_user_budget_id", None) + token = UserAPIKeyAuth( + api_key="sk-test", + token="hashed", + end_user_id="cust-new", + end_user_max_budget=0.1, + end_user_rpm_limit=3, + metadata={"service_account_id": "svc-a", "end_user_budget_id": "svc-a-budget"}, + ) + + reserved_token = await _run_centralized_checks_with_key_end_user_budget( + token, end_user_row=None, budgets={"svc-a-budget": 0.5}, custom_auth=True + ) + + assert reserved_token.end_user_max_budget == 0.1 + assert reserved_token.end_user_rpm_limit == 3 + + +@pytest.mark.asyncio +async def test_centralized_common_checks_fills_a_custom_auth_token_without_a_cap_from_the_key_default(monkeypatch): + monkeypatch.setattr(litellm, "max_end_user_budget_id", None) + token = UserAPIKeyAuth( + api_key="sk-test", + token="hashed", + end_user_id="cust-new", + metadata={"service_account_id": "svc-a", "end_user_budget_id": "svc-a-budget"}, + ) + + reserved_token = await _run_centralized_checks_with_key_end_user_budget( + token, end_user_row=None, budgets={"svc-a-budget": 0.5}, custom_auth=True + ) + + assert reserved_token.end_user_max_budget == 0.5 + + class _RecordingTeamModelBudgetLimiter: def __init__(self): self.calls = []