diff --git a/litellm/models/budget.py b/litellm/models/budget.py index 2bef5d279d6..61123810fd1 100644 --- a/litellm/models/budget.py +++ b/litellm/models/budget.py @@ -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): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a36ef548e2f..4247bf49e71 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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 diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index a5937ea8c7c..78e3ac7bd66 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -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 diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 4028fcaf2ae..528f63064c6 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -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 diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py index 4845dd8a7df..7774f6b543d 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/test_litellm/models/test_models.py @@ -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): diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 3048d4a6a02..348973f2161 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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 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 cabf5e24933..08768f1636f 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 @@ -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() diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 05cad26ce10..0e0025f5194 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -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