scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Merge 6fcd1b5d3b into 3f4dc93b40

Browse source
This commit is contained in:
joshua-berri 2026-11-05 07:30:48 +00:00 • committed by GitHub
commit 240aa70de4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
31 changed files with 682 additions and 291 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,
)
)

2
litellm/accounting/worker.py
View file

@ -54,7 +54,7 @@ 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:
result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only)

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

@ -2489,8 +2489,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)

260
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
@ -2294,6 +2312,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 +2335,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 +2479,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 +2529,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 +2563,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

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

@ -13,7 +13,7 @@ from typing import Final
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.router import Router
from litellm.types.llms.base import LiteLLMBaseModel
@ -80,6 +80,23 @@ def _rewrite_sql(set_clause: str, where_clause: str) -> str:
return f"WITH {ctes} {rows}"
_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 1 AS locked 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"))')
@ -98,9 +115,15 @@ 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
)
async with raw_transactions(prisma_client).tx() as transaction:
_ = await transaction.query_raw(_LOCK_ORGS_LISTING_SQL, old_name)
_ = await transaction.query_raw(_LOCK_TEAMS_LISTING_SQL, old_name)
_ = 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)
touched_rows: Final = await transaction.query_raw(
_APPEND_SQL if old_name_still_backed else _REPLACE_SQL, old_name, new_name
)
touched: Final = tuple(_TouchedRow.model_validate(row) for row in touched_rows)
await evict_and_broadcast(
tuple(cache_key for row in touched for cache_key in _CACHE_KEYS_BY_KIND[row.kind](row)),

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

@ -5528,6 +5528,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()
@ -6828,24 +6842,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
),

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

@ -0,0 +1,160 @@
import asyncio
import subprocess
import sys
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
from pathlib import Path
from typing import Final
import psycopg
import pytest
import pytest_asyncio
from prisma import Prisma
from litellm.proxy._types import DBSpendUpdateTransactions
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
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)

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() == (

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

@ -63,7 +63,7 @@ _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,expected,distinct_models", [(None, 251, 251, False), (None, 501, 500, False), ("200", 251, 200, False), ("10000", 10001, 10000, False), ("3001", 3001, 3001, 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

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,

7
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"):
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)

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]

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

@ -31080,8 +31080,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;
/**