diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 26b65dd48cf..99e201930c1 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 @@ -111,6 +112,7 @@ from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper +from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( get_ui_settings_cached, @@ -3578,6 +3580,63 @@ async def _build_model_max_budget_usage( ) +def _window_max_budget(window: Mapping[str, object]) -> float | None: + """A window's max_budget as a float; None when absent or unparseable.""" + value: Final = window.get("max_budget") + if not isinstance(value, (int, float, str)): + return None + try: + return float(value) + except ValueError: + return None + + +async def _budget_window_usage( + window: Mapping[str, object], api_key_hash: str +) -> tuple[str, Mapping[str, object]] | None: + """ + (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 + window_duration + window_start so a stale-low counter is re-checked against + the LiteLLM_BudgetWindowSpend row instead of a spend-log aggregate. + """ + from litellm.proxy.proxy_server import get_current_spend + + duration: Final = window.get("budget_duration") + if not isinstance(duration, str) or not duration: + return None + 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), + window_entity_type="Key", + window_entity_id=api_key_hash, + window_duration=duration, + window_start=get_budget_window_start(window), + ) + return duration, MappingProxyType({"current_spend": round(spend, 4)}) + + +async def _build_budget_limits_usage( + budget_limits: Sequence[object] | str | None, api_key_hash: str +) -> Mapping[str, Mapping[str, object]] | None: + """ + 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 + 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( "/v2/key/info", tags=["key management"], @@ -3620,7 +3679,6 @@ async def info_key_fn_v2( status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail={"message": "Malformed request. No keys passed in."}, ) - # Resolve key_aliases to tokens so we never pass token=None (unbounded query) tokens_to_query: Final = list(data.keys) if data.keys else [] if data.key_aliases: @@ -3662,6 +3720,13 @@ async def info_key_fn_v2( model_max_budget=model_max_budget, user_api_key_cache=model_max_budget_limiter.dual_cache, ) + if k_token_hash: + 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_usage"] = budget_limits_usage filtered_key_info.append(k_dict) return {"key": data.keys, "info": filtered_key_info} @@ -3698,6 +3763,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, exactly as stored + - budget_limits_usage: dict | None - Current-window spend per budget window, e.g. + {"1h": {"current_spend": 0.0009}}, 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"} @@ -3777,6 +3846,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 _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_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 1c80aa5683f..b5a4204f0af 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 @@ -14142,6 +14142,311 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch): mock_prisma_client.db.query_raw.assert_not_awaited() +@pytest.mark.asyncio +async def test_info_key_fn_reports_budget_limits_usage(monkeypatch): + """ + /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 + + from litellm.proxy._types import LiteLLM_VerificationToken + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + test_key_token = "hashed_token_window_test" + budget_limits = [ + { + "reset_at": "2026-08-15T18:00:00+00:00", + "max_budget": 2.0, + "budget_duration": "1h", + } + ] + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_user_api_key_cache = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + mock_get_current_spend = AsyncMock(return_value=0.73) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) + mock_key_info.token = test_key_token + mock_key_info.object_permission_id = None + mock_key_info.user_id = "user-w" + mock_key_info.team_id = None + mock_key_info.litellm_budget_table = None + mock_key_info.model_dump.return_value = { + "token": test_key_token, + "budget_limits": [dict(w) for w in budget_limits], + "user_id": "user-w", + "team_id": None, + "object_permission_id": None, + "litellm_budget_table": None, + } + mock_key_info.dict.return_value = mock_key_info.model_dump.return_value + + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=mock_key_info + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-test-window-key", + ) + + result = await info_key_fn( + key="sk-test-window-key", + user_api_key_dict=user_api_key_dict, + ) + + assert result["info"]["budget_limits"] == budget_limits + assert result["info"]["budget_limits_usage"] == {"1h": {"current_spend": 0.73}} + + mock_get_current_spend.assert_awaited_once() + call_kwargs = mock_get_current_spend.await_args.kwargs + assert call_kwargs["counter_key"] == f"spend:key:{test_key_token}:window:1h" + assert call_kwargs["max_budget"] == 2.0 + assert call_kwargs["window_entity_type"] == "Key" + assert call_kwargs["window_entity_id"] == test_key_token + assert call_kwargs["window_duration"] == "1h" + assert call_kwargs["window_start"] is not None + + +@pytest.mark.asyncio +async def test_info_key_fn_no_budget_limits_skips_spend_lookup(monkeypatch): + """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 + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + test_key_token = "hashed_token_no_windows" + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_user_api_key_cache = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + mock_get_current_spend = AsyncMock(return_value=0.0) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) + mock_key_info.token = test_key_token + mock_key_info.object_permission_id = None + mock_key_info.user_id = "user-nw" + mock_key_info.team_id = None + mock_key_info.litellm_budget_table = None + mock_key_info.model_dump.return_value = { + "token": test_key_token, + "budget_limits": None, + "user_id": "user-nw", + "team_id": None, + "object_permission_id": None, + "litellm_budget_table": None, + } + mock_key_info.dict.return_value = mock_key_info.model_dump.return_value + + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=mock_key_info + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-test-no-window-key", + ) + + result = await info_key_fn( + key="sk-test-no-window-key", + user_api_key_dict=user_api_key_dict, + ) + + 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_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 + from litellm.proxy.management_endpoints.key_management_endpoints import ( + info_key_fn_v2, + ) + + 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) + mock_user_api_key_cache = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + mock_get_current_spend = AsyncMock(return_value=1.25) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.token = test_key_token + mock_key.user_id = "user-v2-w" + mock_key.team_id = None + mock_key.model_dump.return_value = { + "token": test_key_token, + "budget_limits": [dict(w) for w in budget_limits], + "user_id": "user-v2-w", + "team_id": None, + "litellm_budget_table": None, + } + mock_key.dict.return_value = mock_key.model_dump.return_value + + mock_prisma_client.get_data = AsyncMock(return_value=[mock_key]) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin-v2-w", + ) + + result = await info_key_fn_v2( + data=KeyRequest(keys=[test_key_token]), + user_api_key_dict=user_api_key_dict, + ) + + assert len(result["info"]) == 1 + assert result["info"][0]["budget_limits"] == budget_limits + assert result["info"][0]["budget_limits_usage"] == { + "1h": {"current_spend": 1.25}, + "1d": {"current_spend": 1.25}, + } + assert mock_get_current_spend.await_count == 2 + counter_keys = { + call.kwargs["counter_key"] for call in mock_get_current_spend.await_args_list + } + assert counter_keys == { + f"spend:key:{test_key_token}:window:1h", + f"spend:key:{test_key_token}:window:1d", + } + assert { + call.kwargs["window_duration"] for call in mock_get_current_spend.await_args_list + } == {"1h", "1d"} + + +@pytest.mark.asyncio +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 ( + _build_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 + ) + + raw = json_module.dumps( + [{"budget_duration": "1h", "max_budget": 2.0, "reset_at": None}] + ) + result = await _build_budget_limits_usage(budget_limits=raw, api_key_hash="hash-1") + + assert result == {"1h": {"current_spend": 0.5}} + mock_get_current_spend.assert_awaited_once() + + +@pytest.mark.asyncio +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 ( + _build_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 + ) + + for stored in (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_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 ( + _build_budget_limits_usage, + ) + + mock_get_current_spend = AsyncMock(return_value=0.75) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + result = await _build_budget_limits_usage( + budget_limits=[{"budget_duration": "2d"}], api_key_hash="hash-no-max" + ) + + assert result == {"2d": {"current_spend": 0.75}} + 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" + assert call_kwargs["max_budget"] is None + + +@pytest.mark.asyncio +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 ( + _build_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 + ) + + result = await _build_budget_limits_usage( + budget_limits=[BudgetLimitEntry(budget_duration="7d", max_budget=10.0)], + api_key_hash="hash-2", + ) + + assert result == {"7d": {"current_spend": 1.0}} + 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" + assert call_kwargs["max_budget"] == 10.0 + + @pytest.mark.asyncio async def test_info_key_fn_reads_the_configured_budget_model_key(monkeypatch): """/key/info reads the one counter enforcement reads: the configured budget model. diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index fb81191cfbf..4ad3ca06501 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7710,6 +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, exactly as stored + * - budget_limits_usage: dict | None - Current-window spend per budget window, e.g. + * {"1h": {"current_spend": 0.0009}}, 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"}