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))] == []