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

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

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

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

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 = []