scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix(proxy): align user budget window empty-duration guard with reservation path

Browse source 
Treat an empty-string budget_duration the same as a missing one in
_coerce_budget_limit_window_for_check, matching _get_budget_limit_counters
('if not budget_duration'). Previously a window stored with budget_duration=''
passed the auth check and produced a 'spend:...🪟' counter key that never
accumulated spend, silently disabling the window instead of enforcing it.

Also drop the dead dict(window) branch in _initialize_budget_limits_for_update
in favor of model_dump() for the BudgetLimitEntry path.

Adds regression tests covering the empty-string duration window.
This commit is contained in:
lengkejun 2026-06-02 20:12:39 +08:00 • committed by silencedoctor
parent 96651ac1b6
commit 782f29fa45
10 changed files with 282 additions and 67 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/models/user.py
View file

@ -5,9 +5,10 @@ Canonical definition for ``litellm_usertable``. Re-exported from
``litellm.proxy._types`` for backwards compatibility.
"""
import json
from datetime import datetime
from pydantic import ConfigDict, Field, model_validator
from pydantic import BaseModel, ConfigDict, Field, model_validator
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
from litellm.models.organization_membership import (
@ -52,6 +53,14 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase):
@model_validator(mode="before")
@classmethod
def set_model_info(cls, values):
if isinstance(values, BaseModel):
values = values.model_dump()
budget_limits = values.get("budget_limits")
if isinstance(budget_limits, str):
try:
values["budget_limits"] = json.loads(budget_limits)
except json.JSONDecodeError as exc:
raise ValueError("Field budget_limits should be valid JSON") from exc
if values.get("spend") is None:
values.update({"spend": 0.0})
if values.get("models") is None:

81
litellm/proxy/auth/auth_checks.py
View file

@ -14,11 +14,12 @@ import math
import re
import time
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast
from fastapi import HTTPException, Request, status
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
@ -4323,24 +4324,24 @@ async def _virtual_key_max_budget_check(
)
def _coerce_budget_limit_window_for_check(window: Any) -> Optional[Dict[str, Any]]:
if isinstance(window, dict):
w = window
elif hasattr(window, "model_dump"):
w = window.model_dump()
else:
return None
budget_duration = w.get("budget_duration")
max_budget = w.get("max_budget")
if budget_duration is None or max_budget is None:
return None
class _BudgetLimitWindowForCheck(BaseModel):
budget_duration: str
max_budget: float
reset_at: datetime | str | None = None
_BUDGET_LIMIT_WINDOW_ADAPTER = TypeAdapter(_BudgetLimitWindowForCheck)
def _coerce_budget_limit_window_for_check(window: object) -> _BudgetLimitWindowForCheck | None:
raw_window = window.model_dump() if isinstance(window, BaseModel) else window
try:
max_budget_float = float(max_budget)
except (TypeError, ValueError):
budget_limit = _BudgetLimitWindowForCheck.model_validate(raw_window)
except ValidationError:
return None
if not math.isfinite(max_budget_float):
if not budget_limit.budget_duration or not math.isfinite(budget_limit.max_budget):
return None
return {"budget_duration": budget_duration, "max_budget": max_budget_float}
return budget_limit
async def _virtual_key_multi_budget_check(
@ -4359,26 +4360,30 @@ async def _virtual_key_multi_budget_check(
from litellm.proxy.proxy_server import get_current_spend
for window in valid_token.budget_limits:
w = _coerce_budget_limit_window_for_check(window=window)
if w is None:
for raw_window in valid_token.budget_limits:
try:
window = _BUDGET_LIMIT_WINDOW_ADAPTER.validate_python(raw_window)
except ValidationError:
continue
counter_key = f"spend:key:{valid_token.token}:window:{w['budget_duration']}"
if not window.budget_duration or not math.isfinite(window.max_budget):
continue
w = window
counter_key = f"spend:key:{valid_token.token}:window:{w.budget_duration}"
window_spend = await get_current_spend(
counter_key=counter_key,
fallback_spend=0.0,
max_budget=w["max_budget"],
max_budget=w.max_budget,
window_entity_type="Key",
window_entity_id=valid_token.token,
window_start=get_budget_window_start(w),
)
if window_spend >= w["max_budget"]:
if window_spend >= w.max_budget:
raise litellm.BudgetExceededError(
current_cost=window_spend,
max_budget=w["max_budget"],
max_budget=w.max_budget,
message=(
f"ExceededBudget: Key over {w['budget_duration']} budget. "
f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}"
f"ExceededBudget: Key over {w.budget_duration} budget. "
f"Spend=${window_spend:.4f}, Limit=${w.max_budget:.2f}"
),
entity_type=Litellm_EntityType.KEY.value,
entity_id=valid_token.token,
@ -4386,7 +4391,7 @@ async def _virtual_key_multi_budget_check(
async def _user_multi_budget_check(
user_object: Optional[LiteLLM_UserTable],
user_object: LiteLLM_UserTable | None,
):
"""
Raises BudgetExceededError if any budget window in user_object.budget_limits is exceeded.
@ -4405,22 +4410,22 @@ async def _user_multi_budget_check(
w = _coerce_budget_limit_window_for_check(window=window)
if w is None:
continue
counter_key = f"spend:user:{user_object.user_id}:window:{w['budget_duration']}"
counter_key = f"spend:user:{user_object.user_id}:window:{w.budget_duration}"
window_spend = await get_current_spend(
counter_key=counter_key,
fallback_spend=0.0,
max_budget=w["max_budget"],
max_budget=w.max_budget,
window_entity_type="User",
window_entity_id=user_object.user_id,
window_start=get_budget_window_start(w),
)
if window_spend >= w["max_budget"]:
if window_spend >= w.max_budget:
raise litellm.BudgetExceededError(
current_cost=window_spend,
max_budget=w["max_budget"],
max_budget=w.max_budget,
message=(
f"ExceededBudget: User={user_object.user_id} over {w['budget_duration']} budget. "
f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}"
f"ExceededBudget: User={user_object.user_id} over {w.budget_duration} budget. "
f"Spend=${window_spend:.4f}, Limit=${w.max_budget:.2f}"
),
entity_type=Litellm_EntityType.USER.value,
entity_id=user_object.user_id,
@ -4780,22 +4785,22 @@ async def _team_multi_budget_check(
w = _coerce_budget_limit_window_for_check(window=window)
if w is None:
continue
counter_key = f"spend:team:{team_object.team_id}:window:{w['budget_duration']}"
counter_key = f"spend:team:{team_object.team_id}:window:{w.budget_duration}"
window_spend = await get_current_spend(
counter_key=counter_key,
fallback_spend=0.0,
max_budget=w["max_budget"],
max_budget=w.max_budget,
window_entity_type="Team",
window_entity_id=team_object.team_id,
window_start=get_budget_window_start(w),
)
if window_spend >= w["max_budget"]:
if window_spend >= w.max_budget:
raise litellm.BudgetExceededError(
current_cost=window_spend,
max_budget=w["max_budget"],
max_budget=w.max_budget,
message=(
f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. "
f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}"
f"ExceededBudget: Team={team_object.team_id} over {w.budget_duration} budget. "
f"Spend=${window_spend:.4f}, Limit=${w.max_budget:.2f}"
),
entity_type=Litellm_EntityType.TEAM.value,
entity_id=team_object.team_id,

1
litellm/proxy/auth/user_api_key_auth.py
View file

@ -2696,6 +2696,7 @@ async def _reserve_budget_after_common_checks(
end_user_id=end_user_id,
end_user_object=end_user_object,
apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True,
skip_user_budget_on_team_key=general_settings.get("skip_user_budget_on_team_key") is True,
fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True,
)

11
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -44,11 +44,11 @@ from litellm.repositories.table_repositories import (
TeamMembershipRepository,
)
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.unit_of_work import (
budget_cascade_unit_of_work,
spend_reset_unit_of_work,
)
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import (
VerificationTokenRepository,
)
@ -1193,7 +1193,10 @@ class ResetBudgetJob:
lambda: self.prisma_client.db.query_raw(source.page_query(), cursor, RESET_BUDGET_JOB_BATCH_SIZE),
reason=f"reset_budget_read_{source.retry_subject}_windows_failure",
)
for row in rows:
source_rows: Final = tuple(row for row in rows if source.id_column in row)
if not source_rows:
return None
for row in source_rows:
raw = row["budget_limits"]
if not raw:
continue
@ -1216,9 +1219,9 @@ class ResetBudgetJob:
reason=f"reset_budget_write_{source.retry_subject}_windows_failure",
)
if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE:
if len(source_rows) < RESET_BUDGET_JOB_BATCH_SIZE:
return None
return rows[-1][source.id_column]
return source_rows[-1][source.id_column]
@staticmethod
async def _reset_budget_common(

48
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -1218,24 +1218,32 @@ def _initialize_budget_limits_for_update(
if not new_windows:
return "[]"
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
preserved = _existing_reset_at_by_duration(existing_user_row)
initialized = tuple(
{
**(window if isinstance(window, dict) else window.model_dump()),
"reset_at": preserved.get(
window["budget_duration"] if isinstance(window, dict) else window.budget_duration
)
or get_budget_reset_time(
budget_duration=window["budget_duration"] if isinstance(window, dict) else window.budget_duration
).isoformat(),
}
_initialize_budget_limit_window(window=window, preserved=preserved)
for window in new_windows
)
return json.dumps(initialized)
def _initialize_budget_limit_window(
window: object,
preserved: Mapping[str, str],
) -> dict[str, object]:
window_data = window if isinstance(window, dict) else window.model_dump()
duration = window_data.get("budget_duration")
if not isinstance(duration, str):
raise ValueError("budget_duration must be a string")
validate_budget_duration(duration)
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
return {
**window_data,
"reset_at": preserved.get(duration)
or get_budget_reset_time(budget_duration=duration).isoformat(),
}
def _update_internal_user_params(
data_json: dict,
data: UpdateUserRequest | UpdateUserRequestNoUserIDorEmail,
@ -1372,6 +1380,17 @@ async def _invalidate_user_spend_counter_if_changed(
await _invalidate_spend_counter(counter_key=f"spend:user:{non_default_values['user_id']}")
async def _invalidate_user_cache(user_id: str | None) -> None:
if user_id is None:
return
from litellm.proxy.proxy_server import user_api_key_cache
try:
await user_api_key_cache.async_delete_cache(key=user_id)
except Exception as e: # noqa: BLE001
verbose_proxy_logger.warning("Failed to invalidate user cache for %s: %s", user_id, e)
def _clears_object_permission(user_request: UpdateUserRequest) -> bool:
"""Whether the caller explicitly asked to remove this user's object_permission.
@ -1545,6 +1564,7 @@ async def _update_single_user_helper(
response = await prisma_client.insert_data(data=non_default_values, table_name="user")
if response is not None:
await _invalidate_user_cache(non_default_values.get("user_id"))
await _schedule_user_update_audit_log(
response=response,
existing_user_row=existing_user_row,
@ -1809,7 +1829,7 @@ async def bulk_user_update(
}'
```
"""
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client, user_api_key_cache
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
if prisma_client is None:
raise HTTPException(
@ -1885,9 +1905,7 @@ async def bulk_user_update(
data=non_default_values, # Update all users
)
for user in all_users_in_db:
await user_api_key_cache.async_delete_cache(key=user.user_id)
await user_api_key_cache.async_delete_cache(key=f"{user.user_id}_user_api_key_user_id")
await asyncio.gather(*(_invalidate_user_cache(user.user_id) for user in all_users_in_db))
# Create individual success results
for user in all_users_in_db:

10
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -159,6 +159,7 @@ async def reserve_budget_for_request(
end_user_id: str | None = None,
end_user_object: object = None,
apply_user_budget_to_team_keys: bool = False,
skip_user_budget_on_team_key: bool | None = None,
fail_closed_budget_enforcement: bool = False,
) -> dict | None:
if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
@ -179,6 +180,7 @@ async def reserve_budget_for_request(
end_user_id=end_user_id,
end_user_object=end_user_object,
apply_user_budget_to_team_keys=apply_user_budget_to_team_keys,
skip_user_budget_on_team_key=skip_user_budget_on_team_key,
)
if not counters:
return None
@ -349,6 +351,7 @@ async def _get_budget_counters(
end_user_id: str | None = None,
end_user_object: object = None,
apply_user_budget_to_team_keys: bool = False,
skip_user_budget_on_team_key: bool | None = None,
) -> list[_BudgetCounter]:
counters: Final[list[_BudgetCounter]] = []
@ -415,8 +418,13 @@ async def _get_budget_counters(
entity_id=user_object.user_id,
)
)
should_skip_user_windows = (
skip_user_budget_on_team_key
if skip_user_budget_on_team_key is not None
else False
)
if (
not (is_team_key and skip_user_budget_on_team_key)
not (is_team_key and should_skip_user_windows)
and user_object is not None
and user_object.user_id is not None
):

84
tests/test_litellm/proxy/auth/test_multi_budget_windows.py
View file

@ -2,6 +2,7 @@
Unit tests for multi-budget-window enforcement on API keys.
"""
from datetime import datetime, timezone
from unittest.mock import ANY, AsyncMock, MagicMock, patch
import pytest
@ -10,6 +11,8 @@ from fastapi import Request
import litellm
from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
_coerce_budget_limit_window_for_check,
_team_multi_budget_check,
_user_multi_budget_check,
_virtual_key_multi_budget_check,
common_checks,
@ -73,9 +76,7 @@ async def test_over_first_window_raises():
call_count += 1
return val
with patch(
"litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend
):
with patch("litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend):
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _virtual_key_multi_budget_check(valid_token=token)
@ -105,9 +106,7 @@ async def test_over_second_window_raises():
call_count += 1
return val
with patch(
"litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend
):
with patch("litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend):
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _virtual_key_multi_budget_check(valid_token=token)
@ -193,6 +192,79 @@ async def test_user_budget_limits_over_window_raises():
assert "1d" in str(err)
@pytest.mark.asyncio
async def test_team_budget_limits_use_typed_window():
team = LiteLLM_TeamTable(
team_id="team-budget-window",
budget_limits=[
{"budget_duration": "1d", "max_budget": 10.0, "reset_at": None},
],
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new_callable=AsyncMock,
return_value=1.0,
) as mock_get_spend:
await _team_multi_budget_check(team_object=team)
mock_get_spend.assert_awaited_once_with(
counter_key="spend:team:team-budget-window:window:1d",
fallback_spend=0.0,
max_budget=10.0,
window_entity_type="Team",
window_entity_id="team-budget-window",
window_start=ANY,
)
@pytest.mark.parametrize("empty_duration", ["", None])
def test_coerce_skips_window_with_empty_budget_duration(empty_duration):
"""
Regression: the auth check and the budget-reservation path must agree on
what counts as a missing budget_duration. The reservation path rejects with
`if not budget_duration`, so an empty-string duration must be rejected here
too — otherwise the auth check builds a counter key ending in `:window:` that
never accumulates spend, silently disabling the window instead of enforcing it.
"""
assert _coerce_budget_limit_window_for_check(window={"budget_duration": empty_duration, "max_budget": 10.0}) is None
def test_coerce_preserves_reset_at():
reset_at = datetime(2026, 7, 30, tzinfo=timezone.utc)
window = _coerce_budget_limit_window_for_check(
window={"budget_duration": "1d", "max_budget": 10.0, "reset_at": reset_at}
)
assert window is not None
assert window.reset_at == reset_at
def test_coerce_invalid_reset_at_does_not_disable_window():
window = _coerce_budget_limit_window_for_check(
window={"budget_duration": "1d", "max_budget": 10.0, "reset_at": "invalid"}
)
assert window is not None
assert window.reset_at == "invalid"
@pytest.mark.asyncio
async def test_empty_budget_duration_window_is_skipped_not_queried():
"""A window with an empty-string duration must be skipped entirely, never
turned into a malformed `spend:key:...:window:` counter lookup."""
token = _make_valid_token(budget_limits=[{"budget_duration": "", "max_budget": 1.0, "reset_at": None}])
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new_callable=AsyncMock,
) as mock_get_spend:
await _virtual_key_multi_budget_check(valid_token=token)
mock_get_spend.assert_not_awaited()
def _make_team(team_id: str = "team-xyz") -> LiteLLM_TeamTable:
return LiteLLM_TeamTable(
team_id=team_id,

12
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -2132,9 +2132,15 @@ def _paginating_window_job(monkeypatch, pages_by_table: Dict[str, List[List[Dict
calls: List[Dict[str, Any]] = []
async def fake_query_raw(query: str, *args, **kwargs):
table = "key" if '"LiteLLM_VerificationToken"' in query else "team"
table = (
"key"
if '"LiteLLM_VerificationToken"' in query
else "user"
if '"LiteLLM_UserTable"' in query
else "team"
)
calls.append({"table": table, "sql": query, "cursor": args[0], "limit": args[1]})
pages = remaining[table]
pages = remaining.get(table, [])
return pages.pop(0) if pages else []
prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw)
@ -2273,6 +2279,8 @@ def _cursor_paginating_window_job(monkeypatch, key_rows: List[Dict[str, Any]]):
async def fake_query_raw(query: str, *args, **kwargs):
if '"LiteLLM_TeamTable"' in query:
return []
if '"LiteLLM_UserTable"' in query:
return []
cursor, limit = args[0], args[1]
page = [row for row in ordered if row["token"] > cursor][:limit]
visited.extend(row["token"] for row in page)

87
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -3545,6 +3545,93 @@ async def test_ghsa_wvg4_proxy_admin_can_update_user_budget(mocker):
assert result is not None
@pytest.mark.asyncio
async def test_bulk_update_processed_users_invalidates_each_user_cache(mocker):
from litellm.proxy.management_endpoints.internal_user_endpoints import (
bulk_update_processed_users,
)
mock_prisma_client = mocker.MagicMock()
existing_user = mocker.MagicMock()
existing_user.model_dump.return_value = {
"user_id": "user-1",
"max_budget": 100.0,
}
existing_user.user_id = "user-1"
mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(
return_value=existing_user
)
mock_prisma_client.update_data = mocker.AsyncMock(
side_effect=(
{"user_id": "user-1", "max_budget": 10.0},
{"user_id": "user-2", "max_budget": 20.0},
)
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
mock_user_cache = mocker.MagicMock()
mock_user_cache.async_delete_cache = mocker.AsyncMock()
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mock_user_cache)
response = await bulk_update_processed_users(
users_to_update=[
UpdateUserRequest(user_id="user-1", max_budget=10.0),
UpdateUserRequest(user_id="user-2", max_budget=20.0),
],
user_api_key_dict=UserAPIKeyAuth(
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
),
)
assert response.successful_updates == 2
assert mock_user_cache.async_delete_cache.await_args_list == [
mocker.call(key="user-1"),
mocker.call(key="user-2"),
]
@pytest.mark.asyncio
async def test_bulk_update_all_users_invalidates_each_user_cache(mocker):
from litellm.proxy._types import UpdateUserRequestNoUserIDorEmail
from litellm.proxy.management_endpoints.internal_user_endpoints import (
bulk_user_update,
)
from litellm.types.proxy.management_endpoints.internal_user_endpoints import (
BulkUpdateUserRequest,
)
mock_prisma_client = mocker.MagicMock()
users = [
SimpleNamespace(user_id="user-1", user_email="one@example.com"),
SimpleNamespace(user_id="user-2", user_email="two@example.com"),
]
mock_prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(
return_value=users
)
mock_prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock()
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
mock_user_cache = mocker.MagicMock()
mock_user_cache.async_delete_cache = mocker.AsyncMock()
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mock_user_cache)
response = await bulk_user_update.__wrapped__(
data=BulkUpdateUserRequest(
all_users=True,
user_updates=UpdateUserRequestNoUserIDorEmail(max_budget=10.0),
),
user_api_key_dict=UserAPIKeyAuth(
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
),
)
assert response.successful_updates == 2
mock_user_cache.async_delete_cache.assert_has_awaits(
[mocker.call(key="user-1"), mocker.call(key="user-2")],
any_order=True,
)
@pytest.mark.asyncio
async def test_admin_user_update_spend_invalidates_counter(mocker):
"""A direct /user/update spend change must invalidate the cross-pod

4
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -28079,6 +28079,8 @@ export interface components {
allowed_cache_controls: string[];
/** Budget Duration */
budget_duration?: string | null;
/** Budget Limits */
budget_limits?: components["schemas"]["BudgetLimitEntry"][] | null;
/** Budget Reset At */
budget_reset_at?: string | null;
/** Created At */
@ -28167,6 +28169,8 @@ export interface components {
allowed_cache_controls: string[];
/** Budget Duration */
budget_duration?: string | null;
/** Budget Limits */
budget_limits?: components["schemas"]["BudgetLimitEntry"][] | null;
/** Budget Reset At */
budget_reset_at?: string | null;
/** Created At */