From 0fa9ac3ce0a14c7775201593531dbec016401d98 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 19:18:49 -0800 Subject: [PATCH 1/4] 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 --- litellm/accounting/snapshot.py | 2 + litellm/litellm_core_utils/safe_json_dumps.py | 11 +++- .../accounting_delivery_repository.py | 50 +++++++++------ .../database/test_spendlogs_delivery.py | 63 ++++++++++++++----- tests/unit/accounting/test_projection.py | 12 ++-- .../test_safe_json_dumps.py | 10 ++- 6 files changed, 106 insertions(+), 42 deletions(-) diff --git a/litellm/accounting/snapshot.py b/litellm/accounting/snapshot.py index d5652f79ed0..ab85df08320 100644 --- a/litellm/accounting/snapshot.py +++ b/litellm/accounting/snapshot.py @@ -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, ) ) diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 909e1222ae6..ea26d549326 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -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" diff --git a/litellm/repositories/accounting_delivery_repository.py b/litellm/repositories/accounting_delivery_repository.py index 7dd92aecb2d..262992a48c5 100644 --- a/litellm/repositories/accounting_delivery_repository.py +++ b/litellm/repositories/accounting_delivery_repository.py @@ -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""" @@ -82,7 +87,17 @@ SET attempts = attempts + 1, last_error = $3, next_attempt_at = CURRENT_TIMESTAMP + make_interval(secs => LEAST(300, power(2, LEAST(attempts, 8)))::integer), 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 diff --git a/tests/integration/database/test_spendlogs_delivery.py b/tests/integration/database/test_spendlogs_delivery.py index 19f01a6b902..d1227e85565 100644 --- a/tests/integration/database/test_spendlogs_delivery.py +++ b/tests/integration/database/test_spendlogs_delivery.py @@ -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() == ( diff --git a/tests/unit/accounting/test_projection.py b/tests/unit/accounting/test_projection.py index 17186a89922..048d8b5f01a 100644 --- a/tests/unit/accounting/test_projection.py +++ b/tests/unit/accounting/test_projection.py @@ -521,7 +521,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 +554,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] diff --git a/tests/unit/litellm_core_utils/test_safe_json_dumps.py b/tests/unit/litellm_core_utils/test_safe_json_dumps.py index 1b8bdd21326..740ce4f433a 100644 --- a/tests/unit/litellm_core_utils/test_safe_json_dumps.py +++ b/tests/unit/litellm_core_utils/test_safe_json_dumps.py @@ -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} From bbb85ed48e361c094c2623c45bffe2f395f152c9 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 23:30:32 -0800 Subject: [PATCH 2/4] 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 --- litellm/accounting/application.py | 2 + litellm/accounting/worker.py | 2 +- .../proxy/common_utils/reset_budget_job.py | 5 +- litellm/proxy/db/db_spend_update_writer.py | 258 ++++++++++-------- .../access_group_endpoints.py | 91 +++--- .../internal_user_endpoints.py | 24 +- .../organization_endpoints.py | 29 +- .../access_group_model_sync.py | 10 + .../model_allowlist_rename_sync.py | 31 ++- .../spend_management_endpoints.py | 28 +- litellm/proxy/utils.py | 38 ++- .../accounting_delivery_repository.py | 2 +- litellm/repositories/unit_of_work.py | 42 +++ .../database/test_spendlogs_ledger_parity.py | 2 +- .../proxy/db/test_db_spend_update_writer.py | 34 ++- .../test_internal_user_endpoints.py | 12 +- .../test_model_management_endpoints.py | 7 + .../test_organization_endpoints.py | 7 + tests/unit/repositories/test_unit_of_work.py | 9 +- 19 files changed, 395 insertions(+), 238 deletions(-) diff --git a/litellm/accounting/application.py b/litellm/accounting/application.py index 1e5f23e2997..c911f4e35ce 100644 --- a/litellm/accounting/application.py +++ b/litellm/accounting/application.py @@ -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) ) diff --git a/litellm/accounting/worker.py b/litellm/accounting/worker.py index df71f2da31c..0607c1647e1 100644 --- a/litellm/accounting/worker.py +++ b/litellm/accounting/worker.py @@ -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) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 59fd00cec43..9c433af487e 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -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) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 80f09df6c86..ba59b971a24 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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::::user_id::"; 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()) @@ -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") diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 6558c82f80c..4cbc2155af0 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -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) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index d2aac9cae38..93a429a2fcc 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -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( diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 8a1b9c112dc..4e511aee24a 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -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, diff --git a/litellm/proxy/management_helpers/access_group_model_sync.py b/litellm/proxy/management_helpers/access_group_model_sync.py index 890715bf9c2..20ea2ed0338 100644 --- a/litellm/proxy/management_helpers/access_group_model_sync.py +++ b/litellm/proxy/management_helpers/access_group_model_sync.py @@ -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 diff --git a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py index f1a9faf6b2e..498318cb63d 100644 --- a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py +++ b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py @@ -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)), diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index fcb9eb4a280..0348d8e1f82 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -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", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 71ecae74058..19f6bff8a8a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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 diff --git a/litellm/repositories/accounting_delivery_repository.py b/litellm/repositories/accounting_delivery_repository.py index 262992a48c5..d5bbe891089 100644 --- a/litellm/repositories/accounting_delivery_repository.py +++ b/litellm/repositories/accounting_delivery_repository.py @@ -84,7 +84,7 @@ 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""" diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index f6aad883479..b78f7e5a9b7 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -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 ), diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index e1894aa66bf..78439e4317e 100644 --- a/tests/integration/database/test_spendlogs_ledger_parity.py +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -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 diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index d90c7f6d21d..e2debe05c03 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -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) diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index d8aa0239e2d..cee0c90f80c 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -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, diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 44c43713438..303b0b5e83f 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -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) diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 5e06756f3a9..8d2c6c0054c 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -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) diff --git a/tests/unit/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py index 082a29b7001..049e5275db6 100644 --- a/tests/unit/repositories/test_unit_of_work.py +++ b/tests/unit/repositories/test_unit_of_work.py @@ -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 From 32d53a959bbef44b4276fad1d11240a198e39910 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 23:30:32 -0800 Subject: [PATCH 3/4] 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 --- .../database/test_spend_lock_order.py | 160 ++++++++++++++++++ 1 file changed, 160 insertions(+) create mode 100644 tests/integration/database/test_spend_lock_order.py diff --git a/tests/integration/database/test_spend_lock_order.py b/tests/integration/database/test_spend_lock_order.py new file mode 100644 index 00000000000..00248d99b39 --- /dev/null +++ b/tests/integration/database/test_spend_lock_order.py @@ -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) From 6fcd1b5d3b3b795805547eb6c69966032a5049ec Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 23:30:32 -0800 Subject: [PATCH 4/4] 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 --- litellm/proxy/_types.py | 4 ++-- litellm/proxy/db/db_spend_update_writer.py | 2 +- tests/unit/accounting/test_projection.py | 1 + tests/unit/proxy/common_utils/test_check_batch_cost.py | 6 +++++- tests/unit/proxy/hooks/test_proxy_track_cost_callback.py | 5 ++++- .../proxy/spend_tracking/test_spend_management_endpoints.py | 5 ++++- tests/unit/proxy/test_unit_test_proxy_hooks.py | 5 ++--- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 ++-- 8 files changed, 21 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 25849634764..d7df664f6a9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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( diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ba59b971a24..ba9cae72cf7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -730,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 diff --git a/tests/unit/accounting/test_projection.py b/tests/unit/accounting/test_projection.py index 048d8b5f01a..c78d6bec757 100644 --- a/tests/unit/accounting/test_projection.py +++ b/tests/unit/accounting/test_projection.py @@ -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() diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 92114b00398..eee5d96511b 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -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 diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index c6366a40504..0c7292674a6 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -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() diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index e119acd5ffc..c3e916d581d 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -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", diff --git a/tests/unit/proxy/test_unit_test_proxy_hooks.py b/tests/unit/proxy/test_unit_test_proxy_hooks.py index 6ce7154f2d1..7990ec12037 100644 --- a/tests/unit/proxy/test_unit_test_proxy_hooks.py +++ b/tests/unit/proxy/test_unit_test_proxy_hooks.py @@ -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] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 11ce7f97ad8..2898bd89832 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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; /**