fix(proxy): enforce the personal max_budget on proxy-admin JWTs (#43070)

* fix(proxy): keep the personal budget on a proxy-admin user object

A token carrying the proxy_admin role (JWT admin scope, SSO) replaced the fetched LiteLLM_UserTable row with a bare admin object, so max_budget never reached common_checks and admins spent past their personal budget on /chat/completions. Keep the row and force only user_role.

Co-authored-by: L4XB <lukas.buck@e-mail.de>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): budget denial for a proxy-admin JWT is the 422 ExceededBudget response

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): carry only the personal budget onto a forced proxy-admin user object

Keeping the whole DB row also carried its models list, so a proxy admin
whose row had no-default-models or a restricted list got 403 on personal
requests. Build the admin object as before and copy max_budget and
user_email across. The built-in admin identity (the master key and the UI
username/password login) keeps ignoring its row's budget.

* test(proxy): pin the contrast case in the admin model-list and admin-name budget tests

Each test now also shows the same cached row refusing a non-exempt token, so the
pass case cannot go green on a row the gate never read.

* test(e2e): drive the personal-budget case with a proxy_admin user row and a plain JWT

The Buildkite stack keeps the default admin_allowed_routes, so an admin-scope JWT
gets 403 on /chat/completions before any budget check and the case could never
pass there. A plain token whose sub resolves to a proxy_admin row goes through the
same forced-admin rebuild. Drop the stage-mirror admin_allowed_routes edit, which
the deployed stack never read.

* ci: rerun e2e after an unrelated Gemini streaming spend-row timeout

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: L4XB <lukas.buck@e-mail.de>
Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>

litellm/proxy/auth/user_api_key_auth.py

@@ -3671,12 +3671,23 @@ async def run_centralized_common_checks(
     # may have a non-admin user_role and would otherwise demote the
     # caller. The token is the source of truth for these paths — force
     # the admin user_object whenever the token says PROXY_ADMIN, even
-    # if a DB row was fetched.
+    # if a DB row was fetched. Every admin token except the built-in admin
+    # identity (litellm_proxy_admin_name: the master key and the UI
+    # username/password login) keeps the row's personal budget (max_budget,
+    # and user_email for budget alerts) so common_checks enforces it; every
+    # other field keeps its default.
     if user_api_key_auth_obj.user_role == LitellmUserRoles.PROXY_ADMIN:
+        budget_row: Final = (
+            user_object
+            if isinstance(user_object, LiteLLM_UserTable) and user_api_key_auth_obj.user_id != litellm_proxy_admin_name
+            else None
+        )
         user_object = LiteLLM_UserTable(
             user_id=user_api_key_auth_obj.user_id or litellm_proxy_admin_name,
             user_role=LitellmUserRoles.PROXY_ADMIN,
             spend=user_object.spend if user_object is not None else 0.0,
+            max_budget=budget_row.max_budget if budget_row is not None else None,
+            user_email=budget_row.user_email if budget_row is not None else None,
             object_permission_id=(
                 user_object.object_permission_id if isinstance(user_object, LiteLLM_UserTable) else None
             ),

tests/e2e/coverage_registry/other.yaml

@@ -18,6 +18,7 @@
 - {id: other.auth.jwt.team_header_alias_binds_team, module: other, tier: P0, area: auth, assertions: [team_header_alias_binds_team], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", fail_before_fix: proven, rationale: "x-litellm-team-id carrying the team alias binds and attributes the same team as the team id, so a managed client can pin a stable alias instead of a uuid"}
 - {id: other.auth.jwt.team_header_non_member_alias_denied, module: other, tier: P0, area: auth, assertions: [team_header_non_member_alias_denied], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", rationale: "x-litellm-team-id naming the alias of a team the JWT does not grant is denied 403 with the same body as an unknown value, so the response does not reveal whether that team exists"}
 - {id: other.auth.jwt.team_model_alias_listed_and_routes, module: other, tier: P1, area: auth, assertions: [team_model_alias_listed_and_routes], source: "proxy_server.py model_list / common_utils/model_listing_utils.py alias_listing_entries / LIT-8515", fail_before_fix: proven, rationale: "A team model_aliases name the JWT caller can complete on is also listed by GET /v1/models for that caller, in the OpenAI and the Anthropic (Claude Code) shapes, next to its target, so a managed client can discover the alias it is meant to send"}
+- {id: other.auth.jwt.proxy_admin_personal_budget_enforced, module: other, tier: P0, area: auth, assertions: [proxy_admin_personal_budget_enforced], source: "user_api_key_auth.py _run_centralized_common_checks / auth_checks.py _user_max_budget_check / GH#41226 / LIT-8590", fail_before_fix: proven, rationale: "A JWT whose sub resolves to a proxy_admin user row with a spent max_budget is refused on /chat/completions, the same as a non-admin user; forcing the admin role must not drop the row's budget fields"}
 - {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"}
 - {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"}
 - {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"}

tests/e2e/models.py

@@ -1665,6 +1665,7 @@ class UserNewBody(BaseModel):
     user_role: UserRole
     user_id: str | None = None
     auto_create_key: bool | None = None
+    max_budget: float | None = None
 
 
 class UserNewResponse(BaseModel):

tests/e2e/other/test_jwt_auth_e2e.py

@@ -12,7 +12,7 @@ from e2e_config import CHEAP_OPENAI_MODEL, unique_marker
 from e2e_http import UnauthorizedError, UnknownApiError, unwrap
 from idp import SHORT_LIVED_CLIENT_ID, WRONG_AUDIENCE_CLIENT_ID, Identity
 from lifecycle import ResourceManager
-from models import ChatBody, ChatMessage, TeamNewBody
+from models import ChatBody, ChatMessage, TeamNewBody, UserNewBody
 from other_client import OtherClient
 from pydantic import BaseModel
 from e2e_metadata import Domain, Mode, Provider, Subject, meta
@@ -150,6 +150,38 @@ class TestJwtAuth:
             f"spend row must carry the user from the JWT sub claim {identity.user_id!r}, got {row.user!r}"
         )
 
+    @pytest.mark.covers("other.auth.jwt.proxy_admin_personal_budget_enforced")
+    @meta(
+        Subject(
+            domain=Domain.PROXY_AUTH,
+            models=(CHEAP_OPENAI_MODEL,),
+            mode=Mode.NONSTREAM,
+        )
+    )
+    def test_proxy_admin_token_is_blocked_once_its_personal_budget_is_spent(
+        self, client: OtherClient, resources: ResourceManager
+    ) -> None:
+        identity: Final = client.idp.provision_user(
+            marker=unique_marker(), groups=(), group_ids=(), defer=resources.defer
+        )
+        resources.defer(lambda: client.proxy.delete_user(identity.user_id))
+        unwrap(
+            client.user_new(
+                UserNewBody(
+                    user_id=identity.user_id,
+                    user_email=f"{identity.username}@example.com",
+                    user_role="proxy_admin",
+                    auto_create_key=False,
+                    max_budget=0.0,
+                )
+            )
+        )
+
+        blocked: Final = client.proxy.chat(client.idp.access_token(identity), _ping())
+        assert isinstance(blocked, UnknownApiError), f"proxy_admin JWT must be stopped by its max_budget, got {blocked}"
+        assert (blocked.status_code, "ExceededBudget" in blocked.body) == (422, True), blocked
+        assert f"User={identity.user_id}" in blocked.body, f"budget denial must name the JWT's user: {blocked.body}"
+
     @pytest.mark.covers("other.auth.jwt.invalid_signature_denied")
     @meta(
         Subject(

tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py

@@ -25,6 +25,7 @@ from starlette.datastructures import UploadFile
 import litellm
 import litellm.proxy.proxy_server
 from litellm.caching.dual_cache import DualCache
+from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
 from litellm.proxy.common_utils.user_api_key_cache import (
     NO_TEAM_MEMBERSHIP_SENTINEL,
     UserApiKeyCache,
@@ -4736,6 +4737,98 @@ async def test_centralized_common_checks_runs_for_standard_auth():
             setattr(_proxy_server_mod, k, v)
 
 
+async def _run_checks_for_cached_user(db_user: LiteLLM_UserTable, token: UserAPIKeyAuth) -> None:
+    """Run the centralized gate for `token` on /chat/completions with `db_user` as
+    the cached row its user_id resolves to."""
+    import litellm.proxy.proxy_server as _proxy_server_mod
+    from fastapi import Request
+    from starlette.datastructures import URL
+
+    from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+    from litellm.proxy.utils import ProxyLogging
+
+    request = Request(scope={"type": "http"})
+    request._url = URL(url="/chat/completions")
+    key_cache = UserApiKeyCache()
+    key_cache.set_cache(key=db_user.user_id, value=db_user)
+    attrs = {
+        **_proxy_attrs_for_centralized_checks(),
+        "prisma_client": MagicMock(),
+        "user_api_key_cache": key_cache,
+        "spend_counter_cache": DualCache(),
+        "proxy_logging_obj": ProxyLogging(user_api_key_cache=key_cache),
+    }
+    originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
+    try:
+        for k, v in attrs.items():
+            setattr(_proxy_server_mod, k, v)
+        await run_centralized_common_checks(
+            user_api_key_auth_obj=token,
+            request=request,
+            request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]},
+            route="/chat/completions",
+        )
+    finally:
+        for k, v in originals.items():
+            setattr(_proxy_server_mod, k, v)
+
+
+@pytest.mark.asyncio
+async def test_centralized_common_checks_enforce_personal_budget_on_proxy_admin_token():
+    """A proxy_admin token (JWT admin scope, SSO) stays subject to the DB row's
+    max_budget when the row is forced to the admin role, otherwise the personal
+    budget check in common_checks silently no-ops for admins."""
+    from litellm.proxy.utils import ProxyLogging
+
+    db_user = LiteLLM_UserTable(
+        user_id="admin-user", user_role="internal_user", spend=25.0, max_budget=10.0, user_email="admin@example.com"
+    )
+    token = UserAPIKeyAuth(api_key="sk-test", user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN)
+    with (
+        patch.object(ProxyLogging, "budget_alerts", new_callable=AsyncMock) as budget_alerts,
+        pytest.raises(litellm.BudgetExceededError) as exc_info,
+    ):
+        await _run_checks_for_cached_user(db_user, token)
+    assert (exc_info.value.max_budget, exc_info.value.current_cost) == (10.0, 25.0)
+    user_alerts = [c.kwargs["user_info"] for c in budget_alerts.call_args_list if c.kwargs["type"] == "user_budget"]
+    assert [a.user_email for a in user_alerts] == ["admin@example.com"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("row_models", [["no-default-models"], ["claude-haiku-4-5"]])
+async def test_centralized_common_checks_proxy_admin_token_ignores_row_model_list(row_models):
+    """A proxy_admin token keeps every model when its DB row carries a personal
+    model list, e.g. no-default-models from default_internal_user_params on a
+    user later promoted to admin; only the row's budget carries over. The same
+    row still refuses gpt-4o to a non-admin token."""
+    db_user = LiteLLM_UserTable(user_id="admin-user", user_role="internal_user", models=row_models, spend=1.0)
+    await _run_checks_for_cached_user(
+        db_user, UserAPIKeyAuth(api_key="sk-test", user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN)
+    )
+    with pytest.raises(ProxyException):
+        await _run_checks_for_cached_user(
+            db_user, UserAPIKeyAuth(api_key="sk-test", user_id="admin-user", user_role=LitellmUserRoles.INTERNAL_USER)
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("api_key", [LITELLM_PROXY_MASTER_KEY_ALIAS, "sk-ui-session"])
+async def test_centralized_common_checks_admin_name_identity_ignores_its_row_budget(api_key):
+    """The master key and the UI username/password login both resolve to the
+    litellm_proxy_admin_name row; that row's max_budget does not apply to them,
+    even once its spend is past the budget. The same row under any other admin
+    user_id is refused."""
+    db_user = LiteLLM_UserTable(user_id="admin", user_role="proxy_admin", spend=25.0, max_budget=10.0)
+    await _run_checks_for_cached_user(
+        db_user, UserAPIKeyAuth(api_key=api_key, user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
+    )
+    with pytest.raises(litellm.BudgetExceededError):
+        await _run_checks_for_cached_user(
+            db_user.model_copy(update={"user_id": "admin-user"}),
+            UserAPIKeyAuth(api_key=api_key, user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN),
+        )
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "route",