fix(interactions): settle pending background interaction billing before delete

litellm/interactions/background_cost_polling.py

@@ -9,6 +9,16 @@ interaction). The create call is therefore the only place that can own
 billing: it schedules a poll task that fetches the interaction until it
 reaches a terminal status and logs the final usage as a single success event
 attributed to the original request.
+
+Deleting an interaction makes every subsequent poll fail, which would let a
+caller retrieve the completed output themselves and then delete it before the
+poll task settles, leaving the work unbilled and the budget reservation
+refunded at the poll timeout. ``adelete`` therefore settles any pending poll
+for the interaction before dispatching the delete: it fetches the current
+state with the create's credentials, bills it if it is terminal with usage,
+and releases the reservation otherwise. A settlement gate on the create's
+logging object makes the poll task and the delete path mutually exclusive, so
+the interaction is billed exactly once no matter who settles first.
 """
 
 import asyncio
@@ -65,6 +75,25 @@ def _poll_intervals(initial: float, maximum: float, timeout: float) -> Iterator[
         interval = min(interval * 2, maximum)
 
 
+_SETTLED_KEY = "background_interaction_settled"
+
+
+def _is_settled(logging_obj: "LiteLLMLoggingObj") -> bool:
+    return logging_obj.model_call_details.get(_SETTLED_KEY) is True
+
+
+def _claim_settlement(logging_obj: "LiteLLMLoggingObj") -> bool:
+    """
+    Exactly-once gate between the poll task and the delete-time settlement:
+    both run on the same event loop and neither awaits between reading and
+    setting the flag, so whichever claims first owns billing or release.
+    """
+    if _is_settled(logging_obj):
+        return False
+    logging_obj.model_call_details[_SETTLED_KEY] = True
+    return True
+
+
 async def poll_and_log_background_interaction_cost(
     context: BackgroundInteractionPollContext,
     fetch_interaction: FetchInteraction = _fetch_interaction,
@@ -75,6 +104,8 @@ async def poll_and_log_background_interaction_cost(
         timeout=context.timeout_seconds,
     ):
         await asyncio.sleep(interval)
+        if _is_settled(context.logging_obj):
+            return
         try:
             response = await fetch_interaction(context)
         except Exception as e:  # noqa: BLE001  # any fetch error must not kill the billing poll loop
@@ -86,11 +117,15 @@ async def poll_and_log_background_interaction_cost(
             continue
         if response.status not in _TERMINAL_STATUSES:
             continue
+        if not _claim_settlement(context.logging_obj):
+            return
         if response.usage is not None:
             await context.logging_obj.async_log_background_interaction_completion(result=response)
         else:
             await _release_open_budget_reservation(logging_obj=context.logging_obj)
         return
+    if not _claim_settlement(context.logging_obj):
+        return
     verbose_logger.warning(
         "Gave up cost polling for background interaction %s after %ss; its usage will not be tracked",
         context.interaction_id,
@@ -104,9 +139,10 @@ async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") ->
     The proxy keeps the pre-call budget reservation open for an in-progress
     background interaction so concurrent creates cannot stack past the budget.
     The completion success event reconciles it to the actual cost; when the
-    interaction terminates without billable usage (or polling gives up), no
-    such event fires, so the poller must release the reservation here or the
-    spend counters stay pinned at the estimated cost.
+    interaction terminates without billable usage (or polling gives up, or it
+    is deleted before settling), no such event fires, so whoever claims the
+    settlement must release the reservation here or the spend counters stay
+    pinned at the estimated cost.
     """
     metadata = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details)
     budget_reservation = metadata.get("user_api_key_budget_reservation")
@@ -121,7 +157,21 @@ async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") ->
         verbose_logger.exception("Failed to release budget reservation for an unbilled background interaction")
 
 
-_ACTIVE_POLL_TASKS: set["asyncio.Task[None]"] = set()  # mutable-ok: asyncio requires strong refs to running tasks
+@dataclass(frozen=True, slots=True)
+class _ActiveBackgroundPoll:
+    task: "asyncio.Task[None]"
+    context: BackgroundInteractionPollContext
+
+
+_ACTIVE_POLLS: dict[
+    str, _ActiveBackgroundPoll
+] = {}  # mutable-ok: asyncio requires strong refs to running tasks, and delete settlement looks polls up by interaction id
+
+
+def _discard_poll(interaction_id: str, task: "asyncio.Task[None]") -> None:
+    entry = _ACTIVE_POLLS.get(interaction_id)
+    if entry is not None and entry.task is task:
+        del _ACTIVE_POLLS[interaction_id]
 
 
 def maybe_schedule_background_interaction_cost_polling(
@@ -152,6 +202,35 @@ def maybe_schedule_background_interaction_cost_polling(
         api_base=create_kwargs.get("api_base"),
     )
     task = asyncio.create_task(poll_and_log_background_interaction_cost(context))
-    _ACTIVE_POLL_TASKS.add(task)
-    task.add_done_callback(_ACTIVE_POLL_TASKS.discard)
+    _ACTIVE_POLLS[context.interaction_id] = _ActiveBackgroundPoll(task=task, context=context)
+    task.add_done_callback(
+        lambda finished, interaction_id=context.interaction_id: _discard_poll(interaction_id, finished)
+    )
     return task
+
+
+async def maybe_settle_background_interaction_before_delete(
+    interaction_id: str,
+    fetch_interaction: FetchInteraction = _fetch_interaction,
+) -> None:
+    entry = _ACTIVE_POLLS.get(interaction_id)
+    if entry is None:
+        return
+    context = entry.context
+    try:
+        response = await fetch_interaction(context)
+    except Exception as e:  # noqa: BLE001  # unfetchable pre-delete state settles by releasing the reservation
+        verbose_logger.debug(
+            "Could not fetch background interaction %s before delete, releasing its reservation: %s",
+            interaction_id,
+            e,
+        )
+        if _claim_settlement(context.logging_obj):
+            await _release_open_budget_reservation(logging_obj=context.logging_obj)
+        return
+    if not _claim_settlement(context.logging_obj):
+        return
+    if response.status in _TERMINAL_STATUSES and response.usage is not None:
+        await context.logging_obj.async_log_background_interaction_completion(result=response)
+        return
+    await _release_open_budget_reservation(logging_obj=context.logging_obj)

litellm/interactions/main.py

@@ -41,6 +41,7 @@ import httpx
 import litellm
 from litellm.interactions.background_cost_polling import (
     maybe_schedule_background_interaction_cost_polling,
+    maybe_settle_background_interaction_before_delete,
 )
 from litellm.interactions.http_handler import interactions_http_handler
 from litellm.interactions.utils import (
@@ -474,6 +475,8 @@ async def adelete(
         loop = asyncio.get_event_loop()
         kwargs["adelete_interaction"] = True
 
+        await maybe_settle_background_interaction_before_delete(interaction_id=interaction_id)
+
         func = partial(
             delete,
             interaction_id=interaction_id,

tests/test_litellm/interactions/test_background_cost_polling.py

@@ -5,8 +5,10 @@ from typing import Optional
 import pytest
 
 from litellm.interactions.background_cost_polling import (
+    _SETTLED_KEY,
     BackgroundInteractionPollContext,
     maybe_schedule_background_interaction_cost_polling,
+    maybe_settle_background_interaction_before_delete,
     poll_and_log_background_interaction_cost,
 )
 from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
@@ -221,6 +223,127 @@ async def test_schedule_skips_non_pollable_results(response, create_kwargs):
     assert task is None
 
 
+def _register_poll(logging_obj: LitellmLogging, poll_fetch=None) -> asyncio.Task:
+    import litellm.interactions.background_cost_polling as bg
+
+    if poll_fetch is None:
+        poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False))
+    context = _context(logging_obj)
+    task = asyncio.create_task(poll_and_log_background_interaction_cost(context, fetch_interaction=poll_fetch))
+    bg._ACTIVE_POLLS[context.interaction_id] = bg._ActiveBackgroundPoll(task=task, context=context)
+    task.add_done_callback(lambda finished: bg._discard_poll(context.interaction_id, finished))
+    return task
+
+
+@pytest.mark.asyncio
+async def test_delete_settlement_bills_pending_background_interaction():
+    logging_obj = _logging_obj()
+    task = _register_poll(logging_obj)
+    fetch, calls = _fetch_sequence(_response("completed", with_usage=True))
+
+    await maybe_settle_background_interaction_before_delete(
+        interaction_id="interactions/bg-abc",
+        fetch_interaction=fetch,
+    )
+
+    assert len(calls) == 1
+    assert logging_obj.model_call_details["response_cost"] > 0
+    assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175
+    await asyncio.wait_for(task, timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_delete_settlement_releases_reservation_when_still_in_progress():
+    reservation = _reservation()
+    logging_obj = _logging_obj_with_reservation(reservation)
+    task = _register_poll(logging_obj)
+    fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False))
+
+    await maybe_settle_background_interaction_before_delete(
+        interaction_id="interactions/bg-abc",
+        fetch_interaction=fetch,
+    )
+
+    assert reservation["finalized"] is True
+    assert logging_obj.model_call_details.get("response_cost") is None
+    await asyncio.wait_for(task, timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_delete_settlement_releases_reservation_when_prefetch_fails():
+    reservation = _reservation()
+    logging_obj = _logging_obj_with_reservation(reservation)
+    task = _register_poll(logging_obj)
+    fetch, _ = _fetch_sequence(RuntimeError("interaction already deleted"))
+
+    await maybe_settle_background_interaction_before_delete(
+        interaction_id="interactions/bg-abc",
+        fetch_interaction=fetch,
+    )
+
+    assert reservation["finalized"] is True
+    assert logging_obj.model_call_details.get("response_cost") is None
+    await asyncio.wait_for(task, timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_delete_settlement_ignores_interactions_without_pending_poll():
+    fetch, calls = _fetch_sequence(_response("completed", with_usage=True))
+
+    await maybe_settle_background_interaction_before_delete(
+        interaction_id="interactions/never-polled",
+        fetch_interaction=fetch,
+    )
+
+    assert calls == []
+
+
+@pytest.mark.asyncio
+async def test_delete_settlement_noop_after_poll_task_finished():
+    logging_obj = _logging_obj()
+    poll_fetch, _ = _fetch_sequence(_response("completed", with_usage=True))
+    task = _register_poll(logging_obj, poll_fetch=poll_fetch)
+    await asyncio.wait_for(task, timeout=5)
+    assert logging_obj.model_call_details["response_cost"] > 0
+
+    settle_fetch, settle_calls = _fetch_sequence(_response("completed", with_usage=True))
+    await maybe_settle_background_interaction_before_delete(
+        interaction_id="interactions/bg-abc",
+        fetch_interaction=settle_fetch,
+    )
+
+    assert settle_calls == []
+
+
+@pytest.mark.asyncio
+async def test_delete_settlement_does_not_rebill_when_gate_already_claimed():
+    logging_obj = _logging_obj()
+    logging_obj.model_call_details[_SETTLED_KEY] = True
+    task = _register_poll(logging_obj)
+    fetch, calls = _fetch_sequence(_response("completed", with_usage=True))
+
+    await maybe_settle_background_interaction_before_delete(
+        interaction_id="interactions/bg-abc",
+        fetch_interaction=fetch,
+    )
+
+    assert len(calls) == 1
+    assert logging_obj.model_call_details.get("response_cost") is None
+    await asyncio.wait_for(task, timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_poller_exits_without_billing_once_settled_elsewhere():
+    logging_obj = _logging_obj()
+    logging_obj.model_call_details[_SETTLED_KEY] = True
+    fetch, calls = _fetch_sequence(_response("completed", with_usage=True))
+
+    await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch)
+
+    assert calls == []
+    assert logging_obj.model_call_details.get("response_cost") is None
+
+
 @pytest.mark.asyncio
 async def test_schedule_respects_kill_switch(monkeypatch):
     import litellm.interactions.background_cost_polling as module