import asyncio import json import math import time from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum from types import MappingProxyType from typing import Final, Literal, Protocol, TypeVar, assert_never import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache from litellm.constants import ( GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME, RESET_BUDGET_JOB_BATCH_SIZE, RESET_BUDGET_JOB_LOCK_TTL_SECONDS, RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN, RESET_BUDGET_JOB_NAME, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( DB_RETRY_SAFE_ERROR_TYPES, LiteLLM_BudgetTableFull, LiteLLM_EndUserTable, Litellm_EntityType, LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken, ) from litellm.proxy.common_utils.timezone_utils import ( BudgetResetSettings, compute_budget_reset_at, get_budget_reset_settings, ) from litellm.proxy.common_utils.user_api_key_cache import ( end_user_cache_key, model_access_group_cache_key, model_access_group_spend_counter_key, tag_cache_key, ) from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import SpendLinkedTable from litellm.repositories.table_repositories import ( EndUserRepository, ModelAccessGroupBudgetRepository, TagRepository, TeamMembershipRepository, ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.unit_of_work import ( LinkedSpendResetWrites, budget_cascade_unit_of_work, spend_reset_unit_of_work, ) from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) from litellm.types.services import ServiceTypes _RowT = TypeVar("_RowT") _LINKED_KEYS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"budget_duration": None, "spend": {"gt": 0}}) _SPENT_ROWS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"spend": {"gt": 0}}) class _BudgetLinkedRow(Protocol): @property def spend(self) -> float | None: ... @property def budget_id(self) -> str | None: ... class _TeamMembershipRow(_BudgetLinkedRow, Protocol): @property def user_id(self) -> str: ... @property def team_id(self) -> str: ... class _KeyRow(_BudgetLinkedRow, Protocol): @property def token(self) -> str: ... class _OrgRow(_BudgetLinkedRow, Protocol): @property def organization_id(self) -> str: ... class _TagRow(_BudgetLinkedRow, Protocol): @property def tag_name(self) -> str: ... class _ModelAccessGroupRow(_BudgetLinkedRow, Protocol): @property def access_group_name(self) -> str: ... class _EndUserRow(_BudgetLinkedRow, Protocol): @property def user_id(self) -> str: ... def _rollover_enabled() -> bool: return litellm.budget_rollover is True def _rollover_cap(max_budget: float | None) -> float | None: if max_budget is None or not math.isfinite(max_budget): return None return max_budget def _carried_spend(spend: float | None, cap: float | None) -> float: if cap is None: return 0.0 return max(0.0, (spend or 0.0) - cap) def _row_carried_spend(row: _BudgetLinkedRow, caps: Mapping[str, float]) -> float: if not caps: return 0.0 return _carried_spend(row.spend, caps.get(row.budget_id) if row.budget_id is not None else None) def _team_membership_counter_key(row: _TeamMembershipRow) -> str: return f"spend:team_member:{row.user_id}:{row.team_id}" def _team_membership_cache_keys(row: _TeamMembershipRow) -> tuple[str, ...]: return (f"{row.team_id}_{row.user_id}",) def _key_counter_key(row: _KeyRow) -> str: return f"spend:key:{row.token}" def _key_cache_keys(row: _KeyRow) -> tuple[str, ...]: return (row.token,) def _org_counter_key(row: _OrgRow) -> str: return f"spend:org:{row.organization_id}" def _org_cache_keys(row: _OrgRow) -> tuple[str, ...]: return ( f"org_id:{row.organization_id}", f"org_id:{row.organization_id}:with_budget", ) def _tag_counter_key(row: _TagRow) -> str: return f"spend:tag:{row.tag_name}" def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]: return (tag_cache_key(row.tag_name),) def _model_access_group_counter_key(row: _ModelAccessGroupRow) -> str: return model_access_group_spend_counter_key(row.access_group_name) def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]: return (model_access_group_cache_key(row.access_group_name),) def _enduser_counter_key(row: _EndUserRow) -> str: return f"spend:end_user:{row.user_id}" def _enduser_cache_keys(row: _EndUserRow) -> tuple[str, ...]: return (end_user_cache_key(row.user_id),) def _enduser_carried_spend(row: _EndUserRow, caps: Mapping[str, float]) -> float: if not caps: return 0.0 effective_budget_id = row.budget_id or litellm.max_end_user_budget_id return _carried_spend(row.spend, caps.get(effective_budget_id) if effective_budget_id is not None else None) def _budget_link_where( budget_ids: Sequence[str], extra: Mapping[str, object] = MappingProxyType({}), ) -> dict[str, object]: return {"budget_id": {"in": list(budget_ids)}, **extra} def _queue_budget_linked_resets( writes: LinkedSpendResetWrites, cascade: "_BudgetCascade", extra: Mapping[str, object] = MappingProxyType({}), ) -> None: """Reset one linked table's spend for every expiring tier: tiers with a rollover cap keep spend beyond the cap (decrement preserves writes racing the reset), everything else is zeroed as before. Zero the under-cap rows BEFORE decrementing the over-cap ones: the statements run sequentially in one transaction, so the reverse order lets the zero re-match a row the decrement just moved into the (0, cap] range and erase its carried spend.""" for budget_id, cap in cascade.rollover_caps.items(): writes.queue_spend_zero( where={"budget_id": budget_id, **extra, "spend": {"gt": 0, "lte": cap}} ) # mutable-ok: prisma where filter must be a dict writes.queue_spend_decrement( where={"budget_id": budget_id, **extra, "spend": {"gt": cap}}, amount=cap ) # mutable-ok: prisma where filter must be a dict plain_ids: Final = tuple(bid for bid in cascade.budget_ids if bid not in cascade.rollover_caps) if plain_ids: writes.queue_spend_zero(where=_budget_link_where(plain_ids, extra)) def _queue_enduser_resets(writes: LinkedSpendResetWrites, cascade: "_BudgetCascade") -> None: """End users reset on the budget link like every other gated table, plus a NULL-budget_id branch: rows created implicitly persist no link and ride the default tier (litellm.max_end_user_budget_id). Matching on the link rather than enumerating user ids keeps a statement's bind count proportional to the expiring tiers instead of the customer population, which past ~32,700 dependents exceeds PostgreSQL's per-statement bind ceiling and wedges the cascade permanently (#40564). """ _queue_budget_linked_resets(writes, cascade, extra=_SPENT_ROWS_WHERE) default_budget_id: Final = litellm.max_end_user_budget_id if default_budget_id is None or default_budget_id not in cascade.budget_ids: return cap: Final = cascade.rollover_caps.get(default_budget_id) if cap is None: writes.queue_spend_zero( where={"budget_id": None, **_SPENT_ROWS_WHERE} ) # mutable-ok: prisma where filter must be a dict return writes.queue_spend_zero( where={"budget_id": None, "spend": {"gt": 0, "lte": cap}} ) # mutable-ok: prisma where filter must be a dict writes.queue_spend_decrement( where={"budget_id": None, "spend": {"gt": cap}}, amount=cap ) # mutable-ok: prisma where filter must be a dict @dataclass(frozen=True, slots=True) class _BudgetCascade: """Everything one budget-tier reset touches, resolved before any write.""" budgets: tuple[LiteLLM_BudgetTableFull, ...] = () budget_ids: tuple[str, ...] = () budget_resets: tuple[tuple[str, datetime], ...] = () endusers: tuple[_EndUserRow, ...] = () counter_resets: tuple[tuple[str, float], ...] = () cache_keys: tuple[str, ...] = () rollover_caps: Mapping[str, float] = field(default_factory=lambda: MappingProxyType({})) @dataclass(frozen=True, slots=True) class _BudgetCascadeCommitted: cascade: _BudgetCascade advanced: int @dataclass(frozen=True, slots=True) class _BudgetCascadeFailed: cascade: _BudgetCascade error: Exception _EMPTY_CASCADE: Final = _BudgetCascade() @dataclass(frozen=True, slots=True) class _ChunkOutcome: """One chunk of a reset phase: rows read, and rows whose new budget_reset_at cleared the due cutoff. Anything else is still due and would come straight back on the next fetch, so it is not progress.""" fetched: int advanced: int _NO_PROGRESS: Final = _ChunkOutcome(fetched=0, advanced=0) def _as_utc(moment: datetime) -> datetime: return moment if moment.tzinfo is not None else moment.replace(tzinfo=timezone.utc) def _count_advanced(reset_ats: Iterable[object], cutoff: datetime) -> int: """How many rows the write actually moved past the due cutoff. A budget_duration of "0s" (or one the parser cannot read) resolves to the current time, so the row is written and stays due. Counting it as progress would re-read the same chunk until the per-run cap on every tick. """ utc_cutoff: Final = _as_utc(cutoff) return sum(1 for reset_at in reset_ats if isinstance(reset_at, datetime) and _as_utc(reset_at) > utc_cutoff) def _phase_is_drained(outcome: _ChunkOutcome) -> bool: """A short chunk means the due rows ran out. A full chunk that advanced nothing would be re-read unchanged forever, so it ends the phase too and those rows wait for the next tick.""" return outcome.fetched < RESET_BUDGET_JOB_BATCH_SIZE or outcome.advanced == 0 async def _run_phase_in_chunks(process_chunk: Callable[[], Awaitable[_ChunkOutcome]]) -> None: """Drive one reset phase a chunk at a time, capped so a single run cannot spin unbounded: leftovers are picked up by the next tick.""" for _ in range(RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN): if _phase_is_drained(await process_chunk()): return @dataclass(frozen=True, slots=True) class _LazyJson: """Serialize only if a log record is actually emitted. ``logger.debug("... %s", json.dumps(rows))`` evaluates the dump before the logger decides to drop the record, so a chunk of rows is serialized on the event loop on every tick at any log level. Passing this instead defers the work to the formatter. """ value: object def __str__(self) -> str: return json.dumps(self.value, indent=4, default=str) class _Lease(Enum): """Whether this pod may sweep, and whether it owes a lock release.""" LEADER = "leader" UNGUARDED = "unguarded" FOLLOWER = "follower" async def _write_key_windows(prisma_client: PrismaClient, row_id: str, payload: str) -> None: await VerificationTokenRepository(prisma_client).table.update( where={"token": row_id}, data={"budget_limits": payload}, ) async def _write_team_windows(prisma_client: PrismaClient, row_id: str, payload: str) -> None: await TeamRepository(prisma_client).table.update( where={"team_id": row_id}, data={"budget_limits": payload}, ) @dataclass(frozen=True, slots=True) class _WindowSource: """A table whose rows carry their own per-window budget limits.""" table: str id_column: str entity_type: Litellm_EntityType counter_prefix: str log_subject: str retry_subject: str write: Callable[[PrismaClient, str, str], Awaitable[None]] def page_query(self) -> str: """One keyset page, ordered by the primary key so the cursor never repeats a row. prisma-client-python cannot null-filter a ``Json?`` column (no DbNull / JsonNull sentinel, RobertCraigie/prisma-client-py#714), so the read stays raw SQL; the table and column names are module constants, never input. Writes still go through the ORM. """ return ( f'SELECT {self.id_column}, budget_limits FROM "{self.table}" ' f"WHERE budget_limits IS NOT NULL AND {self.id_column} > $1 " f"ORDER BY {self.id_column} LIMIT $2" ) _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = ( _WindowSource( table="LiteLLM_VerificationToken", id_column="token", entity_type=Litellm_EntityType.KEY, counter_prefix="spend:key", log_subject="keys", retry_subject="key", write=_write_key_windows, ), _WindowSource( table="LiteLLM_TeamTable", id_column="team_id", entity_type=Litellm_EntityType.TEAM, counter_prefix="spend:team", log_subject="teams", retry_subject="team", write=_write_team_windows, ), ) def _budget_cascade_event_metadata(cascade: _BudgetCascade) -> dict[str, object]: return { "num_budgets_found": len(cascade.budgets), "num_endusers_found": len(cascade.endusers), } class ResetBudgetJob: """ Resets the budget for all the keys, users, and teams that need it """ def __init__( self, proxy_logging_obj: ProxyLogging, prisma_client: PrismaClient, reset_settings: BudgetResetSettings | None = None, pod_lock_manager: PodLockManager | None = None, ): self.proxy_logging_obj: ProxyLogging = proxy_logging_obj self.prisma_client: PrismaClient = prisma_client self.reset_settings: BudgetResetSettings = reset_settings or get_budget_reset_settings() self.pod_lock_manager: PodLockManager | None = pod_lock_manager async def _lease_is_held(self, lock_manager: PodLockManager) -> bool: """True only when the lease is readable and someone holds it. An unreadable lock reports as unheld so the caller sweeps rather than skipping; being wrong here costs a duplicate sweep, and the alternative strands every expired budget at its cap. """ if lock_manager.redis_cache is None: return False try: lock_key: Final = lock_manager.get_redis_lock_key(RESET_BUDGET_JOB_NAME) return bool(await lock_manager.redis_cache.async_get_cache(lock_key)) except Exception as exc: # noqa: BLE001 # an unreadable lease must not strand the sweep verbose_proxy_logger.warning("Reset budget job: could not read the reset lease: %s", exc) return False async def _acquire_lease(self) -> _Lease: """Elect one sweeper per tick. Every pod schedules this job, and each one otherwise re-reads the whole due population and writes it back at the same calendar boundary, so a fleet multiplies one sweep's Postgres load by its replica count. A deployment with no Redis-backed lock manager runs unguarded, as it always has. """ lock_manager: Final = self.pod_lock_manager if lock_manager is None or lock_manager.redis_cache is None: return _Lease.UNGUARDED if await lock_manager.acquire_lock( cronjob_id=RESET_BUDGET_JOB_NAME, ttl=RESET_BUDGET_JOB_LOCK_TTL_SECONDS, ): return _Lease.LEADER if await self._lease_is_held(lock_manager): verbose_proxy_logger.debug("Reset budget job: another pod holds the reset lease, skipping this tick") return _Lease.FOLLOWER # acquire_lock reports contention and an unreachable Redis identically, so # treating a failed acquire as contention would skip the sweep on every pod # at once for as long as Redis is down. Sweeping unguarded costs duplicate # work; not sweeping leaves every expired budget pinned at its cap. verbose_proxy_logger.warning( "Reset budget job: could not take the reset lease and no other pod holds it, " "sweeping unguarded rather than skipping the tick" ) return _Lease.UNGUARDED async def reset_budget( self, ): """ Gets all the non-expired keys for a db, which need spend to be reset Resets their spend Updates db Runs on one pod per tick where a Redis lease is available. """ if self.prisma_client is None: return lease: Final = await self._acquire_lease() if lease is _Lease.FOLLOWER: return try: await self.reset_budget_for_litellm_keys() await self.reset_budget_for_litellm_users() await self.reset_budget_for_litellm_teams() await self.reset_budget_for_litellm_budget_table() await self.reset_budget_windows() finally: if lease is _Lease.LEADER and self.pod_lock_manager is not None: await self.pod_lock_manager.release_lock(cronjob_id=RESET_BUDGET_JOB_NAME) async def _with_db_retry(self, operation: Callable[[], Awaitable[_RowT]], *, reason: str) -> _RowT: """Reconnect and retry once on a transport error, so a dropped connection costs one retry instead of the whole tick. """ return await call_with_db_reconnect_retry(self.prisma_client, operation, reason=reason) async def _with_db_write_retry(self, operation: Callable[[], Awaitable[_RowT]], *, reason: str) -> _RowT: """Same, for writes: only replay when the statements provably never reached the database. A reset zeroes spend unconditionally, so replaying an ambiguous commit would erase spend accrued since it landed. """ return await call_with_db_reconnect_retry( self.prisma_client, operation, reason=reason, retry_safe_error_types=DB_RETRY_SAFE_ERROR_TYPES, ) @staticmethod async def _invalidate_spend_counter(counter_key: str, new_spend: float = 0.0) -> None: """Overwrite a spend counter with the post-reset value (0, or the carried overage when budget rollover is enabled) so a DB-row reset takes effect immediately. Call AFTER the DB write commits. Clearing Redis before the DB commit opens a window where get_current_spend reads 0 from Redis while the DB still holds the pre-reset value, allowing bypass. """ try: from litellm.proxy.proxy_server import spend_counter_cache spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=new_spend, ttl=60) if spend_counter_cache.redis_cache is not None: try: await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=new_spend, ttl=60) except Exception as redis_err: verbose_proxy_logger.warning( "Failed to reset spend counter %s in Redis: %s. " "Budget may be over-enforced until counter expires.", counter_key, redis_err, ) except Exception as e: verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e) @staticmethod async def _invalidate_global_proxy_spend_cache() -> None: """Drop the cached global-proxy spend accumulator after the proxy budget aggregate row is reset, so the next auth-time load reads the zeroed row instead of a stale (potentially never-expiring) counter. """ await ResetBudgetJob._invalidate_user_api_key_cache_entry(GLOBAL_PROXY_SPEND_CACHE_KEY) @staticmethod async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None: """Drop a stale management-cache entry so the next read fetches from DB. Tags and end-users are not reseeded by ``SpendCounterReseed.from_db``; for those, when the spend counter expires the budget check falls back to ``cached_obj.spend``. Keys, orgs, and team memberships are reseeded from the DB, but auth still may consult ``user_api_key_cache`` objects whose ``.spend`` field can lag a cross-pod DB reset. Deleting the cache entry forces the next auth-time fetch to reload the zeroed row from Postgres. """ try: from litellm.proxy.proxy_server import user_api_key_cache await user_api_key_cache.async_delete_cache(key=cache_key) except Exception as e: verbose_proxy_logger.warning( "Failed to invalidate user_api_key_cache entry %s: %s", cache_key, e, ) async def _fetch_linked_rows( self, table: SpendLinkedTable[_RowT], where: Mapping[str, object], log_subject: str, ) -> tuple[_RowT, ...]: """Read the rows the cascade will zero, so their counters can be invalidated once the transaction commits.""" try: return tuple( await self._with_db_retry( lambda: table.find_many(where=where), reason=f"reset_budget_read_{log_subject.replace(' ', '_')}_failure", ) ) except Exception as e: verbose_proxy_logger.warning("Failed to fetch %s for counter invalidation: %s", log_subject, e) return () async def _collect_endusers_to_reset(self, budget_ids: Sequence[str]) -> tuple[_EndUserRow, ...]: linked: Final[Sequence[_EndUserRow] | None] = await self._with_db_retry( lambda: self.prisma_client.get_data( table_name="enduser", query_type="find_all", budget_id_list=list(budget_ids), ), reason="reset_budget_read_endusers_failure", ) if litellm.max_end_user_budget_id is None or litellm.max_end_user_budget_id not in budget_ids: return tuple(linked or ()) return (*(linked or ()), *await self._get_endusers_with_no_budget_id()) async def _collect_budget_cascade(self, budgets_to_reset: Sequence[LiteLLM_BudgetTableFull]) -> _BudgetCascade: """Resolve every row the expiring budget tiers gate, before any write. Keys carrying their own budget_duration are left out: they run on their own schedule via reset_budget_for_litellm_keys(), so sweeping them here would reset them twice. """ budget_ids: Final = tuple(b.budget_id for b in budgets_to_reset if b.budget_id is not None) if not budget_ids: return _EMPTY_CASCADE team_memberships: Final[tuple[_TeamMembershipRow, ...]] = await self._fetch_linked_rows( table=TeamMembershipRepository(self.prisma_client).table, where=_budget_link_where(budget_ids), log_subject="team memberships", ) keys: Final[tuple[_KeyRow, ...]] = await self._fetch_linked_rows( table=VerificationTokenRepository(self.prisma_client).table, where=_budget_link_where(budget_ids, _LINKED_KEYS_WHERE), log_subject="keys", ) orgs: Final[tuple[_OrgRow, ...]] = await self._fetch_linked_rows( table=OrganizationRepository(self.prisma_client).table, where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), log_subject="orgs", ) tags: Final[tuple[_TagRow, ...]] = await self._fetch_linked_rows( table=TagRepository(self.prisma_client).table, where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), log_subject="tags", ) model_access_groups: Final[tuple[_ModelAccessGroupRow, ...]] = await self._fetch_linked_rows( table=ModelAccessGroupBudgetRepository(self.prisma_client).table, where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), log_subject="model access groups", ) rollover_caps: Final[Mapping[str, float]] = MappingProxyType( { # mutable-ok: MappingProxyType wraps a one-shot dict comprehension b.budget_id: cap for b in budgets_to_reset if b.budget_id is not None and (cap := _rollover_cap(b.max_budget)) is not None } if _rollover_enabled() else {} # mutable-ok: empty sentinel immediately frozen by MappingProxyType ) endusers: Final[tuple[_EndUserRow, ...]] = await self._collect_endusers_to_reset(budget_ids) return _BudgetCascade( budgets=tuple(budgets_to_reset), budget_ids=budget_ids, budget_resets=tuple( ( b.budget_id, compute_budget_reset_at(budget_duration=b.budget_duration, settings=self.reset_settings), ) for b in budgets_to_reset if b.budget_id is not None and b.budget_duration is not None ), endusers=endusers, counter_resets=( *( (_team_membership_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in team_memberships ), *((_key_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in keys), *((_org_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in orgs), *((_tag_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in tags), *( (_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in model_access_groups ), *((_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps)) for row in endusers), ), rollover_caps=rollover_caps, cache_keys=( *(key for row in team_memberships for key in _team_membership_cache_keys(row)), *(key for row in keys for key in _key_cache_keys(row)), *(key for row in orgs for key in _org_cache_keys(row)), *(key for row in tags for key in _tag_cache_keys(row)), *(key for row in model_access_groups for key in _model_access_group_cache_keys(row)), *(key for row in endusers for key in _enduser_cache_keys(row)), ), ) async def _commit_budget_cascade(self, cascade: _BudgetCascade) -> None: """Zero the gated spend and advance ``budget_reset_at`` in one transaction. Advancing the window on its own hides the tier from every later tick while its dependents stay pinned at the cap for the whole window; batching both means a mid-cascade failure persists nothing and the rows stay due for the next run. """ if not cascade.budget_ids: return await self._with_db_write_retry( lambda: self._commit_budget_cascade_once(cascade), reason="reset_budget_write_budget_cascade_failure", ) async def _commit_budget_cascade_once(self, cascade: _BudgetCascade) -> None: async with budget_cascade_unit_of_work(self.prisma_client.db.batch_) as uow: _queue_budget_linked_resets(uow.team_memberships, cascade) _queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE) _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) _queue_enduser_resets(uow.endusers, cascade) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) async def _invalidate_budget_cascade_caches(self, cascade: _BudgetCascade) -> None: for counter_key, new_spend in cascade.counter_resets: await self._invalidate_spend_counter(counter_key, new_spend=new_spend) for cache_key in cascade.cache_keys: await self._invalidate_user_api_key_cache_entry(cache_key) async def _reset_expired_budget_cascade(self) -> _BudgetCascadeCommitted | _BudgetCascadeFailed: now: Final = datetime.now(timezone.utc) try: budgets_to_reset: Final[Sequence[LiteLLM_BudgetTableFull] | None] = await self._with_db_retry( lambda: self.prisma_client.get_data( table_name="budget", query_type="find_all", reset_at=now, limit=RESET_BUDGET_JOB_BATCH_SIZE, ), reason="reset_budget_read_budgets_failure", ) cascade: Final = await self._collect_budget_cascade(budgets_to_reset or ()) except Exception as e: return _BudgetCascadeFailed(cascade=_EMPTY_CASCADE, error=e) try: await self._commit_budget_cascade(cascade) except Exception as e: return _BudgetCascadeFailed(cascade=cascade, error=e) await self._invalidate_budget_cascade_caches(cascade) return _BudgetCascadeCommitted( cascade=cascade, advanced=_count_advanced( (reset_at for _, reset_at in cascade.budget_resets), cutoff=datetime.now(timezone.utc), ), ) async def reset_budget_for_litellm_budget_table(self) -> None: """ Resets the spend a budget tier gates (end users, team members, keys, orgs, tags, model access groups) and advances the tier's budget_reset_at, atomically. Caches are invalidated only after the transaction commits, so a failed run cannot leave a zeroed counter in front of an un-reset DB row. """ await _run_phase_in_chunks(self._reset_budget_for_litellm_budget_table_chunk) async def _reset_budget_for_litellm_budget_table_chunk(self) -> _ChunkOutcome: start_time: Final = time.time() outcome: Final = await self._reset_expired_budget_cascade() end_time: Final = time.time() match outcome: case _BudgetCascadeCommitted(cascade=cascade, advanced=advanced): asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( service=ServiceTypes.RESET_BUDGET_JOB, duration=end_time - start_time, call_type="reset_budget_budget_table", start_time=start_time, end_time=end_time, event_metadata={ **_budget_cascade_event_metadata(cascade), "num_endusers_updated": len(cascade.endusers), "num_endusers_failed": 0, }, ) ) return _ChunkOutcome(fetched=len(cascade.budgets), advanced=advanced) case _BudgetCascadeFailed(cascade=cascade, error=error): verbose_proxy_logger.exception( "Failed to reset the budget table cascade (team member, enduser, org, tag and model access " "group spend, plus budget_reset_at); nothing was committed and the budgets stay due for the " "next run: %s", error, exc_info=error, ) asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_failure_hook( service=ServiceTypes.RESET_BUDGET_JOB, duration=end_time - start_time, error=error, call_type="reset_budget_endusers", start_time=start_time, end_time=end_time, event_metadata=_budget_cascade_event_metadata(cascade), ) ) return _NO_PROGRESS case _: assert_never(outcome) async def _get_endusers_with_no_budget_id( self, ) -> list[LiteLLM_EndUserTable]: """ Fetch end users that have no explicit budget_id set (NULL) and have accumulated spend > 0. These are implicitly-created end users that rely on the default budget (litellm.max_end_user_budget_id) applied in-memory during auth checks. """ table: Final = EndUserRepository(self.prisma_client).table rows: Final = await self._with_db_retry( lambda: table.find_many( where={ "budget_id": None, "spend": {"gt": 0}, }, ), reason="reset_budget_read_endusers_without_budget_id_failure", ) return [LiteLLM_EndUserTable.model_validate(row.model_dump()) for row in rows] async def _write_key_reset_updates(self, updated_keys: list[LiteLLM_VerificationToken]) -> None: """ Write per-row {spend, budget_reset_at} updates for keys. Avoids the batched full-model update path, which trips prisma.errors.DataError on any row carrying object_permission_id or budget_limits (see #27730). Both fields are rejected by Prisma's update input type for LiteLLM_VerificationToken, and the failure aborts the entire batch — silently leaving spend over the cap and budget_reset_at unchanged forever. """ await self._with_db_write_retry( lambda: self._write_key_reset_updates_once(updated_keys), reason="reset_budget_write_keys_failure", ) async def _write_key_reset_updates_once(self, updated_keys: list[LiteLLM_VerificationToken]) -> None: async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: for k in updated_keys: if k.token is None: continue uow.keys.queue_spend_reset( token=k.token, budget_reset_at=k.budget_reset_at, spend_decrement=k.max_budget if (k.spend or 0.0) > 0.0 else None, ) async def _write_user_reset_updates(self, updated_users: list[LiteLLM_UserTable]) -> None: """ Write per-row {spend, budget_reset_at} updates for users. Mirrors _write_key_reset_updates — avoids the full-model update path that trips Prisma's DataError on rows carrying unrecognised fields (see #27730). """ await self._with_db_write_retry( lambda: self._write_user_reset_updates_once(updated_users), reason="reset_budget_write_users_failure", ) async def _write_user_reset_updates_once(self, updated_users: list[LiteLLM_UserTable]) -> None: async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: for u in updated_users: uow.users.queue_spend_reset( user_id=u.user_id, budget_reset_at=u.budget_reset_at, spend_decrement=u.max_budget if (u.spend or 0.0) > 0.0 else None, ) async def _write_team_reset_updates(self, updated_teams: list[LiteLLM_TeamTable]) -> None: """ Write per-row {spend, budget_reset_at} updates for teams. Mirrors _write_key_reset_updates — avoids the full-model update path that trips Prisma's DataError on rows carrying unrecognised fields (see #27730). """ await self._with_db_write_retry( lambda: self._write_team_reset_updates_once(updated_teams), reason="reset_budget_write_teams_failure", ) async def _write_team_reset_updates_once(self, updated_teams: list[LiteLLM_TeamTable]) -> None: async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: for t in updated_teams: uow.teams.queue_spend_reset( team_id=t.team_id, budget_reset_at=t.budget_reset_at, spend_decrement=t.max_budget if (t.spend or 0.0) > 0.0 else None, ) def _emit_phase_failure( self, call_type: str, error: Exception, start_time: float, end_time: float, event_metadata: dict[str, object], ) -> None: """Report rows that could not be reset without failing the chunk: the rows that did reset are already committed, and raising here would cost the phase every remaining chunk this tick. """ verbose_proxy_logger.error("%s: %s", call_type, error) asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_failure_hook( service=ServiceTypes.RESET_BUDGET_JOB, duration=end_time - start_time, error=error, call_type=call_type, start_time=start_time, end_time=end_time, event_metadata=event_metadata, ) ) async def reset_budget_for_litellm_keys(self) -> None: """ Resets the budget for all the litellm keys Catches Exceptions and logs them """ await _run_phase_in_chunks(self._reset_budget_for_litellm_keys_chunk) async def _reset_budget_for_litellm_keys_chunk(self) -> _ChunkOutcome: now: Final = datetime.utcnow() start_time: Final = time.time() keys_to_reset: list[LiteLLM_VerificationToken] | None = None try: keys_to_reset = await self._with_db_retry( lambda: self.prisma_client.get_data( table_name="key", query_type="find_all", expires=now, reset_at=now, limit=RESET_BUDGET_JOB_BATCH_SIZE, ), reason="reset_budget_read_keys_failure", ) verbose_proxy_logger.debug("Keys to reset %s", _LazyJson(keys_to_reset)) updated_keys: Final[list[LiteLLM_VerificationToken]] = [] failed_keys: Final = [] if keys_to_reset is not None and len(keys_to_reset) > 0: for key in keys_to_reset: try: updated_key = await ResetBudgetJob._reset_budget_for_key( key=key, current_time=now, reset_settings=self.reset_settings, ) if updated_key is not None: updated_keys.append(updated_key) else: failed_keys.append({"key": key, "error": "Returned None without exception"}) except Exception as e: failed_keys.append({"key": key, "error": str(e)}) verbose_proxy_logger.exception("Failed to reset budget for key: %s", key) verbose_proxy_logger.debug("Updated keys %s", _LazyJson(updated_keys)) if updated_keys: await self._write_key_reset_updates(updated_keys=updated_keys) for k in updated_keys: token = getattr(k, "token", None) if token: await self._invalidate_spend_counter(f"spend:key:{token}", new_spend=k.spend or 0.0) end_time = time.time() outcome: Final = _ChunkOutcome( fetched=len(keys_to_reset) if keys_to_reset else 0, advanced=_count_advanced( (k.budget_reset_at for k in updated_keys), cutoff=datetime.now(timezone.utc), ), ) if len(failed_keys) > 0: self._emit_phase_failure( call_type="reset_budget_keys", error=Exception(f"Failed to reset {len(failed_keys)} keys: {json.dumps(failed_keys, default=str)}"), start_time=start_time, end_time=end_time, event_metadata={ "num_keys_found": len(keys_to_reset) if keys_to_reset else 0, }, ) return outcome asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( service=ServiceTypes.RESET_BUDGET_JOB, duration=end_time - start_time, call_type="reset_budget_keys", start_time=start_time, end_time=end_time, event_metadata={ "num_keys_found": len(keys_to_reset) if keys_to_reset else 0, "num_keys_updated": len(updated_keys), "num_keys_failed": len(failed_keys), }, ) ) except Exception as e: end_time = time.time() asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_failure_hook( service=ServiceTypes.RESET_BUDGET_JOB, duration=end_time - start_time, error=e, call_type="reset_budget_keys", start_time=start_time, end_time=end_time, event_metadata={ "num_keys_found": len(keys_to_reset) if keys_to_reset else 0, }, ) ) verbose_proxy_logger.exception("Failed to reset budget for keys: %s", e) return _NO_PROGRESS else: return outcome async def reset_budget_for_litellm_users(self) -> None: """ Resets the budget for all LiteLLM Internal Users if their budget has expired """ await _run_phase_in_chunks(self._reset_budget_for_litellm_users_chunk) async def _reset_budget_for_litellm_users_chunk(self) -> _ChunkOutcome: now: Final = datetime.utcnow() start_time: Final = time.time() users_to_reset: list[LiteLLM_UserTable] | None = None try: users_to_reset = await self._with_db_retry( lambda: self.prisma_client.get_data( table_name="user", query_type="find_all", reset_at=now, limit=RESET_BUDGET_JOB_BATCH_SIZE, ), reason="reset_budget_read_users_failure", ) updated_users: Final[list[LiteLLM_UserTable]] = [] failed_users: Final = [] if users_to_reset is not None and len(users_to_reset) > 0: for user in users_to_reset: try: updated_user = await ResetBudgetJob._reset_budget_for_user( user=user, current_time=now, reset_settings=self.reset_settings, ) if updated_user is not None: updated_users.append(updated_user) else: failed_users.append( { "user": user, "error": "Returned None without exception", } ) except Exception as e: failed_users.append({"user": user, "error": str(e)}) verbose_proxy_logger.exception("Failed to reset budget for user: %s", user) verbose_proxy_logger.debug("Updated users %s", _LazyJson(updated_users)) if updated_users: await self._write_user_reset_updates(updated_users=updated_users) for u in updated_users: user_id = getattr(u, "user_id", None) if user_id: await self._invalidate_spend_counter(f"spend:user:{user_id}", new_spend=u.spend or 0.0) if user_id == LITELLM_PROXY_BUDGET_NAME: await self._invalidate_global_proxy_spend_cache() end_time = time.time() outcome: Final = _ChunkOutcome( fetched=len(users_to_reset) if users_to_reset else 0, advanced=_count_advanced( (u.budget_reset_at for u in updated_users), cutoff=datetime.now(timezone.utc), ), ) if len(failed_users) > 0: self._emit_phase_failure( call_type="reset_budget_users", error=Exception( f"Failed to reset {len(failed_users)} users: {json.dumps(failed_users, default=str)}" ), start_time=start_time, end_time=end_time, event_metadata={ "num_users_found": len(users_to_reset) if users_to_reset else 0, }, ) return outcome asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( service=ServiceTypes.RESET_BUDGET_JOB, duration=end_time - start_time, call_type="reset_budget_users", start_time=start_time, end_time=end_time, event_metadata={ "num_users_found": len(users_to_reset) if users_to_reset else 0, "num_users_updated": len(updated_users), "num_users_failed": len(failed_users), }, ) ) except Exception as e: end_time = time.time() asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_failure_hook( service=ServiceTypes.RESET_BUDGET_JOB, duration=end_time - start_time, error=e, call_type="reset_budget_users", start_time=start_time, end_time=end_time, event_metadata={ "num_users_found": len(users_to_reset) if users_to_reset else 0, }, ) ) verbose_proxy_logger.exception("Failed to reset budget for users: %s", e) return _NO_PROGRESS else: return outcome async def reset_budget_for_litellm_teams(self) -> None: """ Resets the budget for all LiteLLM Internal Teams if their budget has expired """ await _run_phase_in_chunks(self._reset_budget_for_litellm_teams_chunk) async def _reset_budget_for_litellm_teams_chunk(self) -> _ChunkOutcome: now: Final = datetime.utcnow() start_time: Final = time.time() teams_to_reset: list[LiteLLM_TeamTable] | None = None try: teams_to_reset = await self._with_db_retry( lambda: self.prisma_client.get_data( table_name="team", query_type="find_all", reset_at=now, limit=RESET_BUDGET_JOB_BATCH_SIZE, ), reason="reset_budget_read_teams_failure", ) updated_teams: Final[list[LiteLLM_TeamTable]] = [] failed_teams: Final = [] if teams_to_reset is not None and len(teams_to_reset) > 0: for team in teams_to_reset: try: updated_team = await ResetBudgetJob._reset_budget_for_team( team=team, current_time=now, reset_settings=self.reset_settings, ) if updated_team is not None: updated_teams.append(updated_team) else: failed_teams.append( { "team": team, "error": "Returned None without exception", } ) except Exception as e: failed_teams.append({"team": team, "error": str(e)}) verbose_proxy_logger.exception("Failed to reset budget for team: %s", team) verbose_proxy_logger.debug("Updated teams %s", _LazyJson(updated_teams)) if updated_teams: await self._write_team_reset_updates(updated_teams=updated_teams) for t in updated_teams: team_id = getattr(t, "team_id", None) if team_id: await self._invalidate_spend_counter(f"spend:team:{team_id}", new_spend=t.spend or 0.0) end_time = time.time() outcome: Final = _ChunkOutcome( fetched=len(teams_to_reset) if teams_to_reset else 0, advanced=_count_advanced( (t.budget_reset_at for t in updated_teams), cutoff=datetime.now(timezone.utc), ), ) if len(failed_teams) > 0: self._emit_phase_failure( call_type="reset_budget_teams", error=Exception( f"Failed to reset {len(failed_teams)} teams: {json.dumps(failed_teams, default=str)}" ), start_time=start_time, end_time=end_time, event_metadata={ "num_teams_found": len(teams_to_reset) if teams_to_reset else 0, }, ) return outcome asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( service=ServiceTypes.RESET_BUDGET_JOB, duration=end_time - start_time, call_type="reset_budget_teams", start_time=start_time, end_time=end_time, event_metadata={ "num_teams_found": len(teams_to_reset) if teams_to_reset else 0, "num_teams_updated": len(updated_teams), "num_teams_failed": len(failed_teams), }, ) ) except Exception as e: end_time = time.time() asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_failure_hook( service=ServiceTypes.RESET_BUDGET_JOB, duration=end_time - start_time, error=e, call_type="reset_budget_teams", start_time=start_time, end_time=end_time, event_metadata={ "num_teams_found": len(teams_to_reset) if teams_to_reset else 0, }, ) ) verbose_proxy_logger.exception("Failed to reset budget for teams: %s", e) return _NO_PROGRESS else: return outcome @staticmethod async def _reset_expired_window( window: dict, counter_key: str, spend_counter_cache: DualCache, now: datetime, reset_settings: BudgetResetSettings, prisma_client: PrismaClient, entity_type: Litellm_EntityType, entity_id: str, ) -> bool: """Reset a single budget window if expired. Returns True if the window was reset.""" reset_at_str: Final = window.get("reset_at") if not reset_at_str: return False reset_at: Final = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")).replace(tzinfo=None) if reset_at > now: return False new_value: Final = await ResetBudgetJob._window_carried_spend(window, counter_key, spend_counter_cache) spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=new_value) if spend_counter_cache.redis_cache is not None: try: await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=new_value) except Exception as redis_err: verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err) budget_duration: Final = window["budget_duration"] next_reset_at: Final = compute_budget_reset_at(budget_duration=budget_duration, settings=reset_settings) window["reset_at"] = next_reset_at.isoformat() await ResetBudgetJob._roll_window_spend_row( prisma_client=prisma_client, entity_type=entity_type, entity_id=entity_id, budget_duration=budget_duration, next_reset_at=next_reset_at, ) return True @staticmethod async def _roll_window_spend_row( prisma_client: PrismaClient, entity_type: Litellm_EntityType, entity_id: str, budget_duration: str, next_reset_at: datetime, ) -> None: """Move this window's LiteLLM_BudgetWindowSpend row onto the window that just started, so the maintained total the read path uses starts from zero alongside the counter. Best effort: the row is an optimization over aggregating LiteLLM_SpendLogs, so a failure here must not stop the remaining windows from having their counters reset. """ try: window_start: Final = next_reset_at - timedelta(seconds=duration_in_seconds(budget_duration)) except Exception as e: # noqa: BLE001 # duration_in_seconds raises bare exceptions on bad input verbose_proxy_logger.warning("Unparseable budget_duration %s: %s", budget_duration, e) return try: await roll_window_spend_row( prisma_client=prisma_client, entity_type=entity_type.value, entity_id=entity_id, window_duration=budget_duration, new_window_start=window_start, ) except Exception as e: # noqa: BLE001 # the row is best effort; counter resets must still land verbose_proxy_logger.warning( "Failed to roll budget window spend row for %s=%s window=%s: %s", entity_type.value, entity_id, budget_duration, e, ) @staticmethod async def _window_carried_spend( window: Mapping[str, object], counter_key: str, spend_counter_cache: DualCache ) -> float: """Per-window spend lives only in the counter, so the carried overage is read from it before the reset overwrites it.""" if not _rollover_enabled(): return 0.0 window_max: Final = window.get("max_budget") cap: Final = _rollover_cap(window_max) if isinstance(window_max, (int, float)) else None if cap is None: return 0.0 try: current: Final = await spend_counter_cache.async_get_cache(key=counter_key) except Exception as e: # noqa: BLE001 # an unreadable counter falls back to a plain zero reset verbose_proxy_logger.warning("Failed to read spend counter %s for rollover: %s", counter_key, e) return 0.0 if not isinstance(current, (int, float)): return 0.0 return _carried_spend(float(current), cap) async def reset_budget_windows(self) -> None: """ For keys and teams with budget_limits, reset any individual windows where reset_at <= now. Only the expired windows are reset; other windows are untouched. """ from litellm.proxy.proxy_server import spend_counter_cache now: Final = datetime.utcnow() for source in _WINDOW_SOURCES: try: await self._reset_windows_for(source=source, now=now, spend_counter_cache=spend_counter_cache) except Exception as e: verbose_proxy_logger.exception("Failed to reset budget windows for %s: %s", source.log_subject, e) async def _reset_windows_for( self, source: _WindowSource, now: datetime, spend_counter_cache: DualCache, ) -> None: """Walk one table's windowed rows a page at a time, to the end. Paging is what bounds the memory: the previous form pulled every row carrying budget_limits into one result set on every tick, which grows with the deployment's key count and is paid on the event loop. The walk deliberately has no per-run page cap. A cap has to remember where it stopped, and that position cannot live in the process: the lease is released after each sweep, so the next tick can elect a different pod whose own position is unset. It would restart at the first row and never reach the tail, pinning those windows at their cap for good. The cursor strictly advances, so the walk terminates on its own without needing a bound. """ cursor = "" while True: next_cursor = await self._reset_window_page( source=source, cursor=cursor, now=now, spend_counter_cache=spend_counter_cache, ) if next_cursor is None: return cursor = next_cursor async def _reset_window_page( self, source: _WindowSource, cursor: str, now: datetime, spend_counter_cache: DualCache, ) -> str | None: """Reset one page of windows; return the next cursor, or None when drained.""" rows: Final = await self._with_db_retry( lambda: self.prisma_client.db.query_raw(source.page_query(), cursor, RESET_BUDGET_JOB_BATCH_SIZE), reason=f"reset_budget_read_{source.retry_subject}_windows_failure", ) for row in rows: raw = row["budget_limits"] if not raw: continue row_id: str = row[source.id_column] windows: list[dict[str, object]] = raw if isinstance(raw, list) else json.loads(raw) changed = False for window in windows: counter_key = f"{source.counter_prefix}:{row_id}:window:{window['budget_duration']}" if await ResetBudgetJob._reset_expired_window( window, counter_key, spend_counter_cache, now, self.reset_settings, prisma_client=self.prisma_client, entity_type=source.entity_type, entity_id=row_id, ): changed = True if changed: await self._with_db_write_retry( lambda: source.write(self.prisma_client, row_id, json.dumps(windows)), reason=f"reset_budget_write_{source.retry_subject}_windows_failure", ) if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE: return None return rows[-1][source.id_column] @staticmethod async def _reset_budget_common( item: LiteLLM_TeamTable | LiteLLM_UserTable | LiteLLM_VerificationToken, current_time: datetime, item_type: Literal["key", "team", "user"], reset_settings: BudgetResetSettings, ): """ In-place, updates spend=0, and sets budget_reset_at to current_time + budget_duration Common logic for resetting budget for a team, user, or key. Spend-counter invalidation happens in the caller, AFTER the DB write commits. Zeroing the counter here would open a bypass window when the DB write fails: get_current_spend reads 0 from Redis while the DB still holds the pre-reset value, admitting requests past the cap. """ try: item.spend = _carried_spend(item.spend, _rollover_cap(item.max_budget)) if _rollover_enabled() else 0.0 if hasattr(item, "budget_duration") and item.budget_duration is not None: item.budget_reset_at = compute_budget_reset_at( budget_duration=item.budget_duration, settings=reset_settings ) return item except Exception as e: verbose_proxy_logger.exception("Error resetting budget for %s: %s. Item: %s", item_type, e, item) raise e @staticmethod async def _reset_budget_for_team( team: LiteLLM_TeamTable, current_time: datetime, reset_settings: BudgetResetSettings, ) -> LiteLLM_TeamTable | None: await ResetBudgetJob._reset_budget_common( item=team, current_time=current_time, item_type="team", reset_settings=reset_settings, ) return team @staticmethod async def _reset_budget_for_user( user: LiteLLM_UserTable, current_time: datetime, reset_settings: BudgetResetSettings, ) -> LiteLLM_UserTable | None: await ResetBudgetJob._reset_budget_common( item=user, current_time=current_time, item_type="user", reset_settings=reset_settings, ) return user @staticmethod async def _reset_budget_for_key( key: LiteLLM_VerificationToken, current_time: datetime, reset_settings: BudgetResetSettings, ) -> LiteLLM_VerificationToken | None: await ResetBudgetJob._reset_budget_common( item=key, current_time=current_time, item_type="key", reset_settings=reset_settings, ) return key