From 0a31f7b70519ac34cdf19a55c817558269312940 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:13:40 -0700 Subject: [PATCH] fix(interactions): keep the stored settlement gate when registration raises after landing, and fail resumed-poll deletes closed A registration that raised after its row committed moved the poll to a private in-memory gate, so the creating worker billed while the stored row stayed unclaimed for another replica's delete or the next boot to bill again. The row is now read back once and, when it landed, the poll claims through it like every other settler. A worker that resumed the poll after a restart is not the creator, so its delete on a failed pre-delete fetch now fails with the fetch's error instead of releasing and deleting. After a fleet restart every worker holds resumed polls, which left the fail-closed path applying nowhere. --- .../interactions/background_cost_polling.py | 17 +++- .../test_background_cost_polling.py | 91 +++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index 2e10e3d7ada..f1779240436 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -174,6 +174,7 @@ class BackgroundInteractionPollContext: max_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS timeout_seconds: float = BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS store: BackgroundSettlementStore = field(default_factory=InMemoryBackgroundSettlementStore) + resumed: bool = False FetchInteraction: TypeAlias = Callable[[BackgroundInteractionPollContext], Awaitable[InteractionsAPIResponse]] @@ -490,6 +491,12 @@ async def _registered_store( 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 verbose_logger.exception( "Could not durably register background interaction %s; only this process can settle it", pending.interaction_id, @@ -500,6 +507,13 @@ async def _registered_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], @@ -575,7 +589,7 @@ async def maybe_settle_background_interaction_before_delete( store: BackgroundSettlementStore | None = None, ) -> SettlementOutcome | None: entry: Final = _ACTIVE_POLLS.get(interaction_id) - if entry is not None: + if entry is not None and not entry.context.resumed: return await _settle_before_delete(entry.context, await _fetch_before_delete(entry.context, fetch_interaction)) settlement_store: Final = store or _STORE.store pending: Final = await _pending(settlement_store, interaction_id) @@ -633,6 +647,7 @@ def _resumed_context( max_interval_seconds=schedule.max_interval_seconds, timeout_seconds=max(schedule.timeout_seconds - age_seconds, schedule.initial_interval_seconds), store=store, + resumed=True, ) diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/test_litellm/interactions/test_background_cost_polling.py index 22c697e04bc..4534bcbcac2 100644 --- a/tests/test_litellm/interactions/test_background_cost_polling.py +++ b/tests/test_litellm/interactions/test_background_cost_polling.py @@ -848,6 +848,97 @@ class _DownStore: raise RuntimeError("database unavailable") +class _RegistersThenRaises: + def __init__(self): + self.store = InMemoryBackgroundSettlementStore() + + 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): + return await self.store.pending(interaction_id) + + async def is_claimed(self, interaction_id): + return await self.store.is_claimed(interaction_id) + + async def claim(self, interaction_id): + return await self.store.claim(interaction_id) + + async def record_outcome(self, interaction_id, outcome): + return None + + async def unclaimed(self): + return await self.store.unclaimed() + + +@pytest.mark.asyncio +async def test_create_whose_registration_raised_after_landing_still_claims_the_stored_row(): + """ + A registration that raises after its row committed used to move the poll + to a private in-memory gate, so the creating worker billed while the + stored row stayed unclaimed for another replica's delete or the next boot + to bill again. The row that landed is the gate every settler shares. + """ + store = _RegistersThenRaises() + 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, _ = _capturing_fetch(_response("completed", with_usage=True)) + + outcome = await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=store + ) + + assert outcome == "billed" + assert await store.is_claimed("interactions/bg-abc") + assert await store.unclaimed() == () + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +async def test_delete_on_a_worker_that_resumed_the_poll_fails_when_it_cannot_fetch(): + """ + After a restart every worker resumes the unclaimed rows, so none of them + is the creator whose delete may release and delete on a failed fetch. A + resumed worker's delete fails like any other replica's, and its own poll + still bills the interaction once it completes. + """ + store = InMemoryBackgroundSettlementStore() + await store.register( + PendingBackgroundInteraction( + interaction_id="interactions/bg-abc", + custom_llm_provider="gemini", + create_context=_create_context(_logging_obj(litellm_params={"metadata": _create_metadata()}), "gemini"), + created_at=datetime.now(timezone.utc), + ) + ) + responses = [_response("in_progress", with_usage=False)] + + async def poll_fetch(context): + return responses[-1] + + (resumed,) = await resume_unsettled_background_interactions(store, poll_fetch, schedule=FAST_SCHEDULE) + fetch, _ = _fetch_sequence(RuntimeError("Google API key is required")) + + with pytest.raises(RuntimeError, match="Google API key is required"): + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=store + ) + + assert await store.is_claimed("interactions/bg-abc") is False + responses.append(_response("completed", with_usage=True)) + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + + @pytest.mark.asyncio async def test_create_still_settles_on_its_own_replica_when_the_store_is_down(): logging_obj = _logging_obj()