diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index eba16c077b0..2f239c8da84 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1860,18 +1860,23 @@ async def prepare_key_update_data( non_default_values["budget_reset_at"] = key_reset_at non_default_values["budget_duration"] = budget_duration - if "budget_limits" in non_default_values and non_default_values["budget_limits"]: - from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - + if "budget_limits" in non_default_values: raw_windows = non_default_values["budget_limits"] - initialized_windows = [] - for window in raw_windows: - w = window if isinstance(window, dict) else window.model_dump() - w["reset_at"] = get_budget_reset_time( - budget_duration=w["budget_duration"] - ).isoformat() - initialized_windows.append(w) - non_default_values["budget_limits"] = json.dumps(initialized_windows) + if raw_windows: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + initialized_windows = [] + for window in raw_windows: + w = window if isinstance(window, dict) else window.model_dump() + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + non_default_values["budget_limits"] = json.dumps(initialized_windows) + else: + # [] / None clears the field; prisma-client-py has no DbNull + # sentinel for Json? columns, so store the JSON literal null + non_default_values["budget_limits"] = json.dumps(None) if "object_permission" in non_default_values: non_default_values = await _handle_update_object_permission( @@ -2248,14 +2253,18 @@ async def _validate_update_key_data( # - Anyone else (non-PROXY_ADMIN, not the owner, not a team member # on a team key): must pass _check_key_admin_access (PROXY_ADMIN # / key-owner / team-admin / org-admin of the key). - # - max_budget / spend: always require the admin check, even for the - # key owner or a team member (matches the existing admin-only - # budget semantics). + # - 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. _is_budget_change = ( - data.max_budget is not None and data.max_budget != existing_key_row.max_budget - ) or ( - data.spend is not None - and data.spend != getattr(existing_key_row, "spend", None) + (data.max_budget is not None and data.max_budget != existing_key_row.max_budget) + or ( + data.spend is not None + and data.spend != getattr(existing_key_row, "spend", None) + ) + or "budget_limits" in data.model_fields_set ) # Personal-key bypass: the caller both created the key AND still owns it 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 473d61f8a85..046971d033b 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 @@ -1488,6 +1488,65 @@ async def test_prepare_key_update_data_duration_none_never_expires(): assert result["expires"] is None +@pytest.mark.asyncio +@pytest.mark.parametrize("cleared_value", [[], None]) +async def test_prepare_key_update_data_budget_limits_clears_field(cleared_value): + """budget_limits=[] / None must serialize to JSON null, never reach Prisma raw.""" + from litellm.proxy._types import UpdateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + prepare_key_update_data, + ) + + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=["gpt-3.5-turbo"], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", budget_limits=cleared_value) + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert result["budget_limits"] == json.dumps(None) + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_limits_serializes_windows(): + """Non-empty budget_limits stay JSON-encoded with reset_at initialized.""" + from litellm.proxy._types import UpdateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + prepare_key_update_data, + ) + + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=["gpt-3.5-turbo"], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest( + key="test-token", + budget_limits=[{"budget_duration": "1d", "max_budget": 10.0}], + ) + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + windows = json.loads(result["budget_limits"]) + assert isinstance(result["budget_limits"], str) + assert windows[0]["max_budget"] == 10.0 + assert windows[0]["reset_at"] is not None + + @pytest.mark.asyncio async def test_validate_team_id_used_in_service_account_request_requires_team_id(): """ @@ -9685,6 +9744,58 @@ class TestKeyOwnerPrivilegeEscalation: ) mock_check.assert_called_once() + @pytest.mark.asyncio + @pytest.mark.parametrize("cleared_value", [[], None]) + async def test_creator_cannot_clear_own_budget_limits(self, cleared_value): + """Clearing budget_limits is a budget change and requires admin.""" + data = UpdateKeyRequest(key="sk-test", budget_limits=cleared_value) + existing = self._make_existing_key(created_by="creator-123") + 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_admin_can_clear_budget_limits(self): + data = UpdateKeyRequest(key="sk-test", budget_limits=[]) + existing = self._make_existing_key(created_by="someone-else") + auth = UserAPIKeyAuth( + user_id="admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + mock_check = AsyncMock() + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access", + mock_check, + ): + 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_not_called() + @pytest.mark.asyncio async def test_admin_can_update_any_field(self): data = UpdateKeyRequest(key="sk-test", models=["gpt-4"], max_budget=999.0)