scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

fix(proxy): resolve temporary budget grants against the live team default instead of a snapshot

Browse source 
A temporary-only member update no longer clones the team default budget into the private row. The row stores just the temp pair and auth, spend admission and reservation add the active increase to the current shared default, so a later lowering of the default reaches members with an active grant

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-18 02:09:09 +00:00
parent 3926e74e9d
commit 507219e26f
8 changed files with 236 additions and 60 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm/models/budget.py
View file

@ -36,19 +36,20 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
model_config = ConfigDict(protected_namespaces=())
def effective_max_budget(self, now: datetime) -> float | None:
if self.max_budget is None:
return None
def active_temp_budget_increase(self, now: datetime) -> float:
if self.temp_budget_increase is None or self.temp_budget_expiry is None:
return self.max_budget
return 0.0
expiry: Final = (
self.temp_budget_expiry.replace(tzinfo=timezone.utc)
if self.temp_budget_expiry.tzinfo is None
else self.temp_budget_expiry
)
if expiry <= now:
return self.max_budget
return self.max_budget + self.temp_budget_increase
return 0.0 if expiry <= now else self.temp_budget_increase
def effective_max_budget(self, now: datetime) -> float | None:
if self.max_budget is None:
return None
return self.max_budget + self.active_temp_budget_increase(now)
class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable):

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

@ -5325,12 +5325,10 @@ async def _check_team_member_budget(
# Per-member override wins; otherwise fall back to the team-level
# default configured via team.metadata["team_member_budget_id"].
team_member_budget: float | None = None
if (
loaded_membership is not None
and loaded_membership.litellm_budget_table is not None
and loaded_membership.litellm_budget_table.max_budget is not None
):
team_member_budget = loaded_membership.litellm_budget_table.effective_max_budget(now=get_utc_datetime())
member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None
now: Final = get_utc_datetime()
if member_budget_row is not None and member_budget_row.max_budget is not None:
team_member_budget = member_budget_row.effective_max_budget(now=now)
else:
default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id")
if isinstance(default_budget_id, str):
@ -5346,7 +5344,9 @@ async def _check_team_member_budget(
and default_budget.max_budget is not None
and default_budget.max_budget > 0
):
team_member_budget = default_budget.max_budget
team_member_budget = default_budget.max_budget + (
member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0
)
if team_member_budget is not None:
team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0

28
litellm/proxy/management_endpoints/common_utils.py
View file

@ -480,6 +480,8 @@ _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: Final = (
"temp_budget_expiry",
)
_TEMP_BUDGET_FIELDS: Final = frozenset({"temp_budget_increase", "temp_budget_expiry"})
MEMBER_BUDGET_PATCH_FIELDS: Final = MappingProxyType(
{
@ -552,6 +554,8 @@ async def _upsert_budget_and_membership(
``shared_budget_ids`` extends that protection to any other row more than one
membership points at, which a caller patching several members at once has
already counted; a row listed there is cloned rather than written in place.
A patch that only touches the temporary budget pair never copies permanent
limits into a new row, so the member keeps inheriting the live team default.
"""
if not budget_patch:
return
@ -566,6 +570,7 @@ async def _upsert_budget_and_membership(
is_shared_default: Final = existing_budget_id is not None and (
existing_budget_id == team_default_budget_id or existing_budget_id in (shared_budget_ids or frozenset())
)
temp_only: Final = frozenset(write_data) <= _TEMP_BUDGET_FIELDS
async def _disconnect():
await tx.litellm_teammembership.update(
@ -586,30 +591,19 @@ async def _upsert_budget_and_membership(
)
return
seeds_temp_budget: Final = "temp_budget_increase" in write_data or "temp_budget_expiry" in write_data
source_row_id: Final = (
existing_budget_id
if is_shared_default
else team_default_budget_id
if team_default_budget_id is not None and seeds_temp_budget
else None
)
source_row: Final = (
await tx.litellm_budgettable.find_unique(where={"budget_id": source_row_id})
if source_row_id is not None
await tx.litellm_budgettable.find_unique(where={"budget_id": existing_budget_id})
if is_shared_default and not temp_only
else None
)
source: Final[Mapping[str, Any]] = source_row.model_dump() if source_row is not None else MappingProxyType({})
def _seeds(field: str) -> bool:
if field == "max_budget" and source.get(field) == 0 and not is_shared_default:
return False
return _is_set_budget_value(source.get(field))
create_data: Final[dict[str, Any]] = { # mutable-ok: Prisma create payloads are dict-shaped
"created_by": user_api_key_dict.user_id or "",
"updated_by": user_api_key_dict.user_id or "",
**MappingProxyType({f: source[f] for f in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS if _seeds(f)}),
**MappingProxyType(
{f: source[f] for f in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS if _is_set_budget_value(source.get(f))}
),
**write_data,
}
@ -621,7 +615,7 @@ async def _upsert_budget_and_membership(
create_data.pop("budget_reset_at", None)
if not _has_meaningful_budget_limit(create_data):
if existing_budget_id is not None:
if existing_budget_id is not None and not temp_only:
await _disconnect()
return

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

@ -688,16 +688,22 @@ async def _get_team_member_budget_counter(
elif isinstance(cached_team_membership, dict):
team_membership = LiteLLM_TeamMembership(**cached_team_membership)
member_budget_row: Final = team_membership.litellm_budget_table if team_membership is not None else None
now: Final = datetime.now(timezone.utc)
team_member_budget: float | None = None
if team_membership is not None and team_membership.litellm_budget_table is not None:
team_member_budget = team_membership.litellm_budget_table.effective_max_budget(now=datetime.now(timezone.utc))
if member_budget_row is not None and member_budget_row.max_budget is not None:
team_member_budget = member_budget_row.effective_max_budget(now=now)
else:
default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id")
if isinstance(default_budget_id, str):
default_budget: Final = await user_api_key_cache.async_get_cache(
key=f"team_member_default_budget:{default_budget_id}",
)
team_member_budget = _to_float(_get_value(default_budget, "max_budget"))
default_cap: Final = _to_float(_get_value(default_budget, "max_budget"))
if default_cap is not None and default_cap > 0:
team_member_budget = default_cap + (
member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0
)
if team_member_budget is None or team_member_budget <= 0:
return None

9
tests/test_litellm/models/test_models.py
View file

@ -90,6 +90,15 @@ class TestBudget:
assert LiteLLM_BudgetTable(max_budget=100.0).effective_max_budget(now=now) == 100.0
assert LiteLLM_BudgetTable(max_budget=None, temp_budget_increase=50.0).effective_max_budget(now=now) is None
def test_active_temp_budget_increase_is_independent_of_max_budget(self):
now = datetime(2026, 1, 1, tzinfo=timezone.utc)
bare = LiteLLM_BudgetTable(max_budget=None, temp_budget_increase=50.0, temp_budget_expiry=datetime(2100, 1, 1))
assert bare.active_temp_budget_increase(now=now) == 50.0
assert bare.effective_max_budget(now=now) is None
expired = LiteLLM_BudgetTable(max_budget=None, temp_budget_increase=50.0, temp_budget_expiry=now)
assert expired.active_temp_budget_increase(now=now) == 0.0
assert LiteLLM_BudgetTable(max_budget=None).active_temp_budget_increase(now=now) == 0.0
class TestCredentials:
def test_credentials_creation(self):

71
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -8548,3 +8548,74 @@ async def test_team_member_budget_check_temp_budget_increase_extends_cap():
)
assert exc_info.value.current_cost == 150.0
assert exc_info.value.max_budget == 100.0
@pytest.mark.asyncio
@pytest.mark.parametrize(
"default_cap, expiry_offset, spend, expected_cap",
[
(0.4, timedelta(hours=1), 1.0, None),
(0.4, timedelta(hours=-1), 1.0, 0.4),
(0.0, timedelta(hours=1), 1.0, None),
],
)
async def test_team_member_budget_check_adds_temp_increase_to_live_team_default(
default_cap: float, expiry_offset: timedelta, spend: float, expected_cap: float | None
):
"""A member row that carries only the temporary pair inherits the team default
cap live: the increase is added to it while active, the default alone applies
once it expires, and a zero default stays uncapped."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import LiteLLM_TeamMembership
from litellm.proxy.utils import ProxyLogging
cache = DualCache()
await cache.async_set_cache(
key="team_member_default_budget:default-budget-1",
value=LiteLLM_BudgetTable(budget_id="default-budget-1", max_budget=default_cap),
)
team_object = LiteLLM_TeamTable(team_id="test-team", metadata={"team_member_budget_id": "default-budget-1"})
valid_token = UserAPIKeyAuth(token="test-token", user_id="test-user", team_id="test-team")
team_membership = LiteLLM_TeamMembership(
user_id="test-user",
team_id="test-team",
spend=spend,
budget_id="budget-1",
litellm_budget_table=LiteLLM_BudgetTable(
max_budget=None,
temp_budget_increase=1.0,
temp_budget_expiry=datetime.now(timezone.utc) + expiry_offset,
),
)
async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs):
return fallback_spend
with (
patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter
patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch
"litellm.proxy.auth.auth_checks.get_team_membership",
new_callable=AsyncMock,
return_value=team_membership,
),
):
if expected_cap is None:
await _check_team_member_budget(
team_object=team_object,
user_object=LiteLLM_UserTable(user_id="test-user"),
valid_token=valid_token,
prisma_client=MagicMock(),
user_api_key_cache=cache,
proxy_logging_obj=ProxyLogging(user_api_key_cache=None),
)
return
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _check_team_member_budget(
team_object=team_object,
user_object=LiteLLM_UserTable(user_id="test-user"),
valid_token=valid_token,
prisma_client=MagicMock(),
user_api_key_cache=cache,
proxy_logging_obj=ProxyLogging(user_api_key_cache=None),
)
assert exc_info.value.max_budget == expected_cap

98
tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py
View file

@ -233,10 +233,10 @@ async def test_create_from_temp_budget_pair_only(mock_tx, fake_user):
@pytest.mark.asyncio
async def test_create_from_temp_pair_keeps_team_default_cap(mock_tx, fake_user):
async def test_create_from_temp_pair_never_snapshots_team_default(mock_tx, fake_user):
expiry = datetime(2100, 1, 1, tzinfo=timezone.utc)
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, allowed_models=[])
return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, rpm_limit=10)
)
await _upsert_budget_and_membership(
mock_tx,
@ -248,34 +248,80 @@ async def test_create_from_temp_pair_keeps_team_default_cap(mock_tx, fake_user):
team_default_budget_id="team-default-budget-1",
)
mock_tx.litellm_budgettable.find_unique.assert_not_awaited()
data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
assert "max_budget" not in data
assert "rpm_limit" not in data
assert data["temp_budget_increase"] == 1.0
assert data["temp_budget_expiry"] == expiry
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
@pytest.mark.asyncio
async def test_temp_pair_on_shared_default_member_creates_bare_row(mock_tx, fake_user):
expiry = datetime(2100, 1, 1, tzinfo=timezone.utc)
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, rpm_limit=10)
)
await _upsert_budget_and_membership(
mock_tx,
team_id="team-default",
user_id="user-on-default",
existing_budget_id="team-default-budget-1",
user_api_key_dict=fake_user,
budget_patch={"temp_budget_increase": 1.0, "temp_budget_expiry": expiry},
team_default_budget_id="team-default-budget-1",
)
mock_tx.litellm_budgettable.find_unique.assert_not_awaited()
mock_tx.litellm_budgettable.update.assert_not_called()
data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
assert "max_budget" not in data
assert "rpm_limit" not in data
assert data["temp_budget_increase"] == 1.0
assert data["temp_budget_expiry"] == expiry
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
@pytest.mark.asyncio
async def test_clearing_temp_pair_on_shared_default_member_is_noop(mock_tx, fake_user):
await _upsert_budget_and_membership(
mock_tx,
team_id="team-default",
user_id="user-on-default",
existing_budget_id="team-default-budget-1",
user_api_key_dict=fake_user,
budget_patch={"temp_budget_increase": None, "temp_budget_expiry": None},
team_default_budget_id="team-default-budget-1",
)
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_budgettable.update.assert_not_called()
mock_tx.litellm_teammembership.update.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
@pytest.mark.asyncio
async def test_temp_pair_with_permanent_field_still_clones_shared_default(mock_tx, fake_user):
expiry = datetime(2100, 1, 1, tzinfo=timezone.utc)
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, rpm_limit=10)
)
await _upsert_budget_and_membership(
mock_tx,
team_id="team-default",
user_id="user-on-default",
existing_budget_id="team-default-budget-1",
user_api_key_dict=fake_user,
budget_patch={"temp_budget_increase": 1.0, "temp_budget_expiry": expiry, "tpm_limit": 500},
team_default_budget_id="team-default-budget-1",
)
mock_tx.litellm_budgettable.find_unique.assert_awaited_once_with(where={"budget_id": "team-default-budget-1"})
data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
assert data["max_budget"] == 0.4
assert data["temp_budget_increase"] == 1.0
assert data["temp_budget_expiry"] == expiry
assert "allowed_models" not in data
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
@pytest.mark.asyncio
async def test_create_from_temp_pair_skips_zero_team_default_cap(mock_tx, fake_user):
expiry = datetime(2100, 1, 1, tzinfo=timezone.utc)
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(budget_id="team-default-budget-1", max_budget=0, rpm_limit=10)
)
await _upsert_budget_and_membership(
mock_tx,
team_id="team-default",
user_id="user-unlinked",
existing_budget_id=None,
user_api_key_dict=fake_user,
budget_patch={"temp_budget_increase": 1.0, "temp_budget_expiry": expiry},
team_default_budget_id="team-default-budget-1",
)
data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
assert "max_budget" not in data
assert data["rpm_limit"] == 10
assert data["tpm_limit"] == 500
assert data["temp_budget_increase"] == 1.0
mock_tx.litellm_teammembership.upsert.assert_awaited_once()

49
tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
View file

@ -497,3 +497,52 @@ async def test_team_member_reservation_counter_honours_temp_budget_increase(
assert counter is not None
assert counter.max_budget == expected_max_budget
assert counter.fallback_spend == 0.5
@pytest.mark.asyncio
@pytest.mark.parametrize(
"default_cap, expiry_offset, expected_max_budget",
[
(2.0, timedelta(days=1), 3.0),
(2.0, timedelta(days=-1), 2.0),
(0.0, timedelta(days=1), None),
],
)
async def test_team_member_reservation_counter_adds_temp_increase_to_live_team_default(
default_cap: float, expiry_offset: timedelta, expected_max_budget: float | None
) -> None:
user_id: Final = "member-bare"
team_id: Final = "team-bare"
cache: Final = UserApiKeyCache()
await cache.async_set_cache(
key="team_member_default_budget:default-bare",
value=LiteLLM_BudgetTable(budget_id="default-bare", max_budget=default_cap),
)
await cache.async_set_cache(
key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id),
value=LiteLLM_TeamMembership(
user_id=user_id,
team_id=team_id,
spend=0.5,
budget_id="budget-bare",
litellm_budget_table=LiteLLM_BudgetTable(
max_budget=None,
temp_budget_increase=1.0,
temp_budget_expiry=datetime.now(timezone.utc) + expiry_offset,
),
),
)
counter: Final = await _get_team_member_budget_counter(
valid_token=UserAPIKeyAuth(token="hashed", user_id=user_id, team_id=team_id),
team_object=LiteLLM_TeamTable(team_id=team_id, metadata={"team_member_budget_id": "default-bare"}),
user_object=LiteLLM_UserTable(user_id=user_id),
user_api_key_cache=cache,
)
if expected_max_budget is None:
assert counter is None
return
assert counter is not None
assert counter.max_budget == expected_max_budget
assert counter.fallback_spend == 0.5