fix(proxy): do not carry a zero team default cap onto a new member budget row

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-17 21:00:31 +00:00
parent e73b8d49da
commit 664688f337
2 changed files with 25 additions and 6 deletions

View file

@ -578,6 +578,8 @@ async def _upsert_budget_and_membership(
default_budget_dict: Final = default_budget_row.model_dump()
for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS:
value = default_budget_dict.get(field)
if field == "max_budget" and value == 0 and not is_shared_default:
continue
if _is_set_budget_value(value):
create_data[field] = value

View file

@ -219,9 +219,6 @@ async def test_create_seeds_reset_at_and_links(mock_tx, fake_user):
)
# TEST: with no existing budget, a patch carrying only the temporary increase
# pair must still create a budget row and link it, so the fields the 200
# response echoes are actually stored.
@pytest.mark.asyncio
async def test_create_from_temp_budget_pair_only(mock_tx, fake_user):
expiry = datetime(2100, 1, 1, tzinfo=timezone.utc)
@ -243,9 +240,6 @@ async def test_create_from_temp_budget_pair_only(mock_tx, fake_user):
mock_tx.litellm_teammembership.update.assert_not_called()
# TEST: a member with no budget row who falls back to the team default at
# enforcement time must keep that default cap on the new private row, or the
# temporary increase has nothing to add to.
@pytest.mark.asyncio
async def test_create_from_temp_pair_keeps_team_default_cap(mock_tx, fake_user):
expiry = datetime(2100, 1, 1, tzinfo=timezone.utc)
@ -271,6 +265,29 @@ async def test_create_from_temp_pair_keeps_team_default_cap(mock_tx, fake_user):
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["temp_budget_increase"] == 1.0
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
# TEST: clone-on-write when the membership still points at the team's shared
# default budget. Editing this member must fork a private budget instead of
# mutating the shared row, and cloning a duration must seed a fresh reset time.