fix(spend): write a background response's settling read in-process when the spend collector is on (#50957)

* fix(passthrough): bill a background response once however many times it is read

* fix(passthrough): keep the background response flag on spend events sent to the collector

* test(spend): type the background response test helper and its settlement parameter

* fix(proxy): keep a background response's settling read in-process instead of adding a spend event field

An older collector refuses a spend event with a field it does not know and drops the whole charge, so a newer
gateway sending to an older collector lost the background response's cost. The read that settles a background
response now stays in-process, the same way video status polls do, and the spend event keeps its shape

* fix(proxy): keep a batch retrieve in-process when message redaction hides the logged batch

* refactor(proxy): split the is_offloadable_success docstring into one line per reason

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -425,7 +425,7 @@ class _ProxyDBLogger(CustomLogger):
     async def async_log_success_event(
         self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime
     ) -> None:
-        if self.spend_event_producer is None or not is_offloadable_success(response_obj, kwargs.get("call_type")):
+        if self.spend_event_producer is None or not is_offloadable_success(response_obj, kwargs):
             await self._PROXY_track_cost_callback(dict(kwargs), response_obj, start_time, end_time)
             return
         event: Final = build_spend_event(

litellm/proxy/spend_tracking/spend_event.py

@@ -23,10 +23,15 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 import litellm
+from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.db.spend_log_tool_index import response_tool_call_names
 from litellm.types.interactions import InteractionsAPIResponse
 from litellm.types.llms.base import LiteLLMBaseModel
+from litellm.types.passthrough_endpoints.pass_through_endpoints import (
+    SETTLES_BACKGROUND_RESPONSE_KWARG,
+    BackgroundResponseSettlement,
+)
 from litellm.types.utils import LiteLLMBatch, Usage
 from litellm.types.videos.utils import is_video_status_call_type
 
@@ -206,10 +211,19 @@ class SpendEventDecodeError:
     reason: str
 
 
-def is_offloadable_success(response_obj: object, call_type: object) -> bool:
-    """Batch retrieves and interaction polls branch on the concrete response class, and a video status poll on
-    the status and video id the event leaves out, so they stay in-process."""
-    return not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES) and not is_video_status_call_type(call_type)
+def is_offloadable_success(response_obj: object, kwargs: ObjectMapping) -> bool:
+    """A success stays in-process when its cost logic reads something the spend event leaves out.
+
+    Batches and interaction polls need the concrete response class. A batch retrieve needs the polled batch id,
+    which message redaction can hide. A video status poll needs the status and the video id. A background
+    response's settling read needs the settlement flag."""
+    call_type: Final = kwargs.get("call_type")
+    return (
+        not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES)
+        and not is_batch_retrieve_call_type(call_type)
+        and not is_video_status_call_type(call_type)
+        and kwargs.get(SETTLES_BACKGROUND_RESPONSE_KWARG) is not BackgroundResponseSettlement.SETTLES
+    )
 
 
 def _json_fallback(value: object) -> str:

tests/unit/proxy/hooks/test_proxy_track_cost_callback.py

@@ -18,7 +18,9 @@ from fastapi import FastAPI, HTTPException
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.dual_cache import DualCache
+from litellm.litellm_core_utils.batch_utils import POLLED_BATCH_COST_IS_FINAL_KWARG, POLLED_BATCH_ID_KWARG
 from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
+from litellm.litellm_core_utils.redact_messages import perform_redaction
 from litellm.proxy._types import ProxyException, SpendLogsPayload, UserAPIKeyAuth
 from litellm.proxy.collector import SpendEventConsumer
 from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded
@@ -47,6 +49,10 @@ from litellm.proxy.utils import (
 )
 from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys
 from litellm.types.router import Deployment, LiteLLM_Params
+from litellm.types.passthrough_endpoints.pass_through_endpoints import (
+    SETTLES_BACKGROUND_RESPONSE_KWARG,
+    BackgroundResponseSettlement,
+)
 from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage
 from litellm.types.videos.main import VideoObject
 
@@ -3692,17 +3698,41 @@ def _completed_batch() -> LiteLLMBatch:
     )
 
 
+def _redacted_completed_batch() -> object:
+    return perform_redaction({}, _completed_batch())
+
+
 def _completed_video() -> VideoObject:
     return VideoObject(id="video_abc", object="video", status="completed")
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    ("call_type", "polled_job"),
-    [(CallTypes.aretrieve_batch.value, _completed_batch), ("avideo_status", _completed_video)],
+    ("job_kwargs", "polled_job"),
+    [
+        ({"call_type": CallTypes.aretrieve_batch.value}, _completed_batch),
+        (
+            {
+                "call_type": CallTypes.aretrieve_batch.value,
+                POLLED_BATCH_ID_KWARG: "batch_abc",
+                POLLED_BATCH_COST_IS_FINAL_KWARG: True,
+            },
+            _redacted_completed_batch,
+        ),
+        ({"call_type": "avideo_status"}, _completed_video),
+        (
+            {
+                "call_type": "pass_through_endpoint",
+                SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES,
+            },
+            _offload_response,
+        ),
+    ],
+    ids=["batch_retrieve", "redacted_batch_retrieve", "video_status_poll", "background_response_settling_read"],
 )
-async def test_async_log_success_event_keeps_job_polls_in_process(call_type, polled_job):
-    """The spend event leaves out the polled video's status and id, which decide whether the poll charges the job."""
+async def test_async_log_success_event_keeps_job_polls_in_process(job_kwargs, polled_job):
+    """The spend event leaves out the polled video's status and id and a background response's settlement flag,
+    which decide whether the poll charges the job and which spend row it writes."""
     producer = SpendEventProducer(
         address=UnixAddress(path="/nonexistent/spend.sock"),
         on_unavailable="drop",
@@ -3711,7 +3741,7 @@ async def test_async_log_success_event_keeps_job_polls_in_process(call_type, pol
         fallback=_no_fallback,
     )
     logger = ProxyDBLogger(producer)
-    kwargs = {**_offload_kwargs(), "call_type": call_type}
+    kwargs = {**_offload_kwargs(), **job_kwargs}
     completed_job = polled_job()
 
     with (

tests/unit/proxy/spend_tracking/test_spend_event.py

@@ -16,6 +16,10 @@ from litellm.proxy.spend_tracking.spend_event import (
     is_offloadable_success,
     spend_event_callback_args,
 )
+from litellm.types.passthrough_endpoints.pass_through_endpoints import (
+    SETTLES_BACKGROUND_RESPONSE_KWARG,
+    BackgroundResponseSettlement,
+)
 from litellm.types.utils import LiteLLMBatch, ModelResponse, Usage
 
 _BIG_PROMPT: Final = "x" * 20_000
@@ -194,11 +198,23 @@ def test_undecodable_line_is_an_error_value():
     assert isinstance(decode_spend_event(b"not json\n"), SpendEventDecodeError)
 
 
+@pytest.mark.parametrize(
+    ("settlement", "offloadable"),
+    [(BackgroundResponseSettlement.SETTLES, False), (None, True), ("settles", True)],
+    ids=["settling_read_stays_in_process", "other_read_is_offloaded", "lookalike_string_is_offloaded"],
+)
+def test_only_a_background_response_settling_read_stays_in_process(
+    settlement: BackgroundResponseSettlement | str | None, offloadable: bool
+) -> None:
+    kwargs: Final = {"call_type": "pass_through_endpoint", SETTLES_BACKGROUND_RESPONSE_KWARG: settlement}
+    assert is_offloadable_success(None, kwargs) is offloadable
+
+
 def test_batch_retrieves_and_video_status_polls_stay_in_process():
-    assert is_offloadable_success(_response(), "acompletion") is True
-    assert is_offloadable_success(None, None) is True
-    assert is_offloadable_success(None, "avideo_generation") is True
-    assert is_offloadable_success(None, "avideo_status") is False
+    assert is_offloadable_success(_response(), {"call_type": "acompletion"}) is True
+    assert is_offloadable_success(None, {}) is True
+    assert is_offloadable_success(None, {"call_type": "avideo_generation"}) is True
+    assert is_offloadable_success(None, {"call_type": "avideo_status"}) is False
     assert (
         is_offloadable_success(
             LiteLLMBatch(
@@ -210,7 +226,7 @@ def test_batch_retrieves_and_video_status_polls_stay_in_process():
                 object="batch",
                 status="completed",
             ),
-            "aretrieve_batch",
+            {"call_type": "aretrieve_batch"},
         )
         is False
     )