From 84cc725a32d64884f013bb6bdf6c37a58b2c3f39 Mon Sep 17 00:00:00 2001 From: L4XB Date: Tue, 15 Sep 2026 16:50:05 +0200 Subject: [PATCH 1/3] fix(proxy): keep the personal budget on a proxy-admin user object The centralized authorization gate replaces the fetched user with a synthetic admin row whenever the token carries the PROXY_ADMIN role, so common_checks reads the admin role off user_object instead of the token. That rebuild kept only user_id, user_role and spend, so max_budget was lost and the personal-budget check returned early on a None budget. Copy the fetched row and override just the role, which keeps the admin authorization behaviour and leaves the budget fields intact. Closes #41226 --- litellm/proxy/auth/user_api_key_auth.py | 15 +-- .../proxy/auth/test_user_api_key_auth.py | 93 +++++++++++++++++++ 2 files changed, 102 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9491f77ecfc..e0608f747f3 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2724,13 +2724,16 @@ async def _run_centralized_common_checks( # same user_id (e.g. litellm_proxy_admin_name = "default_user_id") # 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. + # the admin role whenever the token says PROXY_ADMIN, while keeping + # the fetched row so its budget fields still reach common_checks. if user_api_key_auth_obj.user_role == LitellmUserRoles.PROXY_ADMIN: - 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, + user_object = ( + LiteLLM_UserTable( + user_id=user_api_key_auth_obj.user_id or litellm_proxy_admin_name, + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + if user_object is None + else user_object.model_copy(update={"user_role": LitellmUserRoles.PROXY_ADMIN}) ) if project_object is not None: diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index c9ae105d982..65509e97327 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -4251,6 +4251,99 @@ async def test_centralized_checks_skip_end_user_lookup_without_a_token_budget(): mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() +def _proxy_admin_world(user_row: LiteLLM_UserTable): + """The proxy globals the centralized gate reads, with ``user_row`` already in the + user cache so get_user_object resolves it without a DB round trip.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + key_cache = UserApiKeyCache() + key_cache.set_cache(key=user_row.user_id, value=user_row) + 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), + } + + @contextmanager + def _world(): + 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) + yield + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + return _world() + + +@pytest.mark.asyncio +async def test_centralized_checks_enforce_personal_budget_for_proxy_admin_token(): + """GH#41226: a proxy admin stays subject to their own personal budget. + + common_checks reads the admin role off user_object rather than off the token, so the + gate substitutes an admin user_object whenever the token says PROXY_ADMIN. Rebuilding + that object from scratch dropped max_budget, and the personal-budget check returns + early on a None budget, so an admin holding an external JWT kept getting completions + after their personal budget was spent. + """ + admin_row = LiteLLM_UserTable( + user_id="admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + spend=25.0, + max_budget=10.0, + ) + token = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) + + with _proxy_admin_world(admin_row): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=_chat_request(), + request_data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, + route="/chat/completions", + ) + + assert exc_info.value.max_budget == 10.0 + assert exc_info.value.current_cost == 25.0 + + +@pytest.mark.asyncio +async def test_centralized_checks_keep_admin_role_when_the_db_row_is_not_admin(): + """The other half of GH#41226: forcing the budget back must not stop the token from + forcing the admin role. A JWT or master-key admin whose DB row carries a non-admin + user_role still has to reach common_checks as PROXY_ADMIN, or admin-only routes would + start rejecting them.""" + demoted_row = LiteLLM_UserTable( + user_id="admin-user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + spend=1.0, + max_budget=500.0, + ) + token = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) + + with _proxy_admin_world(demoted_row): + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=_chat_request(), + request_data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, + route="/chat/completions", + ) + + seen_user_object = mock_checks.await_args.kwargs["user_object"] + assert seen_user_object.user_role == LitellmUserRoles.PROXY_ADMIN + assert seen_user_object.max_budget == 500.0 + + @pytest.mark.asyncio async def test_centralized_common_checks_runs_for_custom_auth_with_flag(): """Custom-auth deployments that opt in via custom_auth_run_common_checks From a6d5f8a2d1d2e8a122849ad25b692e38f0bf4eeb Mon Sep 17 00:00:00 2001 From: L4XB Date: Tue, 15 Sep 2026 17:24:17 +0200 Subject: [PATCH 2/3] test(proxy): declare the guard cell's patch to the test-quality gate The lint job's test-quality budget failed with TQ008 over its ceiling by one, naming the new guard cell. TQ008 asks a patched SDK internal to be replaced by a faked HTTP boundary or an injected collaborator, and neither applies here: the cell exists to observe what the proxy-admin substitution hands common_checks, so that call is the subject rather than incidental wiring, and there is no seam to inject on it. Suppressed with the reason inline, which is the mechanism the gate documents and the form the file already uses four times. The budget file is untouched. scripts/test_quality_gate.py --base c8114ba41f now reports every TQ rule within its ceiling. --- tests/test_litellm/proxy/auth/test_user_api_key_auth.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 65509e97327..b290b3507df 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -4328,7 +4328,7 @@ async def test_centralized_checks_keep_admin_role_when_the_db_row_is_not_admin() token = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) with _proxy_admin_world(demoted_row): - with patch( + with patch( # test-quality-ok: what the substitution hands common_checks IS the subject; no HTTP boundary and no injection seam on that call "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock, ) as mock_checks: From ac542607adc6787e3031f611503ff88685cc94e3 Mon Sep 17 00:00:00 2001 From: L4XB Date: Tue, 15 Sep 2026 23:32:46 +0200 Subject: [PATCH 3/3] test(proxy): assert the admin substitution through behavior, not through a mock The guard cell mocked common_checks and read user_object out of its call kwargs, so it pinned the shape of an internal call rather than anything a caller can see. It now drives an admin-only route with a non-admin DB row, asserts the call is allowed, and asserts the same route rejects the same row when the token is not admin, which is the observable consequence of forcing the role. That removes the TQ008 suppression the previous commit added. Docstrings cut to one line each. scripts/test_quality_gate.py --base c8114ba41f is clean. --- .../proxy/auth/test_user_api_key_auth.py | 56 ++++++++++--------- 1 file changed, 30 insertions(+), 26 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index b290b3507df..f5cdbdeabe4 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -4252,8 +4252,7 @@ async def test_centralized_checks_skip_end_user_lookup_without_a_token_budget(): def _proxy_admin_world(user_row: LiteLLM_UserTable): - """The proxy globals the centralized gate reads, with ``user_row`` already in the - user cache so get_user_object resolves it without a DB round trip.""" + """Same shape as _proxy_attrs_for_centralized_checks, with user_row in the cache.""" import litellm.proxy.proxy_server as _proxy_server_mod from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.utils import ProxyLogging @@ -4284,14 +4283,7 @@ def _proxy_admin_world(user_row: LiteLLM_UserTable): @pytest.mark.asyncio async def test_centralized_checks_enforce_personal_budget_for_proxy_admin_token(): - """GH#41226: a proxy admin stays subject to their own personal budget. - - common_checks reads the admin role off user_object rather than off the token, so the - gate substitutes an admin user_object whenever the token says PROXY_ADMIN. Rebuilding - that object from scratch dropped max_budget, and the personal-budget check returns - early on a None budget, so an admin holding an external JWT kept getting completions - after their personal budget was spent. - """ + """GH#41226: a proxy admin stays subject to their own personal budget.""" admin_row = LiteLLM_UserTable( user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN.value, @@ -4314,11 +4306,18 @@ async def test_centralized_checks_enforce_personal_budget_for_proxy_admin_token( @pytest.mark.asyncio -async def test_centralized_checks_keep_admin_role_when_the_db_row_is_not_admin(): - """The other half of GH#41226: forcing the budget back must not stop the token from - forcing the admin role. A JWT or master-key admin whose DB row carries a non-admin - user_role still has to reach common_checks as PROXY_ADMIN, or admin-only routes would - start rejecting them.""" +def _route_request(route: str): + from fastapi import Request + from starlette.datastructures import URL + + request = Request(scope={"type": "http"}) + request._url = URL(url=route) + return request + + +@pytest.mark.asyncio +async def test_centralized_checks_keep_admin_route_access_when_the_db_row_is_not_admin(): + """GH#41226: an admin token still reaches an admin-only route when its row is not admin.""" demoted_row = LiteLLM_UserTable( user_id="admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -4328,20 +4327,25 @@ async def test_centralized_checks_keep_admin_role_when_the_db_row_is_not_admin() token = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) with _proxy_admin_world(demoted_row): - with patch( # test-quality-ok: what the substitution hands common_checks IS the subject; no HTTP boundary and no injection seam on that call - "litellm.proxy.auth.user_api_key_auth.common_checks", - new_callable=AsyncMock, - ) as mock_checks: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=_route_request("/user/new"), + request_data={}, + route="/user/new", + ) + + non_admin_token = UserAPIKeyAuth( + user_id="admin-user", user_role=LitellmUserRoles.INTERNAL_USER + ) + with pytest.raises(Exception, match="Only proxy admin can be used") as exc_info: await _run_centralized_common_checks( - user_api_key_auth_obj=token, - request=_chat_request(), - request_data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, - route="/chat/completions", + user_api_key_auth_obj=non_admin_token, + request=_route_request("/user/new"), + request_data={}, + route="/user/new", ) - seen_user_object = mock_checks.await_args.kwargs["user_object"] - assert seen_user_object.user_role == LitellmUserRoles.PROXY_ADMIN - assert seen_user_object.max_budget == 500.0 + assert "internal_user" in str(exc_info.value) @pytest.mark.asyncio