From 79c5d480da798fcaa9e0fb7718da124b486a645e Mon Sep 17 00:00:00 2001
From: yucheng-berriai <yucheng@berri.ai>
Date: Mon, 29 Jun 2026 12:31:37 -0700
Subject: [PATCH] fix(proxy): reject non-finite budget_limits windows on
 /key/generate

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.
---
 .../key_management_endpoints.py               | 15 +++--
 .../test_key_management_endpoints.py          | 58 +++++++++++++++++++
 2 files changed, 68 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index 368e411526f9..644d54f01398 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -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:
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 d9fee47aaa7f..f2a4ffae7475 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
@@ -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