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}