From f96af80a0c8ed629b8866f1183306794a42a3509 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 11 Sep 2026 22:49:12 -0700 Subject: [PATCH 1/2] fix(proxy): persist clearing user model budgets --- .../internal_user_endpoints.py | 31 ++++- .../test_internal_user_endpoints.py | 129 +++++++++++++++++- 2 files changed, 150 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index d066f1e9138..ac85e432d8f 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -22,6 +22,7 @@ from typing import Any, Final, Literal, Protocol, cast, overload import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -30,6 +31,7 @@ from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import get_team_object, get_user_object from litellm.proxy.auth.password_policy import validate_password_policy from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.user_api_key_cache import ( object_permission_cache_key, user_object_permission_id_cache_key, @@ -96,6 +98,7 @@ if TYPE_CHECKING: from litellm.proxy.utils import ProxyLogging router: Final = APIRouter() +_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(GenericBudgetConfigType) def _user_table( @@ -1252,6 +1255,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda if k == "max_budget": if "max_budget" in fields_set: non_default_values[k] = v + elif k == "model_max_budget": + if k in fields_set: + _USER_MODEL_BUDGET_ADAPTER.validate_python({} if v is None else v) + non_default_values[k] = {} if v is None else v elif ( v is not None and v @@ -1421,7 +1428,7 @@ async def _update_single_user_helper( Returns the updated user data or raises an exception on failure. """ - from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name, prisma_client + from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name, prisma_client, user_api_key_cache if prisma_client is None: raise Exception("Not connected to DB!") @@ -1464,7 +1471,7 @@ async def _update_single_user_helper( # because `_update_internal_user_params` drops empty values, and `object_permission: {}` is # precisely the clear-my-own-ceiling case this must refuse. _sent_fields: Final = user_request.fields_set() if hasattr(user_request, "fields_set") else set() - _protected_fields: Final = ("max_budget", "soft_budget", "spend", "object_permission") + _protected_fields: Final = ("max_budget", "model_max_budget", "soft_budget", "spend", "object_permission") for _field in _protected_fields: if _field in non_default_values or _field in _sent_fields: raise HTTPException( @@ -1548,6 +1555,12 @@ async def _update_single_user_helper( await _invalidate_user_spend_counter_if_changed(non_default_values) + if "model_max_budget" in non_default_values: + await evict_and_broadcast( + cache_keys=(non_default_values["user_id"],), + user_api_key_cache=user_api_key_cache, + ) + if "object_permission_id" in non_default_values: await _invalidate_cached_user_entitlement( user_id=non_default_values.get("user_id"), @@ -1802,7 +1815,7 @@ async def bulk_user_update( }' ``` """ - from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client, user_api_key_cache if prisma_client is None: raise HTTPException( @@ -1867,9 +1880,19 @@ async def bulk_user_update( # Perform bulk database update await UserRepository(prisma_client).table.update_many( where={}, - data=non_default_values, # Update all users + data=( + {**non_default_values, "model_max_budget": json.dumps(non_default_values["model_max_budget"])} + if "model_max_budget" in non_default_values + else non_default_values + ), ) + if "model_max_budget" in non_default_values: + await evict_and_broadcast( + cache_keys=tuple(user.user_id for user in all_users_in_db), + user_api_key_cache=user_api_key_cache, + ) + # Create individual success results for user in all_users_in_db: results.append( 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 d1d669cae38..2f665f4acd3 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 @@ -1,9 +1,12 @@ import json from datetime import datetime, timezone from types import SimpleNamespace +from typing import Final import pytest from fastapi.testclient import TestClient +from pydantic import ValidationError +from pytest_mock import MockerFixture from litellm.proxy._types import ( @@ -2128,6 +2131,117 @@ def test_update_internal_user_params_keeps_original_max_budget_when_not_provided assert "user_alias" in non_default_values +@pytest.mark.parametrize("cleared_budget", [{}, None], ids=["empty-map", "null"]) +def test_update_internal_user_params_clears_model_budget(cleared_budget: dict[str, object] | None) -> None: + request: Final = UpdateUserRequest(user_id="user-spruce", model_max_budget=cleared_budget) + + update: Final = _update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request) + + assert update == {"user_id": "user-spruce", "model_max_budget": {}} + + +def test_update_internal_user_params_preserves_model_budget_presence_and_neighbors() -> None: + omitted: Final = UpdateUserRequest(user_id="user-spruce", user_alias="Spruce") + assert _update_internal_user_params(data_json=omitted.model_dump(), data=omitted) == { + "user_id": "user-spruce", + "user_alias": "Spruce", + } + + replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}} + request: Final = UpdateUserRequest( + user_id="user-spruce", + model_max_budget=replacement, + max_budget=50, + user_alias=None, + models=[], + allowed_cache_controls=[], + config={}, + ) + assert _update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request) == { + "user_id": "user-spruce", + "model_max_budget": replacement, + "max_budget": 50, + } + + +@pytest.mark.parametrize("invalid_budget", [{"model-spruce": "invalid"}, {"model-spruce": {"budget_limit": "invalid"}}]) +def test_update_internal_user_params_rejects_invalid_model_budget(invalid_budget: dict[str, object]) -> None: + request: Final = UpdateUserRequest(user_id="user-spruce", model_max_budget=invalid_budget) + + with pytest.raises(ValidationError): + _update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request) + + +@pytest.mark.asyncio +async def test_user_model_budget_update_by_email_refreshes_cached_user(mocker: MockerFixture) -> 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 + + saved_user: Final = LiteLLM_UserTable( + user_id="user-spruce", + user_email="spruce@example.test", + model_max_budget={"model-spruce": {"budget_limit": 5, "time_period": "1d"}}, + max_budget=50, + ) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) + 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, + ) + + await _update_single_user_helper( + user_request=UpdateUserRequest(user_email=saved_user.user_email, model_max_budget={}), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert prisma_client.update_data.call_args.kwargs["data"]["model_max_budget"] == {} + assert "max_budget" not in prisma_client.update_data.call_args.kwargs["data"] + 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) + + +@pytest.mark.asyncio +async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocker: MockerFixture) -> 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 bulk_user_update + from litellm.types.proxy.management_endpoints.internal_user_endpoints import BulkUpdateUserRequest + + saved_user: Final = LiteLLM_UserTable(user_id="user-spruce", model_max_budget={"model-spruce": {"budget_limit": 5}}) + prisma_client: Final = mocker.MagicMock() + 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) + 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, + ) + + response: Final = await bulk_user_update( + data=BulkUpdateUserRequest(all_users=True, user_updates={"model_max_budget": None}), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_changed_by=None, + ) + + prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(where={}, data={"model_max_budget": "{}"}) + prisma_client.update_data.assert_not_called() + assert response.successful_updates == 1 + assert response.results[0].updated_user["model_max_budget"] == {} + 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 @@ -3498,7 +3612,11 @@ def test_enforce_user_info_access_blocks_cross_user_lookup(): @pytest.mark.asyncio -async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker): +@pytest.mark.parametrize( + ("budget_field", "budget_value"), + [("max_budget", 999999), ("model_max_budget", {}), ("model_max_budget", None)], +) +async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker, budget_field, budget_value): """Non-admin updating their own record must be blocked from modifying max_budget (self-escalation).""" from fastapi import HTTPException @@ -3508,6 +3626,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker): ) mock_prisma_client = mocker.MagicMock() + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "user-1", "data": {"user_id": "user-1"}}) existing_user = mocker.MagicMock() existing_user.model_dump.return_value = { "user_id": "user-1", @@ -3519,10 +3638,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker): ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - user_request = UpdateUserRequest( - user_id="user-1", - max_budget=999999, - ) + user_request = UpdateUserRequest.model_validate({"user_id": "user-1", budget_field: budget_value}) caller = UserAPIKeyAuth( user_id="user-1", user_role=LitellmUserRoles.INTERNAL_USER, @@ -3533,7 +3649,8 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker): user_request=user_request, user_api_key_dict=caller ) assert exc.value.status_code == 403 - assert "max_budget" in str(exc.value.detail) + assert budget_field in str(exc.value.detail) + mock_prisma_client.update_data.assert_not_called() @pytest.mark.asyncio From f44d5ef101cbdf970ea5451260be0543b2a3d8be Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 11 Sep 2026 22:59:46 -0700 Subject: [PATCH 2/2] fix(proxy): preserve model budget update compatibility --- .../internal_user_endpoints.py | 22 +++++++++++++------ .../test_internal_user_endpoints.py | 17 +++++++++++--- 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index ac85e432d8f..10c11119006 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -22,7 +22,7 @@ from typing import Any, Final, Literal, Protocol, cast, overload import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -88,6 +88,7 @@ from litellm.types.proxy.management_endpoints.scim_v2 import ( SCIM_ENTITLEMENTS_METADATA_KEY, SCIM_ROLES_METADATA_KEY, ) +from litellm.types.utils import BudgetConfig if TYPE_CHECKING: from prisma import models as prisma_models @@ -98,7 +99,8 @@ if TYPE_CHECKING: from litellm.proxy.utils import ProxyLogging router: Final = APIRouter() -_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(GenericBudgetConfigType) +_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig]) +_USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE: Final = 50 def _user_table( @@ -1257,7 +1259,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda non_default_values[k] = v elif k == "model_max_budget": if k in fields_set: - _USER_MODEL_BUDGET_ADAPTER.validate_python({} if v is None else v) + try: + _USER_MODEL_BUDGET_ADAPTER.validate_python({} if v is None else v) + except ValidationError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc non_default_values[k] = {} if v is None else v elif ( v is not None @@ -1888,10 +1893,13 @@ async def bulk_user_update( ) if "model_max_budget" in non_default_values: - await evict_and_broadcast( - cache_keys=tuple(user.user_id for user in all_users_in_db), - user_api_key_cache=user_api_key_cache, - ) + for start in range(0, len(all_users_in_db), _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE): + await asyncio.gather( + *( + evict_and_broadcast(cache_keys=(user.user_id,), user_api_key_cache=user_api_key_cache) + for user in all_users_in_db[start : start + _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE] + ) + ) # Create individual success results for user in all_users_in_db: 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 2f665f4acd3..92b1ab1586d 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 @@ -5,7 +5,7 @@ from typing import Final import pytest from fastapi.testclient import TestClient -from pydantic import ValidationError +from fastapi import HTTPException from pytest_mock import MockerFixture @@ -2147,7 +2147,7 @@ def test_update_internal_user_params_preserves_model_budget_presence_and_neighbo "user_alias": "Spruce", } - replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}} + replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}, "model-birch": 5.0, "model-cedar": 0} request: Final = UpdateUserRequest( user_id="user-spruce", model_max_budget=replacement, @@ -2168,8 +2168,9 @@ def test_update_internal_user_params_preserves_model_budget_presence_and_neighbo def test_update_internal_user_params_rejects_invalid_model_budget(invalid_budget: dict[str, object]) -> None: request: Final = UpdateUserRequest(user_id="user-spruce", model_max_budget=invalid_budget) - with pytest.raises(ValidationError): + with pytest.raises(HTTPException) as exc: _update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request) + assert exc.value.status_code == 400 @pytest.mark.asyncio @@ -2228,6 +2229,16 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke new_callable=mocker.AsyncMock, ) + with pytest.raises(HTTPException) as exc: + await bulk_user_update( + data=BulkUpdateUserRequest(all_users=True, user_updates={"model_max_budget": {"model-spruce": "invalid"}}), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_changed_by=None, + ) + assert exc.value.status_code == 400 + prisma_client.db.litellm_usertable.update_many.assert_not_called() + assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) == saved_user + response: Final = await bulk_user_update( data=BulkUpdateUserRequest(all_users=True, user_updates={"model_max_budget": None}), user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN),