From 0e24bbb8a2c4e250cde60ca92df92ea2ffe3de60 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:56:38 +0000 Subject: [PATCH] fix(responses): bill every round of a streamed MCP auto-execute on one spend row (#47445) * test(responses): cover billing every round of a streamed MCP auto-execute * fix(responses): bill every round of a streamed MCP auto-execute on one spend row * fix(responses): bill every round of a streamed MCP auto-execute on bridged providers A streamed /v1/responses MCP auto-execute on a provider bridged through chat completions (for example Anthropic) billed only round 1, because each round's chat stream dispatched on the shared logging object and the first dispatch won. Bridged rounds now go through the same one-row billing as native rounds. A round is billed as soon as its chat stream ends, before its closing events reach the client. The row carries the final response, usage and cost summed over every round, and the fixed margin once. A round served from the response cache is billed at zero. A follow-up chat stream that breaks bills the paid rounds on its failure row. A single-round stream logs exactly as before. * fix(responses): charge one fixed margin and keep bridged usage fields on streamed MCP rounds A follow-up round that fails after earlier rounds were billed now charges the per-request fixed margin once, and its failure row carries no cost breakdown that would disagree with the summed cost. A bridged multi-round stream logs a plain cache_hit flag like a single-round stream does, and its usage keeps the summed provider usage fields of every round. A new test covers a stream dropped after its event loop closed, which must release its parked log without an unawaited coroutine warning. * fix(responses): price each streamed MCP round on its own and bill every round when a guardrail blocks Each round of a streamed MCP auto-execute is now priced from its own uncached cost, so a provider that states its cost (xAI) charges the stated cost of every round plus the fixed margin once, and the cost breakdown always agrees with the billed cost. Summed usage goes through a JSON-safe dump, so a nested provider field such as xAI's context_details no longer breaks a failed stream's response.failed event. A post-call guardrail that blocks a multi-round bridged stream now drops the parked success log and bills every round on the failure row, and a failed follow-up round's failure row carries the summed breakdown of both rounds. * fix(responses): keep earlier MCP rounds on their success row when a later round fails When a later round of a streamed MCP auto-execute fails, the rounds before it are now logged on their success row, and the failure row bills only the failed round with no second fixed margin. That failure log runs on its own copy of the request's logging state, so its usage and cost never leak into the earlier rounds' row. A post-call guardrail that blocks a native multi-round stream keeps the summed success row beside the $0 failure row, as at the merge base, while a bridged stream still bills every round on its failure row. * fix(responses): keep a failed later MCP round on its own failure row A streamed MCP auto-execute whose later round failed wrote only the earlier rounds' success row. The failure was logged on a copy of the logging state, while the proxy builds its failure row from the request's own state, which by then held the success payload, so the failure row took the success row's id and was dropped. Now the failure is logged on the request's own state as at the merge base, and once a later round starts, the success row for the rounds before it is logged from a copy of the state taken at that moment, so the two rows never mix. * fix(proxy): let an MCP tool call reuse its request's max_parallel_requests slot * test(proxy): create the other-request task in a way Python 3.10 supports and mark its bindings Final * fix(proxy): keep a reused max_parallel slot until the request's last round An MCP tool call that reuses its request's max_parallel_requests slot no longer loses it to the first round's late success callback, so the key stays counted for the rounds still to run. The reused call also stops replacing the request's rate limit headers. * fix(responses): keep an unpriced streamed MCP round unpriced instead of billing it as free * fix(responses): log a bridged MCP stream as a chat completion with every round's usage --- litellm/litellm_core_utils/litellm_logging.py | 48 +- .../litellm_core_utils/streaming_handler.py | 20 +- litellm/proxy/common_request_processing.py | 3 + .../hooks/parallel_request_limiter_v3.py | 67 +- litellm/proxy/utils.py | 18 +- .../streaming_iterator.py | 14 +- .../responses/mcp/mcp_streaming_iterator.py | 104 +- litellm/responses/mcp/round_billing.py | 23 +- litellm/responses/mcp/stream_round_billing.py | 376 +++++++ litellm/responses/streaming_iterator.py | 18 +- .../test_litellm_logging.py | 21 + .../hooks/test_parallel_request_limiter_v3.py | 171 ++++ .../mcp/test_aresponses_api_with_mcp.py | 948 +++++++++++++++++- .../test_streaming_iterator_error_events.py | 81 +- 14 files changed, 1848 insertions(+), 64 deletions(-) create mode 100644 litellm/responses/mcp/stream_round_billing.py diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index eb2e7b8c240..a59672b5cb2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -12,7 +12,7 @@ import subprocess import sys import time import traceback -from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime as dt_object from functools import lru_cache @@ -575,6 +575,27 @@ class UncachedPrice: cost_breakdown: CostBreakdown | None +@dataclass(frozen=True, slots=True) +class BilledRoundsSpend: + usage: Usage + cost: float + cost_breakdown: CostBreakdown | None + + +@dataclass(frozen=True, slots=True) +class BilledRoundsDispatch: + """A parked success log for a stream bridged through chat completions that billed several rounds, and the spend + its failure log bills when a post-call hook rejects the response, as it does for a rejected one-round bridged + stream.""" + + dispatch: Coroutine[object, object, None] + spend: BilledRoundsSpend + + +def _fixed_margin(cost_breakdown: CostBreakdown | None) -> float: + return 0.0 if cost_breakdown is None else cost_breakdown.get("margin_fixed_amount", 0.0) + + def _timestamp_seconds(moment: object) -> float | None: if isinstance(moment, datetime.datetime): return moment.timestamp() @@ -2421,6 +2442,31 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def record_failure_after_earlier_rounds( + self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float + ) -> None: + """Bill a failed round on the failure log after the rounds of its request before it were billed on the success + log. + + The fixed margin is a per-request fee, so it is charged once: the failed round's margin is dropped when the + earlier rounds already paid one.""" + from litellm.responses.mcp.round_billing import breakdown_less_repeated_fixed_margin + + failed_breakdown: Final = None if self.cost_breakdown is earlier.cost_breakdown else self.cost_breakdown + repeated_fixed_margin: Final = ( + 0.0 + if failed_breakdown is None + else min(_fixed_margin(earlier.cost_breakdown), _fixed_margin(failed_breakdown), earlier.cost, failed_cost) + ) + self.model_call_details["combined_usage_object"] = failed_usage + self.model_call_details["response_cost"] = failed_cost - repeated_fixed_margin + self.cost_breakdown = breakdown_less_repeated_fixed_margin(failed_breakdown, repeated_fixed_margin) + + def record_billed_rounds_for_failure(self, spend: BilledRoundsSpend) -> None: + """Bill every round of a fully streamed request on the failure log when a post-call hook rejects it.""" + self.record_partial_usage_for_failure(spend.usage, spend.cost) + self.cost_breakdown = spend.cost_breakdown + def attempt_scoped_copy(self) -> "Logging": """A copy with its own model_call_details, so logging one attempt of a request through it leaves the state the request's other attempts share untouched.""" diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index b210085c5cc..5e43157a55d 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -10,7 +10,7 @@ import traceback from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final, NoReturn, Protocol, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeVar, cast import anyio import httpx @@ -60,6 +60,9 @@ from .exception_mapping_utils import exception_type, with_deployment_api_base from .llm_response_utils.get_api_base import get_api_base from .rules import Rules +if TYPE_CHECKING: + from .litellm_logging import BilledRoundsSpend + # Constants for special delta attribute names AUDIO_ATTRIBUTE: Final = "audio" IMAGE_ATTRIBUTE: Final = "images" @@ -259,6 +262,8 @@ def _provider_hidden_params( class CustomStreamWrapper: + _earlier_rounds_spend: "BilledRoundsSpend | None" = None + @property def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation return self._hidden_params @@ -2399,6 +2404,9 @@ class CustomStreamWrapper: return await asyncio.wait((self._failure_logging,), timeout=timeout) + def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None: + self._earlier_rounds_spend = earlier + def _record_partial_usage_for_failure(self) -> None: """ A stream that breaks mid-flight still billed the provider for the chunks @@ -2425,10 +2433,12 @@ class CustomStreamWrapper: if self.model: partial_response.model = self.model backfill_missing_cache_usage_fields(usage) - self.logging_obj.model_call_details["combined_usage_object"] = usage - self.logging_obj.model_call_details["response_cost"] = ( - self.logging_obj.response_cost_calculator(result=partial_response) or 0.0 - ) + partial_cost: Final = self.logging_obj.response_cost_calculator(result=partial_response) or 0.0 + if self._earlier_rounds_spend is None: + self.logging_obj.model_call_details["combined_usage_object"] = usage + self.logging_obj.model_call_details["response_cost"] = partial_cost + return + self.logging_obj.record_failure_after_earlier_rounds(self._earlier_rounds_spend, usage, partial_cost) except Exception as recover_error: verbose_logger.debug( "could not recover partial usage for interrupted stream: %s", diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index cf616873a2e..6fda5780f63 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -67,6 +67,7 @@ from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, ) from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.litellm_core_utils.llm_response_utils.get_headers import ( @@ -3996,6 +3997,8 @@ class ProxyBaseLLMRequestProcessing: async def _on_deferred_native_stream_complete(*args: object) -> None: match args: + case (BilledRoundsDispatch(dispatch=logging_coroutine),): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine): GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) case (assembled_response, cache_hit): diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 4f2f7f9c256..a8ba8d65785 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -695,6 +695,7 @@ class RequestRateLimiterStash: owner_litellm_call_id: str | None = None rate_limit_response: RateLimitResponse | None = None parallel_slot: ParallelSlotAcquisition | None = None + parallel_slot_reused_at: datetime | None = None parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False) reserved_tokens: int = 0 reserved_model: RateLimitedModel | None = None @@ -791,6 +792,19 @@ def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: return stash +def _ended_before_slot_reuse(ended_at: datetime | None, reused_at: datetime | None) -> bool: + if ended_at is None or reused_at is None or (ended_at.tzinfo is None) != (reused_at.tzinfo is None): + return False + return ended_at <= reused_at + + +def _callback_kwargs_mapping(kwargs: object) -> Mapping[str, object]: + try: + return _REQUEST_RATE_LIMIT_DATA.validate_python(kwargs) + except ValidationError: + return MappingProxyType({}) + + def get_request_stash_for_call(litellm_call_id: str | None) -> RequestRateLimiterStash | None: stash: Final = _request_stash.get() if stash is None: @@ -1652,6 +1666,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): read_only: bool = False, skip_tpm_check: bool = False, parallel_slot_id: str | None = None, + skip_parallel_check: bool = False, ) -> RateLimitResponse: """ Check if any of the rate limit descriptors should be rate limited. @@ -1668,6 +1683,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ``reserve_tpm_tokens`` reservation path should set this to avoid the +1-per-key Lua / in-memory increment double-charging the tokens counter. + skip_parallel_check: If True, ``max_parallel_requests`` gauges are + neither checked nor acquired, for a call that runs inside a + request already holding its slot. ``max_parallel_requests`` descriptors are enforced by the dedicated concurrency-gauge path (``_check_parallel_request_gauges``), never by @@ -1760,7 +1778,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if windowed_response["overall_code"] == "OVER_LIMIT": return windowed_response - if not gauges: + if not gauges or skip_parallel_check: return windowed_response gauge_response: Final = await self._check_parallel_request_gauges( @@ -2023,12 +2041,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span: Span | None, *, in_logging_callback: bool = False, + round_ended_at: datetime | None = None, ) -> None: if stash is None: return async with stash.parallel_slot_release_lock: acquisition: Final = stash.parallel_slot - if acquisition is None: + if acquisition is None or _ended_before_slot_reuse(round_ended_at, stash.parallel_slot_reused_at): return deferred: Final = in_logging_callback and await self._defer_parallel_slot_release( acquisition, parent_otel_span @@ -3947,6 +3966,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): tpm_reservation_amount: int, call_type: str | None = None, endpoint_type: EndpointType = EndpointType.GENERIC, + reuses_request_slot: bool = False, ) -> None: """ Reserve project-scoped ITPM/OTPM tokens (Bedrock Mantle-style @@ -4026,7 +4046,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) stash.reservation_released = True - await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) + if not reuses_request_slot: + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( response=io_response, descriptors=descriptors, @@ -4187,6 +4208,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook") stash: Final = claim_request_stash_for_data(data) + reuses_request_slot: Final = call_type == CallTypes.call_mcp_tool.value and stash.parallel_slot is not None + if reuses_request_slot: + stash.parallel_slot_reused_at = self._get_current_time() ######################################################### # Check if the call type has a specific rate limiter @@ -4231,11 +4255,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # then because _reserve_project_io_tokens_or_raise below charges # them unconditionally and counting them here too would # double-charge every request. - parallel_counter_keys: Final = [ - self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests") - for d in descriptors - if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None - ] + parallel_counter_keys: Final = ( + [] + if reuses_request_slot + else [ + self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests") + for d in descriptors + if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None + ] + ) parallel_slot_id: Final = uuid.uuid4().hex if parallel_counter_keys else None first_pass_descriptors: Final = ( @@ -4250,6 +4278,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, skip_tpm_check=self.tpm_reservation_enabled, parallel_slot_id=parallel_slot_id, + skip_parallel_check=reuses_request_slot, ) if response["overall_code"] == "OVER_LIMIT": @@ -4259,7 +4288,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model=requested_model, ) else: - stash.rate_limit_response = response + if not reuses_request_slot: + stash.rate_limit_response = response if parallel_slot_id is not None: stash.parallel_slot = ParallelSlotAcquisition( slot_id=parallel_slot_id, @@ -4357,7 +4387,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): raw_estimate=estimated_tokens, ) if stash.ptu_ceiling is not None and ptu_estimated_tokens > stash.ptu_ceiling.tpm_limit: - await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) + if not reuses_request_slot: + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) raise PTUShareExceededError(stash.ptu_ceiling, ptu_estimated_tokens) tpm_response: Final = await self.reserve_tpm_tokens( descriptors=descriptors, @@ -4367,7 +4398,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) if tpm_response["overall_code"] == "OVER_LIMIT": - await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) + if not reuses_request_slot: + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( response=tpm_response, descriptors=descriptors, @@ -4418,6 +4450,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): tpm_reservation_amount=tpm_reservation_amount, call_type=call_type, endpoint_type=endpoint_type, + reuses_request_slot=reuses_request_slot, ) def _create_pipeline_operations( @@ -5332,8 +5365,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): try: verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING") - stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) - await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span, in_logging_callback=True) + callback_kwargs: Final = _callback_kwargs_mapping(kwargs) + stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(callback_kwargs)) + await self._release_stashed_parallel_slot( + stash, + litellm_parent_otel_span, + in_logging_callback=True, + round_ended_at=end_time + if isinstance(end_time, datetime) and not callback_kwargs.get("stream") + else None, + ) pipeline_operations: Final = self._build_success_event_pipeline_operations( kwargs=kwargs, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ea1044d4a55..6c9b23860da 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -158,7 +158,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.duration_parser import get_budget_window_start from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int -from litellm.litellm_core_utils.litellm_logging import Logging, deployment_pricing_model_info +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, Logging, deployment_pricing_model_info from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -4732,19 +4732,25 @@ class ProxyLogging: @staticmethod def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool: - """Drop the parked success dispatch for an assembled chat stream that ends in an error - ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead. + """Drop the parked success dispatch for an assembled chat stream or a stream that billed several rounds when it + ends in an error ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead. Returns False when the parked dispatch should still be flushed by the caller.""" logging_obj: Final = request_data.get("litellm_logging_obj") if not isinstance(logging_obj, Logging): return False _args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None) - assembled: Final = _args[0] if _args else None - if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS) or not isinstance(assembled, ModelResponse): + if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS): return False + match _args: + case (ModelResponse() as assembled, *_): + logging_obj.record_assembled_response_for_failure(assembled) + case (BilledRoundsDispatch() as billed,): + billed.dispatch.close() + logging_obj.record_billed_rounds_for_failure(billed.spend) + case _: + return False logging_obj._on_deferred_stream_complete = None logging_obj._deferred_stream_complete_args = None - logging_obj.record_assembled_response_for_failure(assembled) return True async def arelease_max_parallel_requests_on_disconnect( diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 84d4648f75f..4ef68aa072a 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -5,7 +5,8 @@ from itertools import filterfalse from typing import Any, Final, cast import litellm -from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params +from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params, served_from_cache +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend from litellm.main import stream_chunk_builder from litellm.responses.litellm_completion_transformation.custom_tools import ( build_tool_call_item_kwargs, @@ -1363,6 +1364,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return False return not any(getattr(part, "text", None) for part in getattr(item, "content", None) or ()) + def finished_round_event(self) -> ResponseCompletedEvent | ResponseIncompleteEvent | None: + if not isinstance(self.litellm_model_response, ModelResponse): + return None + finished: Final = self._emit_terminal_response_event(self.litellm_model_response) + if finished is not None and served_from_cache(self.litellm_custom_stream_wrapper): + finished.response.hidden_params = {**finished.response.hidden_params, "cache_hit": True} + return finished + + def bill_failure_after_earlier_rounds(self, earlier: BilledRoundsSpend) -> None: + self.litellm_custom_stream_wrapper.bill_failure_after_earlier_rounds(earlier) + def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: """ The response.incomplete event to bill when the client went away mid-stream. diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 73cfba8dca4..030911238ec 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -3,7 +3,9 @@ from typing import TYPE_CHECKING, Any, Final, cast from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.responses.litellm_completion_transformation.streaming_iterator import LiteLLMCompletionStreamingIterator from litellm.responses.mcp.request_context import MCPRequestContext +from litellm.responses.mcp.stream_round_billing import StreamedRoundBilling from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( BaseLiteLLMOpenAIResponseObject, @@ -352,6 +354,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self._round_max_output_index = -1 self._composed_output: list[object] = [] # mutable-ok: grows as each round finishes self._pending_mcp_call_items: list[dict[str, object]] = [] # mutable-ok: grows per executed tool + self._iteration_started = False + self._round_billing: StreamedRoundBilling | None = None + self._bridged_round_billed: LiteLLMCompletionStreamingIterator | None = None def _extract_mcp_headers_from_params(self) -> None: """Extract MCP headers from original request params to pass to tool calls""" @@ -436,7 +441,15 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): return self async def __anext__(self) -> ResponsesAPIStreamingResponse: - chunk: Final = await self._anext_impl() + if not self._iteration_started: + self._iteration_started = True + self._round_billing = self._bill_rounds_on_one_row() + try: + chunk: Final = await self._anext_impl() + except Exception: + if self._round_billing is not None: + self._round_billing.finish() + raise sequence_number: Final = getattr(chunk, "sequence_number", None) if isinstance(sequence_number, int): if sequence_number <= self._last_sequence_number and self._last_sequence_number > 0: @@ -446,6 +459,49 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self._last_sequence_number = max(self._last_sequence_number, sequence_number) return chunk + def __del__(self) -> None: + round_billing: Final[object] = getattr(self, "_round_billing", None) + if isinstance(round_billing, StreamedRoundBilling): + round_billing.release_abandoned() + + def _bill_rounds_on_one_row(self) -> StreamedRoundBilling | None: + if not self.should_auto_execute: + return None + if isinstance(self.base_iterator, LiteLLMCompletionStreamingIterator): + return StreamedRoundBilling(self.base_iterator.litellm_logging_obj) + if isinstance(self.base_iterator, BaseResponsesAPIStreamingIterator): + return StreamedRoundBilling(self.base_iterator.logging_obj) + return None + + def _restore_round_headers(self, logged_event: ResponsesAPIStreamingResponse) -> None: + if isinstance(self.base_iterator, LiteLLMCompletionStreamingIterator): + return + if isinstance(self.base_iterator, BaseResponsesAPIStreamingIterator): + self.base_iterator._restore_provider_response_headers(logged_event) + + def _round_billed_early(self) -> bool: + return self._bridged_round_billed is not None and self._bridged_round_billed is self.base_iterator + + def _bill_finished_bridged_round(self) -> None: + base: Final = self.base_iterator + if self._round_billing is None or not isinstance(base, LiteLLMCompletionStreamingIterator): + return + if self._round_billed_early(): + return + finished: Final = base.finished_round_event() + if finished is None: + return + self._bridged_round_billed = base + round_output: Final = _output_items(finished.response) + self._round_billing.complete_round( + finished, + finished.response, + round_output + if finished.type == ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE + else (*self._composed_output, *round_output), + self._restore_round_headers, + ) + async def _anext_impl(self) -> ResponsesAPIStreamingResponse: """ Phase-based streaming: @@ -617,6 +673,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): follow-up rounds past the items already emitted, and lists every round's items on the final response.completed. """ + self._bill_finished_bridged_round() chunk_type: Final[object] = getattr(chunk, "type", None) if self._round_index > 0 and chunk_type in ( ResponsesAPIStreamEvents.RESPONSE_CREATED, @@ -630,12 +687,21 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if self._output_index_offset: _set_event_field(chunk, "output_index", output_index + self._output_index_offset) + if self.should_auto_execute and chunk_type == ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE: + return self._charge_incomplete_round(chunk) if not (self.should_auto_execute and self._is_response_completed(chunk)): return chunk response_obj: Final[object] = getattr(chunk, "response", None) if isinstance(response_obj, ResponsesAPIResponse): self.collected_response = response_obj + if self._round_billing is not None and not self._round_billed_early(): + self._round_billing.complete_round( + chunk, + response_obj, + (*self._composed_output, *_output_items(response_obj)), + self._restore_round_headers, + ) # Move to tool execution phase after this chunk self.phase = "tool_execution" await self._generate_tool_execution_events() @@ -645,13 +711,33 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if self._follow_up_pending(): self._absorb_round(response_obj) return None - if self._composed_output: - merged_output: Final[list[object]] = [ # mutable-ok: the response model declares output as a list - *self._composed_output, - *_output_items(response_obj), - ] - merged_response: Final = response_obj.model_copy(update={"output": merged_output}) - _set_event_field(chunk, "response", merged_response) + merged_response: Final = ( + response_obj.model_copy(update={"output": [*self._composed_output, *_output_items(response_obj)]}) + if self._composed_output + else response_obj + ) + client_response: Final = ( + merged_response + if self._round_billing is None + else self._round_billing.charge_final_round(chunk, merged_response, self._restore_round_headers) + ) + if client_response is not response_obj: + _set_event_field(chunk, "response", client_response) + return chunk + + def _charge_incomplete_round(self, chunk: ResponsesAPIStreamingResponse) -> ResponsesAPIStreamingResponse: + response_obj: Final[object] = getattr(chunk, "response", None) + if self._round_billing is None or not isinstance(response_obj, ResponsesAPIResponse): + return chunk + if not self._round_billed_early(): + self._round_billing.complete_round( + chunk, response_obj, _output_items(response_obj), self._restore_round_headers + ) + client_response: Final = self._round_billing.charge_final_round( + chunk, response_obj, self._restore_round_headers + ) + if client_response is not response_obj: + _set_event_field(chunk, "response", client_response) return chunk async def _process_base_iterator_chunk(self) -> ResponsesAPIStreamingResponse: @@ -949,6 +1035,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.collected_response = None self._cached_response_id = None self._round_index += 1 + if self._round_billing is not None: + self._round_billing.record_billed_rounds(follow_up_response) except Exception as e: verbose_logger.error("Error creating follow-up iterator: %s", e) diff --git a/litellm/responses/mcp/round_billing.py b/litellm/responses/mcp/round_billing.py index ae34ab31633..93769fbe82a 100644 --- a/litellm/responses/mcp/round_billing.py +++ b/litellm/responses/mcp/round_billing.py @@ -4,7 +4,7 @@ from contextlib import contextmanager from types import MappingProxyType from typing import Final, TypeVar -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from pydantic_core import to_jsonable_python from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params @@ -46,15 +46,18 @@ def _summed_mapping(first: Mapping[str, object], final: Mapping[str, object]) -> return MappingProxyType({key: _summed_entry(key, first, final) for key in keys}) +def _summed_usage(first: BaseModel, final: BaseModel) -> object: + return to_jsonable_python(_summed_mapping(first.model_dump(), final.model_dump()), fallback=dict) + + def merged_round_usage(first: ResponseAPIUsage | None, final: ResponseAPIUsage | None) -> ResponseAPIUsage | None: if first is None or final is None: return final if first is None else first - return ResponseAPIUsage.model_validate(_summed_mapping(first.model_dump(), final.model_dump())) + return ResponseAPIUsage.model_validate(_summed_usage(first, final)) -def _merged_chat_usage(first: Usage, final: Usage) -> Usage: - summed: Final = _summed_mapping(first.model_dump(), final.model_dump()) - return Usage.model_validate(to_jsonable_python(summed, fallback=dict)) +def merged_chat_usage(first: Usage, final: Usage) -> Usage: + return Usage.model_validate(_summed_usage(first, final)) def _chat_usage(round_response: ResponsesAPIResponse | ModelResponse) -> Usage | None: @@ -82,7 +85,7 @@ def _usage_of_both_rounds( final_usage: Final = _chat_usage(final) if first_usage is None or final_usage is None: return final_usage if first_usage is None else first_usage - return _merged_chat_usage(first_usage, final_usage) + return merged_chat_usage(first_usage, final_usage) def _repeated_fixed_margin( @@ -105,6 +108,14 @@ def _less_repeated_fixed_margin(summed: Mapping[str, object], repeated: float) - return MappingProxyType({**summed, **charged_once}) +def breakdown_less_repeated_fixed_margin( + breakdown: CostBreakdown | None, repeated_fixed_margin: float +) -> CostBreakdown | None: + if breakdown is None: + return None + return _COST_BREAKDOWN.validate_python(_less_repeated_fixed_margin(breakdown, repeated_fixed_margin)) + + def summed_cost_breakdown( first: ResponsesAPIResponse | ModelResponse, first_breakdown: CostBreakdown | None, diff --git a/litellm/responses/mcp/stream_round_billing.py b/litellm/responses/mcp/stream_round_billing.py new file mode 100644 index 00000000000..443147230bf --- /dev/null +++ b/litellm/responses/mcp/stream_round_billing.py @@ -0,0 +1,376 @@ +import asyncio +from collections.abc import Callable, Coroutine, Sequence +from dataclasses import dataclass +from datetime import datetime +from typing import Final, Protocol, TypeAlias + +from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, spawn_logging_task +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamingResponse +from litellm.types.utils import CostBreakdown, ModelResponse, Usage + +Parked: TypeAlias = tuple[object, ...] +RestoresHeaders: TypeAlias = Callable[[ResponsesAPIStreamingResponse], None] + +_PARKED_DISPATCH: Final = "_deferred_stream_complete_args" +_DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" +_SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) + + +class _DispatchesSuccessHandlers(Protocol): + async def dispatch_success_handlers( + self, + result: object = None, + start_time: object = None, + end_time: object = None, + cache_hit: object = None, + prefer_async_handlers: bool = False, + ) -> None: ... + + +def _success_dispatcher(logger: LiteLLMLoggingObj) -> _DispatchesSuccessHandlers: + return logger + + +def _event_loop_running() -> bool: + try: + asyncio.get_running_loop() + except RuntimeError: + return False + return True + + +def _priced(response: ResponsesAPIResponse, price: UncachedPrice) -> ResponsesAPIResponse: + priced: Final = response.model_copy() + priced.hidden_params = {**response.hidden_params, "response_cost": price.cost} + return priced + + +def _assembled(parked: Parked | None) -> ModelResponse | None: + match parked: + case (ModelResponse() as assembled, bool()): + return assembled + case _: + return None + + +def _bridged_round_usage(parked: Parked | None) -> Usage | None: + usage: Final[object] = getattr(_assembled(parked), "usage", None) + return usage if isinstance(usage, Usage) else None + + +def _usage_of_bridged_rounds(billed: Usage | None, round_usage: Usage | None, cost: object) -> Usage | None: + if billed is None or round_usage is None: + return None + return merged_chat_usage(billed, round_usage).model_copy(update={"cost": cost}) + + +def _logged_cache_hit(response: ResponsesAPIResponse, chat_usage: Usage | None) -> bool | None: + if chat_usage is not None: + return served_from_cache(response) + return served_from_cache(response) or None + + +@dataclass(frozen=True, slots=True) +class _RoundRow: + logged: ResponsesAPIStreamingResponse | ModelResponse + cache_hit: bool | None + end_time: datetime + + +@dataclass(frozen=True, slots=True) +class _SuccessLog: + logger: LiteLLMLoggingObj + row: _RoundRow + chat_usage: Usage | None + + def carry_bridged_usage(self) -> None: + if self.chat_usage is not None: + self.logger.model_call_details["combined_usage_object"] = self.chat_usage + + +def _round_row( + completed: ResponsesAPIStreamingResponse, + response: ResponsesAPIResponse, + chat_usage: Usage | None, + restore_headers: RestoresHeaders, +) -> _RoundRow: + logged_response: Final = response.model_copy() + logged_response.hidden_params = {**response.hidden_params} + logged: Final = completed.model_copy(update={"response": logged_response}) + restore_headers(logged) + return _RoundRow( + logged, + _logged_cache_hit(response, chat_usage), + datetime.now(), # noqa: DTZ005 # naive to match the logging pipeline start_time + ) + + +@dataclass(frozen=True, slots=True) +class _BilledRounds: + response: ResponsesAPIResponse + breakdown: CostBreakdown | None + chat_usage: Usage | None + rounds: int = 1 + + @property + def stamped_cost(self) -> float | None: + cost: Final = self.response.hidden_params.get("response_cost") + return float(cost) if isinstance(cost, (int, float)) else None + + @property + def cost(self) -> float: + return self.stamped_cost or 0.0 + + @property + def spend(self) -> BilledRoundsSpend: + usage: Final = ( + ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(self.response.usage) + if self.chat_usage is None + else self.chat_usage + ) + return BilledRoundsSpend(usage, self.cost, self.breakdown) + + def with_round( + self, priced: ResponsesAPIResponse, breakdown: CostBreakdown | None, chat_usage: Usage | None + ) -> "_BilledRounds": + merged: Final = billed_for_every_round(self.response, self.breakdown, priced, breakdown) + total: Final = merged.hidden_params.get("response_cost") + charged: Final = merged.model_copy( + update={"usage": None if merged.usage is None else merged.usage.model_copy(update={"cost": total})} + ) + charged.hidden_params = { + **merged.hidden_params, + "cache_hit": served_from_cache(self.response) and served_from_cache(priced), + } + return _BilledRounds( + charged, + summed_cost_breakdown(self.response, self.breakdown, priced, breakdown), + _usage_of_bridged_rounds(self.chat_usage, chat_usage, total), + self.rounds + 1, + ) + + def row( + self, + completed: ResponsesAPIStreamingResponse, + client_output: Sequence[object], + restore_headers: RestoresHeaders, + ) -> _RoundRow: + return _round_row( + completed, + self.response.model_copy(update={"output": list(client_output)}), + self.chat_usage, + restore_headers, + ) + + def bridged_row(self, assembled: ModelResponse | None, logged_id: str) -> _RoundRow | None: + if assembled is None or self.chat_usage is None: + return None + logged: Final = assembled.model_copy(update={"id": logged_id, "usage": self.chat_usage}) + logged.hidden_params = { + **(get_hidden_params(assembled) or {}), + "response_cost": self.stamped_cost, + } + return _RoundRow( + logged, + served_from_cache(self.response), + datetime.now(), # noqa: DTZ005 # naive to match the logging pipeline start_time + ) + + +class StreamedRoundBilling: + """One success row per streamed auto-execute request, billed for every round it ran. + + Each round's iterator parks its success dispatch on the shared logging object rather than sending it, either as a + ready dispatch or, for a stream bridged through chat completions, as the assembled response and its cache hit. + When a round completes, before its tools run, it is folded into the rounds before it and a dispatch for every + round so far stays parked, so a client disconnect or a later round that fails still bills them. Once a later round + starts, that dispatch logs from a copy of the logging state taken then, so the later round's failure, logged on the + request's own state where the proxy reads its failure row, never reaches their success row. Only the latest + dispatch is sent, through the proxy's deferred release when the proxy armed one.""" + + def __init__(self, logging_obj: LiteLLMLoggingObj) -> None: + self._logging_obj: Final = logging_obj + self._billed: _BilledRounds | None = None + self._banked: Parked | None = None + self._banked_row: _RoundRow | None = None + if getattr(logging_obj, _DISPATCH_RELEASE, None) is None: + setattr(logging_obj, _DISPATCH_RELEASE, self._run_parked) + + def complete_round( + self, + completed: ResponsesAPIStreamingResponse, + response: ResponsesAPIResponse, + client_output: Sequence[object], + restore_headers: RestoresHeaders, + ) -> None: + round_parked: Final = self._take_parked() + price: Final = self._price(response, round_parked) + priced: Final = _priced(response, price) + chat_usage: Final = _bridged_round_usage(round_parked) + if self._billed is None: + first_round: Final = _BilledRounds(priced, price.cost_breakdown, chat_usage) + self._billed = first_round + self._bank( + round_parked, + first_round.bridged_row(_assembled(round_parked), priced.id) + or first_round.row(completed, client_output, restore_headers), + ) + return + billed: Final = self._billed.with_round(priced, price.cost_breakdown, chat_usage) + self._close(round_parked) + self._close(self._banked) + self._logging_obj.model_call_details.pop("combined_usage_object", None) + self._logging_obj.model_call_details.pop("response_cost", None) + self._billed = billed + self._logging_obj.cost_breakdown = billed.breakdown + row: Final = billed.bridged_row(_assembled(round_parked), billed.response.id) or billed.row( + completed, client_output, restore_headers + ) + self._bank(None if round_parked is None else (self._dispatch(row, billed, self._logging_obj),), row) + + def record_billed_rounds(self, follow_up: object) -> None: + if self._billed is None or not isinstance(follow_up, BaseResponsesAPIStreamingIterator): + return + follow_up.bill_failure_after_earlier_rounds(self._billed.spend) + self._bank_apart_from_later_rounds(self._billed) + + def charge_final_round( + self, completed: ResponsesAPIStreamingResponse, response: ResponsesAPIResponse, restore_headers: RestoresHeaders + ) -> ResponsesAPIResponse: + parked: Final = self._take_parked() + self._banked = None + if self._billed is None or self._billed.rounds == 1: + if parked is not None: + self._release(parked) + return response + self._close(parked) + charged: Final = response.model_copy(update={"usage": self._billed.response.usage}) + charged.hidden_params = {**self._billed.response.hidden_params} + if parked is not None: + final_row: Final = self._billed.bridged_row(self._latest_assembled(), charged.id) or _round_row( + completed, charged, self._billed.chat_usage, restore_headers + ) + self._release((self._dispatch(final_row, self._billed, self._logging_obj),)) + return charged + + def finish(self) -> None: + parked: Final = self._take_parked() + self._banked = None + if parked is not None: + self._release(parked) + return + self._disarm() + + def release_abandoned(self) -> None: + if not self._owns_release(): + return + if _event_loop_running(): + self.finish() + return + self.abandon() + + def abandon(self) -> None: + self._close(self._take_parked()) + self._close(self._banked) + self._banked = None + self._disarm() + + def _latest_assembled(self) -> ModelResponse | None: + row: Final = self._banked_row + return row.logged if row is not None and isinstance(row.logged, ModelResponse) else None + + def _price(self, response: ResponsesAPIResponse, round_parked: Parked | None) -> UncachedPrice: + if served_from_cache(response): + return _SERVED_FROM_CACHE + match round_parked: + case (ModelResponse() as assembled, bool()): + return self._logging_obj.uncached_price(assembled) + case _: + return self._logging_obj.uncached_price(response) + + def _bank(self, parked: Parked | None, row: _RoundRow) -> None: + self._banked = parked + self._banked_row = row + if parked is not None: + setattr(self._logging_obj, _PARKED_DISPATCH, parked) + + def _bank_apart_from_later_rounds(self, billed: _BilledRounds) -> None: + row: Final = self._banked_row + if row is None: + return + parked: Final = self._take_parked() + if parked is None: + return + 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]: + 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 + async def _log_success(success: _SuccessLog) -> None: + success.carry_bridged_usage() + await _success_dispatcher(success.logger).dispatch_success_handlers( + success.row.logged, + start_time=success.logger.start_time, + end_time=success.row.end_time, + cache_hit=success.row.cache_hit, + prefer_async_handlers=True, + ) + + def _take_parked(self) -> Parked | None: + parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None) + if not isinstance(parked, tuple): + return None + setattr(self._logging_obj, _PARKED_DISPATCH, None) + return parked + + async def _run_parked(self, *parked: object) -> None: + match parked: + case (BilledRoundsDispatch(dispatch=dispatch),): + await dispatch + case (dispatch,) if asyncio.iscoroutine(dispatch): + await dispatch + case (assembled, cache_hit): + await _success_dispatcher(self._logging_obj).dispatch_success_handlers( + assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True + ) + case _: + return + + def _owns_release(self) -> bool: + release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + return release == self._run_parked + + def _release(self, parked: Parked) -> None: + if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, _PARKED_DISPATCH, parked) + return + setattr(self._logging_obj, _DISPATCH_RELEASE, None) + spawn_logging_task( + asyncio.get_running_loop(), self._run_parked(*parked), task_name="Responses MCP stream success logging" + ) + + def _disarm(self) -> None: + if self._owns_release(): + setattr(self._logging_obj, _DISPATCH_RELEASE, None) + + @staticmethod + def _close(parked: Parked | None) -> None: + match parked: + case (BilledRoundsDispatch(dispatch=dispatch),): + dispatch.close() + case (dispatch,) if asyncio.iscoroutine(dispatch): + dispatch.close() + case _: + return diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index c5dc45566a9..bc30da72df9 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -42,6 +42,7 @@ from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend 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 @@ -363,6 +364,7 @@ class BaseResponsesAPIStreamingIterator: """ _pending_logging_tasks: tuple[asyncio.Task[object], ...] = () + _earlier_rounds_spend: BilledRoundsSpend | None = None def __init__( self, @@ -705,18 +707,22 @@ class BaseResponsesAPIStreamingIterator: if usage_obj is None: return try: - self.logging_obj.model_call_details["combined_usage_object"] = ( - ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_obj) - ) + failed_usage: Final = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_obj) except (TypeError, ValueError) as usage_error: verbose_logger.debug( "could not record usage for failed responses stream: %s", usage_error, ) return - self.logging_obj.model_call_details["response_cost"] = ( - self.logging_obj.response_cost_calculator(result=response_obj) or 0.0 - ) + failed_cost: Final = self.logging_obj.response_cost_calculator(result=response_obj) or 0.0 + if self._earlier_rounds_spend is None: + self.logging_obj.model_call_details["combined_usage_object"] = failed_usage + self.logging_obj.model_call_details["response_cost"] = failed_cost + return + self.logging_obj.record_failure_after_earlier_rounds(self._earlier_rounds_spend, failed_usage, failed_cost) + + def bill_failure_after_earlier_rounds(self, earlier: BilledRoundsSpend) -> None: + self._earlier_rounds_spend = earlier def _map_error_event_exception(self, error_obj: object) -> Exception: return _map_stream_error_to_exception( diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 4f8a13c1182..43190852c60 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -42,6 +42,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import ( + BilledRoundsSpend, Logging, StandardLoggingPayloadSetup, _extract_response_obj_and_hidden_params, @@ -6837,6 +6838,26 @@ def test_failure_handler_zeroes_spend_without_recovered_usage(logging_obj): assert payload["total_tokens"] == 0 +def test_failure_after_earlier_rounds_bills_only_the_failed_round_when_it_priced_no_breakdown( + logging_obj: LitellmLogging, +) -> None: + earlier_breakdown: Final = CostBreakdown(total_cost=0.3, margin_fixed_amount=0.1) + logging_obj.cost_breakdown = earlier_breakdown + logging_obj.model_call_details["litellm_params"] = {} + + logging_obj.record_failure_after_earlier_rounds( + BilledRoundsSpend(Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.3, earlier_breakdown), + Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7), + 0.0, + ) + logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "anthropic", "claude"), "tb") + + row: Final = logging_obj.model_call_details["standard_logging_object"] + assert (row["prompt_tokens"], row["completion_tokens"]) == (7, 0) + assert row["response_cost"] == pytest.approx(0.0) + assert row.get("cost_breakdown") is None or row["cost_breakdown"]["total_cost"] == pytest.approx(0.0) + + def test_set_cost_breakdown_stores_reasoning_cost(): """reasoning_cost is stored only when positive, mirroring the cache-cost fields.""" from datetime import datetime diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 8a74f2d50b4..1ec505362dd 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -3,6 +3,7 @@ Unit Tests for the max parallel request limiter v3 for the proxy """ import asyncio +import contextvars import logging import sys import time @@ -54,6 +55,7 @@ from litellm.types.mcp import MCPPreCallRequestObject, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem from litellm.types.utils import ( + CallTypes, EmbeddingResponse, ModelResponse, PromptTokensDetailsWrapper, @@ -4206,6 +4208,175 @@ async def test_completed_responses_post_call_releases_parallel_slot() -> None: ) == 0 +@pytest.mark.asyncio +async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot() -> None: + api_key: Final = hash_token("sk-mcp-auto-execute") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + owner_call_id: Final = "mcp-auto-execute-owner" + + async def in_flight() -> int: + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) + + def mcp_tool_call() -> dict[str, object]: + return { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "Tool: add"}], + "mcp_server_name": "math", + } + + async def pre_call(data: dict[str, object], call_type: str) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + + async def pre_call_from_another_request( + data: dict[str, object], call_type: str + ) -> None: + await contextvars.Context().run(asyncio.create_task, pre_call(data, call_type)) + + await pre_call( + {"model": "gpt-4o-mini", "input": "add 2 and 3", "litellm_call_id": owner_call_id}, + CallTypes.aresponses.value, + ) + await pre_call(mcp_tool_call(), CallTypes.call_mcp_tool.value) + await pre_call(mcp_tool_call(), CallTypes.call_mcp_tool.value) + assert await in_flight() == 1 + + with pytest.raises(HTTPException): + await pre_call_from_another_request( + mcp_tool_call(), CallTypes.call_mcp_tool.value + ) + + await handler.async_log_success_event( + kwargs={"litellm_call_id": owner_call_id}, + response_obj=None, + start_time=None, + end_time=None, + ) + assert await in_flight() == 0 + await pre_call_from_another_request( + {"model": "gpt-4o-mini", "input": "next request"}, + CallTypes.aresponses.value, + ) + assert await in_flight() == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stream", "in_flight_after_first_round_logs"), + [(False, 1), (True, 0)], + ids=["non_stream_round_logged_late_keeps_slot", "stream_logged_at_its_end_releases_slot"], +) +async def test_first_round_logged_after_tool_call_reuse_releases_only_at_request_end( + stream: bool, in_flight_after_first_round_logs: int +) -> None: + reused_at: Final = datetime(2026, 1, 5, 12, 0, 0) + api_key: Final = hash_token(f"sk-mcp-late-round-log-{stream}") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + time_provider=lambda: reused_at, + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + owner_call_id: Final = f"mcp-late-round-log-{stream}" + + async def in_flight() -> int: + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) + + async def pre_call(data: dict[str, object], call_type: str) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + + async def round_logged(ended_at: datetime) -> None: + await handler.async_log_success_event( + kwargs={"litellm_call_id": owner_call_id, "stream": stream}, + response_obj=None, + start_time=None, + end_time=ended_at, + ) + + await pre_call( + { + "model": "claude-haiku-5-5", + "messages": [{"role": "user", "content": "add 2 and 3"}], + "litellm_call_id": owner_call_id, + }, + CallTypes.anthropic_messages.value, + ) + await pre_call( + {"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"}, + CallTypes.call_mcp_tool.value, + ) + + await round_logged(reused_at - timedelta(seconds=1)) + assert await in_flight() == in_flight_after_first_round_logs + + await round_logged(reused_at + timedelta(seconds=1)) + assert await in_flight() == 0 + + +@pytest.mark.asyncio +async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_headers() -> None: + api_key: Final = hash_token("sk-mcp-reuse-headers") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) + data: Final[dict[str, object]] = { + "model": "gpt-4o-mini", + "input": "add 2 and 3", + "litellm_call_id": "mcp-reuse-headers-owner", + } + for request_data, call_type in ( + (data, CallTypes.aresponses.value), + ( + {"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"}, + CallTypes.call_mcp_tool.value, + ), + ): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=request_data, + call_type=call_type, + ) + response: Final = ModelResponse( + id="mcp-reuse-headers", + object="chat.completion", + created=0, + model="gpt-4o-mini", + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + choices=[], + ) + response._hidden_params = {} + + await handler.async_post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + + headers: Final = response._hidden_params["additional_headers"] + assert headers["x-ratelimit-api_key-limit-max_parallel_requests"] == 2 + assert headers["x-ratelimit-api_key-remaining-max_parallel_requests"] == 1 + + @pytest.mark.asyncio async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails() -> None: from unittest.mock import AsyncMock 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 60d4d2e89b5..f0f3fc37ddb 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -1,7 +1,9 @@ import asyncio +import gc import json import logging import sys +import warnings from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager, nullcontext from datetime import datetime @@ -18,9 +20,11 @@ from mcp.types import Tool as MCPTool # Import required modules import litellm +from litellm.exceptions import GuardrailRaisedException, MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.langfuse.langfuse import langfuse_trace_id, langfuse_trace_metadata from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.xai.common_utils import xai_reported_cost_in_usd from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerListOk from litellm.proxy.utils import _failure_fields_to_lift @@ -31,7 +35,7 @@ from litellm.types.llms.openai import ( OpenAIMcpServerTool, ToolParam, ) -from litellm.types.utils import StandardLoggingPayload, Usage +from litellm.types.utils import ModelResponse, StandardLoggingPayload, Usage class MockUserAPIKeyAuth: @@ -1518,25 +1522,39 @@ def _openai_answer_round() -> bytes: ) -def _anthropic_round(message_id: str, block: Mapping[str, object], delta: Mapping[str, object], stop: str) -> bytes: +def _anthropic_message_start( + message_id: str, model: str, input_tokens: int, cache_usage: Mapping[str, object] = MappingProxyType({}) +) -> dict[str, object]: + return { + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": input_tokens, "output_tokens": 1, **cache_usage}, + }, + } + + +def _anthropic_round( + message_id: str, + block: Mapping[str, object], + delta: Mapping[str, object], + stop: str, + model: str = "claude-haiku-4-5", + usage: tuple[int, int] = (5, 5), + cache_usage: Mapping[str, object] = MappingProxyType({}), +) -> bytes: return _sse( ( - { - "type": "message_start", - "message": { - "id": message_id, - "type": "message", - "role": "assistant", - "model": "claude-haiku-4-5", - "content": [], - "stop_reason": None, - "usage": {"input_tokens": 5, "output_tokens": 1}, - }, - }, + _anthropic_message_start(message_id, model, usage[0], cache_usage), {"type": "content_block_start", "index": 0, "content_block": dict(block)}, {"type": "content_block_delta", "index": 0, "delta": dict(delta)}, {"type": "content_block_stop", "index": 0}, - {"type": "message_delta", "delta": {"stop_reason": stop}, "usage": {"output_tokens": 5}}, + {"type": "message_delta", "delta": {"stop_reason": stop}, "usage": {"output_tokens": usage[1]}}, {"type": "message_stop"}, ) ) @@ -1586,6 +1604,7 @@ class _RoundRecorder(CustomLogger): def __init__(self) -> None: super().__init__() self.logged: tuple[StandardLoggingPayload, ...] = () + self.results: tuple[object, ...] = () self.failed: tuple[StandardLoggingPayload, ...] = () self.hooked: tuple[object, ...] = () @@ -1593,6 +1612,7 @@ class _RoundRecorder(CustomLogger): self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.logged = (*self.logged, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + self.results = (*self.results, response_obj) async def async_log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime @@ -1624,12 +1644,15 @@ _ECHO_MCP_TOOL: Final = MappingProxyType( def _serve_responses( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, responses: Sequence[httpx.Response] + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + responses: Sequence[httpx.Response], + url_regex: str = r"https://mcp-rounds\.example\.com/.*responses", ) -> respx.Route: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") litellm.in_memory_llm_clients_cache.flush_cache() - return respx_mock.post(url__regex=r"https://mcp-rounds\.example\.com/.*responses").mock(side_effect=list(responses)) + return respx_mock.post(url__regex=url_regex).mock(side_effect=list(responses)) def _serve_rounds( @@ -2424,3 +2447,894 @@ async def test_streaming_mcp_event_order_and_response_id_consistency( assert lifecycle_ids[0] == lifecycle_ids[1] assert lifecycle_ids[2] != lifecycle_ids[0] assert not [record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR] + + +_FIXED_MARGIN: Final = 0.001 + + +def _streamed(body: Mapping[str, object], **terminal: object) -> httpx.Response: + completed: Final = json.loads(json.dumps({**body, **terminal}, default=dict)) + opened: Final = {**completed, "status": "in_progress", "output": [], "usage": None} + items: Final = tuple(enumerate(completed["output"])) + events: Final = ( + {"type": "response.created", "response": opened}, + {"type": "response.in_progress", "response": opened}, + *({"type": "response.output_item.added", "output_index": index, "item": item} for index, item in items), + *({"type": "response.output_item.done", "output_index": index, "item": item} for index, item in items), + {"type": f"response.{completed['status']}", "response": completed}, + ) + return httpx.Response( + 200, + content=_sse(tuple({**event, "sequence_number": number} for number, event in enumerate(events))), + headers=SSE_HEADERS, + ) + + +class _CutAfterTheFirstEvent(httpx.AsyncByteStream): + async def __aiter__(self): + opened: Final = {**json.loads(json.dumps(_ROUND_TWO_BODY, default=dict)), "status": "in_progress"} + yield _sse(({"type": "response.created", "response": opened, "sequence_number": 0},)) + raise httpx.ReadError("connection reset mid-stream") + + +async def _open_stream_auto_execute(logging_obj: Logging) -> object: + return await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + stream=True, + litellm_logging_obj=logging_obj, + ) + + +async def _drain(stream: object) -> tuple[ResponsesAPIStreamingResponse, ...]: + return tuple([event async for event in cast(Any, stream)]) + + +async def _stream_auto_execute(logging_obj: Logging) -> tuple[ResponsesAPIStreamingResponse, ...]: + with _echo_tool_served("PING"): + return await _drain(await _open_stream_auto_execute(logging_obj)) + + +def _final_response( + events: Sequence[ResponsesAPIStreamingResponse], terminal: str = "response.completed" +) -> ResponsesAPIResponse: + assert _event_type(events[-1]) == terminal + response: Final = getattr(events[-1], "response", None) + assert isinstance(response, ResponsesAPIResponse) + return response + + +def _assert_one_stream_row_bills_both_rounds_once( + logged: tuple[StandardLoggingPayload, ...], client: ResponsesAPIResponse +) -> None: + charged_once: Final = _round_cost(_ROUND_ONE_BODY) + _round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN + assert len(logged) == 1 + row: Final = logged[0] + assert _output_types(row["response"]["output"]) == _output_types(client.output) + assert "message" in _output_types(client.output) + assert "function_call" not in _output_types(client.output) + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["cache_hit"] is None + assert row["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 32 + 64 + assert row["hidden_params"]["additional_headers"]["llm_provider-content-type"] == SSE_HEADERS["content-type"] + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + 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) + + +def _assert_breakdown_agrees_with_the_billed_cost(row: StandardLoggingPayload) -> None: + breakdown: Final = row.get("cost_breakdown") + assert breakdown is None or breakdown["total_cost"] == pytest.approx(row["response_cost"]) + + +def _request_failure_row(logging_obj: Logging, success_row: StandardLoggingPayload) -> Mapping[str, object]: + failure_row: Final = _failure_fields_to_lift(MappingProxyType({"litellm_logging_obj": logging_obj})) + failure_payload: Final = cast(StandardLoggingPayload, failure_row["standard_logging_object"]) + assert failure_payload["status"] == "failure" + assert failure_payload["id"] != success_row["id"] + return failure_row + + +def _assert_billed_tokens(row: Mapping[str, object], tokens: tuple[int, int]) -> None: + usage: Final = row["combined_usage_object"] + assert isinstance(usage, Usage) + assert (usage.prompt_tokens, usage.completion_tokens) == tokens + + +def _assert_one_row_bills(logged: tuple[StandardLoggingPayload, ...], body: Mapping[str, object]) -> None: + usage: Final = cast(Mapping[str, int], body["usage"]) + assert len(logged) == 1 + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == ( + usage["input_tokens"], + usage["output_tokens"], + ) + assert logged[0]["response_cost"] == pytest.approx(_round_cost(body)) + assert logged[0]["response_cost"] > 0 + + +def _assert_one_row_bills_round_one_twice_with_one_fixed_margin(logged: tuple[StandardLoggingPayload, ...]) -> None: + charged_once: Final = 2 * _round_cost(_ROUND_ONE_BODY) - _FIXED_MARGIN + assert len(logged) == 1 + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == (2 * 71, 2 * 19) + assert logged[0]["response_cost"] == pytest.approx(charged_once) + assert logged[0]["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_logs_the_final_response_billed_for_both_rounds_with_one_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-rounds")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once_when_the_proxy_releases( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-rounds-deferred") + + with _echo_tool_served("PING"): + stream: Final = await _open_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + events: Final = await _drain(stream) + await asyncio.sleep(0.2) + logged_before_release: Final = recorder.logged + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0) + + assert logged_before_release == () + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_the_paid_first_round_on_its_success_row_when_the_follow_up_stream_breaks( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.utils import ProxyLogging + + _serve_responses( + respx_mock, + monkeypatch, + (_streamed(_ROUND_ONE_BODY), httpx.Response(200, stream=_CutAfterTheFirstEvent(), headers=SSE_HEADERS)), + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-follow-up-breaks") + + with pytest.raises(litellm.ServiceUnavailableError): + await _stream_auto_execute(logging_obj) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + assert [row["response_cost"] for row in recorder.failed] == [0.0] + assert _request_failure_row(logging_obj, recorder.logged[0]).get("response_cost", 0.0) == 0.0 + + +@pytest.mark.asyncio +async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_the_success_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-rounds-blocked") + blocked: Final = GuardrailRaisedException(guardrail_name="output-policy", message="blocked") + + with _echo_tool_served("PING"): + stream: Final = await _open_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + events: Final = await _drain(stream) + 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) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) + assert [row["response_cost"] for row in recorder.failed] == [0.0] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_the_paid_first_round_once_when_the_follow_up_is_rejected( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + context_window_error: Final = { + "error": { + "message": "Your input exceeds the context window of this model.", + "type": "invalid_request_error", + "param": "input", + "code": "context_length_exceeded", + } + } + _serve_responses( + respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), httpx.Response(400, json=context_window_error)) + ) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-follow-up-rejected")) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + assert _event_type(events[-1]) == "error" + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + assert [row["response_cost"] for row in recorder.failed] == [0.0] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_every_paid_round_with_one_fixed_margin_when_the_third_round_is_rejected( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + ( + _streamed(_ROUND_ONE_BODY), + _streamed(_ROUND_ONE_BODY), + httpx.Response(400, json={"error": {"message": "Your input exceeds the context window of this model."}}), + ), + ) + recorder: Final = _RoundRecorder() + + await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-third-round-rejected")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_round_one_twice_with_one_fixed_margin(recorder.logged) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_every_finished_round_when_the_client_disconnects_during_a_later_tool_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_ONE_BODY))) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-disconnect-in-tool-call") + echo_then_disconnect: Final = AsyncMock( + side_effect=[[{"tool_call_id": "call_echo", "result": "PING", "name": "echo"}], asyncio.CancelledError()] + ) + + with ( + _echo_tool_served("PING"), + patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", echo_then_disconnect), + pytest.raises(asyncio.CancelledError), + ): + await _drain(await _open_stream_auto_execute(logging_obj)) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_round_one_twice_with_one_fixed_margin(recorder.logged) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_answered_without_a_tool_call_logs_its_one_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_TWO_BODY),)) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-one-round")) + await _settle(lambda: len(recorder.logged) > 0) + + assert _output_types(_final_response(events).output) == ("message",) + _assert_one_row_bills(recorder.logged, _ROUND_TWO_BODY) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_both_rounds_with_one_fixed_margin_when_the_final_round_is_incomplete( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + truncated: Final = _streamed( + _ROUND_TWO_BODY, status="incomplete", incomplete_details={"reason": "max_output_tokens"} + ) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), truncated)) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-final-round-incomplete")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events, "response.incomplete")) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_a_failed_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + server_error: Final = {"code": "server_error", "message": "The server had an error processing your request."} + _serve_responses( + respx_mock, + monkeypatch, + (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY, status="failed", error=server_error)), + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-follow-up-failed") + + with pytest.raises(litellm.ServiceUnavailableError): + await _stream_auto_execute(logging_obj) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (149, 5) + assert recorder.failed[0]["response_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN) + _assert_breakdown_agrees_with_the_billed_cost(recorder.failed[0]) + failure_row: Final = _request_failure_row(logging_obj, recorder.logged[0]) + _assert_billed_tokens(failure_row, (149, 5)) + assert failure_row["response_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN) + + +_XAI_MODEL: Final = "xai/grok-4.7" +_XAI_RESPONSES_URL: Final = r"https://api\.x\.ai/v1/responses" + + +def _xai_round(body: Mapping[str, object], cost_in_usd_ticks: int) -> Mapping[str, object]: + usage: Final = cast(Mapping[str, object], body["usage"]) + stated: Final = {**usage, "cost_in_usd_ticks": cost_in_usd_ticks, "context_details": {"cached_tokens": 3}} + return MappingProxyType({**body, "model": "grok-4.7", "usage": stated}) + + +def _xai_round_cost(body: Mapping[str, object]) -> float: + usage: Final = cast(Mapping[str, object], body["usage"]) + stated: Final = {**usage, "cost": xai_reported_cost_in_usd(usage["cost_in_usd_ticks"])} + return litellm.completion_cost( + completion_response=ResponsesAPIResponse.model_validate({**body, "usage": stated}), + model=_XAI_MODEL, + custom_llm_provider="xai", + call_type="aresponses", + ) + + +@pytest.mark.asyncio +async def test_xai_stream_auto_execute_bills_each_round_at_its_own_price_with_one_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"xai": {"fixed_amount": _FIXED_MARGIN}}) + rounds: Final = (_xai_round(_ROUND_ONE_BODY, 12_400_000), _xai_round(_ROUND_TWO_BODY, 8_800_000)) + _serve_responses(respx_mock, monkeypatch, tuple(_streamed(body) for body in rounds), _XAI_RESPONSES_URL) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await litellm.aresponses( + model=_XAI_MODEL, + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="xai-mcp-rounds", + stream=True, + litellm_logging_obj=_rounds_logging_obj(recorder, "mcp-xai-stream-rounds"), + ) + ) + await _settle(lambda: len(recorder.logged) > 0) + + charged_once: Final = _xai_round_cost(rounds[0]) + _xai_round_cost(rounds[1]) - _FIXED_MARGIN + client: Final = _final_response(events) + assert len(recorder.logged) == 1 + assert recorder.logged[0]["response_cost"] == pytest.approx(charged_once) + assert recorder.logged[0]["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert client.usage is not None + assert client.usage.cost == pytest.approx(charged_once) + assert json.loads(events[-1].model_dump_json())["response"]["usage"]["context_details"] == {"cached_tokens": 6} + + +_UNPRICED_MODEL: Final = "litellm-unpriced-mcp-rounds" + + +@pytest.mark.asyncio +async def test_stream_auto_execute_on_an_unpriced_model_logs_both_rounds_as_a_failed_cost_lookup_not_a_free_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + rounds: Final = tuple( + MappingProxyType({**body, "model": _UNPRICED_MODEL}) for body in (_ROUND_ONE_BODY, _ROUND_TWO_BODY) + ) + _serve_rounds(respx_mock, monkeypatch, rounds) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await litellm.aresponses( + model=f"openai/{_UNPRICED_MODEL}", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + stream=True, + litellm_logging_obj=_rounds_logging_obj(recorder, "mcp-stream-unpriced-rounds"), + ) + ) + await _settle(lambda: len(recorder.logged) > 0) + + client: Final = _final_response(events) + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["response_cost_failure_debug_info"] is not None + assert client.usage is not None + assert client.usage.cost is None + + +async def _stop_reading_once_the_follow_up_answers(logging_obj: Logging) -> None: + stream: Final = await _open_stream_auto_execute(logging_obj) + async for event in cast(Any, stream): + if _event_type(event) == "response.output_item.added" and _output_types((event.item,)) == ("message",): + return + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_the_finished_first_round_when_the_caller_stops_reading_the_follow_up( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + await _stop_reading_once_the_follow_up_answers(_rounds_logging_obj(recorder, "mcp-stream-caller-stops")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + + +async def _stream_read_until_the_follow_up_answers(logging_obj: Logging) -> object: + stream: Final = await _open_stream_auto_execute(logging_obj) + async for event in cast(Any, stream): + if _event_type(event) == "response.output_item.added" and _output_types((event.item,)) == ("message",): + return stream + pytest.fail("the follow-up round never answered") + + +_STREAM_BILLING_NAMES: Final = ("StreamedRoundBilling", "dispatch_success_handlers", "StreamingIterator") + + +def _names_stream_billing(text: str) -> bool: + return any(name in text for name in _STREAM_BILLING_NAMES) + + +class _UnraisableRecorder: + def __init__(self) -> None: + self.raised: tuple[str, ...] = () + + def __call__(self, unraisable: "sys.UnraisableHookArgs") -> None: + self.raised = (*self.raised, f"{unraisable.object!r}: {unraisable.exc_value!r}") + + +def test_stream_auto_execute_dropped_after_its_event_loop_closed_discards_its_parked_log_cleanly( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + unraisable: Final = _UnraisableRecorder() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + + with _echo_tool_served("PING"), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + asyncio.run(_stream_read_until_the_follow_up_answers(_rounds_logging_obj(recorder, "mcp-stream-loop-closed"))) + gc.collect() + + assert recorder.logged == () + assert [raised for raised in unraisable.raised if _names_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_stream_billing(str(warning.message))] == [] + + +_BRIDGED_MODEL: Final = "anthropic/claude-haiku-5-5" +_ANTHROPIC_MESSAGES_URL: Final = r"https://api\.anthropic\.com/v1/messages" + + +def _bridged_round(content: bytes) -> httpx.Response: + return httpx.Response(200, content=content, headers=SSE_HEADERS) + + +def _bridged_echo_call_round(cache_usage: Mapping[str, object] = MappingProxyType({})) -> httpx.Response: + return _bridged_round( + _anthropic_round( + "msg_echo_call", + {"type": "tool_use", "id": "call_echo", "name": "echo", "input": {}}, + {"type": "input_json_delta", "partial_json": json.dumps({"value": "PING"})}, + "tool_use", + model="claude-haiku-5-5", + usage=(533, 40), + cache_usage=cache_usage, + ) + ) + + +def _bridged_echo_answer_round( + stop: str = "end_turn", cache_usage: Mapping[str, object] = MappingProxyType({}) +) -> httpx.Response: + return _bridged_round( + _anthropic_round( + "msg_echo_answer", + {"type": "text", "text": ""}, + {"type": "text_delta", "text": "PING"}, + stop, + model="claude-haiku-5-5", + usage=(470, 4), + cache_usage=cache_usage, + ) + ) + + +def _anthropic_prompt_cache_usage(written: int, read: int) -> Mapping[str, object]: + return MappingProxyType( + { + "cache_creation_input_tokens": written, + "cache_read_input_tokens": read, + "cache_creation": {"ephemeral_5m_input_tokens": written, "ephemeral_1h_input_tokens": 0}, + "inference_geo": "us", + "speed": "standard", + } + ) + + +class _CutAfterTheMessageStart(httpx.AsyncByteStream): + async def __aiter__(self): + yield _sse((_anthropic_message_start("msg_echo_answer", "claude-haiku-5-5", 470),)) + raise httpx.ReadError("connection reset mid-stream") + + +def _bridged_round_cost(prompt_tokens: int, completion_tokens: int) -> float: + return litellm.completion_cost( + completion_response=ModelResponse( + model="claude-haiku-5-5", + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ), + model=_BRIDGED_MODEL, + ) + + +async def _open_bridged_stream_auto_execute(logging_obj: Logging) -> object: + return await litellm.aresponses( + model=_BRIDGED_MODEL, + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-ant-mcp-rounds", + stream=True, + litellm_logging_obj=logging_obj, + ) + + +def _assert_logged_as_the_final_chat_answer_with_both_rounds_usage(recorder: _RoundRecorder, answer: str) -> None: + assert len(recorder.results) == 1 + logged: Final = recorder.results[0] + assert isinstance(logged, ModelResponse) + assert logged.choices[0].message.content == answer + usage: Final[object] = getattr(logged, "usage", None) + assert isinstance(usage, Usage) + assert (usage.prompt_tokens, usage.completion_tokens) == (533 + 470, 40 + 4) + + +def _assert_one_bridged_row_bills_both_rounds_once(recorder: _RoundRecorder, client: ResponsesAPIResponse) -> None: + charged_once: Final = _bridged_round_cost(533, 40) + _bridged_round_cost(470, 4) - _FIXED_MARGIN + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + _assert_logged_as_the_final_chat_answer_with_both_rounds_usage(recorder, "PING") + assert row["id"] == client.id + assert "message" in _output_types(client.output) + assert (row["prompt_tokens"], row["completion_tokens"]) == (533 + 470, 40 + 4) + assert row["cache_hit"] is False + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert client.usage is not None + assert (client.usage.input_tokens, client.usage.output_tokens) == (533 + 470, 40 + 4) + assert client.usage.cost == pytest.approx(charged_once) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_logs_the_final_response_billed_for_both_rounds_with_one_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round()), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-stream-rounds")) + ) + await _settle(lambda: len(recorder.logged) > 0) + + assert _output_types(_final_response(events).output) == ("mcp_call", "message") + _assert_one_bridged_row_bills_both_rounds_once(recorder, _final_response(events)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_logs_the_provider_usage_of_both_rounds_summed( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + ( + _bridged_echo_call_round(_anthropic_prompt_cache_usage(written=10, read=5)), + _bridged_echo_answer_round(cache_usage=_anthropic_prompt_cache_usage(written=3, read=7)), + ), + _ANTHROPIC_MESSAGES_URL, + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-provider-usage")) + ) + await _settle(lambda: len(recorder.logged) > 0) + + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + usage: Final = row["metadata"]["usage_object"] + assert "choices" in row["response"] + assert (row["prompt_tokens"], row["completion_tokens"]) == (533 + 10 + 5 + 470 + 3 + 7, 40 + 4) + assert (usage["cache_creation_input_tokens"], usage["cache_read_input_tokens"]) == (10 + 3, 5 + 7) + assert usage["prompt_tokens_details"]["cache_creation_token_details"]["ephemeral_5m_input_tokens"] == 10 + 3 + assert (usage["inference_geo"], usage["speed"]) == ("us", "standard") + assert usage["cost"] == pytest.approx(row["response_cost"]) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_both_rounds_with_one_fixed_margin_when_the_final_round_is_incomplete( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + (_bridged_echo_call_round(), _bridged_echo_answer_round("max_tokens")), + _ANTHROPIC_MESSAGES_URL, + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-stream-incomplete")) + ) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_bridged_row_bills_both_rounds_once(recorder, _final_response(events, "response.incomplete")) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once_on_release( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round()), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-bridged-stream-rounds-deferred") + + with _echo_tool_served("PING"): + stream: Final = await _open_bridged_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + events: Final = await _drain(stream) + await asyncio.sleep(0.2) + logged_before_release: Final = recorder.logged + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0) + + assert logged_before_release == () + _assert_one_bridged_row_bills_both_rounds_once(recorder, _final_response(events)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_blocked_by_a_post_call_guardrail_bills_both_rounds_on_the_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round()), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-bridged-stream-rounds-blocked") + blocked: Final = GuardrailRaisedException(guardrail_name="output-policy", message="blocked") + + with _echo_tool_served("PING"): + stream: Final = await _open_bridged_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _drain(stream) + 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.failed) > 0) + + charged_once: Final = _bridged_round_cost(533, 40) + _bridged_round_cost(470, 4) - _FIXED_MARGIN + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (533 + 470, 40 + 4) + assert recorder.failed[0]["response_cost"] == pytest.approx(charged_once) + assert recorder.failed[0]["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stop", ("end_turn", "max_tokens"), ids=("completed", "incomplete")) +async def test_bridged_stream_auto_execute_bills_both_rounds_when_the_caller_stops_reading_after_the_answer_text( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stop: str +): + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round(stop)), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + + async def stop_reading_once_the_answer_text_is_done() -> None: + stream: Final = await _open_bridged_stream_auto_execute( + _rounds_logging_obj(recorder, "mcp-bridged-stream-caller-stops") + ) + async for event in cast(Any, stream): + if _event_type(event) == "response.output_text.done": + return + + with _echo_tool_served("PING"): + await stop_reading_once_the_answer_text_is_done() + await _settle(lambda: len(recorder.logged) > 0) + + assert len(recorder.logged) == 1 + _assert_logged_as_the_final_chat_answer_with_both_rounds_usage(recorder, "PING") + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (533 + 470, 40 + 4) + assert recorder.logged[0]["response_cost"] == pytest.approx( + _bridged_round_cost(533, 40) + _bridged_round_cost(470, 4) + ) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + (_bridged_echo_call_round(), httpx.Response(200, stream=_CutAfterTheMessageStart(), headers=SSE_HEADERS)), + _ANTHROPIC_MESSAGES_URL, + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-bridged-stream-breaks") + + with _echo_tool_served("PING"), pytest.raises(MidStreamFallbackError): + await _drain(await _open_bridged_stream_auto_execute(logging_obj)) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + assert len(recorder.logged) == 1 + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (533, 40) + assert recorder.logged[0]["response_cost"] == pytest.approx(_bridged_round_cost(533, 40)) + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (470, 0) + assert recorder.failed[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 0) - _FIXED_MARGIN) + _assert_breakdown_agrees_with_the_billed_cost(recorder.failed[0]) + failure_row: Final = _request_failure_row(logging_obj, recorder.logged[0]) + _assert_billed_tokens(failure_row, (470, 0)) + assert failure_row["response_cost"] == pytest.approx(_bridged_round_cost(470, 0) - _FIXED_MARGIN) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_a_broken_follow_up_its_own_fixed_margin_after_a_cached_first_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + _serve_responses( + respx_mock, + monkeypatch, + ( + _bridged_echo_call_round(), + _bridged_echo_answer_round(), + httpx.Response(200, stream=_CutAfterTheMessageStart(), headers=SSE_HEADERS), + ), + _ANTHROPIC_MESSAGES_URL, + ) + warm: Final = _RoundRecorder() + with _echo_tool_served("PING"): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(warm, "mcp-bridged-cache-then-break-warm")) + ) + await _settle(lambda: len(warm.logged) > 0) + + recorder: Final = _RoundRecorder() + with _echo_tool_served("PONG"), pytest.raises(MidStreamFallbackError): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-cache-then-break")) + ) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + assert [row["response_cost"] for row in recorder.logged] == [0.0] + assert len(recorder.failed) == 1 + assert recorder.failed[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 0)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_answered_without_a_tool_call_logs_its_one_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_bridged_echo_answer_round(),), _ANTHROPIC_MESSAGES_URL) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + await _drain(await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-one-round"))) + await _settle(lambda: len(recorder.logged) > 0) + + assert len(recorder.logged) == 1 + assert "choices" in recorder.logged[0]["response"] + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (470, 4) + assert recorder.logged[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 4)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_nothing_for_a_round_served_from_the_response_cache( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + provider: Final = _serve_responses( + respx_mock, + monkeypatch, + (_bridged_echo_call_round(), _bridged_echo_answer_round(), _bridged_echo_answer_round()), + _ANTHROPIC_MESSAGES_URL, + ) + warm: Final = _RoundRecorder() + with _echo_tool_served("PING"): + await _drain(await _open_bridged_stream_auto_execute(_rounds_logging_obj(warm, "mcp-bridged-cache-warm"))) + await _settle(lambda: len(warm.logged) > 0) + + first_cached: Final = _RoundRecorder() + with _echo_tool_served("PONG"): + events: Final = await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(first_cached, "mcp-bridged-cache-first-round")) + ) + await _settle(lambda: len(first_cached.logged) > 0) + + both_cached: Final = _RoundRecorder() + with _echo_tool_served("PING"): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(both_cached, "mcp-bridged-cache-both-rounds")) + ) + await _settle(lambda: len(both_cached.logged) > 0) + + client: Final = _final_response(events) + assert provider.call_count == 3 + assert len(first_cached.logged) == 1 + assert (first_cached.logged[0]["prompt_tokens"], first_cached.logged[0]["completion_tokens"]) == (533 + 470, 40 + 4) + assert first_cached.logged[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 4)) + assert first_cached.logged[0]["cache_hit"] is False + assert [(row["response_cost"], row["cache_hit"]) for row in both_cached.logged] == [(0.0, True)] + assert client.usage is not None + assert client.usage.cost == pytest.approx(_bridged_round_cost(470, 4)) diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index 04d2b9ddf81..7445070bd3a 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -23,7 +23,8 @@ import asyncio import json from datetime import datetime from importlib import import_module -from typing import Final +from collections.abc import Mapping +from typing import Final, cast from unittest.mock import Mock, patch import httpx @@ -31,6 +32,7 @@ import pytest import litellm from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id @@ -44,6 +46,7 @@ from litellm.responses.streaming_iterator import ( _status_code_for_error_fields, ) from litellm.utils import ProviderConfigManager +from litellm.types.utils import StandardLoggingPayload from litellm.types.llms.openai import ( ErrorEvent, ErrorEventError, @@ -523,6 +526,82 @@ def test_handle_logging_failed_response_records_usage_and_cost(): iterator.logging_obj.response_cost_calculator.assert_called_once_with(result=chunk.response) +class _FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failed: tuple[StandardLoggingPayload, ...] = () + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.failed = (*self.failed, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + + +async def _first_failure_row(recorder: _FailureRecorder) -> StandardLoggingPayload: + for _ in range(200): + if recorder.failed: + return recorder.failed[0] + await asyncio.sleep(0.01) + pytest.fail("the failed stream logged no failure row") + + +@pytest.mark.asyncio +async def test_failed_stream_bills_only_its_own_usage_over_an_earlier_attempts_stash(): + recorder: Final = _FailureRecorder() + logging_obj: Final = LiteLLMLoggingObj( + model="gpt-5.6", + messages=[], + stream=True, + call_type="aresponses", + start_time=datetime(2026, 1, 1), + litellm_call_id="responses-failed-over-stash", + function_id="responses-failed-over-stash", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="gpt-5.6", optional_params={}, litellm_params=get_litellm_params(custom_llm_provider="openai") + ) + logging_obj.model_call_details["combined_usage_object"] = litellm.Usage( + prompt_tokens=100, completion_tokens=50, total_tokens=150 + ) + logging_obj.model_call_details["response_cost"] = 0.5 + failed: Final = { + "type": "response.failed", + "sequence_number": 1, + "response": { + "id": "resp_failed_over_stash", + "object": "response", + "created_at": 1790000000, + "status": "failed", + "model": "gpt-5.6", + "output": [], + "error": {"code": "server_error", "message": "boom"}, + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + } + iterator: Final = ResponsesAPIStreamingIterator( + response=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, content=f"data: {json.dumps(failed)}\n\n".encode() + ), + model="gpt-5.6", + responses_api_provider_config=ProviderConfigManager.get_provider_responses_api_config( + provider="openai", model="gpt-5.6" + ), + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + + with pytest.raises(MidStreamFallbackError): + await anext(aiter(iterator)) + + row: Final = await _first_failure_row(recorder) + prompt_cost, completion_cost = litellm.cost_per_token( + model="gpt-5.6", prompt_tokens=10, completion_tokens=5, custom_llm_provider="openai" + ) + assert (row["prompt_tokens"], row["completion_tokens"]) == (10, 5) + assert row["response_cost"] == pytest.approx(prompt_cost + completion_cost) + + def test_handle_logging_failed_response_without_usage_skips_recording(): iterator = _make_iterator() iterator.completed_response = _make_failed_chunk(