From c32692bd9b95361406f4dbcb4c607b8f71287b26 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:55:40 -0700 Subject: [PATCH] fix(mcp): bill every chat MCP stream round once on failure, hang-up, and sync drop When the follow-up call fails, round 1's success row is now logged on a copy of the logging object that does not share the request's guardrail results, so a post-call guardrail that reports on the failed request counts once instead of on both rows A client hang-up during the follow-up now bills round 1 plus the follow-up's partial answer on one success row with round 1's id and the fixed margin counted once, the same way a plain stream's hang-up bills its partial answer The last usage chunk now carries round 1's usage and cost when only round 1 is billed, and a synchronous reader that drops the stream mid-answer bills round 1 on the loop it read with instead of discarding it The deferred release's logged_on argument is typed and every release closure accepts it, and the parked dispatch constants and protocols both MCP billing paths use now live in one module. The unit tests wait on a logging signal instead of a fixed delay and check the billed row through the proxy stream hook instead of private callback state --- litellm/litellm_core_utils/litellm_logging.py | 22 ++ .../litellm_core_utils/streaming_handler.py | 35 +- litellm/proxy/common_request_processing.py | 43 ++- .../responses/mcp/chat_completions_handler.py | 15 + .../mcp/chat_stream_round_billing.py | 107 +++--- .../responses/mcp/parked_stream_dispatch.py | 52 +++ litellm/responses/mcp/stream_round_billing.py | 58 ++-- .../integration/mcp/test_mcp_llm_endpoints.py | 89 ++++- .../proxy/test_common_request_processing.py | 28 ++ .../mcp/test_chat_completions_handler.py | 320 +++++++++++++++--- 10 files changed, 624 insertions(+), 145 deletions(-) create mode 100644 litellm/responses/mcp/parked_stream_dispatch.py diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8b66c27b930..82d2aaac81f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -564,6 +564,11 @@ def _resolve_mantle_region_for_cost( _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"}) _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) _REDACTED_RESULT: Final = TypeAdapter(object) +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" + + +def _without_guardrail_information(metadata: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in metadata.items() if key != _GUARDRAIL_INFORMATION_KEY} def _provider_response_id(source: object) -> str | None: @@ -2494,6 +2499,23 @@ class Logging(LiteLLMLoggingBaseClass): attempt.model_call_details = dict(self.model_call_details) return attempt + def guardrail_scoped_copy(self) -> "Logging": + request_params: Final = _LITELLM_PARAMS.validate_python( + self.litellm_params or {} # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped + ) + litellm_params: Final = { + **request_params, + **{ + key: _without_guardrail_information(_METADATA_MAPPING.validate_python(value)) + for key, value in request_params.items() + if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict) + }, + } + scoped: Final = copy.copy(self) + scoped.litellm_params = litellm_params + scoped.model_call_details = {**self.model_call_details, "litellm_params": litellm_params} + return scoped + def dispatch_attempt_failure_handlers( self, exception: Exception, traceback_exception: str ) -> Coroutine[None, None, None]: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 2348d8e57d6..ede5523eea4 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2458,6 +2458,25 @@ class CustomStreamWrapper: def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None: self._earlier_rounds_spend = earlier + def partial_response(self) -> "ModelResponse | TextCompletionResponse | None": + if not self.chunks: + return None + partial: Final = litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages if isinstance(self.messages, list) else None, + logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, + ) + if partial is None: + return None + usage: Final[object] = getattr(partial, "usage", None) + if not isinstance(usage, Usage): + return None + if self.model: + partial.model = self.model + backfill_missing_cache_usage_fields(usage) + return partial + def _record_partial_usage_for_failure(self) -> None: """ A stream that breaks mid-flight still billed the provider for the chunks @@ -2467,23 +2486,15 @@ class CustomStreamWrapper: moves it into terminal failure usage before opening the next attempt, so a recovered attempt logs its own usage. """ - if self.logging_obj is None or not self.chunks: + if self.logging_obj is None: return try: - partial_response: Final = litellm.stream_chunk_builder( - chunks=self.chunks, - messages=self.messages if isinstance(self.messages, list) else None, - logging_obj=self.logging_obj, - count_prompt_tokens=self.count_prompt_tokens, - ) + partial_response: Final = self.partial_response() if partial_response is None: return - usage: Final = cast(Usage | None, getattr(partial_response, "usage", None)) - if usage is None: + usage: Final[object] = getattr(partial_response, "usage", None) + if not isinstance(usage, Usage): return - if self.model: - partial_response.model = self.model - backfill_missing_cache_usage_fields(usage) 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 diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c22d5a9857d..a36884b90dd 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -137,6 +137,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, ) +from litellm.responses.mcp.parked_stream_dispatch import parked_success from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task from litellm.router import Router from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX @@ -527,6 +528,20 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return await _dispatch_disconnect_billing_event(logging_obj, partial_response) +@runtime_checkable +class _BillsInFlightRoundOnDisconnect(Protocol): + def bill_in_flight_round_on_disconnect(self) -> None: ... + + +def _bill_in_flight_round_on_disconnect(response: object) -> None: + if not isinstance(response, _BillsInFlightRoundOnDisconnect): + return + try: + response.bill_in_flight_round_on_disconnect() + except Exception as e: # noqa: BLE001 # disconnect billing is best-effort; never break stream teardown + verbose_proxy_logger.debug("Failed to bill the in-flight round after client disconnect: %s", e) + + async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, partial_response: object) -> bool: try: await logging_obj.dispatch_success_handlers( @@ -2052,6 +2067,13 @@ class _ClientStreamFailureReport: ) +def _metadata_logged_on(logged_on: LiteLLMLoggingObj) -> Mapping[str, object]: + logged_params: Final = object_mapping(logged_on.model_call_details.get("litellm_params")) or MappingProxyType({}) + return MappingProxyType( + {key: logged_params[key] for key in ("metadata", "litellm_metadata") if key in logged_params} + ) + + class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict, served_model: str | None = None): self.data = data @@ -4065,7 +4087,15 @@ class ProxyBaseLLMRequestProcessing: assembled_response: object, cache_hit: object, logged_on: LiteLLMLoggingObj | None = None ) -> None: await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( - captured_data=_captured_data, + captured_data=_captured_data + if logged_on is None + else { + **_captured_data, + **{ + key: value for key, value in _metadata_logged_on(logged_on).items() if key in _captured_data + }, + "litellm_logging_obj": logged_on, + }, captured_user_api_key_dict=_captured_user_api_key_dict, captured_logging_obj=_captured_logging_obj if logged_on is None else logged_on, assembled_response=assembled_response, @@ -4088,14 +4118,8 @@ class ProxyBaseLLMRequestProcessing: 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): - await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers( - assembled_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - prefer_async_handlers=True, - ) + case _ if (success := parked_success(args)) is not None: + await success.dispatch(_captured_native_logging_obj) case _: verbose_proxy_logger.error( "Deferred stream logging dropped: unexpected stored args shape %s", @@ -4423,6 +4447,7 @@ class ProxyBaseLLMRequestProcessing: client_disconnected, ) if recorded_client_disconnect: + _bill_in_flight_round_on_disconnect(response) deferred_stream_logging_armed: Final = _deferred_stream_logging_is_armed(request_data) ProxyLogging.fire_deferred_stream_logging(request_data) # A disconnect-time success event (the deferred-guardrail flush diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 90438af1487..06308a147fd 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -273,6 +273,18 @@ async def acompletion_with_mcp( if isinstance(round_billing, ChatStreamRoundBilling): round_billing.release_abandoned() + def bill_in_flight_round_on_disconnect(self) -> None: + follow_up: Final[object] = self.follow_up_stream + if ( + self._round_billing is None + or self.follow_up_exhausted + or not isinstance(follow_up, CustomStreamWrapper) + ): + return + partial: Final = follow_up.partial_response() + if isinstance(partial, ModelResponse): + self._round_billing.bill_partial_follow_up(partial) + def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_list_tools to the first chunk.""" from litellm.types.utils import ( @@ -598,6 +610,9 @@ async def acompletion_with_mcp( self._custom_iterator.release_abandoned() await super().aclose() + def bill_in_flight_round_on_disconnect(self) -> None: + self._custom_iterator.bill_in_flight_round_on_disconnect() + def __iter__(self): # For synchronous iteration, create a sync wrapper if self._sync_iterator is None: diff --git a/litellm/responses/mcp/chat_stream_round_billing.py b/litellm/responses/mcp/chat_stream_round_billing.py index 4df79ca13ba..3b0de7a0fde 100644 --- a/litellm/responses/mcp/chat_stream_round_billing.py +++ b/litellm/responses/mcp/chat_stream_round_billing.py @@ -1,39 +1,26 @@ import asyncio -from collections.abc import Callable, Coroutine, Generator +from collections.abc import Generator from contextlib import contextmanager from dataclasses import dataclass, replace from types import MappingProxyType -from typing import Final, Protocol, TypeAlias +from typing import Final, TypeAlias from litellm.litellm_core_utils.asyncify import is_event_loop_running from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.parked_stream_dispatch import ( + DISPATCH_RELEASE, + PARKED_DISPATCH, + SERVED_FROM_CACHE, + DeferredStreamRelease, + parked_success, +) from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown from litellm.responses.streaming_iterator import spawn_logging_task from litellm.types.utils import ModelResponse, ModelResponseStream, Usage Parked: TypeAlias = tuple[object, ...] -_Release: TypeAlias = Callable[..., Coroutine[object, object, 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 @dataclass(frozen=True, slots=True) @@ -73,8 +60,13 @@ class _FirstRound: def client_usage_chunk(self, billed: ModelResponse | None) -> ModelResponseStream | None: if self.usage_chunk is None: return None - merged: Final = merged_chat_usage(self.priced.usage, self.usage_chunk.usage) - usage: Final = merged if billed is None else merged.model_copy(update=MappingProxyType({"cost": _cost(billed)})) + usage: Final = ( + self.priced.usage.model_copy(update=MappingProxyType({"cost": self.priced.price.cost})) + if billed is None + else merged_chat_usage(self.priced.usage, self.usage_chunk.usage).model_copy( + update=MappingProxyType({"cost": _cost(billed)}) + ) + ) return self.usage_chunk.chunk.model_copy(update=MappingProxyType({"usage": usage})) @@ -99,12 +91,14 @@ class ChatStreamRoundBilling: def __init__(self, logging_obj: LiteLLMLoggingObj) -> None: self._logging_obj: Final = logging_obj self._first: _FirstRound | None = None + self._reading_loop: asyncio.AbstractEventLoop | None = None @contextmanager def armed_while_reading(self) -> Generator[None]: - armed: Final = getattr(self._logging_obj, _DISPATCH_RELEASE, None) is None + armed: Final = getattr(self._logging_obj, DISPATCH_RELEASE, None) is None if armed: - setattr(self._logging_obj, _DISPATCH_RELEASE, self._log_success) + setattr(self._logging_obj, DISPATCH_RELEASE, self._log_success) + self._reading_loop = asyncio.get_running_loop() try: yield finally: @@ -135,20 +129,30 @@ class ChatStreamRoundBilling: def fail(self) -> None: self._bank_apart_from_the_follow_up() first: Final = self._first - if first is None or first.banked is None or self._peek_parked() is not first.parked: + if first is None or first.banked is None or not self._holds_round_one_alone(first): self.finish() return self._first = None - setattr(self._logging_obj, _PARKED_DISPATCH, None) - self._log_through_the_release((*first.parked, first.banked)) + setattr(self._logging_obj, PARKED_DISPATCH, None) + self._log_through_the_release((*first.parked, first.banked.guardrail_scoped_copy())) + + def bill_partial_follow_up(self, partial: ModelResponse) -> None: + first: Final = self._first + final_round: Final = _assembled_round((partial, False)) + if first is None or final_round is None or not self._holds_round_one_alone(first): + return + self._first = None + setattr( + self._logging_obj, PARKED_DISPATCH, self._billed_for_both_rounds(first.priced, self._priced(final_round)) + ) def release_abandoned(self) -> None: if not self._owns_release(): return - if is_event_loop_running(): + if is_event_loop_running() or self._idle_reading_loop() is not None: self.finish() return - setattr(self._logging_obj, _PARKED_DISPATCH, None) + setattr(self._logging_obj, PARKED_DISPATCH, None) self._first = None self._disarm() @@ -162,7 +166,7 @@ class ChatStreamRoundBilling: return logged_alone: Final = (*first.parked, banked) self._first = replace(first, logged_alone=logged_alone, banked=banked) - setattr(self._logging_obj, _PARKED_DISPATCH, logged_alone) + setattr(self._logging_obj, PARKED_DISPATCH, logged_alone) def _release_rounds(self) -> ModelResponse | None: parked: Final = self._take_parked() @@ -196,7 +200,7 @@ class ChatStreamRoundBilling: def _priced(self, assembled: _AssembledRound) -> _PricedRound: served: Final = assembled.cache_hit or served_from_cache(assembled.response) - price: Final = _SERVED_FROM_CACHE if served else self._logging_obj.uncached_price(assembled.response) + price: Final = SERVED_FROM_CACHE if served else self._logging_obj.uncached_price(assembled.response) priced: Final = assembled.response.model_copy() set_hidden_params( priced, @@ -205,36 +209,47 @@ class ChatStreamRoundBilling: return _PricedRound(priced, assembled.usage, price, served) def _peek_parked(self) -> Parked | None: - parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None) + parked: Final[Parked | None] = getattr(self._logging_obj, PARKED_DISPATCH, None) return parked if isinstance(parked, tuple) else None def _take_parked(self) -> Parked | None: parked: Final = self._peek_parked() - setattr(self._logging_obj, _PARKED_DISPATCH, None) + setattr(self._logging_obj, PARKED_DISPATCH, None) return parked + def _holds_round_one_alone(self, first: _FirstRound) -> bool: + parked: Final = self._peek_parked() + return parked is first.parked or parked is first.logged_alone + def _owns_release(self) -> bool: - release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + release: Final[object] = getattr(self._logging_obj, DISPATCH_RELEASE, None) return release == self._log_success 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) + if getattr(self._logging_obj, DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, PARKED_DISPATCH, parked) return self._log_through_the_release(parked) def _log_through_the_release(self, parked: Parked) -> None: - armed: Final[_Release | None] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) - setattr(self._logging_obj, _DISPATCH_RELEASE, None) + armed: Final[DeferredStreamRelease | None] = getattr(self._logging_obj, DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) logged: Final = (self._log_success if armed is None else armed)(*parked) + idle_reading_loop: Final = None if is_event_loop_running() else self._idle_reading_loop() + if idle_reading_loop is not None: + idle_reading_loop.run_until_complete(logged) + return spawn_logging_task(asyncio.get_running_loop(), logged, task_name="Chat completions MCP stream success log") + def _idle_reading_loop(self) -> asyncio.AbstractEventLoop | None: + loop: Final = self._reading_loop + return loop if loop is not None and not loop.is_closed() and not loop.is_running() else None + def _disarm(self) -> None: if self._owns_release(): - setattr(self._logging_obj, _DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) - async def _log_success(self, assembled: object, cache_hit: object, banked: object = None) -> None: - logger: Final = banked if isinstance(banked, LiteLLMLoggingObj) else self._logging_obj - await _success_dispatcher(logger).dispatch_success_handlers( - assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True - ) + async def _log_success(self, *parked: object) -> None: + success: Final = parked_success(parked) + if success is not None: + await success.dispatch(self._logging_obj) diff --git a/litellm/responses/mcp/parked_stream_dispatch.py b/litellm/responses/mcp/parked_stream_dispatch.py new file mode 100644 index 00000000000..99adbbbf886 --- /dev/null +++ b/litellm/responses/mcp/parked_stream_dispatch.py @@ -0,0 +1,52 @@ +from collections.abc import Coroutine +from dataclasses import dataclass +from typing import Final, Protocol + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import UncachedPrice + +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: ... + + +class DeferredStreamRelease(Protocol): + def __call__(self, *parked: object) -> Coroutine[object, object, None]: ... + + +def success_dispatcher(logger: LiteLLMLoggingObj) -> DispatchesSuccessHandlers: + return logger + + +@dataclass(frozen=True, slots=True) +class ParkedSuccess: + assembled: object + cache_hit: object + logged_on: LiteLLMLoggingObj | None = None + + async def dispatch(self, request_logger: LiteLLMLoggingObj) -> None: + logger: Final = request_logger if self.logged_on is None else self.logged_on + await success_dispatcher(logger).dispatch_success_handlers( + self.assembled, cache_hit=self.cache_hit, start_time=None, end_time=None, prefer_async_handlers=True + ) + + +def parked_success(parked: tuple[object, ...]) -> ParkedSuccess | None: + match parked: + case (assembled, cache_hit): + return ParkedSuccess(assembled, cache_hit) + case (assembled, cache_hit, LiteLLMLoggingObj() as logged_on): + return ParkedSuccess(assembled, cache_hit, logged_on) + case _: + return None diff --git a/litellm/responses/mcp/stream_round_billing.py b/litellm/responses/mcp/stream_round_billing.py index 443147230bf..5cb256a8ac0 100644 --- a/litellm/responses/mcp/stream_round_billing.py +++ b/litellm/responses/mcp/stream_round_billing.py @@ -2,11 +2,18 @@ import asyncio from collections.abc import Callable, Coroutine, Sequence from dataclasses import dataclass from datetime import datetime -from typing import Final, Protocol, TypeAlias +from typing import Final, 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.parked_stream_dispatch import ( + DISPATCH_RELEASE, + PARKED_DISPATCH, + SERVED_FROM_CACHE, + parked_success, + success_dispatcher, +) 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 @@ -16,25 +23,6 @@ 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: @@ -198,8 +186,8 @@ class StreamedRoundBilling: 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) + if getattr(logging_obj, DISPATCH_RELEASE, None) is None: + setattr(logging_obj, DISPATCH_RELEASE, self._run_parked) def complete_round( self, @@ -286,7 +274,7 @@ class StreamedRoundBilling: def _price(self, response: ResponsesAPIResponse, round_parked: Parked | None) -> UncachedPrice: if served_from_cache(response): - return _SERVED_FROM_CACHE + return SERVED_FROM_CACHE match round_parked: case (ModelResponse() as assembled, bool()): return self._logging_obj.uncached_price(assembled) @@ -297,7 +285,7 @@ class StreamedRoundBilling: self._banked = parked self._banked_row = row if parked is not None: - setattr(self._logging_obj, _PARKED_DISPATCH, parked) + setattr(self._logging_obj, PARKED_DISPATCH, parked) def _bank_apart_from_later_rounds(self, billed: _BilledRounds) -> None: row: Final = self._banked_row @@ -320,7 +308,7 @@ class StreamedRoundBilling: @staticmethod async def _log_success(success: _SuccessLog) -> None: success.carry_bridged_usage() - await _success_dispatcher(success.logger).dispatch_success_handlers( + await success_dispatcher(success.logger).dispatch_success_handlers( success.row.logged, start_time=success.logger.start_time, end_time=success.row.end_time, @@ -329,10 +317,10 @@ class StreamedRoundBilling: ) def _take_parked(self) -> Parked | None: - parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, 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) + setattr(self._logging_obj, PARKED_DISPATCH, None) return parked async def _run_parked(self, *parked: object) -> None: @@ -341,29 +329,27 @@ class StreamedRoundBilling: 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 _ if (success := parked_success(parked)) is not None: + await success.dispatch(self._logging_obj) case _: return def _owns_release(self) -> bool: - release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + 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) + 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) + 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) + setattr(self._logging_obj, DISPATCH_RELEASE, None) @staticmethod def _close(parked: Parked | None) -> None: diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 09dfe020cf5..d54e623f6a1 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -104,11 +104,12 @@ def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway ROUND_USAGE_GUARDRAIL: Final = "mcp-llm-round-usage" +CONTENT_FILTER: Final = f"mcp-llm-content-filter-{uuid.uuid4().hex[:8]}" @pytest.fixture(scope="module") def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: - guardrail: Final[Mapping[str, JsonValue]] = { + round_usage: Final[Mapping[str, JsonValue]] = { "guardrail_name": ROUND_USAGE_GUARDRAIL, "litellm_params": { "guardrail": "integration._support.hook_count_probe.UsageRecordingPostCallGuardrail", @@ -116,7 +117,17 @@ def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[t "default_on": True, }, } - with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), (guardrail,)) as served: + content_filter: Final[Mapping[str, JsonValue]] = { + "guardrail_name": CONTENT_FILTER, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "default_on": False, + "blocked_words": [{"keyword": "forbidden-keeper", "action": "BLOCK"}], + }, + } + guardrails: Final = (round_usage, content_filter) + with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), guardrails) as served: yield served @@ -1153,6 +1164,59 @@ def test_streaming_chat_auto_execute_whose_follow_up_fails_runs_the_post_call_gu ] +def _chat_rows(call_id: str, statuses: Sequence[str]) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + 'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s', + (call_id, "acompletion"), + ), + lambda found: _statuses(found) == sorted(statuses), + seconds=70, + ) + + +def _entries_of(guardrail_name: str, row: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + return tuple( + entry + for entry in _guardrail_information(row["metadata"]) + if isinstance(entry, dict) and entry.get("guardrail_name") == guardrail_name + ) + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_counts_its_post_call_content_filter_once( + guarded_gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(guarded_gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send( + _granted_key(rig), + [AUTO], + stream=True, + stream_options={"include_usage": True}, + guardrails=[CONTENT_FILTER], + ) + assert FOLLOW_UP_OUTAGE in response.text, response.text + rows: Final = _chat_rows(response.headers["x-litellm-call-id"], ("success", "failure")) + filtered: Final = tuple( + (row["status"], entry.get("guardrail_status")) for row in rows for entry in _entries_of(CONTENT_FILTER, row) + ) + assert filtered == (("failure", "guardrail_failed_to_respond"),), rows + (guardrail_id,) = frozenset( + str(entry.get("guardrail_id") or CONTENT_FILTER) + for row in rows + for entry in _entries_of(CONTENT_FILTER, row) + ) + counted: Final = eventually( + lambda: read_rows( + 'SELECT requests_evaluated, flagged_count FROM "LiteLLM_DailyGuardrailMetrics" WHERE guardrail_id=%s', + (guardrail_id,), + ), + lambda found: bool(found), + seconds=70, + ) + assert counted == [{"requests_evaluated": 1, "flagged_count": 1}], counted + + @pytest.mark.parametrize("include_usage", (True, False), ids=("include-usage", "no-usage")) def test_streaming_chat_auto_execute_answered_in_one_round_ends_with_the_usage_it_billed( gateway: Gateway, include_usage: bool @@ -1430,6 +1494,27 @@ def test_streaming_responses_follow_up_broken_mid_stream_bills_round_one_once_an assert float(str(failure["spend"])) == pytest.approx(_price(failed_round)), rows +def test_streaming_chat_follow_up_broken_mid_answer_bills_round_one_and_its_partial_answer_with_one_fixed_margin( + gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_breaks=True) + with _rig(gateway, "chat", upstream="messages_bridge", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert "the sum" in response.text, response.text + assert ANSWER not in response.text, response.text + assert len(rig.upstream_tools()) == 2 + round_one: Final = _price(FIRST_ROUND) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(round_one)) + assert _statuses(rows) == ["failure", "success"], rows + (success,) = (row for row in rows if row["status"] == "success") + (failure,) = (row for row in rows if row["status"] == "failure") + assert (success["prompt_tokens"], success["completion_tokens"]) == FIRST_ROUND, rows + assert float(str(success["spend"])) == pytest.approx(round_one), rows + partial_answer: Final = (int(str(failure["prompt_tokens"])), int(str(failure["completion_tokens"]))) + assert min(partial_answer) > 0, rows + assert float(str(failure["spend"])) == pytest.approx(_price(partial_answer)), rows + + @pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_once( gateway: Gateway, upstream: StreamedUpstream diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index fdcbcbfcd05..4731ea89284 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -9056,6 +9056,34 @@ class TestStreamingClientDisconnectBilling: proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + @pytest.mark.asyncio + async def test_disconnect_whose_in_flight_round_billing_raises_still_closes_the_stream_and_frees_the_slot(self): + import types + + response = await self._start_partial_stream() + closed: Final = AsyncMock() + + def cannot_price_the_in_flight_round() -> None: + raise RuntimeError("the in-flight round could not be priced") + + stream: Final = types.SimpleNamespace( + chunks=[], messages=None, aclose=closed, bill_in_flight_round_on_disconnect=cannot_price_the_in_flight_round + ) + proxy_logging_obj: Final = types.SimpleNamespace(arelease_max_parallel_requests_on_disconnect=AsyncMock()) + + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=stream, + stream_completed=False, + client_disconnected=True, + user_api_key_dict=MagicMock(), + proxy_logging_obj=proxy_logging_obj, + ) + + closed.assert_awaited_once() + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + async def _bill_and_collect_success_event(self, prepare=None, request_data=None): recorder = _RecordingSuccessLogger() original_callbacks = litellm.callbacks diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 23ff0727098..45a6978e952 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1562,6 +1562,7 @@ _CHAT_ROUND_TWO: Final = MappingProxyType( ), } ) +_NEVER_LOGGED_SECONDS: Final = 60 _CONTEXT_WINDOW_ERROR: Final = MappingProxyType( { "error": MappingProxyType( @@ -1581,16 +1582,31 @@ class _ChatRoundRecorder(CustomLogger): super().__init__() self.logged: tuple[StandardLoggingPayload, ...] = () self.failed: tuple[StandardLoggingPayload, ...] = () + self._recorded: Final = asyncio.Condition() async def async_log_success_event( 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"])) + await self._announce() 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"])) + await self._announce() + + async def _announce(self) -> None: + async with self._recorded: + self._recorded.notify_all() + + async def settled(self, logged: int = 0, failed: int = 0) -> None: + async with self._recorded: + await asyncio.wait_for( + self._recorded.wait_for(lambda: len(self.logged) >= logged and len(self.failed) >= failed), + timeout=_NEVER_LOGGED_SECONDS, + ) + await _settled() def _chat_round_cost(body: Mapping[str, object]) -> float: @@ -1606,14 +1622,6 @@ async def _settled() -> None: await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) -async def _settled_once(logged: Callable[[], bool]) -> None: - for _ in range(200): - if logged(): - break - await asyncio.sleep(0.01) - await _settled() - - def _serve_chat_rounds( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response] ) -> None: @@ -1785,7 +1793,9 @@ def _cut_once_the_answer_starts() -> Response: return Response(200, stream=_CutOnceTheAnswerStarts(), headers={"content-type": "text/event-stream"}) -async def _open_chat_echo_stream(logging_obj: Logging) -> CustomStreamWrapper: +async def _open_chat_echo_stream( + logging_obj: Logging, metadata: dict[str, object] | None = None +) -> CustomStreamWrapper: stream: Final = await litellm.acompletion( model="openai/gpt-5.6", messages=[{"role": "user", "content": "Call echo with PING"}], @@ -1795,6 +1805,7 @@ async def _open_chat_echo_stream(logging_obj: Logging) -> CustomStreamWrapper: stream=True, stream_options={"include_usage": True}, litellm_logging_obj=logging_obj, + metadata=metadata, ) assert isinstance(stream, CustomStreamWrapper) return stream @@ -1860,7 +1871,7 @@ async def test_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_ await _settled() logged_before_release: Final = recorder.logged ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - await _settled_once(lambda: len(recorder.logged) > 0) + await recorder.settled(logged=1) assert logged_before_release == () _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) @@ -1888,7 +1899,7 @@ async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_bills_both_r await _read_all(stream) assert ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked) await logging_obj.async_failure_handler(blocked, "blocked by output-policy") - await _settled_once(lambda: len(recorder.failed) > 0) + await recorder.settled(failed=1) charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN assert recorder.logged == () @@ -1910,7 +1921,7 @@ async def test_stream_auto_execute_answered_without_a_tool_call_logs_its_one_rou async for chunk in stream: if chunk.choices and chunk.choices[0].finish_reason is not None: break - await _settled_once(lambda: len(recorder.logged) > 0) + await recorder.settled(logged=1) logged_once_the_answer_finished: Final = recorder.logged after_the_answer: Final = await _read_all(stream) await _settled() @@ -1947,11 +1958,30 @@ async def test_stream_auto_execute_bills_the_finished_first_round_when_the_calle with _echo_tool_served(): await stop_reading_once_the_follow_up_answers() - await _settled_once(lambda: len(recorder.logged) > 0) + await recorder.settled(logged=1) _assert_one_row_bills_round_one(recorder.logged) +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_first_round_breaks_mid_answer_bills_only_its_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_cut_once_the_answer_starts(),)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-first-round-breaks", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + with pytest.raises(litellm.ServiceUnavailableError): + await _read_all(stream) + await recorder.settled(failed=1) + + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert recorder.failed[0]["completion_tokens"] > 0 + + @pytest.mark.asyncio @pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( @@ -1976,7 +2006,7 @@ async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_r await _read_all(stream) if not ProxyLogging._discard_deferred_stream_logging_for_failure(dict(request_data), broken.value): ProxyLogging.fire_deferred_stream_logging(dict(request_data)) - await _settled_once(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + await recorder.settled(logged=1, failed=1) _assert_one_row_bills_round_one(recorder.logged) assert len(recorder.failed) == 1 @@ -2004,23 +2034,67 @@ async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_r @pytest.mark.asyncio -async def test_stream_auto_execute_leaves_no_deferred_log_armed_for_the_proxy_to_orphan_after_its_first_chunk( +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_whose_follow_up_breaks_reports_the_request_guardrail_result_on_the_failure_row_only( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-breaks-guarded", stream=True) + metadata: Final[dict[str, object]] = {} + request_data: Final = {"metadata": metadata, "litellm_logging_obj": logging_obj} + content_filter: Final = CustomGuardrail(guardrail_name="request-content-filter", event_hook="post_call") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj, metadata) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError) as broken: + await _read_all(stream) + content_filter.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=str(broken.value), + request_data=request_data, + guardrail_status="guardrail_failed_to_respond", + ) + if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, broken.value): + ProxyLogging.fire_deferred_stream_logging(request_data) + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert recorder.logged[0]["guardrail_information"] is None + assert [entry["guardrail_name"] for entry in recorder.failed[0]["guardrail_information"] or ()] == [ + "request-content-filter" + ] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_served_through_the_proxy_stream_hook_logs_both_rounds_once_when_it_ends( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) recorder: Final = _ChatRoundRecorder() - logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-first-chunk-read-ahead", stream=True) + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-proxy-stream-hook", stream=True) with _echo_tool_served(): - stream: Final = await _open_chat_echo_stream(logging_obj) - rest: Final = stream.__aiter__() - first_chunk: Final = await rest.__anext__() - armed_after_the_first_chunk: Final = getattr(logging_obj, "_on_deferred_stream_complete", None) - chunks: Final = (first_chunk, *[chunk async for chunk in rest]) - await _settled_once(lambda: len(recorder.logged) > 0) + served: Final = ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=await _open_chat_echo_stream(logging_obj), + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj}, + ) + chunks: Final = tuple([chunk async for chunk in served]) + await recorder.settled(logged=1) - assert armed_after_the_first_chunk is None _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) @@ -2035,17 +2109,152 @@ async def test_stream_auto_execute_whose_follow_up_streams_nothing_still_bills_t with _echo_tool_served(): await _read_all(await _open_chat_echo_stream(logging_obj)) - await _settled_once(lambda: len(recorder.logged) > 0) + await recorder.settled(logged=1) _assert_one_row_bills_round_one(recorder.logged) -async def _stream_read_until_the_follow_up_answers( - logging_obj: Logging, opened: "queue.SimpleQueue[CustomStreamWrapper]" +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_follow_up_cannot_be_assembled_ends_with_the_usage_of_the_round_it_billed( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-unassembled", stream=True) + + def cannot_assemble(*args: object, **kwargs: object) -> ModelResponse: + raise ValueError("the follow-up's chunks cannot be assembled") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + await _read_until(stream, "PING") + monkeypatch.setattr(litellm, "stream_chunk_builder", cannot_assemble) + rest: Final = await _read_all(stream) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + usage: Final = _streamed_usage(rest) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (71, 19, 90) + assert usage.cost == pytest.approx(recorder.logged[0]["response_cost"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_during_the_follow_up_bills_round_one_and_the_partial_answer_on_one_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-mid-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _read_until(stream, "PING") + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + assert row["id"] == _CHAT_ROUND_ONE["id"] + partial_prompt, partial_completion = (row["prompt_tokens"] - 71, row["completion_tokens"] - 19) + assert partial_prompt > 0 and partial_completion > 0 + partial_answer: Final = litellm.completion_cost( + completion_response=ModelResponse( + model="gpt-5.6", + usage=Usage( + prompt_tokens=partial_prompt, + completion_tokens=partial_completion, + total_tokens=partial_prompt + partial_completion, + ), + ), + model="gpt-5.6", + custom_llm_provider="openai", + ) + assert row["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE) + partial_answer - _FIXED_MARGIN) + assert recorder.failed == () + + +async def _read_through_the_usage_chunk( + chunks: AsyncIterator[ModelResponseStream], read: tuple[ModelResponseStream, ...] = () +) -> tuple[ModelResponseStream, ...]: + chunk: Final = await anext(chunks) + if isinstance(getattr(chunk, "usage", None), Usage): + return (*read, chunk) + return await _read_through_the_usage_chunk(chunks, (*read, chunk)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_after_its_whole_answer_bills_both_rounds_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-after-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + chunks: Final = await _read_through_the_usage_chunk(aiter(stream)) + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + assert recorder.failed == () + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bridged_to_messages_bills_round_one_when_its_follow_up_breaks( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-bridged-to-messages", stream=True) + + async def bridged(stream: CustomStreamWrapper) -> AsyncIterator[ModelResponseStream]: + async for chunk in stream: + yield chunk + + with _echo_tool_served(): + relayed: Final = bridged(await _open_chat_echo_stream(logging_obj)) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=relayed, route_type="anthropic_messages", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError): + _ = [chunk async for chunk in relayed] + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + + +def _read_synchronously_until_the_follow_up_answers( + logging_obj: Logging, loop: asyncio.AbstractEventLoop, before_dropping: Callable[[], None] = lambda: None ) -> None: - stream: Final = await _open_chat_echo_stream(logging_obj) - await _read_until(stream, "PING") - opened.put(stream) + stream: Final = loop.run_until_complete(_open_chat_echo_stream(logging_obj)) + for chunk in stream: + if chunk.choices and chunk.choices[0].delta.content == "PING": + before_dropping() + return + pytest.fail("the stream never sent 'PING'") _CHAT_STREAM_BILLING_NAMES: Final = ("MCPStreamingIterator", "ChatStreamRoundBilling", "_log_success") @@ -2063,26 +2272,57 @@ class _UnraisableRecorder: 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( +def test_stream_auto_execute_read_synchronously_and_dropped_mid_answer_bills_the_finished_first_round( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) recorder: Final = _ChatRoundRecorder() unraisable: Final = _UnraisableRecorder() - opened: Final = queue.SimpleQueue[CustomStreamWrapper]() + loop: Final = asyncio.new_event_loop() gc.collect() monkeypatch.setattr(sys, "unraisablehook", unraisable) + asyncio.set_event_loop(loop) - with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - asyncio.run( - _stream_read_until_the_follow_up_answers( - _chat_rounds_logging_obj(recorder, "mcp-chat-stream-loop-closed", stream=True), opened + try: + with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _read_synchronously_until_the_follow_up_answers( + _chat_rounds_logging_obj(recorder, "mcp-chat-stream-sync-dropped", stream=True), loop ) - ) - opened.get_nowait() - gc.collect() + gc.collect() + loop.run_until_complete(recorder.settled(logged=1)) + finally: + asyncio.set_event_loop(None) + loop.close() + + _assert_one_row_bills_round_one(recorder.logged) + assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == [] + + +def test_stream_auto_execute_read_synchronously_and_dropped_after_its_loop_closed_drops_without_errors( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + unraisable: Final = _UnraisableRecorder() + loop: Final = asyncio.new_event_loop() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + asyncio.set_event_loop(loop) + + try: + with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _read_synchronously_until_the_follow_up_answers( + _chat_rounds_logging_obj(recorder, "mcp-chat-stream-sync-dropped-after-close", stream=True), + loop, + before_dropping=loop.close, + ) + gc.collect() + finally: + asyncio.set_event_loop(None) + loop.close() - assert recorder.logged == () assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == [] assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == []