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

fix(proxy): do not reset member budget window on implicit duration

Browse source 
In-place updates only write budget_duration / budget_reset_at when
budget_duration was explicit in /team/member_update. Passing the team's
inherited default duration with explicit=False was recalculating
budget_reset_at and wiping mid-cycle spend (Greptile P1, PR #26779).

Add regression test test_upsert_in_place_omitted_duration_does_not_touch_duration_fields.

Made-with: Cursor
This commit is contained in:
Milan 2026-05-01 01:54:37 +03:00
parent 46995a878b
commit e7950b348d
No known key found for this signature in database
2 changed files with 48 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -388,25 +388,28 @@ def _non_none_budget_limit_fields(
return out
def _apply_budget_duration_to_update_dict(
def _apply_explicit_budget_duration_to_update_dict(
update_data: Dict[str, Any],
budget_duration: Optional[str],
budget_duration_explicit: bool,
) -> None:
if budget_duration_explicit:
if budget_duration is not None:
update_data["budget_duration"] = budget_duration
update_data["budget_reset_at"] = get_budget_reset_time(
budget_duration=budget_duration
)
else:
update_data["budget_duration"] = None
update_data["budget_reset_at"] = None
elif budget_duration is not None:
"""
Persist ``budget_duration`` / ``budget_reset_at`` only when the caller
explicitly included ``budget_duration`` in the request (see
``budget_duration_explicit`` at the /team/member_update layer).
When duration is omitted, inherited team defaults are still passed into
``_upsert_budget_and_membership`` for *create* / clone paths, but in-place
updates must not rewrite ``budget_reset_at`` or the member's cycle is reset
mid-period (Greptile / PR #26779).
"""
if budget_duration is not None:
update_data["budget_duration"] = budget_duration
update_data["budget_reset_at"] = get_budget_reset_time(
budget_duration=budget_duration
)
else:
update_data["budget_duration"] = None
update_data["budget_reset_at"] = None
async def _update_existing_member_budget_in_place(
@ -427,9 +430,8 @@ async def _update_existing_member_budget_in_place(
max_budget, tpm_limit, rpm_limit, allowed_models
),
}
_apply_budget_duration_to_update_dict(
update_data, budget_duration, budget_duration_explicit
)
if budget_duration_explicit:
_apply_explicit_budget_duration_to_update_dict(update_data, budget_duration)
await tx.litellm_budgettable.update(
where={"budget_id": existing_budget_id},
data=update_data,

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

@ -95,6 +95,37 @@ async def test_upsert_with_existing_budget_id_creates_new(mock_tx, fake_user):
mock_tx.litellm_teammembership.update.assert_not_called()
@pytest.mark.asyncio
async def test_upsert_in_place_omitted_duration_does_not_touch_duration_fields(
mock_tx,
fake_user,
):
"""
When budget_duration was omitted from the API (explicit=False), an in-place
update must not write budget_duration / budget_reset_at even if the caller
passes a resolved team-default duration string (Greptile P1, PR #26779).
"""
await _upsert_budget_and_membership(
mock_tx,
team_id="team-no-dur-touch",
user_id="user-no-dur-touch",
max_budget=10.0,
existing_budget_id="priv-budget-1",
user_api_key_dict=fake_user,
team_default_budget_id="team-default-xyz",
budget_duration="30d",
budget_duration_explicit=False,
)
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": "priv-budget-1"},
data={
"max_budget": 10.0,
"updated_by": fake_user.user_id,
},
)
# TEST: create new budget and link membership
@pytest.mark.asyncio
async def test_upsert_create_and_link(mock_tx, fake_user):