fix(proxy): reject non-finite budget_limits windows on /key/generate (#31630)

Enforce that every `budget_limits[*].max_budget` is a finite number;
applies to every caller including proxy admin and runs before the
role / ceiling checks. Six parametrized regression tests cover NaN /
+inf / -inf for both non-admin and admin callers.
This commit is contained in:
yucheng-berri 2026-06-29 14:48:59 -07:00 • committed by GitHub
parent 829bfebe0f
commit be6b28f25e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 68 additions and 5 deletions

View file

@ -576,15 +576,20 @@ def _check_budget_limits_delegation_ceiling(
is_ui_session_team_key: bool,
) -> None:
"""
Enforce the delegation ceiling on every per-window budget entry.
Enforce two invariants on `budget_limits`:
The single-value `max_budget` check upstream guards the all-time budget;
`budget_limits` lets a key carry independent concurrent windows and was
bypassing the ceiling entirely, so a non-admin caller could mint a key
with a window budget far above their own authority.
- Every `budget_limits[*].max_budget` must be a finite number; applies
to every caller including proxy admin.
- Non-admin callers may not set a window above their delegation ceiling.
"""
if not budget_limits:
return
non_finite = next((w for w in budget_limits if not math.isfinite(w.max_budget)), None)
if non_finite is not None:
raise HTTPException(
status_code=400,
detail={"error": (f"budget_limits entry max_budget ({non_finite.max_budget}) must be a finite number.")},
)
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
if is_ui_session_team_key:

View file

@ -12854,6 +12854,64 @@ async def test_budget_limits_admin_unrestricted(monkeypatch):
assert result is not None
@pytest.mark.asyncio
@pytest.mark.parametrize("non_finite", [float("nan"), float("inf"), float("-inf")])
async def test_budget_limits_window_non_finite_rejected_for_non_admin(monkeypatch, non_finite):
"""A non-admin caller submitting a non-finite `budget_limits` window
gets 400. The finite-number invariant applies before role / ceiling
checks."""
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.litellm.default_key_generate_params",
None,
raising=False,
)
caller = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="user-1",
max_budget=10.0,
)
request = GenerateKeyRequest(
budget_limits=[{"budget_duration": "1d", "max_budget": non_finite}],
)
with pytest.raises(HTTPException) as exc_info:
await _common_key_generation_helper(
data=request,
user_api_key_dict=caller,
litellm_changed_by=None,
team_table=None,
)
assert exc_info.value.status_code == 400
assert "finite" in str(exc_info.value.detail)
@pytest.mark.asyncio
@pytest.mark.parametrize("non_finite", [float("nan"), float("inf"), float("-inf")])
async def test_budget_limits_window_non_finite_rejected_for_admin(monkeypatch, non_finite):
"""The finite-number invariant applies to every caller including
proxy admin."""
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.litellm.default_key_generate_params",
None,
raising=False,
)
admin = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-1",
)
request = GenerateKeyRequest(
budget_limits=[{"budget_duration": "1d", "max_budget": non_finite}],
)
with pytest.raises(HTTPException) as exc_info:
await _common_key_generation_helper(
data=request,
user_api_key_dict=admin,
litellm_changed_by=None,
team_table=None,
)
assert exc_info.value.status_code == 400
assert "finite" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_permissions_field_rejected_for_non_admin(monkeypatch):
"""A non-admin caller may not set the `permissions` field on a key