scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-08 08:08:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 198

feat(accounting)!: make the SpendLogs ledger the default accounting path (#49872)

Browse source 
* fix(accounting): keep one undecodable delivery from blocking the ledger

A request whose metadata carried a 401-digit tokens_saved stored that number in
its delivery payload. The query engine could not decode it when the worker
claimed the batch, so every later delivery on the proxy stayed pending

The worker now claims the payload as text and decodes each row in Python. A row
that still cannot be decoded is parked alone while the rest of the batch
applies. The payload builder also clamps integers to the JSON-safe range, so
caller-supplied counts can no longer reach the engine at that size

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): take spend locks in one order so the ledger cannot deadlock

A ledger delivery updates many spend tables in one transaction, so it can
deadlock with any writer that locks the same rows in another order. It took the
team advisory lock after the user and team rows, while team roster writers take
it first

The delivery now takes every team advisory lock before any row, then writes
tables in one canonical order with ids sorted. It also runs under the rollup
lock_timeout, applies 500 deliveries per batch by default, and spreads retries
with jitter

The writers that could still collide with it now follow the same order. The
budget reset cascade, model allowlist rename, global spend reset, bulk user
update, org delete and the delete-by-list paths lock their rows in key order
first, and access group endpoints update teams and keys in one sorted pass

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(accounting): prove a ledger delivery waits instead of deadlocking

Real Postgres tests hold the first lock of a team roster change, an access group
update and a budget reset cascade, start a delivery, and let the other writer
finish. The delivery must wait and then charge once. Before the lock order fix
each case ends in a Postgres deadlock

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(accounting)!: make the SpendLogs ledger the default accounting path

use_spend_logs_ledger now defaults to true, so every request is charged through
LiteLLM_SpendLogDelivery. Setting it to false rolls back to the in-memory and
Redis spend queues, which stay in place. Tests that cover those legacy writers
pin the setting off, and the rest check the queued charge the ledger stores

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(proxy): let the mocked client open a transaction for key and team deletes

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(proxy): take team locks first when a model rename rewrites allowlists

A rename locked team rows before user rows while a roster change holds the team
advisory lock and a user row, so the two could deadlock. The rename now takes the
advisory lock for every team that lists the model, sorted, before any row lock.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): drain a backlog in one worker run, one batch per transaction

One run applied a single fresh batch, so at 500 per batch charges fell behind any
load above that rate. The worker now keeps claiming fresh batches while they come
back full, for up to five seconds, and each batch still commits on its own.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): queue only the charge when a request's spend log is suppressed

With logging off for a request, the ledger still queued the full row, so a failed
flush or a shutdown could leave its prompt and response in Redis. The queue now
holds only the fields the charge needs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* style(proxy): format the rename team listing query

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): remove expired delivery receipts as fast as the worker applies them

The worker now applies several batches per run but still removed one cleanup batch
of applied receipts, so the table could outgrow its retention window. Cleanup now
repeats while each pass removes a full batch, for up to five seconds.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(proxy): retry a model rename when a team starts listing the model mid-lock

The team advisory locks come from a snapshot, so a team that gained the old name
before its row lock was locked without its advisory lock. That attempt now commits
without writing and the next one locks the fresh set of teams.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
joshua-berri 2026-11-05 09:14:10 -08:00 • committed by GitHub
parent dfabf32cd5
commit a12a6cd154
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
32 changed files with 859 additions and 309 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/accounting/application.py
View file

@ -7,6 +7,7 @@ from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_b
from litellm.proxy.db.db_span import db_span
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction
from litellm.proxy.db.model_usage_rollup import queue_model_usage_upsert
from litellm.proxy.db.rollup_lock_timeout import apply_rollup_lock_timeout
from litellm.repositories.accounting_delivery_repository import DeliveryProjection, LedgerTransaction
if TYPE_CHECKING:
@ -25,6 +26,7 @@ async def apply_deliveries(
from litellm.proxy.guardrails.usage_tracking import UsageUnitIncrement, UsageUnitKey, upsert_usage_unit_row
contributions: Final = project_deliveries(deliveries)
await apply_rollup_lock_timeout(transaction)
await writer.apply_accounting_entities(
prisma_client, proxy_logging_obj, transaction, entity_transactions(contributions)
)

2
litellm/accounting/snapshot.py
View file

@ -14,6 +14,7 @@ if TYPE_CHECKING:
from litellm.router import Router
_FIELDS: Final = TypeAdapter(dict[str, object])
_JSON_SAFE_INTEGER: Final = 2**53 - 1
def _mapping(value: object) -> Mapping[str, object] | None:
@ -86,6 +87,7 @@ def build_accounting_payload(
"metadata": accounting_metadata,
},
non_finite_as_null=True,
integer_limit=_JSON_SAFE_INTEGER,
)
)

35
litellm/accounting/worker.py
View file

@ -14,6 +14,7 @@ from litellm.accounting.metrics import record_deliveries
from litellm.accounting.projection import Contribution
from litellm.repositories.accounting_delivery_repository import (
AccountingDeliveryRepository,
DeliveryBatchResult,
DeliveryProjection,
LedgerDatabase,
LedgerTransaction,
@ -34,6 +35,14 @@ class _AccountingAlerts(Protocol):
_JOB: Final = "accounting_delivery"
def _filled(result: DeliveryBatchResult, batch_size: int) -> bool:
return result.applied + len(result.parked) >= batch_size
def _progressed(result: DeliveryBatchResult) -> bool:
return bool(result.applied or result.failed)
async def flush_accounting_deliveries(
writer: "DBSpendUpdateWriter", prisma_client: "PrismaClient", proxy_logging_obj: "ProxyLogging"
) -> None:
@ -54,9 +63,9 @@ async def flush_accounting_deliveries(
)
repository: Final = AccountingDeliveryRepository(database)
started: Final = time.monotonic()
batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "10000"))
batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "500"))
async def run_batch(*, retry_only: bool) -> bool:
async def run_batch(*, retry_only: bool) -> DeliveryBatchResult:
result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only)
if result.applied:
await writer.invalidate_accounting_entities(applied, proxy_logging_obj)
@ -68,17 +77,25 @@ async def flush_accounting_deliveries(
proxy_logging_obj,
message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.",
)
return bool(result.applied or result.failed)
return result
await run_batch(retry_only=False)
fresh_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if not _filled(await run_batch(retry_only=False), batch_size) or time.monotonic() >= fresh_deadline:
break
retry_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if time.monotonic() >= retry_deadline or not await run_batch(retry_only=True):
if time.monotonic() >= retry_deadline or not _progressed(await run_batch(retry_only=True)):
break
retention_days: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1"))
cleanup_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if (
await repository.cleanup_applied(retention_days=retention_days, batch_size=batch_size * 2)
< batch_size * 2
or time.monotonic() >= cleanup_deadline
):
break
await repository.cleanup_applied(
retention_days=int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")),
batch_size=batch_size * 2,
)
metrics: Final = await repository.metrics()
record_deliveries(metrics, time.monotonic() - started)
if metrics.batch_count >= int(

11
litellm/litellm_core_utils/safe_json_dumps.py
View file

@ -19,6 +19,12 @@ def _finite_or_none(value: float | None) -> float | None:
return None if isinstance(value, float) and not math.isfinite(value) else value
def _bounded(value: float | None, integer_limit: int | None) -> float | None:
if integer_limit is None or isinstance(value, bool) or not isinstance(value, int):
return value
return max(-integer_limit, min(integer_limit, value))
def safe_json_structure(
data: object,
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH,
@ -26,12 +32,14 @@ def safe_json_structure(
key: str | None = None,
*,
non_finite_as_null: bool = False,
integer_limit: int | None = None,
) -> object:
"""
Rebuild data out of JSON-native pieces while detecting circular references.
If a circular reference is detected then a marker string is returned.
NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors. With
non_finite_as_null, NaN and infinite floats become null, since jsonb rejects them.
With integer_limit, integers are clamped to that magnitude.
value_transform, when given, is applied to every string leaf (and to the
str() fallback for non-serializable objects) with the mapping key the leaf
@ -51,7 +59,8 @@ def safe_json_structure(
cleaned = obj.replace("\x00", "") if "\x00" in obj else obj
return _transform(key, cleaned)
if isinstance(obj, (int, float, bool, type(None))):
return _finite_or_none(obj) if non_finite_as_null else obj
bounded: Final = _bounded(obj, integer_limit)
return _finite_or_none(bounded) if non_finite_as_null else bounded
# Check for circular reference.
if id(obj) in seen:
return "CircularReference Detected"

4
litellm/proxy/_types.py
View file

@ -2507,8 +2507,8 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
"""
use_spend_logs_ledger: bool = Field(
False,
description="Apply request accounting from stored SpendLogs deliveries. Enable only after all accounting writers are upgraded and ledger activation gates pass.",
True,
description="Apply request accounting from stored SpendLogs deliveries. Set to false to roll back to the previous in-memory and Redis spend queues.",
)
request_tag_mappings: Sequence[RequestTagMapping] | None = Field(

5
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -909,13 +909,14 @@ class ResetBudgetJob:
db_span("reset_budget_cascade", "LiteLLM_BudgetTable"),
budget_cascade_unit_of_work(self._new_batch) as uow,
):
uow.row_locks.queue_row_locks(cascade.budget_ids)
_queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE)
uow.team_memberships.queue_total_spend_repair(cascade.budget_ids)
_queue_budget_linked_resets(uow.team_memberships, cascade)
_queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE)
_queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE)
for budget_id, budget_reset_at in cascade.budget_resets:
uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at)

266
litellm/proxy/db/db_spend_update_writer.py
View file

@ -328,13 +328,13 @@ _SpendTableName = Literal[
"agent_list_transactions",
]
_SPEND_TABLE_COMMIT_ORDER: Final[tuple[_SpendTableName, ...]] = (
"user_list_transactions",
"end_user_list_transactions",
"key_list_transactions",
"team_list_transactions",
"team_member_list_transactions",
"org_list_transactions",
"org_member_list_transactions",
"team_list_transactions",
"user_list_transactions",
"end_user_list_transactions",
"team_member_list_transactions",
"key_list_transactions",
"project_list_transactions",
"tag_list_transactions",
"model_access_group_list_transactions",
@ -563,6 +563,24 @@ SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend,
"""
def _org_members_in_key_order(spend_by_member_key: Mapping[str, float]) -> tuple[tuple[tuple[str, str], float], ...]:
"""Sorted by the decoded (organization_id, user_id), the order every other writer locks these rows in, since
the percent-encoded key does not sort the same way."""
decoded: Final = (
((unquote(key.split("::")[1]), unquote(key.split("::")[3])), cost) for key, cost in spend_by_member_key.items()
)
return tuple(sorted(decoded))
def _charged_team_ids(updates: DBSpendUpdateTransactions) -> tuple[str, ...]:
members: Final = updates.get("team_member_list_transactions") or {}
return tuple(
sorted(
{*(updates.get("team_list_transactions") or {}), *(key.split("::")[1] for key in members)},
)
)
async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_member_key: Mapping[str, float]) -> None:
# key is "team_id::<value>::user_id::<value>"; locks are taken in sorted team_id order like the team endpoints
rows: Final = sorted((key.split("::")[1], key.split("::")[3], cost) for key, cost in spend_by_member_key.items())
@ -712,7 +730,7 @@ class DBSpendUpdateWriter:
proxy_settings: Final = cast( # cast-ok: proxy settings use string keys
Mapping[str, object], general_settings
)
use_ledger: Final = proxy_settings.get("use_spend_logs_ledger") is True
use_ledger: Final = proxy_settings.get("use_spend_logs_ledger", True) is not False
request_fields: Final = _SPEND_METADATA_ADAPTER.validate_python(kwargs or {})
if kwargs is not None:
kwargs[ACCOUNTING_LEDGER_KWARG] = use_ledger # rebind-ok: spend counters leave ledger windows to the worker
@ -863,8 +881,12 @@ class DBSpendUpdateWriter:
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
)
if "accounting_payload" in payload:
from litellm.proxy.utils import accounting_only_spend_log
await self._insert_spend_log_to_db(
payload={**payload, "accounting_log_disabled": disable_spend_logs},
payload=accounting_only_spend_log(payload)
if disable_spend_logs
else {**payload, "accounting_log_disabled": False},
prisma_client=prisma_client,
llm_router=llm_router,
)
@ -2294,6 +2316,10 @@ class DBSpendUpdateWriter:
transaction: _SpendTransaction,
updates: DBSpendUpdateTransactions,
) -> None:
"""Every team advisory lock is taken before the first row lock, the order the team roster writers use, so
applying a batch of charges cannot deadlock with a concurrent roster change."""
for team_id in _charged_team_ids(updates):
_ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id)
await self._commit_spend_updates_to_db(
prisma_client, 0, proxy_logging_obj, updates, outer_transaction=transaction
)
@ -2313,6 +2339,104 @@ class DBSpendUpdateWriter:
"""
from litellm.proxy.utils import ProxyUpdateSpend
### UPDATE ORG TABLE ###
org_list_transactions: Final = db_spend_update_transactions["org_list_transactions"]
verbose_proxy_logger.debug("Org Spend transactions: %s", org_list_transactions)
if org_list_transactions is not None and len(org_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(
prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by org_id for consistent lock ordering across pods to prevent deadlocks.
for org_id, response_cost in sorted(org_list_transactions.items()):
batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists
where={"organization_id": org_id},
data={"spend": {"increment": response_cost}},
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("org_list_transactions")
org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions")
verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions)
if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with (
_spend_update_tx(
prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction
) as transaction,
transaction.batch_() as batcher,
):
for (org_id, member_id), response_cost in _org_members_in_key_order(
org_member_list_transactions
):
batcher.litellm_organizationmembership.update_many(
where={"organization_id": org_id, "user_id": member_id},
data={"spend": {"increment": response_cost}},
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("org_member_list_transactions")
### UPDATE TEAM TABLE ###
team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"]
verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions)
if team_list_transactions is not None and len(team_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(
prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by team_id for consistent lock ordering across pods to prevent deadlocks.
for team_id, response_cost in sorted(team_list_transactions.items()):
verbose_proxy_logger.debug(
"Updating spend for team id=%s by %s", team_id, response_cost
)
batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists
where={"team_id": team_id},
data={"spend": {"increment": response_cost}},
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("team_list_transactions")
### UPDATE USER TABLE ###
user_list_transactions: Final = db_spend_update_transactions["user_list_transactions"]
verbose_proxy_logger.debug("User Spend transactions: %s", user_list_transactions)
@ -2359,75 +2483,6 @@ class DBSpendUpdateWriter:
)
if on_table_committed is not None:
on_table_committed("end_user_list_transactions")
### UPDATE KEY TABLE ###
key_list_transactions: Final = db_spend_update_transactions["key_list_transactions"]
verbose_proxy_logger.debug("KEY Spend transactions: %s", key_list_transactions)
if key_list_transactions is not None and len(key_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(
prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by token for consistent lock ordering across pods to prevent deadlocks.
for token, response_cost in sorted(key_list_transactions.items()):
spend_increment: _SpendIncrement = {"increment": response_cost}
batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists
where={"token": token},
data={
"spend": spend_increment,
"total_spend": spend_increment,
"last_active": datetime.now(timezone.utc),
},
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("key_list_transactions")
### UPDATE TEAM TABLE ###
team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"]
verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions)
if team_list_transactions is not None and len(team_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(
prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by team_id for consistent lock ordering across pods to prevent deadlocks.
for team_id, response_cost in sorted(team_list_transactions.items()):
verbose_proxy_logger.debug(
"Updating spend for team id=%s by %s", team_id, response_cost
)
batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists
where={"team_id": team_id},
data={"spend": {"increment": response_cost}},
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("team_list_transactions")
### UPDATE TEAM Membership TABLE with spend ###
team_member_list_transactions: Final = db_spend_update_transactions["team_member_list_transactions"]
@ -2478,22 +2533,27 @@ class DBSpendUpdateWriter:
elif on_table_committed is not None:
on_table_committed("team_member_list_transactions")
### UPDATE ORG TABLE ###
org_list_transactions: Final = db_spend_update_transactions["org_list_transactions"]
verbose_proxy_logger.debug("Org Spend transactions: %s", org_list_transactions)
if org_list_transactions is not None and len(org_list_transactions.keys()) > 0:
### UPDATE KEY TABLE ###
key_list_transactions: Final = db_spend_update_transactions["key_list_transactions"]
verbose_proxy_logger.debug("KEY Spend transactions: %s", key_list_transactions)
if key_list_transactions is not None and len(key_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(
prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction
prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by org_id for consistent lock ordering across pods to prevent deadlocks.
for org_id, response_cost in sorted(org_list_transactions.items()):
batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists
where={"organization_id": org_id},
data={"spend": {"increment": response_cost}},
# Sort by token for consistent lock ordering across pods to prevent deadlocks.
for token, response_cost in sorted(key_list_transactions.items()):
spend_increment: _SpendIncrement = {"increment": response_cost}
batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists
where={"token": token},
data={
"spend": spend_increment,
"total_spend": spend_increment,
"last_active": datetime.now(timezone.utc),
},
)
break
except Exception as e:
@ -2507,39 +2567,7 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("org_list_transactions")
org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions")
verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions)
if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with (
_spend_update_tx(
prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction
) as transaction,
transaction.batch_() as batcher,
):
for key, response_cost in sorted(org_member_list_transactions.items()):
_, quoted_org_id, _, quoted_user_id = key.split("::")
batcher.litellm_organizationmembership.update_many(
where={"organization_id": unquote(quoted_org_id), "user_id": unquote(quoted_user_id)},
data={"spend": {"increment": response_cost}},
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("org_member_list_transactions")
on_table_committed("key_list_transactions")
### UPDATE PROJECT TABLE ###
project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions")

91
litellm/proxy/management_endpoints/access_group_endpoints.py
View file

@ -347,47 +347,45 @@ async def _cache_access_group_record(record: _AccessGroupRecord) -> None:
# ---------------------------------------------------------------------------
async def _sync_add_access_group_to_teams(tx: _AccessGroupTx, team_ids: list[str], access_group_id: str) -> None:
"""Add access_group_id to each team's access_group_ids (idempotent)."""
for team_id in team_ids:
async def _sync_access_group_on_teams(
tx: _AccessGroupTx, add: Sequence[str], remove: Sequence[str], access_group_id: str
) -> None:
"""Rows are locked in sorted team_id order, the order every spend writer locks team rows in, so this
cannot deadlock with them."""
to_add: Final = frozenset(add)
for team_id in sorted(to_add | frozenset(remove)):
team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id})
if team is not None and access_group_id not in (team.access_group_ids or []):
if team is None:
continue
current = list(team.access_group_ids or [])
if team_id in to_add and access_group_id not in current:
await tx.litellm_teamtable.update(
where={"team_id": team_id}, data={"access_group_ids": [*current, access_group_id]}
)
elif team_id not in to_add and access_group_id in current:
await tx.litellm_teamtable.update(
where={"team_id": team_id},
data={"access_group_ids": list(team.access_group_ids or []) + [access_group_id]},
data={"access_group_ids": list(_without_access_group(current, access_group_id))},
)
async def _sync_remove_access_group_from_teams(tx: _AccessGroupTx, team_ids: list[str], access_group_id: str) -> None:
"""Remove access_group_id from each team's access_group_ids (idempotent)."""
for team_id in team_ids:
team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id})
if team is not None and access_group_id in (team.access_group_ids or []):
await tx.litellm_teamtable.update(
where={"team_id": team_id},
data={"access_group_ids": [ag for ag in (team.access_group_ids or ()) if ag != access_group_id]},
)
async def _sync_add_access_group_to_keys(tx: _AccessGroupTx, key_tokens: list[str], access_group_id: str) -> None:
"""Add access_group_id to each key's access_group_ids (idempotent)."""
for token in key_tokens:
async def _sync_access_group_on_keys(
tx: _AccessGroupTx, add: Sequence[str], remove: Sequence[str], access_group_id: str
) -> None:
"""Rows are locked in sorted token order, the order every spend writer locks key rows in."""
to_add: Final = frozenset(add)
for token in sorted(to_add | frozenset(remove)):
key = await tx.litellm_verificationtoken.find_unique(where={"token": token})
if key is not None and access_group_id not in (key.access_group_ids or []):
if key is None:
continue
current = list(key.access_group_ids or [])
if token in to_add and access_group_id not in current:
await tx.litellm_verificationtoken.update(
where={"token": token},
data={"access_group_ids": list(key.access_group_ids or []) + [access_group_id]},
where={"token": token}, data={"access_group_ids": [*current, access_group_id]}
)
async def _sync_remove_access_group_from_keys(tx: _AccessGroupTx, key_tokens: list[str], access_group_id: str) -> None:
"""Remove access_group_id from each key's access_group_ids (idempotent)."""
for token in key_tokens:
key = await tx.litellm_verificationtoken.find_unique(where={"token": token})
if key is not None and access_group_id in (key.access_group_ids or []):
elif token not in to_add and access_group_id in current:
await tx.litellm_verificationtoken.update(
where={"token": token},
data={"access_group_ids": [ag for ag in (key.access_group_ids or ()) if ag != access_group_id]},
where={"token": token}, data={"access_group_ids": list(_without_access_group(current, access_group_id))}
)
@ -582,8 +580,8 @@ async def create_access_group(
)
# Sync team and key tables to reference the new access group
await _sync_add_access_group_to_teams(tx, data.assigned_team_ids or [], record.access_group_id)
await _sync_add_access_group_to_keys(tx, data.assigned_key_ids or [], record.access_group_id)
await _sync_access_group_on_teams(tx, data.assigned_team_ids or [], (), record.access_group_id)
await _sync_access_group_on_keys(tx, data.assigned_key_ids or [], (), record.access_group_id)
except HTTPException:
raise
except Exception as e:
@ -718,10 +716,8 @@ async def update_access_group(
data=update_data,
)
await _sync_add_access_group_to_teams(tx, teams_to_add, access_group_id)
await _sync_remove_access_group_from_teams(tx, teams_to_remove, access_group_id)
await _sync_add_access_group_to_keys(tx, keys_to_add, access_group_id)
await _sync_remove_access_group_from_keys(tx, keys_to_remove, access_group_id)
await _sync_access_group_on_teams(tx, teams_to_add, teams_to_remove, access_group_id)
await _sync_access_group_on_keys(tx, keys_to_add, keys_to_remove, access_group_id)
except HTTPException:
raise
except Exception as e:
@ -790,25 +786,8 @@ async def delete_access_group(
)
affected_key_tokens = list(all_affected_key_tokens)
# Update teams returned by find_many directly — we already have their data.
for team in teams_with_group:
await tx.litellm_teamtable.update(
where={"team_id": team.team_id},
data={"access_group_ids": [ag for ag in (team.access_group_ids or ()) if ag != access_group_id]},
)
# Use _sync_remove only for out-of-sync teams not found by the hasSome query.
out_of_sync_team_ids: Final = set(existing.assigned_team_ids or []) - {t.team_id for t in teams_with_group}
await _sync_remove_access_group_from_teams(tx, list(out_of_sync_team_ids), access_group_id)
# Update keys returned by find_many directly — we already have their data.
for key in keys_with_group:
await tx.litellm_verificationtoken.update(
where={"token": key.token},
data={"access_group_ids": [ag for ag in (key.access_group_ids or ()) if ag != access_group_id]},
)
# Use _sync_remove only for out-of-sync keys not found by the hasSome query.
out_of_sync_key_tokens: Final = set(existing.assigned_key_ids or []) - {k.token for k in keys_with_group}
await _sync_remove_access_group_from_keys(tx, list(out_of_sync_key_tokens), access_group_id)
await _sync_access_group_on_teams(tx, (), affected_team_ids, access_group_id)
await _sync_access_group_on_keys(tx, (), affected_key_tokens, access_group_id)
detached_agent_ids: Final = await _detach_access_group_from_agents(tx, access_group_id)

24
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -83,6 +83,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa:
handle_update_object_permission_common,
set_object_permission,
)
from litellm.proxy.management_helpers.team_roster_sync import user_tx_db
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
from litellm.proxy.utils import handle_exception_on_proxy, hash_password
from litellm.repositories.organization_repository import OrganizationRepository
@ -125,6 +126,11 @@ _USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig])
_USER_CACHE_INVALIDATION_BATCH_SIZE: Final = 50
_LOCK_ALL_USERS_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_UserTable" WHERE TRUE ORDER BY user_id COLLATE "C" FOR UPDATE'
)
@with_config(ConfigDict(extra="allow", strict=True))
class _AdminTeamAndKeyRows(TypedDict):
keys: ReadOnly[Sequence[object] | None]
@ -2011,14 +2017,16 @@ async def bulk_user_update(
try:
# Perform bulk database update
await UserRepository(prisma_client).table.update_many(
where={},
data=(
{**non_default_values, "model_max_budget": json.dumps(non_default_values["model_max_budget"])}
if "model_max_budget" in non_default_values
else non_default_values
),
)
async with prisma_client.tx() as transaction:
_ = await transaction.query_raw(_LOCK_ALL_USERS_SQL)
_ = await user_tx_db(transaction).update_many(
where={},
data=(
{**non_default_values, "model_max_budget": json.dumps(non_default_values["model_max_budget"])}
if "model_max_budget" in non_default_values
else non_default_values
),
)
for start in range(0, len(all_users_in_db), _USER_CACHE_INVALIDATION_BATCH_SIZE):
await asyncio.gather(

29
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -1047,12 +1047,12 @@ async def delete_organization(
deleted_orgs: Final = []
for organization_id in data.organization_ids:
# delete all teams in the organization
await _table(TeamRepository(prisma_client)).delete_many(where={"organization_id": organization_id})
# delete all members in the organization
await _table(OrganizationMembershipRepository(prisma_client)).delete_many(
where={"organization_id": organization_id}
)
async with prisma_client.tx() as transaction:
_ = await transaction.query_raw(_LOCK_ORG_TEAMS_SQL, organization_id)
_ = await transaction.litellm_teamtable.delete_many(where={"organization_id": organization_id})
async with prisma_client.tx() as transaction:
_ = await transaction.query_raw(_LOCK_ORG_MEMBERS_SQL, organization_id)
_ = await transaction.litellm_organizationmembership.delete_many(where={"organization_id": organization_id})
await _delete_organization_keys(
organization_id=organization_id,
prisma_client=prisma_client,
@ -1074,6 +1074,19 @@ async def delete_organization(
return deleted_orgs
_LOCK_ORG_TEAMS_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE organization_id = $1 ORDER BY team_id COLLATE "C" FOR UPDATE'
)
_LOCK_ORG_MEMBERS_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_OrganizationMembership" WHERE organization_id = $1 '
'ORDER BY organization_id COLLATE "C", user_id COLLATE "C" FOR UPDATE'
)
_LOCK_ORG_KEYS_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE organization_id = $1 '
'ORDER BY token COLLATE "C" FOR UPDATE'
)
async def _delete_organization_keys(
organization_id: str,
prisma_client: PrismaClient,
@ -1087,7 +1100,9 @@ async def _delete_organization_keys(
hashed_tokens=hashed_tokens_to_delete,
prisma_client=prisma_client,
)
await _table(VerificationTokenRepository(prisma_client)).delete_many(where=key_filter)
async with prisma_client.tx() as transaction:
_ = await transaction.query_raw(_LOCK_ORG_KEYS_SQL, organization_id)
_ = await transaction.litellm_verificationtoken.delete_many(where={"organization_id": organization_id})
await delete_cache_key_objects(
hashed_tokens=hashed_tokens_to_delete,
user_api_key_cache=user_api_key_cache,

10
litellm/proxy/management_helpers/access_group_model_sync.py
View file

@ -6,6 +6,7 @@ the arrays alone strands every group on a name nothing serves any more.
"""
from collections.abc import Sequence
from contextlib import AbstractAsyncContextManager
from typing import Final, Protocol
from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
@ -53,6 +54,15 @@ _REMOVE_MODEL_NAME_SQL: Final = (
)
class RawTransactions(Protocol):
def tx(self) -> AbstractAsyncContextManager[RawExecutor]: ...
def raw_transactions(prisma_client: object) -> RawTransactions:
db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client
return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin
def raw_executor(prisma_client: object) -> RawExecutor:
db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client
return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin

62
litellm/proxy/management_helpers/model_allowlist_rename_sync.py
View file

@ -6,18 +6,27 @@ Those allowlists store public model names, not ids, so a deployment rename that
alone denies the new name while the old entry grants a name nothing serves any more.
"""
from collections.abc import Callable
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from pydantic import ConfigDict
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_helpers.access_group_model_sync import raw_executor, still_backed
from litellm.proxy.management_helpers.access_group_model_sync import raw_executor, raw_transactions, still_backed
from litellm.proxy.management_helpers.access_group_team_sync import TEAM_ADVISORY_LOCK_SQL
from litellm.router import Router
from litellm.types.llms.base import LiteLLMBaseModel
class _ListingTeam(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
team_id: str
class _TouchedRow(LiteLLMBaseModel):
kind: str
object_id: str
@ -80,11 +89,56 @@ def _rewrite_sql(set_clause: str, where_clause: str) -> str:
return f"WITH {ctes} {rows}"
_TEAMS_LISTING_SQL: Final = (
'SELECT team_id FROM "LiteLLM_TeamTable" WHERE $1 = ANY("models") ORDER BY team_id COLLATE "C"'
)
_LOCK_ORGS_LISTING_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_OrganizationTable" WHERE $1 = ANY("models") '
'ORDER BY organization_id COLLATE "C" FOR UPDATE'
)
_LOCK_TEAMS_LISTING_SQL: Final = (
'SELECT team_id FROM "LiteLLM_TeamTable" WHERE $1 = ANY("models") ORDER BY team_id COLLATE "C" FOR UPDATE'
)
_LOCK_USERS_LISTING_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_UserTable" WHERE $1 = ANY("models") ORDER BY user_id COLLATE "C" FOR UPDATE'
)
_LOCK_KEYS_LISTING_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE $1 = ANY("models") ORDER BY token COLLATE "C" FOR UPDATE'
)
_LOCK_PROJECTS_LISTING_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_ProjectTable" WHERE $1 = ANY("models") ORDER BY project_id COLLATE "C" FOR UPDATE'
)
_REPLACE_SQL: Final = _rewrite_sql('array_replace(array_remove("models", $2), $1, $2)', '$1 = ANY("models")')
_APPEND_SQL: Final = _rewrite_sql('array_append("models", $2)', '$1 = ANY("models") AND NOT ($2 = ANY("models"))')
def _team_ids(rows: Sequence[object]) -> frozenset[str]:
return frozenset(_ListingTeam.model_validate(row).team_id for row in rows)
async def _rewrite_allowlists(
prisma_client: object, old_name: str, new_name: str, *, append: bool, attempts: int
) -> Sequence[object]:
"""Team advisory locks come from a snapshot, so a team that starts listing the name before its row lock would
be locked without one. That attempt commits having written nothing, and the next one locks the fresh set."""
async with raw_transactions(prisma_client).tx() as transaction:
listed: Final = _team_ids(await transaction.query_raw(_TEAMS_LISTING_SQL, old_name))
for team_id in sorted(listed):
_ = await transaction.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
_ = await transaction.query_raw(_LOCK_ORGS_LISTING_SQL, old_name)
locked: Final = _team_ids(await transaction.query_raw(_LOCK_TEAMS_LISTING_SQL, old_name))
retry: Final = attempts > 1 and not locked <= listed
if not retry:
_ = await transaction.query_raw(_LOCK_USERS_LISTING_SQL, old_name)
_ = await transaction.query_raw(_LOCK_KEYS_LISTING_SQL, old_name)
_ = await transaction.query_raw(_LOCK_PROJECTS_LISTING_SQL, old_name)
return await transaction.query_raw(_APPEND_SQL if append else _REPLACE_SQL, old_name, new_name)
return await _rewrite_allowlists(prisma_client, old_name, new_name, append=append, attempts=attempts - 1)
async def sync_model_allowlists_for_renamed_model(
prisma_client: object,
*,
@ -98,8 +152,8 @@ async def sync_model_allowlists_for_renamed_model(
return
executor: Final = raw_executor(prisma_client)
old_name_still_backed: Final = await still_backed(executor, llm_router, old_name, model_id)
touched_rows: Final = await executor.query_raw(
_APPEND_SQL if old_name_still_backed else _REPLACE_SQL, old_name, new_name
touched_rows: Final = await _rewrite_allowlists(
prisma_client, old_name, new_name, append=old_name_still_backed, attempts=3
)
touched: Final = tuple(_TouchedRow.model_validate(row) for row in touched_rows)
await evict_and_broadcast(

28
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -77,9 +77,6 @@ from litellm.proxy.utils import handle_exception_on_proxy, query_raw_within_stat
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import SpendLogsRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.verification_token_repository import (
VerificationTokenRepository,
)
from litellm.types.proxy.spend_capture_rate import CaptureRateReport, SpendCaptureProvider
if TYPE_CHECKING:
@ -292,24 +289,22 @@ class _TeamTable(Protocol):
async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ...
class _VerificationTokenTable(Protocol):
"""The subset of the Prisma verification token table API this module uses."""
async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ...
def _spend_logs_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_SpendLogs"]:
return SpendLogsRepository(prisma_client).table
_LOCK_ALL_TEAMS_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE TRUE ORDER BY team_id COLLATE "C" FOR UPDATE'
)
_LOCK_ALL_KEYS_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE TRUE ORDER BY token COLLATE "C" FOR UPDATE'
)
def _team_table(prisma_client: PrismaClient) -> _TeamTable:
return TeamRepository(prisma_client).table
def _verification_token_table(prisma_client: PrismaClient) -> _VerificationTokenTable:
return VerificationTokenRepository(prisma_client).table
def _spend_logs_daily_summary_sql(
*,
start_date_iso: str,
@ -3626,8 +3621,11 @@ async def global_spend_reset():
code=status.HTTP_401_UNAUTHORIZED,
)
await _verification_token_table(prisma_client).update_many(data={"spend": 0.0}, where={})
await _team_table(prisma_client).update_many(data={"spend": 0.0}, where={})
async with prisma_client.tx() as transaction:
_ = await transaction.query_raw(_LOCK_ALL_TEAMS_SQL)
_ = await transaction.litellm_teamtable.update_many(data={"spend": 0.0}, where={})
_ = await transaction.query_raw(_LOCK_ALL_KEYS_SQL)
_ = await transaction.litellm_verificationtoken.update_many(data={"spend": 0.0}, where={})
return {
"message": "Spend for all API Keys and Teams reset successfully",

38
litellm/proxy/utils.py
View file

@ -5545,6 +5545,20 @@ class _StaleReadEngine:
return self.wrapper is current and self.generation == current.engine_generation
_LOCK_KEYS_BY_TOKEN_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE token = ANY($1::text[]) '
'ORDER BY token COLLATE "C" FOR UPDATE'
)
_LOCK_KEYS_BY_TEAM_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE team_id = ANY($1::text[]) '
'ORDER BY token COLLATE "C" FOR UPDATE'
)
_LOCK_TEAMS_BY_ID_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE team_id = ANY($1::text[]) '
'ORDER BY team_id COLLATE "C" FOR UPDATE'
)
class PrismaClient:
spend_log_transactions: list = []
_spend_log_transactions_lock = asyncio.Lock()
@ -6845,24 +6859,26 @@ class PrismaClient:
else:
hashed_token = token
hashed_tokens.append(hashed_token)
filter_query: dict[str, object] = {}
if user_id is not None:
filter_query = {"AND": [{"token": {"in": hashed_tokens}}, {"user_id": user_id}]}
else:
filter_query = {"token": {"in": hashed_tokens}}
deleted_tokens: Final[int] = await VerificationTokenRepository(self).table.delete_many(
where=filter_query
)
async with self.tx() as transaction:
_ = await transaction.query_raw(_LOCK_KEYS_BY_TOKEN_SQL, hashed_tokens)
deleted_tokens: Final[int] = await transaction.litellm_verificationtoken.delete_many(
where={"token": {"in": hashed_tokens}, "user_id": user_id}
if user_id is not None
else {"token": {"in": hashed_tokens}}
)
verbose_proxy_logger.debug("deleted_tokens: %s", deleted_tokens)
return {"deleted_keys": deleted_tokens}
elif table_name == "team" and team_id_list is not None and isinstance(team_id_list, list):
# admin only endpoint -> `/team/delete`
await TeamRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}})
async with self.tx() as transaction:
_ = await transaction.query_raw(_LOCK_TEAMS_BY_ID_SQL, team_id_list)
_ = await transaction.litellm_teamtable.delete_many(where={"team_id": {"in": team_id_list}})
return {"deleted_teams": team_id_list}
elif table_name == "key" and team_id_list is not None and isinstance(team_id_list, list):
# admin only endpoint -> `/team/delete`
await VerificationTokenRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}})
async with self.tx() as transaction:
_ = await transaction.query_raw(_LOCK_KEYS_BY_TEAM_SQL, team_id_list)
_ = await transaction.litellm_verificationtoken.delete_many(where={"team_id": {"in": team_id_list}})
except Exception as e:
import traceback

52
litellm/repositories/accounting_delivery_repository.py
View file

@ -5,7 +5,7 @@ from itertools import chain
from types import TracebackType
from typing import TYPE_CHECKING, Final, Literal, Protocol
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
from pydantic import BaseModel, ConfigDict, TypeAdapter
from litellm.repositories.prisma_protocols import TableActions
@ -41,7 +41,10 @@ class _Delivery(BaseModel):
delivery_id: str
attempts: int
accounting_payload: JsonValue
class _ClaimedDelivery(_Delivery):
accounting_payload: str
@dataclass(frozen=True, slots=True)
@ -68,9 +71,11 @@ class DeliveryMetrics(BaseModel):
_DELIVERIES: Final = TypeAdapter(tuple[_Delivery, ...])
_CLAIMED: Final = TypeAdapter(tuple[_ClaimedDelivery, ...])
_PAYLOAD: Final = TypeAdapter(dict[str, object])
_TIMEOUT: Final = timedelta(seconds=60)
_CLAIM: Final = """SELECT delivery_id, attempts, accounting_payload FROM "LiteLLM_SpendLogDelivery"
_CLAIM: Final = """SELECT delivery_id, attempts, accounting_payload::text AS accounting_payload
FROM "LiteLLM_SpendLogDelivery"
WHERE state = 'pending' AND next_attempt_at <= CURRENT_TIMESTAMP
AND (attempts > 0) = $2::boolean
ORDER BY next_attempt_at, delivery_id LIMIT $1::integer FOR UPDATE SKIP LOCKED"""
@ -79,10 +84,20 @@ SET state = 'applied', applied_at = CURRENT_TIMESTAMP, last_error = NULL
WHERE delivery_id = ANY($1::text[]) AND state = 'pending' """
_FAILED: Final = """UPDATE "LiteLLM_SpendLogDelivery"
SET attempts = attempts + 1, last_error = $3,
next_attempt_at = CURRENT_TIMESTAMP + make_interval(secs => LEAST(300, power(2, LEAST(attempts, 8)))::integer),
next_attempt_at = CURRENT_TIMESTAMP + make_interval(secs => LEAST(300, power(2, LEAST(attempts, 8)) * (0.5 + random()))),
state = $4::"AccountingDeliveryState"
WHERE delivery_id = $1 AND state = 'pending' AND attempts = $2::integer
RETURNING delivery_id, attempts, accounting_payload"""
RETURNING delivery_id, attempts"""
_UNDECODABLE: Final = """UPDATE "LiteLLM_SpendLogDelivery"
SET state = 'parked', attempts = attempts + 1, last_error = 'UndecodablePayload'
WHERE delivery_id = ANY($1::text[]) AND state = 'pending'"""
def _decoded(delivery: _ClaimedDelivery) -> DeliveryProjection | None:
try:
return DeliveryProjection(_PAYLOAD.validate_json(delivery.accounting_payload))
except ValueError:
return None
def _is_permanent_failure(error: Exception) -> bool:
@ -138,26 +153,25 @@ class AccountingDeliveryRepository:
try:
async with self._database.tx(timeout=_TIMEOUT) as transaction:
retries: Final = (
_DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, 1, True))
_CLAIMED.validate_python(await transaction.query_raw(_CLAIM, 1, True))
if retry_only is not False
else ()
)
claimed = retries or ( # rebind-ok: retain claimed identities after rollback
_DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, batch_size, False))
rows: Final = retries or (
_CLAIMED.validate_python(await transaction.query_raw(_CLAIM, batch_size, False))
if retry_only is not True
else ()
)
if not claimed:
return DeliveryBatchResult()
await apply(
transaction,
tuple(
DeliveryProjection(_PAYLOAD.validate_python(delivery.accounting_payload))
for delivery in claimed
),
)
await transaction.execute_raw(_APPLIED, tuple(delivery.delivery_id for delivery in claimed))
return DeliveryBatchResult(applied=len(claimed))
decoded: Final = tuple((row, _decoded(row)) for row in rows)
undecodable: Final = tuple(row.delivery_id for row, projection in decoded if projection is None)
applicable: Final = tuple((row, projection) for row, projection in decoded if projection is not None)
if undecodable:
await transaction.execute_raw(_UNDECODABLE, undecodable)
claimed = tuple(row for row, _ in applicable) # rebind-ok: retain claimed identities after rollback
if claimed:
await apply(transaction, tuple(projection for _, projection in applicable))
await transaction.execute_raw(_APPLIED, tuple(row.delivery_id for row in claimed))
return DeliveryBatchResult(applied=len(claimed), parked=undecodable)
except Exception as error: # noqa: BLE001 # roll back all destinations before recording a retry
if not claimed:
raise

42
litellm/repositories/unit_of_work.py
View file

@ -109,6 +109,46 @@ class TeamMembershipSpendResetWrites:
self.batch.execute_raw(_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL, list(budget_ids))
_LOCK_BUDGET_ORGS_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_OrganizationTable" WHERE budget_id = ANY($1::text[]) '
'ORDER BY organization_id COLLATE "C" FOR UPDATE'
)
_LOCK_BUDGET_MEMBERSHIPS_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_TeamMembership" WHERE budget_id = ANY($1::text[]) '
'ORDER BY team_id COLLATE "C", user_id COLLATE "C" FOR UPDATE'
)
_LOCK_BUDGET_KEYS_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE budget_id = ANY($1::text[]) '
'ORDER BY token COLLATE "C" FOR UPDATE'
)
_LOCK_BUDGET_PROJECTS_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_ProjectTable" WHERE budget_id = ANY($1::text[]) '
'ORDER BY project_id COLLATE "C" FOR UPDATE'
)
_LOCK_BUDGET_TAGS_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_TagTable" WHERE budget_id = ANY($1::text[]) '
'ORDER BY tag_name COLLATE "C" FOR UPDATE'
)
_LOCK_BUDGET_ACCESS_GROUPS_SQL: Final = (
'SELECT 1 AS locked FROM "LiteLLM_ModelAccessGroupBudgetTable" WHERE budget_id = ANY($1::text[]) '
'ORDER BY access_group_name COLLATE "C" FOR UPDATE'
)
@dataclass(frozen=True, slots=True)
class BudgetLinkedRowLocks:
batch: PrismaBatch
def queue_row_locks(self, budget_ids: Sequence[str]) -> None:
"""Locks every row a cascade may reset, in the shared row lock order, before the filtered resets run."""
self.batch.execute_raw(_LOCK_BUDGET_ORGS_SQL, list(budget_ids))
self.batch.execute_raw(_LOCK_BUDGET_MEMBERSHIPS_SQL, list(budget_ids))
self.batch.execute_raw(_LOCK_BUDGET_KEYS_SQL, list(budget_ids))
self.batch.execute_raw(_LOCK_BUDGET_PROJECTS_SQL, list(budget_ids))
self.batch.execute_raw(_LOCK_BUDGET_TAGS_SQL, list(budget_ids))
self.batch.execute_raw(_LOCK_BUDGET_ACCESS_GROUPS_SQL, list(budget_ids))
@dataclass(frozen=True, slots=True)
class BudgetWindowWrites:
table: BatchTable
@ -137,6 +177,7 @@ class BudgetCascadeUnitOfWork:
``enduser_spend_reset_unit_of_work`` because one tier can gate millions.
"""
row_locks: BudgetLinkedRowLocks
team_memberships: TeamMembershipSpendResetWrites
keys: LinkedSpendResetWrites
organizations: LinkedSpendResetWrites
@ -163,6 +204,7 @@ async def budget_cascade_unit_of_work(
) -> AsyncGenerator[BudgetCascadeUnitOfWork, None]:
batch = new_batch()
yield BudgetCascadeUnitOfWork(
row_locks=BudgetLinkedRowLocks(batch=batch),
team_memberships=TeamMembershipSpendResetWrites(
resets=LinkedSpendResetWrites(table=batch.litellm_teammembership), batch=batch
),

197
tests/integration/database/test_spend_lock_order.py Normal file
View file

@ -0,0 +1,197 @@
import asyncio
import subprocess
import sys
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
from pathlib import Path
from types import SimpleNamespace
from typing import Final
import psycopg
import pytest
import pytest_asyncio
from prisma import Prisma
from litellm.proxy._types import DBSpendUpdateTransactions
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model
from litellm.repositories import unit_of_work
from tests.integration._support.database import scratch_database
_ROOT: Final = Path(__file__).resolve().parents[3]
_KEY: Final = "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683"
_CHARGES: Final = DBSpendUpdateTransactions(
user_list_transactions={"user": 1.0},
end_user_list_transactions={},
key_list_transactions={_KEY: 1.0},
team_list_transactions={"team": 1.0},
team_member_list_transactions={"team_id::team::user_id::user": 1.0},
org_list_transactions={"org": 1.0},
org_member_list_transactions={"org_id::org::user_id::user": 1.0},
project_list_transactions={"project": 1.0},
tag_list_transactions={"tag": 1.0},
agent_list_transactions={},
model_access_group_list_transactions={"group": 1.0},
)
_SEED: Final = (
"INSERT INTO \"LiteLLM_BudgetTable\" (budget_id, created_by, updated_by) VALUES ('budget', 't', 't')",
'INSERT INTO "LiteLLM_OrganizationTable" (organization_id, organization_alias, budget_id, created_by, updated_by) '
"VALUES ('org', 'org', 'budget', 't', 't')",
"INSERT INTO \"LiteLLM_UserTable\" (user_id) VALUES ('user')",
'INSERT INTO "LiteLLM_TeamTable" (team_id, members_with_roles) '
"""VALUES ('team', '[{"user_id": "user", "role": "user"}]')""",
"INSERT INTO \"LiteLLM_TeamMembership\" (team_id, user_id, budget_id) VALUES ('team', 'user', 'budget')",
"INSERT INTO \"LiteLLM_OrganizationMembership\" (organization_id, user_id) VALUES ('org', 'user')",
f"INSERT INTO \"LiteLLM_VerificationToken\" (token, budget_id) VALUES ('{_KEY}', 'budget')",
"INSERT INTO \"LiteLLM_ProjectTable\" (project_id, budget_id, created_by, updated_by) "
"VALUES ('project', 'budget', 't', 't')",
"INSERT INTO \"LiteLLM_TagTable\" (tag_name, budget_id) VALUES ('tag', 'budget')",
"INSERT INTO \"LiteLLM_ModelAccessGroupBudgetTable\" (access_group_name, budget_id) VALUES ('group', 'budget')",
)
@pytest.fixture(scope="module")
def schema_sql() -> str:
return subprocess.run(
[sys.executable, "-I", "-m", "prisma", "migrate", "diff", "--from-empty", "--to-schema-datamodel",
str(_ROOT / "schema.prisma"), "--script"],
check=True,
capture_output=True,
text=True,
cwd=_ROOT,
).stdout
@pytest.fixture
def database(schema_sql: str) -> Iterator[str]:
with scratch_database("spend_lock_order") as url, psycopg.connect(url, autocommit=True) as connection:
connection.execute(schema_sql)
for statement in _SEED:
connection.execute(statement)
yield url
@pytest_asyncio.fixture(loop_scope="function")
async def ledger(database: str) -> AsyncIterator[Prisma]:
client: Final = Prisma(datasource={"url": database})
await client.connect()
try:
yield client
finally:
await client.disconnect()
async def _apply_charges(client: Prisma) -> None:
async with client.tx() as transaction:
await DBSpendUpdateWriter().apply_accounting_entities(client, None, transaction, _CHARGES)
async def _ledger_waits_for_a_lock(url: str) -> None:
with psycopg.connect(url, autocommit=True) as observer:
for _ in range(200):
if observer.execute("SELECT 1 FROM pg_locks WHERE NOT granted").fetchone():
return
await asyncio.sleep(0.05)
raise AssertionError("the ledger never waited on the other writer")
def _spend(url: str) -> Sequence[object]:
with psycopg.connect(url) as connection:
return connection.execute(
'SELECT (SELECT spend FROM "LiteLLM_UserTable"), (SELECT spend FROM "LiteLLM_TeamTable"), '
'(SELECT spend FROM "LiteLLM_TeamMembership"), (SELECT spend FROM "LiteLLM_VerificationToken"), '
'(SELECT spend FROM "LiteLLM_OrganizationTable"), (SELECT spend FROM "LiteLLM_ProjectTable")'
).fetchone() or ()
_Writer = Callable[[psycopg.Connection[tuple[object, ...]]], None]
def _roster_first(connection: psycopg.Connection[tuple[object, ...]]) -> None:
connection.execute("SELECT pg_advisory_xact_lock(hashtext('team'))")
def _roster_rest(connection: psycopg.Connection[tuple[object, ...]]) -> None:
connection.execute("SELECT 1 FROM \"LiteLLM_UserTable\" WHERE user_id = 'user' FOR UPDATE")
connection.execute("UPDATE \"LiteLLM_TeamTable\" SET members_with_roles = members_with_roles WHERE team_id = 'team'")
def _access_group_first(connection: psycopg.Connection[tuple[object, ...]]) -> None:
connection.execute("UPDATE \"LiteLLM_TeamTable\" SET access_group_ids = ARRAY['ag'] WHERE team_id = 'team'")
def _access_group_rest(connection: psycopg.Connection[tuple[object, ...]]) -> None:
connection.execute(f"UPDATE \"LiteLLM_VerificationToken\" SET access_group_ids = ARRAY['ag'] WHERE token = '{_KEY}'")
def _cascade_first(connection: psycopg.Connection[tuple[object, ...]]) -> None:
connection.execute(unit_of_work._LOCK_BUDGET_ORGS_SQL.replace("$1", "%s"), (["budget"],))
def _cascade_rest(connection: psycopg.Connection[tuple[object, ...]]) -> None:
for statement in (
unit_of_work._LOCK_BUDGET_MEMBERSHIPS_SQL,
unit_of_work._LOCK_BUDGET_KEYS_SQL,
unit_of_work._LOCK_BUDGET_PROJECTS_SQL,
unit_of_work._LOCK_BUDGET_TAGS_SQL,
unit_of_work._LOCK_BUDGET_ACCESS_GROUPS_SQL,
):
connection.execute(statement.replace("$1", "%s"), (["budget"],))
connection.execute("UPDATE \"LiteLLM_TagTable\" SET spend = 0 WHERE budget_id = 'budget'")
@pytest.mark.parametrize(
("first", "rest"),
[(_roster_first, _roster_rest), (_access_group_first, _access_group_rest), (_cascade_first, _cascade_rest)],
ids=["team_roster", "access_group", "budget_cascade"],
)
@pytest.mark.asyncio(loop_scope="function")
async def test_charges_wait_for_a_concurrent_writer_instead_of_deadlocking(
database: str, ledger: Prisma, first: _Writer, rest: _Writer
) -> None:
"""The other writer takes its first lock, the ledger starts and must block holding nothing the other writer
needs next, so the other writer finishes and the charge lands once after it."""
with psycopg.connect(database) as writer:
writer.execute("SET lock_timeout = '10s'")
first(writer)
charges: Final = asyncio.create_task(_apply_charges(ledger))
await _ledger_waits_for_a_lock(database)
await asyncio.to_thread(rest, writer)
writer.commit()
await asyncio.wait_for(charges, timeout=20)
assert _spend(database) == (1.0, 1.0, 1.0, 1.0, 1.0, 1.0)
async def _rename(client: Prisma) -> None:
await sync_model_allowlists_for_renamed_model(
SimpleNamespace(db=client),
model_id="deployment",
old_name="old-model",
new_name="new-model",
llm_router=None,
user_api_key_cache=UserApiKeyCache(),
)
@pytest.mark.asyncio(loop_scope="function")
async def test_a_model_rename_waits_for_a_roster_change_instead_of_deadlocking(database: str, ledger: Prisma) -> None:
"""The roster change holds the team lock and the user row, so the rename must wait before it locks the team
row the roster change updates next."""
with psycopg.connect(database, autocommit=True) as seed:
seed.execute("UPDATE \"LiteLLM_TeamTable\" SET models = ARRAY['old-model']")
seed.execute("UPDATE \"LiteLLM_UserTable\" SET models = ARRAY['old-model']")
with psycopg.connect(database) as writer:
writer.execute("SET lock_timeout = '10s'")
_roster_first(writer)
writer.execute("SELECT 1 FROM \"LiteLLM_UserTable\" WHERE user_id = 'user' FOR UPDATE")
rename: Final = asyncio.create_task(_rename(ledger))
await _ledger_waits_for_a_lock(database)
await asyncio.to_thread(_roster_rest, writer)
writer.commit()
await asyncio.wait_for(rename, timeout=20)
with psycopg.connect(database) as connection:
assert connection.execute(
'SELECT (SELECT models FROM "LiteLLM_TeamTable"), (SELECT models FROM "LiteLLM_UserTable")'
).fetchone() == (["new-model"], ["new-model"])

63
tests/integration/database/test_spendlogs_delivery.py
View file

@ -131,9 +131,9 @@ async def test_duplicate_rows_in_one_flush_keep_first_accounting_payload(deliver
"accounting_log_disabled": mode == "disabled",
}
assert await repository.insert_with_deliveries((first, later), write_spend_logs=mode != "url") == 1
assert delivery_database.execute(
'SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"'
).fetchone() == (first["accounting_payload"],)
assert delivery_database.execute('SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"').fetchone() == (
first["accounting_payload"],
)
assert delivery_database.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == (
[(first["spend"],)] if mode == "local" else []
)
@ -415,11 +415,28 @@ def test_delivery_migration_preserves_unused_old_receipts(state):
assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,)
@pytest.mark.parametrize(
("payload", "huge_state"),
[
(
"jsonb_set(accounting_payload, '{metadata}', ('{\"tokens_saved\": ' || repeat('9', 401) || '}')::jsonb)",
"applied",
),
(
"jsonb_set(accounting_payload, '{metadata}', ('{\"tokens_saved\": ' || repeat('9', 5000) || '}')::jsonb)",
"parked",
),
("'[]'::jsonb", "parked"),
],
ids=["beyond_float", "beyond_python_int_parse", "not_an_object"],
)
@pytest.mark.asyncio
async def test_malformed_payload_is_parked_without_losing_other_deliveries(delivery_database, delivery_client):
_insert(delivery_database, (_row("bad"), _row("good")))
async def test_a_payload_the_engine_cannot_decode_never_blocks_other_deliveries(
delivery_database, delivery_client, payload, huge_state
):
_insert(delivery_database, (_row("huge"), _row("good")))
delivery_database.execute(
'UPDATE "LiteLLM_SpendLogDelivery" SET accounting_payload=$$ [] $$::jsonb WHERE request_id=$$bad$$'
f'UPDATE "LiteLLM_SpendLogDelivery" SET accounting_payload = {payload} WHERE request_id = $$huge$$'
)
async def apply(transaction, deliveries):
@ -428,12 +445,16 @@ async def test_malformed_payload_is_parked_without_losing_other_deliveries(deliv
)
repository = AccountingDeliveryRepository(delivery_client)
assert (await repository.apply_pending(apply, max_attempts=2)).failed == 2
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,)
delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at=CURRENT_TIMESTAMP')
assert (await repository.apply_pending(apply, max_attempts=2)).parked == ("bad",)
assert (await repository.apply_pending(apply, max_attempts=2)).applied == 1
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,)
results = [await repository.apply_pending(apply, max_attempts=1) for _ in range(3)]
_insert(delivery_database, (_row("later"),))
results.append(await repository.apply_pending(apply, max_attempts=1))
applied: Final = 3 if huge_state == "applied" else 2
assert sum(result.applied for result in results) == applied
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5 * applied,)
assert delivery_database.execute(
'SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id'
).fetchall() == [("good", "applied"), ("huge", huge_state), ("later", "applied")]
@pytest.mark.asyncio
@ -463,19 +484,29 @@ async def test_ledger_batch_retrieves_charge_once_in_every_logging_mode(
from litellm.proxy.utils import jsonify_object
completed = {
**_row("batch_completed"), "call_type": "aretrieve_batch", "status": "success", "spend": 0.5,
**_row("batch_completed"),
"call_type": "aretrieve_batch",
"status": "success",
"spend": 0.5,
"metadata": '{"status":"success","private":"request detail"}',
"accounting_payload": {
"request_id": "batch_completed", "call_type": "aretrieve_batch", "status": "success", "spend": 0.5,
"request_id": "batch_completed",
"call_type": "aretrieve_batch",
"status": "success",
"spend": 0.5,
"metadata": {"accounting_savings": {"compression": 0.2}},
},
}
monkeypatch.setattr(proxy_server, "general_settings", {"spend_logs_metadata_fields": {"include": ["status"]}})
client = SimpleNamespace(db=delivery_client, writer_db=delivery_client, jsonify_object=jsonify_object)
writer = DBSpendUpdateWriter()
first = await writer._claim_settle_once_cost_row(payload=completed, prisma_client=client, disable_spend_logs=disabled)
first = await writer._claim_settle_once_cost_row(
payload=completed, prisma_client=client, disable_spend_logs=disabled
)
again = await writer._claim_settle_once_cost_row(
payload={**completed, "startTime": "2026-11-05T12:00:00+00:00"}, prisma_client=client, disable_spend_logs=disabled
payload={**completed, "startTime": "2026-11-05T12:00:00+00:00"},
prisma_client=client,
disable_spend_logs=disabled,
)
assert (first, again) == (True, False)
assert delivery_database.execute('SELECT spend, metadata FROM "LiteLLM_SpendLogs"').fetchall() == (

43
tests/integration/database/test_spendlogs_ledger_parity.py
View file

@ -63,10 +63,10 @@ _DAILY_TABLES: Final = (
)
@pytest.mark.parametrize("batch_size,row_count,expected,distinct_models", [(None, 251, 251, False), ("200", 251, 200, False), (None, 10001, 10000, False), (None, 3001, 3001, True)])
@pytest.mark.parametrize("batch_size,row_count,transactions,distinct_models", [(None, 251, 1, False), (None, 1001, 3, False), ("200", 251, 2, False), ("3001", 3001, 1, True)])
@pytest.mark.asyncio
async def test_worker_applies_the_configured_backlog_in_one_transaction(
parity_database: str, monkeypatch: pytest.MonkeyPatch, batch_size: str | None, row_count: int, expected: int, distinct_models: bool
async def test_worker_drains_a_backlog_in_one_run_one_configured_batch_per_transaction(
parity_database: str, monkeypatch: pytest.MonkeyPatch, batch_size: str | None, row_count: int, transactions: int, distinct_models: bool
):
from litellm.accounting.worker import flush_accounting_deliveries
@ -92,13 +92,37 @@ async def test_worker_applies_the_configured_backlog_in_one_transaction(
)) == row_count - 1
started = time.monotonic()
await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj)
print(f"Ledger batch applied={expected} queued={row_count} seconds={time.monotonic() - started:.3f}")
print(f"Ledger backlog queued={row_count} seconds={time.monotonic() - started:.3f}")
with psycopg.connect(parity_database) as connection:
assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'applied\'').fetchone() == (expected,)
assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\'').fetchone() == (row_count - expected,)
assert connection.execute('SELECT COUNT(DISTINCT applied_at) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (1,)
assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (expected * 0.25,)
assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"').fetchone() == (expected if distinct_models else 1,)
assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'applied\'').fetchone() == (row_count,)
assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\'').fetchone() == (0,)
assert connection.execute('SELECT COUNT(DISTINCT applied_at) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (transactions,)
assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (row_count * 0.25,)
assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"').fetchone() == (row_count if distinct_models else 1,)
finally:
await client.db.disconnect()
@pytest.mark.asyncio
async def test_worker_removes_more_expired_receipts_than_one_cleanup_batch_in_one_run(
parity_database: str, monkeypatch: pytest.MonkeyPatch
):
from litellm.accounting.worker import flush_accounting_deliveries
monkeypatch.setenv("DATABASE_URL", parity_database)
monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "100")
client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj)
with psycopg.connect(parity_database, autocommit=True) as connection:
connection.execute(
'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch, state, applied_at) '
"SELECT 'expired-' || n, '{}'::jsonb, false, 'applied', CURRENT_TIMESTAMP - INTERVAL '2 days' "
"FROM generate_series(1, 501) AS n"
)
await client.db.connect()
try:
await flush_accounting_deliveries(DBSpendUpdateWriter(), client, proxy_server.proxy_logging_obj)
with psycopg.connect(parity_database) as connection:
assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,)
finally:
await client.db.disconnect()
@ -404,6 +428,7 @@ async def _assert_legacy_database_parity(
assert len(rows) == 1
if path == "ledger" and retention == "lean":
assert rows[0]["accounting_log_disabled"] is True
assert not {"messages", "response", "proxy_server_request", "metadata"}.intersection(rows[0])
frozen = rows[0]["accounting_payload"]
if isinstance(frozen, str):
frozen = json.loads(frozen)

13
tests/unit/accounting/test_projection.py
View file

@ -150,6 +150,7 @@ async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest
client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj)
monkeypatch.setattr(proxy_server, "prisma_client", client)
monkeypatch.setattr(proxy_server, "disable_spend_logs", False)
monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False)
monkeypatch.setattr(client, "spend_log_transactions", [])
monkeypatch.setattr(client, "model_usage_transactions", [])
writer: Final = DBSpendUpdateWriter()
@ -521,7 +522,9 @@ def test_stored_payload_keeps_guardrail_billing_but_not_guardrail_content() -> N
assert dict(daily_user.increments)["compression_saved_tokens"] == 7
def test_caller_supplied_savings_that_cannot_be_computed_never_drop_the_charge(monkeypatch: pytest.MonkeyPatch) -> None:
def test_caller_supplied_savings_beyond_json_range_are_clamped_and_never_drop_the_charge(
monkeypatch: pytest.MonkeyPatch,
) -> None:
import litellm
monkeypatch.setitem(
@ -552,6 +555,8 @@ def test_caller_supplied_savings_that_cannot_be_computed_never_drop_the_charge(m
)
restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str))
assert dict(restored["metadata"])["accounting_savings"]["compression"] == 0.0
assert [dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken"] == [
0.5
]
assert dict(restored["metadata"])["compression_savings"] == {"tokens_saved": 2**53 - 1}
assert dict(restored["metadata"])["compression_savings"] == {"tokens_saved": 2**53 - 1}
assert [
dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken"
] == [0.5]

10
tests/unit/litellm_core_utils/test_safe_json_dumps.py
View file

@ -230,9 +230,7 @@ def test_pydantic_base_model():
inner: InnerModel
tags: list
outer = OuterModel(
name="test", inner=InnerModel(value=42, label="hello"), tags=["a", "b"]
)
outer = OuterModel(name="test", inner=InnerModel(value=42, label="hello"), tags=["a", "b"])
# Test a pydantic model at the top level
result = json.loads(safe_dumps(outer))
@ -261,3 +259,9 @@ def test_safe_json_structure_keeps_tuples_and_drops_non_string_keys():
assert structure == {"models": ("A", "B"), "tags": ["X", "Y"], "nested": {"deep": ("C",)}}
assert type(structure["models"]) is tuple
assert json.loads(safe_dumps(data)) == {"models": ["a", "b"], "tags": ["x", "y"], "nested": {"deep": ["c"]}}
def test_integer_limit_clamps_both_signs_and_leaves_bools_and_floats_alone() -> None:
assert safe_json_structure(
{"big": 10**400, "small": -(10**400), "inside": 7, "flag": True, "ratio": 1e300}, integer_limit=100
) == {"big": 100, "small": -100, "inside": 7, "flag": True, "ratio": 1e300}

6
tests/unit/proxy/common_utils/test_check_batch_cost.py
View file

@ -3324,9 +3324,13 @@ class TestMultiPodBatchCostClaim:
@pytest.mark.parametrize("failure_phase", ("before_write", "after_write", "external"))
@pytest.mark.parametrize("collector_enabled", (False, True))
async def test_builtin_accounting_retries_only_before_a_possible_charge(
self, failure_phase: str, collector_enabled: bool
self, failure_phase: str, collector_enabled: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
import base64
from litellm.proxy import proxy_server
monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False)
from datetime import datetime
import httpx

34
tests/unit/proxy/db/test_db_spend_update_writer.py
View file

@ -44,6 +44,15 @@ from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpush
from tests.unit.proxy.db.fake_prisma_engine import engine_call
@pytest.fixture(autouse=True)
def legacy_spend_writers(monkeypatch: pytest.MonkeyPatch) -> None:
"""These tests cover the legacy writers, still reachable as the rollback path with use_spend_logs_ledger off."""
from litellm.proxy import proxy_server
monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False)
@pytest.mark.asyncio
async def test_daily_spend_tracking_with_disabled_spend_logs():
"""
@ -3714,11 +3723,12 @@ async def test_failed_per_entity_increment_from_redis_restores_only_what_may_sti
proxy_logging_obj=proxy_logging_obj,
)
mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once()
restored = mock_redis_update_buffer.restore_transactions_to_redis.call_args.kwargs["db_spend_update_transactions"]
assert restored["user_list_transactions"] is None
assert restored["team_list_transactions"] == {"team-1": 1.5}
assert restored["key_list_transactions"] == ({"key-1": 1.5} if safe_to_resend else None)
restores = mock_redis_update_buffer.restore_transactions_to_redis.await_args_list
assert [
{name: rows for name, rows in call.kwargs["db_spend_update_transactions"].items() if rows}
for call in restores
] == ([{"key_list_transactions": {"key-1": 1.5}}] if safe_to_resend else [])
mock_batcher.litellm_teamtable.update_many.assert_called_once()
mock_batcher.litellm_usertable.update_many.assert_called_once()
@ -3892,13 +3902,13 @@ async def test_commit_spend_updates_to_db_reports_each_completed_table():
)
assert on_table_committed.call_args_list == [
call("user_list_transactions"),
call("end_user_list_transactions"),
call("key_list_transactions"),
call("team_list_transactions"),
call("team_member_list_transactions"),
call("org_list_transactions"),
call("org_member_list_transactions"),
call("team_list_transactions"),
call("user_list_transactions"),
call("end_user_list_transactions"),
call("team_member_list_transactions"),
call("key_list_transactions"),
call("project_list_transactions"),
call("tag_list_transactions"),
call("model_access_group_list_transactions"),
@ -5808,7 +5818,9 @@ async def test_completed_spend_survives_error_log_suppression(
prisma.get_request_status.return_value = status
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "disable_spend_logs", False)
monkeypatch.setattr(proxy_server, "general_settings", {"disable_error_logs": disable_error_logs})
monkeypatch.setattr(
proxy_server, "general_settings", {"disable_error_logs": disable_error_logs, "use_spend_logs_ledger": False}
)
now: Final = datetime(2026, 10, 20, 12, tzinfo=timezone.utc)
result: Final = ModelResponse(
model="gpt-4o-mini", usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)

5
tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -2851,9 +2851,12 @@ async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_requ
live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch
):
"""The window row's one-time seed from the spend log stops at the request start, which keeps the failure row
it already logged from being counted twice."""
it already logged from being counted twice. Only the legacy writers queue this row, so the ledger is off."""
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import proxy_logging_obj
monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False)
window_queue: Final = proxy_logging_obj.db_spend_update_writer.window_spend_update_queue
await window_queue.flush_and_get_aggregated_window_spend_transactions()
key_cache: Final = DualCache()

12
tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -3,7 +3,7 @@ import hashlib
import json
import logging
from collections.abc import AsyncIterator, Mapping, Sequence
from contextlib import asynccontextmanager
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import TYPE_CHECKING, Final, cast
@ -130,6 +130,11 @@ async def test_ui_view_users_proxy_admin_no_org_filter(mocker):
UserWhereCondition = InsensitiveContains | Sequence[Mapping[str, InsensitiveContains]]
@asynccontextmanager
async def _transaction_over(db: object) -> AsyncIterator[SimpleNamespace]:
yield SimpleNamespace(query_raw=AsyncMock(return_value=[]), litellm_usertable=getattr(db, "litellm_usertable"))
def _matches_user_where(row: LiteLLM_UserTableFiltered, where: Mapping[str, UserWhereCondition]) -> bool:
def matches(field: str, condition: UserWhereCondition) -> bool:
if not isinstance(condition, Mapping):
@ -2419,6 +2424,7 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke
prisma_client: Final = mocker.MagicMock()
prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user])
prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1)
prisma_client.tx = lambda: _transaction_over(prisma_client.db)
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
@ -2465,6 +2471,7 @@ async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker:
prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user)
prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user])
prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1)
prisma_client.tx = lambda: _transaction_over(prisma_client.db)
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
@ -2575,6 +2582,9 @@ async def test_user_update_reaches_cached_user_on_every_worker(
self.update_data_payload = data
return {"user_id": user_id, "data": self.updated_user}
def tx(self) -> AbstractAsyncContextManager[SimpleNamespace]:
return _transaction_over(self.db)
saved_user: Final = LiteLLM_UserTable(
user_id="user-spruce",
tpm_limit=100000,

50
tests/unit/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -8567,12 +8567,19 @@ class TestAccessGroupModelSync:
return [{"deployment_count": deployment_count}]
if sql.startswith('UPDATE "LiteLLM_AccessGroupTable"'):
return [{"access_group_id": "ag-1"}]
if sql.startswith(("SELECT 1 AS locked FROM", 'SELECT team_id FROM "LiteLLM_TeamTable"')):
return []
assert sql.startswith("WITH ")
return TestAccessGroupModelSync._ALLOWLIST_ROWS
@contextlib.asynccontextmanager
async def transaction():
yield mock_prisma.db
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.query_raw = AsyncMock(side_effect=query_raw)
mock_prisma.db.tx = MagicMock(side_effect=transaction)
mock_prisma.db.litellm_proxymodeltable = AsyncMock()
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=row)
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=row)
@ -11420,3 +11427,46 @@ async def test_add_new_model_reports_a_database_outage_during_the_insert_as_503(
)
assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection")
@pytest.mark.asyncio
async def test_rename_rewrites_only_in_a_transaction_holding_every_locked_teams_advisory_lock():
"""A team that starts listing the old name between the advisory snapshot and the row lock must not be
rewritten without its advisory lock, or a roster change on that team can still deadlock with the rename."""
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model
snapshots = iter(([{"team_id": "team-1"}], [{"team_id": "team-1"}, {"team_id": "team-2"}]))
transactions: list[list[tuple[str, tuple[str, ...]]]] = []
async def query_raw(sql: str, *params: str) -> list[dict[str, object]]:
if transactions:
transactions[-1].append((sql, params))
if sql.startswith("SELECT COUNT(*)"):
return [{"deployment_count": 0}]
if sql.startswith('SELECT team_id FROM "LiteLLM_TeamTable"') and not sql.endswith("FOR UPDATE"):
return next(snapshots)
if sql.startswith('SELECT team_id FROM "LiteLLM_TeamTable"'):
return [{"team_id": "team-1"}, {"team_id": "team-2"}]
return []
@contextlib.asynccontextmanager
async def transaction():
transactions.append([])
yield db
db = SimpleNamespace(query_raw=query_raw, tx=transaction)
with patch("litellm.proxy.management_helpers.model_allowlist_rename_sync.evict_and_broadcast", AsyncMock()):
await sync_model_allowlists_for_renamed_model(
SimpleNamespace(db=db),
model_id="deployment",
old_name="old-model",
new_name="new-model",
llm_router=None,
user_api_key_cache=UserApiKeyCache(),
)
rewrites = [calls for calls in transactions if any(sql.startswith("WITH ") for sql, _ in calls)]
assert len(rewrites) == 1
advisory = [params[0] for sql, params in rewrites[0] if "pg_advisory_xact_lock" in sql]
assert advisory == ["team-1", "team-2"]

7
tests/unit/proxy/management_endpoints/test_organization_endpoints.py
View file

@ -1,5 +1,6 @@
import asyncio
import json
from contextlib import asynccontextmanager
from types import MappingProxyType, SimpleNamespace
from typing import Final, Mapping, Optional, cast
from unittest.mock import AsyncMock, MagicMock, patch
@ -1560,6 +1561,12 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke
prisma_client.db.litellm_jwtkeymapping = jwt_table
prisma_client.db.litellm_organizationtable.delete = AsyncMock(return_value=MagicMock())
@asynccontextmanager
async def transaction():
yield prisma_client.db
prisma_client.tx = transaction
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)

5
tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -3819,7 +3819,10 @@ class TestSpendLogsPayload:
await _wait_for_mock_call(mock_client)
kwargs = mock_client.call_args.kwargs
payload: SpendLogsPayload = kwargs["payload"]
payload: SpendLogsPayload = {
key: value for key, value in kwargs["payload"].items() if not key.startswith("accounting_")
}
assert kwargs["payload"]["accounting_payload"]["spend"] == kwargs["payload"]["spend"]
expected_payload = SpendLogsPayload(
**{
"request_id": "chatcmpl-34df56d5-4807-45c1-bb99-61e52586b802",

5
tests/unit/proxy/test_unit_test_proxy_hooks.py
View file

@ -11,7 +11,7 @@ import litellm
@pytest.mark.asyncio
async def test_disable_spend_logs():
"""
Test that the spend logs are not written to the database when disable_spend_logs is True
With disable_spend_logs the request is queued only for its charge, flagged so no SpendLogs row is written
"""
# Mock the necessary components
@ -41,5 +41,4 @@ async def test_disable_spend_logs():
org_id="org_id",
kwargs={},
)
# Verify no spend logs were added
assert len(mock_prisma_client.spend_log_transactions) == 0
assert [row["accounting_log_disabled"] for row in mock_prisma_client.spend_log_transactions] == [True]

4
tests/unit/proxy/utils/prisma_and_spend/conftest.py
View file

@ -142,7 +142,9 @@ def mock_prisma_client() -> MagicMock:
client.db.stop_token_refresh_task = AsyncMock()
client.db.query_raw = AsyncMock(return_value=[{"?column?": 1}])
client.db.execute_raw = AsyncMock()
client.db.tx = MagicMock()
client.db.tx = MagicMock(
return_value=MagicMock(__aenter__=AsyncMock(return_value=client.db), __aexit__=AsyncMock(return_value=False))
)
client.db.batch_ = MagicMock()
for table_name in _PRISMA_TABLES:
setattr(client.db, table_name, _make_table_mock())

9
tests/unit/repositories/test_unit_of_work.py
View file

@ -168,7 +168,8 @@ async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch(
return batches[-1]
async with budget_cascade_unit_of_work(_new_batch) as uow:
writes = [getattr(uow, field.name) for field in fields(uow)]
uow.row_locks.queue_row_locks(["budget-1"])
writes = [getattr(uow, field.name) for field in fields(uow) if field.name != "row_locks"]
for write in writes:
if isinstance(write, SpendResetWrites):
write.queue_spend_zero(where={"budget_id": "budget-1"})
@ -177,8 +178,10 @@ async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch(
assert len(batches) == 1, "the cascade must open exactly one transaction"
batch = batches[0]
assert len(batch.calls) == len(writes), "a dependent bound to a batch of its own would not land here"
assert len({call[0] for call in batch.calls}) == len(writes), "two dependents share one table"
table_writes = [call for call in batch.calls if call[0] != "execute_raw"]
assert len(table_writes) == len(writes), "a dependent bound to a batch of its own would not land here"
assert len({call[0] for call in table_writes}) == len(writes), "two dependents share one table"
assert [call[0] for call in batch.calls].count("execute_raw") > 0, "the row locks share the cascade's batch"
assert batch.commit_count == 1

4
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -31345,8 +31345,8 @@ export interface components {
use_google_kms?: boolean | null;
/**
* Use Spend Logs Ledger
* @description Apply request accounting from stored SpendLogs deliveries. Enable only after all accounting writers are upgraded and ledger activation gates pass.
* @default false
* @description Apply request accounting from stored SpendLogs deliveries. Set to false to roll back to the previous in-memory and Redis spend queues.
* @default true
*/
use_spend_logs_ledger: boolean;
/**