From a96e28adee30ecd4a1cec80c6482355432049177 Mon Sep 17 00:00:00 2001 From: IvanShang Date: Fri, 4 Sep 2026 22:11:10 +0800 Subject: [PATCH] fix(proxy): hydrate cached key end-user rate limits --- litellm/proxy/auth/user_api_key_auth.py | 12 ++++ .../proxy/auth/test_user_api_key_auth.py | 67 +++++++++++++++++++ 2 files changed, 79 insertions(+) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 5fb6dad0cd7..adb0214d2d3 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2595,6 +2595,18 @@ async def _run_centralized_common_checks( end_user_object: Final[LiteLLM_EndUserTable | None] = ( None if isinstance(end_user_result, BaseException) else end_user_result ) + if user_api_key_auth_obj.end_user_id is not None and end_user_object is not None: + end_user_params: Final = {"end_user_id": user_api_key_auth_obj.end_user_id} + if end_user_object.litellm_budget_table is not None: + _apply_budget_limits_to_end_user_params( + end_user_params=end_user_params, + budget_info=end_user_object.litellm_budget_table, + end_user_id=user_api_key_auth_obj.end_user_id, + ) + update_valid_token_with_end_user_params( + valid_token=user_api_key_auth_obj, + end_user_params=end_user_params, + ) global_proxy_spend: float | None = None if isinstance(global_spend_result, BaseException) else global_spend_result if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not 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 d44f96d95bf..a38b1010133 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 @@ -4246,6 +4246,73 @@ async def test_centralized_common_checks_propagates_end_user_budget_error(): setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +async def test_centralized_checks_hydrates_cached_key_end_user_rate_limits(): + """A cached virtual key must carry its customer's budget rate limits into the limiter.""" + import litellm.proxy.proxy_server as _proxy_server_mod + + token = UserAPIKeyAuth(api_key="sk-test", end_user_id="customer-1") + end_user_object = LiteLLM_EndUserTable( + user_id="customer-1", + blocked=False, + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(rpm_limit=5, tpm_limit=20), + ) + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=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) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + new_callable=AsyncMock, + return_value=end_user_object, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._reserve_budget_after_common_checks", + new_callable=AsyncMock, + ), + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=_chat_request(), + request_data={"model": "gpt-4o-mini", "user": "customer-1"}, + route="/chat/completions", + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert token.end_user_rpm_limit == 5 + assert token.end_user_tpm_limit == 20 + + from litellm.caching.caching import DualCache + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + ) + from litellm.proxy.utils import InternalUsageCache + + limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + descriptors = limiter._create_rate_limit_descriptors( + user_api_key_dict=token, + data={"model": "gpt-4o-mini"}, + rpm_limit_type=None, + tpm_limit_type=None, + model_has_failures=False, + ) + end_user_descriptor = next(descriptor for descriptor in descriptors if descriptor["key"] == "end_user") + assert end_user_descriptor["value"] == "customer-1" + assert end_user_descriptor["rate_limit"]["requests_per_unit"] == 5 + assert end_user_descriptor["rate_limit"]["tokens_per_unit"] == 20 + + @pytest.mark.asyncio async def test_centralized_common_checks_reserves_request_end_user_budget(): """Regression: reservation runs before user_api_key_auth() copies the