From 55f7b277ee58e981b71648cdcede8012ef4c070c Mon Sep 17 00:00:00 2001 From: Taranum01 Date: Wed, 30 Sep 2026 14:53:14 +0530 Subject: [PATCH 1/7] fix(proxy): rebase end-user model budgets and cover cached virtual keys --- litellm/__init__.py | 5 + litellm/proxy/auth/user_api_key_auth.py | 94 ++- ...t_end_user_model_max_budget_enforcement.py | 660 ++++++++++++++++++ 3 files changed, 749 insertions(+), 10 deletions(-) create mode 100644 tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py diff --git a/litellm/__init__.py b/litellm/__init__.py index fea7a27a5fb..38f51ec3406 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -474,6 +474,11 @@ max_end_user_budget_id: Optional[str] = None # backwards compatibility — arbitrary client-supplied identifiers still # pass through unchanged. validate_end_user_id_in_db: bool = False +# When True, master-key authenticated LLM API requests enforce +# end_user_model_max_budget from the customer budget table. Defaults to False +# for backwards compatibility — master-key callers that act on behalf of +# end-users were previously not subject to this check. +enforce_end_user_model_max_budget_on_master_key: bool = False block_requests_for_models_without_pricing: bool = False disable_end_user_cost_tracking: Optional[bool] = None disable_end_user_cost_tracking_prometheus_only: Optional[bool] = None diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index cb39801a8e0..444508fa21a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2088,6 +2088,13 @@ async def _user_api_key_auth_builder( if _end_user_object is not None: valid_token.end_user_object_permission = _end_user_object.object_permission + await _maybe_enforce_master_key_end_user_model_max_budget( + valid_token=valid_token, + request_data=request_data, + route=route, + request=request, + ) + return valid_token if ( @@ -2150,6 +2157,19 @@ async def _user_api_key_auth_builder( route=route, start_time=start_time, ) + + _user_api_key_obj = update_valid_token_with_end_user_params( + valid_token=_user_api_key_obj, end_user_params=end_user_params + ) + _user_api_key_obj.via_virtual_key = True + + await _maybe_enforce_master_key_end_user_model_max_budget( + valid_token=_user_api_key_obj, + request_data=request_data, + route=route, + request=request, + ) + asyncio.create_task( _cache_key_object( hashed_token=hash_token(master_key), @@ -2159,11 +2179,6 @@ async def _user_api_key_auth_builder( ) ) - _user_api_key_obj = update_valid_token_with_end_user_params( - valid_token=_user_api_key_obj, end_user_params=end_user_params - ) - _user_api_key_obj.via_virtual_key = True - return _user_api_key_obj ## IF it's not a master key @@ -2225,11 +2240,9 @@ async def _user_api_key_auth_builder( raise e # update end-user params on valid token # These can change per request - it's important to update them here - valid_token.end_user_id = end_user_params.get("end_user_id") - valid_token.end_user_tpm_limit = end_user_params.get("end_user_tpm_limit") - valid_token.end_user_rpm_limit = end_user_params.get("end_user_rpm_limit") - valid_token.end_user_tpd_limit = end_user_params.get("end_user_tpd_limit") - valid_token.allowed_model_region = end_user_params.get("allowed_model_region") + valid_token = update_valid_token_with_end_user_params( + valid_token=valid_token, end_user_params=end_user_params + ) return await validate_resolved_virtual_key( request=request, @@ -2282,6 +2295,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e ) if valid_token is not None: + valid_token = update_valid_token_with_end_user_params(valid_token=valid_token, end_user_params=end_user_params) valid_token = _update_key_budget_with_temp_budget_increase(valid_token) user_obj: LiteLLM_UserTable | None = None @@ -3823,6 +3837,66 @@ def _fallback_target_model_name(target: object) -> str | None: return None +def _is_master_key_auth_token(valid_token: UserAPIKeyAuth) -> bool: + return valid_token.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS or valid_token.token == LITELLM_PROXY_MASTER_KEY_ALIAS + + +async def _maybe_enforce_master_key_end_user_model_max_budget( + valid_token: UserAPIKeyAuth, + request_data: dict, + route: str, + request: Request, +) -> None: + if not litellm.enforce_end_user_model_max_budget_on_master_key: + return + if not RouteChecks.is_llm_api_route(route=route): + return + if not _is_master_key_auth_token(valid_token): + return + + await _enforce_end_user_model_max_budget_checks( + valid_token=valid_token, + request_data=request_data, + route=route, + request=request, + ) + + +async def _enforce_end_user_model_max_budget_checks( + valid_token: UserAPIKeyAuth, + request_data: dict, + route: str, + request: Request, +) -> None: + from litellm.proxy.proxy_server import llm_router, model_max_budget_limiter + + end_user_mmb = valid_token.end_user_model_max_budget + if ( + end_user_mmb is None + or not isinstance(end_user_mmb, dict) + or len(end_user_mmb) == 0 + or valid_token.end_user_id is None + ): + return + + current_model = _get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + ) + current_models = _get_model_names_for_budget_checks(model=current_model) + if not current_models: + return + + for model_name in current_models: + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=valid_token.end_user_id, + end_user_model_max_budget=end_user_mmb, + model=model_name, + ) + + async def _run_post_custom_auth_checks( valid_token: UserAPIKeyAuth, request: Request, diff --git a/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py b/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py new file mode 100644 index 00000000000..fc1c2742658 --- /dev/null +++ b/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py @@ -0,0 +1,660 @@ +import pytest +from contextlib import contextmanager +from unittest.mock import AsyncMock, MagicMock, patch + +import litellm +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, + ProxyErrorTypes, + ProxyException, + UserAPIKeyAuth, +) +from litellm.proxy.auth.user_api_key_auth import update_valid_token_with_end_user_params + +MODEL = "google/gemini-2.5-flash-lite" +MODEL_BUDGET = {"max_budget": 1e-05, "budget_duration": "1d"} + + +def _proxy_server_attrs_for_master_key_auth(): + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.delete_cache = MagicMock() + + mock_proxy_logging_obj = 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) + + limiter = AsyncMock() + limiter.is_end_user_within_model_budget = AsyncMock(return_value=None) + + return { + "prisma_client": MagicMock(), + "user_api_key_cache": mock_cache, + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": limiter, + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + }, limiter + + +def _end_user_with_model_budget(): + return LiteLLM_EndUserTable( + user_id="customer-1", + blocked=False, + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(model_max_budget={MODEL: MODEL_BUDGET}), + ) + + +@contextmanager +def _virtual_key_builder_patches(*, resolved_token: UserAPIKeyAuth, cache_hit: bool = False): + async def mock_resolve_key(self, hashed_token: str): + from litellm.proxy.auth.resolvers.store import KeyNotInCacheError + + if self._check_cache_only and not cache_hit: + raise KeyNotInCacheError(hashed_token) + return resolved_token + + with ( + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new=mock_resolve_key, + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth.resolve_and_validate_end_user_id", + new_callable=AsyncMock, + return_value="customer-1", + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + new_callable=AsyncMock, + return_value=_end_user_with_model_budget(), + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth._get_model_from_request_context", + return_value=MODEL, + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth._enforce_key_and_fallback_model_access", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth._virtual_key_max_budget_alert_check", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth._virtual_key_max_budget_check", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth._virtual_key_soft_budget_check", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + ): + yield + + +def test_update_valid_token_applies_end_user_model_max_budget_from_params(): + valid_token = UserAPIKeyAuth(token="test-key") + end_user_params = { + "end_user_id": "customer-1", + "end_user_model_max_budget": {MODEL: MODEL_BUDGET}, + } + + result = update_valid_token_with_end_user_params(valid_token, end_user_params) + + assert result.end_user_id == "customer-1" + assert result.end_user_model_max_budget == end_user_params["end_user_model_max_budget"] + + +@pytest.mark.asyncio +async def test_enforce_end_user_model_max_budget_passes_when_within_budget(): # test-quality-ok: structural assertion of mock wiring by design + from litellm.proxy.auth.user_api_key_auth import _enforce_end_user_model_max_budget_checks + + valid_token = UserAPIKeyAuth( + token="master-key", + end_user_id="customer-1", + end_user_model_max_budget={MODEL: MODEL_BUDGET}, + ) + request = MagicMock() + request_data = {"model": MODEL} + + with patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth._get_model_from_request_context", + return_value=MODEL, + ): + with patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.proxy_server.model_max_budget_limiter.is_end_user_within_model_budget", + new_callable=AsyncMock, + ) as mock_check: + await _enforce_end_user_model_max_budget_checks( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=request, + ) + + mock_check.assert_awaited_once_with( + end_user_id="customer-1", + end_user_model_max_budget=valid_token.end_user_model_max_budget, + model=MODEL, + ) + + +@pytest.mark.asyncio +async def test_enforce_end_user_model_max_budget_raises_when_over_budget(): + from litellm.proxy.auth.user_api_key_auth import _enforce_end_user_model_max_budget_checks + + valid_token = UserAPIKeyAuth( + token="master-key", + end_user_id="customer-1", + end_user_model_max_budget={MODEL: MODEL_BUDGET}, + ) + request = MagicMock() + request_data = {"model": MODEL} + + with patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth._get_model_from_request_context", + return_value=MODEL, + ): + with patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.proxy_server.model_max_budget_limiter.is_end_user_within_model_budget", + new_callable=AsyncMock, + ) as mock_check: + mock_check.side_effect = litellm.BudgetExceededError( + message="Exceeded budget", current_cost=0.0002, max_budget=1e-05 + ) + + with pytest.raises(litellm.BudgetExceededError): + await _enforce_end_user_model_max_budget_checks( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=request, + ) + + mock_check.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_enforce_end_user_model_max_budget_returns_early_when_unconfigured(): # test-quality-ok: structural assertion of mock wiring by design + from litellm.proxy.auth.user_api_key_auth import _enforce_end_user_model_max_budget_checks + + valid_token = UserAPIKeyAuth(token="test-key", end_user_id="customer-1") + request = MagicMock() + request_data = {"model": MODEL} + + with patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.proxy_server.model_max_budget_limiter.is_end_user_within_model_budget", + new_callable=AsyncMock, + ) as mock_check: + await _enforce_end_user_model_max_budget_checks( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=request, + ) + + mock_check.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_master_key_auth_skips_end_user_model_budget_when_flag_disabled(): + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + attrs, limiter = _proxy_server_attrs_for_master_key_auth() + limiter.is_end_user_within_model_budget.side_effect = litellm.BudgetExceededError( + message="Exceeded budget", current_cost=0.0002, max_budget=1e-05 + ) + end_user = LiteLLM_EndUserTable( + user_id="customer-1", + blocked=False, + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(model_max_budget={MODEL: MODEL_BUDGET}), + ) + originals = {k: getattr(proxy_server, k, None) for k in attrs} + flag_original = litellm.enforce_end_user_model_max_budget_on_master_key + litellm.enforce_end_user_model_max_budget_on_master_key = False # test-quality-ok: feature-flag toggle required by design (PR description) + + try: + for k, v in attrs.items(): + setattr(proxy_server, k, v) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/v1/chat/completions") + + with ( + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth.resolve_and_validate_end_user_id", + new_callable=AsyncMock, + return_value="customer-1", + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + new_callable=AsyncMock, + return_value=end_user, + ), + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {attrs['master_key']}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"user": "customer-1", "model": MODEL}, + ) + + assert result.end_user_id == "customer-1" + assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} + limiter.is_end_user_within_model_budget.assert_not_awaited() + finally: + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + for k, v in originals.items(): + setattr(proxy_server, k, v) + + +@pytest.mark.asyncio +async def test_master_key_auth_passes_when_flag_enabled_and_within_budget(): + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + attrs, limiter = _proxy_server_attrs_for_master_key_auth() + originals = {k: getattr(proxy_server, k, None) for k in attrs} + flag_original = litellm.enforce_end_user_model_max_budget_on_master_key + litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) + + try: + for k, v in attrs.items(): + setattr(proxy_server, k, v) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/v1/chat/completions") + + with ( + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth.resolve_and_validate_end_user_id", + new_callable=AsyncMock, + return_value="customer-1", + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + new_callable=AsyncMock, + return_value=_end_user_with_model_budget(), + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth._get_model_from_request_context", + return_value=MODEL, + ), + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {attrs['master_key']}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"user": "customer-1", "model": MODEL}, + ) + + assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} + limiter.is_end_user_within_model_budget.assert_awaited() + finally: + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + for k, v in originals.items(): + setattr(proxy_server, k, v) + + +@pytest.mark.asyncio +async def test_virtual_key_auth_applies_and_enforces_end_user_model_budget(): + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + valid_token = UserAPIKeyAuth(api_key="sk-vk-test", token="hashed-valid") + attrs, limiter = _proxy_server_attrs_for_master_key_auth() + attrs["master_key"] = "sk-different-master" + originals = {k: getattr(proxy_server, k, None) for k in attrs} + + try: + for k, v in attrs.items(): + setattr(proxy_server, k, v) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/v1/chat/completions") + + with _virtual_key_builder_patches(resolved_token=valid_token): + result = await _user_api_key_auth_builder( + request=request, + api_key="Bearer sk-vk-test", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"user": "customer-1", "model": MODEL}, + ) + + assert result.end_user_id == "customer-1" + assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} + limiter.is_end_user_within_model_budget.assert_awaited() + finally: + for k, v in originals.items(): + setattr(proxy_server, k, v) + + +@pytest.mark.asyncio +async def test_master_key_auth_enforces_end_user_model_budget_when_flag_enabled(): + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + attrs, limiter = _proxy_server_attrs_for_master_key_auth() + limiter.is_end_user_within_model_budget.side_effect = litellm.BudgetExceededError( + message="Exceeded budget", current_cost=0.0002, max_budget=1e-05 + ) + end_user = LiteLLM_EndUserTable( + user_id="customer-1", + blocked=False, + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(model_max_budget={MODEL: MODEL_BUDGET}), + ) + originals = {k: getattr(proxy_server, k, None) for k in attrs} + flag_original = litellm.enforce_end_user_model_max_budget_on_master_key + litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) + + try: + for k, v in attrs.items(): + setattr(proxy_server, k, v) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/v1/chat/completions") + + with ( + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth.resolve_and_validate_end_user_id", + new_callable=AsyncMock, + return_value="customer-1", + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + new_callable=AsyncMock, + return_value=end_user, + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth._get_model_from_request_context", + return_value=MODEL, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {attrs['master_key']}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"user": "customer-1", "model": MODEL}, + ) + + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + limiter.is_end_user_within_model_budget.assert_awaited() + finally: + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + for k, v in originals.items(): + setattr(proxy_server, k, v) + + +@pytest.mark.asyncio +async def test_cached_master_key_auth_enforces_end_user_model_budget_when_flag_enabled(): + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as proxy_server + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + cached_master = UserAPIKeyAuth( + api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, + token=LITELLM_PROXY_MASTER_KEY_ALIAS, + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + async def mock_resolve_key(self, hashed_token: str): + from litellm.proxy.auth.resolvers.store import KeyNotInCacheError + + if self._check_cache_only: + return cached_master + raise KeyNotInCacheError(hashed_token) + + attrs, limiter = _proxy_server_attrs_for_master_key_auth() + originals = {k: getattr(proxy_server, k, None) for k in attrs} + flag_original = litellm.enforce_end_user_model_max_budget_on_master_key + litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) + + try: + for k, v in attrs.items(): + setattr(proxy_server, k, v) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/v1/chat/completions") + + with ( + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new=mock_resolve_key, + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth.resolve_and_validate_end_user_id", + new_callable=AsyncMock, + return_value="customer-1", + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + new_callable=AsyncMock, + return_value=_end_user_with_model_budget(), + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth._get_model_from_request_context", + return_value=MODEL, + ), + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {attrs['master_key']}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"user": "customer-1", "model": MODEL}, + ) + + assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} + limiter.is_end_user_within_model_budget.assert_awaited() + finally: + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + for k, v in originals.items(): + setattr(proxy_server, k, v) + + +@pytest.mark.asyncio +async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcement(): + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + cached_admin_key = UserAPIKeyAuth( + api_key="sk-admin-virtual", + token="hashed-admin-virtual", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + async def mock_resolve_key(self, hashed_token: str): + from litellm.proxy.auth.resolvers.store import KeyNotInCacheError + + if self._check_cache_only: + return cached_admin_key + raise KeyNotInCacheError(hashed_token) + + attrs, limiter = _proxy_server_attrs_for_master_key_auth() + limiter.is_end_user_within_model_budget.side_effect = litellm.BudgetExceededError( + message="Exceeded budget", current_cost=0.0002, max_budget=1e-05 + ) + originals = {k: getattr(proxy_server, k, None) for k in attrs} + flag_original = litellm.enforce_end_user_model_max_budget_on_master_key + litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) + + try: + for k, v in attrs.items(): + setattr(proxy_server, k, v) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/v1/chat/completions") + + with ( + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new=mock_resolve_key, + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth.resolve_and_validate_end_user_id", + new_callable=AsyncMock, + return_value="customer-1", + ), + patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + new_callable=AsyncMock, + return_value=_end_user_with_model_budget(), + ), + ): + result = await _user_api_key_auth_builder( + request=request, + api_key="Bearer sk-admin-virtual", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"user": "customer-1", "model": MODEL}, + ) + + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + limiter.is_end_user_within_model_budget.assert_not_awaited() + finally: + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + for k, v in originals.items(): + setattr(proxy_server, k, v) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + [ + "/health/liveliness", + "/health/readiness", + "/key/info", + "/metrics", + ], +) +async def test_master_key_budget_early_return_for_non_llm_routes(route): # test-quality-ok: structural assertion of mock wiring by design + """Branch coverage for ``_maybe_enforce_master_key_end_user_model_max_budget``. + + The flag and master-key guards in the helper are exercised by the + builder-level tests above; this test pins the third guard — the + ``is_llm_api_route`` early return — so a future refactor that + accidentally runs the budget check on health/metrics routes fails + this test loudly. Hits line 2876 of ``user_api_key_auth.py``. + """ + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import ( + _maybe_enforce_master_key_end_user_model_max_budget, + ) + + flag_original = litellm.enforce_end_user_model_max_budget_on_master_key + litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) + + try: + valid_token = UserAPIKeyAuth( + api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, + token=LITELLM_PROXY_MASTER_KEY_ALIAS, + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + with patch( # test-quality-ok: no public seam for proxy internals on this code path yet + "litellm.proxy.proxy_server.model_max_budget_limiter.is_end_user_within_model_budget", + new_callable=AsyncMock, + ) as mock_check: + await _maybe_enforce_master_key_end_user_model_max_budget( + valid_token=valid_token, + request_data={"user": "customer-1", "model": MODEL}, + route=route, + request=MagicMock(), + ) + mock_check.assert_not_awaited() + finally: + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + + +@pytest.mark.asyncio +async def test_cached_virtual_key_auth_applies_and_enforces_end_user_model_budget(): + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + valid_token = UserAPIKeyAuth(api_key="sk-vk-test", token="hashed-valid") + attrs, limiter = _proxy_server_attrs_for_master_key_auth() + attrs["master_key"] = "sk-different-master" + originals = {k: getattr(proxy_server, k, None) for k in attrs} + + try: + for k, v in attrs.items(): + setattr(proxy_server, k, v) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/v1/chat/completions") + + with _virtual_key_builder_patches(resolved_token=valid_token, cache_hit=True): + result = await _user_api_key_auth_builder( + request=request, + api_key="Bearer sk-vk-test", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"user": "customer-1", "model": MODEL}, + ) + + assert result.end_user_id == "customer-1" + assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} + limiter.is_end_user_within_model_budget.assert_awaited() + finally: + for k, v in originals.items(): + setattr(proxy_server, k, v) From 9b363ee0a56f71ad013093cf3017a0ae31b1f403 Mon Sep 17 00:00:00 2001 From: Taranum01 Date: Wed, 30 Sep 2026 14:55:38 +0530 Subject: [PATCH 2/7] ci: register end-user budget regression in current budget shard --- .circleci/scripts/unit_selection.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index 02df32d5eab..6d6d7cedb6f 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -94,6 +94,7 @@ legacy_paths() { echo tests/unit/proxy/test_credential_slot_registry.py echo tests/unit/proxy/test_deprecated_key_grace_period.py ;; proxy-db-budgets) + echo tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py echo tests/unit/proxy/auth/test_default_end_user_budget_simple.py echo tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py echo tests/unit/proxy/test_zero_cost_model_budget_bypass.py ;; From e171a74c9795376421eb0fb04b450379ec8ce5fa Mon Sep 17 00:00:00 2001 From: Instinct Agent Date: Wed, 30 Sep 2026 15:15:05 +0530 Subject: [PATCH 3/7] fix: avoid repeated master-key alias comparisons --- litellm/proxy/auth/user_api_key_auth.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 444508fa21a..6002b157330 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -3838,7 +3838,7 @@ def _fallback_target_model_name(target: object) -> str | None: def _is_master_key_auth_token(valid_token: UserAPIKeyAuth) -> bool: - return valid_token.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS or valid_token.token == LITELLM_PROXY_MASTER_KEY_ALIAS + return LITELLM_PROXY_MASTER_KEY_ALIAS in (valid_token.api_key, valid_token.token) async def _maybe_enforce_master_key_end_user_model_max_budget( From 67d534acee4056fa287977c15c4f6359799637aa Mon Sep 17 00:00:00 2001 From: Instinct Agent Date: Sun, 4 Oct 2026 13:03:18 +0530 Subject: [PATCH 4/7] fix: preserve typed end-user budget limiter boundary --- litellm/proxy/auth/user_api_key_auth.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 6002b157330..cf7963437bc 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -3874,7 +3874,7 @@ async def _enforce_end_user_model_max_budget_checks( if ( end_user_mmb is None or not isinstance(end_user_mmb, dict) - or len(end_user_mmb) == 0 + or not end_user_mmb or valid_token.end_user_id is None ): return @@ -3892,7 +3892,7 @@ async def _enforce_end_user_model_max_budget_checks( for model_name in current_models: await model_max_budget_limiter.is_end_user_within_model_budget( end_user_id=valid_token.end_user_id, - end_user_model_max_budget=end_user_mmb, + end_user_model_max_budget=cast(Mapping[str, object], end_user_mmb), model=model_name, ) From bd5f137664b52d09857843106ee4f8190d0089ce Mon Sep 17 00:00:00 2001 From: Instinct Agent Date: Sun, 4 Oct 2026 13:29:32 +0530 Subject: [PATCH 5/7] fix: retain verified JWT customer during virtual-key budget validation --- litellm/proxy/auth/user_api_key_auth.py | 8 +- ...t_end_user_model_max_budget_enforcement.py | 97 ++++++++++++++++--- 2 files changed, 91 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index cf7963437bc..eef0f84f905 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2295,7 +2295,13 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e ) if valid_token is not None: - valid_token = update_valid_token_with_end_user_params(valid_token=valid_token, end_user_params=end_user_params) + # A JWT auto-registration token already carries the verified customer. + # An absent request customer must not erase it before budget checks. + # Ordinary cached keys still reset request-derived identity each time. + if not valid_token.jwt_claims or end_user_params.get("end_user_id") is not None: + valid_token = update_valid_token_with_end_user_params( + valid_token=valid_token, end_user_params=end_user_params + ) valid_token = _update_key_budget_with_temp_budget_increase(valid_token) user_obj: LiteLLM_UserTable | None = None diff --git a/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py b/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py index fc1c2742658..f0b72a22244 100644 --- a/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py +++ b/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py @@ -230,7 +230,9 @@ async def test_master_key_auth_skips_end_user_model_budget_when_flag_disabled(): ) originals = {k: getattr(proxy_server, k, None) for k in attrs} flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = False # test-quality-ok: feature-flag toggle required by design (PR description) + litellm.enforce_end_user_model_max_budget_on_master_key = ( + False # test-quality-ok: feature-flag toggle required by design (PR description) + ) try: for k, v in attrs.items(): @@ -265,7 +267,9 @@ async def test_master_key_auth_skips_end_user_model_budget_when_flag_disabled(): assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} limiter.is_end_user_within_model_budget.assert_not_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + litellm.enforce_end_user_model_max_budget_on_master_key = ( + flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + ) for k, v in originals.items(): setattr(proxy_server, k, v) @@ -281,7 +285,9 @@ async def test_master_key_auth_passes_when_flag_enabled_and_within_budget(): attrs, limiter = _proxy_server_attrs_for_master_key_auth() originals = {k: getattr(proxy_server, k, None) for k in attrs} flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) + litellm.enforce_end_user_model_max_budget_on_master_key = ( + True # test-quality-ok: feature-flag toggle required by design (PR description) + ) try: for k, v in attrs.items(): @@ -319,7 +325,9 @@ async def test_master_key_auth_passes_when_flag_enabled_and_within_budget(): assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} limiter.is_end_user_within_model_budget.assert_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + litellm.enforce_end_user_model_max_budget_on_master_key = ( + flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + ) for k, v in originals.items(): setattr(proxy_server, k, v) @@ -383,7 +391,9 @@ async def test_master_key_auth_enforces_end_user_model_budget_when_flag_enabled( ) originals = {k: getattr(proxy_server, k, None) for k in attrs} flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) + litellm.enforce_end_user_model_max_budget_on_master_key = ( + True # test-quality-ok: feature-flag toggle required by design (PR description) + ) try: for k, v in attrs.items(): @@ -422,7 +432,9 @@ async def test_master_key_auth_enforces_end_user_model_budget_when_flag_enabled( assert exc_info.value.type == ProxyErrorTypes.budget_exceeded limiter.is_end_user_within_model_budget.assert_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + litellm.enforce_end_user_model_max_budget_on_master_key = ( + flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + ) for k, v in originals.items(): setattr(proxy_server, k, v) @@ -453,7 +465,9 @@ async def test_cached_master_key_auth_enforces_end_user_model_budget_when_flag_e attrs, limiter = _proxy_server_attrs_for_master_key_auth() originals = {k: getattr(proxy_server, k, None) for k in attrs} flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) + litellm.enforce_end_user_model_max_budget_on_master_key = ( + True # test-quality-ok: feature-flag toggle required by design (PR description) + ) try: for k, v in attrs.items(): @@ -495,7 +509,9 @@ async def test_cached_master_key_auth_enforces_end_user_model_budget_when_flag_e assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} limiter.is_end_user_within_model_budget.assert_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + litellm.enforce_end_user_model_max_budget_on_master_key = ( + flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + ) for k, v in originals.items(): setattr(proxy_server, k, v) @@ -528,7 +544,9 @@ async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcemen ) originals = {k: getattr(proxy_server, k, None) for k in attrs} flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) + litellm.enforce_end_user_model_max_budget_on_master_key = ( + True # test-quality-ok: feature-flag toggle required by design (PR description) + ) try: for k, v in attrs.items(): @@ -566,7 +584,9 @@ async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcemen assert result.user_role == LitellmUserRoles.PROXY_ADMIN limiter.is_end_user_within_model_budget.assert_not_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + litellm.enforce_end_user_model_max_budget_on_master_key = ( + flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + ) for k, v in originals.items(): setattr(proxy_server, k, v) @@ -581,7 +601,9 @@ async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcemen "/metrics", ], ) -async def test_master_key_budget_early_return_for_non_llm_routes(route): # test-quality-ok: structural assertion of mock wiring by design +async def test_master_key_budget_early_return_for_non_llm_routes( + route, +): # test-quality-ok: structural assertion of mock wiring by design """Branch coverage for ``_maybe_enforce_master_key_end_user_model_max_budget``. The flag and master-key guards in the helper are exercised by the @@ -597,7 +619,9 @@ async def test_master_key_budget_early_return_for_non_llm_routes(route): # test ) flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) + litellm.enforce_end_user_model_max_budget_on_master_key = ( + True # test-quality-ok: feature-flag toggle required by design (PR description) + ) try: valid_token = UserAPIKeyAuth( @@ -618,7 +642,9 @@ async def test_master_key_budget_early_return_for_non_llm_routes(route): # test ) mock_check.assert_not_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + litellm.enforce_end_user_model_max_budget_on_master_key = ( + flag_original # test-quality-ok: feature-flag toggle required by design (PR description) + ) @pytest.mark.asyncio @@ -658,3 +684,48 @@ async def test_cached_virtual_key_auth_applies_and_enforces_end_user_model_budge finally: for k, v in originals.items(): setattr(proxy_server, k, v) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "jwt_claims,request_params,expected_customer", + [ + ({"sub": "verified-customer"}, {}, "verified-customer"), + ({"sub": "verified-customer"}, {"end_user_id": "request-customer"}, "request-customer"), + (None, {}, None), + ], +) +async def test_virtual_key_validator_preserves_only_verified_jwt_customer( + jwt_claims, request_params, expected_customer +): + from datetime import datetime, timezone + from fastapi import Request + from litellm.proxy.auth.user_api_key_auth import validate_resolved_virtual_key + + token = UserAPIKeyAuth(end_user_id="verified-customer", jwt_claims=jwt_claims) + + class CustomerObserved(Exception): + pass + + def observe_token(valid_token): + assert valid_token.end_user_id == expected_customer + raise CustomerObserved + + # test-quality-ok: stop at the budget-check boundary to inspect identity before DB access + with patch( + "litellm.proxy.auth.user_api_key_auth._update_key_budget_with_temp_budget_increase", + side_effect=observe_token, + ): + with pytest.raises(CustomerObserved): + await validate_resolved_virtual_key( + request=Request(scope={"type": "http"}), + request_data={}, + valid_token=token, + api_key="test-key", + route="/v1/chat/completions", + start_time=datetime.now(timezone.utc), + parent_otel_span=None, + end_user_id=None, + end_user_params=request_params, + _end_user_object=None, + ) From e1560862671628fd32afe1209fe74edaf9aea7e7 Mon Sep 17 00:00:00 2001 From: Instinct Agent Date: Sun, 4 Oct 2026 13:58:07 +0530 Subject: [PATCH 6/7] test: preserve existing inline quality annotations --- ...t_end_user_model_max_budget_enforcement.py | 53 +++++-------------- 1 file changed, 13 insertions(+), 40 deletions(-) diff --git a/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py b/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py index f0b72a22244..fb71917448b 100644 --- a/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py +++ b/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py @@ -230,9 +230,7 @@ async def test_master_key_auth_skips_end_user_model_budget_when_flag_disabled(): ) originals = {k: getattr(proxy_server, k, None) for k in attrs} flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = ( - False # test-quality-ok: feature-flag toggle required by design (PR description) - ) + litellm.enforce_end_user_model_max_budget_on_master_key = False # test-quality-ok: feature-flag toggle required by design (PR description) try: for k, v in attrs.items(): @@ -267,9 +265,7 @@ async def test_master_key_auth_skips_end_user_model_budget_when_flag_disabled(): assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} limiter.is_end_user_within_model_budget.assert_not_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = ( - flag_original # test-quality-ok: feature-flag toggle required by design (PR description) - ) + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) for k, v in originals.items(): setattr(proxy_server, k, v) @@ -285,9 +281,7 @@ async def test_master_key_auth_passes_when_flag_enabled_and_within_budget(): attrs, limiter = _proxy_server_attrs_for_master_key_auth() originals = {k: getattr(proxy_server, k, None) for k in attrs} flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = ( - True # test-quality-ok: feature-flag toggle required by design (PR description) - ) + litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) try: for k, v in attrs.items(): @@ -325,9 +319,7 @@ async def test_master_key_auth_passes_when_flag_enabled_and_within_budget(): assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} limiter.is_end_user_within_model_budget.assert_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = ( - flag_original # test-quality-ok: feature-flag toggle required by design (PR description) - ) + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) for k, v in originals.items(): setattr(proxy_server, k, v) @@ -391,9 +383,7 @@ async def test_master_key_auth_enforces_end_user_model_budget_when_flag_enabled( ) originals = {k: getattr(proxy_server, k, None) for k in attrs} flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = ( - True # test-quality-ok: feature-flag toggle required by design (PR description) - ) + litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) try: for k, v in attrs.items(): @@ -432,9 +422,7 @@ async def test_master_key_auth_enforces_end_user_model_budget_when_flag_enabled( assert exc_info.value.type == ProxyErrorTypes.budget_exceeded limiter.is_end_user_within_model_budget.assert_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = ( - flag_original # test-quality-ok: feature-flag toggle required by design (PR description) - ) + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) for k, v in originals.items(): setattr(proxy_server, k, v) @@ -465,9 +453,7 @@ async def test_cached_master_key_auth_enforces_end_user_model_budget_when_flag_e attrs, limiter = _proxy_server_attrs_for_master_key_auth() originals = {k: getattr(proxy_server, k, None) for k in attrs} flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = ( - True # test-quality-ok: feature-flag toggle required by design (PR description) - ) + litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) try: for k, v in attrs.items(): @@ -509,9 +495,7 @@ async def test_cached_master_key_auth_enforces_end_user_model_budget_when_flag_e assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} limiter.is_end_user_within_model_budget.assert_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = ( - flag_original # test-quality-ok: feature-flag toggle required by design (PR description) - ) + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) for k, v in originals.items(): setattr(proxy_server, k, v) @@ -544,9 +528,7 @@ async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcemen ) originals = {k: getattr(proxy_server, k, None) for k in attrs} flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = ( - True # test-quality-ok: feature-flag toggle required by design (PR description) - ) + litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) try: for k, v in attrs.items(): @@ -584,9 +566,7 @@ async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcemen assert result.user_role == LitellmUserRoles.PROXY_ADMIN limiter.is_end_user_within_model_budget.assert_not_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = ( - flag_original # test-quality-ok: feature-flag toggle required by design (PR description) - ) + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) for k, v in originals.items(): setattr(proxy_server, k, v) @@ -601,9 +581,7 @@ async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcemen "/metrics", ], ) -async def test_master_key_budget_early_return_for_non_llm_routes( - route, -): # test-quality-ok: structural assertion of mock wiring by design +async def test_master_key_budget_early_return_for_non_llm_routes(route): # test-quality-ok: structural assertion of mock wiring by design """Branch coverage for ``_maybe_enforce_master_key_end_user_model_max_budget``. The flag and master-key guards in the helper are exercised by the @@ -619,9 +597,7 @@ async def test_master_key_budget_early_return_for_non_llm_routes( ) flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = ( - True # test-quality-ok: feature-flag toggle required by design (PR description) - ) + litellm.enforce_end_user_model_max_budget_on_master_key = True # test-quality-ok: feature-flag toggle required by design (PR description) try: valid_token = UserAPIKeyAuth( @@ -642,9 +618,7 @@ async def test_master_key_budget_early_return_for_non_llm_routes( ) mock_check.assert_not_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = ( - flag_original # test-quality-ok: feature-flag toggle required by design (PR description) - ) + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original # test-quality-ok: feature-flag toggle required by design (PR description) @pytest.mark.asyncio @@ -685,7 +659,6 @@ async def test_cached_virtual_key_auth_applies_and_enforces_end_user_model_budge for k, v in originals.items(): setattr(proxy_server, k, v) - @pytest.mark.asyncio @pytest.mark.parametrize( "jwt_claims,request_params,expected_customer", From f50232e3b2e1bcb078a2de8528be1b4eafc43691 Mon Sep 17 00:00:00 2001 From: Taranum Wasu Date: Wed, 7 Oct 2026 02:56:48 +0530 Subject: [PATCH 7/7] test(proxy): run end-user budget regressions in auth CI --- .github/workflows/test-unit-proxy-db.yml | 1 + ...t_end_user_model_max_budget_enforcement.py | 71 ++++++------------- 2 files changed, 23 insertions(+), 49 deletions(-) diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index c4397c2067b..d32d8d55ed9 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -82,6 +82,7 @@ jobs: test-path: >- tests/unit/proxy/auth/test_auth_checks.py tests/unit/proxy/auth/test_user_api_key_auth.py + tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py tests/unit/proxy/test_credential_slot_registry.py tests/unit/proxy/test_deprecated_key_grace_period.py workers: 4 diff --git a/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py b/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py index f0b72a22244..72411e345e7 100644 --- a/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py +++ b/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py @@ -211,7 +211,7 @@ async def test_enforce_end_user_model_max_budget_returns_early_when_unconfigured @pytest.mark.asyncio -async def test_master_key_auth_skips_end_user_model_budget_when_flag_disabled(): +async def test_master_key_auth_skips_end_user_model_budget_when_flag_disabled(monkeypatch): from fastapi import Request from starlette.datastructures import URL @@ -229,10 +229,7 @@ async def test_master_key_auth_skips_end_user_model_budget_when_flag_disabled(): litellm_budget_table=LiteLLM_BudgetTable(model_max_budget={MODEL: MODEL_BUDGET}), ) originals = {k: getattr(proxy_server, k, None) for k in attrs} - flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = ( - False # test-quality-ok: feature-flag toggle required by design (PR description) - ) + monkeypatch.setattr(litellm, "enforce_end_user_model_max_budget_on_master_key", False) try: for k, v in attrs.items(): @@ -267,15 +264,13 @@ async def test_master_key_auth_skips_end_user_model_budget_when_flag_disabled(): assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} limiter.is_end_user_within_model_budget.assert_not_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = ( - flag_original # test-quality-ok: feature-flag toggle required by design (PR description) - ) + pass for k, v in originals.items(): setattr(proxy_server, k, v) @pytest.mark.asyncio -async def test_master_key_auth_passes_when_flag_enabled_and_within_budget(): +async def test_master_key_auth_passes_when_flag_enabled_and_within_budget(monkeypatch): from fastapi import Request from starlette.datastructures import URL @@ -284,10 +279,7 @@ async def test_master_key_auth_passes_when_flag_enabled_and_within_budget(): attrs, limiter = _proxy_server_attrs_for_master_key_auth() originals = {k: getattr(proxy_server, k, None) for k in attrs} - flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = ( - True # test-quality-ok: feature-flag toggle required by design (PR description) - ) + monkeypatch.setattr(litellm, "enforce_end_user_model_max_budget_on_master_key", True) try: for k, v in attrs.items(): @@ -325,9 +317,7 @@ async def test_master_key_auth_passes_when_flag_enabled_and_within_budget(): assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} limiter.is_end_user_within_model_budget.assert_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = ( - flag_original # test-quality-ok: feature-flag toggle required by design (PR description) - ) + pass for k, v in originals.items(): setattr(proxy_server, k, v) @@ -372,7 +362,7 @@ async def test_virtual_key_auth_applies_and_enforces_end_user_model_budget(): @pytest.mark.asyncio -async def test_master_key_auth_enforces_end_user_model_budget_when_flag_enabled(): +async def test_master_key_auth_enforces_end_user_model_budget_when_flag_enabled(monkeypatch): from fastapi import Request from starlette.datastructures import URL @@ -390,10 +380,7 @@ async def test_master_key_auth_enforces_end_user_model_budget_when_flag_enabled( litellm_budget_table=LiteLLM_BudgetTable(model_max_budget={MODEL: MODEL_BUDGET}), ) originals = {k: getattr(proxy_server, k, None) for k in attrs} - flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = ( - True # test-quality-ok: feature-flag toggle required by design (PR description) - ) + monkeypatch.setattr(litellm, "enforce_end_user_model_max_budget_on_master_key", True) try: for k, v in attrs.items(): @@ -432,15 +419,13 @@ async def test_master_key_auth_enforces_end_user_model_budget_when_flag_enabled( assert exc_info.value.type == ProxyErrorTypes.budget_exceeded limiter.is_end_user_within_model_budget.assert_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = ( - flag_original # test-quality-ok: feature-flag toggle required by design (PR description) - ) + pass for k, v in originals.items(): setattr(proxy_server, k, v) @pytest.mark.asyncio -async def test_cached_master_key_auth_enforces_end_user_model_budget_when_flag_enabled(): +async def test_cached_master_key_auth_enforces_end_user_model_budget_when_flag_enabled(monkeypatch): from fastapi import Request from starlette.datastructures import URL @@ -464,10 +449,7 @@ async def test_cached_master_key_auth_enforces_end_user_model_budget_when_flag_e attrs, limiter = _proxy_server_attrs_for_master_key_auth() originals = {k: getattr(proxy_server, k, None) for k in attrs} - flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = ( - True # test-quality-ok: feature-flag toggle required by design (PR description) - ) + monkeypatch.setattr(litellm, "enforce_end_user_model_max_budget_on_master_key", True) try: for k, v in attrs.items(): @@ -509,15 +491,13 @@ async def test_cached_master_key_auth_enforces_end_user_model_budget_when_flag_e assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} limiter.is_end_user_within_model_budget.assert_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = ( - flag_original # test-quality-ok: feature-flag toggle required by design (PR description) - ) + pass for k, v in originals.items(): setattr(proxy_server, k, v) @pytest.mark.asyncio -async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcement(): +async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcement(monkeypatch): from fastapi import Request from starlette.datastructures import URL @@ -543,10 +523,7 @@ async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcemen message="Exceeded budget", current_cost=0.0002, max_budget=1e-05 ) originals = {k: getattr(proxy_server, k, None) for k in attrs} - flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = ( - True # test-quality-ok: feature-flag toggle required by design (PR description) - ) + monkeypatch.setattr(litellm, "enforce_end_user_model_max_budget_on_master_key", True) try: for k, v in attrs.items(): @@ -584,9 +561,7 @@ async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcemen assert result.user_role == LitellmUserRoles.PROXY_ADMIN limiter.is_end_user_within_model_budget.assert_not_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = ( - flag_original # test-quality-ok: feature-flag toggle required by design (PR description) - ) + pass for k, v in originals.items(): setattr(proxy_server, k, v) @@ -602,8 +577,9 @@ async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcemen ], ) async def test_master_key_budget_early_return_for_non_llm_routes( + monkeypatch, route, -): # test-quality-ok: structural assertion of mock wiring by design +): """Branch coverage for ``_maybe_enforce_master_key_end_user_model_max_budget``. The flag and master-key guards in the helper are exercised by the @@ -618,10 +594,7 @@ async def test_master_key_budget_early_return_for_non_llm_routes( _maybe_enforce_master_key_end_user_model_max_budget, ) - flag_original = litellm.enforce_end_user_model_max_budget_on_master_key - litellm.enforce_end_user_model_max_budget_on_master_key = ( - True # test-quality-ok: feature-flag toggle required by design (PR description) - ) + monkeypatch.setattr(litellm, "enforce_end_user_model_max_budget_on_master_key", True) try: valid_token = UserAPIKeyAuth( @@ -634,17 +607,17 @@ async def test_master_key_budget_early_return_for_non_llm_routes( "litellm.proxy.proxy_server.model_max_budget_limiter.is_end_user_within_model_budget", new_callable=AsyncMock, ) as mock_check: - await _maybe_enforce_master_key_end_user_model_max_budget( + result = await _maybe_enforce_master_key_end_user_model_max_budget( valid_token=valid_token, request_data={"user": "customer-1", "model": MODEL}, route=route, request=MagicMock(), ) + assert result is None + assert valid_token.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS mock_check.assert_not_awaited() finally: - litellm.enforce_end_user_model_max_budget_on_master_key = ( - flag_original # test-quality-ok: feature-flag toggle required by design (PR description) - ) + pass @pytest.mark.asyncio