diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 02bfa5c0ca8..c43b6ddf06a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -20,6 +20,7 @@ import secrets import traceback from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, cast import fastapi @@ -3590,9 +3591,12 @@ def _window_max_budget(window: Mapping[str, object]) -> float | None: return None -async def _budget_window_with_usage(window: Mapping[str, object], api_key_hash: str) -> Mapping[str, object]: +async def _budget_window_usage( + window: Mapping[str, object], api_key_hash: str +) -> tuple[str, Mapping[str, object]] | None: """ - Copy of a budget window with current-window spend attached. + (budget_duration, usage entry) for one budget window; None when the window + has no budget_duration to key it by. Reads the same cross-pod counter (spend:key:{hashed_token}:window:{budget_duration}) that _virtual_key_multi_budget_check enforces against, passing the same @@ -3603,34 +3607,41 @@ async def _budget_window_with_usage(window: Mapping[str, object], api_key_hash: duration: Final = window.get("budget_duration") if not isinstance(duration, str) or not duration: - return window + return None + max_budget: Final = _window_max_budget(window) spend: Final = await get_current_spend( counter_key=f"spend:key:{api_key_hash}:window:{duration}", fallback_spend=0.0, - max_budget=_window_max_budget(window), + max_budget=max_budget, window_entity_type="Key", window_entity_id=api_key_hash, window_duration=duration, window_start=get_budget_window_start(window), ) - return {**window, "current_spend": round(spend, 4)} # mutable-ok: per-window response copy, built once per window + return duration, MappingProxyType( + { + "current_spend": round(spend, 4), + "budget_limit": max_budget, + "reset_at": window.get("reset_at"), + } + ) -async def _budget_limits_with_usage( +async def _build_budget_limits_usage( budget_limits: Sequence[object] | str | None, api_key_hash: str -) -> tuple[Mapping[str, object], ...] | None: +) -> Mapping[str, Mapping[str, object]] | None: """ - budget_limits as window dicts with current-window spend attached; None when - the key has no windows so the caller keeps the stored value. + Current-window spend per budget window, keyed by budget_duration, reported + next to the stored budget_limits (which is returned untouched). None when + the key has no windows, so the field only appears on keys that have them. """ windows: Final = _budget_limit_windows(budget_limits) if not windows: return None - return tuple( - await asyncio.gather( - *(_budget_window_with_usage(window=window, api_key_hash=api_key_hash) for window in windows) - ) + usages: Final = await asyncio.gather( + *(_budget_window_usage(window=window, api_key_hash=api_key_hash) for window in windows) ) + return MappingProxyType({duration: usage for duration, usage in (u for u in usages if u is not None)}) @router.post( @@ -3717,12 +3728,12 @@ async def info_key_fn_v2( user_api_key_cache=model_max_budget_limiter.dual_cache, ) if k_token_hash: - budget_limits_usage = await _budget_limits_with_usage( + budget_limits_usage = await _build_budget_limits_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 + k_dict["budget_limits_usage"] = budget_limits_usage filtered_key_info.append(k_dict) return {"key": data.keys, "info": filtered_key_info} @@ -3759,9 +3770,10 @@ async def info_key_fn( - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} - model_max_budget_usage: dict | None - Current-window spend per model, present only when the key has per-model budgets - - budget_limits: list | None - Concurrent budget windows. Each entry includes - current_spend: spend accumulated in the window so far (read from the same cross-pod - spend counter the budget enforcement uses) + - budget_limits: list | None - Concurrent budget windows, exactly as stored + - budget_limits_usage: dict | None - Current-window spend per budget window, keyed by + budget_duration, present only when the key has budget windows (read from the same + cross-pod spend counter the budget enforcement uses) - models: list - Model_name's the key is allowed to call - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"} @@ -3841,12 +3853,12 @@ async def info_key_fn( model_max_budget=model_max_budget, user_api_key_cache=model_max_budget_limiter.dual_cache, ) - budget_limits_usage: Final = await _budget_limits_with_usage( + budget_limits_usage: Final = await _build_budget_limits_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 + key_info["budget_limits_usage"] = 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 425365c8560..f6e57607717 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 @@ -14143,11 +14143,11 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch): @pytest.mark.asyncio -async def test_info_key_fn_budget_limits_includes_current_spend(monkeypatch): +async def test_info_key_fn_reports_budget_limits_usage(monkeypatch): """ - /key/info should attach current_spend to each budget_limits window, read from - the same spend counter (spend:key:{token}:window:{duration}) that budget - enforcement uses. + /key/info reports current-window spend per budget window under budget_limits_usage, + keyed by budget_duration and read from the same counter enforcement uses, while + budget_limits itself comes back exactly as stored. """ from unittest.mock import AsyncMock, MagicMock @@ -14204,11 +14204,14 @@ async def test_info_key_fn_budget_limits_includes_current_spend(monkeypatch): user_api_key_dict=user_api_key_dict, ) - windows = result["info"]["budget_limits"] - assert len(windows) == 1 - assert windows[0]["current_spend"] == 0.73 - assert windows[0]["max_budget"] == 2.0 - assert windows[0]["budget_duration"] == "1h" + assert result["info"]["budget_limits"] == budget_limits + assert result["info"]["budget_limits_usage"] == { + "1h": { + "current_spend": 0.73, + "budget_limit": 2.0, + "reset_at": "2026-08-15T18:00:00+00:00", + } + } mock_get_current_spend.assert_awaited_once() call_kwargs = mock_get_current_spend.await_args.kwargs @@ -14222,7 +14225,7 @@ async def test_info_key_fn_budget_limits_includes_current_spend(monkeypatch): @pytest.mark.asyncio async def test_info_key_fn_no_budget_limits_skips_spend_lookup(monkeypatch): - """Keys without budget_limits should not trigger window spend lookups.""" + """Keys without budget windows get no budget_limits_usage field and trigger no spend lookup.""" from unittest.mock import AsyncMock, MagicMock from litellm.proxy._types import LiteLLM_VerificationToken @@ -14272,12 +14275,13 @@ async def test_info_key_fn_no_budget_limits_skips_spend_lookup(monkeypatch): ) assert result["info"]["budget_limits"] is None + assert "budget_limits_usage" not in result["info"] mock_get_current_spend.assert_not_awaited() @pytest.mark.asyncio -async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): - """/v2/key/info should attach current_spend to each budget_limits window.""" +async def test_info_key_fn_v2_reports_budget_limits_usage(monkeypatch): + """/v2/key/info reports budget_limits_usage per window and leaves budget_limits as stored.""" from unittest.mock import AsyncMock, MagicMock from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken @@ -14286,6 +14290,18 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): ) test_key_token = "hashed_token_v2_window_test" + budget_limits = [ + { + "reset_at": "2026-08-15T18:00:00+00:00", + "max_budget": 2.0, + "budget_duration": "1h", + }, + { + "reset_at": "2026-08-16T00:00:00+00:00", + "max_budget": 20.0, + "budget_duration": "1d", + }, + ] mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -14304,18 +14320,7 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): mock_key.team_id = None mock_key.model_dump.return_value = { "token": test_key_token, - "budget_limits": [ - { - "reset_at": "2026-08-15T18:00:00+00:00", - "max_budget": 2.0, - "budget_duration": "1h", - }, - { - "reset_at": "2026-08-16T00:00:00+00:00", - "max_budget": 20.0, - "budget_duration": "1d", - }, - ], + "budget_limits": [dict(w) for w in budget_limits], "user_id": "user-v2-w", "team_id": None, "litellm_budget_table": None, @@ -14335,10 +14340,19 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): ) assert len(result["info"]) == 1 - windows = result["info"][0]["budget_limits"] - assert len(windows) == 2 - assert windows[0]["current_spend"] == 1.25 - assert windows[1]["current_spend"] == 1.25 + assert result["info"][0]["budget_limits"] == budget_limits + assert result["info"][0]["budget_limits_usage"] == { + "1h": { + "current_spend": 1.25, + "budget_limit": 2.0, + "reset_at": "2026-08-15T18:00:00+00:00", + }, + "1d": { + "current_spend": 1.25, + "budget_limit": 20.0, + "reset_at": "2026-08-16T00:00:00+00:00", + }, + } assert mock_get_current_spend.await_count == 2 counter_keys = { call.kwargs["counter_key"] for call in mock_get_current_spend.await_args_list @@ -14353,13 +14367,13 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): @pytest.mark.asyncio -async def test_budget_limits_with_usage_json_string_input(monkeypatch): - """budget_limits stored as a JSON string should be parsed and annotated.""" +async def test_build_budget_limits_usage_json_string_input(monkeypatch): + """budget_limits stored as a JSON string is parsed and reported per window.""" import json as json_module from unittest.mock import AsyncMock from litellm.proxy.management_endpoints.key_management_endpoints import ( - _budget_limits_with_usage, + _build_budget_limits_usage, ) mock_get_current_spend = AsyncMock(return_value=0.5) @@ -14370,26 +14384,19 @@ async def test_budget_limits_with_usage_json_string_input(monkeypatch): 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") + result = await _build_budget_limits_usage(budget_limits=raw, api_key_hash="hash-1") - assert list(result) == [ - { - "budget_duration": "1h", - "max_budget": 2.0, - "reset_at": None, - "current_spend": 0.5, - } - ] + assert result == {"1h": {"current_spend": 0.5, "budget_limit": 2.0, "reset_at": None}} mock_get_current_spend.assert_awaited_once() @pytest.mark.asyncio -async def test_budget_limits_with_usage_empty_windows_keep_stored_value(monkeypatch): - """A key with no windows (None, [], or "[]") returns None so /key/info keeps the stored value; no spend lookup runs.""" +async def test_build_budget_limits_usage_empty_windows_returns_none(monkeypatch): + """A key with no windows (None, [], or "[]") returns None so the field is left off; no spend lookup runs.""" from unittest.mock import AsyncMock from litellm.proxy.management_endpoints.key_management_endpoints import ( - _budget_limits_with_usage, + _build_budget_limits_usage, ) mock_get_current_spend = AsyncMock(return_value=0.0) @@ -14398,17 +14405,17 @@ async def test_budget_limits_with_usage_empty_windows_keep_stored_value(monkeypa ) for stored in (None, [], "[]"): - assert await _budget_limits_with_usage(budget_limits=stored, api_key_hash="hash-1") is None + assert await _build_budget_limits_usage(budget_limits=stored, api_key_hash="hash-1") is None mock_get_current_spend.assert_not_awaited() @pytest.mark.asyncio -async def test_budget_limits_with_usage_window_without_max_budget(monkeypatch): - """A window with only budget_duration still gets current_spend, read without a budget ceiling.""" +async def test_build_budget_limits_usage_window_without_max_budget(monkeypatch): + """A window with only budget_duration still reports current_spend, read without a budget ceiling.""" from unittest.mock import AsyncMock from litellm.proxy.management_endpoints.key_management_endpoints import ( - _budget_limits_with_usage, + _build_budget_limits_usage, ) mock_get_current_spend = AsyncMock(return_value=0.75) @@ -14416,11 +14423,11 @@ async def test_budget_limits_with_usage_window_without_max_budget(monkeypatch): "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend ) - result = await _budget_limits_with_usage( + result = await _build_budget_limits_usage( budget_limits=[{"budget_duration": "2d"}], api_key_hash="hash-no-max" ) - assert list(result) == [{"budget_duration": "2d", "current_spend": 0.75}] + assert result == {"2d": {"current_spend": 0.75, "budget_limit": None, "reset_at": None}} call_kwargs = mock_get_current_spend.await_args.kwargs assert call_kwargs["counter_key"] == "spend:key:hash-no-max:window:2d" assert call_kwargs["window_duration"] == "2d" @@ -14428,13 +14435,13 @@ async def test_budget_limits_with_usage_window_without_max_budget(monkeypatch): @pytest.mark.asyncio -async def test_budget_limits_with_usage_pydantic_windows(monkeypatch): - """BudgetLimitEntry windows (the shape UserAPIKeyAuth carries) are dumped to dicts and annotated.""" +async def test_build_budget_limits_usage_pydantic_windows(monkeypatch): + """BudgetLimitEntry windows (the shape UserAPIKeyAuth carries) are dumped to dicts and reported.""" from unittest.mock import AsyncMock from litellm.models.team import BudgetLimitEntry from litellm.proxy.management_endpoints.key_management_endpoints import ( - _budget_limits_with_usage, + _build_budget_limits_usage, ) mock_get_current_spend = AsyncMock(return_value=1.0) @@ -14442,14 +14449,12 @@ async def test_budget_limits_with_usage_pydantic_windows(monkeypatch): "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend ) - result = await _budget_limits_with_usage( + result = await _build_budget_limits_usage( budget_limits=[BudgetLimitEntry(budget_duration="7d", max_budget=10.0)], api_key_hash="hash-2", ) - assert list(result) == [ - {"budget_duration": "7d", "max_budget": 10.0, "reset_at": None, "current_spend": 1.0} - ] + assert result == {"7d": {"current_spend": 1.0, "budget_limit": 10.0, "reset_at": None}} call_kwargs = mock_get_current_spend.await_args.kwargs assert call_kwargs["counter_key"] == "spend:key:hash-2:window:7d" assert call_kwargs["window_duration"] == "7d" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9b5adec516d..7b6f428c5a4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7710,9 +7710,10 @@ export interface paths { * - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} * - model_max_budget_usage: dict | None - Current-window spend per model, present only when * the key has per-model budgets - * - budget_limits: list | None - Concurrent budget windows. Each entry includes - * current_spend: spend accumulated in the window so far (read from the same cross-pod - * spend counter the budget enforcement uses) + * - budget_limits: list | None - Concurrent budget windows, exactly as stored + * - budget_limits_usage: dict | None - Current-window spend per budget window, keyed by + * budget_duration, present only when the key has budget windows (read from the same + * cross-pod spend counter the budget enforcement uses) * - models: list - Model_name's the key is allowed to call * - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits * - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"}