diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index fe7b85258d0..05de65f6601 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -358,6 +358,7 @@ async def _upsert_budget_and_membership( allowed_models: Optional[List[str]] = None, team_default_budget_id: Optional[str] = None, budget_duration: Optional[str] = None, + budget_duration_explicit: bool = False, ): """ Helper function to Create/Update or Delete the budget within the team membership @@ -377,8 +378,15 @@ async def _upsert_budget_and_membership( member's budget does not mutate the shared default (and therefore every other member who still points at it). budget_duration: Budget reset period for the team member (e.g. '30d', '1mo') + budget_duration_explicit: True when the HTTP caller included ``budget_duration`` + in the request body (so ``None`` means explicit lifetime). When False, + ``None`` means duration was not supplied for this write path. + + If max_budget, tpm_limit, rpm_limit, and allowed_models are all None, and there + is no explicit ``budget_duration`` in the request (``budget_duration_explicit`` + is False) with a resolved ``budget_duration`` of None, the user's budget is + removed from the team membership. - If max_budget, tpm_limit, rpm_limit, allowed_models, and budget_duration are all None, the user's budget is removed from the team membership. If any of these values exist, a budget is updated or created and linked to the team membership. """ if ( @@ -387,6 +395,7 @@ async def _upsert_budget_and_membership( and rpm_limit is None and allowed_models is None and budget_duration is None + and not budget_duration_explicit ): # disconnect the budget since all limits are None await tx.litellm_teammembership.update( @@ -415,7 +424,16 @@ async def _upsert_budget_and_membership( update_data["rpm_limit"] = rpm_limit if allowed_models is not None: update_data["allowed_models"] = allowed_models - if budget_duration is not 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: update_data["budget_duration"] = budget_duration update_data["budget_reset_at"] = get_budget_reset_time( budget_duration=budget_duration @@ -468,12 +486,20 @@ async def _upsert_budget_and_membership( create_data["rpm_limit"] = rpm_limit if allowed_models is not None: create_data["allowed_models"] = allowed_models - if budget_duration is not None: + if budget_duration_explicit: + if budget_duration is not None: + create_data["budget_duration"] = budget_duration + else: + create_data.pop("budget_duration", None) + create_data.pop("budget_reset_at", None) + elif budget_duration is not None: create_data["budget_duration"] = budget_duration bd = create_data.get("budget_duration") if bd is not None: create_data["budget_reset_at"] = get_budget_reset_time(budget_duration=bd) + else: + create_data.pop("budget_reset_at", None) new_budget = await tx.litellm_budgettable.create( data=create_data, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 01e71e5c4ee..9f7c50dfdec 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2723,7 +2723,8 @@ async def team_member_update( ### resolve effective budget_duration # - Explicit value (including null) takes precedence # - If omitted, inherit the team's configured team_member_budget_duration - if "budget_duration" in data.model_fields_set: + budget_duration_explicit = "budget_duration" in data.model_fields_set + if budget_duration_explicit: effective_budget_duration = data.budget_duration else: if team_default_budget_id: @@ -2750,6 +2751,7 @@ async def team_member_update( allowed_models=data.allowed_models, team_default_budget_id=team_default_budget_id, budget_duration=effective_budget_duration, + budget_duration_explicit=budget_duration_explicit, ) ### update team member role 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 bc5d6562f7e..bfe63eef797 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 @@ -423,6 +423,96 @@ async def test_upsert_clones_when_pointing_at_shared_default(mock_tx, fake_user) ) +@pytest.mark.asyncio +async def test_upsert_explicit_lifetime_clears_duration_on_private_budget( + mock_tx, fake_user +): + """Explicit ``budget_duration: null`` clears periodic duration on private row.""" + await _upsert_budget_and_membership( + mock_tx, + team_id="team-lifetime", + user_id="user-lifetime", + max_budget=None, + existing_budget_id="private-budget-periodic", + user_api_key_dict=fake_user, + team_default_budget_id="team-default-budget-1", + budget_duration=None, + budget_duration_explicit=True, + ) + + mock_tx.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "private-budget-periodic"}, + data={ + "updated_by": fake_user.user_id, + "budget_duration": None, + "budget_reset_at": None, + }, + ) + mock_tx.litellm_budgettable.create.assert_not_called() + mock_tx.litellm_teammembership.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_upsert_clone_explicit_lifetime_strips_cloned_duration( + mock_tx, fake_user +): + """Clone-from-default: explicit null must not keep default's budget_duration.""" + shared_default_id = "team-default-budget-1" + default_row = MagicMock() + default_row.model_dump.return_value = { + "budget_id": shared_default_id, + "max_budget": 200.0, + "soft_budget": None, + "max_parallel_requests": None, + "tpm_limit": 500, + "rpm_limit": None, + "model_max_budget": None, + "budget_duration": "1d", + "allowed_models": [], + } + mock_tx.litellm_budgettable.find_unique = AsyncMock(return_value=default_row) + + await _upsert_budget_and_membership( + mock_tx, + team_id="team-shared-lt", + user_id="user-shared-lt", + max_budget=50.0, + existing_budget_id=shared_default_id, + user_api_key_dict=fake_user, + team_default_budget_id=shared_default_id, + budget_duration=None, + budget_duration_explicit=True, + ) + + mock_tx.litellm_budgettable.create.assert_awaited_once_with( + data={ + "created_by": fake_user.user_id, + "updated_by": fake_user.user_id, + "max_budget": 50.0, + "tpm_limit": 500, + }, + include={"team_membership": True}, + ) + + +@pytest.mark.asyncio +async def test_upsert_explicit_null_only_does_not_disconnect(mock_tx, fake_user): + """Explicit lifetime with no other fields must not disconnect membership budget.""" + await _upsert_budget_and_membership( + mock_tx, + team_id="team-explicit-null", + user_id="user-explicit-null", + max_budget=None, + existing_budget_id=None, + user_api_key_dict=fake_user, + budget_duration=None, + budget_duration_explicit=True, + ) + + mock_tx.litellm_teammembership.update.assert_not_called() + mock_tx.litellm_budgettable.create.assert_awaited_once() + + # TEST: when team default exists but member already has their own budget, in-place update @pytest.mark.asyncio async def test_upsert_updates_in_place_when_member_has_private_budget( diff --git a/tests/test_litellm/proxy/test_team_member_update.py b/tests/test_litellm/proxy/test_team_member_update.py index 3cd1aae5fca..eab4d223af5 100644 --- a/tests/test_litellm/proxy/test_team_member_update.py +++ b/tests/test_litellm/proxy/test_team_member_update.py @@ -105,6 +105,7 @@ async def test_team_member_update_explicit_budget_duration(): mock_upsert.assert_awaited_once() assert mock_upsert.call_args.kwargs["budget_duration"] == "30d" + assert mock_upsert.call_args.kwargs["budget_duration_explicit"] is True # Should NOT have fetched the team-level budget row prisma.db.litellm_budgettable.find_unique.assert_not_awaited() @@ -139,6 +140,7 @@ async def test_team_member_update_explicit_null_budget_duration(): mock_upsert.assert_awaited_once() assert mock_upsert.call_args.kwargs["budget_duration"] is None + assert mock_upsert.call_args.kwargs["budget_duration_explicit"] is True # Should NOT have fetched the team-level budget row prisma.db.litellm_budgettable.find_unique.assert_not_awaited() @@ -174,6 +176,7 @@ async def test_team_member_update_inherits_team_budget_duration(): mock_upsert.assert_awaited_once() assert mock_upsert.call_args.kwargs["budget_duration"] == "30d" + assert mock_upsert.call_args.kwargs["budget_duration_explicit"] is False prisma.db.litellm_budgettable.find_unique.assert_awaited_once_with( where={"budget_id": "team-bud-1"} ) @@ -208,6 +211,7 @@ async def test_team_member_update_no_team_budget_duration_defaults_to_none(): mock_upsert.assert_awaited_once() assert mock_upsert.call_args.kwargs["budget_duration"] is None + assert mock_upsert.call_args.kwargs["budget_duration_explicit"] is False prisma.db.litellm_budgettable.find_unique.assert_not_awaited()