diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index 5f217054c22..d8272db7fbb 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -152,12 +152,26 @@ def _poll_intervals(initial: float, maximum: float, timeout: float) -> Iterator[ _SETTLED_KEY = "background_interaction_settled" +_SETTLED_OUTCOME_KEY = "background_interaction_settled_outcome" +_SETTLEMENT_OUTCOMES: tuple[SettlementOutcome, ...] = ("billed", "released", "abandoned", "error") def _is_settled(logging_obj: "LiteLLMLoggingObj") -> bool: return logging_obj.model_call_details.get(_SETTLED_KEY) is True +def _stash_settlement_outcome(logging_obj: "LiteLLMLoggingObj", outcome: SettlementOutcome) -> None: + logging_obj.model_call_details[_SETTLED_OUTCOME_KEY] = outcome + + +def _stashed_settlement_outcome(logging_obj: "LiteLLMLoggingObj") -> SettlementOutcome: + stashed = logging_obj.model_call_details.get(_SETTLED_OUTCOME_KEY) + for outcome in _SETTLEMENT_OUTCOMES: + if stashed == outcome: + return outcome + return "released" + + def _claim_settlement(logging_obj: "LiteLLMLoggingObj") -> bool: """ Exactly-once gate between the poll task and the delete-time settlement: @@ -198,7 +212,7 @@ def build_settlement_context( return BackgroundSettlementContext( interaction_id=context.interaction_id, custom_llm_provider=context.custom_llm_provider, - model=str(context.logging_obj.model_call_details.get("model") or context.logging_obj.model), + model=str(context.logging_obj.model or context.logging_obj.model_call_details.get("model")), model_group=model_group if isinstance(model_group, str) else None, litellm_call_id=str(context.logging_obj.litellm_call_id), litellm_trace_id=context.logging_obj.model_call_details.get("litellm_trace_id"), @@ -258,9 +272,11 @@ async def _record_outcome_best_effort( async def _claim_across_gates( context: BackgroundInteractionPollContext, store: Optional[BackgroundSettlementStore], + intended_outcome: SettlementOutcome, ) -> bool: if not _claim_settlement(context.logging_obj): return False + _stash_settlement_outcome(context.logging_obj, intended_outcome) if store is not None and not await _claim_in_store(store, context.interaction_id): _finalize_reservation_locally(context.logging_obj) return False @@ -287,6 +303,7 @@ async def _settle_claimed( ) await _release_open_budget_reservation(logging_obj=context.logging_obj) outcome = "error" + _stash_settlement_outcome(context.logging_obj, outcome) if store is not None: await _record_outcome_best_effort(store, context.interaction_id, outcome) @@ -312,6 +329,14 @@ async def poll_and_log_background_interaction_cost( if configured_store is not None and await _persist_pending_settlement(context, configured_store) else None ) + if active_store is not None and _is_settled(context.logging_obj): + if await _claim_in_store(active_store, context.interaction_id): + await _record_outcome_best_effort( + active_store, + context.interaction_id, + _stashed_settlement_outcome(context.logging_obj), + ) + return for interval in _poll_intervals( initial=context.initial_interval_seconds, maximum=context.max_interval_seconds, @@ -335,11 +360,12 @@ async def poll_and_log_background_interaction_cost( continue if response.status not in _TERMINAL_STATUSES: continue - if not await _claim_across_gates(context, active_store): + intended_outcome: SettlementOutcome = "billed" if response.usage is not None else "released" + if not await _claim_across_gates(context, active_store, intended_outcome): return await _settle_claimed(context=context, response=response, store=active_store) return - if not await _claim_across_gates(context, active_store): + if not await _claim_across_gates(context, active_store, "abandoned"): return verbose_logger.warning( "Gave up cost polling for background interaction %s after %ss; its usage will not be tracked", @@ -445,15 +471,16 @@ async def maybe_settle_background_interaction_before_delete( interaction_id, e, ) - if not await _claim_across_gates(context, active_store): + if not await _claim_across_gates(context, active_store, "released"): return await _release_open_budget_reservation(logging_obj=context.logging_obj) if active_store is not None: await _record_outcome_best_effort(active_store, interaction_id, "released") return - if not await _claim_across_gates(context, active_store): + terminal_with_usage = response.status in _TERMINAL_STATUSES and response.usage is not None + if not await _claim_across_gates(context, active_store, "billed" if terminal_with_usage else "released"): return - if response.status in _TERMINAL_STATUSES and response.usage is not None: + if terminal_with_usage: await _settle_claimed(context=context, response=response, store=active_store) return await _release_open_budget_reservation(logging_obj=context.logging_obj) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5ce57db230f..bc22342afd4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7904,7 +7904,6 @@ class ProxyStartupEvent: ) pass - ### BACKGROUND INTERACTION SETTLEMENT ### if BACKGROUND_INTERACTION_COST_POLLING_ENABLED: from litellm.interactions.background_cost_polling import set_settlement_store from litellm.proxy.spend_tracking.background_settlement import ( diff --git a/litellm/proxy/spend_tracking/background_settlement.py b/litellm/proxy/spend_tracking/background_settlement.py index 0209efab3df..6905f3ba950 100644 --- a/litellm/proxy/spend_tracking/background_settlement.py +++ b/litellm/proxy/spend_tracking/background_settlement.py @@ -27,6 +27,7 @@ from litellm.interactions.background_cost_polling import ( _TERMINAL_STATUSES, BackgroundSettlementContext, SettlementOutcome, + _release_open_budget_reservation, ) from litellm.types.interactions import InteractionsAPIResponse @@ -120,7 +121,7 @@ def rebuild_logging_for_settlement(context: BackgroundSettlementContext) -> "Lit messages=[{"role": "user", "content": f""}], stream=False, call_type=context.call_type, - start_time=datetime.now(timezone.utc), + start_time=datetime.now(), litellm_call_id=context.litellm_call_id, function_id=str(uuid.uuid4()), litellm_trace_id=context.litellm_trace_id, @@ -155,6 +156,34 @@ async def _release_persisted_reservation(context: BackgroundSettlementContext) - ) +async def _claim_row_best_effort(store: SettlementRowStore, interaction_id: str) -> bool: + try: + return await store.claim(interaction_id) + except Exception: # noqa: BLE001 # an unreachable store defers to the sweep instead of risking a double bill + verbose_proxy_logger.warning( + "Failed to claim settlement for background interaction %s; deferring to the settlement sweep", + interaction_id, + exc_info=True, + ) + return False + + +async def _record_row_outcome_best_effort( + store: SettlementRowStore, + interaction_id: str, + outcome: SettlementOutcome, +) -> None: + try: + await store.record_outcome(interaction_id, outcome) + except Exception: # noqa: BLE001 # the outcome column is observability, never worth failing settlement over + verbose_proxy_logger.warning( + "Failed to record settlement outcome %s for background interaction %s", + outcome, + interaction_id, + exc_info=True, + ) + + async def settle_claimed_row( row: PendingSettlementRow, response: InteractionsAPIResponse, @@ -163,20 +192,30 @@ async def settle_claimed_row( if response.status in _TERMINAL_STATUSES and response.usage is not None: try: logging_obj = rebuild_logging_for_settlement(row.context) - await logging_obj.async_log_background_interaction_completion(result=response) - await store.record_outcome(row.interaction_id, "billed") + except Exception: # noqa: BLE001 # an unbuildable logging context settles by releasing the reservation + verbose_proxy_logger.exception( + "Could not rebuild a billable logging context for background interaction %s; " + "its spend may be under-tracked", + row.interaction_id, + ) + await _release_persisted_reservation(row.context) + await _record_row_outcome_best_effort(store, row.interaction_id, "error") return + try: + await logging_obj.async_log_background_interaction_completion(result=response) except Exception: # noqa: BLE001 # a billing failure after winning the claim must be surfaced, not retried verbose_proxy_logger.exception( "Billing failed after claiming settlement for background interaction %s; " "its spend may be under-tracked", row.interaction_id, ) - await _release_persisted_reservation(row.context) - await store.record_outcome(row.interaction_id, "error") + await _release_open_budget_reservation(logging_obj=logging_obj) + await _record_row_outcome_best_effort(store, row.interaction_id, "error") return + await _record_row_outcome_best_effort(store, row.interaction_id, "billed") + return await _release_persisted_reservation(row.context) - await store.record_outcome(row.interaction_id, "released") + await _record_row_outcome_best_effort(store, row.interaction_id, "released") async def settle_row_before_delete( @@ -192,12 +231,12 @@ async def settle_row_before_delete( row.interaction_id, e, ) - if not await store.claim(row.interaction_id): + if not await _claim_row_best_effort(store, row.interaction_id): return await _release_persisted_reservation(row.context) - await store.record_outcome(row.interaction_id, "released") + await _record_row_outcome_best_effort(store, row.interaction_id, "released") return - if not await store.claim(row.interaction_id): + if not await _claim_row_best_effort(store, row.interaction_id): return await settle_claimed_row(row=row, response=response, store=store) @@ -211,7 +250,13 @@ async def sweep_pending_settlements( now = datetime.now(timezone.utc) rows = await store.list_due(older_than=now - timedelta(seconds=min_age_seconds), limit=limit) for row in rows: - await _sweep_row(row=row, store=store, fetch=fetch, now=now) + try: + await _sweep_row(row=row, store=store, fetch=fetch, now=now) + except Exception: # noqa: BLE001 # one bad row must not stop the sweep from settling the rest + verbose_proxy_logger.exception( + "Settlement sweep failed for background interaction %s; leaving it for the next cycle", + row.interaction_id, + ) async def _sweep_row( @@ -221,7 +266,7 @@ async def _sweep_row( now: datetime, ) -> None: if row.timeout_at <= now: - if not await store.claim(row.interaction_id): + if not await _claim_row_best_effort(store, row.interaction_id): return verbose_proxy_logger.warning( "Abandoning settlement for background interaction %s past its %s timeout; its usage will not be tracked", @@ -229,7 +274,7 @@ async def _sweep_row( row.timeout_at, ) await _release_persisted_reservation(row.context) - await store.record_outcome(row.interaction_id, "abandoned") + await _record_row_outcome_best_effort(store, row.interaction_id, "abandoned") return try: response = await fetch(row.context) @@ -242,7 +287,7 @@ async def _sweep_row( return if response.status not in _TERMINAL_STATUSES: return - if not await store.claim(row.interaction_id): + if not await _claim_row_best_effort(store, row.interaction_id): return await settle_claimed_row(row=row, response=response, store=store) @@ -331,7 +376,16 @@ class PrismaBackgroundSettlementStore: return tuple(row for row in parsed if row is not None) async def settle_pending_before_delete(self, interaction_id: str) -> None: - row = await self.get_pending(interaction_id) + try: + row = await self.get_pending(interaction_id) + except Exception: # noqa: BLE001 # a store failure must not fail the caller's delete; the sweep settles later + verbose_proxy_logger.warning( + "Could not read the pending settlement for background interaction %s before delete; " + "deferring to the settlement sweep", + interaction_id, + exc_info=True, + ) + return if row is None: return await settle_row_before_delete(row=row, store=self, fetch=self.fetch) diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/test_litellm/interactions/test_background_cost_polling.py index 19a4fcc3595..473a39ab20e 100644 --- a/tests/test_litellm/interactions/test_background_cost_polling.py +++ b/tests/test_litellm/interactions/test_background_cost_polling.py @@ -524,6 +524,48 @@ async def test_poller_exits_early_on_fetch_error_once_settled_elsewhere(): assert store.outcomes == {} +@pytest.mark.asyncio +async def test_poll_task_closes_row_when_delete_settled_before_persist(): + import litellm.interactions.background_cost_polling as bg + + store = _InMemorySettlementStore() + logging_obj = _logging_obj() + context = _context(logging_obj) + placeholder = asyncio.create_task(asyncio.sleep(0)) + bg._ACTIVE_POLLS[context.interaction_id] = bg._ActiveBackgroundPoll(task=placeholder, context=context) + delete_fetch, _ = _fetch_sequence(_response("completed", with_usage=True)) + await maybe_settle_background_interaction_before_delete( + interaction_id=context.interaction_id, + fetch_interaction=delete_fetch, + store=store, + ) + bg._ACTIVE_POLLS.pop(context.interaction_id, None) + await placeholder + assert logging_obj.model_call_details["response_cost"] > 0 + assert store.rows == {} + + poll_fetch, poll_calls = _fetch_sequence(_response("completed", with_usage=True)) + await poll_and_log_background_interaction_cost(context, fetch_interaction=poll_fetch, store=store) + + assert poll_calls == [] + assert store.rows["interactions/bg-abc"]["status"] == "settled" + assert store.outcomes["interactions/bg-abc"] == "billed" + + +@pytest.mark.asyncio +async def test_persisted_model_keeps_provider_scoped_form_when_precall_overwrites_it(): + store = _InMemorySettlementStore() + logging_obj = _logging_obj() + logging_obj.update_environment_variables(litellm_params={}, optional_params={}, model="gemini/gemini-2.5-flash") + logging_obj._pre_call(input="hi", api_key=None, model="gemini-2.5-flash") + fetch, _ = _fetch_sequence(_response("completed", with_usage=True)) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch, store=store) + + assert logging_obj.model_call_details["model"] == "gemini-2.5-flash" + assert store.contexts["interactions/bg-abc"].model == "gemini/gemini-2.5-flash" + + @pytest.mark.asyncio async def test_delete_registry_miss_delegates_to_store(): store = _InMemorySettlementStore() diff --git a/tests/test_litellm/proxy/spend_tracking/test_background_settlement.py b/tests/test_litellm/proxy/spend_tracking/test_background_settlement.py index 570b9019f19..8ba0073caf6 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_background_settlement.py +++ b/tests/test_litellm/proxy/spend_tracking/test_background_settlement.py @@ -12,6 +12,7 @@ from litellm.interactions.background_cost_polling import ( ) from litellm.proxy.spend_tracking.background_settlement import ( PendingSettlementRow, + PrismaBackgroundSettlementStore, _parse_row, _resolve_settlement_credentials, rebuild_logging_for_settlement, @@ -135,6 +136,12 @@ async def test_rebuilt_logging_bills_with_original_attribution_and_request_id(): assert payload["spend"] > 0 +def test_rebuilt_logging_start_time_is_naive_for_duration_math(): + logging_obj = rebuild_logging_for_settlement(_settlement_context()) + + assert logging_obj.start_time.tzinfo is None + + def test_rebuilt_logging_carries_reservation_for_reconcile(): context = _settlement_context(reservation=_reservation()) logging_obj = rebuild_logging_for_settlement(context) @@ -245,6 +252,67 @@ async def test_delete_settlement_noop_when_claim_lost(): assert store.outcomes == {} +class _RaisingRowStore: + def __init__(self) -> None: + self.outcomes: dict[str, SettlementOutcome] = {} + + async def claim(self, interaction_id: str) -> bool: + raise RuntimeError("settlement db unavailable") + + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: + raise RuntimeError("settlement db unavailable") + + async def list_due(self, older_than: datetime, limit: int) -> tuple[PendingSettlementRow, ...]: + return () + + +@pytest.mark.asyncio +async def test_delete_settlement_swallows_store_claim_errors_and_defers_to_sweep(): + reservation = _reservation() + row = _row(context=_settlement_context(reservation=reservation)) + store = _RaisingRowStore() + fetch, calls = _fetch_returning(_response("completed", with_usage=True)) + + await settle_row_before_delete(row=row, store=store, fetch=fetch) + + assert len(calls) == 1 + assert store.outcomes == {} + assert reservation.finalized is False + + +class _RecordOutcomeRaisingStore: + async def claim(self, interaction_id: str) -> bool: + return True + + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: + raise RuntimeError("settlement db unavailable") + + async def list_due(self, older_than: datetime, limit: int) -> tuple[PendingSettlementRow, ...]: + return () + + +@pytest.mark.asyncio +async def test_record_outcome_failure_after_billing_does_not_propagate_or_release(): + row = _row() + store = _RecordOutcomeRaisingStore() + fetch, _ = _fetch_returning(_response("completed", with_usage=True)) + + await settle_row_before_delete(row=row, store=store, fetch=fetch) + + +class _BrokenPrismaClient: + @property + def db(self): + raise RuntimeError("database unreachable") + + +@pytest.mark.asyncio +async def test_settle_pending_before_delete_swallows_store_read_errors(): + store = PrismaBackgroundSettlementStore(prisma_client=_BrokenPrismaClient()) + + await store.settle_pending_before_delete(INTERACTION_ID) + + class _FakeRouter: def __init__(self, deployment: Optional[dict] = None, error: Optional[Exception] = None) -> None: self.deployment = deployment