diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 0b8baf92217..9a5d9faba38 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2557,7 +2557,39 @@ async def _validate_update_key_data( ) }, ) - if data.temp_budget_increase is not None and _delegation_ceiling is not None: + # Determine the effective temp_budget_increase after this update. + # Either the caller is setting a new one, or a previously persisted + # one in the key's metadata still applies (it is added on top of + # max_budget at request time by _update_key_budget_with_temp_budget_increase). + _effective_temp_increase: float | None = data.temp_budget_increase + if ( + _effective_temp_increase is None + and existing_key_row.metadata is not None + ): + _persisted_increase: Final = existing_key_row.metadata.get( + "temp_budget_increase" + ) + if _persisted_increase is not None: + try: + _persisted_expiry: Final = datetime.fromisoformat( + existing_key_row.metadata.get("temp_budget_expiry", "") + ) + if _persisted_expiry > datetime.now(timezone.utc): + _effective_temp_increase = float(_persisted_increase) + except (ValueError, TypeError): + pass + + if _effective_temp_increase is not None and _delegation_ceiling is not None: + if not math.isfinite(_effective_temp_increase): + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"temp_budget_increase must be a finite number. " + f"Received: {_effective_temp_increase}" + ) + }, + ) # temp_budget_increase is applied on top of the key's max_budget at # request time (user_api_key_auth._update_key_budget_with_temp_budget_increase), # so the effective budget must stay under the caller's ceiling too. @@ -2566,15 +2598,15 @@ async def _validate_update_key_data( ) if ( _effective_max_budget is not None - and _effective_max_budget + data.temp_budget_increase > _delegation_ceiling + and _effective_max_budget + _effective_temp_increase > _delegation_ceiling ): raise HTTPException( status_code=400, detail={ "error": ( f"max_budget plus temp_budget_increase " - f"({_effective_max_budget} + {data.temp_budget_increase} = " - f"{_effective_max_budget + data.temp_budget_increase}) cannot exceed " + f"({_effective_max_budget} + {_effective_temp_increase} = " + f"{_effective_max_budget + _effective_temp_increase}) cannot exceed " f"the caller's own max_budget ({_delegation_ceiling})." ) }, 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 310fe99fd6a..17f4e98a981 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 @@ -906,6 +906,65 @@ async def test_update_key_non_admin_cannot_raise_temp_budget_increase_above_ceil assert "cannot exceed the caller's own max_budget" in str(exc.value.detail) +@pytest.mark.asyncio +async def test_update_key_persisted_temp_increase_blocks_max_budget_raise(): + """A key with a persisted (non-expired) temp_budget_increase of 60 in its + metadata must not be raised from 30 to 80 by a caller capped at 100, + because the effective budget (80 + 60 = 140) exceeds the ceiling.""" + from datetime import datetime, timedelta, timezone + + from litellm.proxy._types import UpdateKeyRequest + + row = _update_key_ceiling_existing_row(30) + row.metadata = { + "temp_budget_increase": 60, + "temp_budget_expiry": ( + datetime.now(timezone.utc) + timedelta(days=1) + ).isoformat(), + } + data = UpdateKeyRequest(key="sk-alice-personal", max_budget=80) + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data=data, + existing_key_row=row, + user_api_key_dict=_update_key_ceiling_caller(100), + llm_router=None, + premium_user=True, + prisma_client=None, + user_api_key_cache=MagicMock(), + ) + assert exc.value.status_code == 400 + assert "cannot exceed the caller's own max_budget" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_update_key_nan_temp_budget_increase_rejected(): + """NaN temp_budget_increase must be rejected: any comparison with NaN + is False, so the ceiling check would silently pass.""" + from datetime import datetime, timedelta, timezone + + from litellm.proxy._types import UpdateKeyRequest + + data = UpdateKeyRequest( + key="sk-alice-personal", + max_budget=50, + temp_budget_increase=float("nan"), + temp_budget_expiry=datetime.now(timezone.utc) + timedelta(days=1), + ) + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data=data, + existing_key_row=_update_key_ceiling_existing_row(50), + user_api_key_dict=_update_key_ceiling_caller(100), + llm_router=None, + premium_user=True, + prisma_client=None, + user_api_key_cache=MagicMock(), + ) + assert exc.value.status_code == 400 + assert "finite" in str(exc.value.detail) + + @pytest.mark.asyncio async def test_update_key_proxy_admin_can_raise_max_budget_above_ceiling(): """PROXY_ADMIN is not bound by the delegation ceiling."""