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(