fix(proxy): close persisted-temp-increase and NaN bypasses in /key/update delegation ceiling

Two review findings on the delegation-ceiling guard:

1. Persisted temp_budget_increase bypass: when a caller updates only
   max_budget without setting temp_budget_increase, the key row's
   metadata may still carry a non-expired persisted increase. The
   effective budget (new max_budget + retained increase) could exceed
   the caller's ceiling. Now reads the persisted value from metadata
   (checking expiry) and includes it in the check.

2. NaN temp_budget_increase bypass: any comparison with NaN is False,
   so _effective_max_budget + float('nan') > ceiling silently passed.
   Added a math.isfinite guard that rejects non-finite values with a
   400 before the comparison runs.

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})."
                         )
                     },

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."""