From 31130036c0013b8c0ad34e6cffbd58b24ee8b78a Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Mon, 17 Aug 2026 16:21:31 +0200 Subject: [PATCH] refactor: build budget window usage without in-place mutation Replace _attach_budget_limits_usage, which rewrote the caller's key_info dict, with _budget_limits_with_usage returning a new list. Callers assign the result once. Keeps the response shape and spend-counter read path identical while following the repo's no-mutation rule. --- .../key_management_endpoints.py | 119 +++++++++++------- .../test_key_management_endpoints.py | 83 ++++++------ 2 files changed, 121 insertions(+), 81 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 51b9ebf91d8..f6f739f3eae 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3517,23 +3517,44 @@ async def _build_model_max_budget_usage( return result -def _budget_window_to_dict(window: object) -> dict | None: +def _budget_window_to_dict(window: object) -> Mapping[str, object] | 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): + model_dump: Final = getattr(window, "model_dump", None) + if not callable(model_dump): + return None + try: + dumped: Final = model_dump() + except Exception: # noqa: BLE001 # model_dump implementations can raise arbitrary errors + return None + return dumped if isinstance(dumped, dict) else None + + +def _coerce_budget_limits(budget_limits: object) -> Sequence[object] | None: + """Coerce budget_limits to a sequence of windows, parsing JSON strings; None when unusable.""" + if isinstance(budget_limits, str): try: - dumped: Any = model_dump() - except Exception: # noqa: BLE001 + parsed: Final = json.loads(budget_limits) + except (TypeError, ValueError): + return None + return parsed if isinstance(parsed, list) else None + return budget_limits if isinstance(budget_limits, list) else None + + +def _parse_window_max_budget(value: object) -> float | None: + """Coerce a window's max_budget to float; None when absent or unparseable.""" + if isinstance(value, (int, float, str)): + try: + return float(value) + except (TypeError, ValueError): 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: +async def _budget_window_with_usage(window: Mapping[str, object], api_key_hash: str) -> Mapping[str, object]: """ - Attach current-window spend to each entry in key_info["budget_limits"], in place. + Return a copy of a budget window with current-window spend attached. Per-window spend is not persisted in the DB; it lives in the cross-pod spend counters (spend:key:{hashed_token}:window:{budget_duration}) that @@ -3544,39 +3565,43 @@ async def _attach_budget_limits_usage(key_info: dict, api_key_hash: str) -> None """ from litellm.proxy.proxy_server import get_current_spend - budget_limits: Any = key_info.get("budget_limits") - if isinstance(budget_limits, str): - try: - budget_limits = json.loads(budget_limits) - except (TypeError, ValueError): - return - key_info["budget_limits"] = budget_limits - if not isinstance(budget_limits, list): - return + duration: Final = window.get("budget_duration") + if not duration: + return dict(window) # mutable-ok: per-window response copy, built once per window + spend: Final = await get_current_spend( + counter_key=f"spend:key:{api_key_hash}:window:{duration}", + fallback_spend=0.0, + max_budget=_parse_window_max_budget(window.get("max_budget")), + window_entity_type="Key", + window_entity_id=api_key_hash, + window_start=get_budget_window_start(window), + ) + return {**window, "current_spend": round(spend, 4)} # mutable-ok: per-window response copy, built once per window - for idx, window in enumerate(budget_limits): - w = _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: - continue - try: - max_budget = float(max_budget) if max_budget is not None else None - except (TypeError, ValueError): - max_budget = None - counter_key = f"spend:key:{api_key_hash}:window:{duration}" - spend = await get_current_spend( - counter_key=counter_key, - fallback_spend=0.0, - max_budget=max_budget, - window_entity_type="Key", - window_entity_id=api_key_hash, - window_start=get_budget_window_start(w), - ) - w["current_spend"] = round(spend, 4) + +async def _budget_limits_entry_with_usage(window: object, api_key_hash: str) -> object: + """Return the window as an enriched dict when dict-coercible; the original entry otherwise.""" + coerced: Final = _budget_window_to_dict(window) + if not coerced: + return window + return await _budget_window_with_usage(window=coerced, api_key_hash=api_key_hash) + + +async def _budget_limits_with_usage(budget_limits: object, api_key_hash: str) -> Sequence[object] | None: + """ + Return budget_limits as window dicts with current-window spend attached. + + None when budget_limits is not a usable (possibly JSON-encoded) list; the + caller keeps the original value then. Entries that are not dict-coercible + are preserved as-is. + """ + windows: Final = _coerce_budget_limits(budget_limits) + if windows is None: + return None + return [ # mutable-ok: entries are awaited, so they cannot be built inside a frozen wrapper + await _budget_limits_entry_with_usage(window=window, api_key_hash=api_key_hash) + for window in windows + ] @router.post( @@ -3661,7 +3686,12 @@ async def info_key_fn_v2( user_api_key_cache=user_api_key_cache, ) if k_token_hash: - await _attach_budget_limits_usage(key_info=k_dict, api_key_hash=k_token_hash) + budget_limits_usage = await _budget_limits_with_usage( + budget_limits=k_dict.get("budget_limits"), + api_key_hash=k_token_hash, + ) + if budget_limits_usage is not None: + k_dict["budget_limits"] = budget_limits_usage filtered_key_info.append(k_dict) return {"key": data.keys, "info": filtered_key_info} @@ -3777,7 +3807,12 @@ async def info_key_fn( model_max_budget=model_max_budget, user_api_key_cache=user_api_key_cache, ) - await _attach_budget_limits_usage(key_info=key_info, api_key_hash=key_token_hash) + budget_limits_usage: Final = await _budget_limits_with_usage( + budget_limits=key_info.get("budget_limits"), + api_key_hash=key_token_hash, + ) + if budget_limits_usage is not None: + key_info["budget_limits"] = budget_limits_usage # Attach object_permission if object_permission_id is set key_info = await attach_object_permission_to_dict(key_info, prisma_client) 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 5df7c1cb28d..1f3b9042ed7 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 @@ -13996,13 +13996,13 @@ 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): +async def test_budget_limits_with_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, + _budget_limits_with_usage, ) mock_get_current_spend = AsyncMock(return_value=0.5) @@ -14010,24 +14010,29 @@ async def test_attach_budget_limits_usage_json_string_input(monkeypatch): "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") + raw = json_module.dumps( + [{"budget_duration": "1h", "max_budget": 2.0, "reset_at": None}] + ) + result = await _budget_limits_with_usage(budget_limits=raw, api_key_hash="hash-1") - assert isinstance(key_info["budget_limits"], list) - assert key_info["budget_limits"][0]["current_spend"] == 0.5 + assert result == [ + { + "budget_duration": "1h", + "max_budget": 2.0, + "reset_at": None, + "current_spend": 0.5, + } + ] + mock_get_current_spend.assert_awaited_once() @pytest.mark.asyncio -async def test_attach_budget_limits_usage_skips_unusable_inputs(monkeypatch): +async def test_budget_limits_with_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, + _budget_limits_with_usage, ) mock_get_current_spend = AsyncMock(return_value=0.0) @@ -14035,25 +14040,18 @@ async def test_attach_budget_limits_usage_skips_unusable_inputs(monkeypatch): "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") + # invalid JSON string and non-list values return None: callers keep the original + assert await _budget_limits_with_usage(budget_limits="{not json", api_key_hash="hash-1") is None + assert await _budget_limits_with_usage(budget_limits={"budget_duration": "1h"}, api_key_hash="hash-1") is None # 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") + windows = [ + {}, + {"max_budget": 2.0}, + {"budget_duration": "1h", "max_budget": "not-a-number"}, + 42, + ] + result = await _budget_limits_with_usage(budget_limits=windows, api_key_hash="hash-1") # only the well-formed window (with unparseable max_budget coerced to None) # triggers a spend lookup @@ -14061,17 +14059,22 @@ async def test_attach_budget_limits_usage_skips_unusable_inputs(monkeypatch): 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 + assert result is not None + assert result[0] == {} + assert result[1] == {"max_budget": 2.0} + assert result[2] == {"budget_duration": "1h", "max_budget": "not-a-number", "current_spend": 0.0} + assert result[3] == 42 + # input is not mutated + assert windows[2] == {"budget_duration": "1h", "max_budget": "not-a-number"} @pytest.mark.asyncio -async def test_attach_budget_limits_usage_pydantic_windows(monkeypatch): - """Window objects with model_dump() are converted to dicts in place.""" +async def test_budget_limits_with_usage_pydantic_windows(monkeypatch): + """Window objects with model_dump() are converted to dicts; failing windows pass through.""" from unittest.mock import AsyncMock, MagicMock from litellm.proxy.management_endpoints.key_management_endpoints import ( - _attach_budget_limits_usage, + _budget_limits_with_usage, ) mock_get_current_spend = AsyncMock(return_value=1.0) @@ -14088,13 +14091,15 @@ async def test_attach_budget_limits_usage_pydantic_windows(monkeypatch): 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") + result = await _budget_limits_with_usage( + budget_limits=[good_window, bad_window], 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 + assert result is not None + assert isinstance(result[0], dict) + assert result[0]["current_spend"] == 1.0 + assert result[1] is bad_window mock_get_current_spend.assert_awaited_once()