diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 9c06e3afc18..9e3071f020e 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -8,7 +8,14 @@ from datetime import datetime, timedelta, timezone from enum import Enum from itertools import chain from types import MappingProxyType -from typing import Final, Generic, Literal, Protocol, TypeVar +from typing import ( + Final, + Generic, + Literal, + Protocol, + TypeVar, + cast, # noqa: TID251 # get_data returns validated rows selected by table_name +) from typing_extensions import assert_never @@ -94,6 +101,7 @@ _RowT = TypeVar("_RowT") class _RowReset(Generic[_RowT]): row: _RowT spend_decrement: float + expected_reset_at: datetime | None _LINKED_KEYS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"budget_duration": None, "spend": {"gt": 0}}) @@ -1050,6 +1058,8 @@ class ResetBudgetJob: token=k.row.token, budget_reset_at=k.row.budget_reset_at, spend_decrement=k.spend_decrement, + expected_reset_at=k.expected_reset_at, + spend_after_reset=float(k.row.spend or 0.0), ) async def _write_user_reset_updates(self, updated_users: Sequence[_RowReset[LiteLLM_UserTable]]) -> None: @@ -1075,6 +1085,8 @@ class ResetBudgetJob: user_id=u.row.user_id, budget_reset_at=u.row.budget_reset_at, spend_decrement=u.spend_decrement, + expected_reset_at=u.expected_reset_at, + spend_after_reset=float(u.row.spend or 0.0), ) async def _write_team_reset_updates(self, updated_teams: Sequence[_RowReset[LiteLLM_TeamTable]]) -> None: @@ -1100,6 +1112,8 @@ class ResetBudgetJob: team_id=t.row.team_id, budget_reset_at=t.row.budget_reset_at, spend_decrement=t.spend_decrement, + expected_reset_at=t.expected_reset_at, + spend_after_reset=float(t.row.spend or 0.0), ) def _emit_phase_failure( @@ -1156,7 +1170,12 @@ class ResetBudgetJob: if keys_to_reset is not None and len(keys_to_reset) > 0: for key in keys_to_reset: try: - pre_reset_spend = float(key.spend or 0.0) + pre_reset_spend, expected_reset_at = ( + float(key.spend or 0.0), + cast( # cast-ok: get_data's key branch returns validated key rows + LiteLLM_VerificationToken, key + ).budget_reset_at, + ) updated_key = await ResetBudgetJob._reset_budget_for_key( key=key, current_time=now, @@ -1167,6 +1186,7 @@ class ResetBudgetJob: _RowReset( row=updated_key, spend_decrement=pre_reset_spend - float(updated_key.spend or 0.0), + expected_reset_at=expected_reset_at, ) ) else: @@ -1263,7 +1283,12 @@ class ResetBudgetJob: if users_to_reset is not None and len(users_to_reset) > 0: for user in users_to_reset: try: - pre_reset_spend = float(user.spend or 0.0) + pre_reset_spend, expected_reset_at = ( + float(user.spend or 0.0), + cast( # cast-ok: get_data's user branch returns validated user rows + LiteLLM_UserTable, user + ).budget_reset_at, + ) updated_user = await ResetBudgetJob._reset_budget_for_user( user=user, current_time=now, @@ -1274,6 +1299,7 @@ class ResetBudgetJob: _RowReset( row=updated_user, spend_decrement=pre_reset_spend - float(updated_user.spend or 0.0), + expected_reset_at=expected_reset_at, ) ) else: @@ -1378,7 +1404,12 @@ class ResetBudgetJob: if teams_to_reset is not None and len(teams_to_reset) > 0: for team in teams_to_reset: try: - pre_reset_spend = float(team.spend or 0.0) + pre_reset_spend, expected_reset_at = ( + float(team.spend or 0.0), + cast( # cast-ok: get_data's team branch returns validated team rows + LiteLLM_TeamTable, team + ).budget_reset_at, + ) updated_team = await ResetBudgetJob._reset_budget_for_team( team=team, current_time=now, @@ -1389,6 +1420,7 @@ class ResetBudgetJob: _RowReset( row=updated_team, spend_decrement=pre_reset_spend - float(updated_team.spend or 0.0), + expected_reset_at=expected_reset_at, ) ) else: diff --git a/litellm/proxy/common_utils/timezone_utils.py b/litellm/proxy/common_utils/timezone_utils.py index f7907675fb0..7a40931f50d 100644 --- a/litellm/proxy/common_utils/timezone_utils.py +++ b/litellm/proxy/common_utils/timezone_utils.py @@ -1,4 +1,5 @@ from datetime import datetime, time, timezone +from typing import Final from pydantic import ConfigDict @@ -81,6 +82,15 @@ def get_budget_reset_time(budget_duration: str) -> datetime: return compute_budget_reset_at(budget_duration, get_budget_reset_settings()) +def keeps_due_budget_reset( + requested_duration: str, stored_duration: str | None, stored_reset_at: datetime | None +) -> bool: + if stored_reset_at is None or requested_duration != stored_duration: + return False + stored_utc: Final = stored_reset_at if stored_reset_at.tzinfo else stored_reset_at.replace(tzinfo=timezone.utc) + return stored_utc <= datetime.now(timezone.utc) + + def _is_persistable_budget_duration(budget_duration: str) -> bool: from litellm.litellm_core_utils.duration_parser import duration_in_seconds diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index f88ac445a0b..21ac82754b4 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -51,6 +51,7 @@ from litellm.proxy.auth.password_policy import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.openai_error_payload import openai_shaped_proxy_exception, proxy_exception_for +from litellm.proxy.common_utils.timezone_utils import keeps_due_budget_reset from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, object_permission_cache_key, @@ -1576,6 +1577,16 @@ async def _update_single_user_helper( if existing_user_row is not None: existing_user_row = LiteLLM_UserTable.model_validate(existing_user_row.model_dump(exclude_none=True)) + if ( + "budget_reset_at" in non_default_values + and isinstance(non_default_values["budget_duration"], str) + and keeps_due_budget_reset( + non_default_values["budget_duration"], + existing_user_row.budget_duration, + existing_user_row.budget_reset_at, + ) + ): + del non_default_values["budget_reset_at"] # Prevent budget self-escalation (GHSA-wvg4-6222-3q4r): non-admin callers # must not be able to raise their own budget/spend fields. diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 038c1bdf5a6..b59606a8dcf 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -84,7 +84,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( ) from litellm.proxy.common_utils.openai_error_payload import openai_error_param, proxy_exception_for from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys -from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time, keeps_due_budget_reset from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache, is_user_key_cache_key from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage @@ -3116,11 +3116,11 @@ async def prepare_key_update_data( non_default_values["budget_duration"] = None non_default_values["budget_reset_at"] = None elif isinstance(budget_duration, str) and len(budget_duration) > 0: - from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - - key_reset_at: Final = get_budget_reset_time(budget_duration=budget_duration) - non_default_values["budget_reset_at"] = key_reset_at non_default_values["budget_duration"] = budget_duration + if not keeps_due_budget_reset( + budget_duration, existing_key_row.budget_duration, existing_key_row.budget_reset_at + ): + non_default_values["budget_reset_at"] = get_budget_reset_time(budget_duration=budget_duration) if "budget_limits" in non_default_values: requested_windows: Final = _budget_limit_windows(data.budget_limits) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 210aa1aaa6f..f14f3d93b79 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -135,6 +135,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time, keeps_due_budget_reset from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, team_membership_auth_cache_key, @@ -2774,7 +2775,7 @@ async def update_team( ) # Check budget_duration and budget_reset_at - _set_budget_reset_at(data, updated_kv) + _set_budget_reset_at(data, updated_kv, existing_team) _team_member_fields_in_request: Final = { field @@ -3002,19 +3003,16 @@ async def patch_team( raise handle_exception_on_proxy(e) -def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None: +def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict, existing_team: LiteLLM_TeamTable) -> None: """Set budget_reset_at in updated_kv if budget_duration is provided.""" - if data.budget_duration is not None: - from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - - reset_at: Final = get_budget_reset_time(budget_duration=data.budget_duration) - updated_kv["budget_reset_at"] = reset_at + if data.budget_duration is not None and not keeps_due_budget_reset( + data.budget_duration, existing_team.budget_duration, existing_team.budget_reset_at + ): + updated_kv["budget_reset_at"] = get_budget_reset_time(budget_duration=data.budget_duration) elif "budget_duration" in updated_kv and updated_kv["budget_duration"] is None: updated_kv["budget_reset_at"] = None if data.budget_limits is not None and len(data.budget_limits) > 0: - from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - initialized_windows: Final = [] for window in data.budget_limits: w = window if isinstance(window, dict) else window.model_dump() diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index b78f7e5a9b7..950b98e7b64 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -18,7 +18,8 @@ carrying fields the update input type rejects (see #27730). from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass -from datetime import datetime +from datetime import datetime, timezone +from math import isfinite, ulp from typing import Final, Protocol, runtime_checkable from litellm.repositories.prisma_protocols import BatchTable, PrismaBatch @@ -29,13 +30,79 @@ def _spend_reset_data(budget_reset_at: datetime | None, spend_decrement: float) return {"spend": spend, "budget_reset_at": budget_reset_at} +_SPEND_RESET_SET: Final = ( + " SET spend = CASE WHEN ABS(spend - $3::text::double precision) <= $5::text::double precision" + " THEN 0 ELSE GREATEST(spend - $3::text::double precision, 0) END, budget_reset_at = $2::timestamp," + " updated_at = (NOW() AT TIME ZONE 'UTC')" +) +_SPEND_RESET_WHERE: Final = " AND budget_reset_at IS NOT DISTINCT FROM $4::timestamp" +_KEY_SPEND_RESET_SQL: Final = ( + 'UPDATE "LiteLLM_VerificationToken"' + _SPEND_RESET_SET + " WHERE token = $1" + _SPEND_RESET_WHERE +) +_USER_SPEND_RESET_SQL: Final = ( + 'UPDATE "LiteLLM_UserTable"' + _SPEND_RESET_SET + " WHERE user_id = $1" + _SPEND_RESET_WHERE +) +_TEAM_SPEND_RESET_SQL: Final = ( + 'UPDATE "LiteLLM_TeamTable"' + _SPEND_RESET_SET + " WHERE team_id = $1" + _SPEND_RESET_WHERE +) + + +def _reset_timestamp(value: datetime | None) -> str | None: + if value is None: + return None + reset_at: Final = (value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value).astimezone(timezone.utc) + return reset_at.replace(tzinfo=None).isoformat(timespec="microseconds") + + +def _queue_guarded_spend_reset( + batch: PrismaBatch, + query: str, + identity: str, + budget_reset_at: datetime | None, + spend_decrement: float, + expected_reset_at: datetime | None, + spend_after_reset: float, +) -> bool: + if budget_reset_at is None or spend_decrement == 0 or not isfinite(spend_decrement): + return False + # Prisma rounds Float reads, so exact subtraction can still leave a few ULPs. + tolerance: Final = 4 * ulp(spend_decrement) if spend_after_reset == 0.0 else 0.0 + batch.execute_raw( + query, + identity, + _reset_timestamp(budget_reset_at), + repr(spend_decrement), + _reset_timestamp(expected_reset_at), + str(tolerance), + ) + return True + + @dataclass(frozen=True, slots=True) class KeySpendResetWrites: table: BatchTable + batch: PrismaBatch - def queue_spend_reset(self, token: str, budget_reset_at: datetime | None, spend_decrement: float) -> None: - self.table.update( - where={"token": token}, + def queue_spend_reset( + self, + token: str, + budget_reset_at: datetime | None, + spend_decrement: float, + expected_reset_at: datetime | None, + spend_after_reset: float = 0.0, + ) -> None: + if _queue_guarded_spend_reset( + self.batch, + _KEY_SPEND_RESET_SQL, + token, + budget_reset_at, + spend_decrement, + expected_reset_at, + spend_after_reset, + ): + return + self.table.update_many( + where={"token": token, "budget_reset_at": expected_reset_at}, data=_spend_reset_data(budget_reset_at, spend_decrement), ) @@ -43,10 +110,28 @@ class KeySpendResetWrites: @dataclass(frozen=True, slots=True) class UserSpendResetWrites: table: BatchTable + batch: PrismaBatch - def queue_spend_reset(self, user_id: str, budget_reset_at: datetime | None, spend_decrement: float) -> None: - self.table.update( - where={"user_id": user_id}, + def queue_spend_reset( + self, + user_id: str, + budget_reset_at: datetime | None, + spend_decrement: float, + expected_reset_at: datetime | None, + spend_after_reset: float = 0.0, + ) -> None: + if _queue_guarded_spend_reset( + self.batch, + _USER_SPEND_RESET_SQL, + user_id, + budget_reset_at, + spend_decrement, + expected_reset_at, + spend_after_reset, + ): + return + self.table.update_many( + where={"user_id": user_id, "budget_reset_at": expected_reset_at}, data=_spend_reset_data(budget_reset_at, spend_decrement), ) @@ -54,10 +139,28 @@ class UserSpendResetWrites: @dataclass(frozen=True, slots=True) class TeamSpendResetWrites: table: BatchTable + batch: PrismaBatch - def queue_spend_reset(self, team_id: str, budget_reset_at: datetime | None, spend_decrement: float) -> None: - self.table.update( - where={"team_id": team_id}, + def queue_spend_reset( + self, + team_id: str, + budget_reset_at: datetime | None, + spend_decrement: float, + expected_reset_at: datetime | None, + spend_after_reset: float = 0.0, + ) -> None: + if _queue_guarded_spend_reset( + self.batch, + _TEAM_SPEND_RESET_SQL, + team_id, + budget_reset_at, + spend_decrement, + expected_reset_at, + spend_after_reset, + ): + return + self.table.update_many( + where={"team_id": team_id, "budget_reset_at": expected_reset_at}, data=_spend_reset_data(budget_reset_at, spend_decrement), ) @@ -191,9 +294,9 @@ class BudgetCascadeUnitOfWork: async def spend_reset_unit_of_work(new_batch: Callable[[], PrismaBatch]) -> AsyncGenerator[SpendResetUnitOfWork, None]: batch = new_batch() yield SpendResetUnitOfWork( - keys=KeySpendResetWrites(table=batch.litellm_verificationtoken), - users=UserSpendResetWrites(table=batch.litellm_usertable), - teams=TeamSpendResetWrites(table=batch.litellm_teamtable), + keys=KeySpendResetWrites(table=batch.litellm_verificationtoken, batch=batch), + users=UserSpendResetWrites(table=batch.litellm_usertable, batch=batch), + teams=TeamSpendResetWrites(table=batch.litellm_teamtable, batch=batch), ) await batch.commit() diff --git a/tests/integration/spend/test_reset_budget_leader_election.py b/tests/integration/spend/test_reset_budget_leader_election.py index 58c018c4181..2cd46abfc52 100644 --- a/tests/integration/spend/test_reset_budget_leader_election.py +++ b/tests/integration/spend/test_reset_budget_leader_election.py @@ -1,16 +1,25 @@ +import asyncio import json import os +from collections.abc import Generator +from contextlib import contextmanager +from datetime import datetime, timedelta, timezone +from hashlib import sha256 from pathlib import Path from types import MappingProxyType -from typing import Final +from typing import Final, Literal import psycopg import pytest from integration._support.client import Gateway, eventually from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import OwnedState, owned_proxy +from prisma import Prisma +from psycopg import sql from redis import Redis +from litellm.repositories.unit_of_work import spend_reset_unit_of_work + pytestmark: Final = pytest.mark.shares_worker_state RESET_LEASE_KEY: Final = "cronjob_lock:reset_budget_job" @@ -20,6 +29,21 @@ FAST_RESET_TICK: Final = MappingProxyType( ) +@contextmanager +def _peer_reset_lease() -> Generator["Redis[bytes]", None, None]: + cache: Final["Redis[bytes]"] = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) + with cache: + eventually( + lambda: cache.set(RESET_LEASE_KEY, PEER_POD_LEASE, ex=300, nx=True), + lambda claimed: claimed is True, + seconds=60, + ) + try: + yield cache + finally: + cache.delete(RESET_LEASE_KEY) + + def _team_spend(team: str) -> float: rows: Final = read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team,)) assert len(rows) == 1, rows @@ -38,25 +62,26 @@ def _make_team_budget_due(team: str, spend: float) -> None: @pytest.mark.covers("spend.budget_reset.one_pod_sweeps_per_tick") +@pytest.mark.parametrize("starting_spend", (0.5, 40.080000000000005, 1009.9442791267694)) def test_reset_sweep_skips_ticks_while_another_pod_holds_the_lease_and_resumes_after_release( - gateway: Gateway, tmp_path: Path + gateway: Gateway, tmp_path: Path, starting_spend: float ) -> None: with ( gateway.scenario() as scenario, Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache, ): model: Final = scenario.model(input_cost_per_token=0.0, output_cost_per_token=0.0) - team: Final = scenario.team(max_budget=1.0, budget_duration="30d") + team: Final = scenario.team(max_budget=2000000.0, budget_duration="30d") key: Final = scenario.key(team_id=team) - _make_team_budget_due(team, spend=0.5) - assert _team_spend(team) == 0.5 + _make_team_budget_due(team, spend=starting_spend) + assert _team_spend(team) == starting_spend eventually( lambda: cache.set(RESET_LEASE_KEY, PEER_POD_LEASE, ex=120, nx=True), lambda claimed: claimed is True, seconds=60, ) try: - with owned_proxy(gateway, tmp_path, FAST_RESET_TICK) as replica: + with owned_proxy(gateway, tmp_path, FAST_RESET_TICK, workers=2) as replica: response: Final = replica.request( "POST", "/v1/chat/completions", @@ -65,13 +90,272 @@ def test_reset_sweep_skips_ticks_while_another_pod_holds_the_lease_and_resumes_a ) assert response.status_code == 200, response.text held: Final = eventually( - lambda: _team_spend(team), lambda spend: spend != 0.5, seconds=8, return_last_on_timeout=True + lambda: _team_spend(team), + lambda spend: spend != starting_spend, + seconds=8, + return_last_on_timeout=True, + ) + assert held == starting_spend, ( + f"team {team} was swept while another pod held the reset lease: spend={held}" ) - assert held == 0.5, f"team {team} was swept while another pod held the reset lease: spend={held}" assert cache.get(RESET_LEASE_KEY) == PEER_POD_LEASE.encode() cache.delete(RESET_LEASE_KEY) - swept: Final = eventually(lambda: _team_spend(team), lambda spend: spend == 0.0, seconds=15) + swept: Final = eventually( + lambda: _team_spend(team), lambda spend: spend == 0.0, seconds=15, return_last_on_timeout=True + ) + info: Final = replica.request("GET", f"/team/info?team_id={team}") + assert info.status_code == 200, info.text + assert info.json()["team_info"]["spend"] == 0.0 assert swept == 0.0 eventually(lambda: cache.get(RESET_LEASE_KEY), lambda value: value is None, seconds=15) finally: cache.delete(RESET_LEASE_KEY) + + +def _blocked_reset_updates(table: str) -> int: + rows: Final = read_rows( + "SELECT count(*) AS blocked FROM pg_stat_activity WHERE datname = current_database() " + "AND wait_event_type = 'Lock' AND query LIKE %s AND query LIKE '%%UPDATE%%'", + (f'%"{table}"%',), + ) + return int(str(rows[0]["blocked"])) + + +@pytest.mark.parametrize("entity", ("key", "user", "team")) +@pytest.mark.parametrize( + ("workers", "spend", "increment", "rollover", "reset_timezone"), + ( + (3, 40.0, 0.25, False, "UTC"), + (1, 40.080000000000005, 0.0, False, "UTC"), + (1, 1234567.8901234567, 0.0, False, "UTC"), + (1, 1009.9442791267694, 0.0, False, "UTC"), + (3, 40.000000000000014, 0.0, True, "UTC"), + (1, 0.30000000000000004, 0.0, False, "America/Los_Angeles"), + ), + ids=( + "overlap-preserves-new-spend", + "fractional-roundoff", + "large-roundoff", + "four-ulp-roundoff", + "rollover-keeps-small-carry", + "non-utc-roundoff", + ), +) +def test_scheduled_reset_preserves_new_spend_and_never_leaves_roundoff( + gateway: Gateway, + tmp_path: Path, + entity: Literal["key", "user", "team"], + workers: int, + spend: float, + increment: float, + rollover: bool, + reset_timezone: str, +) -> None: + max_budget: Final = 40.0 if rollover else 2000000.0 + expected_spend: Final = spend - max_budget + increment if rollover else increment + config: Final = tmp_path / "without-redis.yaml" + config.write_text( + "model_list: []\ngeneral_settings:\n" + " master_key: os.environ/LITELLM_MASTER_KEY\n" + " database_url: os.environ/DATABASE_URL\n" + " store_model_in_db: true\n" + "litellm_settings:\n cache: false\n" + f" budget_rollover: {str(rollover).lower()}\n" + f" timezone: {reset_timezone}\n" + ) + with gateway.scenario() as scenario, _peer_reset_lease() as peer: + identity: Final = ( + sha256(scenario.key(max_budget=max_budget, budget_duration="1d").encode()).hexdigest() + if entity == "key" + else scenario.user(max_budget=max_budget, budget_duration="1d") + if entity == "user" + else scenario.team(max_budget=max_budget, budget_duration="1d") + ) + table, column = { + "key": ("LiteLLM_VerificationToken", "token"), + "user": ("LiteLLM_UserTable", "user_id"), + "team": ("LiteLLM_TeamTable", "team_id"), + }[entity] + with psycopg.connect(os.environ["DATABASE_URL"]) as gate: + gate.execute( + sql.SQL( + "UPDATE {} SET spend = %s, budget_reset_at = now() - interval '1 day'," + " updated_at = now() - interval '1 day' WHERE {} = %s" + ).format(sql.Identifier(table), sql.Identifier(column)), + (spend, identity), + ) + gate.commit() + gate.execute( + sql.SQL("SELECT 1 FROM {} WHERE {} = %s FOR UPDATE").format( + sql.Identifier(table), sql.Identifier(column) + ), + (identity,), + ) + with owned_proxy( + gateway, + tmp_path, + FAST_RESET_TICK, + config=config, + remove_environment=("REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_URL"), + workers=workers, + state=OwnedState.SHARED_WITH_WORKER, + ) as replica: + try: + blocked: Final = eventually( + lambda: _blocked_reset_updates(table), lambda count: count == workers, seconds=30 + ) + assert blocked == workers + gate.execute( + sql.SQL("UPDATE {} SET spend = spend + %s WHERE {} = %s").format( + sql.Identifier(table), sql.Identifier(column) + ), + (increment, identity), + ) + gate.commit() + eventually(lambda: _blocked_reset_updates(table), lambda count: count == 0, seconds=30) + with psycopg.connect(os.environ["DATABASE_URL"]) as observer: + result: Final = observer.execute( + sql.SQL( + "SELECT spend, budget_reset_at > now(), updated_at > now() - interval '1 hour'" + " FROM {} WHERE {} = %s" + ).format(sql.Identifier(table), sql.Identifier(column)), + (identity,), + ).fetchall() + parameter, info = { + "key": ("key", "info"), + "user": ("user_id", "user_info"), + "team": ("team_id", "team_info"), + }[entity] + response: Final = replica.request("GET", f"/{entity}/info?{parameter}={identity}") + assert response.status_code == 200, response.text + assert response.json()[info]["spend"] == pytest.approx(expected_spend, rel=1e-14, abs=0.0) + assert result == [(expected_spend, True, True)], result + assert peer.get(RESET_LEASE_KEY) == PEER_POD_LEASE.encode() + finally: + gate.rollback() + + +@pytest.mark.parametrize("entity", ("key", "user", "team")) +def test_stale_zero_reset_cannot_erase_another_workers_rollover(gateway: Gateway, entity: str) -> None: + carry: Final = 1.4210854715202004e-14 + previous_reset: Final = datetime(2020, 1, 1, tzinfo=timezone.utc) + next_reset: Final = datetime.now(timezone.utc).replace(microsecond=0) + timedelta(days=1) + with gateway.scenario() as scenario: + identity: Final = ( + sha256(scenario.key(max_budget=40.0, budget_duration="1d").encode()).hexdigest() + if entity == "key" + else scenario.user(max_budget=40.0, budget_duration="1d") + if entity == "user" + else scenario.team(max_budget=40.0, budget_duration="1d") + ) + table, column = { + "key": ("LiteLLM_VerificationToken", "token"), + "user": ("LiteLLM_UserTable", "user_id"), + "team": ("LiteLLM_TeamTable", "team_id"), + }[entity] + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute( + sql.SQL("UPDATE {} SET spend = %s, budget_reset_at = %s WHERE {} = %s").format( + sql.Identifier(table), sql.Identifier(column) + ), + (40.0 + carry, previous_reset, identity), + ) + + async def commit_winner_then_stale_reset() -> None: + async with Prisma() as client: + for spend_after_reset in (carry, 0.0): + async with spend_reset_unit_of_work(client.batch_) as writes: + writer: Final = {"key": writes.keys, "user": writes.users, "team": writes.teams}[entity] + writer.queue_spend_reset(identity, next_reset, 40.0, previous_reset, spend_after_reset) + assert read_rows( + sql.SQL("SELECT spend FROM {} WHERE {} = %s") + .format(sql.Identifier(table), sql.Identifier(column)) + .as_string(), + (identity,), + ) == [{"spend": carry}] + + asyncio.run(commit_winner_then_stale_reset()) + parameter, info = { + "key": ("key", "info"), + "user": ("user_id", "user_info"), + "team": ("team_id", "team_info"), + }[entity] + response: Final = gateway.request("GET", f"/{entity}/info?{parameter}={identity}") + assert response.status_code == 200, response.text + assert response.json()[info]["spend"] == pytest.approx(carry, rel=1e-14, abs=0.0) + + +@pytest.mark.parametrize("entity", ("user", "team")) +def test_reset_starts_the_window_of_a_row_that_never_had_a_reset_time(gateway: Gateway, entity: str) -> None: + next_reset: Final = datetime.now(timezone.utc).replace(microsecond=0) + timedelta(days=1) + with gateway.scenario() as scenario: + identity: Final = ( + scenario.user(max_budget=100.0, budget_duration="1d") + if entity == "user" + else scenario.team(max_budget=100.0, budget_duration="1d") + ) + table, column, info = { + "user": ("LiteLLM_UserTable", "user_id", "user_info"), + "team": ("LiteLLM_TeamTable", "team_id", "team_info"), + }[entity] + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute( + sql.SQL("UPDATE {} SET spend = 40, budget_reset_at = NULL WHERE {} = %s").format( + sql.Identifier(table), sql.Identifier(column) + ), + (identity,), + ) + + async def reset() -> None: + async with Prisma() as client: + async with spend_reset_unit_of_work(client.batch_) as writes: + (writes.users if entity == "user" else writes.teams).queue_spend_reset( + identity, next_reset, 40.0, None + ) + + asyncio.run(reset()) + response: Final = gateway.request("GET", f"/{entity}/info?{column}={identity}") + assert response.status_code == 200, response.text + assert response.json()[info]["spend"] == 0.0 + assert response.json()[info]["budget_reset_at"] is not None + + +@pytest.mark.parametrize("entity", ("key", "user", "team")) +@pytest.mark.parametrize(("spend_decrement", "spend_after_reset"), ((40.0, 0.0), (30.0, 10.0))) +def test_reset_after_a_manual_spend_reset_never_goes_negative( + gateway: Gateway, entity: str, spend_decrement: float, spend_after_reset: float +) -> None: + previous_reset: Final = datetime(2020, 1, 1, tzinfo=timezone.utc) + next_reset: Final = datetime.now(timezone.utc).replace(microsecond=0) + timedelta(days=1) + with gateway.scenario() as scenario: + identity: Final = ( + sha256(scenario.key(max_budget=100.0, budget_duration="1d").encode()).hexdigest() + if entity == "key" + else scenario.user(max_budget=100.0, budget_duration="1d") + if entity == "user" + else scenario.team(max_budget=100.0, budget_duration="1d") + ) + table, column, parameter, info = { + "key": ("LiteLLM_VerificationToken", "token", "key", "info"), + "user": ("LiteLLM_UserTable", "user_id", "user_id", "user_info"), + "team": ("LiteLLM_TeamTable", "team_id", "team_id", "team_info"), + }[entity] + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute( + sql.SQL("UPDATE {} SET spend = 0, budget_reset_at = %s WHERE {} = %s").format( + sql.Identifier(table), sql.Identifier(column) + ), + (previous_reset, identity), + ) + + async def reset_with_the_spend_read_before_the_manual_reset() -> None: + async with Prisma() as client: + async with spend_reset_unit_of_work(client.batch_) as writes: + {"key": writes.keys, "user": writes.users, "team": writes.teams}[entity].queue_spend_reset( + identity, next_reset, spend_decrement, previous_reset, spend_after_reset + ) + + asyncio.run(reset_with_the_spend_read_before_the_manual_reset()) + response: Final = gateway.request("GET", f"/{entity}/info?{parameter}={identity}") + assert response.status_code == 200, response.text + assert response.json()[info]["spend"] == 0.0 diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index aaa686f2a98..f0ce610be01 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -77,6 +77,33 @@ class MockTable: return {"count": 1} +def _raw_reset_write(query: str, args: Sequence[object]) -> dict[str, object] | None: + if "CASE WHEN ABS(spend -" not in query: + return None + table, column = { + "LiteLLM_VerificationToken": ("key", "token"), + "LiteLLM_UserTable": ("user", "user_id"), + "LiteLLM_TeamTable": ("team", "team_id"), + }[query.split('"')[1]] + identity, reset_at, decrement, expected_reset_at, _ = args + assert isinstance(identity, str) and isinstance(reset_at, str) and isinstance(decrement, str) + assert expected_reset_at is None or isinstance(expected_reset_at, str) + return { + "table": table, + "op": "update_many", + "where": { + column: identity, + "budget_reset_at": datetime.fromisoformat(expected_reset_at).replace(tzinfo=timezone.utc) + if expected_reset_at is not None + else None, + }, + "data": { + "spend": {"decrement": float(decrement)}, + "budget_reset_at": datetime.fromisoformat(reset_at).replace(tzinfo=timezone.utc), + }, + } + + class MockBatcher: """Captures the writes queued on one `db.batch_()` and whether it committed. @@ -115,7 +142,9 @@ class MockBatcher: self.litellm_endusertable = _Table("enduser", self) def execute_raw(self, query: str, *args: object) -> None: - self.calls.append({"table": "raw", "op": "execute_raw", "query": query, "args": args}) + self.calls.append( + _raw_reset_write(query, args) or {"table": "raw", "op": "execute_raw", "query": query, "args": args} + ) async def commit(self): self.committed = True @@ -283,15 +312,15 @@ def test_write_key_reset_updates_skips_none_token_and_still_writes_the_rest(rese asyncio.run( reset_budget_job._write_key_reset_updates( - updated_keys=[_RowReset(row=k, spend_decrement=(k.spend or 0.0)) for k in keys] + updated_keys=[_RowReset(row=k, spend_decrement=(k.spend or 0.0), expected_reset_at=reset_at) for k in keys] ) ) assert _batch_writes(mock_prisma_client, "key") == [ { "table": "key", - "op": "update", - "where": {"token": "tok-ok"}, + "op": "update_many", + "where": {"token": "tok-ok", "budget_reset_at": reset_at}, "data": {"spend": {"decrement": 0.0}, "budget_reset_at": reset_at}, } ] @@ -323,7 +352,7 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client): key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"] assert len(key_writes) == 1 write = key_writes[0] - assert write["where"] == {"token": "tok-key-1"} + assert write["where"] == {"token": "tok-key-1", "budget_reset_at": now} assert write["data"]["spend"] == {"decrement": 100.0} assert write["data"]["budget_reset_at"] > now assert set(write["data"].keys()) == {"spend", "budget_reset_at"} @@ -403,7 +432,7 @@ def test_reset_budget_for_user(reset_budget_job, mock_prisma_client): user_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "user"] assert len(user_writes) == 1 write = user_writes[0] - assert write["where"] == {"user_id": "uid-1"} + assert write["where"] == {"user_id": "uid-1", "budget_reset_at": now} assert write["data"]["spend"] == {"decrement": 200.0} assert write["data"]["budget_reset_at"] > now assert set(write["data"].keys()) == {"spend", "budget_reset_at"} @@ -432,7 +461,7 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): team_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "team"] assert len(team_writes) == 1 write = team_writes[0] - assert write["where"] == {"team_id": "tid-1"} + assert write["where"] == {"team_id": "tid-1", "budget_reset_at": now} assert write["data"]["spend"] == {"decrement": 500.0} assert write["data"]["budget_reset_at"] > now assert set(write["data"].keys()) == {"spend", "budget_reset_at"} @@ -663,9 +692,13 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): ("user", {"user_id": "uid-all-1"}, 200.0), ("team", {"team_id": "tid-all-1"}, 500.0), ]: - writes = _batch_writes(mock_prisma_client, table_name, op="update") + writes = [ + write + for write in _batch_writes(mock_prisma_client, table_name, op="update_many") + if "budget_reset_at" in write["where"] + ] assert len(writes) == 1, f"expected 1 {table_name} write, got {len(writes)}" - assert writes[0]["where"] == where + assert writes[0]["where"] == {**where, "budget_reset_at": now} assert writes[0]["data"]["spend"] == {"decrement": decrement} assert set(writes[0]["data"].keys()) == {"spend", "budget_reset_at"} @@ -1839,7 +1872,6 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma assert evicted == {f"end_user_id:cust-{i:06d}" for i in range(population)} - def test_enduser_invalidation_reports_a_page_read_failure_instead_of_a_clean_finish( mock_prisma_client, monkeypatch ): @@ -2379,8 +2411,8 @@ def test_key_reset_walks_the_due_rows_one_chunk_at_a_time(monkeypatch): assert _fetch_limits(client, "key") == [2, 2] assert len(client.db.batchers) == 2 assert all(batcher.committed for batcher in client.db.batchers) - assert [len(batcher.calls) for batcher in client.db.batchers] == [2, 1] - assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update")] == ["k1", "k2", "k3"] + assert [sum(call["op"] == "update_many" for call in batcher.calls) for batcher in client.db.batchers] == [2, 1] + assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update_many")] == ["k1", "k2", "k3"] def test_key_reset_stops_after_a_chunk_shorter_than_the_batch_size(monkeypatch): @@ -2453,7 +2485,7 @@ def test_user_and_team_resets_are_chunked_too(monkeypatch, phase, table_name, ro assert client.fetches_by_table[table_name] == 2 assert _fetch_limits(client, table_name) == [2, 2] assert len(client.db.batchers) == 2 - assert len(_batch_writes(client, table_name, op="update")) == 3 + assert len(_batch_writes(client, table_name, op="update_many")) == 3 def test_budget_table_reset_walks_chunks_until_it_runs_dry(monkeypatch): @@ -2510,7 +2542,7 @@ def test_key_reset_stops_when_the_new_reset_time_is_not_in_the_future(monkeypatc asyncio.run(job.reset_budget_for_litellm_keys()) assert client.fetches_by_table["key"] == 1 - assert len(_batch_writes(client, "key", op="update")) == 2 + assert len(_batch_writes(client, "key", op="update_many")) == 2 def test_budget_table_reset_stops_when_the_new_window_is_not_in_the_future(monkeypatch): @@ -2577,7 +2609,7 @@ def test_key_reset_keeps_paging_when_some_rows_in_a_chunk_fail(monkeypatch): _run_and_drain_hooks(job.reset_budget_for_litellm_keys) assert client.fetches_by_table["key"] == 2 - assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update")] == ["k1", "k2"] + assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update_many")] == ["k1", "k2"] assert [call["call_type"] for call in logging_obj.service_logging_obj.failure_calls] == ["reset_budget_keys"] assert set(logging_obj.service_logging_obj.failure_calls[0]["event_metadata"]) == {"num_keys_found"} assert [call["call_type"] for call in logging_obj.service_logging_obj.success_calls] == ["reset_budget_keys"] @@ -2602,7 +2634,7 @@ def test_user_and_team_chunks_report_progress_despite_a_failed_row( _run_and_drain_hooks(getattr(job, phase)) assert client.fetches_by_table[table_name] == 2 - assert len(_batch_writes(client, table_name, op="update")) == 2 + assert len(_batch_writes(client, table_name, op="update_many")) == 2 assert [call["call_type"] for call in logging_obj.service_logging_obj.failure_calls] == [call_type] @@ -3025,7 +3057,7 @@ def test_transient_transport_error_on_read_reconnects_and_still_resets(phase, ta asyncio.run(getattr(job, phase)()) assert client.reconnect_reasons == [reason] - assert len(_batch_writes(client, table_name, op="update")) == 1 + assert len(_batch_writes(client, table_name, op="update_many")) == 1 @pytest.mark.parametrize( @@ -3049,7 +3081,7 @@ def test_connect_error_on_write_reconnects_and_commits(phase, table_name, reason assert client.reconnect_reasons == [reason] assert client.commit_attempts == 2 - assert len(_batch_writes(client, table_name, op="update")) == 1 + assert len(_batch_writes(client, table_name, op="update_many")) == 1 @pytest.mark.parametrize("ambiguous_error_name", ["ReadError", "ReadTimeout"]) @@ -3081,7 +3113,7 @@ def test_ambiguous_transport_error_on_read_still_retries(ambiguous_error_name): asyncio.run(job.reset_budget_for_litellm_keys()) assert client.reconnect_reasons == ["reset_budget_read_keys_failure"] - assert len(_batch_writes(client, "key", op="update")) == 1 + assert len(_batch_writes(client, "key", op="update_many")) == 1 def test_transport_error_on_budget_cascade_read_reconnects_and_commits(): @@ -3318,7 +3350,15 @@ def rollover_enabled(monkeypatch): ], ) def test_direct_reset_carries_overage_when_rollover_enabled( - rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch, run_phase, table, id_field, id_value, row_factory + rollover_enabled, + reset_budget_job, + mock_prisma_client, + monkeypatch, + run_phase, + table, + id_field, + id_value, + row_factory, ): """spend=150 against max_budget=100 must decrement by the cap (leaving 50) rather than zero the row, and the spend counter must be seeded with 50.""" @@ -3330,7 +3370,7 @@ def test_direct_reset_carries_overage_when_rollover_enabled( writes = _batch_writes(mock_prisma_client, table) assert len(writes) == 1 - assert writes[0]["where"] == {id_field: id_value} + assert writes[0]["where"] == {id_field: id_value, "budget_reset_at": now} assert writes[0]["data"]["spend"] == {"decrement": 100.0} assert writes[0]["data"]["budget_reset_at"] > now counter_prefix = {"key": "spend:key", "user": "spend:user", "team": "spend:team"}[table] @@ -3652,7 +3692,7 @@ def test_reset_decrement_preserves_spend_landed_after_read( writes = _batch_writes(mock_prisma_client, table) assert len(writes) == 1 - assert writes[0]["where"] == {id_field: id_value} + assert writes[0]["where"] == {id_field: id_value, "budget_reset_at": now} assert writes[0]["data"]["spend"] == {"decrement": 5.0} assert writes[0]["data"]["budget_reset_at"] > now assert _apply_spend_payload(db_spend=5.4, spend_field=writes[0]["data"]["spend"]) == pytest.approx(0.4) @@ -3756,6 +3796,7 @@ async def test_a_budget_window_write_renders_a_postgres_update_span_for_its_tabl assert update.await_count == 1 assert await postgres_span_names() == (span_name,) + def _attrify(d: dict): """ Wrap a dict so that attribute access (`.token`, `.user_id`, `.team_id`, @@ -3775,7 +3816,8 @@ def _attrify(d: dict): def __setattr__(self, k, v): self[k] = v - return _AttrDict(d) + return _AttrDict({"budget_reset_at": None, **d}) + def _wire_batcher_for_test(prisma_client, fail_commit=False): """ @@ -3834,12 +3876,19 @@ def _wire_batcher_for_test(prisma_client, fail_commit=False): batcher.litellm_organizationtable = _Table("org") batcher.litellm_tagtable = _Table("tag") batcher.litellm_endusertable = _Table("enduser") + def execute_raw(query: str, *args: object) -> None: + write: Final = _raw_reset_write(query, args) + if write is not None: + queued.append(write) + + batcher.execute_raw = execute_raw batcher.commit = commit return batcher prisma_client.db.batch_ = MagicMock(side_effect=make_batcher) return batch_calls + def _wire_cascade_reads_for_test(prisma_client, endusers=()): """ The budget tier's cascade reads the rows it is about to zero, so their @@ -3860,6 +3909,7 @@ def _wire_cascade_reads_for_test(prisma_client, endusers=()): getattr(prisma_client.db, table).find_many = AsyncMock(return_value=[]) prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=list(endusers)) + @pytest.mark.asyncio async def test_reset_budget_keys_partial_failure(): """ @@ -3919,9 +3969,7 @@ async def test_reset_budget_keys_partial_failure(): # Simulate successful reset modification key["spend"] = 0.0 # Compute a new reset time based on the budget duration - key["budget_reset_at"] = ( - current_time + timedelta(seconds=key["budget_duration"]) - ).isoformat() + key["budget_reset_at"] = current_time + timedelta(seconds=key["budget_duration"]) return key with patch.object( @@ -3957,6 +4005,7 @@ async def test_reset_budget_keys_partial_failure(): for call in failure_hook_calls ) + @pytest.mark.asyncio async def test_reset_budget_users_partial_failure(): """ @@ -4007,9 +4056,7 @@ async def test_reset_budget_users_partial_failure(): raise Exception("Simulated failure for user1") else: user["spend"] = 0.0 - user["budget_reset_at"] = ( - current_time + timedelta(seconds=user["budget_duration"]) - ).isoformat() + user["budget_reset_at"] = current_time + timedelta(seconds=user["budget_duration"]) return user with patch.object( @@ -4038,6 +4085,7 @@ async def test_reset_budget_users_partial_failure(): for call in failure_hook_calls ) + @pytest.mark.asyncio async def test_reset_budget_endusers_cascade_failure_is_all_or_nothing(): """ @@ -4251,12 +4299,12 @@ async def test_reset_budget_continues_other_categories_on_failure(frozen_reset_c assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1"]}, "spend": {"gt": 0}} assert enduser_writes[0]["data"] == {"spend": 0} - key_writes: Final = [c for c in batch_calls if c["table"] == "key" and c["op"] == "update"] + key_writes: Final = [c for c in batch_calls if c["table"] == "key" and "token" in c["where"]] user_writes: Final = [c for c in batch_calls if c["table"] == "user"] team_writes: Final = [c for c in batch_calls if c["table"] == "team"] assert len(key_writes) == 2 assert len(user_writes) == 1 - assert user_writes[0]["where"] == {"user_id": "user2"} + assert user_writes[0]["where"] == {"user_id": "user2", "budget_reset_at": None} assert len(team_writes) == 2 for c in key_writes + user_writes + team_writes: assert set(c["data"].keys()) == {"spend", "budget_reset_at"} @@ -4300,9 +4348,7 @@ async def test_reset_budget_teams_partial_failure(): raise Exception("Simulated failure for team1") else: team["spend"] = 0.0 - team["budget_reset_at"] = ( - current_time + timedelta(seconds=team["budget_duration"]) - ).isoformat() + team["budget_reset_at"] = current_time + timedelta(seconds=team["budget_duration"]) return team with patch.object( @@ -4315,7 +4361,7 @@ async def test_reset_budget_teams_partial_failure(): prisma_client.update_data.assert_not_awaited() team_writes = [c for c in batch_calls if c["table"] == "team"] assert len(team_writes) == 1 - assert team_writes[0]["where"] == {"team_id": "team2"} + assert team_writes[0]["where"] == {"team_id": "team2", "budget_reset_at": None} assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"} assert team_writes[0]["data"]["spend"] == {"decrement": pre_reset_spend} @@ -4327,6 +4373,7 @@ async def test_reset_budget_teams_partial_failure(): for call in failure_hook_calls ) + @pytest.mark.asyncio async def test_service_logger_keys_success(): """ @@ -4355,9 +4402,7 @@ async def test_service_logger_keys_success(): async def fake_reset_key(key, current_time, reset_settings=None): key["spend"] = 0.0 - key["budget_reset_at"] = ( - current_time + timedelta(seconds=key["budget_duration"]) - ).isoformat() + key["budget_reset_at"] = current_time + timedelta(seconds=key["budget_duration"]) return key with patch.object( @@ -4386,6 +4431,7 @@ async def test_service_logger_keys_success(): # Failure hook should not be executed. proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called() + @pytest.mark.asyncio async def test_service_logger_keys_failure(): """ @@ -4412,9 +4458,7 @@ async def test_service_logger_keys_failure(): if key["id"] == "key1": raise Exception("Simulated failure for key1") key["spend"] = 0.0 - key["budget_reset_at"] = ( - current_time + timedelta(seconds=key["budget_duration"]) - ).isoformat() + key["budget_reset_at"] = current_time + timedelta(seconds=key["budget_duration"]) return key with patch.object( @@ -4448,6 +4492,7 @@ async def test_service_logger_keys_failure(): # Success hook should not be called. proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called() + @pytest.mark.asyncio async def test_service_logger_users_success(): """ @@ -4476,9 +4521,7 @@ async def test_service_logger_users_success(): async def fake_reset_user(user, current_time, reset_settings=None): user["spend"] = 0.0 - user["budget_reset_at"] = ( - current_time + timedelta(seconds=user["budget_duration"]) - ).isoformat() + user["budget_reset_at"] = current_time + timedelta(seconds=user["budget_duration"]) return user with patch.object( @@ -4504,6 +4547,7 @@ async def test_service_logger_users_success(): assert event_metadata.get("num_users_failed") == 0 proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called() + @pytest.mark.asyncio async def test_service_logger_users_failure(): """ @@ -4529,9 +4573,7 @@ async def test_service_logger_users_failure(): if user["id"] == "user1": raise Exception("Simulated failure for user1") user["spend"] = 0.0 - user["budget_reset_at"] = ( - current_time + timedelta(seconds=user["budget_duration"]) - ).isoformat() + user["budget_reset_at"] = current_time + timedelta(seconds=user["budget_duration"]) return user with patch.object( @@ -4562,6 +4604,7 @@ async def test_service_logger_users_failure(): assert "users_found" not in event_metadata proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called() + @pytest.mark.asyncio async def test_service_logger_teams_success(): """ @@ -4590,9 +4633,7 @@ async def test_service_logger_teams_success(): async def fake_reset_team(team, current_time, reset_settings=None): team["spend"] = 0.0 - team["budget_reset_at"] = ( - current_time + timedelta(seconds=team["budget_duration"]) - ).isoformat() + team["budget_reset_at"] = current_time + timedelta(seconds=team["budget_duration"]) return team with patch.object( @@ -4618,6 +4659,7 @@ async def test_service_logger_teams_success(): assert event_metadata.get("num_teams_failed") == 0 proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called() + @pytest.mark.asyncio async def test_service_logger_teams_failure(): """ @@ -4643,9 +4685,7 @@ async def test_service_logger_teams_failure(): if team["id"] == "team1": raise Exception("Simulated failure for team1") team["spend"] = 0.0 - team["budget_reset_at"] = ( - current_time + timedelta(seconds=team["budget_duration"]) - ).isoformat() + team["budget_reset_at"] = current_time + timedelta(seconds=team["budget_duration"]) return team with patch.object( @@ -4676,6 +4716,7 @@ async def test_service_logger_teams_failure(): assert "teams_found" not in event_metadata proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called() + @pytest.mark.asyncio async def test_service_logger_endusers_success(): """ @@ -4740,6 +4781,7 @@ async def test_service_logger_endusers_success(): assert event_metadata.get("num_endusers_failed") == 0 proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called() + @pytest.mark.asyncio async def test_service_logger_endusers_failure(): """ @@ -4806,6 +4848,7 @@ async def test_service_logger_endusers_failure(): assert "budgets_found" not in event_metadata proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called() + @pytest.mark.asyncio async def test_reset_budget_for_litellm_team_members_called(): """ diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index b7be80da31a..ea36ffad044 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2459,6 +2459,54 @@ async def test_user_status_update_refreshes_cached_user( broadcast.assert_awaited_once_with(cache_key=saved_user.user_id) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_fields", "role_default_duration", "stored_reset_at", "moves_reset"), + [ + ({"budget_duration": "30d"}, None, datetime(2020, 1, 1, tzinfo=timezone.utc), False), + ({"budget_duration": "30d"}, None, datetime(2099, 1, 1, tzinfo=timezone.utc), True), + ({"budget_duration": "7d"}, None, datetime(2020, 1, 1, tzinfo=timezone.utc), True), + ({"user_role": LitellmUserRoles.INTERNAL_USER}, "30d", datetime(2020, 1, 1, tzinfo=timezone.utc), False), + ], +) +async def test_user_update_keeps_a_due_budget_reset_when_the_duration_is_resent( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + request_fields: dict[str, object], + role_default_duration: str | None, + stored_reset_at: datetime, + moves_reset: bool, +) -> None: + import litellm + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper + + monkeypatch.setattr(litellm, "internal_user_budget_duration", role_default_duration) + saved_user: Final = LiteLLM_UserTable( + user_id="user-spruce", + budget_duration="30d", + budget_reset_at=stored_reset_at, + ) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) + prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( # test-quality-ok: keep the Redis publication boundary offline + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=mocker.AsyncMock, + ) + + await _update_single_user_helper( + user_request=UpdateUserRequest(user_id=saved_user.user_id, **request_fields), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + written: Final = prisma_client.update_data.call_args.kwargs["data"] + reset_at: Final = written.get("budget_reset_at") + assert (isinstance(reset_at, datetime) and reset_at > datetime.now(timezone.utc)) is moves_reset + assert ("budget_reset_at" in written) is moves_reset + + @pytest.mark.asyncio async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocker: MockerFixture) -> None: from litellm.proxy._types import LiteLLM_UserTable diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 2e18a0c1099..0ebd6a3ff3a 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -2109,6 +2109,34 @@ async def test_update_key_end_user_budget_id_folds_into_metadata_and_survives_om assert untouched["metadata"]["end_user_budget_id"] == "svc-a-budget" +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stored_duration", "stored_reset_at", "moves_reset"), + [ + ("30d", datetime(2020, 1, 1, tzinfo=timezone.utc), False), + ("30d", datetime(2020, 1, 1), False), + ("30d", datetime(2099, 1, 1, tzinfo=timezone.utc), True), + ("7d", datetime(2020, 1, 1, tzinfo=timezone.utc), True), + ("30d", None, True), + ], +) +async def test_update_key_keeps_a_due_budget_reset_when_the_duration_is_resent( + stored_duration: str, stored_reset_at: datetime | None, moves_reset: bool +) -> None: + existing_key: Final = LiteLLM_VerificationToken( + token="hashed", budget_duration=stored_duration, budget_reset_at=stored_reset_at + ) + + updated: Final = await prepare_key_update_data( + data=UpdateKeyRequest(key="sk-1", budget_duration="30d"), existing_key_row=existing_key + ) + + assert updated["budget_duration"] == "30d" + reset_at: Final = updated.get("budget_reset_at") + assert (isinstance(reset_at, datetime) and reset_at > datetime.now(timezone.utc)) is moves_reset + assert ("budget_reset_at" in updated) is moves_reset + + @pytest.mark.asyncio async def test_update_key_clears_end_user_budget_id_with_empty_string(): """Sending an empty end_user_budget_id detaches the key default without touching any budget row, diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 0881c902a9d..7e54f0893c0 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -13155,7 +13155,7 @@ def test_set_budget_reset_at_clears_when_budget_duration_null(): data = UpdateTeamRequest(team_id="test-team", budget_duration=None) updated_kv = {"team_id": "test-team", "budget_duration": None} - _set_budget_reset_at(data, updated_kv) + _set_budget_reset_at(data, updated_kv, LiteLLM_TeamTable(team_id="test-team")) assert "budget_reset_at" in updated_kv assert updated_kv["budget_reset_at"] is None @@ -13172,7 +13172,7 @@ def test_set_budget_reset_at_noop_when_budget_duration_not_sent(): data = UpdateTeamRequest(team_id="test-team") updated_kv = {"team_id": "test-team"} - _set_budget_reset_at(data, updated_kv) + _set_budget_reset_at(data, updated_kv, LiteLLM_TeamTable(team_id="test-team")) assert "budget_reset_at" not in updated_kv @@ -13188,12 +13188,41 @@ def test_set_budget_reset_at_sets_value_when_budget_duration_provided(): data = UpdateTeamRequest(team_id="test-team", budget_duration="30d") updated_kv = {"team_id": "test-team", "budget_duration": "30d"} - _set_budget_reset_at(data, updated_kv) + _set_budget_reset_at(data, updated_kv, LiteLLM_TeamTable(team_id="test-team")) assert "budget_reset_at" in updated_kv assert updated_kv["budget_reset_at"] is not None +@pytest.mark.parametrize( + ("stored_duration", "stored_reset_at", "moves_reset"), + [ + ("30d", datetime(2020, 1, 1, tzinfo=timezone.utc), False), + ("30d", datetime(2020, 1, 1), False), + ("30d", datetime(2099, 1, 1, tzinfo=timezone.utc), True), + ("7d", datetime(2020, 1, 1, tzinfo=timezone.utc), True), + ("30d", None, True), + ], +) +def test_set_budget_reset_at_keeps_a_due_reset_when_the_duration_is_resent( + stored_duration: str, stored_reset_at: datetime | None, moves_reset: bool +) -> None: + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _set_budget_reset_at + + updated_kv: Final[dict[str, object]] = {"team_id": "test-team", "budget_duration": "30d"} + + _set_budget_reset_at( + UpdateTeamRequest(team_id="test-team", budget_duration="30d"), + updated_kv, + LiteLLM_TeamTable(team_id="test-team", budget_duration=stored_duration, budget_reset_at=stored_reset_at), + ) + + reset_at: Final = updated_kv.get("budget_reset_at") + assert (isinstance(reset_at, datetime) and reset_at > datetime.now(timezone.utc)) is moves_reset + assert ("budget_reset_at" in updated_kv) is moves_reset + + @pytest.mark.asyncio async def test_clear_team_member_budget_duration_calls_update_budget(): """ diff --git a/tests/unit/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py index 049e5275db6..5aea22b00cf 100644 --- a/tests/unit/repositories/test_unit_of_work.py +++ b/tests/unit/repositories/test_unit_of_work.py @@ -52,16 +52,24 @@ async def test_updates_across_tables_share_one_batch_and_commit_once(): reset_at = datetime(2026, 8, 3, 12, 0, tzinfo=timezone.utc) async with spend_reset_unit_of_work(lambda: batch) as uow: - uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=reset_at, spend_decrement=1.5) - uow.users.queue_spend_reset(user_id="user-1", budget_reset_at=reset_at, spend_decrement=2.5) - uow.teams.queue_spend_reset(team_id="team-1", budget_reset_at=None, spend_decrement=0.0) + uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=reset_at, spend_decrement=1.5, expected_reset_at=None) + uow.users.queue_spend_reset( + user_id="user-1", budget_reset_at=reset_at, spend_decrement=2.5, expected_reset_at=None + ) + uow.teams.queue_spend_reset(team_id="team-1", budget_reset_at=None, spend_decrement=0.0, expected_reset_at=None) assert batch.commit_count == 0 assert batch.commit_count == 1 - assert batch.calls == [ - ("litellm_verificationtoken", {"token": "tok-1"}, {"spend": {"decrement": 1.5}, "budget_reset_at": reset_at}), - ("litellm_usertable", {"user_id": "user-1"}, {"spend": {"decrement": 2.5}, "budget_reset_at": reset_at}), - ("litellm_teamtable", {"team_id": "team-1"}, {"spend": {"decrement": 0.0}, "budget_reset_at": None}), + assert [call[2]["args"][:4] for call in batch.calls if call[0] == "execute_raw"] == [ + ("tok-1", "2026-08-03T12:00:00.000000", "1.5", None), + ("user-1", "2026-08-03T12:00:00.000000", "2.5", None), + ] + assert [call for call in batch.calls if call[0] != "execute_raw"] == [ + ( + "litellm_teamtable.update_many", + {"team_id": "team-1", "budget_reset_at": None}, + {"spend": {"decrement": 0.0}, "budget_reset_at": None}, + ), ] @@ -70,7 +78,7 @@ async def test_raising_inside_block_skips_commit(): async def _blow_up_mid_transaction(): async with spend_reset_unit_of_work(lambda: batch) as uow: - uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=None, spend_decrement=0.0) + uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=None, spend_decrement=0.0, expected_reset_at=None) raise RuntimeError("boom") with pytest.raises(RuntimeError, match="boom"): @@ -133,7 +141,11 @@ async def test_enduser_page_commits_on_its_own_batch(): {"user_id": {"in": ["enduser-1", "enduser-2"]}, "spend": {"gt": 0}}, {"spend": 0}, ), - ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-3"]}, "spend": {"gt": 5.0}}, {"spend": {"decrement": 5.0}}), + ( + "litellm_endusertable.update_many", + {"user_id": {"in": ["enduser-3"]}, "spend": {"gt": 5.0}}, + {"spend": {"decrement": 5.0}}, + ), ]