diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 116a9e6ff42..546078cce09 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -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 diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index 3b7bd054874..fbd6ef51f05 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -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.