From 8b63b7045c9300d6d01f79a6954409dbb369080d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 15 Jul 2026 10:06:18 -0700 Subject: [PATCH] feat(interactions): durable cross-pod settlement for background interaction billing Background interaction billing previously lived only in the creating pod's memory: a DELETE routed to another pod or a proxy restart meant the completed provider work was never billed and the budget reservation was refunded at the poll timeout. The poll task now persists a settlement context row (virtual key attribution, model, original request id, budget reservation) in a new LiteLLM_BackgroundInteractionSettlementTable so any pod can settle. A conditional update flipping status pending to settled is the cross-pod exactly-once claim; every outcome (bill, release, timeout give-up) must win it before touching spend counters, and a losing pod finalizes its local reservation copy without touching shared counters. The delete route settles registry misses through the store by rebuilding a billable logging context, resolving provider credentials through the router by model group. A per-pod sweep job resumes pending settlements after restarts and surfaces timed-out entries as abandoned rows instead of silently reconciling to zero. Without a registered store (SDK usage) behavior is unchanged. --- .../migration.sql | 16 + .../litellm_proxy_extras/schema.prisma | 13 + litellm/constants.py | 3 + .../interactions/background_cost_polling.py | 253 ++++++++++++- litellm/proxy/proxy_server.py | 23 ++ litellm/proxy/schema.prisma | 13 + .../spend_tracking/background_settlement.py | 337 ++++++++++++++++++ schema.prisma | 13 + .../test_background_cost_polling.py | 229 ++++++++++++ .../test_background_settlement.py | 320 +++++++++++++++++ 10 files changed, 1206 insertions(+), 14 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260715120000_add_background_interaction_settlement_table/migration.sql create mode 100644 litellm/proxy/spend_tracking/background_settlement.py create mode 100644 tests/test_litellm/proxy/spend_tracking/test_background_settlement.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260715120000_add_background_interaction_settlement_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260715120000_add_background_interaction_settlement_table/migration.sql new file mode 100644 index 00000000000..98c80090372 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260715120000_add_background_interaction_settlement_table/migration.sql @@ -0,0 +1,16 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_BackgroundInteractionSettlementTable" ( + "interaction_id" TEXT NOT NULL, + "status" TEXT NOT NULL DEFAULT 'pending', + "outcome" TEXT, + "context" JSONB NOT NULL, + "claimed_by" TEXT, + "timeout_at" TIMESTAMP(3) NOT NULL, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_BackgroundInteractionSettlementTable_pkey" PRIMARY KEY ("interaction_id") +); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_BackgroundInteractionSettlement_status_created_idx" ON "LiteLLM_BackgroundInteractionSettlementTable"("status", "created_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index a23cecc3911..9117aaacdff 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -950,6 +950,19 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t @@index([team_id, created_at(sort: Desc)]) } +model LiteLLM_BackgroundInteractionSettlementTable { + interaction_id String @id + status String @default("pending") + outcome String? + context Json + claimed_by String? + timeout_at DateTime + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@index([status, created_at], map: "LiteLLM_BackgroundInteractionSettlement_status_created_idx") +} + model LiteLLM_ManagedVectorStoreTable { id String @id @default(uuid()) unified_resource_id String @unique // The base64 encoded unified vector store ID diff --git a/litellm/constants.py b/litellm/constants.py index be880b72854..aa006420c59 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1468,6 +1468,9 @@ BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS = float( ) _background_interaction_cost_polling_env = os.getenv("BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true").lower() BACKGROUND_INTERACTION_COST_POLLING_ENABLED = _background_interaction_cost_polling_env == "true" +BACKGROUND_SETTLEMENT_SWEEP_INTERVAL_SECONDS = int(os.getenv("BACKGROUND_SETTLEMENT_SWEEP_INTERVAL_SECONDS", 300)) +BACKGROUND_SETTLEMENT_SWEEP_MIN_AGE_SECONDS = float(os.getenv("BACKGROUND_SETTLEMENT_SWEEP_MIN_AGE_SECONDS", 60)) +MAX_BACKGROUND_SETTLEMENTS_PER_SWEEP = int(os.getenv("MAX_BACKGROUND_SETTLEMENTS_PER_SWEEP", 100)) PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)) PROXY_BATCH_WRITE_AT = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10 diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index ed30d2573bd..5f217054c22 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -19,11 +19,23 @@ 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. + +The poll registry and the create's logging object only live in the pod that +served the create, so a delete routed to another pod or a proxy restart would +leave the completed work unbilled. When a durable settlement store is +registered (the proxy registers a Prisma-backed one at startup), the poll +task persists enough create context for any pod to settle, and every outcome +must win the store's cross-pod compare-and-swap claim before touching spend +counters; the in-memory gate remains the same-pod fast path and the only gate +when no store is registered (SDK usage). """ import asyncio from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterator, Optional +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterator, Literal, Optional, Protocol + +from pydantic import BaseModel from litellm._logging import verbose_logger from litellm.constants import ( @@ -40,6 +52,70 @@ if TYPE_CHECKING: _TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete", "budget_exceeded"}) +SettlementOutcome = Literal["billed", "released", "abandoned", "error"] + + +class SettlementReservationEntry(BaseModel): + counter_key: str + entity_type: str + entity_id: str + reserved_cost: float + applied_adjustment: float = 0.0 + + +class SettlementReservation(BaseModel): + reserved_cost: float + entries: list[SettlementReservationEntry] + finalized: bool + input_cost: float + + +class SettlementKeyAttribution(BaseModel): + user_api_key: Optional[str] = None + user_api_key_user_id: Optional[str] = None + user_api_key_team_id: Optional[str] = None + user_api_key_org_id: Optional[str] = None + user_api_key_alias: Optional[str] = None + user_api_key_team_alias: Optional[str] = None + user_api_key_user_email: Optional[str] = None + user_api_key_end_user_id: Optional[str] = None + + +class BackgroundSettlementContext(BaseModel): + interaction_id: str + custom_llm_provider: str + model: str + model_group: Optional[str] = None + litellm_call_id: str + litellm_trace_id: Optional[str] = None + call_type: str + attribution: SettlementKeyAttribution + budget_reservation: Optional[SettlementReservation] = None + + +class BackgroundSettlementStore(Protocol): + async def persist_pending(self, context: BackgroundSettlementContext, timeout_at: datetime) -> None: ... + + async def claim(self, interaction_id: str) -> bool: ... + + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: ... + + async def is_pending(self, interaction_id: str) -> bool: ... + + async def settle_pending_before_delete(self, interaction_id: str) -> None: ... + + +_SETTLEMENT_STORE: Optional[BackgroundSettlementStore] = None + + +def get_settlement_store() -> Optional[BackgroundSettlementStore]: + return _SETTLEMENT_STORE + + +def set_settlement_store(store: Optional[BackgroundSettlementStore]) -> None: + global _SETTLEMENT_STORE + _SETTLEMENT_STORE = store + @dataclass(frozen=True, slots=True) class BackgroundInteractionPollContext: @@ -94,10 +170,148 @@ def _claim_settlement(logging_obj: "LiteLLMLoggingObj") -> bool: return True +def _get_reservation_dict(logging_obj: "LiteLLMLoggingObj") -> Optional[dict]: + metadata = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) + budget_reservation = metadata.get("user_api_key_budget_reservation") + return budget_reservation if isinstance(budget_reservation, dict) else None + + +def _finalize_reservation_locally(logging_obj: "LiteLLMLoggingObj") -> None: + """ + Losing the durable store's claim means another pod owns (or already + performed) the billing or release for this interaction. The local copy of + the reservation dict must be marked finalized without touching the shared + counters, or a later reconcile of this copy would re-apply its full + adjustment on top of the winner's. + """ + budget_reservation = _get_reservation_dict(logging_obj) + if budget_reservation is not None: + budget_reservation["finalized"] = True + + +def build_settlement_context( + context: BackgroundInteractionPollContext, +) -> BackgroundSettlementContext: + metadata = get_litellm_metadata_from_kwargs(kwargs=context.logging_obj.model_call_details) + reservation_dict = metadata.get("user_api_key_budget_reservation") + model_group = metadata.get("deployment_model_name") + 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_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"), + call_type=str(context.logging_obj.call_type), + attribution=SettlementKeyAttribution.model_validate(metadata), + budget_reservation=( + SettlementReservation.model_validate(reservation_dict) if isinstance(reservation_dict, dict) else None + ), + ) + + +async def _persist_pending_settlement( + context: BackgroundInteractionPollContext, + store: BackgroundSettlementStore, +) -> bool: + try: + settlement_context = build_settlement_context(context) + timeout_at = datetime.now(timezone.utc) + timedelta(seconds=context.timeout_seconds) + await store.persist_pending(context=settlement_context, timeout_at=timeout_at) + return True + except Exception: # noqa: BLE001 # a failed persist degrades to in-memory-only settlement, never kills the poll + verbose_logger.warning( + "Failed to persist pending settlement for background interaction %s; settlement will be in-memory only", + context.interaction_id, + exc_info=True, + ) + return False + + +async def _claim_in_store(store: BackgroundSettlementStore, 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_logger.exception( + "Failed to claim settlement for background interaction %s; deferring to the settlement sweep", + interaction_id, + ) + return False + + +async def _record_outcome_best_effort( + store: BackgroundSettlementStore, + 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_logger.warning( + "Failed to record settlement outcome %s for background interaction %s", + outcome, + interaction_id, + exc_info=True, + ) + + +async def _claim_across_gates( + context: BackgroundInteractionPollContext, + store: Optional[BackgroundSettlementStore], +) -> bool: + if not _claim_settlement(context.logging_obj): + return False + if store is not None and not await _claim_in_store(store, context.interaction_id): + _finalize_reservation_locally(context.logging_obj) + return False + return True + + +async def _settle_claimed( + context: BackgroundInteractionPollContext, + response: InteractionsAPIResponse, + store: Optional[BackgroundSettlementStore], +) -> None: + outcome: SettlementOutcome + try: + if response.usage is not None: + await context.logging_obj.async_log_background_interaction_completion(result=response) + outcome = "billed" + else: + await _release_open_budget_reservation(logging_obj=context.logging_obj) + outcome = "released" + except Exception: # noqa: BLE001 # a billing failure after winning the claim must be surfaced, not retried + verbose_logger.exception( + "Billing failed after claiming settlement for background interaction %s; its spend may be under-tracked", + context.interaction_id, + ) + await _release_open_budget_reservation(logging_obj=context.logging_obj) + outcome = "error" + if store is not None: + await _record_outcome_best_effort(store, context.interaction_id, outcome) + + +async def _no_longer_pending_in_store( + store: BackgroundSettlementStore, + interaction_id: str, +) -> bool: + try: + return not await store.is_pending(interaction_id) + except Exception: # noqa: BLE001 # an unreachable store must not change the poll loop's behavior + return False + + async def poll_and_log_background_interaction_cost( context: BackgroundInteractionPollContext, fetch_interaction: FetchInteraction = _fetch_interaction, + store: Optional[BackgroundSettlementStore] = None, ) -> None: + configured_store = store if store is not None else get_settlement_store() + active_store = ( + configured_store + if configured_store is not None and await _persist_pending_settlement(context, configured_store) + else None + ) for interval in _poll_intervals( initial=context.initial_interval_seconds, maximum=context.max_interval_seconds, @@ -109,6 +323,10 @@ async def poll_and_log_background_interaction_cost( try: response = await fetch_interaction(context) except Exception as e: # noqa: BLE001 # any fetch error must not kill the billing poll loop + if active_store is not None and await _no_longer_pending_in_store(active_store, context.interaction_id): + if _claim_settlement(context.logging_obj): + _finalize_reservation_locally(context.logging_obj) + return verbose_logger.debug( "Background interaction cost poll for %s failed, will retry: %s", context.interaction_id, @@ -117,14 +335,11 @@ async def poll_and_log_background_interaction_cost( continue if response.status not in _TERMINAL_STATUSES: continue - if not _claim_settlement(context.logging_obj): + if not await _claim_across_gates(context, active_store): 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) + await _settle_claimed(context=context, response=response, store=active_store) return - if not _claim_settlement(context.logging_obj): + if not await _claim_across_gates(context, active_store): return verbose_logger.warning( "Gave up cost polling for background interaction %s after %ss; its usage will not be tracked", @@ -132,6 +347,8 @@ async def poll_and_log_background_interaction_cost( context.timeout_seconds, ) await _release_open_budget_reservation(logging_obj=context.logging_obj) + if active_store is not None: + await _record_outcome_best_effort(active_store, context.interaction_id, "abandoned") async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> None: @@ -144,9 +361,8 @@ async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> 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") - if not isinstance(budget_reservation, dict): + budget_reservation = _get_reservation_dict(logging_obj) + if budget_reservation is None: return from litellm.proxy.spend_tracking.budget_reservation import release_budget_reservation @@ -212,9 +428,13 @@ def maybe_schedule_background_interaction_cost_polling( async def maybe_settle_background_interaction_before_delete( interaction_id: str, fetch_interaction: FetchInteraction = _fetch_interaction, + store: Optional[BackgroundSettlementStore] = None, ) -> None: + active_store = store if store is not None else get_settlement_store() entry = _ACTIVE_POLLS.get(interaction_id) if entry is None: + if active_store is not None: + await active_store.settle_pending_before_delete(interaction_id) return context = entry.context try: @@ -225,12 +445,17 @@ async def maybe_settle_background_interaction_before_delete( interaction_id, e, ) - if _claim_settlement(context.logging_obj): - await _release_open_budget_reservation(logging_obj=context.logging_obj) + if not await _claim_across_gates(context, active_store): + 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 _claim_settlement(context.logging_obj): + if not await _claim_across_gates(context, active_store): return if response.status in _TERMINAL_STATUSES and response.usage is not None: - await context.logging_obj.async_log_background_interaction_completion(result=response) + await _settle_claimed(context=context, response=response, store=active_store) 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") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a0b07ece4dd..5ce57db230f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -224,6 +224,8 @@ from litellm.constants import ( APSCHEDULER_MAX_INSTANCES, APSCHEDULER_MISFIRE_GRACE_TIME, APSCHEDULER_REPLACE_EXISTING, + BACKGROUND_INTERACTION_COST_POLLING_ENABLED, + BACKGROUND_SETTLEMENT_SWEEP_INTERVAL_SECONDS, DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_MODEL_CREATED_AT_TIME, @@ -7902,6 +7904,27 @@ 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 ( + PrismaBackgroundSettlementStore, + sweep_pending_settlements, + ) + + settlement_store = PrismaBackgroundSettlementStore(prisma_client=prisma_client) + set_settlement_store(settlement_store) + scheduler.add_job( + sweep_pending_settlements, + "interval", + seconds=BACKGROUND_SETTLEMENT_SWEEP_INTERVAL_SECONDS + random.randint(0, 30), + kwargs={"store": settlement_store}, + id="background_interaction_settlement_sweep_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + verbose_proxy_logger.info("Background interaction settlement sweep job scheduled successfully") + ### CHECK RESPONSES COST ### if llm_router is not None and PROXY_BATCH_POLLING_ENABLED: try: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index a23cecc3911..9117aaacdff 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -950,6 +950,19 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t @@index([team_id, created_at(sort: Desc)]) } +model LiteLLM_BackgroundInteractionSettlementTable { + interaction_id String @id + status String @default("pending") + outcome String? + context Json + claimed_by String? + timeout_at DateTime + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@index([status, created_at], map: "LiteLLM_BackgroundInteractionSettlement_status_created_idx") +} + model LiteLLM_ManagedVectorStoreTable { id String @id @default(uuid()) unified_resource_id String @unique // The base64 encoded unified vector store ID diff --git a/litellm/proxy/spend_tracking/background_settlement.py b/litellm/proxy/spend_tracking/background_settlement.py new file mode 100644 index 00000000000..0209efab3df --- /dev/null +++ b/litellm/proxy/spend_tracking/background_settlement.py @@ -0,0 +1,337 @@ +""" +Prisma-backed durable settlement for background interaction billing. + +The creating pod's poll task persists a settlement context row here so that a +delete routed to any pod, or a sweep after a pod restart, can rebuild a +billable logging object and settle the interaction exactly once. The row's +``status`` column is the cross-pod claim: whoever flips ``pending`` to +``settled`` via a conditional update owns billing or release, and the +``outcome`` column records what the winner did (``billed``, ``released``, +``abandoned``, or ``error``) for audit of anything that could not be billed. +""" + +import json +import uuid +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Awaitable, Callable, Optional, Protocol + +from pydantic import BaseModel + +from litellm._logging import verbose_proxy_logger +from litellm.constants import ( + BACKGROUND_SETTLEMENT_SWEEP_MIN_AGE_SECONDS, + MAX_BACKGROUND_SETTLEMENTS_PER_SWEEP, +) +from litellm.interactions.background_cost_polling import ( + _TERMINAL_STATUSES, + BackgroundSettlementContext, + SettlementOutcome, +) +from litellm.types.interactions import InteractionsAPIResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy.utils import PrismaClient + + +class PendingSettlementRow(BaseModel): + interaction_id: str + context: BackgroundSettlementContext + timeout_at: datetime + + +class SettlementRowStore(Protocol): + async def claim(self, interaction_id: str) -> bool: ... + + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: ... + + async def list_due(self, older_than: datetime, limit: int) -> tuple[PendingSettlementRow, ...]: ... + + +SettlementFetch = Callable[[BackgroundSettlementContext], Awaitable[InteractionsAPIResponse]] + + +class DeploymentResolver(Protocol): + async def async_get_available_deployment(self, model: str, request_kwargs: dict) -> dict: ... + + +def _get_proxy_router() -> Optional[DeploymentResolver]: + from litellm.proxy.proxy_server import llm_router + + return llm_router + + +@dataclass(frozen=True, slots=True) +class _SettlementCredentials: + api_key: Optional[str] + api_base: Optional[str] + + +async def _resolve_settlement_credentials( + context: BackgroundSettlementContext, + router: Optional[DeploymentResolver] = None, +) -> _SettlementCredentials: + if context.model_group is None: + return _SettlementCredentials(api_key=None, api_base=None) + active_router = router if router is not None else _get_proxy_router() + if active_router is None: + return _SettlementCredentials(api_key=None, api_base=None) + try: + deployment = await active_router.async_get_available_deployment( + model=context.model_group, + request_kwargs={}, + ) + except Exception: # noqa: BLE001 # a failed deployment lookup falls back to provider env credentials + verbose_proxy_logger.warning( + "Could not resolve a deployment for %s to settle background interaction %s; " + "falling back to provider environment credentials", + context.model_group, + context.interaction_id, + exc_info=True, + ) + return _SettlementCredentials(api_key=None, api_base=None) + litellm_params = deployment.get("litellm_params") or {} + return _SettlementCredentials( + api_key=litellm_params.get("api_key"), + api_base=litellm_params.get("api_base"), + ) + + +async def fetch_interaction_for_settlement( + context: BackgroundSettlementContext, + router: Optional[DeploymentResolver] = None, +) -> InteractionsAPIResponse: + from litellm.interactions import aget + + credentials = await _resolve_settlement_credentials(context, router=router) + return await aget( + interaction_id=context.interaction_id, + custom_llm_provider=context.custom_llm_provider, + **{"api_key": credentials.api_key, "api_base": credentials.api_base, "no-log": True}, + ) + + +def rebuild_logging_for_settlement(context: BackgroundSettlementContext) -> "LiteLLMLoggingObj": + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model=context.model, + messages=[{"role": "user", "content": f""}], + stream=False, + call_type=context.call_type, + start_time=datetime.now(timezone.utc), + litellm_call_id=context.litellm_call_id, + function_id=str(uuid.uuid4()), + litellm_trace_id=context.litellm_trace_id, + ) + attribution = {key: value for key, value in context.attribution.model_dump().items() if value is not None} + reservation = context.budget_reservation.model_dump() if context.budget_reservation is not None else None + metadata = attribution if reservation is None else {**attribution, "user_api_key_budget_reservation": reservation} + logging_obj.update_environment_variables( + litellm_params={ + "litellm_call_id": context.litellm_call_id, + "proxy_server_request": {"headers": {"user-agent": "litellm-background-settlement"}}, + "metadata": metadata, + }, + optional_params={}, + model=context.model, + custom_llm_provider=context.custom_llm_provider, + ) + return logging_obj + + +async def _release_persisted_reservation(context: BackgroundSettlementContext) -> None: + if context.budget_reservation is None: + return + from litellm.proxy.spend_tracking.budget_reservation import release_budget_reservation + + try: + await release_budget_reservation(budget_reservation=context.budget_reservation.model_dump()) + except Exception: # noqa: BLE001 # a failed release must not fail settlement; counters expire via TTL + verbose_proxy_logger.exception( + "Failed to release persisted budget reservation for background interaction %s", + context.interaction_id, + ) + + +async def settle_claimed_row( + row: PendingSettlementRow, + response: InteractionsAPIResponse, + store: SettlementRowStore, +) -> None: + 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") + return + 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") + return + await _release_persisted_reservation(row.context) + await store.record_outcome(row.interaction_id, "released") + + +async def settle_row_before_delete( + row: PendingSettlementRow, + store: SettlementRowStore, + fetch: SettlementFetch, +) -> None: + try: + response = await fetch(row.context) + except Exception as e: # noqa: BLE001 # unfetchable pre-delete state settles by releasing the reservation + verbose_proxy_logger.debug( + "Could not fetch background interaction %s before delete, releasing its reservation: %s", + row.interaction_id, + e, + ) + if not await store.claim(row.interaction_id): + return + await _release_persisted_reservation(row.context) + await store.record_outcome(row.interaction_id, "released") + return + if not await store.claim(row.interaction_id): + return + await settle_claimed_row(row=row, response=response, store=store) + + +async def sweep_pending_settlements( + store: SettlementRowStore, + fetch: SettlementFetch = fetch_interaction_for_settlement, + min_age_seconds: float = BACKGROUND_SETTLEMENT_SWEEP_MIN_AGE_SECONDS, + limit: int = MAX_BACKGROUND_SETTLEMENTS_PER_SWEEP, +) -> None: + 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) + + +async def _sweep_row( + row: PendingSettlementRow, + store: SettlementRowStore, + fetch: SettlementFetch, + now: datetime, +) -> None: + if row.timeout_at <= now: + if not await store.claim(row.interaction_id): + return + verbose_proxy_logger.warning( + "Abandoning settlement for background interaction %s past its %s timeout; its usage will not be tracked", + row.interaction_id, + row.timeout_at, + ) + await _release_persisted_reservation(row.context) + await store.record_outcome(row.interaction_id, "abandoned") + return + try: + response = await fetch(row.context) + except Exception as e: # noqa: BLE001 # a sweep fetch error must leave the row pending for a later cycle + verbose_proxy_logger.debug( + "Settlement sweep could not fetch background interaction %s, leaving it pending: %s", + row.interaction_id, + e, + ) + return + if response.status not in _TERMINAL_STATUSES: + return + if not await store.claim(row.interaction_id): + return + await settle_claimed_row(row=row, response=response, store=store) + + +def _parse_row(interaction_id: str, context_value: object, timeout_at: datetime) -> Optional[PendingSettlementRow]: + try: + raw = json.loads(context_value) if isinstance(context_value, str) else context_value + return PendingSettlementRow( + interaction_id=interaction_id, + context=BackgroundSettlementContext.model_validate(raw), + timeout_at=timeout_at, + ) + except Exception: # noqa: BLE001 # an unparseable row is surfaced and skipped rather than crashing the sweep + verbose_proxy_logger.exception( + "Could not parse persisted settlement context for background interaction %s", + interaction_id, + ) + return None + + +class PrismaBackgroundSettlementStore: + def __init__( + self, + prisma_client: "PrismaClient", + fetch: SettlementFetch = fetch_interaction_for_settlement, + claimed_by: Optional[str] = None, + ) -> None: + self.prisma_client = prisma_client + self.fetch = fetch + self.claimed_by = claimed_by or str(uuid.uuid4()) + + @property + def _table(self): + return self.prisma_client.db.litellm_backgroundinteractionsettlementtable + + async def persist_pending(self, context: BackgroundSettlementContext, timeout_at: datetime) -> None: + context_json = context.model_dump_json() + await self._table.upsert( + where={"interaction_id": context.interaction_id}, + data={ + "create": { + "interaction_id": context.interaction_id, + "context": context_json, + "timeout_at": timeout_at, + }, + "update": {"context": context_json, "timeout_at": timeout_at}, + }, + ) + + async def claim(self, interaction_id: str) -> bool: + updated = await self._table.update_many( + where={"interaction_id": interaction_id, "status": "pending"}, + data={"status": "settled", "claimed_by": self.claimed_by}, + ) + if updated > 0: + return True + row = await self._table.find_unique(where={"interaction_id": interaction_id}) + return row is None + + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: + await self._table.update_many( + where={"interaction_id": interaction_id}, + data={"outcome": outcome}, + ) + + async def is_pending(self, interaction_id: str) -> bool: + row = await self._table.find_unique(where={"interaction_id": interaction_id}) + return row is not None and row.status == "pending" + + async def get_pending(self, interaction_id: str) -> Optional[PendingSettlementRow]: + row = await self._table.find_unique(where={"interaction_id": interaction_id}) + if row is None or row.status != "pending": + return None + return _parse_row(interaction_id=row.interaction_id, context_value=row.context, timeout_at=row.timeout_at) + + async def list_due(self, older_than: datetime, limit: int) -> tuple[PendingSettlementRow, ...]: + rows = await self._table.find_many( + where={"status": "pending", "created_at": {"lt": older_than}}, + take=limit, + order={"created_at": "asc"}, + ) + parsed = ( + _parse_row(interaction_id=row.interaction_id, context_value=row.context, timeout_at=row.timeout_at) + for row in rows + ) + 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) + if row is None: + return + await settle_row_before_delete(row=row, store=self, fetch=self.fetch) diff --git a/schema.prisma b/schema.prisma index a23cecc3911..9117aaacdff 100644 --- a/schema.prisma +++ b/schema.prisma @@ -950,6 +950,19 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t @@index([team_id, created_at(sort: Desc)]) } +model LiteLLM_BackgroundInteractionSettlementTable { + interaction_id String @id + status String @default("pending") + outcome String? + context Json + claimed_by String? + timeout_at DateTime + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@index([status, created_at], map: "LiteLLM_BackgroundInteractionSettlement_status_created_idx") +} + model LiteLLM_ManagedVectorStoreTable { id String @id @default(uuid()) unified_resource_id String @unique // The base64 encoded unified vector store ID diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/test_litellm/interactions/test_background_cost_polling.py index 64436ebe4a7..19a4fcc3595 100644 --- a/tests/test_litellm/interactions/test_background_cost_polling.py +++ b/tests/test_litellm/interactions/test_background_cost_polling.py @@ -1,5 +1,6 @@ import asyncio import time +from datetime import datetime, timezone from typing import Optional import pytest @@ -7,6 +8,9 @@ import pytest from litellm.interactions.background_cost_polling import ( _SETTLED_KEY, BackgroundInteractionPollContext, + BackgroundSettlementContext, + SettlementOutcome, + get_settlement_store, maybe_schedule_background_interaction_cost_polling, maybe_settle_background_interaction_before_delete, poll_and_log_background_interaction_cost, @@ -344,6 +348,231 @@ async def test_poller_exits_without_billing_once_settled_elsewhere(): assert logging_obj.model_call_details.get("response_cost") is None +class _InMemorySettlementStore: + def __init__(self, persist_error: Optional[Exception] = None) -> None: + self.rows: dict[str, dict] = {} + self.contexts: dict[str, BackgroundSettlementContext] = {} + self.outcomes: dict[str, SettlementOutcome] = {} + self.settle_before_delete_calls: list[str] = [] + self.persist_error = persist_error + + async def persist_pending(self, context: BackgroundSettlementContext, timeout_at: datetime) -> None: + if self.persist_error is not None: + raise self.persist_error + existing = self.rows.get(context.interaction_id) + status = existing["status"] if existing is not None else "pending" + self.rows[context.interaction_id] = {"status": status, "timeout_at": timeout_at} + self.contexts[context.interaction_id] = context + + async def claim(self, interaction_id: str) -> bool: + row = self.rows.get(interaction_id) + if row is None: + return True + if row["status"] != "pending": + return False + row["status"] = "settled" + return True + + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: + self.outcomes[interaction_id] = outcome + + async def is_pending(self, interaction_id: str) -> bool: + row = self.rows.get(interaction_id) + return row is not None and row["status"] == "pending" + + async def settle_pending_before_delete(self, interaction_id: str) -> None: + self.settle_before_delete_calls.append(interaction_id) + + +def _attributed_litellm_params(reservation: Optional[dict] = None) -> dict: + metadata = { + "user_api_key": "hashed-virtual-key", + "user_api_key_user_id": "user-123", + "user_api_key_team_id": "team-456", + "deployment_model_name": "gemini-3-flash-preview", + } + if reservation is not None: + metadata["user_api_key_budget_reservation"] = reservation + return {"metadata": metadata} + + +def test_no_settlement_store_registered_by_default(): + assert get_settlement_store() is None + + +@pytest.mark.asyncio +async def test_poller_persists_settlement_context_with_attribution_and_reservation(): + store = _InMemorySettlementStore() + reservation = _reservation() + logging_obj = _logging_obj(litellm_params=_attributed_litellm_params(reservation)) + fetch, _ = _fetch_sequence(_response("completed", with_usage=True)) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch, store=store) + + context = store.contexts["interactions/bg-abc"] + assert context.attribution.user_api_key == "hashed-virtual-key" + assert context.attribution.user_api_key_user_id == "user-123" + assert context.attribution.user_api_key_team_id == "team-456" + assert context.model == "gemini-2.5-flash" + assert context.model_group == "gemini-3-flash-preview" + assert context.custom_llm_provider == "gemini" + assert context.call_type == "acreate_interaction" + assert context.litellm_call_id == "bg-interactions-call-id" + assert context.budget_reservation is not None + assert context.budget_reservation.reserved_cost == reservation["reserved_cost"] + assert store.rows["interactions/bg-abc"]["timeout_at"] > datetime.now(timezone.utc) + + +@pytest.mark.asyncio +async def test_poller_records_billed_outcome_so_sweep_cannot_rebill(): + store = _InMemorySettlementStore() + logging_obj = _logging_obj() + 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["response_cost"] > 0 + assert store.rows["interactions/bg-abc"]["status"] == "settled" + assert store.outcomes["interactions/bg-abc"] == "billed" + + +def _store_with_row_settled_elsewhere(interaction_id: str = "interactions/bg-abc") -> _InMemorySettlementStore: + store = _InMemorySettlementStore() + store.rows[interaction_id] = {"status": "settled", "timeout_at": datetime.now(timezone.utc)} + return store + + +@pytest.mark.asyncio +async def test_poller_skips_billing_and_finalizes_locally_when_store_claim_lost(): + store = _store_with_row_settled_elsewhere() + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + 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.get("response_cost") is None + assert reservation["finalized"] is True + assert store.outcomes == {} + + +@pytest.mark.asyncio +async def test_timeout_give_up_does_not_release_when_store_claim_lost(): + store = _store_with_row_settled_elsewhere() + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False)) + + await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=0.01), + fetch_interaction=fetch, + store=store, + ) + + assert reservation["finalized"] is True + assert store.outcomes == {} + assert logging_obj.model_call_details.get("response_cost") is None + + +@pytest.mark.asyncio +async def test_timeout_give_up_records_abandoned_outcome_when_claim_won(): + store = _InMemorySettlementStore() + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + + await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=0.01), + fetch_interaction=_fetch_sequence(_response("in_progress", with_usage=False))[0], + store=store, + ) + + assert reservation["finalized"] is True + assert store.rows["interactions/bg-abc"]["status"] == "settled" + assert store.outcomes["interactions/bg-abc"] == "abandoned" + + +@pytest.mark.asyncio +async def test_poller_continues_in_memory_when_persist_fails(): + store = _InMemorySettlementStore(persist_error=RuntimeError("db unavailable")) + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence(_response("completed", with_usage=True)) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch, store=store) + + assert len(calls) == 1 + assert logging_obj.model_call_details["response_cost"] > 0 + assert store.rows == {} + assert store.outcomes == {} + + +@pytest.mark.asyncio +async def test_poller_exits_early_on_fetch_error_once_settled_elsewhere(): + store = _store_with_row_settled_elsewhere() + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, calls = _fetch_sequence(RuntimeError("interaction deleted on another pod")) + + await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=1.0), + fetch_interaction=fetch, + store=store, + ) + + assert len(calls) == 1 + assert reservation["finalized"] is True + assert logging_obj.model_call_details.get("response_cost") is None + assert store.outcomes == {} + + +@pytest.mark.asyncio +async def test_delete_registry_miss_delegates_to_store(): + store = _InMemorySettlementStore() + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/created-on-another-pod", + store=store, + ) + + assert store.settle_before_delete_calls == ["interactions/created-on-another-pod"] + + +@pytest.mark.asyncio +async def test_delete_registry_hit_records_billed_outcome_in_store(): + store = _InMemorySettlementStore() + logging_obj = _logging_obj() + task = _register_poll(logging_obj) + fetch, _ = _fetch_sequence(_response("completed", with_usage=True)) + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=fetch, + store=store, + ) + + assert logging_obj.model_call_details["response_cost"] > 0 + assert store.outcomes["interactions/bg-abc"] == "billed" + await asyncio.wait_for(task, timeout=5) + + +def test_settlement_reservation_shape_matches_budget_reservation_module(): + from litellm.interactions.background_cost_polling import SettlementReservation + from litellm.proxy.spend_tracking.budget_reservation import _BudgetCounter, _counter_to_reservation_entry + + entry = _counter_to_reservation_entry( + counter=_BudgetCounter( + counter_key="spend:key:hashed-virtual-key", + max_budget=10.0, + fallback_spend=0.0, + entity_type="key", + entity_id="hashed-virtual-key", + ), + reserved_cost=0.05, + ) + reservation = {"reserved_cost": 0.05, "entries": [entry], "finalized": False, "input_cost": 0.001} + + assert SettlementReservation.model_validate(reservation).model_dump() == reservation + + @pytest.mark.asyncio async def test_schedule_respects_kill_switch(monkeypatch): import litellm.interactions.background_cost_polling as module diff --git a/tests/test_litellm/proxy/spend_tracking/test_background_settlement.py b/tests/test_litellm/proxy/spend_tracking/test_background_settlement.py new file mode 100644 index 00000000000..570b9019f19 --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_background_settlement.py @@ -0,0 +1,320 @@ +import asyncio +from datetime import datetime, timedelta, timezone +from typing import Optional + +import pytest + +from litellm.interactions.background_cost_polling import ( + BackgroundSettlementContext, + SettlementKeyAttribution, + SettlementOutcome, + SettlementReservation, +) +from litellm.proxy.spend_tracking.background_settlement import ( + PendingSettlementRow, + _parse_row, + _resolve_settlement_credentials, + rebuild_logging_for_settlement, + settle_row_before_delete, + sweep_pending_settlements, +) +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.types.interactions import InteractionsAPIResponse + +USAGE_BLOCK = { + "total_tokens": 175, + "total_input_tokens": 100, + "input_tokens_by_modality": [{"modality": "text", "tokens": 100}], + "total_cached_tokens": 0, + "total_output_tokens": 50, + "output_tokens_by_modality": [{"modality": "text", "tokens": 50}], + "total_tool_use_tokens": 0, + "total_thought_tokens": 25, +} + +INTERACTION_ID = "interactions/bg-foreign" + + +def _settlement_context( + reservation: Optional[SettlementReservation] = None, + model_group: Optional[str] = None, +) -> BackgroundSettlementContext: + return BackgroundSettlementContext( + interaction_id=INTERACTION_ID, + custom_llm_provider="gemini", + model="gemini-2.5-flash", + model_group=model_group, + litellm_call_id="original-create-call-id", + call_type="acreate_interaction", + attribution=SettlementKeyAttribution( + user_api_key="hashed-virtual-key", + user_api_key_user_id="user-123", + user_api_key_team_id="team-456", + ), + budget_reservation=reservation, + ) + + +def _reservation() -> SettlementReservation: + return SettlementReservation(reserved_cost=0.05, entries=[], finalized=False, input_cost=0.001) + + +def _row( + context: Optional[BackgroundSettlementContext] = None, + timeout_at: Optional[datetime] = None, +) -> PendingSettlementRow: + return PendingSettlementRow( + interaction_id=INTERACTION_ID, + context=context if context is not None else _settlement_context(), + timeout_at=timeout_at if timeout_at is not None else datetime.now(timezone.utc) + timedelta(hours=1), + ) + + +def _response(status: str, with_usage: bool) -> InteractionsAPIResponse: + return InteractionsAPIResponse( + id=INTERACTION_ID, + model="gemini-2.5-flash", + status=status, + steps=[], + usage=dict(USAGE_BLOCK) if with_usage else None, + ) + + +def _fetch_returning(item): + calls = [] + + async def fetch(context: BackgroundSettlementContext) -> InteractionsAPIResponse: + calls.append(context.interaction_id) + if isinstance(item, Exception): + raise item + return item + + return fetch, calls + + +class _FakeRowStore: + def __init__(self, rows: tuple[PendingSettlementRow, ...]) -> None: + self.rows = {row.interaction_id: row for row in rows} + self.status = {row.interaction_id: "pending" for row in rows} + self.outcomes: dict[str, SettlementOutcome] = {} + self.claims_won = 0 + + async def claim(self, interaction_id: str) -> bool: + if self.status.get(interaction_id) != "pending": + return False + self.status[interaction_id] = "settled" + self.claims_won += 1 + return True + + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: + self.outcomes[interaction_id] = outcome + + async def list_due(self, older_than: datetime, limit: int) -> tuple[PendingSettlementRow, ...]: + due = tuple(row for row_id, row in self.rows.items() if self.status[row_id] == "pending") + return due[:limit] + + +@pytest.mark.asyncio +async def test_rebuilt_logging_bills_with_original_attribution_and_request_id(): + logging_obj = rebuild_logging_for_settlement(_settlement_context()) + response = _response("completed", with_usage=True) + + await logging_obj.async_log_background_interaction_completion(result=response) + + assert logging_obj.model_call_details["response_cost"] > 0 + payload = get_logging_payload( + kwargs=logging_obj.model_call_details, + response_obj=response, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + ) + assert payload["request_id"] == INTERACTION_ID + assert payload["user"] == "user-123" + assert payload["team_id"] == "team-456" + assert payload["api_key"] + assert payload["spend"] > 0 + + +def test_rebuilt_logging_carries_reservation_for_reconcile(): + context = _settlement_context(reservation=_reservation()) + logging_obj = rebuild_logging_for_settlement(context) + + from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs + + metadata = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) + assert metadata["user_api_key_budget_reservation"]["reserved_cost"] == 0.05 + assert metadata["user_api_key_budget_reservation"]["finalized"] is False + + +@pytest.mark.asyncio +async def test_sweep_bills_terminal_row_exactly_once_across_concurrent_sweeps(): + store = _FakeRowStore(rows=(_row(),)) + fetch, calls = _fetch_returning(_response("completed", with_usage=True)) + + await asyncio.gather( + sweep_pending_settlements(store=store, fetch=fetch), + sweep_pending_settlements(store=store, fetch=fetch), + ) + + assert store.claims_won == 1 + assert store.outcomes == {INTERACTION_ID: "billed"} + assert store.status[INTERACTION_ID] == "settled" + assert len(calls) >= 1 + + +@pytest.mark.asyncio +async def test_sweep_abandons_row_past_timeout_without_fetching(): + store = _FakeRowStore(rows=(_row(timeout_at=datetime.now(timezone.utc) - timedelta(seconds=1)),)) + fetch, calls = _fetch_returning(_response("completed", with_usage=True)) + + await sweep_pending_settlements(store=store, fetch=fetch) + + assert calls == [] + assert store.outcomes == {INTERACTION_ID: "abandoned"} + assert store.status[INTERACTION_ID] == "settled" + + +@pytest.mark.asyncio +async def test_sweep_leaves_row_pending_on_fetch_error(): + store = _FakeRowStore(rows=(_row(),)) + fetch, calls = _fetch_returning(RuntimeError("provider unavailable")) + + await sweep_pending_settlements(store=store, fetch=fetch) + + assert len(calls) == 1 + assert store.status[INTERACTION_ID] == "pending" + assert store.outcomes == {} + + +@pytest.mark.asyncio +async def test_sweep_leaves_row_pending_while_interaction_still_running(): + store = _FakeRowStore(rows=(_row(),)) + fetch, calls = _fetch_returning(_response("in_progress", with_usage=False)) + + await sweep_pending_settlements(store=store, fetch=fetch) + + assert len(calls) == 1 + assert store.status[INTERACTION_ID] == "pending" + assert store.outcomes == {} + + +@pytest.mark.asyncio +async def test_delete_settlement_releases_when_interaction_not_terminal(): + row = _row(context=_settlement_context(reservation=_reservation())) + store = _FakeRowStore(rows=(row,)) + fetch, _ = _fetch_returning(_response("in_progress", with_usage=False)) + + await settle_row_before_delete(row=row, store=store, fetch=fetch) + + assert store.outcomes == {INTERACTION_ID: "released"} + assert store.status[INTERACTION_ID] == "settled" + + +@pytest.mark.asyncio +async def test_delete_settlement_bills_terminal_interaction(): + row = _row() + store = _FakeRowStore(rows=(row,)) + fetch, _ = _fetch_returning(_response("completed", with_usage=True)) + + await settle_row_before_delete(row=row, store=store, fetch=fetch) + + assert store.outcomes == {INTERACTION_ID: "billed"} + + +@pytest.mark.asyncio +async def test_delete_settlement_releases_on_fetch_error(): + row = _row(context=_settlement_context(reservation=_reservation())) + store = _FakeRowStore(rows=(row,)) + fetch, _ = _fetch_returning(RuntimeError("interaction already deleted")) + + await settle_row_before_delete(row=row, store=store, fetch=fetch) + + assert store.outcomes == {INTERACTION_ID: "released"} + assert store.status[INTERACTION_ID] == "settled" + + +@pytest.mark.asyncio +async def test_delete_settlement_noop_when_claim_lost(): + row = _row() + store = _FakeRowStore(rows=(row,)) + store.status[INTERACTION_ID] = "settled" + fetch, _ = _fetch_returning(_response("completed", with_usage=True)) + + await settle_row_before_delete(row=row, store=store, fetch=fetch) + + assert store.outcomes == {} + + +class _FakeRouter: + def __init__(self, deployment: Optional[dict] = None, error: Optional[Exception] = None) -> None: + self.deployment = deployment + self.error = error + self.requested_models: list[str] = [] + + async def async_get_available_deployment(self, model: str, request_kwargs: dict) -> dict: + self.requested_models.append(model) + if self.error is not None: + raise self.error + assert self.deployment is not None + return self.deployment + + +@pytest.mark.asyncio +async def test_credentials_resolved_from_router_deployment(): + router = _FakeRouter(deployment={"litellm_params": {"api_key": "resolved-key", "api_base": "https://resolved"}}) + + credentials = await _resolve_settlement_credentials( + _settlement_context(model_group="gemini-3-flash-preview"), + router=router, + ) + + assert credentials.api_key == "resolved-key" + assert credentials.api_base == "https://resolved" + assert router.requested_models == ["gemini-3-flash-preview"] + + +@pytest.mark.asyncio +async def test_credentials_fall_back_to_env_when_router_lookup_fails(): + router = _FakeRouter(error=RuntimeError("no deployments available")) + + credentials = await _resolve_settlement_credentials( + _settlement_context(model_group="gemini-3-flash-preview"), + router=router, + ) + + assert credentials.api_key is None + assert credentials.api_base is None + + +@pytest.mark.asyncio +async def test_credentials_fall_back_to_env_without_model_group(): + credentials = await _resolve_settlement_credentials(_settlement_context(model_group=None)) + + assert credentials.api_key is None + assert credentials.api_base is None + + +def test_parse_row_roundtrips_persisted_context_json(): + context = _settlement_context(reservation=_reservation(), model_group="gemini-3-flash-preview") + timeout_at = datetime.now(timezone.utc) + + row = _parse_row( + interaction_id=INTERACTION_ID, + context_value=context.model_dump_json(), + timeout_at=timeout_at, + ) + + assert row is not None + assert row.context == context + assert row.timeout_at == timeout_at + + +def test_parse_row_returns_none_for_corrupt_context(): + assert ( + _parse_row( + interaction_id=INTERACTION_ID, + context_value="{not json", + timeout_at=datetime.now(timezone.utc), + ) + is None + )