From 48bde68781c20df4d98915cc970eb65b23e343fe Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:47:14 +0000 Subject: [PATCH 1/3] fix(key_generate): use user's budget for UI session personal keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 13 ++-- .../test_key_management_endpoints.py | 77 +++++++++++++++++++ 2 files changed, 84 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 802a7c3e469..4f4e56d7418 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1233,11 +1233,10 @@ async def _common_key_generation_helper( # Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller # cannot grant a key a higher budget than their own authority. - is_ui_session_team_key = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID and _requested_team_id is not None - # Session tokens (lite login) carry max_budget=None to avoid a per-session - # LLM spend cap, but that None must not be read as "unlimited delegation - # authority". A personal key (no team) has no team-budget enforcement at - # request time, so a session token cannot delegate any budget for one. + # Session tokens (lite login) use their session max_budget for team keys, but + # personal keys are capped by user_max_budget when it is available. + is_ui_session_token: Final = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID + is_ui_session_team_key = is_ui_session_token and _requested_team_id is not None if ( user_api_key_dict.is_session_token and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value @@ -1255,7 +1254,9 @@ async def _common_key_generation_helper( }, ) delegation_ceiling: Final = ( - user_api_key_dict.max_budget + user_api_key_dict.user_max_budget + if is_ui_session_token and user_api_key_dict.user_max_budget is not None + else user_api_key_dict.max_budget if user_api_key_dict.max_budget is not None else (team_table.max_budget if user_api_key_dict.is_session_token and team_table is not None else None) ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index cc0a7631b59..809c4183e13 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -15509,6 +15509,83 @@ async def test_ghsa_q775_ui_session_token_personal_key_still_capped(): assert "cannot exceed" in msg.lower() +@pytest.mark.asyncio +async def test_ui_session_token_personal_key_ceiling_is_user_budget(): + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + data = GenerateKeyRequest(max_budget=100) + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-ui-session", + user_id="user-1", + team_id=UI_SESSION_TOKEN_TEAM_ID, + max_budget=1.0, + user_max_budget=500.0, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), # test-quality-ok: helper reads proxy_server.prisma_client directly + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: helper reads proxy_server.user_api_key_cache directly + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: helper reads proxy_server.llm_router directly + patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: helper reads proxy_server.premium_user directly + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"), # test-quality-ok: helper reads proxy_server.litellm_proxy_admin_name directly + patch( # test-quality-ok: helper has no dependency injection seam for key persistence + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key, + ): + mock_generate_key.return_value = {"key": "sk-test-key", "token_id": "token-id"} + try: + await _common_key_generation_helper( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + team_table=None, + ) + except (HTTPException, ProxyException) as err: + msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", "")) + assert "cannot exceed" not in msg.lower() + + +@pytest.mark.asyncio +async def test_ui_session_token_personal_key_above_user_budget_rejected(): + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + data = GenerateKeyRequest(max_budget=600) + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-ui-session", + user_id="user-1", + team_id=UI_SESSION_TOKEN_TEAM_ID, + max_budget=1.0, + user_max_budget=500.0, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), # test-quality-ok: helper reads proxy_server.prisma_client directly + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: helper reads proxy_server.user_api_key_cache directly + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: helper reads proxy_server.llm_router directly + patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: helper reads proxy_server.premium_user directly + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"), # test-quality-ok: helper reads proxy_server.litellm_proxy_admin_name directly + patch( # test-quality-ok: helper has no dependency injection seam for key persistence + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key, + ): + mock_generate_key.return_value = {"key": "sk-test-key", "token_id": "token-id"} + with pytest.raises((HTTPException, ProxyException)) as exc_info: + await _common_key_generation_helper( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + team_table=None, + ) + err = exc_info.value + code = getattr(err, "status_code", None) or getattr(err, "code", None) + msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", "")) + assert str(code) == "400" + assert "cannot exceed" in msg.lower() + assert "500.0" in msg + + @pytest.mark.asyncio async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption(): """ From ee7d2b50946b4d84f249e7119c346b253abc8bef Mon Sep 17 00:00:00 2001 From: jesus-berri Date: Fri, 18 Sep 2026 14:31:34 -0700 Subject: [PATCH 2/3] Update litellm/proxy/management_endpoints/key_management_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/key_management_endpoints.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 4f4e56d7418..f60e68bcfe7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1233,8 +1233,7 @@ async def _common_key_generation_helper( # Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller # cannot grant a key a higher budget than their own authority. - # Session tokens (lite login) use their session max_budget for team keys, but - # personal keys are capped by user_max_budget when it is available. + # UI session personal keys are capped by user_max_budget when it is available. is_ui_session_token: Final = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID is_ui_session_team_key = is_ui_session_token and _requested_team_id is not None if ( From 2e23c2d6536b883e2a37d3aa4dcd5e5647f8c040 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 20 Sep 2026 00:50:18 +0000 Subject: [PATCH 3/3] fix(user_update): evict cached user on max_budget change so the personal key ceiling refreshes on every worker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../internal_user_endpoints.py | 5 ++- .../test_internal_user_endpoints.py | 43 +++++++++++++++++++ 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index ba7a3309a90..0e5028d797b 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -101,6 +101,7 @@ if TYPE_CHECKING: router: Final = APIRouter() _USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig]) _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE: Final = 50 +_USER_BUDGET_CACHE_FIELDS: Final = frozenset({"max_budget", "model_max_budget"}) def _user_table( @@ -1561,7 +1562,7 @@ async def _update_single_user_helper( await _invalidate_user_spend_counter_if_changed(non_default_values) - if "model_max_budget" in non_default_values: + if not _USER_BUDGET_CACHE_FIELDS.isdisjoint(non_default_values): await evict_and_broadcast( cache_keys=(non_default_values["user_id"],), user_api_key_cache=user_api_key_cache, @@ -1892,7 +1893,7 @@ async def bulk_user_update( ), ) - if "model_max_budget" in non_default_values: + if not _USER_BUDGET_CACHE_FIELDS.isdisjoint(non_default_values): for start in range(0, len(all_users_in_db), _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE): await asyncio.gather( *( diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 0d8b19345f1..7655e6f80ff 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2269,6 +2269,49 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke broadcast.assert_awaited_once_with(cache_key=saved_user.user_id) +@pytest.mark.asyncio +@pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"]) +async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: MockerFixture, all_users: bool) -> None: + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper, bulk_user_update + from litellm.types.proxy.management_endpoints.internal_user_endpoints import BulkUpdateUserRequest + + saved_user: Final = LiteLLM_UserTable(user_id="user-spruce", max_budget=500.0) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) + prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user]) + prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) + prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) + prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=mocker.AsyncMock, + ) + admin: Final = UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN) + + if all_users: + await bulk_user_update( + data=BulkUpdateUserRequest(all_users=True, user_updates={"max_budget": 50.0}), + user_api_key_dict=admin, + litellm_changed_by=None, + ) + prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(where={}, data={"max_budget": 50.0}) + else: + await _update_single_user_helper( + user_request=UpdateUserRequest(user_id=saved_user.user_id, max_budget=50.0), + user_api_key_dict=admin, + ) + assert prisma_client.update_data.call_args.kwargs["data"]["max_budget"] == 50.0 + + assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) is None + broadcast.assert_awaited_once_with(cache_key=saved_user.user_id) + + def test_generate_request_base_validator(): """ Test that GenerateRequestBase validator converts empty string to None for max_budget