fix(interactions): bill the completed response a poll already saw when its claim only answers at the deadline

litellm/interactions/background_cost_polling.py

@@ -332,7 +332,9 @@ async def poll_and_log_background_interaction_cost(
     context: BackgroundInteractionPollContext,
     fetch_interaction: FetchInteraction = fetch_background_interaction,
 ) -> SettlementOutcome | None:
-    last_seen_status: str | None = None  # rebind-ok: the give-up log names the status the poll last saw
+    last_response: InteractionsAPIResponse | None = (
+        None  # rebind-ok: the give-up path settles from, or names, what the poll last saw
+    )
     for interval in _poll_intervals(
         initial=context.initial_interval_seconds,
         maximum=context.max_interval_seconds,
@@ -350,7 +352,7 @@ async def poll_and_log_background_interaction_cost(
                 e,
             )
             continue
-        last_seen_status = response.status
+        last_response = response
         if response.status not in _TERMINAL_STATUSES:
             continue
         if (claimed := await _claim(context)) is None:
@@ -360,14 +362,16 @@ async def poll_and_log_background_interaction_cost(
         return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=response))
     if not await _claim(context):
         return None
-    if last_seen_status is not None and last_seen_status not in _POLLABLE_STATUSES:
+    if last_response is not None and last_response.status in _TERMINAL_STATUSES:
+        return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=last_response))
+    if last_response is not None and last_response.status not in _POLLABLE_STATUSES:
         verbose_logger.error(
             "Gave up cost polling for background interaction %s after %ss: its last status %r is in neither "
             "the pollable nor the terminal set, so this proxy never learned how to settle it and its usage "
             "will not be tracked",
             context.interaction_id,
             context.timeout_seconds,
-            last_seen_status,
+            last_response.status,
         )
     else:
         verbose_logger.warning(

tests/test_litellm/interactions/test_background_cost_polling.py

@@ -724,6 +724,53 @@ async def test_restart_resumes_only_the_rows_no_replica_claimed():
     assert await store.is_claimed("interactions/bg-orphaned")
 
 
+class _ClaimAnswersOnlyAfterTheLastFetch:
+    def __init__(self):
+        self.store = InMemoryBackgroundSettlementStore()
+        self.fetches = 0
+        self.fetches_at_last_claim = -1
+
+    async def fetch(self, context):
+        self.fetches += 1
+        return _response("completed", with_usage=True)
+
+    async def register(self, pending):
+        await self.store.register(pending)
+
+    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):
+        if self.fetches != self.fetches_at_last_claim:
+            self.fetches_at_last_claim = self.fetches
+            raise RuntimeError("database unavailable")
+        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_poller_bills_the_completed_response_it_saw_when_the_claim_only_answers_at_the_deadline():
+    logging_obj = _logging_obj()
+    store = _ClaimAnswersOnlyAfterTheLastFetch()
+
+    outcome = await poll_and_log_background_interaction_cost(
+        _context(logging_obj, timeout_seconds=0.01, store=store),
+        fetch_interaction=store.fetch,
+    )
+
+    assert store.fetches >= 2
+    assert outcome == "billed"
+    assert logging_obj.model_call_details["response_cost"] > 0
+
+
 class _DownStore:
     async def register(self, pending):
         raise RuntimeError("database unavailable")