From d30c702c5cfe98d0f2c6832ce5ef5bb0ca8019de Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Wed, 28 Oct 2026 16:11:51 -0700 Subject: [PATCH] feat(accounting): add durable event ledger writer (#48883) * feat(accounting): add durable event ledger writer * fix(accounting): validate delegated transaction managers locally * docs(accounting): clarify ledger integration contracts * fix(accounting): return ledger conflicts as values and keep cleanup on the kind index Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(accounting): import assert_never from typing_extensions for Python 3.10 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(accounting): annotate ledger integration test locals as Final Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(accounting): finish ledger test conventions --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/accounting_ledger.py | 269 ++++++++++++ .../database/test_accounting_ledger.py | 412 +++++++++++++++++- tests/unit/proxy/db/test_accounting_ledger.py | 103 +++++ 3 files changed, 781 insertions(+), 3 deletions(-) create mode 100644 litellm/proxy/db/accounting_ledger.py create mode 100644 tests/unit/proxy/db/test_accounting_ledger.py diff --git a/litellm/proxy/db/accounting_ledger.py b/litellm/proxy/db/accounting_ledger.py new file mode 100644 index 00000000000..91f8307283e --- /dev/null +++ b/litellm/proxy/db/accounting_ledger.py @@ -0,0 +1,269 @@ +from __future__ import annotations + +import hashlib +import json +from collections.abc import AsyncGenerator, Callable +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from pydantic import AwareDatetime, ConfigDict, JsonValue, TypeAdapter +from typing_extensions import assert_never + +from litellm.accounting import CostEvent, DeliveryTarget, EventKind, InboundUsageEvent, OutboundUsageEvent +from litellm.proxy.db.db_span import db_span +from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper + +if TYPE_CHECKING: + from prisma import Prisma + + from litellm.proxy.utils import PrismaClient + +AccountingEvent: TypeAlias = InboundUsageEvent | OutboundUsageEvent | CostEvent +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_INSTANT: Final[TypeAdapter[datetime]] = TypeAdapter(AwareDatetime) +_STRINGS: Final = TypeAdapter(tuple[dict[str, str], ...]) +_TIMEOUT: Final = timedelta(seconds=10) + + +class AccountingConflictError(ValueError): + pass + + +_Transaction: TypeAlias = AbstractAsyncContextManager["Prisma"] + + +@dataclass(frozen=True, slots=True) +class _AcceptedEvent: + event_id: str + kind: EventKind + source_event_id: str | None + occurred_at: datetime + payload: str + payload_hash: str + targets: tuple[str, ...] + schema_version: int + version: int + + +def _json(value: object) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def _prepare( + event: AccountingEvent, + targets: tuple[DeliveryTarget, ...], + envelope: str, + schema_version: int, + version: int, +) -> _AcceptedEvent: + if type(schema_version) is not int or type(version) is not int or min(schema_version, version) < 1: + raise ValueError("Accounting versions must be positive integers") + validated: Final = type(event).model_validate_json(event.model_dump_json()) + payload: Final = _JSON.validate_json(envelope) + if payload.get("event") != validated.model_dump(mode="json"): + raise ValueError("Accounting envelope must contain the matching public event") + destinations: Final = tuple(sorted(frozenset(DeliveryTarget(target).value for target in targets))) + serialized: Final = _json(payload) + identity: Final = _json((schema_version, version, destinations, payload)) + kind: Final = ( + EventKind.COST + if isinstance(validated, CostEvent) + else EventKind.INBOUND_USAGE + if isinstance(validated, InboundUsageEvent) + else EventKind.OUTBOUND_USAGE + ) + return _AcceptedEvent( + event_id=validated.event_id, + kind=kind, + source_event_id=validated.usage_event_id if isinstance(validated, CostEvent) else None, + occurred_at=validated.occurred_at if isinstance(validated, CostEvent) else validated.ended_at, + payload=serialized, + payload_hash=hashlib.sha256(identity.encode()).hexdigest(), + targets=destinations, + schema_version=schema_version, + version=version, + ) + + +_INSERT_EVENT: Final = """ +INSERT INTO "LiteLLM_AccountingEvent" + (event_id, kind, source_event_id, schema_version, payload_hash, occurred_at, payload) +VALUES ($1, $2::"AccountingEventKind", $3, $4::int, $5, $6::timestamp, $7::jsonb) +ON CONFLICT (event_id) DO NOTHING RETURNING event_id +""" +_INSERT_DELIVERIES: Final = """ +INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) +SELECT $1, value::"AccountingDeliveryTarget", $3::int FROM jsonb_array_elements_text($2::jsonb) +""" +_LOCK_SOURCE: Final = """ +SELECT event_id FROM "LiteLLM_AccountingEvent" +WHERE event_id = $1 AND kind = 'outbound_usage' FOR UPDATE +""" +_EXISTING_COST: Final = """ +SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE kind = 'cost' AND source_event_id = $1 +""" +_LOCK_EVENT: Final = """ +SELECT payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 FOR UPDATE +""" +_COMPLETED_FAMILY: Final = """ +root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage') +AND root.recorded_at < $1::timestamp +AND (root.kind = 'inbound_usage' OR EXISTS ( + SELECT 1 FROM "LiteLLM_AccountingEvent" AS cost WHERE cost.source_event_id = root.event_id AND cost.kind = 'cost' +)) +AND NOT EXISTS ( + SELECT 1 FROM "LiteLLM_AccountingEvent" AS member + WHERE (member.event_id = root.event_id OR (member.kind = 'cost' AND member.source_event_id = root.event_id)) + AND (member.recorded_at >= $1::timestamp + OR EXISTS (SELECT 1 FROM "LiteLLM_AccountingDelivery" AS delivery + WHERE delivery.event_id = member.event_id AND delivery.state != 'applied')) +) +""" +_RETENTION_CANDIDATES: Final = ( + 'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE ' + + _COMPLETED_FAMILY + + " ORDER BY root.recorded_at, root.event_id LIMIT $2::int FOR UPDATE OF root SKIP LOCKED" +) +_ELIGIBLE_GROUP: Final = ( + 'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE root.event_id = $2 AND ' + _COMPLETED_FAMILY +) +_DELETE_DELIVERIES: Final = """ +DELETE FROM "LiteLLM_AccountingDelivery" WHERE event_id IN ( + SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1) +) +""" +_DELETE_EVENTS: Final = """ +DELETE FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1) +""" + + +_Conflict: TypeAlias = Literal["missing_source", "competing_cost", "identity_conflict"] + + +async def _accept(db: Prisma, record: _AcceptedEvent) -> Literal["inserted", "duplicate"] | _Conflict: + if record.source_event_id is not None: + source: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_SOURCE, record.source_event_id)) + if not source: + return "missing_source" + costs: Final = _STRINGS.validate_python(tuple(await db.query_raw(_EXISTING_COST, record.source_event_id))) + if any(cost["event_id"] != record.event_id for cost in costs): + return "competing_cost" + inserted: Final = _STRINGS.validate_python( + await db.query_raw( + _INSERT_EVENT, + record.event_id, + record.kind.value, + record.source_event_id, + record.schema_version, + record.payload_hash, + record.occurred_at.astimezone(timezone.utc).isoformat(), + record.payload, + ) + ) + if not inserted: + existing: Final = _STRINGS.validate_python(tuple(await db.query_raw(_LOCK_EVENT, record.event_id))) + if not existing or existing[0]["payload_hash"] != record.payload_hash: + return "identity_conflict" + return "duplicate" + await db.execute_raw(_INSERT_DELIVERIES, record.event_id, _json(record.targets), record.version) + return "inserted" + + +async def _delete_completed(db: Prisma, event_id: str, cutoff: str) -> int: + if not _STRINGS.validate_python(await db.query_raw(_ELIGIBLE_GROUP, cutoff, event_id)): + return 0 + await db.execute_raw(_DELETE_DELIVERIES, event_id) + return await db.execute_raw(_DELETE_EVENTS, event_id) + + +class AccountingLedger: + def __init__(self, transaction: Callable[[], _Transaction]) -> None: + self._transaction: Final = transaction + + @classmethod + def for_client(cls, client: PrismaClient) -> AccountingLedger: + from prisma.client import TransactionManager + + managers: Final = TypeAdapter(TransactionManager, config=ConfigDict(arbitrary_types_allowed=True)) + + @asynccontextmanager + async def transaction() -> AsyncGenerator[Prisma]: + primary: Final = writer_wrapper(client.db) + manager: Final = managers.validate_python(primary.tx(timeout=_TIMEOUT)) + async with db_span("accounting_ledger", "LiteLLM_AccountingEvent"), manager as db: + await db.execute_raw("SET LOCAL statement_timeout = 5000") + await db.execute_raw("SET LOCAL lock_timeout = 1000") + yield db + + return cls(transaction) + + async def insert_event( + self, + event: AccountingEvent, + targets: tuple[DeliveryTarget, ...], + *, + envelope: str, + schema_version: int = 1, + version: int = 1, + ) -> Literal["inserted", "duplicate"]: + """Accept T5/T6's serialized envelope, containing its public event under `event`. + + The immutable acceptance hash includes the target set and delivery version. + Changing either for an accepted event conflicts. T7 must provide a separate + add-target or new-version workflow that preserves receipts and prevents + duplicate contributions; redelivery cannot silently create or reset work. + + The Prisma adapter opens one interactive transaction per event, including + two SET LOCAL statements and a source-row lock for costs. I/O is async, + but awaiting this method inline adds commit latency per event, potentially + several times per request. T5 owns bounded asynchronous capture and an + explicit batching strategy; looping over this method does not batch commits. + Success acknowledges commit, independently of HTTP response acknowledgement. + """ + record: Final = _prepare(event, targets, envelope, schema_version, version) + async with self._transaction() as db: + result: Final = await _accept(db, record) + match result: + case "inserted" | "duplicate": + return result + case "missing_source": + raise AccountingConflictError("Cost source must be a durably accepted outbound usage event") + case "competing_cost": + raise AccountingConflictError("Usage event already has an authoritative cost") + case "identity_conflict": + raise AccountingConflictError("Event identity conflicts with its accepted content or delivery set") + case _: + assert_never(result) + + async def purge_completed( + self, + *, + redelivery_before: datetime | None = None, + recovery_before: datetime | None = None, + limit: int = 100, + ) -> int: + """Delete at most `limit` completed families once both retention owners permit it. + + T8 must supply a coherent recovery watermark including outstanding holds. + Omitting either watermark disables deletion. Each usage row and its cost + remain together so a retained charge never loses its source or vice versa. + Applied receipts are deleted with the family, ending its retained dedup + protection. The caller must guarantee events recorded before + redelivery_before cannot be redelivered through supported paths. + T8's recovery_before must also protect checkpoints, + outstanding holds and unfinished settlements; age alone is insufficient. + """ + if type(limit) is not int or not 1 <= limit <= 1000: + raise ValueError("Accounting retention batch size must be between 1 and 1000") + if redelivery_before is None or recovery_before is None: + return 0 + cutoff: Final = min( + _INSTANT.validate_python(redelivery_before).astimezone(timezone.utc), + _INSTANT.validate_python(recovery_before).astimezone(timezone.utc), + ).isoformat() + async with self._transaction() as db: + roots: Final = _STRINGS.validate_python(tuple(await db.query_raw(_RETENTION_CANDIDATES, cutoff, limit))) + deleted: Final = tuple([await _delete_completed(db, root["event_id"], cutoff) for root in roots]) + return sum(deleted) diff --git a/tests/integration/database/test_accounting_ledger.py b/tests/integration/database/test_accounting_ledger.py index 089ed872dee..0c2ceb62eaf 100644 --- a/tests/integration/database/test_accounting_ledger.py +++ b/tests/integration/database/test_accounting_ledger.py @@ -1,16 +1,48 @@ +from __future__ import annotations + +import asyncio +import json import os import subprocess import sys -from collections.abc import Iterator +from collections.abc import AsyncGenerator, AsyncIterator, Iterator, Mapping, Sequence +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from decimal import Decimal +from itertools import chain from pathlib import Path -from typing import Final +from typing import TYPE_CHECKING, Final, TypeAlias import psycopg import pytest +import pytest_asyncio from psycopg import sql +from pydantic import JsonValue -from litellm.accounting import DeliveryState, DeliveryTarget, EventKind +from litellm.accounting import ( + CostEvent, + DeliveryState, + DeliveryTarget, + EventKind, + InboundUsageEvent, + Metric, + MetricName, + OutboundUsageEvent, + Outcome, + PricingSnapshot, +) +from litellm.proxy.db import accounting_ledger +from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger +from litellm.types.utils import CallTypes from tests.integration._support.database import scratch_database +from tests.integration._support.prisma_clients import connected_proxy_prisma + +if TYPE_CHECKING: + from prisma import Prisma + + from litellm.proxy.utils import PrismaClient + +_Writer: TypeAlias = tuple["PrismaClient", AccountingLedger] _SCHEMA: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/schema.prisma" _EVENT_INSERT: Final = ( @@ -126,3 +158,377 @@ def test_new_delivery_defaults_to_pending_with_no_attempts(ledger: psycopg.Conne (DeliveryTarget.KEY_SPEND.value,), ).fetchone() assert row == ("pending", 0) + + +_AT: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) +_CUTOFF: Final = datetime(2026, 2, 1, tzinfo=timezone.utc) +_TARGETS: Final = (DeliveryTarget.USER_SPEND, DeliveryTarget.TEAM_SPEND) + + +def _usage(event_id: str = "usage", *, inbound: bool = False) -> InboundUsageEvent | OutboundUsageEvent: + fields: Final = dict( + event_id=event_id, + request_id="request", + call_type=CallTypes.completion, + started_at=_AT, + ended_at=_AT, + outcome=Outcome.SUCCEEDED, + dimensions=(), + ) + if inbound: + return InboundUsageEvent(**fields, metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),)) + return OutboundUsageEvent(**fields, attempt_id=event_id, retry_of_attempt_id=None, metrics=None) + + +def _cost(source: str = "usage", event_id: str = "cost", amount: str = "0.125") -> CostEvent: + return CostEvent( + event_id=event_id, + usage_event_id=source, + occurred_at=_AT, + dimensions=(), + charges=(Metric(name=MetricName.SPEND, value=Decimal(amount), unit="USD"),), + snapshot=PricingSnapshot( + calculator_version="fixture", + model="fixture", + provider=None, + pricing_source="cost_map", + priced_model_key="fixture", + total_cost=Decimal(amount), + ), + ) + + +def _envelope(event: AccountingEvent, *, context: Mapping[str, JsonValue] | None = None) -> str: + return json.dumps({"event": event.model_dump(mode="json"), "context": context or {}}) + + +async def _insert( + store: AccountingLedger, + event: AccountingEvent, + targets: tuple[DeliveryTarget, ...] = _TARGETS, + *, + context: Mapping[str, JsonValue] | None = None, +) -> str: + return await store.insert_event(event, targets, envelope=_envelope(event, context=context)) + + +def _rows(database_url: str, statement: str, parameters: tuple[object, ...] = ()) -> Sequence[tuple[object, ...]]: + with psycopg.connect(database_url) as connection: + cursor: Final = connection.execute(statement, parameters) + return cursor.fetchall() if cursor.description is not None else [] + + +def _complete_and_age(database_url: str) -> None: + with psycopg.connect(database_url) as connection: + connection.execute('UPDATE "LiteLLM_AccountingEvent" SET recorded_at = %s', (_AT,)) + connection.execute("UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'applied', applied_at = %s", (_AT,)) + + +@pytest_asyncio.fixture(loop_scope="function") +async def writer(ledger_database: str, monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[_Writer]: + monkeypatch.setenv("DATABASE_URL", ledger_database) + monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) + _rows(ledger_database, 'TRUNCATE "LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent"') + async with connected_proxy_prisma(ledger_database) as client: + yield client, AccountingLedger.for_client(client) + + +@pytest.mark.asyncio +async def test_event_and_complete_delivery_set_commit_together(writer: _Writer, ledger_database: str) -> None: + _, store = writer + event: Final = _usage() + assert await _insert(store, event, context={"requested_region": "eu", "quantity": "1.25"}) == "inserted" + rows: Final = _rows(ledger_database, 'SELECT kind::text, source_event_id, payload FROM "LiteLLM_AccountingEvent"') + assert rows == [ + ("outbound_usage", None, json.loads(_envelope(event, context={"requested_region": "eu", "quantity": "1.25"}))) + ] + deliveries: Final = _rows( + ledger_database, 'SELECT target::text, state::text, attempts FROM "LiteLLM_AccountingDelivery"' + ) + assert set(deliveries) == {(target.value, "pending", 0) for target in _TARGETS} + + +@pytest.mark.asyncio +async def test_concurrent_duplicates_preserve_completed_deliveries(writer: _Writer, ledger_database: str) -> None: + _, store = writer + event: Final = _usage(inbound=True) + results: Final = await asyncio.gather(*(_insert(store, event) for _ in range(8))) + assert results.count("inserted") == 1 and results.count("duplicate") == 7 + _complete_and_age(ledger_database) + before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') + reordered: Final = json.dumps(json.loads(_envelope(event)), sort_keys=True, indent=2) + assert await store.insert_event(event, tuple(reversed(_TARGETS)) + _TARGETS, envelope=reordered) == "duplicate" + assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') == before + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["payload", "context", "targets", "schema_version", "version"]) +async def test_conflicting_acceptance_cannot_overwrite_or_add_work( + writer: _Writer, ledger_database: str, change: str +) -> None: + _, store = writer + event: Final = _usage() + await _insert(store, event) + altered: Final = event.model_copy(update={"outcome": Outcome.FAILED}) if change == "payload" else event + targets: Final = (DeliveryTarget.KEY_SPEND,) if change == "targets" else _TARGETS + context: Final = {"override": "0.5"} if change == "context" else {} + options: Final = {change: 2} if change in ("schema_version", "version") else {} + with pytest.raises(AccountingConflictError, match="conflicts"): + await store.insert_event(altered, targets, envelope=_envelope(altered, context=context), **options) + assert _rows(ledger_database, 'SELECT payload FROM "LiteLLM_AccountingEvent"') == [(json.loads(_envelope(event)),)] + assert set(_rows(ledger_database, 'SELECT target::text FROM "LiteLLM_AccountingDelivery"')) == { + (t.value,) for t in _TARGETS + } + + +@pytest.mark.asyncio +async def test_competing_cost_identities_accept_only_one_authoritative_charge( + writer: _Writer, ledger_database: str +) -> None: + _, store = writer + await _insert(store, _usage()) + results: Final = await asyncio.gather( + _insert(store, _cost()), _insert(store, _cost(event_id="other-cost")), return_exceptions=True + ) + assert sum(result == "inserted" for result in results) == 1 + assert sum(isinstance(result, AccountingConflictError) for result in results) == 1 + accepted: Final = _rows( + ledger_database, "SELECT event_id, payload FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'" + ) + assert len(accepted) == 1 + event: Final = CostEvent.model_validate_json(json.dumps(accepted[0][1]["event"])) + assert await _insert(store, event) == "duplicate" + assert event.snapshot.total_cost == Decimal("0.125") + with pytest.raises(AccountingConflictError): + await _insert(store, _cost(event_id=event.event_id, amount="0.5")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("source", ["missing", "inbound", "cost"]) +async def test_cost_requires_an_existing_outbound_source(writer: _Writer, ledger_database: str, source: str) -> None: + _, store = writer + if source == "inbound": + await _insert(store, _usage("source", inbound=True)) + if source == "cost": + await _insert(store, _usage()) + await _insert(store, _cost(event_id="source")) + with pytest.raises(AccountingConflictError, match="outbound usage"): + await _insert(store, _cost(source="source", event_id="invalid")) + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'invalid'") == [ + (0,) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "cancellation"]) +async def test_interruption_rolls_back_or_retry_recognizes_commit( + writer: _Writer, ledger_database: str, failure: str +) -> None: + client, store = writer + + @asynccontextmanager + async def interrupted() -> AsyncGenerator[Prisma]: + async with client.writer_db.tx() as tx: + yield tx + if failure == "before_commit": + raise RuntimeError("injected pre-commit failure") + if failure == "cancellation": + raise asyncio.CancelledError + raise RuntimeError("injected lost acknowledgement") + + error: Final = asyncio.CancelledError if failure == "cancellation" else RuntimeError + with pytest.raises(error): + await _insert(AccountingLedger(interrupted), _usage()) + expected: Final = 1 if failure == "after_commit" else 0 + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(expected,)] + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(expected * len(_TARGETS),)] + assert await _insert(store, _usage()) == ("duplicate" if expected else "inserted") + + +@pytest.mark.asyncio +async def test_delivery_insert_failure_rolls_back_event(writer: _Writer, ledger_database: str) -> None: + from prisma.errors import RawQueryError + + _, store = writer + with psycopg.connect(ledger_database) as connection: + connection.execute( + 'ALTER TABLE "LiteLLM_AccountingDelivery" ADD CONSTRAINT reject_delivery CHECK (version < 1)' + ) + try: + with pytest.raises(RawQueryError, match="reject_delivery"): + await _insert(store, _usage()) + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] + finally: + with psycopg.connect(ledger_database) as connection: + connection.execute('ALTER TABLE "LiteLLM_AccountingDelivery" DROP CONSTRAINT reject_delivery') + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blocked", ["pending", "parked", "unpriced", "recent_cost", "recovery"]) +async def test_retention_preserves_incomplete_or_recovery_needed_families( + writer: _Writer, ledger_database: str, blocked: str +) -> None: + _, store = writer + await _insert(store, _usage()) + if blocked != "unpriced": + await _insert(store, _cost()) + _complete_and_age(ledger_database) + with psycopg.connect(ledger_database) as connection: + if blocked in ("pending", "parked"): + connection.execute( + "UPDATE \"LiteLLM_AccountingDelivery\" SET state = %s WHERE event_id = 'cost'", (blocked,) + ) + if blocked == "recent_cost": + connection.execute( + "UPDATE \"LiteLLM_AccountingEvent\" SET recorded_at = %s WHERE event_id = 'cost'", (_CUTOFF,) + ) + before: Final = _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') + assert ( + await store.purge_completed( + redelivery_before=_CUTOFF, recovery_before=_AT if blocked == "recovery" else _CUTOFF + ) + == 0 + ) + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == before + + +@pytest.mark.asyncio +async def test_retention_removes_completed_family_and_receipts_together(writer: _Writer, ledger_database: str) -> None: + _, store = writer + await _insert(store, _usage()) + await _insert(store, _cost()) + _complete_and_age(ledger_database) + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)] + + +@pytest.mark.asyncio +async def test_retention_respects_batch_bound_and_skips_locked_families(writer: _Writer, ledger_database: str) -> None: + client, store = writer + for name in ("a", "b", "c"): + await _insert(store, _usage(name, inbound=True)) + _complete_and_age(ledger_database) + async with client.writer_db.tx() as tx: + await tx.query_raw("SELECT event_id FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'a' FOR UPDATE") + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1 + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ + ("a",), + ("c",), + ] + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 + + +@pytest.mark.asyncio +async def test_retention_makes_progress_past_pending_and_unpriced_events(writer: _Writer, ledger_database: str) -> None: + _, store = writer + await _insert(store, _usage("a", inbound=True)) + await _insert(store, _usage("b")) + await _insert(store, _usage("c", inbound=True)) + _complete_and_age(ledger_database) + _rows(ledger_database, "UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'pending' WHERE event_id = 'a'") + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1 + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ + ("a",), + ("b",), + ] + + +@pytest.mark.asyncio +async def test_zero_cost_and_empty_delivery_sets_are_durable(writer: _Writer, ledger_database: str) -> None: + _, store = writer + await _insert(store, _usage(), targets=()) + event: Final = _cost(amount="0") + assert await _insert(store, event, targets=()) == "inserted" + assert await _insert(store, event, targets=()) == "duplicate" + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)] + _complete_and_age(ledger_database) + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 + + +@pytest.mark.asyncio +async def test_retention_cannot_delete_source_while_cost_acceptance_is_uncommitted( + writer: _Writer, ledger_database: str +) -> None: + client, store = writer + await _insert(store, _usage()) + _complete_and_age(ledger_database) + inserted: Final = asyncio.Event() + release: Final = asyncio.Event() + + @asynccontextmanager + async def paused_commit() -> AsyncGenerator[Prisma]: + async with client.writer_db.tx() as tx: + yield tx + inserted.set() + await release.wait() + + task: Final = asyncio.create_task(_insert(AccountingLedger(paused_commit), _cost())) + try: + await asyncio.wait_for(inserted.wait(), timeout=5) + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 0 + finally: + release.set() + await task + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ + ("cost",), + ("usage",), + ] + + +@pytest.mark.asyncio +async def test_retention_failure_rolls_back_event_and_receipt_deletion(writer: _Writer, ledger_database: str) -> None: + client, store = writer + await _insert(store, _usage(inbound=True)) + _complete_and_age(ledger_database) + + @asynccontextmanager + async def interrupted() -> AsyncGenerator[Prisma]: + async with client.writer_db.tx() as tx: + yield tx + raise RuntimeError("injected cleanup failure") + + with pytest.raises(RuntimeError, match="cleanup failure"): + await AccountingLedger(interrupted).purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)] + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)] + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 1 + + +@pytest.mark.asyncio +async def test_ledger_reads_and_writes_use_primary_when_replica_is_configured( + writer: _Writer, ledger_database: str, monkeypatch: pytest.MonkeyPatch +) -> None: + with scratch_database("replica") as replica_url: + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", replica_url) + async with connected_proxy_prisma(ledger_database) as client: + store: Final = AccountingLedger.for_client(client) + await _insert(store, _usage()) + assert await _insert(store, _usage()) == "duplicate" + assert await _insert(store, _cost()) == "inserted" + _complete_and_age(ledger_database) + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 + assert _rows(replica_url, "SELECT to_regclass('public.\"LiteLLM_AccountingEvent\"')") == [(None,)] + + +def _index_conditions(plan: object) -> tuple[tuple[object, object], ...]: + if isinstance(plan, dict): + own: Final = ((plan.get("Index Name"), plan.get("Index Cond")),) if "Index Name" in plan else () + children: Final = (plan.get("Plan"), *plan.get("Plans", ())) + return own + tuple(chain.from_iterable(_index_conditions(child) for child in children if child is not None)) + return tuple(chain.from_iterable(_index_conditions(item) for item in plan)) if isinstance(plan, list) else () + + +@pytest.mark.parametrize("statement", [accounting_ledger._DELETE_EVENTS, accounting_ledger._DELETE_DELIVERIES]) +def test_family_cleanup_lookups_use_the_kind_source_index(ledger_database: str, statement: str) -> None: + explained: Final = "EXPLAIN (FORMAT JSON) " + statement.replace("$1", "%s") + with psycopg.connect(ledger_database) as connection: + connection.execute("SET enable_seqscan = off") + plan: Final = connection.execute(explained, ("usage",) * statement.count("$1")).fetchone() + assert plan is not None + conditions: Final = _index_conditions(plan[0]) + assert any( + name == "LiteLLM_AccountingEvent_kind_source_event_id_idx" and "source_event_id" in str(condition) + for name, condition in conditions + ), conditions diff --git a/tests/unit/proxy/db/test_accounting_ledger.py b/tests/unit/proxy/db/test_accounting_ledger.py new file mode 100644 index 00000000000..3916fe869e6 --- /dev/null +++ b/tests/unit/proxy/db/test_accounting_ledger.py @@ -0,0 +1,103 @@ +import json +from datetime import datetime, timezone +from decimal import Decimal +from typing import Final, Literal, NoReturn + +import pytest + +from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome +from litellm.proxy.db.accounting_ledger import AccountingLedger +from litellm.types.utils import CallTypes + + +def _no_transaction() -> NoReturn: + raise AssertionError("Invalid accounting input must not open a database transaction") + + +@pytest.fixture +def event() -> InboundUsageEvent: + return InboundUsageEvent( + event_id="caller", + request_id="caller", + call_type=CallTypes.completion, + started_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + outcome=Outcome.SUCCEEDED, + dimensions=(), + metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("envelope", ["{}", "[]", '{"event":null}', '{"event":{}}', "not json"]) +async def test_invalid_envelopes_fail_before_database_access(event: InboundUsageEvent, envelope: str) -> None: + with pytest.raises(ValueError, match=r"Accounting envelope|validation error"): + await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [0, -1, True, 1.5]) +@pytest.mark.parametrize("field", ["version", "schema_version"]) +async def test_invalid_versions_fail_before_database_access( + event: InboundUsageEvent, version: int | float, field: Literal["version", "schema_version"] +) -> None: + with pytest.raises(ValueError, match="positive integers"): + await AccountingLedger(_no_transaction).insert_event( + event, + (), + envelope="{}", + **{field: version}, # pyright: ignore[reportArgumentType] # Exercise invalid runtime versions + ) + + +@pytest.mark.asyncio +async def test_constructed_invalid_events_are_revalidated(event: InboundUsageEvent) -> None: + invalid: Final = event.model_copy(update={"event_id": ""}) + with pytest.raises(ValueError, match="event_id"): + await AccountingLedger(_no_transaction).insert_event(invalid, (), envelope="{}") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limit", [0, -1, 1001, True, 1.5]) +async def test_retention_rejects_invalid_batch_bounds(limit: int | float) -> None: + with pytest.raises(ValueError, match="batch size"): + await AccountingLedger(_no_transaction).purge_completed( + limit=limit # pyright: ignore[reportArgumentType] # Exercise invalid runtime batch sizes + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("missing", ["both", "redelivery", "recovery"]) +async def test_retention_requires_both_safety_watermarks(missing: Literal["both", "redelivery", "recovery"]) -> None: + cutoff: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + deleted: Final = await AccountingLedger(_no_transaction).purge_completed( + redelivery_before=cutoff if missing == "recovery" else None, + recovery_before=cutoff if missing == "redelivery" else None, + ) + assert deleted == 0 + + +@pytest.mark.asyncio +async def test_retention_rejects_ambiguous_naive_watermarks() -> None: + with pytest.raises(ValueError, match="timezone"): + await AccountingLedger(_no_transaction).purge_completed( + redelivery_before=datetime(2026, 1, 1), recovery_before=datetime(2026, 1, 1, tzinfo=timezone.utc) + ) + + +@pytest.mark.asyncio +async def test_nonfinite_context_is_rejected_before_database_access(event: InboundUsageEvent) -> None: + envelope: Final = json.dumps({"event": event.model_dump(mode="json"), "context": {"price": float("nan")}}) + with pytest.raises(ValueError, match="Out of range float"): + await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope) + + +@pytest.mark.asyncio +async def test_unknown_delivery_target_is_rejected_before_database_access(event: InboundUsageEvent) -> None: + envelope: Final = json.dumps({"event": event.model_dump(mode="json")}) + with pytest.raises(ValueError, match="not a valid DeliveryTarget"): + await AccountingLedger(_no_transaction).insert_event( + event, + ("unknown",), # pyright: ignore[reportArgumentType] # Exercise an unknown runtime target + envelope=envelope, + )