fix(responses): preserve native stream spend on guardrail rejection (#50286)

* fix(responses): preserve native stream spend on guardrail rejection

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(responses): type deferred logging dispatch once

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(responses): expect guarded rounds on one failure row

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* style(tests): format guarded billing regression

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* style(tests): preserve existing MCP unit formatting

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/responses/mcp/stream_round_billing.py

@@ -1,5 +1,5 @@
 import asyncio
-from collections.abc import Callable, Coroutine, Sequence
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 from datetime import datetime
 from typing import Final, TypeAlias
@@ -297,12 +297,8 @@ class StreamedRoundBilling:
         self._close(parked)
         self._bank((self._dispatch(row, billed, self._logging_obj.attempt_scoped_copy()),), row)
 
-    def _dispatch(
-        self, row: _RoundRow, billed: _BilledRounds, logger: LiteLLMLoggingObj
-    ) -> BilledRoundsDispatch | Coroutine[object, object, None]:
+    def _dispatch(self, row: _RoundRow, billed: _BilledRounds, logger: LiteLLMLoggingObj) -> BilledRoundsDispatch:
         dispatch: Final = self._log_success(_SuccessLog(logger, row, billed.chat_usage))
-        if billed.chat_usage is None:
-            return dispatch
         return BilledRoundsDispatch(dispatch, billed.spend)
 
     @staticmethod

litellm/responses/streaming_iterator.py

@@ -53,7 +53,7 @@ from litellm.litellm_core_utils.internal_call_metadata import (
     requests_background,
     without_budget_reservation,
 )
-from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend
+from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
 from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers
@@ -638,6 +638,24 @@ class BaseResponsesAPIStreamingIterator:
             self._handle_failure(e)
             raise
 
+    def _deferred_logging_dispatch(
+        self, logging_coroutine: Coroutine[object, object, None], logging_response: object
+    ) -> BilledRoundsDispatch | Coroutine[object, object, None]:
+        response: Final[object] = getattr(logging_response, "response", None)
+        if not isinstance(response, ResponsesAPIResponse):
+            return logging_coroutine
+        price: Final = (
+            UncachedPrice(0.0, None)
+            if self._completed_response_cache_hit
+            else self.logging_obj.uncached_price(response)
+        )
+        spend: Final = BilledRoundsSpend(
+            ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(response.usage),
+            price.cost or 0.0,
+            price.cost_breakdown,
+        )
+        return BilledRoundsDispatch(logging_coroutine, spend)
+
     def _log_completed_response(self, *, is_async: bool) -> None:
         if self._completed_response_logged:
             return
@@ -662,7 +680,9 @@ class BaseResponsesAPIStreamingIterator:
             if deferred_dispatch_armed:
                 # End-of-stream guardrail scans write guardrail_information after
                 # the terminal event; dispatching now would snapshot metadata early.
-                self.logging_obj._deferred_stream_complete_args = (logging_coroutine,)
+                self.logging_obj._deferred_stream_complete_args = (
+                    self._deferred_logging_dispatch(logging_coroutine, logging_response),
+                )
             else:
                 spawn_logging_task(
                     asyncio.get_running_loop(), logging_coroutine, task_name="Responses stream success logging"

tests/integration/mcp/test_mcp_llm_endpoints.py

@@ -42,7 +42,6 @@ from openai.types.chat import ChatCompletionMessageParam
 from openai.types.responses.tool_param import Mcp
 from pydantic import JsonValue, TypeAdapter
 
-
 pytestmark: Final = [pytest.mark.xdist_group("mcp_llm_private_gateway"), pytest.mark.shares_worker_state]
 
 
@@ -589,7 +588,11 @@ def _upstream_model(surface: Surface) -> str:
     }[surface]
 
 
-PRICES: Final[Mapping[str, JsonValue]] = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002, "num_retries": 0}
+PRICES: Final[Mapping[str, JsonValue]] = {
+    "input_cost_per_token": 0.001,
+    "output_cost_per_token": 0.002,
+    "num_retries": 0,
+}
 
 
 @contextmanager
@@ -917,7 +920,9 @@ def test_an_open_server_in_another_gateway_cannot_widen_this_gateways_tools(gate
             register_mcp(scenario, local, "open" + uuid.uuid4().hex[:8], allow_all_keys=True)
             listed, called = official_client_outcomes(shared, key, "/mcp", f"{foreign_alias}-add", ADD)
             assert set(listed.tools) == {
-                f"{foreign_alias}-add", f"{foreign_alias}-multiply", f"{foreign_alias}-fail"
+                f"{foreign_alias}-add",
+                f"{foreign_alias}-multiply",
+                f"{foreign_alias}-fail",
             }, listed.tools
             assert called.ok and called.text == str(ADD["a"] + ADD["b"]), called
             assert [call["body"]["params"]["name"] for call in _peer_add_calls(foreign)] == ["add"]
@@ -1528,6 +1533,7 @@ def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_
         one_round_rows: Final = _model_rows(
             one_round.headers["x-litellm-call-id"], _logs_a_failure_billing(_price(FIRST_ROUND) + FIXED_MARGIN)
         )
+        assert _statuses(one_round_rows) == ["failure"], one_round_rows
         rig.wire.drain()
         response: Final = rig.send(key, [AUTO], stream=True, guardrails=[STREAM_BLOCK])
         assert STREAM_BLOCK_DETAIL in response.text, response.text

tests/unit/responses/mcp/test_aresponses_api_with_mcp.py

@@ -2722,7 +2722,7 @@ async def test_stream_auto_execute_bills_a_follow_up_cut_after_streamed_text_wit
 
 
 @pytest.mark.asyncio
-async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_the_success_row(
+async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_one_failure_row(
     respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
 ):
     from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
@@ -2743,10 +2743,22 @@ async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_r
     if not ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked):
         ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
     await logging_obj.async_failure_handler(blocked, "blocked by output-policy")
-    await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0)
+    await _settle(lambda: len(recorder.failed) > 0)
 
-    _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events))
-    assert [row["response_cost"] for row in recorder.failed] == [0.0]
+    assert recorder.logged == ()
+    assert len(recorder.failed) == 1
+    row: Final = recorder.failed[0]
+    charged_once: Final = _round_cost(_ROUND_ONE_BODY) + _round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN
+    assert row["status"] == "failure"
+    assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5)
+    assert row["response_cost"] == pytest.approx(charged_once)
+    assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once)
+    client: Final = _final_response(events)
+    assert "message" in _output_types(client.output)
+    assert "function_call" not in _output_types(client.output)
+    assert client.usage is not None
+    assert (client.usage.input_tokens, client.usage.output_tokens) == (71 + 149, 19 + 5)
+    assert client.usage.cost == pytest.approx(charged_once)
 
 
 @pytest.mark.asyncio