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/60] 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/60] 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/60] 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 6f24d85490f03c8af4f6b94e93aff0ffa2240c1b Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 25 Aug 2026 18:26:57 +0000 Subject: [PATCH 04/60] fix(proxy): grant users with empty models list direct access in model listing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 4 +- .../test_team_model_name_translation.py | 55 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0abcdeaf3f6..3d3169eeab9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12382,8 +12382,10 @@ def get_direct_access_models( The 'all-proxy-models' sentinel grants direct access to every non-team deployment, mirroring how get_key_models expands it for the key/team path. + An empty models list means unrestricted access at call time (see + can_user_call_model), so it resolves the same way. """ - if SpecialModelNames.all_proxy_models.value in user_db_object.models: + if not user_db_object.models or SpecialModelNames.all_proxy_models.value in user_db_object.models: return llm_router.get_model_ids(exclude_team_models=True) return [ diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 2f4018b55ab..cb1521f10ae 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -1507,6 +1507,61 @@ def test_get_direct_access_models_resolves_explicit_model_names(): router.get_model_list.assert_called_once_with(model_name="gpt-4o") +def test_get_direct_access_models_empty_models_grants_all_non_team_models(): + """An empty user.models list means unrestricted access at call time + (can_user_call_model), so the listing must resolve it like 'all-proxy-models' + instead of returning nothing. Regression for a user with models=[] and no + teams seeing an empty Models+Endpoints page.""" + router = MagicMock() + router.get_model_ids.return_value = ["global-id-1", "global-id-2"] + + user = LiteLLM_UserTable(user_id="u", models=[], teams=[]) + + result = ps.get_direct_access_models(user_db_object=user, llm_router=router) + + assert result == ["global-id-1", "global-id-2"] + router.get_model_ids.assert_called_once_with(exclude_team_models=True) + router.get_model_list.assert_not_called() + + +@pytest.mark.asyncio +async def test_populate_team_access_grants_empty_models_user_direct_access(monkeypatch): + """An internal user with models=[] and no teams can call every non-team model, + so the Models+Endpoints page must list them instead of rendering empty.""" + global_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + + router = MagicMock() + router.get_model_ids.return_value = ["global-id-1"] + + user_row = LiteLLM_UserTable( + user_id="u", + user_role=LitellmUserRoles.INTERNAL_USER.value, + models=[], + teams=[], + ) + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + + monkeypatch.setattr(ps, "get_all_team_models", AsyncMock(return_value={})) + + caller = UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.INTERNAL_USER, team_models=[]) + + populated = await ps._populate_team_access_on_models( + user_api_key_dict=caller, + prisma_client=prisma_client, + llm_router=router, + all_models=[global_row], + ) + visible = ps._filter_models_to_user_accessible(populated) + + assert [m["model_info"]["id"] for m in visible] == ["global-id-1"] + assert visible[0]["model_info"]["direct_access"] is True + + @pytest.mark.asyncio async def test_populate_team_access_grants_all_proxy_models_user_direct_access( monkeypatch, From bd2b93c6889614f04986fcd9d94067e061a20567 Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 25 Aug 2026 20:19:47 +0000 Subject: [PATCH 05/60] chore(proxy): drop redundant docstring note in get_direct_access_models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3d3169eeab9..f536f521d57 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12382,8 +12382,6 @@ def get_direct_access_models( The 'all-proxy-models' sentinel grants direct access to every non-team deployment, mirroring how get_key_models expands it for the key/team path. - An empty models list means unrestricted access at call time (see - can_user_call_model), so it resolves the same way. """ if not user_db_object.models or SpecialModelNames.all_proxy_models.value in user_db_object.models: return llm_router.get_model_ids(exclude_team_models=True) From f81408535447db7cf6be81caa48c1447965f2288 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 27 Aug 2026 01:34:33 +0000 Subject: [PATCH 06/60] feat(mcp): bulk-import Anthropic MCP connectors via API and admin UI Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 1147 ++++++++++++++++- .../mcp_connector_import.py | 168 +++ .../mcp_management_endpoints.py | 117 ++ .../test_mcp_connector_import.py | 157 +++ .../test_mcp_management_endpoints.py | 140 ++ .../_components/ImportMCPServers.tsx | 134 ++ .../_components/importConnectorConfig.test.ts | 74 ++ .../_components/importConnectorConfig.ts | 64 + .../mcp-servers/_components/mcp_servers.tsx | 19 +- .../src/components/networking.tsx | 9 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 486 ++++++- 11 files changed, 2506 insertions(+), 9 deletions(-) create mode 100644 litellm/proxy/management_endpoints/mcp_connector_import.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ImportMCPServers.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/importConnectorConfig.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/importConnectorConfig.ts diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 0e3cbbfd560..5e0b5ba3faf 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -14727,6 +14727,17 @@ "title": "Args", "type": "array" }, + "audience": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Audience" + }, "auth_type": { "anyOf": [ { @@ -14738,7 +14749,11 @@ "authorization", "oauth2", "aws_sigv4", - "token" + "token", + "oauth2_token_exchange", + "oauth2_id_jag", + "true_passthrough", + "oauth_delegate" ], "type": "string" }, @@ -14793,6 +14808,17 @@ ], "title": "Command" }, + "connected_app_reachable": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Connected App Reachable" + }, "created_at": { "anyOf": [ { @@ -14826,6 +14852,22 @@ } ] }, + "dcr_bridge": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Dcr Bridge" + }, + "delegate_auth_to_upstream": { + "default": false, + "title": "Delegate Auth To Upstream", + "type": "boolean" + }, "description": { "anyOf": [ { @@ -14844,6 +14886,20 @@ "title": "Env", "type": "object" }, + "env_vars": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/MCPEnvVar" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Env Vars" + }, "extra_headers": { "items": { "type": "string" @@ -14889,6 +14945,17 @@ "title": "Is Byok", "type": "boolean" }, + "issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Issuer" + }, "last_health_check": { "anyOf": [ { @@ -14901,6 +14968,17 @@ ], "title": "Last Health Check" }, + "max_concurrent_requests": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Max Concurrent Requests" + }, "mcp_access_groups": { "items": { "type": "string" @@ -14920,6 +14998,26 @@ ], "title": "Mcp Info" }, + "oauth2_flow": { + "anyOf": [ + { + "enum": [ + "client_credentials", + "authorization_code" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Oauth2 Flow" + }, + "oauth_passthrough": { + "default": false, + "title": "Oauth Passthrough", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -15023,6 +15121,17 @@ "description": "Health status: 'healthy', 'unhealthy', 'unknown'", "title": "Status" }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" + }, "submitted_at": { "anyOf": [ { @@ -15063,6 +15172,39 @@ "title": "Teams", "type": "array" }, + "timeout": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Timeout" + }, + "token_exchange_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Endpoint" + }, + "token_exchange_profile": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Profile" + }, "token_url": { "anyOf": [ { @@ -15153,8 +15295,236 @@ "title": "LiteLLM_MCPServerTable", "type": "object" }, + "MCPConnectorEntry": { + "properties": { + "args": { + "items": { + "type": "string" + }, + "title": "Args", + "type": "array" + }, + "authorization_token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Authorization Token" + }, + "command": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Command" + }, + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "env": { + "additionalProperties": { + "type": "string" + }, + "title": "Env", + "type": "object" + }, + "headers": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Headers" + }, + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Name" + }, + "type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Type" + }, + "url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Url" + } + }, + "title": "MCPConnectorEntry", + "type": "object" + }, + "MCPConnectorImportFailure": { + "properties": { + "error": { + "title": "Error", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "name", + "error" + ], + "title": "MCPConnectorImportFailure", + "type": "object" + }, + "MCPConnectorImportRequest": { + "properties": { + "mcp_servers": { + "anyOf": [ + { + "additionalProperties": { + "$ref": "#/components/schemas/MCPConnectorEntry" + }, + "type": "object" + }, + { + "items": { + "$ref": "#/components/schemas/MCPConnectorEntry" + }, + "type": "array" + } + ], + "title": "Mcp Servers" + } + }, + "required": [ + "mcp_servers" + ], + "title": "MCPConnectorImportRequest", + "type": "object" + }, + "MCPConnectorImportResponse": { + "properties": { + "errors": { + "items": { + "$ref": "#/components/schemas/MCPConnectorImportFailure" + }, + "title": "Errors", + "type": "array" + }, + "imported": { + "items": { + "$ref": "#/components/schemas/MCPConnectorImportResult" + }, + "title": "Imported", + "type": "array" + }, + "skipped": { + "items": { + "$ref": "#/components/schemas/MCPConnectorImportSkipped" + }, + "title": "Skipped", + "type": "array" + } + }, + "required": [ + "imported", + "skipped", + "errors" + ], + "title": "MCPConnectorImportResponse", + "type": "object" + }, + "MCPConnectorImportResult": { + "properties": { + "alias": { + "title": "Alias", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "server_id": { + "title": "Server Id", + "type": "string" + } + }, + "required": [ + "name", + "server_id", + "alias" + ], + "title": "MCPConnectorImportResult", + "type": "object" + }, + "MCPConnectorImportSkipped": { + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "reason": { + "title": "Reason", + "type": "string" + } + }, + "required": [ + "name", + "reason" + ], + "title": "MCPConnectorImportSkipped", + "type": "object" + }, "MCPCredentials": { "properties": { + "audience": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Audience" + }, "auth_value": { "anyOf": [ { @@ -15243,6 +15613,17 @@ ], "title": "Aws Session Token" }, + "client_assertion_signing_alg": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Assertion Signing Alg" + }, "client_id": { "anyOf": [ { @@ -15254,6 +15635,28 @@ ], "title": "Client Id" }, + "client_private_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Private Key" + }, + "client_private_key_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Private Key Id" + }, "client_secret": { "anyOf": [ { @@ -15265,6 +15668,42 @@ ], "title": "Client Secret" }, + "id_jag_resource": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Resource" + }, + "id_jag_resource_token_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Resource Token Endpoint" + }, + "redirect_uris": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Redirect Uris" + }, "scopes": { "anyOf": [ { @@ -15278,11 +15717,113 @@ } ], "title": "Scopes" + }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" + }, + "token_endpoint_auth_method": { + "anyOf": [ + { + "enum": [ + "client_secret_basic", + "client_secret_post" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Endpoint Auth Method" + }, + "token_exchange_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Endpoint" + }, + "token_exchange_profile": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Profile" + }, + "upstream_resource": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Upstream Resource" } }, "title": "MCPCredentials", "type": "object" }, + "MCPEnvVar": { + "description": "One environment variable for an MCP server.\n\nVariables can be interpolated into ``static_headers`` using ``${NAME}``\nsyntax. ``scope=global`` values are stored on the server. ``scope=user``\nvalues are stored per-user in ``LiteLLM_MCPUserEnvVars`` and supplied by\neach user.", + "properties": { + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "name": { + "title": "Name", + "type": "string" + }, + "scope": { + "$ref": "#/components/schemas/MCPEnvVarScope", + "default": "global" + }, + "value": { + "default": "", + "title": "Value", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "MCPEnvVar", + "type": "object" + }, + "MCPEnvVarScope": { + "description": "Scope for an MCP server environment variable.\n\n- ``global``: value is provided by the admin and used for all users.\n- ``user``: each user must provide their own value via the per-user\n env-var endpoint. The admin-supplied ``value`` is treated as a\n placeholder/hint and is not used at request time.", + "enum": [ + "global", + "user" + ], + "title": "MCPEnvVarScope", + "type": "string" + }, "MCPOAuthUserCredentialRequest": { "description": "Stores a user's OAuth2 token for an OpenAPI MCP server.", "properties": { @@ -15537,6 +16078,112 @@ "title": "MCPUserCredentialResponse", "type": "object" }, + "MCPUserEnvVarSpec": { + "description": "Describes one per-user env var slot for the calling user.\n\nStored values are write-only: the status only reports whether a value\n``is_set`` and never echoes the decrypted secret back to the client.", + "properties": { + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "is_set": { + "default": false, + "title": "Is Set", + "type": "boolean" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "MCPUserEnvVarSpec", + "type": "object" + }, + "MCPUserEnvVarsRequest": { + "description": "Payload for storing the calling user's per-user env var values.", + "properties": { + "values": { + "additionalProperties": { + "type": "string" + }, + "title": "Values", + "type": "object" + } + }, + "required": [ + "values" + ], + "title": "MCPUserEnvVarsRequest", + "type": "object" + }, + "MCPUserEnvVarsStatus": { + "description": "Per-user env var status for a single MCP server.", + "properties": { + "alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Alias" + }, + "missing_count": { + "default": 0, + "title": "Missing Count", + "type": "integer" + }, + "required": { + "items": { + "$ref": "#/components/schemas/MCPUserEnvVarSpec" + }, + "title": "Required", + "type": "array" + }, + "server_id": { + "title": "Server Id", + "type": "string" + }, + "server_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Server Name" + }, + "setup_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Setup Url" + } + }, + "required": [ + "server_id" + ], + "title": "MCPUserEnvVarsStatus", + "type": "object" + }, "MakeMCPServersPublicRequest": { "properties": { "mcp_server_ids": { @@ -15604,6 +16251,17 @@ "title": "Args", "type": "array" }, + "audience": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Audience" + }, "auth_type": { "anyOf": [ { @@ -15615,7 +16273,11 @@ "authorization", "oauth2", "aws_sigv4", - "token" + "token", + "oauth2_token_exchange", + "oauth2_id_jag", + "true_passthrough", + "oauth_delegate" ], "type": "string" }, @@ -15680,6 +16342,22 @@ } ] }, + "dcr_bridge": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Dcr Bridge" + }, + "delegate_auth_to_upstream": { + "default": false, + "title": "Delegate Auth To Upstream", + "type": "boolean" + }, "description": { "anyOf": [ { @@ -15698,6 +16376,20 @@ "title": "Env", "type": "object" }, + "env_vars": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/MCPEnvVar" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Env Vars" + }, "extra_headers": { "anyOf": [ { @@ -15728,6 +16420,28 @@ "title": "Is Byok", "type": "boolean" }, + "issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Issuer" + }, + "max_concurrent_requests": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Max Concurrent Requests" + }, "mcp_access_groups": { "items": { "type": "string" @@ -15762,6 +16476,11 @@ ], "title": "Oauth2 Flow" }, + "oauth_passthrough": { + "default": false, + "title": "Oauth Passthrough", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -15831,6 +16550,17 @@ ], "title": "Static Headers" }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" + }, "submitted_at": { "anyOf": [ { @@ -15856,6 +16586,39 @@ "description": "Server-managed: set by the endpoint; caller values are overridden.", "title": "Submitted By" }, + "timeout": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Timeout" + }, + "token_exchange_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Endpoint" + }, + "token_exchange_profile": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Profile" + }, "token_url": { "anyOf": [ { @@ -16008,6 +16771,17 @@ "title": "Args", "type": "array" }, + "audience": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Audience" + }, "auth_type": { "anyOf": [ { @@ -16019,7 +16793,11 @@ "authorization", "oauth2", "aws_sigv4", - "token" + "token", + "oauth2_token_exchange", + "oauth2_id_jag", + "true_passthrough", + "oauth_delegate" ], "type": "string" }, @@ -16084,6 +16862,22 @@ } ] }, + "dcr_bridge": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Dcr Bridge" + }, + "delegate_auth_to_upstream": { + "default": false, + "title": "Delegate Auth To Upstream", + "type": "boolean" + }, "description": { "anyOf": [ { @@ -16102,6 +16896,20 @@ "title": "Env", "type": "object" }, + "env_vars": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/MCPEnvVar" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Env Vars" + }, "extra_headers": { "anyOf": [ { @@ -16132,6 +16940,28 @@ "title": "Is Byok", "type": "boolean" }, + "issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Issuer" + }, + "max_concurrent_requests": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Max Concurrent Requests" + }, "mcp_access_groups": { "items": { "type": "string" @@ -16151,6 +16981,26 @@ ], "title": "Mcp Info" }, + "oauth2_flow": { + "anyOf": [ + { + "enum": [ + "client_credentials", + "authorization_code" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Oauth2 Flow" + }, + "oauth_passthrough": { + "default": false, + "title": "Oauth Passthrough", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -16213,6 +17063,50 @@ ], "title": "Static Headers" }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" + }, + "timeout": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Timeout" + }, + "token_exchange_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Endpoint" + }, + "token_exchange_profile": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Profile" + }, "token_url": { "anyOf": [ { @@ -16331,6 +17225,13 @@ }, "ValidationError": { "properties": { + "ctx": { + "title": "Context", + "type": "object" + }, + "input": { + "title": "Input" + }, "loc": { "items": { "anyOf": [ @@ -16600,6 +17501,18 @@ "description": "Filter MCP servers by team scope. When provided, returns only servers the team has access to plus globally available (allow_all_keys) servers. Used by the Create Key UI to show team-scoped MCP servers.", "title": "Team Id" } + }, + { + "description": "Annotate each returned server with connected_app_reachable: whether a connected app authorized by the calling user (a gateway OAuth session) is served this server on the aggregate MCP endpoint.", + "in": "query", + "name": "connected_app_view", + "required": false, + "schema": { + "default": false, + "description": "Annotate each returned server with connected_app_reachable: whether a connected app authorized by the calling user (a gateway OAuth session) is served this server on the aggregate MCP endpoint.", + "title": "Connected App View", + "type": "boolean" + } } ], "responses": { @@ -16827,6 +17740,53 @@ ] } }, + "/v1/mcp/server/import": { + "post": { + "description": "Bulk-import MCP connectors from Anthropic mcpServers or mcp_servers JSON", + "operationId": "import_mcp_servers_v1_mcp_server_import_post", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPConnectorImportRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPConnectorImportResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Import Mcp Servers", + "tags": [ + "mcp_management" + ] + } + }, "/v1/mcp/server/oauth/session": { "post": { "description": "Temporarily cache an MCP server in memory without writing to the database", @@ -17438,6 +18398,156 @@ ] } }, + "/v1/mcp/server/{server_id}/user-env-vars": { + "delete": { + "description": "Clear the calling user's per-user MCP env var values for this server.", + "operationId": "clear_mcp_user_env_vars_v1_mcp_server__server_id__user_env_vars_delete", + "parameters": [ + { + "in": "path", + "name": "server_id", + "required": true, + "schema": { + "title": "Server Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPUserEnvVarsStatus" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Clear Mcp User Env Vars", + "tags": [ + "mcp_management" + ] + }, + "get": { + "description": "Return the calling user's per-user MCP env var status for this server.", + "operationId": "get_mcp_user_env_vars_v1_mcp_server__server_id__user_env_vars_get", + "parameters": [ + { + "in": "path", + "name": "server_id", + "required": true, + "schema": { + "title": "Server Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPUserEnvVarsStatus" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Mcp User Env Vars", + "tags": [ + "mcp_management" + ] + }, + "post": { + "description": "Store the calling user's per-user MCP env var values for this server. Submitted values are merged over any previously stored values, so you only send the fields you want to set or change; a variable omitted (or sent empty) keeps its stored value. Use DELETE to clear all stored values.", + "operationId": "store_mcp_user_env_vars_v1_mcp_server__server_id__user_env_vars_post", + "parameters": [ + { + "in": "path", + "name": "server_id", + "required": true, + "schema": { + "title": "Server Id", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPUserEnvVarsRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPUserEnvVarsStatus" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Store Mcp User Env Vars", + "tags": [ + "mcp_management" + ] + } + }, "/v1/mcp/tools": { "get": { "description": "Get all MCP tools available for the current key, including those from access groups", @@ -17746,6 +18856,37 @@ "mcp_management" ] } + }, + "/v1/mcp/user-env-vars/status": { + "get": { + "description": "Per-user MCP env var status across every server the user can access. Used by the dashboard to highlight servers with missing per-user vars.", + "operationId": "list_mcp_user_env_var_status_v1_mcp_user_env_vars_status_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/MCPUserEnvVarsStatus" + }, + "title": "Response List Mcp User Env Var Status V1 Mcp User Env Vars Status Get", + "type": "array" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "List Mcp User Env Var Status", + "tags": [ + "mcp_management" + ] + } } } }, diff --git a/litellm/proxy/management_endpoints/mcp_connector_import.py b/litellm/proxy/management_endpoints/mcp_connector_import.py new file mode 100644 index 00000000000..d576131ae1b --- /dev/null +++ b/litellm/proxy/management_endpoints/mcp_connector_import.py @@ -0,0 +1,168 @@ +""" +Convert Anthropic MCP connector definitions into LiteLLM MCP server create requests. + +Two interchange shapes are accepted: +- the ``mcpServers`` mapping used by Claude Desktop / Claude Code config files +- the ``mcp_servers`` array used by the Anthropic Messages API MCP connector +""" + +import re +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, ValidationError + +from litellm.proxy._types import MCPApprovalStatus, NewMCPServerRequest +from litellm.types.mcp import MCPAuth, MCPCredentials, MCPTransport + + +class MCPConnectorEntry(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + name: str | None = None + type: str | None = None + url: str | None = None + authorization_token: str | None = Field( + default=None, validation_alias=AliasChoices("authorization_token", "authorizationToken") + ) + headers: Mapping[str, str] | None = None + command: str | None = None + args: tuple[str, ...] = Field(default_factory=tuple) + env: Mapping[str, str] = Field(default_factory=dict) + description: str | None = None + + +class MCPConnectorImportRequest(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + mcp_servers: Mapping[str, MCPConnectorEntry] | tuple[MCPConnectorEntry, ...] = Field( + validation_alias=AliasChoices("mcp_servers", "mcpServers") + ) + + +@dataclass(frozen=True, slots=True) +class ConvertedConnector: + name: str + request: NewMCPServerRequest + + +@dataclass(frozen=True, slots=True) +class ConnectorConversionError: + name: str + error: str + + +class MCPConnectorImportResult(BaseModel): + name: str + server_id: str + alias: str + + +class MCPConnectorImportSkipped(BaseModel): + name: str + reason: str + + +class MCPConnectorImportFailure(BaseModel): + name: str + error: str + + +class MCPConnectorImportResponse(BaseModel): + imported: tuple[MCPConnectorImportResult, ...] + skipped: tuple[MCPConnectorImportSkipped, ...] + errors: tuple[MCPConnectorImportFailure, ...] + + +_INVALID_SERVER_NAME_CHARS: Final = re.compile(r"[^A-Za-z0-9_]") + + +def sanitize_connector_name(name: str) -> str: + sanitized: Final = re.sub(r"_+", "_", _INVALID_SERVER_NAME_CHARS.sub("_", name.strip())).strip("_") + return sanitized + + +_SSE_TYPES: Final = frozenset({"sse"}) +_URL_TYPES: Final = frozenset({"url", "http", "streamable_http", "streamable-http", "sse", ""}) + + +def _convert_entry(name: str, entry: MCPConnectorEntry) -> ConvertedConnector | ConnectorConversionError: + sanitized_name: Final = sanitize_connector_name(name) + if not sanitized_name: + return ConnectorConversionError(name=name, error="Connector name is empty after sanitization.") + + if entry.url and entry.command: + return ConnectorConversionError(name=name, error="Connector cannot have both a url and a command.") + + if entry.command: + try: + stdio_request: Final = NewMCPServerRequest( + server_name=sanitized_name, + alias=sanitized_name, + description=entry.description, + approval_status=MCPApprovalStatus.active, + transport=MCPTransport.stdio, + command=entry.command, + args=entry.args, + env=entry.env, + ) + except ValidationError as e: + return ConnectorConversionError(name=name, error=_first_validation_message(e)) + return ConvertedConnector(name=name, request=stdio_request) + + if not entry.url: + return ConnectorConversionError(name=name, error="Connector must have either a url or a command.") + + entry_type: Final = (entry.type or "").lower() + if entry_type not in _URL_TYPES: + return ConnectorConversionError(name=name, error=f"Unsupported connector type '{entry.type}'.") + + transport: Final = MCPTransport.sse if entry_type in _SSE_TYPES else MCPTransport.http + credentials: Final = _bearer_credentials(entry.authorization_token) + try: + remote_request: Final = NewMCPServerRequest( + server_name=sanitized_name, + alias=sanitized_name, + description=entry.description, + approval_status=MCPApprovalStatus.active, + transport=transport, + url=entry.url, + auth_type=MCPAuth.bearer_token if entry.authorization_token else MCPAuth.none, + credentials=credentials, + static_headers=entry.headers, + ) + except ValidationError as e: + return ConnectorConversionError(name=name, error=_first_validation_message(e)) + return ConvertedConnector(name=name, request=remote_request) + + +def _bearer_credentials(token: str | None) -> MCPCredentials | None: + if not token: + return None + credentials: Final[MCPCredentials] = {"auth_value": token} + return credentials + + +def _first_validation_message(error: ValidationError) -> str: + messages: Final = tuple(str(detail.get("msg", "")) for detail in error.errors()) + return messages[0] if messages else str(error) + + +def convert_connector_entries( + payload: MCPConnectorImportRequest, +) -> tuple[ConvertedConnector | ConnectorConversionError, ...]: + servers: Final = payload.mcp_servers + if isinstance(servers, Mapping): + return tuple(_convert_entry(name, entry) for name, entry in servers.items()) + return tuple( + _convert_entry(entry.name or "", entry) if entry.name else _named_entry_error(index, entry) + for index, entry in enumerate(servers) + ) + + +def _named_entry_error(index: int, entry: MCPConnectorEntry) -> ConnectorConversionError: + return ConnectorConversionError( + name=entry.url or f"entry {index}", + error="Connector entries in list form must have a name.", + ) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 54a591a5e1a..6583d5405f8 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -133,6 +133,7 @@ if MCP_AVAILABLE: delete_mcp_server, delete_user_credential, delete_user_env_vars, + get_all_mcp_servers, get_all_mcp_servers_for_user, get_draft_mcp_server, get_mcp_server, @@ -199,6 +200,16 @@ if MCP_AVAILABLE: populate_request_with_path_params, ) from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view + from litellm.proxy.management_endpoints.mcp_connector_import import ( + ConnectorConversionError, + ConvertedConnector, + MCPConnectorImportFailure, + MCPConnectorImportRequest, + MCPConnectorImportResponse, + MCPConnectorImportResult, + MCPConnectorImportSkipped, + convert_connector_entries, + ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.types.mcp import ( MCP_ADMIN_CONFIG_CREDENTIAL_KEYS, @@ -1607,6 +1618,112 @@ if MCP_AVAILABLE: return _redact_mcp_credentials(new_mcp_server) + @router.post( + "/server/import", + description="Bulk-import MCP connectors from Anthropic mcpServers or mcp_servers JSON", + dependencies=(Depends(user_api_key_auth),), + response_model=MCPConnectorImportResponse, + status_code=status.HTTP_200_OK, + ) + @management_endpoint_wrapper + async def import_mcp_servers( + payload: MCPConnectorImportRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection + ): + """ + Bulk-import MCP connectors. Accepts the Claude Desktop / Claude Code + ``mcpServers`` mapping or the Anthropic Messages API ``mcp_servers`` + array, creates each entry as a LiteLLM MCP server, and returns + per-entry results so partial imports are visible to the caller. + """ + prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") + + if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict + "error": "User does not have permission to import mcp servers. You can only import mcp servers if you are a PROXY_ADMIN." + }, + ) + + conversions: Final = convert_connector_entries(payload) + existing_servers: Final = await get_all_mcp_servers(prisma_client) + existing_names: Final = frozenset( + name for server in existing_servers for name in (server.alias, server.server_name) if name + ) + + def _classify( + index: int, conversion: ConvertedConnector | ConnectorConversionError + ) -> ConvertedConnector | ConnectorConversionError | MCPConnectorImportSkipped: + if isinstance(conversion, ConnectorConversionError): + return conversion + alias: Final = conversion.request.alias or "" + if alias in existing_names: + return MCPConnectorImportSkipped( + name=conversion.name, reason=f"An MCP server named '{alias}' already exists." + ) + earlier_aliases: Final = frozenset( + earlier.request.alias or "" + for earlier in conversions[:index] + if isinstance(earlier, ConvertedConnector) + ) + if alias in earlier_aliases: + return MCPConnectorImportSkipped( + name=conversion.name, reason=f"Duplicate connector name '{alias}' in the import payload." + ) + return conversion + + async def _create( + conversion: ConvertedConnector, + ) -> MCPConnectorImportResult | MCPConnectorImportFailure: + try: + validate_and_normalize_mcp_server_payload(conversion.request) + except HTTPException as e: + error_text: Final = ( + str(e.detail.get("error", e.detail)) if isinstance(e.detail, dict) else str(e.detail) + ) + return MCPConnectorImportFailure(name=conversion.name, error=error_text) + try: + created: Final = await create_mcp_server( + prisma_client, + conversion.request, + touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, + ) + except Exception as e: # noqa: BLE001 # any create failure must become a per-entry error, not a 500 + verbose_proxy_logger.exception("Error importing mcp server %s: %s", conversion.name, e) + return MCPConnectorImportFailure(name=conversion.name, error=str(e)) + return MCPConnectorImportResult( + name=conversion.name, server_id=created.server_id, alias=created.alias or "" + ) + + classified: Final = tuple(_classify(index, conversion) for index, conversion in enumerate(conversions)) + outcomes: Final = tuple( + [ + await _create(entry) if isinstance(entry, ConvertedConnector) else entry for entry in classified + ] # mutable-ok: await is illegal in a generator expression here + ) + + imported: Final = tuple(entry for entry in outcomes if isinstance(entry, MCPConnectorImportResult)) + if imported: + # Best-effort registry refresh, mirroring add_mcp_server: rows are + # already committed, so a refresh failure must not surface as a 500. + try: + await global_mcp_server_manager.reload_servers_from_database() + except Exception as e: # noqa: BLE001 # rows are committed; a refresh failure must not surface as a 500 + verbose_proxy_logger.exception("MCP connector import committed but registry refresh failed: %s", e) + + return MCPConnectorImportResponse( + imported=imported, + skipped=tuple(entry for entry in outcomes if isinstance(entry, MCPConnectorImportSkipped)), + errors=tuple( + MCPConnectorImportFailure(name=entry.name, error=entry.error) + if isinstance(entry, ConnectorConversionError) + else entry + for entry in outcomes + if isinstance(entry, (ConnectorConversionError, MCPConnectorImportFailure)) + ), + ) + @router.post( "/server/oauth/session", description="Temporarily cache an MCP server in memory without writing to the database", diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py new file mode 100644 index 00000000000..9f2874bcbf3 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py @@ -0,0 +1,157 @@ +import pytest + +from litellm.proxy.management_endpoints.mcp_connector_import import ( + ConnectorConversionError, + ConvertedConnector, + MCPConnectorImportRequest, + convert_connector_entries, + sanitize_connector_name, +) +from litellm.types.mcp import MCPAuth, MCPTransport + + +def _single(payload: dict) -> ConvertedConnector | ConnectorConversionError: + results = convert_connector_entries(MCPConnectorImportRequest.model_validate(payload)) + assert len(results) == 1 + return results[0] + + +class TestSanitizeConnectorName: + @pytest.mark.parametrize( + "raw,expected", + [ + ("my-server", "my_server"), + (" spaced name ", "spaced_name"), + ("already_ok", "already_ok"), + ("a.b.c", "a_b_c"), + ("---", ""), + ], + ) + def test_sanitizes_to_mcp_safe_names(self, raw, expected): + assert sanitize_connector_name(raw) == expected + + +class TestConvertMcpServersMapping: + def test_url_connector_with_authorization_token(self): + result = _single( + { + "mcpServers": { + "github-mcp": { + "url": "https://api.example.com/mcp", + "authorization_token": "secret-token", + "headers": {"X-Env": "prod"}, + "description": "GitHub connector", + } + } + } + ) + assert isinstance(result, ConvertedConnector) + assert result.request.server_name == "github_mcp" + assert result.request.alias == "github_mcp" + assert result.request.transport == MCPTransport.http + assert result.request.url == "https://api.example.com/mcp" + assert result.request.auth_type == MCPAuth.bearer_token + assert result.request.credentials == {"auth_value": "secret-token"} + assert result.request.static_headers == {"X-Env": "prod"} + assert result.request.description == "GitHub connector" + + def test_camel_case_authorization_token_alias(self): + result = _single( + {"mcpServers": {"srv": {"url": "https://x.example/mcp", "authorizationToken": "tok"}}} + ) + assert isinstance(result, ConvertedConnector) + assert result.request.credentials == {"auth_value": "tok"} + + def test_url_connector_without_token_uses_no_auth(self): + result = _single({"mcpServers": {"open": {"url": "https://open.example/mcp"}}}) + assert isinstance(result, ConvertedConnector) + assert result.request.auth_type == MCPAuth.none + assert result.request.credentials is None + + def test_sse_type_maps_to_sse_transport(self): + result = _single({"mcpServers": {"legacy": {"type": "sse", "url": "https://sse.example/mcp"}}}) + assert isinstance(result, ConvertedConnector) + assert result.request.transport == MCPTransport.sse + + def test_stdio_connector(self): + result = _single( + { + "mcpServers": { + "local": { + "command": "npx", + "args": ["-y", "@example/mcp-server"], + "env": {"API_KEY": "value"}, + } + } + } + ) + assert isinstance(result, ConvertedConnector) + assert result.request.transport == MCPTransport.stdio + assert result.request.command == "npx" + assert result.request.args == ["-y", "@example/mcp-server"] + assert result.request.env == {"API_KEY": "value"} + + def test_disallowed_stdio_command_returns_error(self): + result = _single({"mcpServers": {"evil": {"command": "rm", "args": ["-rf", "/"]}}}) + assert isinstance(result, ConnectorConversionError) + assert "not in the allowed commands list" in result.error + + def test_unsupported_type_returns_error(self): + result = _single({"mcpServers": {"ws": {"type": "websocket", "url": "wss://x.example"}}}) + assert isinstance(result, ConnectorConversionError) + assert "Unsupported connector type" in result.error + + def test_missing_url_and_command_returns_error(self): + result = _single({"mcpServers": {"empty": {}}}) + assert isinstance(result, ConnectorConversionError) + assert "either a url or a command" in result.error + + def test_url_and_command_together_returns_error(self): + result = _single({"mcpServers": {"both": {"url": "https://x.example/mcp", "command": "npx"}}}) + assert isinstance(result, ConnectorConversionError) + assert "both a url and a command" in result.error + + def test_name_empty_after_sanitization_returns_error(self): + result = _single({"mcpServers": {"---": {"url": "https://x.example/mcp"}}}) + assert isinstance(result, ConnectorConversionError) + assert "empty after sanitization" in result.error + + +class TestConvertMcpServersList: + def test_anthropic_messages_api_list_shape(self): + result = _single( + { + "mcp_servers": [ + { + "type": "url", + "url": "https://mcp.example.com/sse", + "name": "deepwiki", + "authorization_token": "tok", + } + ] + } + ) + assert isinstance(result, ConvertedConnector) + assert result.request.server_name == "deepwiki" + assert result.request.transport == MCPTransport.http + assert result.request.credentials == {"auth_value": "tok"} + + def test_list_entry_without_name_returns_error(self): + result = _single({"mcp_servers": [{"type": "url", "url": "https://x.example/mcp"}]}) + assert isinstance(result, ConnectorConversionError) + assert "must have a name" in result.error + + def test_partial_conversion_preserves_per_entry_results(self): + results = convert_connector_entries( + MCPConnectorImportRequest.model_validate( + { + "mcpServers": { + "good": {"url": "https://good.example/mcp"}, + "bad": {"type": "websocket", "url": "wss://bad.example"}, + } + } + ) + ) + assert len(results) == 2 + assert isinstance(results[0], ConvertedConnector) + assert isinstance(results[1], ConnectorConversionError) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 0d639e1cb6a..7b3580f166a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -6766,3 +6766,143 @@ class TestConnectedAppViewAnnotation: assert all(server.connected_app_reachable is None for server in result) reload_mock.assert_not_awaited() + + +class TestImportMCPServers: + """Bulk connector import must be admin-only and report per-entry outcomes.""" + + @staticmethod + def _import_patches(existing_servers, create_mock, mock_manager): + return ( + patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers", + AsyncMock(return_value=existing_servers), + ), + patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", + create_mock, + ), + patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ) + + @pytest.mark.asyncio + async def test_non_admin_is_rejected(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + MCPConnectorImportRequest, + import_mcp_servers, + ) + + payload = MCPConnectorImportRequest.model_validate( + {"mcpServers": {"srv": {"url": "https://x.example/mcp"}}} + ) + caller = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.INTERNAL_USER) + + with patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ): + with pytest.raises(HTTPException) as exc_info: + await import_mcp_servers(payload=payload, user_api_key_dict=caller) + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_import_reports_imported_skipped_and_errors(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + MCPConnectorImportRequest, + import_mcp_servers, + ) + + payload = MCPConnectorImportRequest.model_validate( + { + "mcpServers": { + "new-server": {"url": "https://new.example/mcp", "authorization_token": "tok"}, + "existing": {"url": "https://existing.example/mcp"}, + "broken": {"type": "websocket", "url": "wss://x.example"}, + } + } + ) + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") + existing = generate_mock_mcp_server_db_record(server_id="existing-1", alias="existing") + created = generate_mock_mcp_server_db_record(server_id="created-1", alias="new_server") + create_mock = AsyncMock(return_value=created) + mock_manager = MagicMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ExitStack() as stack: + for p in self._import_patches([existing], create_mock, mock_manager): + stack.enter_context(p) + result = await import_mcp_servers(payload=payload, user_api_key_dict=admin) + + assert [entry.name for entry in result.imported] == ["new-server"] + assert result.imported[0].server_id == "created-1" + assert [entry.name for entry in result.skipped] == ["existing"] + assert "already exists" in result.skipped[0].reason + assert [entry.name for entry in result.errors] == ["broken"] + create_mock.assert_awaited_once() + sent_request = create_mock.await_args[0][1] + assert sent_request.credentials == {"auth_value": "tok"} + mock_manager.reload_servers_from_database.assert_awaited_once() + + @pytest.mark.asyncio + async def test_duplicate_names_within_payload_are_skipped(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + MCPConnectorImportRequest, + import_mcp_servers, + ) + + payload = MCPConnectorImportRequest.model_validate( + { + "mcp_servers": [ + {"type": "url", "url": "https://a.example/mcp", "name": "dup srv"}, + {"type": "url", "url": "https://b.example/mcp", "name": "dup-srv"}, + ] + } + ) + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") + created = generate_mock_mcp_server_db_record(server_id="created-1", alias="dup_srv") + create_mock = AsyncMock(return_value=created) + mock_manager = MagicMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ExitStack() as stack: + for p in self._import_patches([], create_mock, mock_manager): + stack.enter_context(p) + result = await import_mcp_servers(payload=payload, user_api_key_dict=admin) + + assert len(result.imported) == 1 + assert len(result.skipped) == 1 + assert "Duplicate connector name" in result.skipped[0].reason + create_mock.assert_awaited_once() + + @pytest.mark.asyncio + async def test_no_imports_skips_registry_refresh(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + MCPConnectorImportRequest, + import_mcp_servers, + ) + + payload = MCPConnectorImportRequest.model_validate( + {"mcpServers": {"existing": {"url": "https://existing.example/mcp"}}} + ) + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") + existing = generate_mock_mcp_server_db_record(server_id="existing-1", alias="existing") + create_mock = AsyncMock() + mock_manager = MagicMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ExitStack() as stack: + for p in self._import_patches([existing], create_mock, mock_manager): + stack.enter_context(p) + result = await import_mcp_servers(payload=payload, user_api_key_dict=admin) + + assert result.imported == () + create_mock.assert_not_awaited() + mock_manager.reload_servers_from_database.assert_not_awaited() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ImportMCPServers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ImportMCPServers.tsx new file mode 100644 index 00000000000..d04160bca3e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ImportMCPServers.tsx @@ -0,0 +1,134 @@ +import React, { useState } from "react"; +import { Button } from "@/components/ui/button"; +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { Textarea } from "@/components/ui/textarea"; +import { Alert, AlertTitle } from "@/components/shared/Alert"; +import { importMCPServers } from "@/components/networking"; +import { toast } from "@/lib/toast"; +import { MCPConnectorImportResponse, parseConnectorConfig } from "./importConnectorConfig"; + +interface ImportMCPServersProps { + accessToken: string; + open: boolean; + onClose: () => void; + onImported: () => void; +} + +const PLACEHOLDER = `{ + "mcpServers": { + "my_server": { + "url": "https://example.com/mcp", + "authorization_token": "..." + } + } +}`; + +const ImportMCPServers: React.FC = ({ accessToken, open, onClose, onImported }) => { + const [configText, setConfigText] = useState(""); + const [parseError, setParseError] = useState(null); + const [isImporting, setIsImporting] = useState(false); + const [result, setResult] = useState(null); + + const handleClose = () => { + setConfigText(""); + setParseError(null); + setResult(null); + onClose(); + }; + + const handleImport = async () => { + const parsed = parseConnectorConfig(configText); + if (!parsed.ok) { + setParseError(parsed.error); + return; + } + setParseError(null); + setIsImporting(true); + try { + const response = (await importMCPServers(accessToken, parsed.payload)) as MCPConnectorImportResponse; + setResult(response); + if (response.imported.length > 0) { + toast.success(`Imported ${response.imported.length} MCP server${response.imported.length === 1 ? "" : "s"}`); + onImported(); + } + } catch (error) { + console.error("Failed to import MCP servers:", error); + setParseError("Import request failed. Check the proxy logs for details."); + } finally { + setIsImporting(false); + } + }; + + return ( + !isOpen && handleClose()}> + + + Import MCP Connectors + +
+

+ Paste an Anthropic connector configuration: the mcpServers mapping from a Claude Desktop / + Claude Code config file, or the mcp_servers array from the Anthropic Messages API. +

+