From ed5fe048ff652242dcc0241b8a81ee253f81d7e2 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 09:19:36 -0700 Subject: [PATCH] fix(proxy/auth): apply max_end_user_budget_id to end users without a row on cached keys (#48704) * test(proxy/auth): cover default budgets on cached virtual keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy/auth): use in-memory cache in cached budget tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy/auth): apply max_end_user_budget_id to end users without a row on cached keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy/auth): share cached end-user parameter update Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 3 +- .../test_user_api_key_auth_request_flow.py | 91 +++++++++++++++++++ 2 files changed, 93 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index cb490b55352..99edfb940d7 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -3257,8 +3257,9 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e if valid_token.budget_throttle_pct is not None: valid_token_dict["budget_throttle_pct"] = valid_token.budget_throttle_pct - if _end_user_object is not None: + if _end_user_object is not None or litellm.max_end_user_budget_id is not None: valid_token_dict.update(end_user_params) + if _end_user_object is not None: valid_token_dict["end_user_object_permission"] = _end_user_object.object_permission # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions 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 2d9c781b0d1..ae7168ed1e1 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 @@ -1537,6 +1537,97 @@ async def test_cached_proxy_admin_key_sets_via_virtual_key_marker(): setattr(_proxy_server_mod, attr, val) +@pytest.mark.parametrize( + ("default_budget_id", "expected_limits"), + [ + pytest.param("proxy-default-end-user-budget", (0, 7), id="default-configured"), + pytest.param(None, (None, None), id="no-default"), + ], +) +@pytest.mark.asyncio +async def test_cached_virtual_key_applies_proxy_default_budget_to_new_end_user( + monkeypatch: pytest.MonkeyPatch, + default_budget_id: str | None, + expected_limits: tuple[int | None, int | None], +) -> None: + """Cached virtual keys apply configured limits to new end users.""" + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy.proxy_server import hash_token + + end_user_id: Final = "cached-virtual-key-end-user" + monkeypatch.setattr(litellm, "max_end_user_budget_id", default_budget_id) + + api_key: Final = "sk-cached-default-end-user-budget-test" + hashed_key: Final = hash_token(api_key) + cached_token: Final = UserAPIKeyAuth( + api_key=api_key, + token=hashed_key, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + budget: Final = LiteLLM_BudgetTable( + budget_id="proxy-default-end-user-budget", + rpm_limit=0, + tpm_limit=7, + ) + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=budget) + + user_api_key_cache: Final = UserApiKeyCache() + + mock_proxy_logging_obj: Final = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + import litellm.proxy.proxy_server as _proxy_server_mod + + monkeypatch.setattr(_proxy_server_mod, "prisma_client", mock_prisma_client) + monkeypatch.setattr(_proxy_server_mod, "user_api_key_cache", user_api_key_cache) + monkeypatch.setattr(_proxy_server_mod, "proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr(_proxy_server_mod, "master_key", "sk-master-key") + monkeypatch.setattr(_proxy_server_mod, "general_settings", {}) + monkeypatch.setattr(_proxy_server_mod, "llm_model_list", []) + monkeypatch.setattr(_proxy_server_mod, "llm_router", None) + monkeypatch.setattr(_proxy_server_mod, "open_telemetry_logger", None) + monkeypatch.setattr(_proxy_server_mod, "model_max_budget_limiter", MagicMock()) + monkeypatch.setattr(_proxy_server_mod, "user_custom_auth", None) + monkeypatch.setattr(_proxy_server_mod, "jwt_handler", None) + monkeypatch.setattr(_proxy_server_mod, "litellm_proxy_admin_name", "admin") + + request: Final = Request( + scope={ + "type": "http", + "headers": [(b"x-litellm-end-user-id", end_user_id.encode())], + } + ) + request._url = URL(url="/chat/completions") + + with patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, + return_value=cached_token, + ): + result: Final = await user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert (result.end_user_rpm_limit, result.end_user_tpm_limit) == expected_limits + if default_budget_id is not None: + assert result.end_user_id == end_user_id + + @pytest.mark.asyncio async def test_master_key_auth_sets_via_virtual_key_marker(): """Master-key requests must also be stamped by overwrite_user_with_key_hash;