Merge 89d7a3890f into 252c71c0b2

.github/workflows/test-unit-proxy-db.yml

@@ -173,6 +173,7 @@ jobs:
           - test-group: budgets
             test-path: >-
               tests/proxy_unit_tests/test_default_end_user_budget_simple.py
+              tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py
               tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
               tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py
             workers: 4

litellm/__init__.py

@@ -469,6 +469,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

litellm/proxy/auth/user_api_key_auth.py

@@ -1870,6 +1870,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 (
@@ -1932,6 +1939,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),
@@ -1941,11 +1961,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
@@ -2007,10 +2022,13 @@ 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.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
+            )
+            # update key budget with temp budget increase
+            valid_token = _update_key_budget_with_temp_budget_increase(
+                valid_token
+            )  # updating it here, allows all downstream reporting / checks to use the updated budget
 
         if valid_token is not None:
             valid_token = _update_key_budget_with_temp_budget_increase(valid_token)
@@ -2246,20 +2264,12 @@ async def _user_api_key_auth_builder(
                         )
 
                     # Check 5b. End-user model max budget
-                    end_user_mmb: Final = valid_token.end_user_model_max_budget
-                    if (
-                        end_user_mmb is not None
-                        and isinstance(end_user_mmb, dict)
-                        and len(end_user_mmb) > 0
-                        and current_models
-                        and valid_token.end_user_id is not None
-                    ):
-                        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,
-                            )
+                    await _enforce_end_user_model_max_budget_checks(
+                        valid_token=valid_token,
+                        request_data=request_data,
+                        route=route,
+                        request=request,
+                    )
 
             # Check 6: Additional Common Checks across jwt + key auth
             if valid_token.team_id is not None:
@@ -3336,6 +3346,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,

tests/benchmarks/conftest.py

@@ -6,8 +6,20 @@ virtual CPU. Work deferred to litellm's shared logging executor would therefore
 be attributed to whichever benchmark the valgrind scheduler resumes it under,
 flipping results between runs. Running the executor inline keeps each
 benchmark's cost self-contained and deterministic.
+
+The benchmarks also disable Python's cyclic garbage collector for the
+duration of the measurement window. CPython's GC is non-deterministic and
+runs on its own clock; a collection triggered mid-benchmark inflates the
+per-call instruction count in a way that depends on when (and whether) the
+collector happened to fire rather than on anything the code under test does.
+CodSpeed's per-iteration measurement is small enough (~hundreds of
+microseconds) that this noise dominates the signal for the multi-turn
+benchmark. ``mock_response`` already isolates the benchmarks from any
+real network I/O, so the synthetic allocations here have no live-tracing
+implications: deferring GC until the session ends is safe.
 """
 
+import gc
 from collections.abc import Callable, Iterator
 from concurrent.futures import Future
 from typing import ParamSpec, TypeVar
@@ -34,3 +46,21 @@ def inline_logging_executor() -> Iterator[None]:
     executor.submit = _submit_inline
     yield
     del executor.submit
+
+
+@pytest.fixture(autouse=True, scope="session")
+def disable_gc_during_benchmarks() -> Iterator[None]:
+    """Disable CPython's cyclic GC for the duration of the benchmark session.
+
+    CodSpeed counts instructions per measured iteration; a GC that happens to
+    run mid-iteration shows up as a deterministic-looking inflation that flips
+    between runs (because ``gc.collect()`` fires on its own clock). ``mock_response``
+    keeps the SDK from allocating anything that escapes the benchmark loop, so
+    the deferred collection at session teardown stays bounded.
+    """
+    gc.disable()
+    try:
+        yield
+    finally:
+        gc.enable()
+        gc.collect()

tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py

@@ -0,0 +1,621 @@
+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):
+    async def mock_resolve_key(self, hashed_token: str):
+        from litellm.proxy.auth.resolvers.store import KeyNotInCacheError
+
+        if self._check_cache_only:
+            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)