From 3b1ab1908c5657b45bb2a0c6a6b24c5c3339ebb7 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 4 Aug 2026 16:56:37 -0700 Subject: [PATCH 01/20] feat(proxy): add LiteLLM_BudgetWindowSpend table for per-window budget spend Multi-window budgets (budget_limits on keys/teams) currently keep window spend only in cache. Every cold or expired counter recomputes the window by aggregating LiteLLM_SpendLogs, which has no usable index for that query and saturates the DB on large tables (#35766). This adds a LiteLLM_BudgetWindowSpend table holding one row per configured window, keyed (entity_type, entity_id, window_duration), with window_start identifying the period the spend belongs to. Follow-up PRs maintain these rows from the spend update writer and move window budget enforcement reads onto them. --- .../migration.sql | 13 +++++++++++++ .../litellm_proxy_extras/schema.prisma | 12 ++++++++++++ litellm/proxy/schema.prisma | 12 ++++++++++++ schema.prisma | 12 ++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 5 files changed, 50 insertions(+), 1 deletion(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql new file mode 100644 index 00000000000..45cc927a328 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql @@ -0,0 +1,13 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" ( + "entity_type" TEXT NOT NULL, + "entity_id" TEXT NOT NULL, + "window_duration" TEXT NOT NULL, + "window_start" TIMESTAMP(3) NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_BudgetWindowSpend_pkey" PRIMARY KEY ("entity_type","entity_id","window_duration") +); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 17339541fd9..85401b82676 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -645,6 +645,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 17339541fd9..85401b82676 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -645,6 +645,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/schema.prisma b/schema.prisma index 17339541fd9..85401b82676 100644 --- a/schema.prisma +++ b/schema.prisma @@ -645,6 +645,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9133bfb5cf4..fd2155701c2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23744,7 +23744,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; + user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; }; /** * DefaultTeamSSOParams From 187e0fab605c63dcb73fedfe756c6c80ceb396ec Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 4 Aug 2026 18:25:00 -0700 Subject: [PATCH 02/20] feat(proxy): maintain per-window budget spend rows in the spend writer Multi-window budgets (budget_limits on keys and teams) enforce off Redis counters with a 60s TTL. Every cold counter, and every authoritative floor check, aggregates LiteLLM_SpendLogs with a range scan over startTime; that table has no index on api_key or team_id, so the scan lands on the highest volume table in the schema and saturates the connection pool (#35766). Every other budget feature reads a maintained running spend value instead. This gives windows the same shape by keeping LiteLLM_BudgetWindowSpend up to date: one row per configured window, whose window_start rolls forward in place. The cost callback already iterates a key's and a team's windows with the window start computed and the actual cost in hand, so it enqueues there, onto a new WindowSpendUpdateQueue. Increments are enqueued even when the cache increment is skipped for a reserved counter, since the reservation only pre-charged an estimate and the row still owes the actual cost. Windows with no reset_at slide with wall clock and cannot be represented by a single row, so they are left to the read path's existing aggregate. The queue flushes alongside the daily spend queues, through the Redis buffer when one is configured (only the pod-lock winner commits) and directly otherwise. A flush selects the primary keys that already exist, seeds the ones that do not from LiteLLM_SpendLogs so a new row cannot undercount spend that predates it, and applies one batch of upserts ordered by primary key. An increment at or behind the stored window_start adds into the row, matching how in-flight requests carry into a window after a reset; a newer one rolls the window and starts from that increment. The conflict arm adds only the increment, never the seeded base, so two pods seeding the same new window cannot double count it. The seed excludes the requests its own batch is about to apply. Spend logs are drained by a separate monitor that fires on a ~2s poll whenever anything is queued, while window increments flush on the much slower batch tick, so by seed time the batch's log rows are normally already in the table; counting them in the aggregate and again in the increments made a fresh row land at exactly twice the true spend. Each increment therefore carries the LiteLLM_SpendLogs request_id it was recorded under, which update_database now returns rather than having the callback re-derive it (a cache hit appends time.time() to that id, so a second derivation would not match). The reset job rolls each expired window's row alongside the counter it zeroes, conditional on the stored window_start still being behind the new one so a pod that already rolled it is not clobbered. --- litellm/constants.py | 1 + .../proxy/common_utils/reset_budget_job.py | 65 ++- .../proxy/db/budget_window_spend_writer.py | 266 +++++++++ litellm/proxy/db/db_spend_update_writer.py | 65 ++- .../redis_update_buffer.py | 46 +- .../window_spend_update_queue.py | 160 ++++++ .../proxy/hooks/proxy_track_cost_callback.py | 3 +- litellm/proxy/proxy_server.py | 80 ++- litellm/types/services.py | 2 + .../common_utils/test_reset_budget_job.py | 140 +++++ .../test_redis_update_buffer.py | 147 ++++- .../test_window_spend_update_queue.py | 249 ++++++++ .../db/test_budget_window_spend_writer.py | 535 ++++++++++++++++++ .../proxy/db/test_db_spend_update_writer.py | 259 +++++++++ .../hooks/test_proxy_track_cost_callback.py | 60 +- tests/test_litellm/proxy/test_proxy_server.py | 258 +++++++++ 16 files changed, 2317 insertions(+), 19 deletions(-) create mode 100644 litellm/proxy/db/budget_window_spend_writer.py create mode 100644 litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py create mode 100644 tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py create mode 100644 tests/test_litellm/proxy/db/test_budget_window_spend_writer.py 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",) From 124d08d59234443fb441c3baba6b23d4a5dd1756 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 4 Aug 2026 17:39:51 -0700 Subject: [PATCH 03/20] perf(proxy): read budget-window spend from the maintained window table Per-window budget enforcement aggregated LiteLLM_SpendLogs on every cold counter and on every 5s authoritative floor check. SpendLogs has no index on api_key or team_id, so each check range-scanned the highest-volume table. Reads now hit the LiteLLM_BudgetWindowSpend row by primary key and only fall back to the aggregate when no row exists for the window being enforced. A row is current when its window_start is at or past the caller's expected start, so a pod holding a stale reset_at trusts a window another pod already rolled instead of summing the previous window back in. window_duration is threaded from the budget_limits entry through to the read rather than parsed back out of the counter key. The reader never writes rows. --- litellm/proxy/auth/auth_checks.py | 2 + litellm/proxy/db/spend_counter_reseed.py | 134 ++++++++-- litellm/proxy/proxy_server.py | 15 +- .../spend_tracking/budget_reservation.py | 3 + litellm/repositories/table_repositories.py | 4 + .../proxy/db/test_spend_counter_reseed.py | 250 ++++++++++++++++++ .../proxy/proxy_server/test_spend_counters.py | 79 ++++++ tests/test_litellm/proxy/test_proxy_server.py | 8 + 8 files changed, 478 insertions(+), 17 deletions(-) create mode 100644 tests/test_litellm/proxy/db/test_spend_counter_reseed.py diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5ef4eb471ad..3e197695710 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3710,6 +3710,7 @@ async def _virtual_key_multi_budget_check( max_budget=w["max_budget"], window_entity_type="Key", window_entity_id=valid_token.token, + window_duration=str(w["budget_duration"]), window_start=get_budget_window_start(w), ) if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: @@ -4083,6 +4084,7 @@ async def _team_multi_budget_check( max_budget=w["max_budget"], window_entity_type="Team", window_entity_id=team_object.team_id, + window_duration=str(w["budget_duration"]), window_start=get_budget_window_start(w), ) if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 72e4111e0f4..1f4fdef156c 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -14,14 +14,18 @@ memory in long-lived deployments. import asyncio from collections import OrderedDict -from datetime import datetime +from collections.abc import Mapping +from datetime import datetime, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final, Optional from litellm._logging import verbose_proxy_logger from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import Litellm_EntityType from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( + BudgetWindowSpendRepository, SpendLogsRepository, TeamMembershipRepository, ) @@ -36,6 +40,25 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +_WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType( + { + "Key": Litellm_EntityType.KEY.value, + "Team": Litellm_EntityType.TEAM.value, + } +) + +_WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + { + "Key": "api_key", + "Team": "team_id", + } +) + + +def _as_utc(value: datetime) -> datetime: + return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) + + class SpendCounterReseed: """ Reseeds spend counters from the authoritative DB and warms the cache, @@ -205,6 +228,92 @@ class SpendCounterReseed: raise return current_value + @staticmethod + async def window_from_table( + prisma_client: Optional["PrismaClient"], + entity_type: str, + entity_id: str, + window_duration: str, + expected_window_start: datetime, + ) -> float | None: + """ + Read the maintained per-window spend row by primary key. + + Returns the row's spend only when the row belongs to the window the + caller is enforcing, i.e. ``row.window_start >= expected_window_start``. + A row at or past the expected start was rolled by a pod whose reset_at + was at least as fresh as this caller's, so it is trusted; an older row + means the window boundary was crossed and nothing has rolled the row + yet, so its spend belongs to a previous window. + + Returns None for a missing, stale or unreadable row so the caller falls + back to the spend-logs aggregate. ``entity_type`` is the counter-facing + label ("Key"/"Team"); anything else has no row and returns None. + """ + if prisma_client is None: + return None + row_entity_type: Final = _WINDOW_SPEND_ENTITY_TYPES.get(entity_type) + if row_entity_type is None: + return None + + try: + row: Final = await BudgetWindowSpendRepository(prisma_client).table.find_unique( + where={ + "entity_type_entity_id_window_duration": { + "entity_type": row_entity_type, + "entity_id": entity_id, + "window_duration": window_duration, + } + } + ) + except Exception: + verbose_proxy_logger.exception( + "SpendCounterReseed.window_from_table: failed for %s=%s window=%s", + entity_type, + entity_id, + window_duration, + ) + return None + + if row is None: + return None + if _as_utc(row.window_start) < _as_utc(expected_window_start): + return None + return float(row.spend or 0.0) + + @staticmethod + async def window_from_db( + prisma_client: Optional["PrismaClient"], + entity_type: str, + entity_id: str, + window_duration: str | None, + window_start: datetime, + ) -> float | None: + """ + Authoritative window spend: the maintained row first, falling back to + the spend-logs aggregate only when no current row exists. + + The aggregate range-scans an unindexed table, so it must stay a + transitional path (window configured before the row existed) rather + than a steady-state read. + """ + if window_duration is not None: + from_table: Final = await SpendCounterReseed.window_from_table( + prisma_client=prisma_client, + entity_type=entity_type, + entity_id=entity_id, + window_duration=window_duration, + expected_window_start=window_start, + ) + if from_table is not None: + return from_table + return await SpendCounterReseed.window_from_spend_logs( + prisma_client=prisma_client, + entity_type=entity_type, + entity_id=entity_id, + window_start=window_start, + ) + @staticmethod async def window_from_spend_logs( prisma_client: Optional["PrismaClient"], @@ -215,20 +324,13 @@ class SpendCounterReseed: if prisma_client is None: return None - if entity_type == "Key": - group_field = "api_key" - where = { - "api_key": entity_id, - "startTime": {"gte": window_start}, - } - elif entity_type == "Team": - group_field = "team_id" - where = { - "team_id": entity_id, - "startTime": {"gte": window_start}, - } - else: + group_field: Final = _WINDOW_SPEND_LOG_FIELDS.get(entity_type) + if group_field is None: return None + where: Final = { + group_field: entity_id, + "startTime": {"gte": window_start}, + } try: response: Final = await SpendLogsRepository(prisma_client).table.group_by( @@ -258,6 +360,7 @@ class SpendCounterReseed: counter_key: str, entity_type: str, entity_id: str, + window_duration: str | None, window_start: datetime, ) -> float | None: lock: Final = await SpendCounterReseed._get_lock(counter_key) @@ -276,10 +379,11 @@ class SpendCounterReseed: if val is not None: return float(val) - window_spend: Final = await SpendCounterReseed.window_from_spend_logs( + window_spend: Final = await SpendCounterReseed.window_from_db( prisma_client=prisma_client, entity_type=entity_type, entity_id=entity_id, + window_duration=window_duration, window_start=window_start, ) if window_spend is None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b3b59d72a53..0a017a50697 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2184,6 +2184,7 @@ async def get_current_spend( max_budget: float | None = None, window_entity_type: str | None = None, window_entity_id: str | None = None, + window_duration: str | None = None, window_start: datetime | None = None, fallback_authoritative: bool = False, ) -> float: @@ -2208,7 +2209,8 @@ async def get_current_spend( runs and a key can leak spend past ``max_budget`` indefinitely. The authoritative source depends on the counter: primary key/team/user/org counters read the DB row; per-window counters (``window_start`` supplied) - aggregate spend logs; end-user/tag counters have no DB row, so the caller's + read the maintained window-spend row and only aggregate spend logs when + that row is missing or stale; end-user/tag counters have no DB row, so the caller's ``fallback_spend`` (loaded fresh in auth) is authoritative. The DB read is skipped for healthy primary counters (counter at or above recorded spend) and cached in-process for a few seconds, so a persistently stale counter @@ -2233,6 +2235,7 @@ async def get_current_spend( counter_key=counter_key, window_entity_type=window_entity_type, window_entity_id=window_entity_id, + window_duration=window_duration, window_start=window_start, ) if authoritative is not None: @@ -2314,6 +2317,7 @@ async def _authoritative_floor_spend( counter_key: str, window_entity_type: str | None = None, window_entity_id: str | None = None, + window_duration: str | None = None, window_start: datetime | None = None, ) -> float | None: marker_key: Final = f"spend_db_floor:{counter_key}" @@ -2328,10 +2332,11 @@ async def _authoritative_floor_spend( and window_entity_id is not None and window_start is not None ): - db_spend = await SpendCounterReseed.window_from_spend_logs( + db_spend = await SpendCounterReseed.window_from_db( prisma_client=prisma_client, entity_type=window_entity_type, entity_id=window_entity_id, + window_duration=window_duration, window_start=window_start, ) if db_spend is None: @@ -2456,6 +2461,7 @@ async def increment_spend_counters( counter_key=key_window_counter, entity_type="Key", entity_id=hashed_token, + window_duration=duration, window_start=key_window_start, increment=cost, ) @@ -2497,6 +2503,7 @@ async def increment_spend_counters( counter_key=team_window_counter, entity_type="Team", entity_id=scope_team_id, + window_duration=duration, window_start=team_window_start, increment=cost, ) @@ -2749,6 +2756,7 @@ async def _init_and_increment_window_spend_counter( counter_key: str, entity_type: str, entity_id: str, + window_duration: str | None, window_start: datetime | None, increment: float, ): @@ -2763,6 +2771,7 @@ async def _init_and_increment_window_spend_counter( counter_key=counter_key, entity_type=entity_type, entity_id=entity_id, + window_duration=window_duration, window_start=window_start, ) if initialized is False: @@ -2808,6 +2817,7 @@ async def _ensure_window_spend_counter_initialized( counter_key: str, entity_type: str, entity_id: str, + window_duration: str | None, window_start: datetime, ) -> bool: is_warm: Final = await _is_spend_counter_cache_warm(counter_key=counter_key) @@ -2820,6 +2830,7 @@ async def _ensure_window_spend_counter_initialized( counter_key=counter_key, entity_type=entity_type, entity_id=entity_id, + window_duration=window_duration, window_start=window_start, ) if window_spend is None: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 907c4ac7344..de82d3cfd60 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -37,6 +37,7 @@ class _BudgetCounter: entity_id: str source_cache_key: str | None = None spend_log_entity_id: str | None = None + window_duration: str | None = None window_start: datetime | None = None @@ -633,6 +634,7 @@ def _get_budget_limit_counters( entity_type=entity_type, entity_id=f"{entity_id}:{budget_duration}", spend_log_entity_id=entity_id, + window_duration=str(budget_duration), window_start=window_start, ) ) @@ -676,6 +678,7 @@ async def _reserve_counter( counter_key=counter.counter_key, entity_type=counter.entity_type, entity_id=counter.spend_log_entity_id, + window_duration=counter.window_duration, window_start=counter.window_start, ) if initialized is False: diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index af8be986831..d3c79bb4f6f 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -62,6 +62,10 @@ class SpendLogsRepository(PrismaTableRepository): table_name = "litellm_spendlogs" +class BudgetWindowSpendRepository(PrismaTableRepository): + table_name = "litellm_budgetwindowspend" + + class ClaudeCodePluginRepository(PrismaTableRepository): table_name = "litellm_claudecodeplugintable" diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py new file mode 100644 index 00000000000..816f9ae72f4 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -0,0 +1,250 @@ +"""Window-spend reads in ``SpendCounterReseed``. + +The maintained ``LiteLLM_BudgetWindowSpend`` row replaces a per-request +``LiteLLM_SpendLogs`` range scan, so these pin *when* the aggregate is still +allowed to run: only when the row is missing or belongs to an older window. +""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace + +import pytest + +from litellm.caching.dual_cache import DualCache +from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + +WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) + + +class _FakeWindowSpendTable: + def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None: + self._row = row + self._error = error + self.where_clauses: list[dict] = [] + + async def find_unique(self, where: dict): + self.where_clauses.append(where) + if self._error is not None: + raise self._error + return self._row + + +class _FakeSpendLogsTable: + def __init__(self, total: float) -> None: + self._total = total + self.call_count = 0 + + async def group_by(self, by: list[str], where: dict, sum: dict): + self.call_count += 1 + return [{by[0]: where.get(by[0]), "_sum": {"spend": self._total}}] + + +class _FakePrismaClient: + def __init__( + self, + row: SimpleNamespace | None = None, + spend_logs_total: float = 0.0, + error: Exception | None = None, + ) -> None: + self.db = SimpleNamespace( + litellm_budgetwindowspend=_FakeWindowSpendTable(row=row, error=error), + litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total), + ) + + +def _row(window_start: datetime, spend: float) -> SimpleNamespace: + return SimpleNamespace(window_start=window_start, spend=spend) + + +@pytest.mark.asyncio +async def test_window_from_table_reads_row_by_primary_key(): + """The lookup must use the table's own entity_type values ("key"), not the + "Key"/"Team" labels the counter keys and spend-log aggregates use.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result == 4.5 + assert prisma.db.litellm_budgetwindowspend.where_clauses == [ + { + "entity_type_entity_id_window_duration": { + "entity_type": "key", + "entity_id": "tok-1", + "window_duration": "30d", + } + } + ] + + +@pytest.mark.asyncio +async def test_window_from_table_maps_team_entity_type(): + prisma = _FakePrismaClient(row=_row(WINDOW_START, 9.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Team", + entity_id="team-1", + window_duration="1d", + expected_window_start=WINDOW_START, + ) + + assert result == 9.0 + inner = prisma.db.litellm_budgetwindowspend.where_clauses[0]["entity_type_entity_id_window_duration"] + assert inner["entity_type"] == "team" + + +@pytest.mark.asyncio +async def test_window_from_table_trusts_row_newer_than_expected_window(): + """Regression: a pod holding a stale ``reset_at`` computes an expected start + behind a window another pod already rolled. Trusting only an exact match + would make it re-add the previous window's spend to the current one.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START + timedelta(days=1), 2.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result == 2.0 + + +@pytest.mark.asyncio +async def test_window_from_table_rejects_row_from_previous_window(): + prisma = _FakePrismaClient(row=_row(WINDOW_START - timedelta(seconds=1), 99.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_window_from_table_treats_naive_row_timestamp_as_utc(): + """The column is ``timestamp(3)``, so a driver that hands back a naive value + must still compare against the tz-aware expected start.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START.replace(tzinfo=None), 3.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result == 3.0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prisma, entity_type", + [ + (_FakePrismaClient(row=None), "Key"), + (_FakePrismaClient(row=_row(WINDOW_START, 1.0)), "User"), + (_FakePrismaClient(error=RuntimeError("connection reset")), "Key"), + (None, "Key"), + ], +) +async def test_window_from_table_returns_none_without_a_usable_row(prisma, entity_type): + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type=entity_type, + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_window_from_db_prefers_the_row_over_the_spend_logs_aggregate(): + """The aggregate range-scans an unindexed table; a current row must keep it + from running at all.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + window_start=WINDOW_START, + ) + + assert result == 4.5 + assert prisma.db.litellm_spendlogs.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "row", + [None, _row(WINDOW_START - timedelta(seconds=1), 99.0)], + ids=["missing_row", "previous_window_row"], +) +async def test_window_from_db_falls_back_to_spend_logs(row): + prisma = _FakePrismaClient(row=row, spend_logs_total=7.25) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + window_start=WINDOW_START, + ) + + assert result == 7.25 + assert prisma.db.litellm_spendlogs.call_count == 1 + + +@pytest.mark.asyncio +async def test_window_from_db_without_a_duration_skips_the_row_lookup(): + """Callers that cannot name the window (no PK) keep the pre-table behavior.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=7.25) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration=None, + window_start=WINDOW_START, + ) + + assert result == 7.25 + assert prisma.db.litellm_budgetwindowspend.where_clauses == [] + + +@pytest.mark.asyncio +async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0) + cache = DualCache() + counter_key = "spend:key:tok-1:window:30d" + + result = await SpendCounterReseed.coalesced_window( + prisma_client=prisma, + spend_counter_cache=cache, + counter_key=counter_key, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + window_start=WINDOW_START, + ) + + assert result == 4.5 + assert cache.in_memory_cache.get_cache(key=counter_key) == 4.5 + assert prisma.db.litellm_spendlogs.call_count == 0 diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 51980342a1d..fb3de990deb 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -271,6 +271,81 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): ) +def _make_window_spend_prisma(row=None, spend_logs_total=0.0): + prisma = MagicMock() + prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=row) + prisma.db.litellm_spendlogs.group_by = AsyncMock( + return_value=[{"api_key": "tok", "_sum": {"spend": spend_logs_total}}] + ) + return prisma + + +@pytest.mark.asyncio +async def test_get_current_spend_floors_window_against_maintained_row(monkeypatch): + """The floor re-check runs every few seconds per pod, so the window branch + must read the maintained row and leave the unindexed spend-logs scan alone.""" + from datetime import timezone + from types import SimpleNamespace + + window_start = datetime(2026, 1, 1, tzinfo=timezone.utc) + fake_prisma = _make_window_spend_prisma( + row=SimpleNamespace(window_start=window_start, spend=15.0), + spend_logs_total=100.0, + ) + fake_cache = _make_spend_counter_cache(redis_get_value=2.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + + counter_key = "spend:key:tok:window:7d" + result = await ps.get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + max_budget=10.0, + window_entity_type="Key", + window_entity_id="tok", + window_duration="7d", + window_start=window_start, + ) + + assert result == 15.0 + fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited() + fake_cache.redis_cache.async_set_max.assert_awaited_once_with( + key=counter_key, value=15.0 + ) + + +@pytest.mark.asyncio +async def test_get_current_spend_floors_window_against_logs_when_row_stale(monkeypatch): + """A row left behind at a crossed window boundary must not be read as the + current window's spend; the aggregate stays the fallback.""" + from datetime import timedelta, timezone + from types import SimpleNamespace + + window_start = datetime(2026, 1, 8, tzinfo=timezone.utc) + fake_prisma = _make_window_spend_prisma( + row=SimpleNamespace( + window_start=window_start - timedelta(days=7), spend=999.0 + ), + spend_logs_total=15.0, + ) + fake_cache = _make_spend_counter_cache(redis_get_value=2.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + + result = await ps.get_current_spend( + counter_key="spend:key:tok:window:7d", + fallback_spend=0.0, + max_budget=10.0, + window_entity_type="Key", + window_entity_id="tok", + window_duration="7d", + window_start=window_start, + ) + + assert result == 15.0 + fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once() + + @pytest.mark.asyncio async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypatch): """With fail_closed_budget_enforcement on, an admit decision backed only by a @@ -895,6 +970,7 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=datetime(2024, 1, 1), increment=5.0, ) @@ -922,6 +998,7 @@ async def test_init_and_increment_window_spend_counter_missing_window_start_inva counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=None, increment=5.0, ) @@ -1059,6 +1136,7 @@ async def test_ensure_window_spend_counter_initialized_warm_returns_true(monkeyp counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=datetime(2024, 1, 1), ) @@ -1091,6 +1169,7 @@ async def test_ensure_window_spend_counter_initialized_db_failure_invalid_return counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=datetime(2024, 1, 1), ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7e7c2e01fcc..bb3e3d84949 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7643,6 +7643,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): counter_cache = DualCache() window_start = datetime.now(timezone.utc) - timedelta(hours=1) fake_prisma = MagicMock() + fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[{"api_key": "key-window", "_sum": {"spend": 2.25}}] ) @@ -7657,6 +7658,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): counter_key="spend:key:key-window:window:1h", entity_type="Key", entity_id="key-window", + window_duration="1h", window_start=window_start, increment=0.5, ) @@ -7766,6 +7768,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() + fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[{"api_key": "key-window-stale-local", "_sum": {"spend": 2.25}}] ) @@ -7780,6 +7783,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): counter_key=counter_key, entity_type="Key", entity_id="key-window-stale-local", + window_duration="1h", window_start=window_start, increment=0.5, ) @@ -7830,6 +7834,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() + fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[ {"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}} @@ -7846,6 +7851,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() counter_key=counter_key, entity_type="Key", entity_id="key-window-concurrent-seed", + window_duration="1h", window_start=window_start, increment=0.5, ) @@ -7880,6 +7886,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): counter_key="spend:key:key-invalid-window:window:not-a-duration", entity_type="Key", entity_id="key-invalid-window", + window_duration="not-a-duration", window_start=None, increment=0.5, ) @@ -7912,6 +7919,7 @@ async def test_window_spend_counter_does_not_seed_zero_when_db_unavailable(): counter_key=counter_key, entity_type="Key", entity_id="key-window-db-unavailable", + window_duration="1h", window_start=datetime.now(timezone.utc) - timedelta(hours=1), ) From 83cedfd2b23e5d6a13c3e2dcf8491b13a2e289de Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 10:51:38 -0700 Subject: [PATCH 04/20] refactor(proxy): type the budget window reset_at passed to the window spend enqueue --- litellm/proxy/proxy_server.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b3b59d72a53..89fe315fbd4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2449,6 +2449,7 @@ async def increment_spend_counters( return for window in key_budget_limits: duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration + key_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at key_window_counter = f"spend:key:{hashed_token}:window:{duration}" key_window_start = get_budget_window_start(window) if key_window_counter not in reserved_counter_keys: @@ -2462,7 +2463,7 @@ async def increment_spend_counters( await _enqueue_window_spend_row_update( entity_type=Litellm_EntityType.KEY, entity_id=hashed_token, - window=window, + reset_at=key_window_reset_at, window_duration=duration, window_start=key_window_start, increment=cost, @@ -2490,6 +2491,7 @@ async def increment_spend_counters( return for window in team_budget_limits: duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration + team_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at team_window_counter = f"spend:team:{scope_team_id}:window:{duration}" team_window_start = get_budget_window_start(window) if team_window_counter not in reserved_counter_keys: @@ -2503,7 +2505,7 @@ async def increment_spend_counters( await _enqueue_window_spend_row_update( entity_type=Litellm_EntityType.TEAM, entity_id=scope_team_id, - window=window, + reset_at=team_window_reset_at, window_duration=duration, window_start=team_window_start, increment=cost, @@ -2696,7 +2698,7 @@ async def _init_and_increment_spend_counter( async def _enqueue_window_spend_row_update( entity_type: Litellm_EntityType, entity_id: str, - window: Any, + reset_at: datetime | str | None, window_duration: str, window_start: datetime | None, increment: float, @@ -2719,10 +2721,7 @@ async def _enqueue_window_spend_row_update( 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: + if window_start is None or not reset_at: return try: await proxy_logging_obj.db_spend_update_writer.window_spend_update_queue.add_update( From 041cae82801574d68a56da97437c31e38fc1847e Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 11:05:57 -0700 Subject: [PATCH 05/20] fix(proxy): bound the window spend seed exclusion to the batch's own start time request_id can be chosen by the client through x-litellm-call-id, so an unbounded NOT (request_id = ANY(batch)) let a replayed old id drop that id's historical LiteLLM_SpendLogs row from the one-time seed while its increment still landed. The increment now carries the request start, the batch keeps the earliest one, and the seed only excludes ids whose startTime is at or after it. --- .../proxy/db/budget_window_spend_writer.py | 57 +++++- .../window_spend_update_queue.py | 14 ++ .../proxy/hooks/proxy_track_cost_callback.py | 1 + litellm/proxy/proxy_server.py | 11 +- .../test_redis_update_buffer.py | 2 + .../test_window_spend_update_queue.py | 28 +++ .../db/test_budget_window_spend_writer.py | 168 ++++++++++++------ .../hooks/test_proxy_track_cost_callback.py | 4 +- tests/test_litellm/proxy/test_proxy_server.py | 26 +++ 9 files changed, 245 insertions(+), 66 deletions(-) diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index 0308651f4e3..02ff9c4f944 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -65,13 +65,23 @@ _ROLL_WINDOW_SPEND_SQL: Final = ( _SEED_FROM_SPEND_LOGS_KEY_SQL: Final = ( 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " - "AND NOT (request_id = ANY($3::text[]))" + "AND NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" ) _SEED_FROM_SPEND_LOGS_TEAM_SQL: Final = ( 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " - "AND NOT (request_id = ANY($3::text[]))" + "AND NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" +) + +_SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL: Final = ( + 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" +) + +_SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL: Final = ( + 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" ) _UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60) @@ -92,6 +102,7 @@ class WindowSpendLogsAggregate(Protocol): entity_id: str, window_start: datetime, exclude_request_ids: Sequence[str], + exclude_started_at: datetime | None, ) -> float | None: ... @@ -101,6 +112,7 @@ async def spend_logs_total_excluding( entity_id: str, window_start: datetime, exclude_request_ids: Sequence[str], + exclude_started_at: datetime | None, ) -> float | None: """LiteLLM_SpendLogs spend for one entity since window_start, minus the requests already accounted for by the increments being flushed. @@ -110,22 +122,43 @@ async def spend_logs_total_excluding( by the time a window row is seeded its batch's log rows are normally already in the table. Counting them in the seed and again in the increment is what made a fresh row land at twice the true spend. + + The exclusion is bounded to rows that started at or after the batch's + earliest request. request_id can be chosen by the client + (x-litellm-call-id), so an unbounded exclusion would let a replayed old id + erase a historical row from the seed while its increment still lands. + Without a known start the batch's ids are not excluded at all: that can + only over-count once, which enforcement tolerates, whereas under-counting + is a budget bypass. """ if entity_type == Litellm_EntityType.KEY.value: - rows = await prisma_client.db.query_raw( - _SEED_FROM_SPEND_LOGS_KEY_SQL, entity_id, window_start, tuple(exclude_request_ids) - ) + bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_KEY_SQL, _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL 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) - ) + bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_TEAM_SQL, _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL else: return None + rows: Final = ( + await prisma_client.db.query_raw(unbounded_sql, entity_id, window_start) + if exclude_started_at is None or not exclude_request_ids + else await prisma_client.db.query_raw( + bounded_sql, + entity_id, + window_start, + tuple(exclude_request_ids), + _exclusion_lower_bound(exclude_started_at), + ) + ) if not rows: return 0.0 return float(rows[0].get("total") or 0.0) +def _exclusion_lower_bound(started_at: datetime) -> datetime: + """LiteLLM_SpendLogs.startTime is TIMESTAMP(3); floor to the second so a + millisecond rounding of the batch's own earliest row cannot slip under it.""" + return to_naive_utc(started_at).replace(microsecond=0) + + def _primary_key(transaction: WindowSpendTransaction) -> tuple[str, str, str]: return ( transaction["entity_type"], @@ -168,10 +201,18 @@ async def _seed_base_for_missing_row( entity_id=transaction["entity_id"], window_start=datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc), exclude_request_ids=transaction["request_ids"], + exclude_started_at=_transaction_started_at(transaction), ) return float(base or 0.0) +def _transaction_started_at(transaction: WindowSpendTransaction) -> datetime | None: + started_at: Final = transaction.get("started_at") + if started_at is None: + return None + return datetime.fromisoformat(started_at).replace(tzinfo=timezone.utc) + + def _upsert_params( transaction: WindowSpendTransaction, seed_base: float, 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 index d2ad2aa61d2..8b3b2bf2050 100644 --- a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py @@ -30,6 +30,11 @@ class WindowSpendTransaction(TypedDict): own ~2s poll and will usually have persisted these rows before the window queue flushes; without the exclusion the seed and the increment would each count them. + + started_at is the earliest request start in the batch. The seed only + subtracts a request_id whose LiteLLM_SpendLogs.startTime is at or after it, + so a client that replays an old id through x-litellm-call-id cannot make the + seed drop the historical row that id already paid for. """ entity_type: str @@ -38,6 +43,7 @@ class WindowSpendTransaction(TypedDict): window_start: str spend: float request_ids: Sequence[str] + started_at: str | None def to_naive_utc(value: datetime) -> datetime: @@ -65,6 +71,7 @@ def build_window_spend_transaction( window_start: datetime, spend: float, request_id: str | None = None, + started_at: datetime | None = None, ) -> WindowSpendTransaction: return WindowSpendTransaction( entity_type=entity_type, @@ -73,6 +80,9 @@ def build_window_spend_transaction( window_start=to_naive_utc(window_start).isoformat(timespec="microseconds"), spend=spend, request_ids=() if request_id is None else (request_id,), + started_at=None + if started_at is None + else to_naive_utc(started_at.astimezone(timezone.utc)).isoformat(timespec="microseconds"), ) @@ -80,6 +90,9 @@ def _merge_window_spend_transactions( payloads: tuple[WindowSpendTransaction, ...], ) -> WindowSpendTransaction: first: Final = payloads[0] + started_ats: Final = tuple( + started_at for payload in payloads if (started_at := payload.get("started_at")) is not None + ) return WindowSpendTransaction( entity_type=first["entity_type"], entity_id=first["entity_id"], @@ -87,6 +100,7 @@ def _merge_window_spend_transactions( window_start=first["window_start"], spend=math.fsum(payload["spend"] for payload in payloads), request_ids=tuple(sorted(frozenset(chain.from_iterable(payload["request_ids"] for payload in payloads)))), + started_at=min(started_ats) if started_ats else None, ) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 58bc9cfb301..fed7bff08ff 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -535,6 +535,7 @@ async def _update_database_and_spend_counters( end_user_id=end_user_id, tags=request_tags, request_id=spend_log_request_id, + request_started_at=start_time, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 89fe315fbd4..9a28c274f92 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2396,6 +2396,7 @@ async def increment_spend_counters( end_user_id: str | None = None, tags: list[str] | None = None, request_id: str | None = None, + request_started_at: datetime | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -2468,6 +2469,7 @@ async def increment_spend_counters( window_start=key_window_start, increment=cost, request_id=request_id, + request_started_at=request_started_at, ) async def _team_scope(scope_team_id: str) -> None: @@ -2510,6 +2512,7 @@ async def increment_spend_counters( window_start=team_window_start, increment=cost, request_id=request_id, + request_started_at=request_started_at, ) async def _team_member_scope(scope_user_id: str, scope_team_id: str) -> None: @@ -2703,14 +2706,15 @@ async def _enqueue_window_spend_row_update( window_start: datetime | None, increment: float, request_id: str | None, + request_started_at: datetime | None, ) -> None: """Queue this request's cost against the LiteLLM_BudgetWindowSpend row for the window, so enforcement can read a maintained total instead of aggregating LiteLLM_SpendLogs. - request_id is the LiteLLM_SpendLogs id this cost was recorded under; the - flush uses it to keep the one-time seed from counting a request that its - increment already covers. + request_id is the LiteLLM_SpendLogs id this cost was recorded under and + request_started_at its startTime; the flush uses them to keep the one-time + seed from counting a request that its increment already covers. Enqueued even when the cache increment was skipped for a reserved counter: the reservation only pre-charged the counter, and the row still owes the @@ -2732,6 +2736,7 @@ async def _enqueue_window_spend_row_update( window_start=window_start, spend=increment, request_id=request_id, + started_at=request_started_at, ) ) except Exception as e: # noqa: BLE001 # spend tracking must never fail the cost callback 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 e4a653ddb51..d884a9becaf 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 @@ -442,6 +442,7 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend( window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), spend=1.25, request_id="req-1", + started_at=datetime(2026, 8, 10, 12, 0, 0, tzinfo=timezone.utc), ) ) @@ -466,6 +467,7 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend( "window_start": "2026-08-01T00:00:00.000000", "spend": 1.25, "request_ids": ["req-1"], + "started_at": "2026-08-10T12:00:00.000000", }] 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 index 880cc5044f9..25c685f6f29 100644 --- 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 @@ -24,6 +24,7 @@ def _txn( duration: str = "30d", entity_type: str = "key", request_id: str | None = None, + started_at: datetime | None = None, ): return build_window_spend_transaction( entity_type=entity_type, @@ -32,6 +33,7 @@ def _txn( window_start=window_start, spend=spend, request_id=request_id, + started_at=started_at, ) @@ -47,9 +49,35 @@ def test_build_window_spend_transaction_stores_naive_utc_iso(): "window_start": "2026-08-02T00:00:00.000000", "spend": 1.0, "request_ids": ("req-1",), + "started_at": None, } +def test_build_window_spend_transaction_stores_started_at_as_naive_utc_iso(): + """started_at is compared against LiteLLM_SpendLogs.startTime, which the + spend log writer stores after converting the request start to UTC.""" + non_utc = datetime(2026, 8, 10, 8, 30, 15, 123456, tzinfo=timezone(timedelta(hours=-4))) + + assert _txn("k1", WINDOW_A, 1.0, started_at=non_utc)["started_at"] == "2026-08-10T12:30:15.123456" + + +@pytest.mark.asyncio +async def test_aggregation_keeps_the_earliest_started_at_of_the_batch(): + """The seed bounds its request-id exclusion at the batch's earliest start, + so a later start must never win the merge.""" + queue = WindowSpendUpdateQueue() + earliest = datetime(2026, 8, 10, 12, 0, 0, tzinfo=timezone.utc) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-2", started_at=earliest + timedelta(seconds=5))) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-1", started_at=earliest)) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-3")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 1 + assert aggregated[0]["started_at"] == "2026-08-10T12:00:00.000000" + assert aggregated[0]["request_ids"] == ("req-1", "req-2", "req-3") + + def test_to_naive_utc_leaves_naive_values_alone(): naive = datetime(2026, 8, 1, 12, 0) assert to_naive_utc(naive) == naive 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 index c05e3d410f9..4d7d86130fb 100644 --- a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py +++ b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py @@ -2,7 +2,7 @@ import math import os import sys from contextlib import asynccontextmanager -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Any sys.path.insert(0, os.path.abspath("../../../..")) @@ -20,6 +20,8 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WINDOW_A = datetime(2026, 8, 1, tzinfo=timezone.utc) WINDOW_B = datetime(2026, 8, 31, tzinfo=timezone.utc) +BATCH_STARTED_AT = datetime(2026, 8, 10, 12, 0, 0, 250_000, tzinfo=timezone.utc) +BEFORE_BATCH = BATCH_STARTED_AT - timedelta(hours=1) ENTITY_TYPE, ENTITY_ID, WINDOW_DURATION, WINDOW_START, INSERT_SPEND, INCREMENT, NOW = range(7) @@ -85,6 +87,7 @@ class _RecordingAggregate: entity_id: str, window_start: datetime, exclude_request_ids: Any, + exclude_started_at: datetime | None, ) -> float | None: self.calls.append( { @@ -92,16 +95,18 @@ class _RecordingAggregate: "entity_id": entity_id, "window_start": window_start, "exclude_request_ids": tuple(exclude_request_ids), + "exclude_started_at": exclude_started_at, } ) return self.value class _SpendLogsFake: - """Sums the LiteLLM_SpendLogs rows it holds, honouring the request-id - exclusion exactly as the real aggregate's NOT (request_id = ANY(...)) does.""" + """Sums the LiteLLM_SpendLogs rows (request_id, spend, startTime) it holds, + honouring the exclusion exactly as the real aggregate's + NOT (request_id = ANY(...) AND startTime >= bound) does.""" - def __init__(self, rows: tuple[tuple[str, float], ...]) -> None: + def __init__(self, rows: tuple[tuple[str, float, datetime], ...]) -> None: self.rows = rows async def __call__( @@ -111,9 +116,26 @@ class _SpendLogsFake: entity_id: str, window_start: datetime, exclude_request_ids: Any, + exclude_started_at: datetime | None, ) -> float | None: - excluded = frozenset(exclude_request_ids) - return math.fsum(spend for request_id, spend in self.rows if request_id not in excluded) + excluded = frozenset(exclude_request_ids) if exclude_started_at is not None else frozenset() + return math.fsum( + spend + for request_id, spend, started_at in self.rows + if not (request_id in excluded and started_at >= exclude_started_at) + ) + + +def _batch(request_ids: tuple[str, ...], spend: float, started_at: datetime | None = BATCH_STARTED_AT) -> dict: + return { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": spend, + "request_ids": request_ids, + "started_at": None if started_at is None else started_at.replace(tzinfo=None).isoformat(timespec="microseconds"), + } def _existing(entity_type: str, entity_id: str, window_duration: str) -> dict[str, str]: @@ -321,7 +343,9 @@ async def test_unknown_entity_type_contributes_no_seed(): 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): + async def no_such_column( + prisma_client, entity_type, entity_id, window_start, exclude_request_ids, exclude_started_at + ): return None await commit_window_spend_updates( @@ -338,7 +362,7 @@ async def test_unknown_entity_type_contributes_no_seed(): 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): + async def unavailable(prisma_client, entity_type, entity_id, window_start, exclude_request_ids, exclude_started_at): return None await commit_window_spend_updates( @@ -374,26 +398,32 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o @pytest.mark.asyncio -async def test_seed_receives_the_batch_request_ids_to_exclude(): +async def test_seed_receives_the_batch_request_ids_and_earliest_start_to_exclude(): db = _FakeDB(existing_rows=[]) aggregate = _RecordingAggregate(value=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=( - { - "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"), - }, - ), + transactions=(_batch(("req-1", "req-2", "req-3"), 3.0),), spend_logs_aggregate=aggregate, ) assert aggregate.calls[0]["exclude_request_ids"] == ("req-1", "req-2", "req-3") + assert aggregate.calls[0]["exclude_started_at"] == BATCH_STARTED_AT + + +@pytest.mark.asyncio +async def test_seed_passes_no_start_bound_when_the_batch_has_none(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("req-1",), 1.0, started_at=None),), + spend_logs_aggregate=aggregate, + ) + + assert aggregate.calls[0]["exclude_started_at"] is None @pytest.mark.asyncio @@ -404,21 +434,16 @@ async def test_new_row_is_not_double_counted_when_the_batch_logs_already_flushed 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)), + rows=( + ("req-1", 0.000047, BATCH_STARTED_AT), + ("req-2", 0.000047, BATCH_STARTED_AT + timedelta(seconds=1)), + ("req-3", 0.000047, BATCH_STARTED_AT + timedelta(seconds=2)), + ), ) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=( - { - "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"), - }, - ), + transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),), spend_logs_aggregate=already_flushed, ) @@ -430,20 +455,30 @@ async def test_new_row_is_not_double_counted_when_the_batch_logs_already_flushed 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))) + spend_logs = _SpendLogsFake(rows=(("older", 0.5, BEFORE_BATCH), ("req-1", 0.000047, BATCH_STARTED_AT))) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=( - { - "entity_type": "key", - "entity_id": "k1", - "window_duration": "30d", - "window_start": "2026-08-01T00:00:00.000000", - "spend": 0.000047, - "request_ids": ("req-1",), - }, - ), + transactions=(_batch(("req-1",), 0.000047),), + spend_logs_aggregate=spend_logs, + ) + + (_, params), = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(0.500047) + + +@pytest.mark.asyncio +async def test_replayed_request_id_cannot_erase_historical_spend_from_the_seed(): + """request_id can be chosen by the client via x-litellm-call-id. A request + that replays an id from before this batch writes no new LiteLLM_SpendLogs + row (the insert skips duplicates), so the seed must keep counting the + historical row that id belongs to; only its increment is new.""" + db = _FakeDB(existing_rows=[]) + spend_logs = _SpendLogsFake(rows=(("replayed", 0.5, BEFORE_BATCH),)) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("replayed",), 0.000047),), spend_logs_aggregate=spend_logs, ) @@ -460,16 +495,7 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): 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"), - }, - ), + transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),), spend_logs_aggregate=nothing_flushed, ) @@ -482,7 +508,9 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): "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): +async def test_seed_aggregate_sql_excludes_the_request_ids_only_within_the_batch_start_bound( + entity_type, expected_column +): db = _FakeDB(existing_rows=[{"total": 1.25}]) total = await spend_logs_total_excluding( @@ -491,19 +519,49 @@ async def test_seed_aggregate_sql_excludes_the_request_ids_by_parameter(entity_t entity_id="e1", window_start=WINDOW_A, exclude_request_ids=("req-1", "req-2"), + exclude_started_at=BATCH_STARTED_AT, ) assert total == pytest.approx(1.25) (query, params), = db.query_raw_calls normalized = " ".join(query.split()) assert expected_column in normalized - assert "NOT (request_id = ANY($3::text[]))" in normalized + assert "NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" in normalized assert 'FROM "LiteLLM_SpendLogs"' in normalized - assert params == ("e1", WINDOW_A, ("req-1", "req-2")) + # startTime is TIMESTAMP(3): the bound is floored to the second so the + # batch's own earliest row cannot round under it. + assert params == ("e1", WINDOW_A, ("req-1", "req-2"), datetime(2026, 8, 10, 12, 0, 0)) # The ids are bound, never spliced into the statement. assert "req-1" not in query +@pytest.mark.asyncio +@pytest.mark.parametrize( + "exclude_request_ids, exclude_started_at", + [(("req-1",), None), ((), BATCH_STARTED_AT)], +) +async def test_seed_aggregate_excludes_nothing_without_both_ids_and_a_start_bound( + exclude_request_ids, exclude_started_at +): + """Ids without a start bound would reopen the replayed-id hole, so the + seed counts everything instead; at worst that over-counts one batch.""" + db = _FakeDB(existing_rows=[{"total": 1.25}]) + + total = await spend_logs_total_excluding( + prisma_client=_FakePrismaClient(db), + entity_type="key", + entity_id="e1", + window_start=WINDOW_A, + exclude_request_ids=exclude_request_ids, + exclude_started_at=exclude_started_at, + ) + + assert total == pytest.approx(1.25) + (query, params), = db.query_raw_calls + assert "request_id" not in query + assert params == ("e1", WINDOW_A) + + @pytest.mark.asyncio async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs_column(): db = _FakeDB(existing_rows=[]) @@ -514,6 +572,7 @@ async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs entity_id="u1", window_start=WINDOW_A, exclude_request_ids=(), + exclude_started_at=None, ) assert total is None @@ -530,6 +589,7 @@ async def test_seed_aggregate_treats_an_entity_with_no_rows_as_zero(): entity_id="k-unknown", window_start=WINDOW_A, exclude_request_ids=(), + exclude_started_at=None, ) assert total == 0.0 diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 0c0b4ac7889..9ee79caec5b 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 @@ -445,6 +445,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda ) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} + start_time = datetime.now() await _update_database_and_spend_counters( proxy_logging_obj=proxy_logging_obj, @@ -456,7 +457,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda org_id="test_org_id", kwargs={}, completion_response=None, - start_time=datetime.now(), + start_time=start_time, end_time=datetime.now(), response_cost=0.2, budget_reservation=budget_reservation, @@ -474,6 +475,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda end_user_id="test_end_user_id", tags=["tag-a"], request_id="chatcmpl-abc123", + request_started_at=start_time, ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7e7c2e01fcc..f11d2e00f72 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11234,6 +11234,32 @@ async def test_window_spend_row_carries_the_spend_log_request_id(): assert enqueued[0]["request_ids"] == ("chatcmpl-abc123",) +@pytest.mark.asyncio +async def test_window_spend_row_carries_the_request_start_time(): + """The seed only excludes a batch id whose LiteLLM_SpendLogs.startTime is at + or after this, so it must be the same start the spend log was written with.""" + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", + team_id=None, + user_id=None, + response_cost=0.25, + request_id="chatcmpl-abc123", + request_started_at=datetime(2026, 8, 10, 12, 0, 0, 500_000, tzinfo=timezone.utc), + ) + enqueued = await _drain(queue) + + assert enqueued[0]["started_at"] == "2026-08-10T12:00:00.500000" + + @pytest.mark.asyncio async def test_window_spend_row_without_a_request_id_excludes_nothing(): from litellm.proxy.proxy_server import increment_spend_counters From a6273bb33292c63c9ba689e39644112715dafea6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 11:36:59 -0700 Subject: [PATCH 06/20] chore(migrations): drop the generated comment from the budget window spend migration --- .../20260804162853_add_budget_window_spend_table/migration.sql | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql index 45cc927a328..c3018006adb 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql @@ -1,4 +1,3 @@ --- CreateTable CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" ( "entity_type" TEXT NOT NULL, "entity_id" TEXT NOT NULL, From 0cd89148e313653274f9dacaf0bc1f0706e4c5f5 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 4 Aug 2026 16:56:37 -0700 Subject: [PATCH 07/20] feat(proxy): add LiteLLM_BudgetWindowSpend table for per-window budget spend Multi-window budgets (budget_limits on keys/teams) currently keep window spend only in cache. Every cold or expired counter recomputes the window by aggregating LiteLLM_SpendLogs, which has no usable index for that query and saturates the DB on large tables (#35766). This adds a LiteLLM_BudgetWindowSpend table holding one row per configured window, keyed (entity_type, entity_id, window_duration), with window_start identifying the period the spend belongs to. Follow-up PRs maintain these rows from the spend update writer and move window budget enforcement reads onto them. --- .../migration.sql | 13 +++++++++++++ .../litellm_proxy_extras/schema.prisma | 12 ++++++++++++ litellm/proxy/schema.prisma | 12 ++++++++++++ schema.prisma | 12 ++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 5 files changed, 50 insertions(+), 1 deletion(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql new file mode 100644 index 00000000000..45cc927a328 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql @@ -0,0 +1,13 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" ( + "entity_type" TEXT NOT NULL, + "entity_id" TEXT NOT NULL, + "window_duration" TEXT NOT NULL, + "window_start" TIMESTAMP(3) NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_BudgetWindowSpend_pkey" PRIMARY KEY ("entity_type","entity_id","window_duration") +); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 2bb850139a2..7818e3acd56 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -649,6 +649,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 2bb850139a2..7818e3acd56 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -649,6 +649,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/schema.prisma b/schema.prisma index 2bb850139a2..7818e3acd56 100644 --- a/schema.prisma +++ b/schema.prisma @@ -649,6 +649,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index eaa05ddc005..558a54c01fb 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25945,7 +25945,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; + user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; }; /** * DefaultTeamSSOParams From 01b7610d0a091ec3c86b8e9645570872e1af66dc Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 11:36:59 -0700 Subject: [PATCH 08/20] chore(migrations): drop the generated comment from the budget window spend migration --- .../20260804162853_add_budget_window_spend_table/migration.sql | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql index 45cc927a328..c3018006adb 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql @@ -1,4 +1,3 @@ --- CreateTable CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" ( "entity_type" TEXT NOT NULL, "entity_id" TEXT NOT NULL, From cfe5e37e9566cd86046d07d420ce4326b3581d5a Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 12:27:49 -0700 Subject: [PATCH 09/20] chore(ui): drop unrelated schema.d.ts enum reorder from the window spend schema branch --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 558a54c01fb..eaa05ddc005 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25945,7 +25945,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; + user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; }; /** * DefaultTeamSSOParams From 6ef42f5991dba70076b7e5b9247ff62f25155f49 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 12:27:57 -0700 Subject: [PATCH 10/20] style(proxy): satisfy ANN204 and SIM117 in the budget window spend writer --- litellm/proxy/db/budget_window_spend_writer.py | 16 +++++++++------- .../window_spend_update_queue.py | 2 +- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index 02ff9c4f944..d8c2c7c2695 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -275,13 +275,15 @@ async def commit_window_spend_updates( 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 with ( + prisma_client.db.tx(timeout=_UPSERT_TRANSACTION_TIMEOUT) as db_transaction, + db_transaction.batch_() as batcher, + ): + for transaction, seed_base in zip(ordered, seed_bases): + batcher.execute_raw( + _UPSERT_WINDOW_SPEND_SQL, + *_upsert_params(transaction=transaction, seed_base=seed_base, now=now), + ) async def roll_window_spend_row( 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 index 8b3b2bf2050..2e7366ce85f 100644 --- a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py @@ -121,7 +121,7 @@ class WindowSpendUpdateQueue(BaseUpdateQueue): ) """ - def __init__(self): + def __init__(self) -> None: super().__init__() self.update_queue: asyncio.Queue[tuple[WindowSpendTransaction, ...]] = asyncio.Queue( maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE From 12aea29bad6d1807fd0d040217385e92648d9265 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 12:28:59 -0700 Subject: [PATCH 11/20] chore(ui): keep schema.d.ts in sync with staging on the window spend writer branch --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 558a54c01fb..eaa05ddc005 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25945,7 +25945,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; + user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; }; /** * DefaultTeamSSOParams From ea5a9f6ef01d5d63a2ef0363e38cf8a00e16ed34 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 12:38:26 -0700 Subject: [PATCH 12/20] style(proxy): mark WindowSpendTransaction fields ReadOnly for the LIT012 gate --- .../window_spend_update_queue.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) 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 index 2e7366ce85f..04dea66165e 100644 --- a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py @@ -13,6 +13,8 @@ from datetime import datetime, timezone from itertools import chain, groupby from typing import Final, TypedDict +from typing_extensions import ReadOnly + from litellm._logging import verbose_proxy_logger from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE from litellm.proxy.db.db_transaction_queue.base_update_queue import BaseUpdateQueue @@ -37,13 +39,13 @@ class WindowSpendTransaction(TypedDict): seed drop the historical row that id already paid for. """ - entity_type: str - entity_id: str - window_duration: str - window_start: str - spend: float - request_ids: Sequence[str] - started_at: str | None + entity_type: ReadOnly[str] + entity_id: ReadOnly[str] + window_duration: ReadOnly[str] + window_start: ReadOnly[str] + spend: ReadOnly[float] + request_ids: ReadOnly[Sequence[str]] + started_at: ReadOnly[str | None] def to_naive_utc(value: datetime) -> datetime: From f962a1188e2edde92722127e3d690ad9cdcb5c1a Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 12:42:06 -0700 Subject: [PATCH 13/20] style(proxy): justify the blanket except in window_from_table for the BLE001 gate --- litellm/proxy/db/spend_counter_reseed.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 4513300a12c..7b3c261036e 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -266,7 +266,7 @@ class SpendCounterReseed: } } ) - except Exception: + except Exception: # noqa: BLE001 # any read failure (DB, stale prisma client) must degrade to the aggregate path verbose_proxy_logger.exception( "SpendCounterReseed.window_from_table: failed for %s=%s window=%s", entity_type, From 56dd4e06accc47a83b031c241f532c3b9a07ce1b Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 13:29:59 -0700 Subject: [PATCH 14/20] test(proxy): satisfy the test-quality gate for the window spend writer tests --- .../test_window_spend_update_queue.py | 12 +-- .../db/test_budget_window_spend_writer.py | 43 ++++----- .../proxy/db/test_db_spend_update_writer.py | 88 +++++++------------ 3 files changed, 58 insertions(+), 85 deletions(-) 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 index 25c685f6f29..b1ecda57afa 100644 --- 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 @@ -1,10 +1,6 @@ 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 ( @@ -207,9 +203,7 @@ def test_aggregation_survives_the_redis_json_round_trip(): [(_txn("k1", WINDOW_A, 1.0),), (_txn("k1", WINDOW_B, 2.0),)] ) - reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( - [json.loads(json.dumps(aggregated))] - ) + reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions([json.loads(json.dumps(aggregated))]) assert reloaded == aggregated @@ -269,9 +263,7 @@ def test_request_ids_survive_the_redis_json_round_trip(): [(_txn("k1", WINDOW_A, 1.0, request_id="req-1"),)] ) - reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( - [json.loads(json.dumps(aggregated))] - ) + 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 index 4d7d86130fb..a849317c930 100644 --- a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py +++ b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py @@ -1,12 +1,8 @@ import math -import os -import sys from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone from typing import Any -sys.path.insert(0, os.path.abspath("../../../..")) - import pytest from litellm.proxy.db.budget_window_spend_writer import ( @@ -134,7 +130,9 @@ def _batch(request_ids: tuple[str, ...], spend: float, started_at: datetime | No "window_start": "2026-08-01T00:00:00.000000", "spend": spend, "request_ids": request_ids, - "started_at": None if started_at is None else started_at.replace(tzinfo=None).isoformat(timespec="microseconds"), + "started_at": None + if started_at is None + else started_at.replace(tzinfo=None).isoformat(timespec="microseconds"), } @@ -171,7 +169,7 @@ async def test_missing_row_is_seeded_from_spend_logs_once(): assert aggregate.calls[0]["entity_id"] == "k1" assert aggregate.calls[0]["window_start"] == WINDOW_A - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[ENTITY_TYPE] == "key" assert params[ENTITY_ID] == "k1" assert params[WINDOW_DURATION] == "30d" @@ -194,7 +192,7 @@ async def test_existing_row_is_never_reseeded(): ) assert aggregate.calls == [] - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(1.0) assert params[INCREMENT] == pytest.approx(1.0) @@ -233,7 +231,7 @@ async def test_insert_spend_and_increment_differ_only_when_a_row_is_seeded(): spend_logs_aggregate=aggregate, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(9.25) assert params[INCREMENT] == pytest.approx(0.25) @@ -249,7 +247,7 @@ async def test_upsert_sql_adds_for_a_current_window_and_replaces_for_a_newer_one transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), ) - (query, _), = db.batcher.calls + ((query, _),) = db.batcher.calls normalized = " ".join(query.split()) assert ( 'spend = CASE WHEN "LiteLLM_BudgetWindowSpend".window_start >= EXCLUDED.window_start ' @@ -270,7 +268,7 @@ async def test_upsert_never_interpolates_values_into_the_sql(): spend_logs_aggregate=aggregate, ) - (query, params), = db.batcher.calls + ((query, params),) = db.batcher.calls assert "DROP TABLE" not in query assert params[ENTITY_ID] == "'; DROP TABLE x; --" @@ -293,7 +291,10 @@ async def test_upserts_are_ordered_by_primary_key_then_window_start(): spend_logs_aggregate=aggregate, ) - ordered = [(params[ENTITY_TYPE], params[ENTITY_ID], params[WINDOW_DURATION], params[WINDOW_START]) for _, params in db.batcher.calls] + 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)), @@ -316,7 +317,7 @@ async def test_existing_row_lookup_sends_every_primary_key_as_array_params(): spend_logs_aggregate=aggregate, ) - (query, params), = db.query_raw_calls + ((query, params),) = db.query_raw_calls assert "unnest($1::text[], $2::text[], $3::text[])" in query assert params == (("key", "team"), ("k1", "t1"), ("30d", "7d")) @@ -354,7 +355,7 @@ async def test_unknown_entity_type_contributes_no_seed(): spend_logs_aggregate=no_such_column, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(1.0) @@ -371,7 +372,7 @@ async def test_unavailable_spend_logs_aggregate_seeds_zero_rather_than_failing() spend_logs_aggregate=unavailable, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(1.0) @@ -389,7 +390,7 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o new_window_start=WINDOW_B, ) - (query, params), = db.execute_raw_calls + ((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 @@ -447,7 +448,7 @@ async def test_new_row_is_not_double_counted_when_the_batch_logs_already_flushed spend_logs_aggregate=already_flushed, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(0.000141) @@ -463,7 +464,7 @@ async def test_new_row_still_covers_spend_that_predates_the_batch(): spend_logs_aggregate=spend_logs, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(0.500047) @@ -482,7 +483,7 @@ async def test_replayed_request_id_cannot_erase_historical_spend_from_the_seed() spend_logs_aggregate=spend_logs, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(0.500047) @@ -499,7 +500,7 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): spend_logs_aggregate=nothing_flushed, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(0.000141) @@ -523,7 +524,7 @@ async def test_seed_aggregate_sql_excludes_the_request_ids_only_within_the_batch ) assert total == pytest.approx(1.25) - (query, params), = db.query_raw_calls + ((query, params),) = db.query_raw_calls normalized = " ".join(query.split()) assert expected_column in normalized assert "NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" in normalized @@ -557,7 +558,7 @@ async def test_seed_aggregate_excludes_nothing_without_both_ids_and_a_start_boun ) assert total == pytest.approx(1.25) - (query, params), = db.query_raw_calls + ((query, params),) = db.query_raw_calls assert "request_id" not in query assert params == ("e1", WINDOW_A) 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 8165dbbe3d1..e07d326e8be 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 @@ -67,9 +67,7 @@ async def test_daily_spend_tracking_with_disabled_spend_logs(): assert db_writer.add_spend_log_transaction_to_daily_user_transaction.called # Verify the payload passed to add_spend_log_transaction_to_daily_user_transaction - call_args = ( - db_writer.add_spend_log_transaction_to_daily_user_transaction.call_args[1] - ) + call_args = db_writer.add_spend_log_transaction_to_daily_user_transaction.call_args[1] assert "payload" in call_args assert call_args["payload"]["spend"] == 0.1 assert call_args["payload"]["model"] == "gpt-4" @@ -409,7 +407,7 @@ async def test_update_daily_spend_sorting(): # fields, but entity_id is sufficient to test sorting. daily_spend_transactions = { f"test_key_{i}": { - "user_id": f"user{60-i}", # user60 ... user11, reverse order + "user_id": f"user{60 - i}", # user60 ... user11, reverse order "date": "2024-01-01", "api_key": "test-api-key", "model": "gpt-4", @@ -988,9 +986,9 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i transaction_dict = call[1]["update"] # Each transaction should have one key with the format tag_date_api_key_model_provider for key, transaction in transaction_dict.items(): - assert ( - transaction["request_id"] == request_id - ), f"request_id should be {request_id} but got {transaction.get('request_id')}" + assert transaction["request_id"] == request_id, ( + f"request_id should be {request_id} but got {transaction.get('request_id')}" + ) @pytest.mark.asyncio @@ -1216,21 +1214,15 @@ async def test_add_spend_log_transaction_to_daily_agent_transaction_calls_common } writer.daily_agent_spend_update_queue.add_update = AsyncMock() - original_common_helper = ( - writer._common_add_spend_log_transaction_to_daily_transaction - ) - writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock( - wraps=original_common_helper - ) + original_common_helper = writer._common_add_spend_log_transaction_to_daily_transaction + writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock(wraps=original_common_helper) await writer.add_spend_log_transaction_to_daily_agent_transaction( payload=payload, prisma_client=mock_prisma, ) - assert ( - writer._common_add_spend_log_transaction_to_daily_transaction.await_count == 1 - ) + assert writer._common_add_spend_log_transaction_to_daily_transaction.await_count == 1 @pytest.mark.asyncio @@ -1385,6 +1377,7 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): Test that when batch upsert fails, the exception is properly re-raised after logging. This ensures that error handling continues to work correctly upstream. """ + def raise_connection_lost(): raise ValueError("Database connection lost") @@ -1565,9 +1558,7 @@ async def test_update_database_creates_single_task(): patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), - patch( - "litellm.proxy.db.db_spend_update_writer.asyncio.create_task" - ) as mock_create_task, + patch("litellm.proxy.db.db_spend_update_writer.asyncio.create_task") as mock_create_task, ): await db_writer.update_database( token="test-token", @@ -1666,9 +1657,7 @@ async def test_daily_agent_receives_deepcopied_payload(): db_writer._update_agent_db = AsyncMock() db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() - db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock( - side_effect=capture_agent_payload - ) + db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock(side_effect=capture_agent_payload) db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock() @@ -1730,8 +1719,8 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock() # Return all-None tuple (no data to commit); the pipeline yields 6 slots - mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = ( - AsyncMock(return_value=(None, None, None, None, None, None, None)) + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, None) ) db_writer.redis_update_buffer = mock_redis_update_buffer @@ -2159,9 +2148,7 @@ async def test_update_database_does_not_deepcopy_on_request_path(): db_writer._update_org_db = AsyncMock() db_writer._update_tag_db = AsyncMock() db_writer._update_agent_db = AsyncMock() - db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock( - side_effect=capture_batch_payload - ) + db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock(side_effect=capture_batch_payload) db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() @@ -2253,9 +2240,7 @@ async def test_spend_update_path_never_queries_user_cache_with_none_user_id(): db_writer = DBSpendUpdateWriter() strict_redis_backed_cache = MagicMock() - strict_redis_backed_cache.async_get_cache = AsyncMock( - side_effect=DataError("Invalid input of type: 'NoneType'") - ) + strict_redis_backed_cache.async_get_cache = AsyncMock(side_effect=DataError("Invalid input of type: 'NoneType'")) with ( patch.object(litellm, "max_budget", 0), @@ -2383,8 +2368,7 @@ async def test_daily_transaction_carries_compression_saved_tokens(): cache_write_cost = model_info.get("cache_creation_input_token_cost") or input_cost assert transaction["compression_savings_spend"] == pytest.approx(7600 * input_cost) assert transaction["prompt_caching_savings_spend"] == pytest.approx( - 40 * max(input_cost - cache_read_cost, 0.0) - - 15 * (cache_write_cost - input_cost) + 40 * max(input_cost - cache_read_cost, 0.0) - 15 * (cache_write_cost - input_cost) ) assert transaction["compression_savings_spend"] > 0 assert transaction["prompt_caching_savings_spend"] > 0 @@ -2472,9 +2456,7 @@ class _WindowSpendFakePrisma: def _window_spend_upserts(db): - return [ - params for query, params in db.batcher.calls if "LiteLLM_BudgetWindowSpend" in query - ] + return [params for query, params in db.batcher.calls if "LiteLLM_BudgetWindowSpend" in query] @pytest.mark.asyncio @@ -2492,9 +2474,7 @@ async def test_window_spend_queue_is_flushed_without_redis_buffer(): ) ) db = _WindowSpendFakeDB( - existing_rows=[ - {"entity_type": "key", "entity_id": "hashed-token", "window_duration": "30d"} - ] + existing_rows=[{"entity_type": "key", "entity_id": "hashed-token", "window_duration": "30d"}] ) await db_writer._commit_spend_updates_to_db_without_redis_buffer( @@ -2554,9 +2534,7 @@ async def test_window_spend_transactions_from_redis_are_committed_by_the_lock_wi 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"}] - ) + 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), @@ -2627,9 +2605,12 @@ async def test_update_database_returns_the_spend_log_request_id(): 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"), + patch.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam + "litellm.proxy.proxy_server", + disable_spend_logs=False, + prisma_client=MagicMock(), + litellm_proxy_budget_name="test-budget", + ) ): request_id = await db_writer.update_database( token="test-token", @@ -2655,10 +2636,13 @@ 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( + patch.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam + "litellm.proxy.proxy_server", + disable_spend_logs=False, + prisma_client=MagicMock(), + litellm_proxy_budget_name="test-budget", + ), + patch( # test-quality-ok: the payload builder is called by name inside update_database; no injection seam "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", side_effect=Exception("payload boom"), ), @@ -2979,9 +2963,7 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey "call_type, expects_flush", [("aresponses", True), ("responses", True), ("acompletion", False)], ) -async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls( - call_type: str, expects_flush: bool -): +async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(call_type: str, expects_flush: bool): """ A `previous_response_id` chained straight off the previous turn reads the DB, so a Responses row cannot sit in this worker's queue until the monitor's next poll. @@ -3011,9 +2993,7 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls( pytest.param("", True, id="injected-before-a-deployment-was-chosen"), ], ) -async def test_caching_savings_are_attributed_to_the_deployment_that_was_injected( - injected_deployment, attributed -): +async def test_caching_savings_are_attributed_to_the_deployment_that_was_injected(injected_deployment, attributed): """Retries, same-group failover and cross-model-group fallbacks all reuse one metadata bucket and one litellm_call_id, so a marker written by the leg that injected is visible to every sibling and nothing request-scoped can tell them apart. From a5f47a271ad982052ac7057a28491a7ea15ad8b8 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 13:55:36 -0700 Subject: [PATCH 15/20] fix(proxy): re-queue budget window spend increments when the commit fails Budget enforcement trusts a current LiteLLM_BudgetWindowSpend row without reconciling it against LiteLLM_SpendLogs, so an increment dropped after a failed commit let the entity spend past its window limit after the next counter reseed. Failed increments now go back on the in-memory queue, or back to the Redis buffer, and retry on the next scheduler tick like every other spend category. --- litellm/proxy/db/db_spend_update_writer.py | 44 +++--- .../redis_update_buffer.py | 2 + .../test_redis_update_buffer.py | 140 +++++++++--------- .../proxy/db/test_db_spend_update_writer.py | 65 ++++++-- 4 files changed, 151 insertions(+), 100 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 193f41901cd..1ce7a959fec 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -936,6 +936,7 @@ class DBSpendUpdateWriter: "daily_org_spend_update_transactions": daily_org_spend_update_transactions, "daily_end_user_spend_update_transactions": daily_end_user_spend_update_transactions, "daily_agent_spend_update_transactions": daily_agent_spend_update_transactions, + "window_spend_update_transactions": window_spend_update_transactions, } if db_spend_update_transactions is not None: @@ -1008,6 +1009,7 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, window_spend_transactions=window_spend_update_transactions, ) + uncommitted.pop("window_spend_update_transactions", None) except Exception as e: spend_log_error( "Spend tracking - failed to commit spend updates from Redis to DB. " @@ -1129,10 +1131,20 @@ class DBSpendUpdateWriter: 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, - ) + try: + await DBSpendUpdateWriter._commit_window_spend_updates( + prisma_client=prisma_client, + window_spend_transactions=window_spend_update_transactions, + ) + except Exception as e: # noqa: BLE001 # the increments go back on the queue; the rest of the flush must run + spend_log_error( + "Spend tracking - failed to commit budget window spend updates. " + "Re-queued %d window increments for retry on next tick. Error: %s", + len(window_spend_update_transactions), + str(e), + exc=e, + ) + await self.window_spend_update_queue.update_queue.put(window_spend_update_transactions) ################## Tool Registry Upserts ################## await self._flush_tool_discovery_queue(prisma_client=prisma_client) @@ -1206,27 +1218,19 @@ class DBSpendUpdateWriter: """ 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. + Raises on failure so the caller re-queues the increments: budget + enforcement trusts a current row without reconciling it against + LiteLLM_SpendLogs, so a dropped increment would let the entity spend + past its window limit after the next counter reseed. """ from litellm.proxy.db.budget_window_spend_writer import ( commit_window_spend_updates, ) - 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, - ) + await commit_window_spend_updates( + prisma_client=prisma_client, + transactions=window_spend_transactions, + ) async def _drain_and_commit_daily_tag_spend_from_redis( self, 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 8c53c03031b..fa096fa0bf2 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -456,6 +456,7 @@ class RedisUpdateBuffer: daily_end_user_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, daily_agent_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, daily_tag_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + window_spend_update_transactions: Sequence[WindowSpendTransaction] | None = None, ) -> None: """ Re-push transactions that were popped from Redis but not committed to the DB. @@ -477,6 +478,7 @@ class RedisUpdateBuffer: (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY), (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY), (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY), + (window_spend_update_transactions, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY), ) rpush_list: Final = tuple( 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 78aa10929e8..504654e103a 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -1,4 +1,5 @@ import json +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -22,9 +23,7 @@ def redis_update_buffer(mock_redis_cache): @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_uses_pipeline( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache): """ Verify store_in_memory_spend_updates_in_redis calls async_rpush_pipeline once with the correct operations and skips empty queues. @@ -33,35 +32,29 @@ async def test_store_in_memory_spend_updates_uses_pipeline( # Create mock queues - only 3 of 6 have data spend_update_queue = AsyncMock() - spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = ( - AsyncMock(return_value={"key_list_transactions": {"key1": 1.0}}) + spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock( + return_value={"key_list_transactions": {"key1": 1.0}} ) daily_spend_queue = AsyncMock() - daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={"user_key1": {"spend": 1.0}}) + daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={"user_key1": {"spend": 1.0}} ) daily_team_queue = AsyncMock() - daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={"team_key1": {"spend": 2.0}}) + daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={"team_key1": {"spend": 2.0}} ) # Empty queues daily_org_queue = AsyncMock() - daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={}) - ) + daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) daily_end_user_queue = AsyncMock() - daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value=None) - ) + daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value=None) daily_agent_queue = AsyncMock() - daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={}) - ) + daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) await redis_update_buffer.store_in_memory_spend_updates_in_redis( spend_update_queue=spend_update_queue, @@ -82,9 +75,7 @@ async def test_store_in_memory_spend_updates_uses_pipeline( @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_restores_on_rpush_failure( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_restores_on_rpush_failure(redis_update_buffer, mock_redis_cache): """ If async_rpush_pipeline raises, the already-drained transactions must be put back into the in-memory queues so the next scheduler tick retries. @@ -98,9 +89,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( SpendUpdateQueue, ) - mock_redis_cache.async_rpush_pipeline = AsyncMock( - side_effect=ConnectionError("redis went away") - ) + mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away")) spend_queue = SpendUpdateQueue() daily_user_queue = DailySpendUpdateQueue() @@ -145,16 +134,12 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( # After restore, the main spend queue should hold one item per # (entity_type, entity_id) pair with the aggregated cost - restored_spend = ( - await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() - ) + restored_spend = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() assert restored_spend["key_list_transactions"] == {"key-abc": 1.5} assert restored_spend["team_list_transactions"] == {"team-xyz": 2.5} # Daily user queue should hold the same aggregated dict - restored_daily = ( - await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions() - ) + restored_daily = await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions() assert restored_daily == { "user1_day_model": { "spend": 1.0, @@ -165,9 +150,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_all_empty_returns_early( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_update_buffer, mock_redis_cache): """ When all queues are empty, pipeline should never be called. """ @@ -175,13 +158,9 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early( # All queues return empty empty_queue = AsyncMock() - empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock( - return_value={} - ) + empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={}) empty_daily_queue = AsyncMock() - empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={}) - ) + empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) await redis_update_buffer.store_in_memory_spend_updates_in_redis( spend_update_queue=empty_queue, @@ -196,9 +175,7 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early( @pytest.mark.asyncio -async def test_get_all_transactions_from_redis_buffer_pipeline( - redis_update_buffer, mock_redis_cache -): +async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buffer, mock_redis_cache): """ Verify get_all_transactions_from_redis_buffer_pipeline correctly parses and aggregates results from async_lpop_pipeline. @@ -287,9 +264,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( 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"] - ] + 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 @@ -302,9 +277,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): @pytest.mark.asyncio -async def test_restore_transactions_to_redis_pushes_only_provided( - redis_update_buffer, mock_redis_cache -): +async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_buffer, mock_redis_cache): """ restore_transactions_to_redis re-pushes only the transaction sets it was given, to their matching buffer keys, so uncommitted spend can be retried. @@ -338,9 +311,42 @@ async def test_restore_transactions_to_redis_pushes_only_provided( @pytest.mark.asyncio -async def test_restore_transactions_to_redis_noop_when_empty( - redis_update_buffer, mock_redis_cache -): +async def test_restored_window_spend_transactions_drain_back_unchanged(redis_update_buffer, mock_redis_cache): + """A window commit that fails after the destructive lpop must be re-pushed + in the store path's encoding, so the next drain returns the same increments.""" + from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, + ) + + window_transactions = ( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=3.0, + request_id="req-1", + started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc), + ), + ) + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) + + await redis_update_buffer.restore_transactions_to_redis(window_spend_update_transactions=window_transactions) + + rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] + assert [op["key"] for op in rpush_list] == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY] + + mock_redis_cache.async_lpop_pipeline = AsyncMock( + return_value=[None, None, None, None, None, None, list(rpush_list[0]["values"])] + ) + drained = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + + assert drained[6] == window_transactions + + +@pytest.mark.asyncio +async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer, mock_redis_cache): """Nothing to restore -> no Redis call.""" mock_redis_cache.async_rpush_pipeline = AsyncMock() await redis_update_buffer.restore_transactions_to_redis() @@ -348,15 +354,11 @@ async def test_restore_transactions_to_redis_noop_when_empty( @pytest.mark.asyncio -async def test_restore_transactions_to_redis_swallows_redis_error( - redis_update_buffer, mock_redis_cache -): +async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_buffer, mock_redis_cache): """A Redis failure during restore must not propagate to the caller's finally block.""" from redis.exceptions import RedisError - mock_redis_cache.async_rpush_pipeline = AsyncMock( - side_effect=RedisError("redis down") - ) + mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=RedisError("redis down")) await redis_update_buffer.restore_transactions_to_redis( db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}}, @@ -472,9 +474,7 @@ def test_get_transaction_buffer_redis_cache_parses_string_flag(monkeypatch): @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_pushes_budget_window_spend( - redis_update_buffer, mock_redis_cache -): +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 @@ -519,15 +519,17 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend( assert len(rpush_list) == 1 assert rpush_list[0]["key"] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY pushed = json.loads(rpush_list[0]["values"][0]) - assert pushed == [{ - "entity_type": "key", - "entity_id": "hashed-token", - "window_duration": "30d", - "window_start": "2026-08-01T00:00:00.000000", - "spend": 1.25, - "request_ids": ["req-1"], - "started_at": "2026-08-10T12:00:00.000000", - }] + assert pushed == [ + { + "entity_type": "key", + "entity_id": "hashed-token", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 1.25, + "request_ids": ["req-1"], + "started_at": "2026-08-10T12:00:00.000000", + } + ] @pytest.mark.asyncio 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 e07d326e8be..d28cf8c9c6a 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2569,19 +2569,21 @@ async def test_window_spend_transactions_are_not_committed_without_the_pod_lock( @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.""" +async def test_failed_window_spend_commit_requeues_the_increments_and_continues_the_flush(): + """Budget enforcement trusts a current window row without reconciling it + against LiteLLM_SpendLogs, so a dropped increment would let the key spend + past its limit after the next reseed. The increments must go back on the + queue, and the tool registry flush must still run.""" db_writer = DBSpendUpdateWriter() - 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, - ) + transaction = build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=0.5, + request_id="req-1", ) + await db_writer.window_spend_update_queue.add_update(transaction) db = _WindowSpendFakeDB() db.query_raw = AsyncMock(side_effect=Exception("connection reset")) db_writer._flush_tool_discovery_queue = AsyncMock() @@ -2593,6 +2595,47 @@ async def test_failed_window_spend_commit_does_not_abort_the_rest_of_the_flush() ) db_writer._flush_tool_discovery_queue.assert_called_once() + requeued = await db_writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() + assert requeued == (transaction,) + + +@pytest.mark.asyncio +async def test_failed_window_spend_commit_from_redis_is_restored_to_redis(): + """The Redis drain is destructive, so a failed window commit has to push + the popped increments back exactly like the other spend categories.""" + db_writer = DBSpendUpdateWriter() + window_transactions = ( + build_window_spend_transaction( + entity_type="team", + entity_id="team-1", + window_duration="7d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=2.0, + request_id="req-1", + ), + ) + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, window_transactions) + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + db = _WindowSpendFakeDB() + db.query_raw = AsyncMock(side_effect=Exception("connection reset")) + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + assert _window_spend_upserts(db) == [] + mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once_with( + window_spend_update_transactions=window_transactions + ) + db_writer.pod_lock_manager.release_lock.assert_awaited_once() @pytest.mark.asyncio From 085d5f5d4e0ba321eecf43f677c7b2f16a01dfd7 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 14:18:41 -0700 Subject: [PATCH 16/20] docs(proxy): state that the cross-pod window seed gap over-counts, never under-counts --- litellm/proxy/db/budget_window_spend_writer.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index d8c2c7c2695..f9188f95cfd 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -11,9 +11,13 @@ 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. +include spend logs whose increments are still queued on another pod, and those +increments are added again when that pod flushes. That is bounded by a single +flush interval, happens at most once per window row, and only ever over-counts: +the seed never omits spend, because every increment not yet in the row still +reaches it on its own pod's next flush. A row therefore lags real spend by at +most one flush interval of queued increments, the same lag the SpendLogs +aggregate it replaces (and every other spend column) already has. """ from collections.abc import Sequence From d28621685a6b0038b52d52dc43701fce0cc71dc1 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 4 Aug 2026 16:56:37 -0700 Subject: [PATCH 17/20] feat(proxy): add LiteLLM_BudgetWindowSpend table for per-window budget spend Multi-window budgets (budget_limits on keys/teams) currently keep window spend only in cache. Every cold or expired counter recomputes the window by aggregating LiteLLM_SpendLogs, which has no usable index for that query and saturates the DB on large tables (#35766). This adds a LiteLLM_BudgetWindowSpend table holding one row per configured window, keyed (entity_type, entity_id, window_duration), with window_start identifying the period the spend belongs to. Follow-up PRs maintain these rows from the spend update writer and move window budget enforcement reads onto them. --- .../migration.sql | 13 +++++++++++++ .../litellm_proxy_extras/schema.prisma | 12 ++++++++++++ litellm/proxy/schema.prisma | 12 ++++++++++++ schema.prisma | 12 ++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 5 files changed, 50 insertions(+), 1 deletion(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql new file mode 100644 index 00000000000..45cc927a328 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql @@ -0,0 +1,13 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" ( + "entity_type" TEXT NOT NULL, + "entity_id" TEXT NOT NULL, + "window_duration" TEXT NOT NULL, + "window_start" TIMESTAMP(3) NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_BudgetWindowSpend_pkey" PRIMARY KEY ("entity_type","entity_id","window_duration") +); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 2bb850139a2..7818e3acd56 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -649,6 +649,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 2bb850139a2..7818e3acd56 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -649,6 +649,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/schema.prisma b/schema.prisma index 2bb850139a2..7818e3acd56 100644 --- a/schema.prisma +++ b/schema.prisma @@ -649,6 +649,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6c5319f0361..c1a85a297c9 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26073,7 +26073,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; + user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; }; /** * DefaultTeamSSOParams From a923502132a05e0ea52f3430d0abd97de66048b5 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 11:36:59 -0700 Subject: [PATCH 18/20] chore(migrations): drop the generated comment from the budget window spend migration --- .../20260804162853_add_budget_window_spend_table/migration.sql | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql index 45cc927a328..c3018006adb 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql @@ -1,4 +1,3 @@ --- CreateTable CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" ( "entity_type" TEXT NOT NULL, "entity_id" TEXT NOT NULL, From 774954d19a578061ec43abd74d4da6d1a76b58ea Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 12:27:49 -0700 Subject: [PATCH 19/20] chore(ui): drop unrelated schema.d.ts enum reorder from the window spend schema branch --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index c1a85a297c9..6c5319f0361 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26073,7 +26073,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; + user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; }; /** * DefaultTeamSSOParams From ce96db5a6154d597a3689a7a6999d787c033ea9d Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 16:36:04 -0700 Subject: [PATCH 20/20] test(proxy): pass the window spend args in the access group requeue test --- tests/test_litellm/proxy/db/test_model_access_group_spend.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_litellm/proxy/db/test_model_access_group_spend.py b/tests/test_litellm/proxy/db/test_model_access_group_spend.py index 79dffc29e7e..d2d079bb0e4 100644 --- a/tests/test_litellm/proxy/db/test_model_access_group_spend.py +++ b/tests/test_litellm/proxy/db/test_model_access_group_spend.py @@ -414,12 +414,14 @@ async def test_redis_buffer_requeues_access_group_transactions_as_queue_items(): daily_org_spend_update_transactions=None, daily_end_user_spend_update_transactions=None, daily_agent_spend_update_transactions=None, + window_spend_update_transactions=None, spend_update_queue=queue, daily_spend_update_queue=daily_queue, daily_team_spend_update_queue=daily_queue, daily_org_spend_update_queue=daily_queue, daily_end_user_spend_update_queue=daily_queue, daily_agent_spend_update_queue=daily_queue, + window_spend_update_queue=None, ) updates = await _drain(queue)