mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
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.
This commit is contained in:
parent
e15b37a18e
commit
8382e549e4
2 changed files with 75 additions and 11 deletions
|
|
@ -1861,18 +1861,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(
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue