From faf4ce39850048aed7eb671b589b9a0d49bb7793 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 15:47:24 -0800 Subject: [PATCH] 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 | 8 ++----- litellm/responses/streaming_iterator.py | 24 +++++++++++++++++-- .../integration/mcp/test_mcp_llm_endpoints.py | 12 +++++++--- .../mcp/test_aresponses_api_with_mcp.py | 20 ++++++++++++---- 4 files changed, 49 insertions(+), 15 deletions(-) diff --git a/litellm/responses/mcp/stream_round_billing.py b/litellm/responses/mcp/stream_round_billing.py index 5cb256a8ac0..be331a91475 100644 --- a/litellm/responses/mcp/stream_round_billing.py +++ b/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 diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 249c6370b86..288e4a7119d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/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" diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index f6743a8751a..0c97ae62d07 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/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 diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 037004ab999..a39bc09cee6 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/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