diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 4fa6a1d0571..455afdefe7a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2267,6 +2267,11 @@ async def prepare_key_update_data( if "duration" in non_default_values: duration: Final = non_default_values.pop("duration") + if duration is None or isinstance(duration, str): + # Reject garbage here so bulk/team-bulk/regenerate callers that + # skip _validate_update_key_data still 400 instead of 500ing + # inside duration math below. + validate_key_duration(duration) if duration is None or duration == "-1": # Set expires to None to indicate the key never expires non_default_values["expires"] = None @@ -2277,6 +2282,9 @@ async def prepare_key_update_data( if "budget_duration" in non_default_values: budget_duration: Final = non_default_values.pop("budget_duration") + if budget_duration is None or isinstance(budget_duration, str): + # Same write-boundary guarantee as duration above. + validate_budget_duration(budget_duration) if budget_duration is None: non_default_values["budget_duration"] = None non_default_values["budget_reset_at"] = None 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 65cc23ea67f..f06f40c98dd 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 @@ -1796,6 +1796,32 @@ async def test_update_service_account_works_with_team_id(): await prepare_key_update_data(data=data, existing_key_row=existing_key) +@pytest.mark.asyncio +async def test_update_key_garbage_duration_400s_instead_of_500(): + """Regression: bulk/team-bulk/regenerate skip _validate_update_key_data, so + prepare_key_update_data itself must reject garbage durations (#39711).""" + data = UpdateKeyRequest(key="sk-1", duration="not-a-duration") + existing_key = LiteLLM_VerificationToken(token="hashed") + + with pytest.raises(HTTPException) as exc_info: + await prepare_key_update_data(data=data, existing_key_row=existing_key) + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("duration", ["7d", None, "-1"]) +async def test_update_key_valid_durations_still_flow_through(duration): + data = UpdateKeyRequest(key="sk-1", duration=duration) + existing_key = LiteLLM_VerificationToken(token="hashed") + + updated = await prepare_key_update_data(data=data, existing_key_row=existing_key) + + if duration in (None, "-1"): + assert updated["expires"] is None + else: + assert updated["expires"] is not None + + @pytest.mark.asyncio @pytest.mark.parametrize("flag_value", [True, False]) async def test_update_key_enable_prompt_caching_folds_into_metadata(flag_value):