From ebc8e56aece36798cd358da2e1dca56dd77a0c7f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:04:12 -0700 Subject: [PATCH] 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 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 Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 13 ++- tests/e2e/coverage_registry/other.yaml | 1 + tests/e2e/models.py | 1 + tests/e2e/other/test_jwt_auth_e2e.py | 34 ++++++- .../test_user_api_key_auth_request_flow.py | 93 +++++++++++++++++++ 5 files changed, 140 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2231706b282..cdb3e685c20 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/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 ), diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 2c8bf4309ae..25fd76c9786 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/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"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 42a59af45d2..5b1c5778c10 100644 --- a/tests/e2e/models.py +++ b/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): diff --git a/tests/e2e/other/test_jwt_auth_e2e.py b/tests/e2e/other/test_jwt_auth_e2e.py index 6da3924c284..1fb2e2b1e9f 100644 --- a/tests/e2e/other/test_jwt_auth_e2e.py +++ b/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( diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 492ab408fc8..a0441981447 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/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",