diff --git a/litellm/constants.py b/litellm/constants.py index c292cb53e29..cc6db6c10cc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -289,6 +289,7 @@ REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_spend_update_buffer" REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer" REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer" +REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth LITELLM_ASYNCIO_QUEUE_MAXSIZE: Final = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f6345eae06d..5703c6cd5e8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4826,6 +4826,7 @@ async def _virtual_key_multi_budget_check( max_budget=w["max_budget"], window_entity_type="Key", window_entity_id=valid_token.token, + window_duration=str(w["budget_duration"]), window_start=get_budget_window_start(w), ) if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: @@ -5199,6 +5200,7 @@ async def _team_multi_budget_check( max_budget=w["max_budget"], window_entity_type="Team", window_entity_id=team_object.team_id, + window_duration=str(w["budget_duration"]), window_start=get_budget_window_start(w), ) if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 55539ded02b..1682cf12f4e 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -4,7 +4,7 @@ import math import time from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass, field -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from enum import Enum from types import MappingProxyType from typing import Final, Literal, Protocol, TypeVar, assert_never @@ -20,10 +20,12 @@ from litellm.constants import ( 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, @@ -38,6 +40,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( 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 @@ -347,6 +350,7 @@ class _WindowSource: table: str id_column: str + entity_type: Litellm_EntityType counter_prefix: str log_subject: str retry_subject: str @@ -371,6 +375,7 @@ _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", @@ -379,6 +384,7 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = ( _WindowSource( table="LiteLLM_TeamTable", id_column="team_id", + entity_type=Litellm_EntityType.TEAM, counter_prefix="spend:team", log_subject="teams", retry_subject="team", @@ -1241,6 +1247,9 @@ class ResetBudgetJob: 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") @@ -1256,11 +1265,56 @@ class ResetBudgetJob: 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) - window["reset_at"] = compute_budget_reset_at( - budget_duration=window["budget_duration"], settings=reset_settings - ).isoformat() + 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 @@ -1356,6 +1410,9 @@ class ResetBudgetJob: 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: diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py new file mode 100644 index 00000000000..f9188f95cfd --- /dev/null +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -0,0 +1,313 @@ +""" +Writer for LiteLLM_BudgetWindowSpend. + +The table holds one row per configured budget window whose window_start rolls +forward in place, so budget enforcement can read a maintained running total +instead of aggregating LiteLLM_SpendLogs every time a window counter goes cold +(issue #35766). Raw SQL rather than the Prisma upsert helper because the +conditional roll cannot be expressed through the query builder. + +Seeding a row that does not exist yet reads LiteLLM_SpendLogs once, excluding +the requests whose increments are in the same batch so neither source counts +them twice. One gap survives that exclusion: without the Redis transaction +buffer every pod flushes its own increments, so a row seeded by one pod can +include spend logs whose increments are still queued on another pod, and those +increments are added again when that pod flushes. That is bounded by a single +flush interval, happens at most once per window row, and only ever over-counts: +the seed never omits spend, because every increment not yet in the row still +reaches it on its own pod's next flush. A row therefore lags real spend by at +most one flush interval of queued increments, the same lag the SpendLogs +aggregate it replaces (and every other spend column) already has. +""" + +from collections.abc import Sequence +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Final, Protocol + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendTransaction, + to_naive_utc, + window_spend_group_key, +) + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +_SELECT_EXISTING_ROWS_SQL: Final = ( + 'SELECT entity_type, entity_id, window_duration FROM "LiteLLM_BudgetWindowSpend" ' + "WHERE (entity_type, entity_id, window_duration) " + "IN (SELECT * FROM unnest($1::text[], $2::text[], $3::text[]))" +) + +_UPSERT_WINDOW_SPEND_SQL: Final = ( + 'INSERT INTO "LiteLLM_BudgetWindowSpend" ' + "(entity_type, entity_id, window_duration, window_start, spend, created_at, updated_at) " + "VALUES ($1, $2, $3, ($4::timestamptz AT TIME ZONE 'UTC'), $5, " + "($7::timestamptz AT TIME ZONE 'UTC'), ($7::timestamptz AT TIME ZONE 'UTC')) " + "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE SET " + "spend = CASE " + 'WHEN "LiteLLM_BudgetWindowSpend".window_start >= EXCLUDED.window_start ' + 'THEN "LiteLLM_BudgetWindowSpend".spend + $6 ' + "ELSE EXCLUDED.spend " + "END, " + 'window_start = GREATEST("LiteLLM_BudgetWindowSpend".window_start, EXCLUDED.window_start), ' + "updated_at = ($7::timestamptz AT TIME ZONE 'UTC')" +) + +_ROLL_WINDOW_SPEND_SQL: Final = ( + 'UPDATE "LiteLLM_BudgetWindowSpend" SET ' + "window_start = ($4::timestamptz AT TIME ZONE 'UTC'), " + "spend = 0, " + "updated_at = ($5::timestamptz AT TIME ZONE 'UTC') " + "WHERE entity_type = $1 AND entity_id = $2 AND window_duration = $3 " + "AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')" +) + +_SEED_FROM_SPEND_LOGS_KEY_SQL: Final = ( + 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " + "AND NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" +) + +_SEED_FROM_SPEND_LOGS_TEAM_SQL: Final = ( + 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " + "AND NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" +) + +_SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL: Final = ( + 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" +) + +_SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL: Final = ( + 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" +) + +_UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60) + + +class WindowSpendLogsAggregate(Protocol): + """Sums LiteLLM_SpendLogs for one entity since window_start, ignoring the + requests whose ids are handed in. + + Injected so the flush can be exercised without a database and so the + expensive aggregate stays swappable. + """ + + async def __call__( + self, + prisma_client: "PrismaClient", + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Sequence[str], + exclude_started_at: datetime | None, + ) -> float | None: ... + + +async def spend_logs_total_excluding( + prisma_client: "PrismaClient", + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Sequence[str], + exclude_started_at: datetime | None, +) -> float | None: + """LiteLLM_SpendLogs spend for one entity since window_start, minus the + requests already accounted for by the increments being flushed. + + The spend log writer drains its own queue on a ~2s poll whenever anything + is queued, while window increments flush on the much slower batch tick, so + by the time a window row is seeded its batch's log rows are normally + already in the table. Counting them in the seed and again in the increment + is what made a fresh row land at twice the true spend. + + The exclusion is bounded to rows that started at or after the batch's + earliest request. request_id can be chosen by the client + (x-litellm-call-id), so an unbounded exclusion would let a replayed old id + erase a historical row from the seed while its increment still lands. + Without a known start the batch's ids are not excluded at all: that can + only over-count once, which enforcement tolerates, whereas under-counting + is a budget bypass. + """ + if entity_type == Litellm_EntityType.KEY.value: + bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_KEY_SQL, _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL + elif entity_type == Litellm_EntityType.TEAM.value: + bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_TEAM_SQL, _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL + else: + return None + rows: Final = ( + await prisma_client.db.query_raw(unbounded_sql, entity_id, window_start) + if exclude_started_at is None or not exclude_request_ids + else await prisma_client.db.query_raw( + bounded_sql, + entity_id, + window_start, + tuple(exclude_request_ids), + _exclusion_lower_bound(exclude_started_at), + ) + ) + if not rows: + return 0.0 + return float(rows[0].get("total") or 0.0) + + +def _exclusion_lower_bound(started_at: datetime) -> datetime: + """LiteLLM_SpendLogs.startTime is TIMESTAMP(3); floor to the second so a + millisecond rounding of the batch's own earliest row cannot slip under it.""" + return to_naive_utc(started_at).replace(microsecond=0) + + +def _primary_key(transaction: WindowSpendTransaction) -> tuple[str, str, str]: + return ( + transaction["entity_type"], + transaction["entity_id"], + transaction["window_duration"], + ) + + +async def _existing_primary_keys( + prisma_client: "PrismaClient", + transactions: tuple[WindowSpendTransaction, ...], +) -> frozenset[tuple[str, str, str]]: + rows: Final = await prisma_client.db.query_raw( + _SELECT_EXISTING_ROWS_SQL, + tuple(transaction["entity_type"] for transaction in transactions), + tuple(transaction["entity_id"] for transaction in transactions), + tuple(transaction["window_duration"] for transaction in transactions), + ) + return frozenset((row["entity_type"], row["entity_id"], row["window_duration"]) for row in rows or ()) + + +async def _seed_base_for_missing_row( + prisma_client: "PrismaClient", + transaction: WindowSpendTransaction, + existing_primary_keys: frozenset[tuple[str, str, str]], + spend_logs_aggregate: WindowSpendLogsAggregate, +) -> float: + """Spend already recorded for a window that has no row yet. + + This is the LiteLLM_SpendLogs aggregate the window counter reseed runs on + every cold counter today, but here it runs once per window lifetime and off + the request path, and it excludes this batch's own requests so they are + counted by their increments alone. + """ + if _primary_key(transaction) in existing_primary_keys: + return 0.0 + base: Final = await spend_logs_aggregate( + prisma_client=prisma_client, + entity_type=transaction["entity_type"], + entity_id=transaction["entity_id"], + window_start=datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc), + exclude_request_ids=transaction["request_ids"], + exclude_started_at=_transaction_started_at(transaction), + ) + return float(base or 0.0) + + +def _transaction_started_at(transaction: WindowSpendTransaction) -> datetime | None: + started_at: Final = transaction.get("started_at") + if started_at is None: + return None + return datetime.fromisoformat(started_at).replace(tzinfo=timezone.utc) + + +def _upsert_params( + transaction: WindowSpendTransaction, + seed_base: float, + now: datetime, +) -> tuple[str, str, str, datetime, float, float, datetime]: + """$5 is what a brand new row starts at (pre-existing spend plus this + increment); $6 is the increment alone, which is all an already-current row + may add. They are equal for every row that already existed, so a row is + never seeded twice when two pods flush the same new window.""" + increment: Final = float(transaction["spend"]) + return ( + transaction["entity_type"], + transaction["entity_id"], + transaction["window_duration"], + datetime.fromisoformat(transaction["window_start"]), + seed_base + increment, + increment, + now, + ) + + +async def commit_window_spend_updates( + prisma_client: "PrismaClient", + transactions: Sequence[WindowSpendTransaction], + spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_total_excluding, +) -> None: + """Apply aggregated window increments to LiteLLM_BudgetWindowSpend. + + An increment at or behind the row's window_start adds into the row (this is + how in-flight requests that raced a reset carry into the new window); an + increment ahead of it rolls the window and starts from that increment. + + Statements are ordered by primary key so concurrent pods take row locks in + the same order, with window_start breaking ties so an older window is + applied before the roll that supersedes it. + """ + if not transactions: + return + + ordered: Final = tuple(sorted(transactions, key=window_spend_group_key)) + existing_primary_keys: Final = await _existing_primary_keys( + prisma_client=prisma_client, + transactions=ordered, + ) + seed_bases: Final = tuple( + [ + await _seed_base_for_missing_row( + prisma_client=prisma_client, + transaction=transaction, + existing_primary_keys=existing_primary_keys, + spend_logs_aggregate=spend_logs_aggregate, + ) + for transaction in ordered + ] + ) + + now: Final = to_naive_utc(datetime.now(timezone.utc)) + verbose_proxy_logger.debug( + "Spend tracking - committing %d budget window spend upserts over %d existing rows", + len(ordered), + len(existing_primary_keys), + ) + async with ( + prisma_client.db.tx(timeout=_UPSERT_TRANSACTION_TIMEOUT) as db_transaction, + db_transaction.batch_() as batcher, + ): + for transaction, seed_base in zip(ordered, seed_bases): + batcher.execute_raw( + _UPSERT_WINDOW_SPEND_SQL, + *_upsert_params(transaction=transaction, seed_base=seed_base, now=now), + ) + + +async def roll_window_spend_row( + prisma_client: "PrismaClient", + entity_type: str, + entity_id: str, + window_duration: str, + new_window_start: datetime, +) -> None: + """Move a row onto the window that just started and zero its spend. + + Conditional on the stored window_start still being behind the new one so a + pod that already rolled the row (or increments that arrived under the new + window) are not clobbered. + """ + await prisma_client.db.execute_raw( + _ROLL_WINDOW_SPEND_SQL, + entity_type, + entity_id, + window_duration, + to_naive_utc(new_window_start), + to_naive_utc(datetime.now(timezone.utc)), + ) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index dc18a57c10c..641c07914d9 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -56,6 +56,10 @@ from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdate from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( ToolDiscoveryQueue, ) +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendTransaction, + WindowSpendUpdateQueue, +) from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING from litellm.proxy.spend_tracking.compression_savings import ( extract_compression_saved_tokens, @@ -195,6 +199,7 @@ class DBSpendUpdateWriter: self.daily_agent_spend_update_queue = DailySpendUpdateQueue() self.daily_org_spend_update_queue = DailySpendUpdateQueue() self.daily_tag_spend_update_queue = DailySpendUpdateQueue() + self.window_spend_update_queue = WindowSpendUpdateQueue() async def update_database( # LiteLLM management object fields @@ -210,7 +215,11 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, - ): + ) -> str | None: + """Returns the LiteLLM_SpendLogs request_id this call was recorded + under, so the caller can tell the budget-window writer which log rows + its increments already cover. None when the payload could not be built. + """ from litellm.proxy.proxy_server import ( disable_spend_logs, litellm_proxy_budget_name, @@ -227,7 +236,7 @@ class DBSpendUpdateWriter: team_id, ) if ProxyUpdateSpend.disable_spend_updates() is True: - return + return None if token is not None and isinstance(token, str) and token.startswith("sk-"): hashed_token = hash_token(token=token) else: @@ -301,6 +310,7 @@ class DBSpendUpdateWriter: ) verbose_proxy_logger.debug("Runs spend update on all tables") + return payload.get("request_id") except Exception: spend_log_error( "Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue " @@ -313,6 +323,7 @@ class DBSpendUpdateWriter: org_id, end_user_id, ) + return None async def _enqueue_tool_usage_transaction( self, @@ -999,6 +1010,7 @@ class DBSpendUpdateWriter: daily_org_spend_update_queue=self.daily_org_spend_update_queue, daily_end_user_spend_update_queue=self.daily_end_user_spend_update_queue, daily_agent_spend_update_queue=self.daily_agent_spend_update_queue, + window_spend_update_queue=self.window_spend_update_queue, ) # Only commit from redis to db if this pod is the leader @@ -1017,6 +1029,7 @@ class DBSpendUpdateWriter: daily_org_spend_update_transactions, daily_end_user_spend_update_transactions, daily_agent_spend_update_transactions, + window_spend_update_transactions, ) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() uncommitted = { # mutable-ok: drives which popped categories still need re-queuing @@ -1026,6 +1039,7 @@ class DBSpendUpdateWriter: "daily_org_spend_update_transactions": daily_org_spend_update_transactions, "daily_end_user_spend_update_transactions": daily_end_user_spend_update_transactions, "daily_agent_spend_update_transactions": daily_agent_spend_update_transactions, + "window_spend_update_transactions": window_spend_update_transactions, } if db_spend_update_transactions is not None: @@ -1095,6 +1109,12 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_agent_spend_update_transactions, ) uncommitted.pop("daily_agent_spend_update_transactions", None) + if window_spend_update_transactions is not None: + await DBSpendUpdateWriter._commit_window_spend_updates( + prisma_client=prisma_client, + window_spend_transactions=window_spend_update_transactions, + ) + uncommitted.pop("window_spend_update_transactions", None) except Exception as e: spend_log_error( "Spend tracking - failed to commit spend updates from Redis to DB. " @@ -1210,6 +1230,27 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_agent_spend_update_transactions, ) + ################## Budget Window Spend Update Transactions ################## + # Aggregate all in memory budget window spend transactions and commit to db + window_spend_update_transactions: Final = ( + await self.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() + ) + + try: + await DBSpendUpdateWriter._commit_window_spend_updates( + prisma_client=prisma_client, + window_spend_transactions=window_spend_update_transactions, + ) + except Exception as e: # noqa: BLE001 # the increments go back on the queue; the rest of the flush must run + spend_log_error( + "Spend tracking - failed to commit budget window spend updates. " + "Re-queued %d window increments for retry on next tick. Error: %s", + len(window_spend_update_transactions), + str(e), + exc=e, + ) + await self.window_spend_update_queue.update_queue.put(window_spend_update_transactions) + ################## Tool Registry Upserts ################## await self._flush_tool_discovery_queue(prisma_client=prisma_client) @@ -1274,6 +1315,28 @@ class DBSpendUpdateWriter: cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, ) + @staticmethod + async def _commit_window_spend_updates( + prisma_client: PrismaClient, + window_spend_transactions: Sequence[WindowSpendTransaction], + ) -> None: + """ + Commit per-budget-window spend increments to LiteLLM_BudgetWindowSpend. + + Raises on failure so the caller re-queues the increments: budget + enforcement trusts a current row without reconciling it against + LiteLLM_SpendLogs, so a dropped increment would let the entity spend + past its window limit after the next counter reseed. + """ + from litellm.proxy.db.budget_window_spend_writer import ( + commit_window_spend_updates, + ) + + await commit_window_spend_updates( + prisma_client=prisma_client, + transactions=window_spend_transactions, + ) + async def _drain_and_commit_daily_tag_spend_from_redis( self, prisma_client: PrismaClient, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 39b2111fb22..4dd23270bf8 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -23,6 +23,7 @@ from litellm.constants import ( REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY, + REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( @@ -42,6 +43,10 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, ) from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendTransaction, + WindowSpendUpdateQueue, +) from litellm.secret_managers.main import str_to_bool from litellm.types.caching import ( RedisPipelineLpopOperation, @@ -182,6 +187,7 @@ class RedisUpdateBuffer: daily_org_spend_update_queue: DailySpendUpdateQueue, daily_end_user_spend_update_queue: DailySpendUpdateQueue, daily_agent_spend_update_queue: DailySpendUpdateQueue, + window_spend_update_queue: WindowSpendUpdateQueue | None = None, ): """ Stores the in-memory spend updates to Redis @@ -250,6 +256,11 @@ class RedisUpdateBuffer: daily_agent_spend_update_transactions: Final = ( await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) + window_spend_update_transactions: Final = ( + await window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() + if window_spend_update_queue is not None + else () + ) verbose_proxy_logger.debug("ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions) verbose_proxy_logger.debug("ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions) @@ -286,6 +297,11 @@ class RedisUpdateBuffer: REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE, ), + ( + window_spend_update_transactions, + REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + ServiceTypes.REDIS_WINDOW_SPEND_UPDATE_QUEUE, + ), ] rpush_list: Final[list[RedisPipelineRpushOperation]] = [] @@ -326,12 +342,14 @@ class RedisUpdateBuffer: daily_org_spend_update_transactions=daily_org_spend_update_transactions, daily_end_user_spend_update_transactions=daily_end_user_spend_update_transactions, daily_agent_spend_update_transactions=daily_agent_spend_update_transactions, + window_spend_update_transactions=window_spend_update_transactions, spend_update_queue=spend_update_queue, daily_spend_update_queue=daily_spend_update_queue, daily_team_spend_update_queue=daily_team_spend_update_queue, daily_org_spend_update_queue=daily_org_spend_update_queue, daily_end_user_spend_update_queue=daily_end_user_spend_update_queue, daily_agent_spend_update_queue=daily_agent_spend_update_queue, + window_spend_update_queue=window_spend_update_queue, ) return @@ -351,12 +369,14 @@ class RedisUpdateBuffer: daily_org_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, daily_end_user_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, daily_agent_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, + window_spend_update_transactions: tuple[WindowSpendTransaction, ...] | None, spend_update_queue: SpendUpdateQueue, daily_spend_update_queue: DailySpendUpdateQueue, daily_team_spend_update_queue: DailySpendUpdateQueue, daily_org_spend_update_queue: DailySpendUpdateQueue, daily_end_user_spend_update_queue: DailySpendUpdateQueue, daily_agent_spend_update_queue: DailySpendUpdateQueue, + window_spend_update_queue: WindowSpendUpdateQueue | None, ) -> None: """ Put drained-but-unpushed transactions back into in-memory queues. @@ -430,6 +450,9 @@ class RedisUpdateBuffer: if daily_txns: await daily_queue.update_queue.put(daily_txns) + if window_spend_update_transactions and window_spend_update_queue is not None: + await window_spend_update_queue.update_queue.put(window_spend_update_transactions) + async def restore_transactions_to_redis( self, db_spend_update_transactions: DBSpendUpdateTransactions | None = None, @@ -439,6 +462,7 @@ class RedisUpdateBuffer: daily_end_user_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, daily_agent_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, daily_tag_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + window_spend_update_transactions: Sequence[WindowSpendTransaction] | None = None, ) -> None: """ Re-push transactions that were popped from Redis but not committed to the DB. @@ -460,6 +484,7 @@ class RedisUpdateBuffer: (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY), (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY), (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY), + (window_spend_update_transactions, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY), ) rpush_list: Final = tuple( @@ -577,20 +602,22 @@ class RedisUpdateBuffer: dict[str, DailyOrganizationSpendTransaction] | None, dict[str, DailyEndUserSpendTransaction] | None, dict[str, DailyAgentSpendTransaction] | None, + tuple[WindowSpendTransaction, ...] | None, ]: """ - Drains the main 6 Redis buffer queues in a single pipeline round-trip. + Drains the main 7 Redis buffer queues in a single pipeline round-trip. - Returns a 6-tuple of parsed results in this order: + Returns a 7-tuple of parsed results in this order: 0: DBSpendUpdateTransactions 1: daily user spend 2: daily team spend 3: daily org spend 4: daily end-user spend 5: daily agent spend + 6: budget window spend """ if self.redis_cache is None: - return None, None, None, None, None, None + return None, None, None, None, None, None, None lpop_list: Final[list[RedisPipelineLpopOperation]] = [ RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), @@ -614,12 +641,16 @@ class RedisUpdateBuffer: key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ), + RedisPipelineLpopOperation( + key=REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + ), ] raw_results: Final = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list) # Pad with None if pipeline returned fewer results than expected - while len(raw_results) < 6: + while len(raw_results) < 7: raw_results.append(None) # Slot 0: DBSpendUpdateTransactions @@ -640,6 +671,14 @@ class RedisUpdateBuffer: aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(list_of_daily) daily_results.append(aggregated) + window_spend: Final = ( + WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( + tuple(json.loads(transaction) for transaction in raw_results[6]) + ) + if raw_results[6] is not None + else None + ) + return ( db_spend, cast(dict[str, DailyUserSpendTransaction] | None, daily_results[0]), @@ -647,6 +686,7 @@ class RedisUpdateBuffer: cast(dict[str, DailyOrganizationSpendTransaction] | None, daily_results[2]), cast(dict[str, DailyEndUserSpendTransaction] | None, daily_results[3]), cast(dict[str, DailyAgentSpendTransaction] | None, daily_results[4]), + window_spend, ) async def store_in_memory_daily_tag_spend_updates_in_redis( diff --git a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py new file mode 100644 index 00000000000..04dea66165e --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py @@ -0,0 +1,176 @@ +""" +In memory buffer for per-budget-window spend increments. + +Kept separate from SpendUpdateQueue: an increment is only meaningful together +with the window it landed in, so two increments for the same entity must not be +merged when their window_start differs. +""" + +import asyncio +import math +from collections.abc import Sequence +from datetime import datetime, timezone +from itertools import chain, groupby +from typing import Final, TypedDict + +from typing_extensions import ReadOnly + +from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE +from litellm.proxy.db.db_transaction_queue.base_update_queue import BaseUpdateQueue + + +class WindowSpendTransaction(TypedDict): + """One increment for a single (entity, budget window) pair. + + window_start is an ISO-8601 string rather than a datetime so the + transaction survives the JSON round trip through the Redis buffer. + + request_ids carries the LiteLLM_SpendLogs ids this spend came from. The + one-time seed for a window that has no row yet subtracts them from its + LiteLLM_SpendLogs aggregate, because the spend log writer flushes on its + own ~2s poll and will usually have persisted these rows before the window + queue flushes; without the exclusion the seed and the increment would each + count them. + + started_at is the earliest request start in the batch. The seed only + subtracts a request_id whose LiteLLM_SpendLogs.startTime is at or after it, + so a client that replays an old id through x-litellm-call-id cannot make the + seed drop the historical row that id already paid for. + """ + + entity_type: ReadOnly[str] + entity_id: ReadOnly[str] + window_duration: ReadOnly[str] + window_start: ReadOnly[str] + spend: ReadOnly[float] + request_ids: ReadOnly[Sequence[str]] + started_at: ReadOnly[str | None] + + +def to_naive_utc(value: datetime) -> datetime: + """LiteLLM_BudgetWindowSpend.window_start is TIMESTAMP(3), which holds naive UTC.""" + if value.tzinfo is None: + return value + return value.astimezone(timezone.utc).replace(tzinfo=None) + + +def window_spend_group_key(transaction: WindowSpendTransaction) -> tuple[str, str, str, str]: + """Identity of a window increment: the row's primary key plus the window it + belongs to. Two increments only aggregate when all four match.""" + return ( + transaction["entity_type"], + transaction["entity_id"], + transaction["window_duration"], + transaction["window_start"], + ) + + +def build_window_spend_transaction( + entity_type: str, + entity_id: str, + window_duration: str, + window_start: datetime, + spend: float, + request_id: str | None = None, + started_at: datetime | None = None, +) -> WindowSpendTransaction: + return WindowSpendTransaction( + entity_type=entity_type, + entity_id=entity_id, + window_duration=window_duration, + window_start=to_naive_utc(window_start).isoformat(timespec="microseconds"), + spend=spend, + request_ids=() if request_id is None else (request_id,), + started_at=None + if started_at is None + else to_naive_utc(started_at.astimezone(timezone.utc)).isoformat(timespec="microseconds"), + ) + + +def _merge_window_spend_transactions( + payloads: tuple[WindowSpendTransaction, ...], +) -> WindowSpendTransaction: + first: Final = payloads[0] + started_ats: Final = tuple( + started_at for payload in payloads if (started_at := payload.get("started_at")) is not None + ) + return WindowSpendTransaction( + entity_type=first["entity_type"], + entity_id=first["entity_id"], + window_duration=first["window_duration"], + window_start=first["window_start"], + spend=math.fsum(payload["spend"] for payload in payloads), + request_ids=tuple(sorted(frozenset(chain.from_iterable(payload["request_ids"] for payload in payloads)))), + started_at=min(started_ats) if started_ats else None, + ) + + +class WindowSpendUpdateQueue(BaseUpdateQueue): + """ + In memory buffer for budget-window spend increments committed to + LiteLLM_BudgetWindowSpend. + + Add an update with the payload built by build_window_spend_transaction: + window_spend_update_queue.add_update( + build_window_spend_transaction( + entity_type="key", + entity_id="", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=0.02, + ) + ) + """ + + def __init__(self) -> None: + super().__init__() + self.update_queue: asyncio.Queue[tuple[WindowSpendTransaction, ...]] = asyncio.Queue( + maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE + ) + + async def add_update(self, update: WindowSpendTransaction) -> None: + """Enqueue an update.""" + verbose_proxy_logger.debug("Adding budget window spend update to queue: %s", update) + await self.update_queue.put((update,)) + if self.update_queue.qsize() >= self.MAX_SIZE_IN_MEMORY_QUEUE: + verbose_proxy_logger.warning( + "Budget window spend update queue is full. Aggregating all entries in queue to concatenate entries." + ) + await self.aggregate_queue_updates() + + async def aggregate_queue_updates(self) -> None: + """Collapse everything currently queued into a single aggregated update.""" + updates: Final = await self.flush_all_updates_from_in_memory_queue() + await self.update_queue.put(WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(updates)) + + async def flush_and_get_aggregated_window_spend_transactions( + self, + ) -> tuple[WindowSpendTransaction, ...]: + """Drain the queue and return the increments aggregated per window.""" + updates: Final = await self.flush_all_updates_from_in_memory_queue() + if len(updates) > 0: + verbose_proxy_logger.info( + "Spend tracking - flushed %d budget window spend update batches from in-memory queue", + len(updates), + ) + return WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(updates) + + @staticmethod + def get_aggregated_window_spend_transactions( + updates: Sequence[Sequence[WindowSpendTransaction]], + ) -> tuple[WindowSpendTransaction, ...]: + """Sum spend per (entity_type, entity_id, window_duration, window_start). + + Increments belonging to different windows stay separate even when they + share a primary key, so a window boundary crossed mid-tick does not fold + the new window's spend into the previous window's total. + + The result is ordered by that same key, which is the order the flush + needs: primary key first for cross-pod lock ordering, then window_start + so an older window is applied before the roll that supersedes it. + """ + ordered: Final = tuple(sorted(chain.from_iterable(updates), key=window_spend_group_key)) + return tuple( + _merge_window_spend_transactions(tuple(group)) for _, group in groupby(ordered, key=window_spend_group_key) + ) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index deb9cd5ae25..7b3c261036e 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -14,14 +14,18 @@ memory in long-lived deployments. import asyncio from collections import OrderedDict -from datetime import datetime +from collections.abc import Mapping +from datetime import datetime, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final, Optional from litellm._logging import verbose_proxy_logger from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import Litellm_EntityType from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( + BudgetWindowSpendRepository, SpendLogsRepository, TeamMembershipRepository, ) @@ -36,6 +40,25 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +_WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType( + { + "Key": Litellm_EntityType.KEY.value, + "Team": Litellm_EntityType.TEAM.value, + } +) + +_WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + { + "Key": "api_key", + "Team": "team_id", + } +) + + +def _as_utc(value: datetime) -> datetime: + return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) + + class SpendCounterReseed: """ Reseeds spend counters from the authoritative DB and warms the cache, @@ -205,6 +228,92 @@ class SpendCounterReseed: raise return current_value + @staticmethod + async def window_from_table( + prisma_client: Optional["PrismaClient"], + entity_type: str, + entity_id: str, + window_duration: str, + expected_window_start: datetime, + ) -> float | None: + """ + Read the maintained per-window spend row by primary key. + + Returns the row's spend only when the row belongs to the window the + caller is enforcing, i.e. ``row.window_start >= expected_window_start``. + A row at or past the expected start was rolled by a pod whose reset_at + was at least as fresh as this caller's, so it is trusted; an older row + means the window boundary was crossed and nothing has rolled the row + yet, so its spend belongs to a previous window. + + Returns None for a missing, stale or unreadable row so the caller falls + back to the spend-logs aggregate. ``entity_type`` is the counter-facing + label ("Key"/"Team"); anything else has no row and returns None. + """ + if prisma_client is None: + return None + row_entity_type: Final = _WINDOW_SPEND_ENTITY_TYPES.get(entity_type) + if row_entity_type is None: + return None + + try: + row: Final = await BudgetWindowSpendRepository(prisma_client).table.find_unique( + where={ + "entity_type_entity_id_window_duration": { + "entity_type": row_entity_type, + "entity_id": entity_id, + "window_duration": window_duration, + } + } + ) + except Exception: # noqa: BLE001 # any read failure (DB, stale prisma client) must degrade to the aggregate path + verbose_proxy_logger.exception( + "SpendCounterReseed.window_from_table: failed for %s=%s window=%s", + entity_type, + entity_id, + window_duration, + ) + return None + + if row is None: + return None + if _as_utc(row.window_start) < _as_utc(expected_window_start): + return None + return float(row.spend or 0.0) + + @staticmethod + async def window_from_db( + prisma_client: Optional["PrismaClient"], + entity_type: str, + entity_id: str, + window_duration: str | None, + window_start: datetime, + ) -> float | None: + """ + Authoritative window spend: the maintained row first, falling back to + the spend-logs aggregate only when no current row exists. + + The aggregate range-scans an unindexed table, so it must stay a + transitional path (window configured before the row existed) rather + than a steady-state read. + """ + if window_duration is not None: + from_table: Final = await SpendCounterReseed.window_from_table( + prisma_client=prisma_client, + entity_type=entity_type, + entity_id=entity_id, + window_duration=window_duration, + expected_window_start=window_start, + ) + if from_table is not None: + return from_table + return await SpendCounterReseed.window_from_spend_logs( + prisma_client=prisma_client, + entity_type=entity_type, + entity_id=entity_id, + window_start=window_start, + ) + @staticmethod async def window_from_spend_logs( prisma_client: Optional["PrismaClient"], @@ -215,20 +324,13 @@ class SpendCounterReseed: if prisma_client is None: return None - if entity_type == "Key": - group_field = "api_key" - where = { - "api_key": entity_id, - "startTime": {"gte": window_start}, - } - elif entity_type == "Team": - group_field = "team_id" - where = { - "team_id": entity_id, - "startTime": {"gte": window_start}, - } - else: + group_field: Final = _WINDOW_SPEND_LOG_FIELDS.get(entity_type) + if group_field is None: return None + where: Final = { + group_field: entity_id, + "startTime": {"gte": window_start}, + } try: response: Final = await SpendLogsRepository(prisma_client).table.group_by( @@ -258,6 +360,7 @@ class SpendCounterReseed: counter_key: str, entity_type: str, entity_id: str, + window_duration: str | None, window_start: datetime, ) -> float | None: lock: Final = await SpendCounterReseed._get_lock(counter_key) @@ -276,10 +379,11 @@ class SpendCounterReseed: if val is not None: return float(val) - window_spend: Final = await SpendCounterReseed.window_from_spend_logs( + window_spend: Final = await SpendCounterReseed.window_from_db( prisma_client=prisma_client, entity_type=entity_type, entity_id=entity_id, + window_duration=window_duration, window_start=window_start, ) if window_spend is None: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index d0e4f548c99..f02901f0e97 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -587,7 +587,7 @@ async def _update_database_and_spend_counters( model_access_groups: Sequence[str] | None = None, ) -> None: try: - await proxy_logging_obj.db_spend_update_writer.update_database( + spend_log_request_id = await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, @@ -623,6 +623,8 @@ async def _update_database_and_spend_counters( budget_reservation=budget_reservation, end_user_id=end_user_id, tags=request_tags, + request_id=spend_log_request_id, + request_started_at=start_time, model_access_groups=model_access_groups, ) except Exception: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2bad742ad02..654fe4a3f2e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -398,6 +398,9 @@ from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( SPEND_LOG_CLEANUP_BOUND_SETTINGS, SpendLogCleanup, ) +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, +) from litellm.proxy.db.exception_handler import ( PrismaDBExceptionHandler, call_with_db_reconnect_retry, @@ -2433,6 +2436,7 @@ async def get_current_spend( max_budget: float | None = None, window_entity_type: str | None = None, window_entity_id: str | None = None, + window_duration: str | None = None, window_start: datetime | None = None, fallback_authoritative: bool = False, ) -> float: @@ -2457,7 +2461,8 @@ async def get_current_spend( runs and a key can leak spend past ``max_budget`` indefinitely. The authoritative source depends on the counter: primary key/team/user/org counters read the DB row; per-window counters (``window_start`` supplied) - aggregate spend logs; end-user/tag counters have no DB row, so the caller's + read the maintained window-spend row and only aggregate spend logs when + that row is missing or stale; end-user/tag counters have no DB row, so the caller's ``fallback_spend`` (loaded fresh in auth) is authoritative. The DB read is skipped for healthy primary counters (counter at or above recorded spend) and cached in-process for a few seconds, so a persistently stale counter @@ -2482,6 +2487,7 @@ async def get_current_spend( counter_key=counter_key, window_entity_type=window_entity_type, window_entity_id=window_entity_id, + window_duration=window_duration, window_start=window_start, ) if authoritative is not None: @@ -2563,6 +2569,7 @@ async def _authoritative_floor_spend( counter_key: str, window_entity_type: str | None = None, window_entity_id: str | None = None, + window_duration: str | None = None, window_start: datetime | None = None, ) -> float | None: marker_key: Final = f"spend_db_floor:{counter_key}" @@ -2577,10 +2584,11 @@ async def _authoritative_floor_spend( and window_entity_id is not None and window_start is not None ): - db_spend = await SpendCounterReseed.window_from_spend_logs( + db_spend = await SpendCounterReseed.window_from_db( prisma_client=prisma_client, entity_type=window_entity_type, entity_id=window_entity_id, + window_duration=window_duration, window_start=window_start, ) if db_spend is None: @@ -2650,6 +2658,8 @@ async def increment_spend_counters( budget_reservation: dict | None = None, end_user_id: str | None = None, tags: list[str] | None = None, + request_id: str | None = None, + request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, ): """ @@ -2704,15 +2714,28 @@ async def increment_spend_counters( return for window in key_budget_limits: duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration + key_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at key_window_counter = f"spend:key:{hashed_token}:window:{duration}" + key_window_start = get_budget_window_start(window) if key_window_counter not in reserved_counter_keys: await _init_and_increment_window_spend_counter( counter_key=key_window_counter, entity_type="Key", entity_id=hashed_token, - window_start=get_budget_window_start(window), + window_duration=duration, + window_start=key_window_start, increment=cost, ) + await _enqueue_window_spend_row_update( + entity_type=Litellm_EntityType.KEY, + entity_id=hashed_token, + reset_at=key_window_reset_at, + window_duration=duration, + window_start=key_window_start, + increment=cost, + request_id=request_id, + request_started_at=request_started_at, + ) async def _team_scope(scope_team_id: str) -> None: team_counter_key: Final = f"spend:team:{scope_team_id}" @@ -2735,15 +2758,28 @@ async def increment_spend_counters( return for window in team_budget_limits: duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration + team_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at team_window_counter = f"spend:team:{scope_team_id}:window:{duration}" + team_window_start = get_budget_window_start(window) if team_window_counter not in reserved_counter_keys: await _init_and_increment_window_spend_counter( counter_key=team_window_counter, entity_type="Team", entity_id=scope_team_id, - window_start=get_budget_window_start(window), + window_duration=duration, + window_start=team_window_start, increment=cost, ) + await _enqueue_window_spend_row_update( + entity_type=Litellm_EntityType.TEAM, + entity_id=scope_team_id, + reset_at=team_window_reset_at, + window_duration=duration, + window_start=team_window_start, + increment=cost, + request_id=request_id, + request_started_at=request_started_at, + ) async def _team_member_scope(scope_user_id: str, scope_team_id: str) -> None: team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}" @@ -2962,10 +2998,62 @@ async def _init_and_increment_spend_counter( await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) +async def _enqueue_window_spend_row_update( + entity_type: Litellm_EntityType, + entity_id: str, + reset_at: datetime | str | None, + window_duration: str, + window_start: datetime | None, + increment: float, + request_id: str | None, + request_started_at: datetime | None, +) -> None: + """Queue this request's cost against the LiteLLM_BudgetWindowSpend row for + the window, so enforcement can read a maintained total instead of + aggregating LiteLLM_SpendLogs. + + request_id is the LiteLLM_SpendLogs id this cost was recorded under and + request_started_at its startTime; the flush uses them to keep the one-time + seed from counting a request that its increment already covers. + + Enqueued even when the cache increment was skipped for a reserved counter: + the reservation only pre-charged the counter, and the row still owes the + actual cost. + + Windows with no reset_at slide with wall clock, so their window_start moves + on every request and no single row can represent them. Those are left to + the read path's LiteLLM_SpendLogs fallback rather than rewritten per + request. + """ + if window_start is None or not reset_at: + return + try: + await proxy_logging_obj.db_spend_update_writer.window_spend_update_queue.add_update( + build_window_spend_transaction( + entity_type=entity_type.value, + entity_id=entity_id, + window_duration=window_duration, + window_start=window_start, + spend=increment, + request_id=request_id, + started_at=request_started_at, + ) + ) + except Exception as e: # noqa: BLE001 # spend tracking must never fail the cost callback + verbose_proxy_logger.debug( + "Unable to enqueue budget window spend update for %s=%s window=%s: %s", + entity_type.value, + entity_id, + window_duration, + e, + ) + + async def _init_and_increment_window_spend_counter( counter_key: str, entity_type: str, entity_id: str, + window_duration: str | None, window_start: datetime | None, increment: float, ): @@ -2980,6 +3068,7 @@ async def _init_and_increment_window_spend_counter( counter_key=counter_key, entity_type=entity_type, entity_id=entity_id, + window_duration=window_duration, window_start=window_start, ) if initialized is False: @@ -3025,6 +3114,7 @@ async def _ensure_window_spend_counter_initialized( counter_key: str, entity_type: str, entity_id: str, + window_duration: str | None, window_start: datetime, ) -> bool: is_warm: Final = await _is_spend_counter_cache_warm(counter_key=counter_key) @@ -3037,6 +3127,7 @@ async def _ensure_window_spend_counter_initialized( counter_key=counter_key, entity_type=entity_type, entity_id=entity_id, + window_duration=window_duration, window_start=window_start, ) if window_spend is None: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 46086711923..8b2a5dd9312 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -46,6 +46,7 @@ class _BudgetCounter: entity_id: str source_cache_key: str | None = None spend_log_entity_id: str | None = None + window_duration: str | None = None window_start: datetime | None = None @@ -709,6 +710,7 @@ def _get_budget_limit_counters( entity_type=entity_type, entity_id=f"{entity_id}:{budget_duration}", spend_log_entity_id=entity_id, + window_duration=str(budget_duration), window_start=window_start, ) ) @@ -752,6 +754,7 @@ async def _reserve_counter( counter_key=counter.counter_key, entity_type=counter.entity_type, entity_id=counter.spend_log_entity_id, + window_duration=counter.window_duration, window_start=counter.window_start, ) if initialized is False: diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index fac142bb017..18cf884f267 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -68,6 +68,10 @@ class SpendLogsRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogs table_name = "litellm_spendlogs" +class BudgetWindowSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_BudgetWindowSpend"]): + table_name = "litellm_budgetwindowspend" + + class ClaudeCodePluginRepository(PrismaTableRepository["prisma_models.LiteLLM_ClaudeCodePluginTable"]): table_name = "litellm_claudecodeplugintable" diff --git a/litellm/types/services.py b/litellm/types/services.py index 74f908548d5..c558f6fb9d2 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -40,6 +40,8 @@ class ServiceTypes(str, enum.Enum): # spend update queue - current spend of key, user, team IN_MEMORY_SPEND_UPDATE_QUEUE = "in_memory_spend_update_queue" REDIS_SPEND_UPDATE_QUEUE = "redis_spend_update_queue" + # budget window spend queue - per-window spend of key, team + REDIS_WINDOW_SPEND_UPDATE_QUEUE = "redis_window_spend_update_queue" class ServiceConfig(TypedDict): diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index eb713b6f3c2..03b05bd9d87 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -833,6 +833,7 @@ def _make_reset_budget_windows_job( raise AssertionError(f"Unexpected query_raw call: {query}") prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw) + prisma_client.db.execute_raw = AsyncMock(return_value=1) prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=None) prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None) @@ -904,6 +905,145 @@ def test_reset_budget_windows_resets_expired_key_window(monkeypatch): spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0) +def _window_spend_rolls(prisma_client): + return [ + call.args + for call in prisma_client.db.execute_raw.await_args_list + if "LiteLLM_BudgetWindowSpend" in call.args[0] + ] + + +def test_reset_budget_windows_rolls_the_key_window_spend_row(monkeypatch): + """The maintained per-window total has to start the new window at zero + alongside the counter, or enforcement keeps reading the old window's spend.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=5)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-expired", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + rolls = _window_spend_rolls(prisma_client) + assert len(rolls) == 1 + query, entity_type, entity_id, window_duration, new_window_start, _updated_at = rolls[0] + assert (entity_type, entity_id, window_duration) == ("key", "sk-expired", "1d") + assert "spend = 0" in " ".join(query.split()) + + # window_start is the start of the window that just began: new reset_at minus the duration. + written_windows = json.loads( + prisma_client.db.litellm_verificationtoken.update.await_args.kwargs["data"]["budget_limits"] + ) + new_reset_at = datetime.fromisoformat(written_windows[0]["reset_at"].replace("Z", "+00:00")).replace(tzinfo=None) + assert new_window_start == pytest.approx( + new_reset_at - timedelta(days=1), + abs=timedelta(seconds=1), + ) + + +def test_reset_budget_windows_roll_is_conditional_on_an_older_stored_window(monkeypatch): + """Another pod may already have rolled the row; clobbering it would drop + spend that landed under the new window.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=5)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-expired", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + query = " ".join(_window_spend_rolls(prisma_client)[0][0].split()) + assert "AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')" in query + + +def test_reset_budget_windows_rolls_the_team_window_spend_row(monkeypatch): + now = datetime.utcnow() + expired = (now - timedelta(minutes=1)).isoformat() + "Z" + + team_rows = [ + { + "team_id": "team-expired", + "budget_limits": [{"budget_duration": "30d", "reset_at": expired}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=[], team_rows=team_rows) + + asyncio.run(job.reset_budget_windows()) + + rolls = _window_spend_rolls(prisma_client) + assert len(rolls) == 1 + assert rolls[0][1:4] == ("team", "team-expired", "30d") + + +def test_reset_budget_windows_does_not_roll_an_unexpired_window(monkeypatch): + now = datetime.utcnow() + future = (now + timedelta(hours=1)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-future", + "budget_limits": [{"budget_duration": "1d", "reset_at": future}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + assert _window_spend_rolls(prisma_client) == [] + + +def test_reset_budget_windows_rolls_only_the_expired_window_of_a_key(monkeypatch): + now = datetime.utcnow() + key_rows = [ + { + "token": "sk-mixed", + "budget_limits": [ + {"budget_duration": "1d", "reset_at": (now - timedelta(minutes=5)).isoformat() + "Z"}, + {"budget_duration": "30d", "reset_at": (now + timedelta(days=2)).isoformat() + "Z"}, + ], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + rolls = _window_spend_rolls(prisma_client) + assert [roll[3] for roll in rolls] == ["1d"] + + +def test_reset_budget_windows_survives_a_failed_window_spend_roll(monkeypatch): + """The row is an optimization over aggregating LiteLLM_SpendLogs; a DB + failure there must not stop the counter reset from being persisted.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=5)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-expired", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired}], + } + ] + job, prisma_client, spend_counter_cache = _make_reset_budget_windows_job( + monkeypatch, key_rows=key_rows, team_rows=[] + ) + prisma_client.db.execute_raw = AsyncMock(side_effect=Exception("connection reset")) + + asyncio.run(job.reset_budget_windows()) + + prisma_client.db.litellm_verificationtoken.update.assert_awaited_once() + spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0) + + def test_reset_budget_windows_skips_unexpired_key_window(monkeypatch): """If `reset_at` is in the future, no write should happen for that key.""" now = datetime.utcnow() diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index fb0c994a476..504654e103a 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -1,4 +1,5 @@ import json +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -22,9 +23,7 @@ def redis_update_buffer(mock_redis_cache): @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_uses_pipeline( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache): """ Verify store_in_memory_spend_updates_in_redis calls async_rpush_pipeline once with the correct operations and skips empty queues. @@ -33,35 +32,29 @@ async def test_store_in_memory_spend_updates_uses_pipeline( # Create mock queues - only 3 of 6 have data spend_update_queue = AsyncMock() - spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = ( - AsyncMock(return_value={"key_list_transactions": {"key1": 1.0}}) + spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock( + return_value={"key_list_transactions": {"key1": 1.0}} ) daily_spend_queue = AsyncMock() - daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={"user_key1": {"spend": 1.0}}) + daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={"user_key1": {"spend": 1.0}} ) daily_team_queue = AsyncMock() - daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={"team_key1": {"spend": 2.0}}) + daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={"team_key1": {"spend": 2.0}} ) # Empty queues daily_org_queue = AsyncMock() - daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={}) - ) + daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) daily_end_user_queue = AsyncMock() - daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value=None) - ) + daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value=None) daily_agent_queue = AsyncMock() - daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={}) - ) + daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) await redis_update_buffer.store_in_memory_spend_updates_in_redis( spend_update_queue=spend_update_queue, @@ -82,9 +75,7 @@ async def test_store_in_memory_spend_updates_uses_pipeline( @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_restores_on_rpush_failure( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_restores_on_rpush_failure(redis_update_buffer, mock_redis_cache): """ If async_rpush_pipeline raises, the already-drained transactions must be put back into the in-memory queues so the next scheduler tick retries. @@ -98,9 +89,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( SpendUpdateQueue, ) - mock_redis_cache.async_rpush_pipeline = AsyncMock( - side_effect=ConnectionError("redis went away") - ) + mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away")) spend_queue = SpendUpdateQueue() daily_user_queue = DailySpendUpdateQueue() @@ -145,16 +134,12 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( # After restore, the main spend queue should hold one item per # (entity_type, entity_id) pair with the aggregated cost - restored_spend = ( - await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() - ) + restored_spend = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() assert restored_spend["key_list_transactions"] == {"key-abc": 1.5} assert restored_spend["team_list_transactions"] == {"team-xyz": 2.5} # Daily user queue should hold the same aggregated dict - restored_daily = ( - await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions() - ) + restored_daily = await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions() assert restored_daily == { "user1_day_model": { "spend": 1.0, @@ -165,9 +150,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_all_empty_returns_early( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_update_buffer, mock_redis_cache): """ When all queues are empty, pipeline should never be called. """ @@ -175,13 +158,9 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early( # All queues return empty empty_queue = AsyncMock() - empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock( - return_value={} - ) + empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={}) empty_daily_queue = AsyncMock() - empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={}) - ) + empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) await redis_update_buffer.store_in_memory_spend_updates_in_redis( spend_update_queue=empty_queue, @@ -196,14 +175,13 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early( @pytest.mark.asyncio -async def test_get_all_transactions_from_redis_buffer_pipeline( - redis_update_buffer, mock_redis_cache -): +async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buffer, mock_redis_cache): """ Verify get_all_transactions_from_redis_buffer_pipeline correctly parses and aggregates results from async_lpop_pipeline. """ - # Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories + # Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories, + # slot 6 = budget window spend db_spend_json = json.dumps( { "key_list_transactions": {"key1": 1.0, "key2": 2.0}, @@ -217,6 +195,18 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( ) daily_user_json = json.dumps({"user_key1": {"spend": 1.0, "api_requests": 1}}) daily_team_json = json.dumps({"team_key1": {"spend": 2.0, "api_requests": 2}}) + window_spend_json = json.dumps( + [ + { + "entity_type": "key", + "entity_id": "hashed-token", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 3.0, + "request_ids": ["req-1"], + } + ] + ) mock_redis_cache.async_lpop_pipeline = AsyncMock( return_value=[ @@ -226,13 +216,30 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( None, # slot 3: daily org (empty) None, # slot 4: daily end-user (empty) None, # slot 5: daily agent (empty) + [window_spend_json, window_spend_json], # slot 6: budget window spend ] ) result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() - assert len(result) == 6 - db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent = result + assert len(result) == 7 + ( + db_spend, + daily_user, + daily_team, + daily_org, + daily_end_user, + daily_agent, + window_spend, + ) = result + + # Budget window spend from two pods is summed per window, not overwritten, + # and both pods' request ids reach the seed exclusion. + assert window_spend is not None + assert len(window_spend) == 1 + assert window_spend[0]["spend"] == 6.0 + assert window_spend[0]["entity_id"] == "hashed-token" + assert window_spend[0]["request_ids"] == ("req-1",) # Verify db spend was parsed correctly assert db_spend is not None @@ -255,6 +262,10 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( # Verify pipeline was called once with correct keys mock_redis_cache.async_lpop_pipeline.assert_called_once() + from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + + popped_keys = [op["key"] for op in mock_redis_cache.async_lpop_pipeline.call_args.kwargs["lpop_list"]] + assert popped_keys[6] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY @pytest.mark.asyncio @@ -262,13 +273,11 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): """When redis_cache is None, should return all Nones""" buffer = RedisUpdateBuffer(redis_cache=None) result = await buffer.get_all_transactions_from_redis_buffer_pipeline() - assert result == (None, None, None, None, None, None) + assert result == (None, None, None, None, None, None, None) @pytest.mark.asyncio -async def test_restore_transactions_to_redis_pushes_only_provided( - redis_update_buffer, mock_redis_cache -): +async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_buffer, mock_redis_cache): """ restore_transactions_to_redis re-pushes only the transaction sets it was given, to their matching buffer keys, so uncommitted spend can be retried. @@ -302,9 +311,42 @@ async def test_restore_transactions_to_redis_pushes_only_provided( @pytest.mark.asyncio -async def test_restore_transactions_to_redis_noop_when_empty( - redis_update_buffer, mock_redis_cache -): +async def test_restored_window_spend_transactions_drain_back_unchanged(redis_update_buffer, mock_redis_cache): + """A window commit that fails after the destructive lpop must be re-pushed + in the store path's encoding, so the next drain returns the same increments.""" + from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, + ) + + window_transactions = ( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=3.0, + request_id="req-1", + started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc), + ), + ) + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) + + await redis_update_buffer.restore_transactions_to_redis(window_spend_update_transactions=window_transactions) + + rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] + assert [op["key"] for op in rpush_list] == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY] + + mock_redis_cache.async_lpop_pipeline = AsyncMock( + return_value=[None, None, None, None, None, None, list(rpush_list[0]["values"])] + ) + drained = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + + assert drained[6] == window_transactions + + +@pytest.mark.asyncio +async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer, mock_redis_cache): """Nothing to restore -> no Redis call.""" mock_redis_cache.async_rpush_pipeline = AsyncMock() await redis_update_buffer.restore_transactions_to_redis() @@ -312,15 +354,11 @@ async def test_restore_transactions_to_redis_noop_when_empty( @pytest.mark.asyncio -async def test_restore_transactions_to_redis_swallows_redis_error( - redis_update_buffer, mock_redis_cache -): +async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_buffer, mock_redis_cache): """A Redis failure during restore must not propagate to the caller's finally block.""" from redis.exceptions import RedisError - mock_redis_cache.async_rpush_pipeline = AsyncMock( - side_effect=RedisError("redis down") - ) + mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=RedisError("redis down")) await redis_update_buffer.restore_transactions_to_redis( db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}}, @@ -433,3 +471,108 @@ def test_get_transaction_buffer_redis_cache_parses_string_flag(monkeypatch): mock_redis_cache.assert_called_once() assert result is mock_redis_cache.return_value + + +@pytest.mark.asyncio +async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_update_buffer, mock_redis_cache): + """The budget window queue has to ride the same rpush as the daily queues, + otherwise multi-pod deployments never persist per-window spend.""" + from datetime import datetime, timezone + + from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + build_window_spend_transaction, + ) + + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) + + empty_queue = AsyncMock() + empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={}) + empty_daily_queue = AsyncMock() + empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) + + window_queue = WindowSpendUpdateQueue() + await window_queue.add_update( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=1.25, + request_id="req-1", + started_at=datetime(2026, 8, 10, 12, 0, 0, tzinfo=timezone.utc), + ) + ) + + await redis_update_buffer.store_in_memory_spend_updates_in_redis( + spend_update_queue=empty_queue, + daily_spend_update_queue=empty_daily_queue, + daily_team_spend_update_queue=empty_daily_queue, + daily_org_spend_update_queue=empty_daily_queue, + daily_end_user_spend_update_queue=empty_daily_queue, + daily_agent_spend_update_queue=empty_daily_queue, + window_spend_update_queue=window_queue, + ) + + rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] + assert len(rpush_list) == 1 + assert rpush_list[0]["key"] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + pushed = json.loads(rpush_list[0]["values"][0]) + assert pushed == [ + { + "entity_type": "key", + "entity_id": "hashed-token", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 1.25, + "request_ids": ["req-1"], + "started_at": "2026-08-10T12:00:00.000000", + } + ] + + +@pytest.mark.asyncio +async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpush_failure( + redis_update_buffer, mock_redis_cache +): + """The window queue is drained before the rpush, so a Redis hiccup would + silently drop per-window spend without the restore.""" + from datetime import datetime, timezone + + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + build_window_spend_transaction, + ) + + mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away")) + + empty_queue = AsyncMock() + empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={}) + empty_daily_queue = AsyncMock() + empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) + + window_queue = WindowSpendUpdateQueue() + await window_queue.add_update( + build_window_spend_transaction( + entity_type="team", + entity_id="team-1", + window_duration="7d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=4.0, + ) + ) + + await redis_update_buffer.store_in_memory_spend_updates_in_redis( + spend_update_queue=empty_queue, + daily_spend_update_queue=empty_daily_queue, + daily_team_spend_update_queue=empty_daily_queue, + daily_org_spend_update_queue=empty_daily_queue, + daily_end_user_spend_update_queue=empty_daily_queue, + daily_agent_spend_update_queue=empty_daily_queue, + window_spend_update_queue=window_queue, + ) + + restored = await window_queue.flush_and_get_aggregated_window_spend_transactions() + assert [payload["spend"] for payload in restored] == [4.0] + assert [payload["entity_id"] for payload in restored] == ["team-1"] diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py new file mode 100644 index 00000000000..b1ecda57afa --- /dev/null +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py @@ -0,0 +1,269 @@ +import json +from datetime import datetime, timedelta, timezone + +import pytest + +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + build_window_spend_transaction, + to_naive_utc, +) + +WINDOW_A = datetime(2026, 8, 1, tzinfo=timezone.utc) +WINDOW_B = datetime(2026, 8, 31, tzinfo=timezone.utc) + + +def _txn( + entity_id: str, + window_start: datetime, + spend: float, + duration: str = "30d", + entity_type: str = "key", + request_id: str | None = None, + started_at: datetime | None = None, +): + return build_window_spend_transaction( + entity_type=entity_type, + entity_id=entity_id, + window_duration=duration, + window_start=window_start, + spend=spend, + request_id=request_id, + started_at=started_at, + ) + + +def test_build_window_spend_transaction_stores_naive_utc_iso(): + """window_start rides the Redis buffer as a string and lands in a naive-UTC + TIMESTAMP(3) column, so a non-UTC input must be converted, not truncated.""" + non_utc = datetime(2026, 8, 1, 20, 0, tzinfo=timezone(timedelta(hours=-4))) + + assert _txn("k1", non_utc, 1.0, request_id="req-1") == { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-02T00:00:00.000000", + "spend": 1.0, + "request_ids": ("req-1",), + "started_at": None, + } + + +def test_build_window_spend_transaction_stores_started_at_as_naive_utc_iso(): + """started_at is compared against LiteLLM_SpendLogs.startTime, which the + spend log writer stores after converting the request start to UTC.""" + non_utc = datetime(2026, 8, 10, 8, 30, 15, 123456, tzinfo=timezone(timedelta(hours=-4))) + + assert _txn("k1", WINDOW_A, 1.0, started_at=non_utc)["started_at"] == "2026-08-10T12:30:15.123456" + + +@pytest.mark.asyncio +async def test_aggregation_keeps_the_earliest_started_at_of_the_batch(): + """The seed bounds its request-id exclusion at the batch's earliest start, + so a later start must never win the merge.""" + queue = WindowSpendUpdateQueue() + earliest = datetime(2026, 8, 10, 12, 0, 0, tzinfo=timezone.utc) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-2", started_at=earliest + timedelta(seconds=5))) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-1", started_at=earliest)) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-3")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 1 + assert aggregated[0]["started_at"] == "2026-08-10T12:00:00.000000" + assert aggregated[0]["request_ids"] == ("req-1", "req-2", "req-3") + + +def test_to_naive_utc_leaves_naive_values_alone(): + naive = datetime(2026, 8, 1, 12, 0) + assert to_naive_utc(naive) == naive + + +@pytest.mark.asyncio +async def test_aggregation_sums_increments_within_one_window(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.5)) + await queue.add_update(_txn("k1", WINDOW_A, 2.25)) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 1 + assert aggregated[0]["spend"] == pytest.approx(3.75) + + +@pytest.mark.asyncio +async def test_aggregation_keeps_different_windows_of_same_entity_separate(): + """Merging across windows would fold spend from a window that already + rolled into the new window's total, over-counting the new window.""" + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) + await queue.add_update(_txn("k1", WINDOW_B, 2.0)) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 2 + assert {payload["window_start"]: payload["spend"] for payload in aggregated} == { + "2026-08-01T00:00:00.000000": 1.0, + "2026-08-31T00:00:00.000000": 2.0, + } + + +@pytest.mark.asyncio +async def test_aggregation_keeps_durations_entities_and_types_separate(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0, duration="30d")) + await queue.add_update(_txn("k1", WINDOW_A, 2.0, duration="7d")) + await queue.add_update(_txn("k2", WINDOW_A, 4.0, duration="30d")) + await queue.add_update(_txn("k1", WINDOW_A, 8.0, duration="30d", entity_type="team")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 4 + assert sorted(payload["spend"] for payload in aggregated) == [1.0, 2.0, 4.0, 8.0] + + +@pytest.mark.asyncio +async def test_aggregation_orders_by_primary_key_then_window_start(): + """The flush relies on this order: primary key first for cross-pod lock + ordering, then window_start so an older window is applied before the roll + that supersedes it.""" + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("t1", WINDOW_A, 1.0, entity_type="team")) + await queue.add_update(_txn("k2", WINDOW_B, 1.0)) + await queue.add_update(_txn("k2", WINDOW_A, 1.0)) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, duration="7d")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert [ + (payload["entity_type"], payload["entity_id"], payload["window_duration"], payload["window_start"]) + for payload in aggregated + ] == [ + ("key", "k1", "7d", "2026-08-01T00:00:00.000000"), + ("key", "k2", "30d", "2026-08-01T00:00:00.000000"), + ("key", "k2", "30d", "2026-08-31T00:00:00.000000"), + ("team", "t1", "30d", "2026-08-01T00:00:00.000000"), + ] + + +@pytest.mark.asyncio +async def test_aggregation_does_not_collide_on_entity_ids_containing_a_separator(): + """entity_id is free-form (team ids are user supplied), so grouping must not + depend on a flattened string key.""" + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("a:30d:2026-08-01T00:00:00.000000:b", WINDOW_A, 1.0)) + await queue.add_update(_txn("b", WINDOW_A, 2.0)) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 2 + + +@pytest.mark.asyncio +async def test_flush_empties_the_queue(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) + + assert await queue.flush_and_get_aggregated_window_spend_transactions() != () + assert await queue.flush_and_get_aggregated_window_spend_transactions() == () + + +@pytest.mark.asyncio +async def test_aggregate_queue_updates_collapses_in_place(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) + await queue.add_update(_txn("k1", WINDOW_A, 2.0)) + await queue.add_update(_txn("k1", WINDOW_B, 4.0)) + + await queue.aggregate_queue_updates() + + assert queue.update_queue.qsize() == 1 + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + assert sorted(payload["spend"] for payload in aggregated) == [3.0, 4.0] + + +@pytest.mark.asyncio +async def test_aggregation_does_not_mutate_the_queued_payloads(): + """The same payload can be re-aggregated after a failed Redis push, so + aggregation must not accumulate into the caller's object.""" + queue = WindowSpendUpdateQueue() + update = _txn("k1", WINDOW_A, 1.0) + await queue.add_update(update) + await queue.add_update(_txn("k1", WINDOW_A, 2.0)) + + await queue.flush_and_get_aggregated_window_spend_transactions() + + assert update["spend"] == 1.0 + + +def test_aggregation_survives_the_redis_json_round_trip(): + """The Redis buffer stores transactions as JSON, so the aggregated shape + must reload into an equivalent aggregation.""" + aggregated = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( + [(_txn("k1", WINDOW_A, 1.0),), (_txn("k1", WINDOW_B, 2.0),)] + ) + + reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions([json.loads(json.dumps(aggregated))]) + + assert reloaded == aggregated + + +@pytest.mark.asyncio +async def test_aggregation_unions_the_request_ids_of_merged_increments(): + """The seed excludes exactly the requests its batch already covers, so every + merged increment's id has to survive aggregation.""" + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-1")) + await queue.add_update(_txn("k1", WINDOW_A, 2.0, request_id="req-2")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 1 + assert aggregated[0]["request_ids"] == ("req-1", "req-2") + + +@pytest.mark.asyncio +async def test_request_ids_stay_with_their_own_window(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) + await queue.add_update(_txn("k1", WINDOW_B, 2.0, request_id="req-b")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert {payload["window_start"]: payload["request_ids"] for payload in aggregated} == { + "2026-08-01T00:00:00.000000": ("req-a",), + "2026-08-31T00:00:00.000000": ("req-b",), + } + + +@pytest.mark.asyncio +async def test_request_ids_are_deduplicated_and_ordered(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-b")) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert aggregated[0]["request_ids"] == ("req-a", "req-b") + + +@pytest.mark.asyncio +async def test_increment_without_a_request_id_carries_no_exclusion(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert aggregated[0]["request_ids"] == () + + +def test_request_ids_survive_the_redis_json_round_trip(): + aggregated = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( + [(_txn("k1", WINDOW_A, 1.0, request_id="req-1"),)] + ) + + reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions([json.loads(json.dumps(aggregated))]) + + assert reloaded[0]["request_ids"] == ("req-1",) + assert reloaded[0]["spend"] == 1.0 diff --git a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py new file mode 100644 index 00000000000..a849317c930 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py @@ -0,0 +1,596 @@ +import math +from contextlib import asynccontextmanager +from datetime import datetime, timedelta, timezone +from typing import Any + +import pytest + +from litellm.proxy.db.budget_window_spend_writer import ( + commit_window_spend_updates, + roll_window_spend_row, + spend_logs_total_excluding, +) +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, +) + +WINDOW_A = datetime(2026, 8, 1, tzinfo=timezone.utc) +WINDOW_B = datetime(2026, 8, 31, tzinfo=timezone.utc) +BATCH_STARTED_AT = datetime(2026, 8, 10, 12, 0, 0, 250_000, tzinfo=timezone.utc) +BEFORE_BATCH = BATCH_STARTED_AT - timedelta(hours=1) + +ENTITY_TYPE, ENTITY_ID, WINDOW_DURATION, WINDOW_START, INSERT_SPEND, INCREMENT, NOW = range(7) + + +class _FakeBatcher: + def __init__(self) -> None: + self.calls: list[tuple[str, tuple[Any, ...]]] = [] + + def execute_raw(self, query: str, *args: Any) -> None: + self.calls.append((query, args)) + + +class _FakeDB: + """Stands in for prisma_client.db; records every statement it is handed.""" + + def __init__(self, existing_rows: list[dict[str, str]] | None = None) -> None: + self.existing_rows = existing_rows or [] + self.query_raw_calls: list[tuple[str, tuple[Any, ...]]] = [] + self.execute_raw_calls: list[tuple[str, tuple[Any, ...]]] = [] + self.batcher = _FakeBatcher() + self.committed = False + + async def query_raw(self, query: str, *args: Any) -> list[dict[str, str]]: + self.query_raw_calls.append((query, args)) + return self.existing_rows + + async def execute_raw(self, query: str, *args: Any) -> int: + self.execute_raw_calls.append((query, args)) + return 1 + + @asynccontextmanager + async def _tx(self): + yield self + + def tx(self, timeout: Any = None): + return self._tx() + + @asynccontextmanager + async def _batch(self): + yield self.batcher + self.committed = True + + def batch_(self): + return self._batch() + + +class _FakePrismaClient: + def __init__(self, db: _FakeDB) -> None: + self.db = db + + +class _RecordingAggregate: + """Stands in for the LiteLLM_SpendLogs seed aggregate.""" + + def __init__(self, value: float = 5.0) -> None: + self.value = value + self.calls: list[dict[str, Any]] = [] + + async def __call__( + self, + prisma_client: Any, + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Any, + exclude_started_at: datetime | None, + ) -> float | None: + self.calls.append( + { + "entity_type": entity_type, + "entity_id": entity_id, + "window_start": window_start, + "exclude_request_ids": tuple(exclude_request_ids), + "exclude_started_at": exclude_started_at, + } + ) + return self.value + + +class _SpendLogsFake: + """Sums the LiteLLM_SpendLogs rows (request_id, spend, startTime) it holds, + honouring the exclusion exactly as the real aggregate's + NOT (request_id = ANY(...) AND startTime >= bound) does.""" + + def __init__(self, rows: tuple[tuple[str, float, datetime], ...]) -> None: + self.rows = rows + + async def __call__( + self, + prisma_client: Any, + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Any, + exclude_started_at: datetime | None, + ) -> float | None: + excluded = frozenset(exclude_request_ids) if exclude_started_at is not None else frozenset() + return math.fsum( + spend + for request_id, spend, started_at in self.rows + if not (request_id in excluded and started_at >= exclude_started_at) + ) + + +def _batch(request_ids: tuple[str, ...], spend: float, started_at: datetime | None = BATCH_STARTED_AT) -> dict: + return { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": spend, + "request_ids": request_ids, + "started_at": None + if started_at is None + else started_at.replace(tzinfo=None).isoformat(timespec="microseconds"), + } + + +def _existing(entity_type: str, entity_id: str, window_duration: str) -> dict[str, str]: + return {"entity_type": entity_type, "entity_id": entity_id, "window_duration": window_duration} + + +@pytest.mark.asyncio +async def test_no_transactions_touches_no_database(): + db = _FakeDB() + + await commit_window_spend_updates(prisma_client=_FakePrismaClient(db), transactions=()) + + assert db.query_raw_calls == [] + assert db.batcher.calls == [] + + +@pytest.mark.asyncio +async def test_missing_row_is_seeded_from_spend_logs_once(): + """A row created mid-window would undercount everything spent before it + existed, so a brand new primary key inserts the SpendLogs total plus this + increment.""" + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=5.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=aggregate, + ) + + assert len(aggregate.calls) == 1 + assert aggregate.calls[0]["entity_type"] == "key" + assert aggregate.calls[0]["entity_id"] == "k1" + assert aggregate.calls[0]["window_start"] == WINDOW_A + + ((_, params),) = db.batcher.calls + assert params[ENTITY_TYPE] == "key" + assert params[ENTITY_ID] == "k1" + assert params[WINDOW_DURATION] == "30d" + assert params[WINDOW_START] == datetime(2026, 8, 1) + assert params[INSERT_SPEND] == pytest.approx(6.0) + assert params[INCREMENT] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_existing_row_is_never_reseeded(): + """The seed is a full LiteLLM_SpendLogs scan; running it for a row that is + already maintained would both cost a scan and double count.""" + db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) + aggregate = _RecordingAggregate(value=5.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=aggregate, + ) + + assert aggregate.calls == [] + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(1.0) + assert params[INCREMENT] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_seed_runs_only_for_the_primary_keys_that_are_missing(): + db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) + aggregate = _RecordingAggregate(value=5.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("team", "t1", "30d", WINDOW_A, 2.0), + ), + spend_logs_aggregate=aggregate, + ) + + assert [call["entity_id"] for call in aggregate.calls] == ["t1"] + assert [call["entity_type"] for call in aggregate.calls] == ["team"] + by_entity = {params[ENTITY_ID]: params for _, params in db.batcher.calls} + assert by_entity["k1"][INSERT_SPEND] == pytest.approx(1.0) + assert by_entity["t1"][INSERT_SPEND] == pytest.approx(7.0) + + +@pytest.mark.asyncio +async def test_insert_spend_and_increment_differ_only_when_a_row_is_seeded(): + """The conflict arm adds the increment alone so two pods that both seed the + same new window cannot add the SpendLogs base twice.""" + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=9.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 0.25),), + spend_logs_aggregate=aggregate, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(9.25) + assert params[INCREMENT] == pytest.approx(0.25) + + +@pytest.mark.asyncio +async def test_upsert_sql_adds_for_a_current_window_and_replaces_for_a_newer_one(): + """The CASE is the whole contract: an increment at or behind the stored + window_start accumulates, a newer one restarts the window.""" + db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), + ) + + ((query, _),) = db.batcher.calls + normalized = " ".join(query.split()) + assert ( + 'spend = CASE WHEN "LiteLLM_BudgetWindowSpend".window_start >= EXCLUDED.window_start ' + 'THEN "LiteLLM_BudgetWindowSpend".spend + $6 ELSE EXCLUDED.spend END' in normalized + ) + assert 'window_start = GREATEST("LiteLLM_BudgetWindowSpend".window_start, EXCLUDED.window_start)' in normalized + assert "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE SET" in normalized + + +@pytest.mark.asyncio +async def test_upsert_never_interpolates_values_into_the_sql(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "'; DROP TABLE x; --", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=aggregate, + ) + + ((query, params),) = db.batcher.calls + assert "DROP TABLE" not in query + assert params[ENTITY_ID] == "'; DROP TABLE x; --" + + +@pytest.mark.asyncio +async def test_upserts_are_ordered_by_primary_key_then_window_start(): + """Cross-pod lock ordering, plus an older window must be applied before the + roll that supersedes it or the roll would be undone.""" + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + build_window_spend_transaction("team", "t1", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("key", "k2", "30d", WINDOW_B, 1.0), + build_window_spend_transaction("key", "k2", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("key", "k1", "7d", WINDOW_A, 1.0), + ), + spend_logs_aggregate=aggregate, + ) + + ordered = [ + (params[ENTITY_TYPE], params[ENTITY_ID], params[WINDOW_DURATION], params[WINDOW_START]) + for _, params in db.batcher.calls + ] + assert ordered == [ + ("key", "k1", "7d", datetime(2026, 8, 1)), + ("key", "k2", "30d", datetime(2026, 8, 1)), + ("key", "k2", "30d", datetime(2026, 8, 31)), + ("team", "t1", "30d", datetime(2026, 8, 1)), + ] + + +@pytest.mark.asyncio +async def test_existing_row_lookup_sends_every_primary_key_as_array_params(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("team", "t1", "7d", WINDOW_A, 1.0), + ), + spend_logs_aggregate=aggregate, + ) + + ((query, params),) = db.query_raw_calls + assert "unnest($1::text[], $2::text[], $3::text[])" in query + assert params == (("key", "team"), ("k1", "t1"), ("30d", "7d")) + + +@pytest.mark.asyncio +async def test_all_upserts_are_committed_in_one_transaction(): + db = _FakeDB(existing_rows=[_existing("key", "k1", "30d"), _existing("key", "k2", "30d")]) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("key", "k2", "30d", WINDOW_A, 2.0), + ), + ) + + assert len(db.batcher.calls) == 2 + assert db.committed is True + + +@pytest.mark.asyncio +async def test_unknown_entity_type_contributes_no_seed(): + """Only key and team windows have a LiteLLM_SpendLogs column to aggregate; + anything else starts from its increment alone.""" + db = _FakeDB(existing_rows=[]) + + async def no_such_column( + prisma_client, entity_type, entity_id, window_start, exclude_request_ids, exclude_started_at + ): + return None + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("user", "u1", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=no_such_column, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_unavailable_spend_logs_aggregate_seeds_zero_rather_than_failing(): + db = _FakeDB(existing_rows=[]) + + async def unavailable(prisma_client, entity_type, entity_id, window_start, exclude_request_ids, exclude_started_at): + return None + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=unavailable, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_older(): + """Unconditional zeroing would wipe increments a pod already applied under + the new window.""" + db = _FakeDB() + + await roll_window_spend_row( + prisma_client=_FakePrismaClient(db), + entity_type="team", + entity_id="t1", + window_duration="30d", + new_window_start=WINDOW_B, + ) + + ((query, params),) = db.execute_raw_calls + normalized = " ".join(query.split()) + assert "SET window_start = ($4::timestamptz AT TIME ZONE 'UTC'), spend = 0" in normalized + assert "WHERE entity_type = $1 AND entity_id = $2 AND window_duration = $3" in normalized + assert "AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')" in normalized + assert params[:4] == ("team", "t1", "30d", datetime(2026, 8, 31)) + + +@pytest.mark.asyncio +async def test_seed_receives_the_batch_request_ids_and_earliest_start_to_exclude(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("req-1", "req-2", "req-3"), 3.0),), + spend_logs_aggregate=aggregate, + ) + + assert aggregate.calls[0]["exclude_request_ids"] == ("req-1", "req-2", "req-3") + assert aggregate.calls[0]["exclude_started_at"] == BATCH_STARTED_AT + + +@pytest.mark.asyncio +async def test_seed_passes_no_start_bound_when_the_batch_has_none(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("req-1",), 1.0, started_at=None),), + spend_logs_aggregate=aggregate, + ) + + assert aggregate.calls[0]["exclude_started_at"] is None + + +@pytest.mark.asyncio +async def test_new_row_is_not_double_counted_when_the_batch_logs_already_flushed(): + """The spend log writer drains on a ~2s poll while window increments flush + on the ~10s batch tick, so a new row is normally seeded from a table that + already holds this batch's rows. Counting them in both places is what made + a fresh row land at exactly twice the true spend.""" + db = _FakeDB(existing_rows=[]) + already_flushed = _SpendLogsFake( + rows=( + ("req-1", 0.000047, BATCH_STARTED_AT), + ("req-2", 0.000047, BATCH_STARTED_AT + timedelta(seconds=1)), + ("req-3", 0.000047, BATCH_STARTED_AT + timedelta(seconds=2)), + ), + ) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),), + spend_logs_aggregate=already_flushed, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(0.000141) + + +@pytest.mark.asyncio +async def test_new_row_still_covers_spend_that_predates_the_batch(): + """The exclusion must not throw away the pre-existing spend the seed is for.""" + db = _FakeDB(existing_rows=[]) + spend_logs = _SpendLogsFake(rows=(("older", 0.5, BEFORE_BATCH), ("req-1", 0.000047, BATCH_STARTED_AT))) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("req-1",), 0.000047),), + spend_logs_aggregate=spend_logs, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(0.500047) + + +@pytest.mark.asyncio +async def test_replayed_request_id_cannot_erase_historical_spend_from_the_seed(): + """request_id can be chosen by the client via x-litellm-call-id. A request + that replays an id from before this batch writes no new LiteLLM_SpendLogs + row (the insert skips duplicates), so the seed must keep counting the + historical row that id belongs to; only its increment is new.""" + db = _FakeDB(existing_rows=[]) + spend_logs = _SpendLogsFake(rows=(("replayed", 0.5, BEFORE_BATCH),)) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("replayed",), 0.000047),), + spend_logs_aggregate=spend_logs, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(0.500047) + + +@pytest.mark.asyncio +async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): + """The other side of the race: rows absent from the aggregate are still + counted exactly once, by their increment.""" + db = _FakeDB(existing_rows=[]) + nothing_flushed = _SpendLogsFake(rows=()) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),), + spend_logs_aggregate=nothing_flushed, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(0.000141) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entity_type, expected_column", + [("key", "api_key = $1"), ("team", "team_id = $1")], +) +async def test_seed_aggregate_sql_excludes_the_request_ids_only_within_the_batch_start_bound( + entity_type, expected_column +): + db = _FakeDB(existing_rows=[{"total": 1.25}]) + + total = await spend_logs_total_excluding( + prisma_client=_FakePrismaClient(db), + entity_type=entity_type, + entity_id="e1", + window_start=WINDOW_A, + exclude_request_ids=("req-1", "req-2"), + exclude_started_at=BATCH_STARTED_AT, + ) + + assert total == pytest.approx(1.25) + ((query, params),) = db.query_raw_calls + normalized = " ".join(query.split()) + assert expected_column in normalized + assert "NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" in normalized + assert 'FROM "LiteLLM_SpendLogs"' in normalized + # startTime is TIMESTAMP(3): the bound is floored to the second so the + # batch's own earliest row cannot round under it. + assert params == ("e1", WINDOW_A, ("req-1", "req-2"), datetime(2026, 8, 10, 12, 0, 0)) + # The ids are bound, never spliced into the statement. + assert "req-1" not in query + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "exclude_request_ids, exclude_started_at", + [(("req-1",), None), ((), BATCH_STARTED_AT)], +) +async def test_seed_aggregate_excludes_nothing_without_both_ids_and_a_start_bound( + exclude_request_ids, exclude_started_at +): + """Ids without a start bound would reopen the replayed-id hole, so the + seed counts everything instead; at worst that over-counts one batch.""" + db = _FakeDB(existing_rows=[{"total": 1.25}]) + + total = await spend_logs_total_excluding( + prisma_client=_FakePrismaClient(db), + entity_type="key", + entity_id="e1", + window_start=WINDOW_A, + exclude_request_ids=exclude_request_ids, + exclude_started_at=exclude_started_at, + ) + + assert total == pytest.approx(1.25) + ((query, params),) = db.query_raw_calls + assert "request_id" not in query + assert params == ("e1", WINDOW_A) + + +@pytest.mark.asyncio +async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs_column(): + db = _FakeDB(existing_rows=[]) + + total = await spend_logs_total_excluding( + prisma_client=_FakePrismaClient(db), + entity_type="user", + entity_id="u1", + window_start=WINDOW_A, + exclude_request_ids=(), + exclude_started_at=None, + ) + + assert total is None + assert db.query_raw_calls == [] + + +@pytest.mark.asyncio +async def test_seed_aggregate_treats_an_entity_with_no_rows_as_zero(): + db = _FakeDB(existing_rows=[]) + + total = await spend_logs_total_excluding( + prisma_client=_FakePrismaClient(db), + entity_type="key", + entity_id="k-unknown", + window_start=WINDOW_A, + exclude_request_ids=(), + exclude_started_at=None, + ) + + assert total == 0.0 diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index b1f647bdd3d..d28cf8c9c6a 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -4,8 +4,8 @@ import json import re - from collections.abc import Callable +from contextlib import asynccontextmanager from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, call, patch @@ -15,6 +15,9 @@ from redis.exceptions import DataError import litellm from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, +) @pytest.mark.asyncio @@ -64,9 +67,7 @@ async def test_daily_spend_tracking_with_disabled_spend_logs(): assert db_writer.add_spend_log_transaction_to_daily_user_transaction.called # Verify the payload passed to add_spend_log_transaction_to_daily_user_transaction - call_args = ( - db_writer.add_spend_log_transaction_to_daily_user_transaction.call_args[1] - ) + call_args = db_writer.add_spend_log_transaction_to_daily_user_transaction.call_args[1] assert "payload" in call_args assert call_args["payload"]["spend"] == 0.1 assert call_args["payload"]["model"] == "gpt-4" @@ -406,7 +407,7 @@ async def test_update_daily_spend_sorting(): # fields, but entity_id is sufficient to test sorting. daily_spend_transactions = { f"test_key_{i}": { - "user_id": f"user{60-i}", # user60 ... user11, reverse order + "user_id": f"user{60 - i}", # user60 ... user11, reverse order "date": "2024-01-01", "api_key": "test-api-key", "model": "gpt-4", @@ -985,9 +986,9 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i transaction_dict = call[1]["update"] # Each transaction should have one key with the format tag_date_api_key_model_provider for key, transaction in transaction_dict.items(): - assert ( - transaction["request_id"] == request_id - ), f"request_id should be {request_id} but got {transaction.get('request_id')}" + assert transaction["request_id"] == request_id, ( + f"request_id should be {request_id} but got {transaction.get('request_id')}" + ) @pytest.mark.asyncio @@ -1213,21 +1214,15 @@ async def test_add_spend_log_transaction_to_daily_agent_transaction_calls_common } writer.daily_agent_spend_update_queue.add_update = AsyncMock() - original_common_helper = ( - writer._common_add_spend_log_transaction_to_daily_transaction - ) - writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock( - wraps=original_common_helper - ) + original_common_helper = writer._common_add_spend_log_transaction_to_daily_transaction + writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock(wraps=original_common_helper) await writer.add_spend_log_transaction_to_daily_agent_transaction( payload=payload, prisma_client=mock_prisma, ) - assert ( - writer._common_add_spend_log_transaction_to_daily_transaction.await_count == 1 - ) + assert writer._common_add_spend_log_transaction_to_daily_transaction.await_count == 1 @pytest.mark.asyncio @@ -1382,6 +1377,7 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): Test that when batch upsert fails, the exception is properly re-raised after logging. This ensures that error handling continues to work correctly upstream. """ + def raise_connection_lost(): raise ValueError("Database connection lost") @@ -1562,9 +1558,7 @@ async def test_update_database_creates_single_task(): patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), - patch( - "litellm.proxy.db.db_spend_update_writer.asyncio.create_task" - ) as mock_create_task, + patch("litellm.proxy.db.db_spend_update_writer.asyncio.create_task") as mock_create_task, ): await db_writer.update_database( token="test-token", @@ -1663,9 +1657,7 @@ async def test_daily_agent_receives_deepcopied_payload(): db_writer._update_agent_db = AsyncMock() db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() - db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock( - side_effect=capture_agent_payload - ) + db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock(side_effect=capture_agent_payload) db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock() @@ -1727,8 +1719,8 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock() # Return all-None tuple (no data to commit); the pipeline yields 6 slots - mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = ( - AsyncMock(return_value=(None, None, None, None, None, None)) + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, None) ) db_writer.redis_update_buffer = mock_redis_update_buffer @@ -1782,7 +1774,7 @@ async def test_commit_with_redis_requeues_all_on_db_failure(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( - return_value=(db_spend, daily_user, None, None, None, None) + return_value=(db_spend, daily_user, None, None, None, None, None) ) mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() db_writer.redis_update_buffer = mock_redis_update_buffer @@ -1837,7 +1829,7 @@ async def test_commit_with_redis_only_requeues_failed_category(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( - return_value=(db_spend, daily_user, None, None, None, None) + return_value=(db_spend, daily_user, None, None, None, None, None) ) mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() db_writer.redis_update_buffer = mock_redis_update_buffer @@ -1885,7 +1877,7 @@ async def test_commit_with_redis_no_requeue_on_success(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( - return_value=(db_spend, None, None, None, None, None) + return_value=(db_spend, None, None, None, None, None, None) ) mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() db_writer.redis_update_buffer = mock_redis_update_buffer @@ -2156,9 +2148,7 @@ async def test_update_database_does_not_deepcopy_on_request_path(): db_writer._update_org_db = AsyncMock() db_writer._update_tag_db = AsyncMock() db_writer._update_agent_db = AsyncMock() - db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock( - side_effect=capture_batch_payload - ) + db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock(side_effect=capture_batch_payload) db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() @@ -2250,9 +2240,7 @@ async def test_spend_update_path_never_queries_user_cache_with_none_user_id(): db_writer = DBSpendUpdateWriter() strict_redis_backed_cache = MagicMock() - strict_redis_backed_cache.async_get_cache = AsyncMock( - side_effect=DataError("Invalid input of type: 'NoneType'") - ) + strict_redis_backed_cache.async_get_cache = AsyncMock(side_effect=DataError("Invalid input of type: 'NoneType'")) with ( patch.object(litellm, "max_budget", 0), @@ -2380,8 +2368,7 @@ async def test_daily_transaction_carries_compression_saved_tokens(): cache_write_cost = model_info.get("cache_creation_input_token_cost") or input_cost assert transaction["compression_savings_spend"] == pytest.approx(7600 * input_cost) assert transaction["prompt_caching_savings_spend"] == pytest.approx( - 40 * max(input_cost - cache_read_cost, 0.0) - - 15 * (cache_write_cost - input_cost) + 40 * max(input_cost - cache_read_cost, 0.0) - 15 * (cache_write_cost - input_cost) ) assert transaction["compression_savings_spend"] > 0 assert transaction["prompt_caching_savings_spend"] > 0 @@ -2421,6 +2408,304 @@ async def test_daily_transaction_compression_saved_tokens_zero_when_absent(): assert transaction["prompt_caching_savings_spend"] == 0 +# --------------------------------------------------------------------------- +# Budget window spend flush (LiteLLM_BudgetWindowSpend) +# --------------------------------------------------------------------------- + + +class _WindowSpendFakeBatcher: + def __init__(self): + self.calls = [] + + def execute_raw(self, query, *args): + self.calls.append((query, args)) + + +class _WindowSpendFakeDB: + """Minimal prisma_client.db that records the raw statements it is handed.""" + + def __init__(self, existing_rows=None): + self.existing_rows = existing_rows or [] + self.query_raw_calls = [] + self.batcher = _WindowSpendFakeBatcher() + + async def query_raw(self, query, *args): + self.query_raw_calls.append((query, args)) + if "LiteLLM_BudgetWindowSpend" in query: + return self.existing_rows + return [] + + @asynccontextmanager + async def _tx(self): + yield self + + def tx(self, timeout=None): + return self._tx() + + @asynccontextmanager + async def _batch(self): + yield self.batcher + + def batch_(self): + return self._batch() + + +class _WindowSpendFakePrisma: + def __init__(self, db): + self.db = db + + +def _window_spend_upserts(db): + return [params for query, params in db.batcher.calls if "LiteLLM_BudgetWindowSpend" in query] + + +@pytest.mark.asyncio +async def test_window_spend_queue_is_flushed_without_redis_buffer(): + """The in-memory window queue must reach the DB on the same scheduler tick + as the other spend queues when the Redis buffer is off.""" + db_writer = DBSpendUpdateWriter() + await db_writer.window_spend_update_queue.add_update( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=0.5, + ) + ) + db = _WindowSpendFakeDB( + existing_rows=[{"entity_type": "key", "entity_id": "hashed-token", "window_duration": "30d"}] + ) + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + upserts = _window_spend_upserts(db) + assert len(upserts) == 1 + assert upserts[0][0] == "key" + assert upserts[0][1] == "hashed-token" + assert upserts[0][2] == "30d" + assert upserts[0][5] == pytest.approx(0.5) + assert db_writer.window_spend_update_queue.update_queue.qsize() == 0 + + +@pytest.mark.asyncio +async def test_window_spend_queue_is_handed_to_the_redis_buffer(): + """Multi-pod deployments buffer through Redis, so the window queue has to + ride the same rpush path as the daily queues.""" + db_writer = DBSpendUpdateWriter() + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, None) + ) + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + stored = mock_redis_update_buffer.store_in_memory_spend_updates_in_redis.call_args[1] + assert stored["window_spend_update_queue"] is db_writer.window_spend_update_queue + + +@pytest.mark.asyncio +async def test_window_spend_transactions_from_redis_are_committed_by_the_lock_winner(): + db_writer = DBSpendUpdateWriter() + window_transactions = ( + build_window_spend_transaction( + entity_type="team", + entity_id="team-1", + window_duration="7d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=2.0, + ), + ) + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, window_transactions) + ) + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + db = _WindowSpendFakeDB(existing_rows=[{"entity_type": "team", "entity_id": "team-1", "window_duration": "7d"}]) + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + upserts = _window_spend_upserts(db) + assert len(upserts) == 1 + assert upserts[0][:3] == ("team", "team-1", "7d") + assert upserts[0][5] == pytest.approx(2.0) + + +@pytest.mark.asyncio +async def test_window_spend_transactions_are_not_committed_without_the_pod_lock(): + """Every pod buffers to Redis but only the lock winner may drain it.""" + db_writer = DBSpendUpdateWriter() + mock_redis_update_buffer = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=False) + db = _WindowSpendFakeDB() + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline.assert_not_called() + assert _window_spend_upserts(db) == [] + + +@pytest.mark.asyncio +async def test_failed_window_spend_commit_requeues_the_increments_and_continues_the_flush(): + """Budget enforcement trusts a current window row without reconciling it + against LiteLLM_SpendLogs, so a dropped increment would let the key spend + past its limit after the next reseed. The increments must go back on the + queue, and the tool registry flush must still run.""" + db_writer = DBSpendUpdateWriter() + transaction = build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=0.5, + request_id="req-1", + ) + await db_writer.window_spend_update_queue.add_update(transaction) + db = _WindowSpendFakeDB() + db.query_raw = AsyncMock(side_effect=Exception("connection reset")) + db_writer._flush_tool_discovery_queue = AsyncMock() + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + db_writer._flush_tool_discovery_queue.assert_called_once() + requeued = await db_writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() + assert requeued == (transaction,) + + +@pytest.mark.asyncio +async def test_failed_window_spend_commit_from_redis_is_restored_to_redis(): + """The Redis drain is destructive, so a failed window commit has to push + the popped increments back exactly like the other spend categories.""" + db_writer = DBSpendUpdateWriter() + window_transactions = ( + build_window_spend_transaction( + entity_type="team", + entity_id="team-1", + window_duration="7d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=2.0, + request_id="req-1", + ), + ) + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, window_transactions) + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + db = _WindowSpendFakeDB() + db.query_raw = AsyncMock(side_effect=Exception("connection reset")) + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + assert _window_spend_upserts(db) == [] + mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once_with( + window_spend_update_transactions=window_transactions + ) + db_writer.pod_lock_manager.release_lock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_update_database_returns_the_spend_log_request_id(): + """The budget-window seed excludes the log rows its increments already + cover, so the caller needs the id this call was recorded under. It cannot + be re-derived: cache hits append time.time() to the id.""" + db_writer = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._enqueue_tool_usage_transaction = AsyncMock() + + with ( + patch.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam + "litellm.proxy.proxy_server", + disable_spend_logs=False, + prisma_client=MagicMock(), + litellm_proxy_budget_name="test-budget", + ) + ): + request_id = await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id="test-team", + org_id=None, + kwargs={"model": "gpt-4", "custom_llm_provider": "openai", "litellm_call_id": "call-xyz"}, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.1, + ) + await asyncio.sleep(0) + + assert request_id is not None + # Same id the spend log row was queued under. + assert request_id == db_writer._insert_spend_log_to_db.call_args[1]["payload"]["request_id"] + + +@pytest.mark.asyncio +async def test_update_database_returns_none_when_the_payload_cannot_be_built(): + db_writer = DBSpendUpdateWriter() + + with ( + patch.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam + "litellm.proxy.proxy_server", + disable_spend_logs=False, + prisma_client=MagicMock(), + litellm_proxy_budget_name="test-budget", + ), + patch( # test-quality-ok: the payload builder is called by name inside update_database; no injection seam + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + side_effect=Exception("payload boom"), + ), + ): + request_id = await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id="test-team", + org_id=None, + kwargs={"model": "gpt-4"}, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.1, + ) + + assert request_id is None + + @pytest.mark.asyncio async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at(): """Spend flushes must leave settings_updated_at alone, or it decays into @@ -2721,9 +3006,7 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey "call_type, expects_flush", [("aresponses", True), ("responses", True), ("acompletion", False)], ) -async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls( - call_type: str, expects_flush: bool -): +async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(call_type: str, expects_flush: bool): """ A `previous_response_id` chained straight off the previous turn reads the DB, so a Responses row cannot sit in this worker's queue until the monitor's next poll. @@ -2753,9 +3036,7 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls( pytest.param("", True, id="injected-before-a-deployment-was-chosen"), ], ) -async def test_caching_savings_are_attributed_to_the_deployment_that_was_injected( - injected_deployment, attributed -): +async def test_caching_savings_are_attributed_to_the_deployment_that_was_injected(injected_deployment, attributed): """Retries, same-group failover and cross-model-group fallbacks all reuse one metadata bucket and one litellm_call_id, so a marker written by the leg that injected is visible to every sibling and nothing request-scoped can tell them apart. diff --git a/tests/test_litellm/proxy/db/test_model_access_group_spend.py b/tests/test_litellm/proxy/db/test_model_access_group_spend.py index 79dffc29e7e..d2d079bb0e4 100644 --- a/tests/test_litellm/proxy/db/test_model_access_group_spend.py +++ b/tests/test_litellm/proxy/db/test_model_access_group_spend.py @@ -414,12 +414,14 @@ async def test_redis_buffer_requeues_access_group_transactions_as_queue_items(): daily_org_spend_update_transactions=None, daily_end_user_spend_update_transactions=None, daily_agent_spend_update_transactions=None, + window_spend_update_transactions=None, spend_update_queue=queue, daily_spend_update_queue=daily_queue, daily_team_spend_update_queue=daily_queue, daily_org_spend_update_queue=daily_queue, daily_end_user_spend_update_queue=daily_queue, daily_agent_spend_update_queue=daily_queue, + window_spend_update_queue=None, ) updates = await _drain(queue) diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py new file mode 100644 index 00000000000..816f9ae72f4 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -0,0 +1,250 @@ +"""Window-spend reads in ``SpendCounterReseed``. + +The maintained ``LiteLLM_BudgetWindowSpend`` row replaces a per-request +``LiteLLM_SpendLogs`` range scan, so these pin *when* the aggregate is still +allowed to run: only when the row is missing or belongs to an older window. +""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace + +import pytest + +from litellm.caching.dual_cache import DualCache +from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + +WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) + + +class _FakeWindowSpendTable: + def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None: + self._row = row + self._error = error + self.where_clauses: list[dict] = [] + + async def find_unique(self, where: dict): + self.where_clauses.append(where) + if self._error is not None: + raise self._error + return self._row + + +class _FakeSpendLogsTable: + def __init__(self, total: float) -> None: + self._total = total + self.call_count = 0 + + async def group_by(self, by: list[str], where: dict, sum: dict): + self.call_count += 1 + return [{by[0]: where.get(by[0]), "_sum": {"spend": self._total}}] + + +class _FakePrismaClient: + def __init__( + self, + row: SimpleNamespace | None = None, + spend_logs_total: float = 0.0, + error: Exception | None = None, + ) -> None: + self.db = SimpleNamespace( + litellm_budgetwindowspend=_FakeWindowSpendTable(row=row, error=error), + litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total), + ) + + +def _row(window_start: datetime, spend: float) -> SimpleNamespace: + return SimpleNamespace(window_start=window_start, spend=spend) + + +@pytest.mark.asyncio +async def test_window_from_table_reads_row_by_primary_key(): + """The lookup must use the table's own entity_type values ("key"), not the + "Key"/"Team" labels the counter keys and spend-log aggregates use.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result == 4.5 + assert prisma.db.litellm_budgetwindowspend.where_clauses == [ + { + "entity_type_entity_id_window_duration": { + "entity_type": "key", + "entity_id": "tok-1", + "window_duration": "30d", + } + } + ] + + +@pytest.mark.asyncio +async def test_window_from_table_maps_team_entity_type(): + prisma = _FakePrismaClient(row=_row(WINDOW_START, 9.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Team", + entity_id="team-1", + window_duration="1d", + expected_window_start=WINDOW_START, + ) + + assert result == 9.0 + inner = prisma.db.litellm_budgetwindowspend.where_clauses[0]["entity_type_entity_id_window_duration"] + assert inner["entity_type"] == "team" + + +@pytest.mark.asyncio +async def test_window_from_table_trusts_row_newer_than_expected_window(): + """Regression: a pod holding a stale ``reset_at`` computes an expected start + behind a window another pod already rolled. Trusting only an exact match + would make it re-add the previous window's spend to the current one.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START + timedelta(days=1), 2.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result == 2.0 + + +@pytest.mark.asyncio +async def test_window_from_table_rejects_row_from_previous_window(): + prisma = _FakePrismaClient(row=_row(WINDOW_START - timedelta(seconds=1), 99.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_window_from_table_treats_naive_row_timestamp_as_utc(): + """The column is ``timestamp(3)``, so a driver that hands back a naive value + must still compare against the tz-aware expected start.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START.replace(tzinfo=None), 3.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result == 3.0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prisma, entity_type", + [ + (_FakePrismaClient(row=None), "Key"), + (_FakePrismaClient(row=_row(WINDOW_START, 1.0)), "User"), + (_FakePrismaClient(error=RuntimeError("connection reset")), "Key"), + (None, "Key"), + ], +) +async def test_window_from_table_returns_none_without_a_usable_row(prisma, entity_type): + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type=entity_type, + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_window_from_db_prefers_the_row_over_the_spend_logs_aggregate(): + """The aggregate range-scans an unindexed table; a current row must keep it + from running at all.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + window_start=WINDOW_START, + ) + + assert result == 4.5 + assert prisma.db.litellm_spendlogs.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "row", + [None, _row(WINDOW_START - timedelta(seconds=1), 99.0)], + ids=["missing_row", "previous_window_row"], +) +async def test_window_from_db_falls_back_to_spend_logs(row): + prisma = _FakePrismaClient(row=row, spend_logs_total=7.25) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + window_start=WINDOW_START, + ) + + assert result == 7.25 + assert prisma.db.litellm_spendlogs.call_count == 1 + + +@pytest.mark.asyncio +async def test_window_from_db_without_a_duration_skips_the_row_lookup(): + """Callers that cannot name the window (no PK) keep the pre-table behavior.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=7.25) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration=None, + window_start=WINDOW_START, + ) + + assert result == 7.25 + assert prisma.db.litellm_budgetwindowspend.where_clauses == [] + + +@pytest.mark.asyncio +async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0) + cache = DualCache() + counter_key = "spend:key:tok-1:window:30d" + + result = await SpendCounterReseed.coalesced_window( + prisma_client=prisma, + spend_counter_cache=cache, + counter_key=counter_key, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + window_start=WINDOW_START, + ) + + assert result == 4.5 + assert cache.in_memory_cache.get_cache(key=counter_key) == 4.5 + assert prisma.db.litellm_spendlogs.call_count == 0 diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 5160dd3431e..2a540f4f522 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -567,9 +567,12 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re @pytest.mark.asyncio async def test_update_database_and_spend_counters_updates_counters_after_db_update(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( + return_value="chatcmpl-abc123" + ) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} + start_time = datetime.now() await _update_database_and_spend_counters( proxy_logging_obj=proxy_logging_obj, @@ -581,7 +584,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda org_id="test_org_id", kwargs={}, completion_response=None, - start_time=datetime.now(), + start_time=start_time, end_time=datetime.now(), response_cost=0.2, budget_reservation=budget_reservation, @@ -599,6 +602,8 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda budget_reservation=budget_reservation, end_user_id="test_end_user_id", tags=["tag-a"], + request_id="chatcmpl-abc123", + request_started_at=start_time, model_access_groups=("premium",), ) @@ -1879,6 +1884,61 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( ) +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_forwards_the_spend_log_request_id(): + """The budget-window flush excludes the log rows its increments already + cover. That only works if the id update_database recorded the row under is + handed to the counter update, so this seam is load-bearing.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( + return_value="chatcmpl-abc123" + ) + increment_spend_counters = AsyncMock() + + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + assert increment_spend_counters.await_args.kwargs["request_id"] == "chatcmpl-abc123" + + +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_forwards_a_missing_request_id_as_none(): + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(return_value=None) + increment_spend_counters = AsyncMock() + + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + assert increment_spend_counters.await_args.kwargs["request_id"] is None + + class _FakeDeploymentLookup: """Deployment lookup returning the access groups each deployment declares.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 51980342a1d..fb3de990deb 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -271,6 +271,81 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): ) +def _make_window_spend_prisma(row=None, spend_logs_total=0.0): + prisma = MagicMock() + prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=row) + prisma.db.litellm_spendlogs.group_by = AsyncMock( + return_value=[{"api_key": "tok", "_sum": {"spend": spend_logs_total}}] + ) + return prisma + + +@pytest.mark.asyncio +async def test_get_current_spend_floors_window_against_maintained_row(monkeypatch): + """The floor re-check runs every few seconds per pod, so the window branch + must read the maintained row and leave the unindexed spend-logs scan alone.""" + from datetime import timezone + from types import SimpleNamespace + + window_start = datetime(2026, 1, 1, tzinfo=timezone.utc) + fake_prisma = _make_window_spend_prisma( + row=SimpleNamespace(window_start=window_start, spend=15.0), + spend_logs_total=100.0, + ) + fake_cache = _make_spend_counter_cache(redis_get_value=2.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + + counter_key = "spend:key:tok:window:7d" + result = await ps.get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + max_budget=10.0, + window_entity_type="Key", + window_entity_id="tok", + window_duration="7d", + window_start=window_start, + ) + + assert result == 15.0 + fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited() + fake_cache.redis_cache.async_set_max.assert_awaited_once_with( + key=counter_key, value=15.0 + ) + + +@pytest.mark.asyncio +async def test_get_current_spend_floors_window_against_logs_when_row_stale(monkeypatch): + """A row left behind at a crossed window boundary must not be read as the + current window's spend; the aggregate stays the fallback.""" + from datetime import timedelta, timezone + from types import SimpleNamespace + + window_start = datetime(2026, 1, 8, tzinfo=timezone.utc) + fake_prisma = _make_window_spend_prisma( + row=SimpleNamespace( + window_start=window_start - timedelta(days=7), spend=999.0 + ), + spend_logs_total=15.0, + ) + fake_cache = _make_spend_counter_cache(redis_get_value=2.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + + result = await ps.get_current_spend( + counter_key="spend:key:tok:window:7d", + fallback_spend=0.0, + max_budget=10.0, + window_entity_type="Key", + window_entity_id="tok", + window_duration="7d", + window_start=window_start, + ) + + assert result == 15.0 + fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once() + + @pytest.mark.asyncio async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypatch): """With fail_closed_budget_enforcement on, an admit decision backed only by a @@ -895,6 +970,7 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=datetime(2024, 1, 1), increment=5.0, ) @@ -922,6 +998,7 @@ async def test_init_and_increment_window_spend_counter_missing_window_start_inva counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=None, increment=5.0, ) @@ -1059,6 +1136,7 @@ async def test_ensure_window_spend_counter_initialized_warm_returns_true(monkeyp counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=datetime(2024, 1, 1), ) @@ -1091,6 +1169,7 @@ async def test_ensure_window_spend_counter_initialized_db_failure_invalid_return counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=datetime(2024, 1, 1), ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 71d4d5f4874..1de3ed6e56d 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import importlib import json import os @@ -7816,6 +7817,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): counter_cache = DualCache() window_start = datetime.now(timezone.utc) - timedelta(hours=1) fake_prisma = MagicMock() + fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[{"api_key": "key-window", "_sum": {"spend": 2.25}}] ) @@ -7830,6 +7832,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): counter_key="spend:key:key-window:window:1h", entity_type="Key", entity_id="key-window", + window_duration="1h", window_start=window_start, increment=0.5, ) @@ -7933,6 +7936,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() + fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[{"api_key": "key-window-stale-local", "_sum": {"spend": 2.25}}] ) @@ -7947,6 +7951,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): counter_key=counter_key, entity_type="Key", entity_id="key-window-stale-local", + window_duration="1h", window_start=window_start, increment=0.5, ) @@ -7995,6 +8000,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() + fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[{"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}}] ) @@ -8009,6 +8015,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() counter_key=counter_key, entity_type="Key", entity_id="key-window-concurrent-seed", + window_duration="1h", window_start=window_start, increment=0.5, ) @@ -8041,6 +8048,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): counter_key="spend:key:key-invalid-window:window:not-a-duration", entity_type="Key", entity_id="key-invalid-window", + window_duration="not-a-duration", window_start=None, increment=0.5, ) @@ -8068,6 +8076,7 @@ async def test_window_spend_counter_does_not_seed_zero_when_db_unavailable(): counter_key=counter_key, entity_type="Key", entity_id="key-window-db-unavailable", + window_duration="1h", window_start=datetime.now(timezone.utc) - timedelta(hours=1), ) @@ -11200,6 +11209,289 @@ def test_startup_is_silent_when_mock_testing_params_disabled(caplog): assert MOCK_TESTING_CONFIG_KEY not in caplog.text +# --------------------------------------------------------------------------- +# Budget window spend row enqueue (LiteLLM_BudgetWindowSpend writer) +# --------------------------------------------------------------------------- + + +@contextlib.contextmanager +def _window_spend_enqueue_env(cached_objects: dict): + """Point increment_spend_counters at throwaway caches and a real + WindowSpendUpdateQueue, and hand back the queue to inspect.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + ) + import litellm.proxy.proxy_server as ps + + user_api_key_cache = MagicMock() + user_api_key_cache.async_get_cache = AsyncMock(side_effect=lambda key, **_: cached_objects.get(key)) + + queue = WindowSpendUpdateQueue() + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.window_spend_update_queue = queue + + originals = ( + ps.user_api_key_cache, + ps.spend_counter_cache, + ps.prisma_client, + ps.proxy_logging_obj, + ) + ps.user_api_key_cache = user_api_key_cache + ps.spend_counter_cache = DualCache() + ps.prisma_client = None + ps.proxy_logging_obj = proxy_logging_obj + try: + yield queue + finally: + ( + ps.user_api_key_cache, + ps.spend_counter_cache, + ps.prisma_client, + ps.proxy_logging_obj, + ) = originals + + +async def _drain(queue): + return list(await queue.flush_and_get_aggregated_window_spend_transactions()) + + +@pytest.mark.asyncio +async def test_key_window_spend_row_is_enqueued_with_the_actual_cost(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert len(enqueued) == 1 + assert enqueued[0]["entity_type"] == "key" + assert enqueued[0]["entity_id"] == "hashed-token" + assert enqueued[0]["window_duration"] == "30d" + assert enqueued[0]["spend"] == pytest.approx(0.25) + assert enqueued[0]["window_start"] == (reset_at - timedelta(days=30)).astimezone(timezone.utc).replace( + tzinfo=None + ).isoformat(timespec="microseconds") + + +@pytest.mark.asyncio +async def test_team_window_spend_row_is_enqueued(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=3) + team_obj = MagicMock() + team_obj.budget_limits = [ + {"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"team_id:team-1": team_obj}) as queue: + await increment_spend_counters( + token=None, team_id="team-1", user_id=None, response_cost=1.5 + ) + enqueued = await _drain(queue) + + assert len(enqueued) == 1 + assert enqueued[0]["entity_type"] == "team" + assert enqueued[0]["entity_id"] == "team-1" + assert enqueued[0]["window_duration"] == "7d" + assert enqueued[0]["spend"] == pytest.approx(1.5) + + +@pytest.mark.asyncio +async def test_window_spend_row_is_enqueued_even_when_the_counter_was_reserved(): + """A reservation only pre-charged the cache counter with an estimate; the + row still owes the actual cost, so the enqueue must not be skipped.""" + from litellm.proxy.proxy_server import increment_spend_counters + import litellm.proxy.spend_tracking.budget_reservation as br + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + reservation = { + "entries": [ + {"counter_key": "spend:key:hashed-token", "reserved": 1.0}, + {"counter_key": "spend:key:hashed-token:window:30d", "reserved": 1.0}, + ] + } + + original_reconcile = br.reconcile_budget_reservation + br.reconcile_budget_reservation = AsyncMock(return_value=None) + try: + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", + team_id=None, + user_id=None, + response_cost=0.25, + budget_reservation=reservation, + ) + enqueued = await _drain(queue) + finally: + br.reconcile_budget_reservation = original_reconcile + + assert len(enqueued) == 1 + assert enqueued[0]["spend"] == pytest.approx(0.25) + + +@pytest.mark.asyncio +async def test_sliding_window_without_reset_at_is_not_enqueued(): + """Windows with no reset_at slide with wall clock, so window_start moves on + every request and no single row can represent them; the read path keeps + using its LiteLLM_SpendLogs fallback instead.""" + from litellm.proxy.proxy_server import increment_spend_counters + + key_obj = MagicMock() + key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0}] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert enqueued == [] + + +@pytest.mark.asyncio +async def test_each_configured_window_gets_its_own_row_enqueue(): + from litellm.proxy.proxy_server import increment_spend_counters + + now = datetime.now(timezone.utc) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "1d", "max_budget": 5.0, "reset_at": (now + timedelta(hours=5)).isoformat()}, + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": (now + timedelta(days=10)).isoformat()}, + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert sorted(item["window_duration"] for item in enqueued) == ["1d", "30d"] + assert all(item["spend"] == pytest.approx(0.25) for item in enqueued) + + +@pytest.mark.asyncio +async def test_no_window_spend_row_enqueued_without_budget_limits(): + from litellm.proxy.proxy_server import increment_spend_counters + + key_obj = MagicMock() + key_obj.budget_limits = None + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert enqueued == [] + + +@pytest.mark.asyncio +async def test_window_spend_row_carries_the_spend_log_request_id(): + """The flush excludes these ids from its one-time seed, so the id threaded + here has to be the same one the LiteLLM_SpendLogs row was written under.""" + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", + team_id=None, + user_id=None, + response_cost=0.25, + request_id="chatcmpl-abc123", + ) + enqueued = await _drain(queue) + + assert enqueued[0]["request_ids"] == ("chatcmpl-abc123",) + + +@pytest.mark.asyncio +async def test_window_spend_row_carries_the_request_start_time(): + """The seed only excludes a batch id whose LiteLLM_SpendLogs.startTime is at + or after this, so it must be the same start the spend log was written with.""" + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", + team_id=None, + user_id=None, + response_cost=0.25, + request_id="chatcmpl-abc123", + request_started_at=datetime(2026, 8, 10, 12, 0, 0, 500_000, tzinfo=timezone.utc), + ) + enqueued = await _drain(queue) + + assert enqueued[0]["started_at"] == "2026-08-10T12:00:00.500000" + + +@pytest.mark.asyncio +async def test_window_spend_row_without_a_request_id_excludes_nothing(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert enqueued[0]["request_ids"] == () + + +@pytest.mark.asyncio +async def test_team_window_spend_row_carries_the_request_id(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=3) + team_obj = MagicMock() + team_obj.budget_limits = [ + {"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"team_id:team-1": team_obj}) as queue: + await increment_spend_counters( + token=None, + team_id="team-1", + user_id=None, + response_cost=1.5, + request_id="chatcmpl-team", + ) + enqueued = await _drain(queue) + + assert enqueued[0]["request_ids"] == ("chatcmpl-team",) + + def _mock_startup_prisma_client(health_check_error=None, connect_error=None): client = MagicMock() client.connect = AsyncMock(side_effect=connect_error)