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] 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}