diff --git a/litellm/constants.py b/litellm/constants.py index 0c7316455d6..2c8c2918a30 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -270,6 +270,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/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 450dc12a9ff..0fb82c75274 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -2,16 +2,18 @@ import asyncio import json import time from collections.abc import Callable, Mapping, Sequence -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Final, Literal, Protocol, TypeVar import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( LiteLLM_BudgetTableFull, LiteLLM_EndUserTable, + Litellm_EntityType, LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken, @@ -21,6 +23,7 @@ from litellm.proxy.common_utils.timezone_utils import ( compute_budget_reset_at, get_budget_reset_settings, ) +from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( @@ -782,6 +785,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") @@ -796,11 +802,56 @@ class ResetBudgetJob: await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=0.0) 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, + ) + async def reset_budget_windows(self) -> None: """ For keys and teams with budget_limits, reset any individual windows where @@ -836,6 +887,9 @@ class ResetBudgetJob: spend_counter_cache, now, self.reset_settings, + prisma_client=self.prisma_client, + entity_type=Litellm_EntityType.KEY, + entity_id=row["token"], ): changed = True if changed: @@ -865,6 +919,9 @@ class ResetBudgetJob: spend_counter_cache, now, self.reset_settings, + prisma_client=self.prisma_client, + entity_type=Litellm_EntityType.TEAM, + entity_id=row["team_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..0308651f4e3 --- /dev/null +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -0,0 +1,266 @@ +""" +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 +miss increments still queued on another. That is bounded by a single flush +interval, happens at most once per window row, and errs toward under-counting +for that interval only. +""" + +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[]))" +) + +_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[]))" +) + +_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], + ) -> 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], +) -> 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. + """ + if entity_type == Litellm_EntityType.KEY.value: + rows = await prisma_client.db.query_raw( + _SEED_FROM_SPEND_LOGS_KEY_SQL, entity_id, window_start, tuple(exclude_request_ids) + ) + elif entity_type == Litellm_EntityType.TEAM.value: + rows = await prisma_client.db.query_raw( + _SEED_FROM_SPEND_LOGS_TEAM_SQL, entity_id, window_start, tuple(exclude_request_ids) + ) + else: + return None + if not rows: + return 0.0 + return float(rows[0].get("total") or 0.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"], + ) + return float(base or 0.0) + + +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: + async with 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 4e5c5daaa32..4e0137e33c6 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,7 +12,7 @@ import os import random import time import traceback -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload @@ -50,6 +50,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, @@ -132,6 +136,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 @@ -147,7 +152,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, @@ -164,7 +173,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: @@ -232,6 +241,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 " @@ -244,6 +254,7 @@ class DBSpendUpdateWriter: org_id, end_user_id, ) + return None async def _enqueue_tool_usage_transaction( self, @@ -831,6 +842,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 @@ -847,6 +859,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() if db_spend_update_transactions is not None: @@ -906,6 +919,11 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_agent_spend_update_transactions, ) + 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, + ) except Exception as e: spend_log_error( "Spend tracking - failed to commit spend updates from Redis to DB. " @@ -1016,6 +1034,17 @@ 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() + ) + + await DBSpendUpdateWriter._commit_window_spend_updates( + prisma_client=prisma_client, + window_spend_transactions=window_spend_update_transactions, + ) + ################## Tool Registry Upserts ################## await self._flush_tool_discovery_queue(prisma_client=prisma_client) @@ -1086,6 +1115,36 @@ 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. + + Failures are logged rather than raised: these rows exist so budget + enforcement can stop aggregating LiteLLM_SpendLogs, and the read path + falls back to that aggregate, so a failed commit must not abort the + entity and daily spend commits that share this scheduler tick. + """ + from litellm.proxy.db.budget_window_spend_writer import ( + commit_window_spend_updates, + ) + + try: + await commit_window_spend_updates( + prisma_client=prisma_client, + transactions=window_spend_transactions, + ) + except Exception as e: # noqa: BLE001 # any DB failure here must stay contained to the window rows + spend_log_error( + "Spend tracking - failed to commit budget window spend updates. %d window increments lost. Error: %s", + len(window_spend_transactions), + str(e), + exc=e, + ) + async def _flush_tool_discovery_queue( 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 d64280efa8c..d5786aa7088 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -19,6 +19,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 ( @@ -38,6 +39,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, @@ -130,6 +135,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 @@ -196,6 +202,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) @@ -232,6 +243,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]] = [] @@ -272,12 +288,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 @@ -297,12 +315,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. @@ -372,6 +392,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) + @staticmethod def _number_of_transactions_to_store_in_redis( db_spend_update_transactions: DBSpendUpdateTransactions, @@ -468,20 +491,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), @@ -505,12 +530,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 @@ -531,6 +560,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]), @@ -538,6 +575,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..d2ad2aa61d2 --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py @@ -0,0 +1,160 @@ +""" +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 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. + """ + + entity_type: str + entity_id: str + window_duration: str + window_start: str + spend: float + request_ids: Sequence[str] + + +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, +) -> 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,), + ) + + +def _merge_window_spend_transactions( + payloads: tuple[WindowSpendTransaction, ...], +) -> WindowSpendTransaction: + first: Final = payloads[0] + 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)))), + ) + + +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): + 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/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3a4914d56a8..58bc9cfb301 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -498,7 +498,7 @@ async def _update_database_and_spend_counters( request_tags: list[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, @@ -534,6 +534,7 @@ 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, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fb9c4e67aad..b3b59d72a53 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -349,6 +349,9 @@ from litellm.proxy.config_resolvers.alerting import ( from litellm.proxy.container_endpoints.endpoints import router as container_router from litellm.proxy.credential_endpoints.endpoints import router as credential_router from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import 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, @@ -2392,6 +2395,7 @@ 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, ): """ Atomically increment spend counters for budget enforcement. @@ -2446,14 +2450,24 @@ async def increment_spend_counters( for window in key_budget_limits: duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration 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_start=key_window_start, increment=cost, ) + await _enqueue_window_spend_row_update( + entity_type=Litellm_EntityType.KEY, + entity_id=hashed_token, + window=window, + window_duration=duration, + window_start=key_window_start, + increment=cost, + request_id=request_id, + ) async def _team_scope(scope_team_id: str) -> None: team_counter_key: Final = f"spend:team:{scope_team_id}" @@ -2477,14 +2491,24 @@ async def increment_spend_counters( for window in team_budget_limits: duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration 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_start=team_window_start, increment=cost, ) + await _enqueue_window_spend_row_update( + entity_type=Litellm_EntityType.TEAM, + entity_id=scope_team_id, + window=window, + window_duration=duration, + window_start=team_window_start, + increment=cost, + request_id=request_id, + ) 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}" @@ -2669,6 +2693,58 @@ 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, + window: Any, + window_duration: str, + window_start: datetime | None, + increment: float, + request_id: str | 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; the + flush uses it 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: + return + reset_at: Final = window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + if 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, + ) + ) + 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, diff --git a/litellm/types/services.py b/litellm/types/services.py index f43494a0a09..60cfc35c0a5 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -42,6 +42,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 f04d6f3cf5a..c7d31972d87 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 @@ -1082,6 +1082,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) @@ -1153,6 +1154,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 33372e7794a..e4a653ddb51 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 @@ -208,7 +208,8 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( 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}, @@ -222,6 +223,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=[ @@ -231,13 +244,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 @@ -260,6 +290,12 @@ 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 @@ -267,7 +303,7 @@ 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) def test_validate_redis_transaction_buffer_raises_without_redis(): @@ -374,3 +410,106 @@ 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", + ) + ) + + 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"], + }] + + +@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..880cc5044f9 --- /dev/null +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py @@ -0,0 +1,249 @@ +import json +import os +import sys +from datetime import datetime, timedelta, timezone + +sys.path.insert(0, os.path.abspath("../../../../..")) + +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, +): + 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, + ) + + +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",), + } + + +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..c05e3d410f9 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py @@ -0,0 +1,535 @@ +import math +import os +import sys +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from typing import Any + +sys.path.insert(0, os.path.abspath("../../../..")) + +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) + +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, + ) -> float | None: + self.calls.append( + { + "entity_type": entity_type, + "entity_id": entity_id, + "window_start": window_start, + "exclude_request_ids": tuple(exclude_request_ids), + } + ) + return self.value + + +class _SpendLogsFake: + """Sums the LiteLLM_SpendLogs rows it holds, honouring the request-id + exclusion exactly as the real aggregate's NOT (request_id = ANY(...)) does.""" + + def __init__(self, rows: tuple[tuple[str, float], ...]) -> None: + self.rows = rows + + async def __call__( + self, + prisma_client: Any, + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Any, + ) -> float | None: + excluded = frozenset(exclude_request_ids) + return math.fsum(spend for request_id, spend in self.rows if request_id not in excluded) + + +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): + 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): + 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_to_exclude(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 3.0, + "request_ids": ("req-1", "req-2", "req-3"), + }, + ), + spend_logs_aggregate=aggregate, + ) + + assert aggregate.calls[0]["exclude_request_ids"] == ("req-1", "req-2", "req-3") + + +@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), ("req-2", 0.000047), ("req-3", 0.000047)), + ) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 0.000141, + "request_ids": ("req-1", "req-2", "req-3"), + }, + ), + 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), ("req-1", 0.000047))) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 0.000047, + "request_ids": ("req-1",), + }, + ), + 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=( + { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 0.000141, + "request_ids": ("req-1", "req-2", "req-3"), + }, + ), + 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_by_parameter(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"), + ) + + 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[]))" in normalized + assert 'FROM "LiteLLM_SpendLogs"' in normalized + assert params == ("e1", WINDOW_A, ("req-1", "req-2")) + # The ids are bound, never spliced into the statement. + assert "req-1" not in query + + +@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=(), + ) + + 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=(), + ) + + 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 191080e3a48..bd55703408e 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 @@ -9,6 +9,7 @@ sys.path.insert( ) # Adds the parent directory to the system path +from contextlib import asynccontextmanager from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, call, patch @@ -18,6 +19,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 @@ -2236,3 +2240,258 @@ async def test_daily_transaction_compression_saved_tokens_zero_when_absent(): assert transaction["compression_saved_tokens"] == 0 assert transaction["compression_savings_spend"] == 0 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_does_not_abort_the_rest_of_the_flush(): + """Window rows are an optimization over aggregating LiteLLM_SpendLogs, so a + failure must not take the tool registry flush down with it.""" + 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() + 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() + + +@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("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.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("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch( + "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 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 f289148101a..0c0b4ac7889 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 @@ -440,7 +440,9 @@ 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": []} @@ -471,6 +473,7 @@ 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", ) @@ -1263,3 +1266,58 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( 1 if expect_spend_log else 0 ) + + +@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 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index c7aa376a2f7..7e7c2e01fcc 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 @@ -11016,3 +11017,260 @@ def test_startup_is_silent_when_mock_testing_params_disabled(caplog): ProxyStartupEvent._warn_if_mock_testing_params_enabled(general_settings={}) 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_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",)