From 2ba9b7b2c3118013c0b2a7ed962559faf1d44f4e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:43:12 -0700 Subject: [PATCH] fix(interactions): settle an unverified registration through the durable claim A create whose settlement-store write raised no longer bills through a private in-memory gate that a later boot's resume cannot see. The claim asks the durable store first and falls back to the local gate only when the store answers that no row exists, and a missing settlement table reads as no rows so a replica without the migration still settles in process. --- .../interactions/background_cost_polling.py | 58 +++++--- .../background_interaction_settlement.py | 30 +++- .../test_background_cost_polling.py | 128 +++++++++++++++++- .../test_background_interaction_settlement.py | 44 ++++++ 4 files changed, 234 insertions(+), 26 deletions(-) diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index f1779240436..f5d025bb785 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -485,35 +485,57 @@ def _track_poll( return task +@dataclass(frozen=True, slots=True) +class _UnverifiedRegistrationStore: + """ + Store of a create whose registration raised, so whether its row landed is + unknown until the durable store answers. The settlement claim asks it + first, and only an interaction it reports as never stored settles through + the local gate, which no other process can reach. + """ + + durable: BackgroundSettlementStore + local: InMemoryBackgroundSettlementStore = field(default_factory=InMemoryBackgroundSettlementStore) + + async def register(self, pending: PendingBackgroundInteraction) -> None: + await self.durable.register(pending) + + async def pending(self, interaction_id: str) -> PendingBackgroundInteraction | None: + return await self.durable.pending(interaction_id) + + async def is_claimed(self, interaction_id: str) -> bool: + return await self.local.is_claimed(interaction_id) or await self.durable.is_claimed(interaction_id) + + async def claim(self, interaction_id: str) -> bool: + if await self.durable.claim(interaction_id): + return True + if await self.durable.is_claimed(interaction_id): + return False + return await self.local.claim(interaction_id) + + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: + if await self.local.is_claimed(interaction_id): + return + await self.durable.record_outcome(interaction_id, outcome) + + async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: + return await self.durable.unclaimed() + + async def _registered_store( store: BackgroundSettlementStore, pending: PendingBackgroundInteraction ) -> BackgroundSettlementStore: try: await store.register(pending) - except Exception: # noqa: BLE001 # a store outage must not fail the create; the poll settles from this process - if await _registration_landed(store, pending.interaction_id): - verbose_logger.exception( - "Registering background interaction %s raised although its row landed; it settles through the store", - pending.interaction_id, - ) - return store + except Exception: # noqa: BLE001 # a store outage must not fail the create; the claim learns if the row landed verbose_logger.exception( - "Could not durably register background interaction %s; only this process can settle it", + "Could not durably register background interaction %s; its settlement claim decides whether the row landed", pending.interaction_id, ) - fallback: Final = InMemoryBackgroundSettlementStore() - await fallback.register(pending) - return fallback + return _UnverifiedRegistrationStore(durable=store) return store -async def _registration_landed(store: BackgroundSettlementStore, interaction_id: str) -> bool: - try: - return await store.pending(interaction_id) is not None or await store.is_claimed(interaction_id) - except Exception: # noqa: BLE001 # the row cannot be read either, so the poll settles from this process - return False - - async def maybe_schedule_background_interaction_cost_polling( response: object, create_kwargs: Mapping[str, object], diff --git a/litellm/proxy/spend_tracking/background_interaction_settlement.py b/litellm/proxy/spend_tracking/background_interaction_settlement.py index 3403ee74189..d3b116dfbd3 100644 --- a/litellm/proxy/spend_tracking/background_interaction_settlement.py +++ b/litellm/proxy/spend_tracking/background_interaction_settlement.py @@ -5,7 +5,7 @@ from collections.abc import Awaitable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol, TypeVar from pydantic import ValidationError from typing_extensions import ReadOnly, TypedDict @@ -88,6 +88,16 @@ def _settlement_table(prisma_client: "PrismaClient") -> _SettlementTableActions: _CLEARED_CREATE_CONTEXT: Final[Mapping[str, object]] = MappingProxyType({}) +_T = TypeVar("_T") + + +async def _read_from_a_table_that_may_not_exist(query: Awaitable[_T], when_missing: _T) -> _T: + from prisma.errors import TableNotFoundError # noqa: PLC0415 # local import: prisma may be ungenerated at load + + try: + return await query + except TableNotFoundError: + return when_missing def _json(data: Mapping[str, object]) -> object: @@ -135,22 +145,30 @@ class PrismaBackgroundSettlementStore: ) async def pending(self, interaction_id: str) -> PendingBackgroundInteraction | None: - row: Final = await self.table.find_unique(where=_RowKey(interaction_id=interaction_id)) + row: Final = await self._row(interaction_id) if row is None or row.claimed_at is not None: return None return next(iter(_pending_row(row)), None) async def is_claimed(self, interaction_id: str) -> bool: - row: Final = await self.table.find_unique(where=_RowKey(interaction_id=interaction_id)) + row: Final = await self._row(interaction_id) return row is not None and row.claimed_at is not None async def claim(self, interaction_id: str) -> bool: - claimed_rows: Final = await self.table.update_many( - data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by), - where=_UnclaimedRowKey(interaction_id=interaction_id, claimed_at=None), + claimed_rows: Final = await _read_from_a_table_that_may_not_exist( + self.table.update_many( + data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by), + where=_UnclaimedRowKey(interaction_id=interaction_id, claimed_at=None), + ), + when_missing=0, ) return claimed_rows == 1 + async def _row(self, interaction_id: str) -> _SettlementRow | None: + return await _read_from_a_table_that_may_not_exist( + self.table.find_unique(where=_RowKey(interaction_id=interaction_id)), when_missing=None + ) + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: await self.table.update_many( data=_Outcome( diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/test_litellm/interactions/test_background_cost_polling.py index 4534bcbcac2..9dc71a4fa1b 100644 --- a/tests/test_litellm/interactions/test_background_cost_polling.py +++ b/tests/test_litellm/interactions/test_background_cost_polling.py @@ -939,8 +939,103 @@ async def test_delete_on_a_worker_that_resumed_the_poll_fails_when_it_cannot_fet assert await asyncio.wait_for(resumed, timeout=5) == "billed" +class _LandsThenGoesDown: + """Register commits the row and loses its acknowledgement; every read fails until the store recovers.""" + + def __init__(self): + self.store = InMemoryBackgroundSettlementStore() + self.down = True + + async def register(self, pending): + await self.store.register(pending) + raise RuntimeError("connection reset after the row was committed") + + async def pending(self, interaction_id): + self._answer() + return await self.store.pending(interaction_id) + + async def is_claimed(self, interaction_id): + self._answer() + return await self.store.is_claimed(interaction_id) + + async def claim(self, interaction_id): + self._answer() + return await self.store.claim(interaction_id) + + async def record_outcome(self, interaction_id, outcome): + return None + + async def unclaimed(self): + self._answer() + return await self.store.unclaimed() + + def _answer(self): + if self.down: + raise RuntimeError("database unavailable") + + +class _TableLessStore: + """A replica whose database never got the settlement table: writes fail and reads see no rows.""" + + async def register(self, pending): + raise RuntimeError("the settlement table does not exist") + + async def pending(self, interaction_id): + return None + + async def is_claimed(self, interaction_id): + return False + + async def claim(self, interaction_id): + return False + + async def record_outcome(self, interaction_id, outcome): + raise RuntimeError("the settlement table does not exist") + + async def unclaimed(self): + raise RuntimeError("the settlement table does not exist") + + @pytest.mark.asyncio -async def test_create_still_settles_on_its_own_replica_when_the_store_is_down(): +async def test_create_whose_registration_and_read_back_both_failed_bills_once_through_the_landed_row(): + """ + A registration that raised and could not be read back used to give the + creator a private in-memory gate, so it billed while the stored row stayed + unclaimed for the next boot to resume and bill again. With the durable + state unknown, the claim waits for the store and settles through the row. + """ + store = _LandsThenGoesDown() + logging_obj = _logging_obj(litellm_params={"metadata": _create_metadata()}) + poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False)) + task = await maybe_schedule_background_interaction_cost_polling( + response=_response("in_progress", with_usage=False), + create_kwargs={"litellm_logging_obj": logging_obj}, + custom_llm_provider="gemini", + store=store, + fetch_interaction=poll_fetch, + ) + fetch, calls = _fetch_sequence(_response("completed", with_usage=True), _response("completed", with_usage=True)) + + while_down = await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=store + ) + store.down = False + recovered = await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=store + ) + + assert (while_down, recovered) == (None, "billed") + assert len(calls) == 2 + assert await store.is_claimed("interactions/bg-abc") + assert await store.unclaimed() == () + assert await resume_unsettled_background_interactions(store, poll_fetch, schedule=FAST_SCHEDULE) == () + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +async def test_create_whose_store_never_answers_is_not_billed_through_a_private_gate(): logging_obj = _logging_obj() poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False)) task = await maybe_schedule_background_interaction_cost_polling( @@ -959,8 +1054,37 @@ async def test_create_still_settles_on_its_own_replica_when_the_store_is_down(): store=_DownStore(), ) - assert outcome == "billed" + assert outcome is None assert len(calls) == 1 + assert "response_cost" not in logging_obj.model_call_details + assert not task.done() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +async def test_create_on_a_replica_without_the_settlement_table_still_settles_in_process(): + logging_obj = _logging_obj() + poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False)) + task = await maybe_schedule_background_interaction_cost_polling( + response=_response("in_progress", with_usage=False), + create_kwargs={"litellm_logging_obj": logging_obj}, + custom_llm_provider="gemini", + store=_TableLessStore(), + fetch_interaction=poll_fetch, + ) + fetch, calls = _fetch_sequence(_response("completed", with_usage=True), _response("completed", with_usage=True)) + + outcome = await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=_TableLessStore() + ) + again = await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=_TableLessStore() + ) + + assert (outcome, again) == ("billed", None) + assert len(calls) == 2 assert logging_obj.model_call_details["response_cost"] > 0 task.cancel() with pytest.raises(asyncio.CancelledError): diff --git a/tests/test_litellm/proxy/spend_tracking/test_background_interaction_settlement.py b/tests/test_litellm/proxy/spend_tracking/test_background_interaction_settlement.py index 0d026f260d3..21384279bbb 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_background_interaction_settlement.py +++ b/tests/test_litellm/proxy/spend_tracking/test_background_interaction_settlement.py @@ -167,6 +167,50 @@ async def test_claim_is_won_by_exactly_one_settler(): assert table.rows["interactions/bg-1"].claimed_by == "replica-b:1" +class _MissingSettlementTable: + """Prisma's per-model actions against a database whose migration for this table was held back.""" + + async def create(self, *, data): + raise self._missing() + + async def find_unique(self, *, where): + raise self._missing() + + async def find_many(self, *, where): + raise self._missing() + + async def update_many(self, *, data, where): + raise self._missing() + + def _missing(self): + from prisma.errors import TableNotFoundError + + return TableNotFoundError( + { + "user_facing_error": { + "error_code": "P2021", + "meta": {"table": "public.LiteLLM_BackgroundInteractionSettlement"}, + "message": "The table does not exist in the current database.", + } + } + ) + + +@pytest.mark.asyncio +async def test_a_missing_table_holds_no_rows_and_takes_no_registration(): + from prisma.errors import TableNotFoundError + + store = PrismaBackgroundSettlementStore(table=_MissingSettlementTable(), claimed_by="replica-a:1") + + with pytest.raises(TableNotFoundError): + await store.register(_pending("interactions/bg-1")) + assert await store.pending("interactions/bg-1") is None + assert await store.is_claimed("interactions/bg-1") is False + assert await store.claim("interactions/bg-1") is False + with pytest.raises(TableNotFoundError): + await store.unclaimed() + + @pytest.mark.asyncio async def test_unclaimed_skips_claimed_and_unreadable_rows(): table = _FakeSettlementTable(