mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
fix(key): allow /key/update to clear budget_limits with [] or null (#30085)
* Fix /key/update rejecting budget_limits clear requests with HTTP 400 Sending budget_limits: [] or null to /key/update returned HTTP 400, so once a key had budget windows the last one could never be removed. prepare_key_update_data only json.dumps'd budget_limits when the value was truthy, so [] and None passed through raw to the Prisma Json? column; jsonify_object only serializes dicts, and prisma-client-py has no DbNull sentinel for Json? writes, so Prisma rejected both shapes. Serialize the clear case explicitly as the JSON literal null, matching how memory_endpoints encodes metadata for the same column type. Truthy values keep the existing reset_at window initialization path. Fixes #30067. * Require admin access for budget_limits changes on /key/update Clearing budget_limits via [] or null is a budget mutation, but _validate_update_key_data only counted max_budget and spend as budget changes before deciding whether to skip _check_key_admin_access. A non-admin key owner or a team member with /key/update could therefore remove a key's per-window spend caps without admin authorization. Treat any explicit budget_limits value in the request (set, change, or clear) as a budget change so it gates through the same admin check as max_budget. model_fields_set is used because an explicit null is indistinguishable from an omitted field by value alone.
This commit is contained in:
Filippo Menghi
GitHub
parent
e89bd5c2e8
commit
93688ea7c1
2 changed files with 138 additions and 18 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue