fix: address pr comments

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -3517,6 +3517,20 @@ async def _build_model_max_budget_usage(
     return result
 
 
+def _budget_window_to_dict(window: object) -> dict | None:
+    """Coerce a budget_limits entry to a dict; None when the entry is unusable."""
+    if isinstance(window, dict):
+        return window
+    model_dump = getattr(window, "model_dump", None)
+    if callable(model_dump):
+        try:
+            dumped: Any = model_dump()
+        except Exception:  # noqa: BLE001
+            return None
+        return dumped if isinstance(dumped, dict) else None
+    return None
+
+
 async def _attach_budget_limits_usage(key_info: dict, api_key_hash: str) -> None:
     """
     Attach current-window spend to each entry in key_info["budget_limits"], in place.
@@ -3541,17 +3555,10 @@ async def _attach_budget_limits_usage(key_info: dict, api_key_hash: str) -> None
         return
 
     for idx, window in enumerate(budget_limits):
-        w: dict | None = None
-        if isinstance(window, dict):
-            w = window
-        elif hasattr(window, "model_dump"):
-            try:
-                w = window.model_dump()
-            except Exception:  # noqa: BLE001
-                continue
-            budget_limits[idx] = w
+        w: Final = _budget_window_to_dict(window)
         if not w:
             continue
+        budget_limits[idx] = w
         duration: Any = w.get("budget_duration")
         max_budget: Any = w.get("max_budget")
         if not duration:

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -13995,6 +13995,109 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch):
     }
 
 
+@pytest.mark.asyncio
+async def test_attach_budget_limits_usage_json_string_input(monkeypatch):
+    """budget_limits stored as a JSON string should be parsed and annotated."""
+    import json as json_module
+    from unittest.mock import AsyncMock
+
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _attach_budget_limits_usage,
+    )
+
+    mock_get_current_spend = AsyncMock(return_value=0.5)
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend
+    )
+
+    key_info = {
+        "budget_limits": json_module.dumps(
+            [{"budget_duration": "1h", "max_budget": 2.0, "reset_at": None}]
+        )
+    }
+    await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1")
+
+    assert isinstance(key_info["budget_limits"], list)
+    assert key_info["budget_limits"][0]["current_spend"] == 0.5
+
+
+@pytest.mark.asyncio
+async def test_attach_budget_limits_usage_skips_unusable_inputs(monkeypatch):
+    """Invalid JSON strings, non-list values, and malformed windows are skipped."""
+    from unittest.mock import AsyncMock
+
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _attach_budget_limits_usage,
+    )
+
+    mock_get_current_spend = AsyncMock(return_value=0.0)
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend
+    )
+
+    # invalid JSON string
+    key_info = {"budget_limits": "{not json"}
+    await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1")
+    assert key_info["budget_limits"] == "{not json"
+
+    # non-list value
+    key_info = {"budget_limits": {"budget_duration": "1h"}}
+    await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1")
+
+    # windows that are falsy, missing budget_duration, or not dict-like
+    key_info = {
+        "budget_limits": [
+            {},
+            {"max_budget": 2.0},
+            {"budget_duration": "1h", "max_budget": "not-a-number"},
+            42,
+        ]
+    }
+    await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1")
+
+    # only the well-formed window (with unparseable max_budget coerced to None)
+    # triggers a spend lookup
+    mock_get_current_spend.assert_awaited_once()
+    call_kwargs = mock_get_current_spend.await_args.kwargs
+    assert call_kwargs["counter_key"] == "spend:key:hash-1:window:1h"
+    assert call_kwargs["max_budget"] is None
+    assert key_info["budget_limits"][2]["current_spend"] == 0.0
+    assert key_info["budget_limits"][3] == 42
+
+
+@pytest.mark.asyncio
+async def test_attach_budget_limits_usage_pydantic_windows(monkeypatch):
+    """Window objects with model_dump() are converted to dicts in place."""
+    from unittest.mock import AsyncMock, MagicMock
+
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _attach_budget_limits_usage,
+    )
+
+    mock_get_current_spend = AsyncMock(return_value=1.0)
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend
+    )
+
+    good_window = MagicMock()
+    good_window.model_dump.return_value = {
+        "budget_duration": "7d",
+        "max_budget": 10.0,
+        "reset_at": None,
+    }
+    bad_window = MagicMock()
+    bad_window.model_dump.side_effect = ValueError("boom")
+
+    key_info = {"budget_limits": [good_window, bad_window]}
+    await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-2")
+
+    # good window converted to dict and annotated; failing window left as-is
+    assert isinstance(key_info["budget_limits"][0], dict)
+    assert key_info["budget_limits"][0]["current_spend"] == 1.0
+    assert key_info["budget_limits"][1] is bad_window
+    mock_get_current_spend.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch):
     """Cached spend for 'gpt-4o' matches budget key 'openai/gpt-4o' via suffix match."""