From e3c6258a7657e8648e40fa4b3994fad52f7bb8d3 Mon Sep 17 00:00:00 2001 From: Justin Israel Date: Tue, 6 Oct 2026 01:39:01 +0000 Subject: [PATCH 1/2] fix(proxy): require admin to clear max_budget on /key/update An explicit max_budget of null removes the key's cap, but the owner bypass only treated a non-null, changed max_budget as a budget change. A non-admin key owner could clear their own key budget without the admin check. Gate on the field being sent and different from the stored value instead, so the dashboard form echoing a stored null still counts as no change --- .../key_management_endpoints.py | 8 +-- .../test_key_management_endpoints.py | 53 +++++++++++++++++++ 2 files changed, 57 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2ee29f894b5..b6bf2a9b26f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3165,16 +3165,16 @@ async def _validate_update_key_data( # / key-owner / team-admin / org-admin of the key). # - max_budget / spend / budget_limits: always require the admin # check, even for the key owner or a team member (matches the - # existing admin-only budget semantics). budget_limits uses - # model_fields_set because an explicit null/[] clears the field - # and must gate the same as setting or changing it. + # existing admin-only budget semantics). max_budget and + # budget_limits use model_fields_set because an explicit null + # clears the field and must gate the same as setting or changing it. # - spend gates on presence alone (not a value diff): the DB spend # lags the live cross-pod counter, so letting an "unchanged" spend # through the non-admin path would let a key owner / team member # overwrite the live counter below real usage and silently weaken # enforcement. _is_budget_change: Final = ( - (data.max_budget is not None and data.max_budget != existing_key_row.max_budget) + ("max_budget" in data.model_fields_set and data.max_budget != existing_key_row.max_budget) or data.spend is not None or "budget_limits" in data.model_fields_set or "soft_budget" in data.model_fields_set diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index ce8f05e879e..48c6bddea1d 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -12386,6 +12386,59 @@ class TestKeyOwnerPrivilegeEscalation: ) mock_check.assert_called_once() + @pytest.mark.asyncio + async def test_creator_cannot_clear_own_max_budget(self): + """An explicit max_budget null removes the key's cap, so it is a budget change.""" + data = UpdateKeyRequest(key="sk-test", max_budget=None) + existing = self._make_existing_key(created_by="creator-123") + existing.max_budget = 10.0 + auth = self._make_auth(user_id="creator-123") + + mock_check = AsyncMock( + side_effect=HTTPException(status_code=403, detail="Not authorized") + ) + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access", + mock_check, + ): + with pytest.raises(HTTPException): + await _validate_update_key_data( + data=data, + existing_key_row=existing, + user_api_key_dict=auth, + llm_router=None, + premium_user=False, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + mock_check.assert_called_once() + + @pytest.mark.asyncio + async def test_creator_null_max_budget_on_uncapped_key_needs_no_admin(self): + """The dashboard form always sends max_budget, null when empty, so echoing the stored null is not a change.""" + data = UpdateKeyRequest(key="sk-test", max_budget=None, key_alias="renamed") + existing = self._make_existing_key(created_by="creator-123") + existing.max_budget = None + auth = self._make_auth(user_id="creator-123") + + mock_check = AsyncMock( + side_effect=HTTPException(status_code=403, detail="Not authorized") + ) + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access", + mock_check, + ): + result = await _validate_update_key_data( + data=data, + existing_key_row=existing, + user_api_key_dict=auth, + llm_router=None, + premium_user=False, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + assert result is None + @pytest.mark.asyncio @pytest.mark.parametrize("cleared_value", [[], None]) async def test_creator_cannot_clear_own_budget_limits(self, cleared_value): From 3b862be4706e017c969a41f20d3cc872be40fce9 Mon Sep 17 00:00:00 2001 From: Justin Israel Date: Tue, 6 Oct 2026 16:50:26 +1300 Subject: [PATCH 2/2] test(proxy): cover clearing max_budget through /key/update end to end --- .../test_key_management_endpoints.py | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 48c6bddea1d..eb8b27c96dc 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -3253,6 +3253,46 @@ async def test_update_key_with_key_and_alias_selects_by_key(monkeypatch): assert result["key"] == "sk-test-key" +@pytest.mark.asyncio +async def test_update_key_creator_clearing_max_budget_is_denied_and_cap_is_kept(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + key_in_db = LiteLLM_VerificationToken( + token="hashed-creator-key", + user_id="creator-123", + created_by="creator-123", + max_budget=10.0, + ) + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=key_in_db + ) + mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock( + return_value=None + ) + mock_prisma_client.update_data = AsyncMock() + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + + creator = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-creator", + user_id="creator-123", + ) + + with pytest.raises(ProxyException) as exc_info: + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key="sk-test-key", max_budget=None), + user_api_key_dict=creator, + litellm_changed_by=None, + ) + + assert str(exc_info.value.code) == "403" + mock_prisma_client.update_data.assert_not_called() + + @pytest.mark.asyncio async def test_block_key_existing_key_succeeds(monkeypatch): """