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.

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

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(

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):

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(