diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index ba88464d652..6a53143b994 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -97,6 +97,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/litellm/__init__.py b/litellm/__init__.py index 487d5da1663..51b9fc6319a 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 65f9b09a4d3..331aa6f6d5c 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,13 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e ) if valid_token is not None: + # 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 @@ -3816,6 +3836,66 @@ async def _enforce_key_and_fallback_model_access( ) +def _is_master_key_auth_token(valid_token: UserAPIKeyAuth) -> bool: + 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( + 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 not end_user_mmb + 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=cast(Mapping[str, object], 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..72411e345e7 --- /dev/null +++ b/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py @@ -0,0 +1,704 @@ +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(monkeypatch): + 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} + monkeypatch.setattr(litellm, "enforce_end_user_model_max_budget_on_master_key", False) + + 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: + 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(monkeypatch): + 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} + monkeypatch.setattr(litellm, "enforce_end_user_model_max_budget_on_master_key", True) + + 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: + pass + 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(monkeypatch): + 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} + monkeypatch.setattr(litellm, "enforce_end_user_model_max_budget_on_master_key", True) + + 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: + 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(monkeypatch): + 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} + monkeypatch.setattr(litellm, "enforce_end_user_model_max_budget_on_master_key", True) + + 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: + 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(monkeypatch): + 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} + monkeypatch.setattr(litellm, "enforce_end_user_model_max_budget_on_master_key", True) + + 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: + pass + 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( + monkeypatch, + route, +): + """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, + ) + + monkeypatch.setattr(litellm, "enforce_end_user_model_max_budget_on_master_key", True) + + 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: + 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: + pass + + +@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) + + +@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, + )