From 160a0b72cff4dbbc41bd5de1130966a207f0c17d Mon Sep 17 00:00:00 2001 From: simonp9 <349053457+simonp9@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:34:10 -0700 Subject: [PATCH] fix(mcp): bill both calls of a streamed chat completions MCP auto-execute (#48899) * fix(mcp): bill both rounds of a streamed chat completions MCP auto-execute A streamed /v1/chat/completions request with an MCP tool set to require_approval "never" runs two model calls on one logging object, and the final-stream dedupe kept only the first round's spend row. The final usage chunk also showed only the second round. The MCP stream iterator now holds the first round's assembled response, prices it, and logs one success row that carries the usage and cost of both rounds, with the fixed margin charged once. The final usage chunk shows the same summed usage. A single-round stream still logs its one row when the answer finishes. When the follow-up round fails, the first round is logged from a copy of the logging object taken before the follow-up opened, and the failure row carries only the follow-up's partial usage with no second fixed margin. With proxy post-call guardrails armed, only (assembled, cache_hit) tuples are parked for the proxy, and the deferred log is armed only while the iterator reads, so the proxy's orphaned-callback cleanup never sees it after the Router's first-chunk read * fix(mcp): send the usage chunk of a streamed chat MCP request that ends in one round When a streamed chat completions MCP request ends in its first round (the model answers without a tool call, or no tool results come back), the iterator drained the inner stream to fire its spend log and dropped the usage chunk that stream produced, so a client that asked for include_usage got no usage chunk at all. The drained usage chunk is now held and sent to the client after the final content chunk and before [DONE], the same place a plain stream sends it, so the last usage chunk matches the billed row. A two-round request still drops the first round's usage chunk and sends the merged one from the follow-up, and a stream without include_usage still sends none * fix(mcp): hold round one's log for the guardrail, bill the chunk cost once, and keep round one's id When the follow-up call of a streamed chat MCP auto-execute fails, round one's success log now goes through the proxy's deferred guardrail release instead of straight to the handlers, so a post-call guardrail runs on it and the row carries it. Without a proxy release it logs directly as before The last usage chunk now takes its cost from the combined row's billed total, so a fixed cost margin is counted once in the chunk, as it is on the row. The follow-up's usage chunk is held until the follow-up ends to make that possible The combined row now keeps round one's id, the stream's first chunk id, which matches how a single-round stream is logged * 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 * test(mcp): drive the unassembled follow-up through an upstream audio delta that is not base64 The test replaced litellm.stream_chunk_builder with a stub that raises. Round 2 now streams an audio delta whose data is not valid base64, so the real builder raises, the stream parks only best-effort usage, and the client's last usage chunk must carry round 1's tokens and cost * fix(mcp): count the fixed margin once in the combined stream row's usage cost * fix(mcp): bill round 2's partial answer when a client hangs up on a routed stream On the proxy the Router wraps the MCP stream, and when a client hangs up during round 2 the router's own cleanup closes that stream before the proxy's disconnect cleanup runs. The close logged round 1 alone, so round 2's partial answer was never billed. The MCP stream now bills round 1 plus round 2's partial answer as soon as the hang-up's cancellation reaches its read, before anything closes it * fix(mcp): count the follow-up's own messages in round 2's partial prompt The follow-up stream shares the request's logger, so it took the client's original messages as its prompt. When round 2 ended early, from a hang-up or a broken upstream, its partial usage counted only those messages and left out the assistant tool call and the tool results the follow-up actually sent. The follow-up stream now counts the messages the follow-up call sent * fix(mcp): bill round 2's partial when a routed stream's client hangs up between reads On a proxy the stream the disconnect cleanup holds is the router's fallback wrapper, so the cleanup never found the MCP stream's in-flight billing and released round 1 alone when the client hung up while the proxy sat at a yield. The round billing now registers its in-flight biller on the request's logging object, next to the parked release, and the cleanup looks it up there, so any wrapper chain bills round 1 plus round 2's partial once * fix(mcp): send the MCP stream usage chunk under the billed row id and count the partial off the loop A streamed chat completion that runs a LiteLLM-hosted MCP tool bills both model calls on one row under round 1's chatcmpl id, but the usage chunk it sent the client was round 2's held chunk and kept round 2's id, so a caller could not join the usage chunk to its row. The usage chunk now carries round 1's id in every path that sends it, with or without a deferred post-call guardrail, while its usage and cost and every content chunk's id stay as they were When the client hangs up during round 2, assembling round 2's partial answer counts the follow-up messages, tool result included, with the local tokenizer. That ran on the event loop and stalled every other request on the worker for about a second on a large tool result. The assembly now runs on a worker thread through asyncify from both the reader's cancellation and the proxy's disconnect cleanup, shielded from cancellation, and round 1 is checked again after the await so a release that happened meanwhile is never billed twice * test(mcp): wait on the loop with a hang guard instead of a 300 ms deadline * fix(mcp): give a proxy-rejected follow-up's failure row round 1's own cost breakdown * refactor(streaming): fold the partial response's empty guards into one check --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 22 + .../litellm_core_utils/streaming_handler.py | 35 +- litellm/proxy/common_request_processing.py | 45 +- litellm/proxy/utils.py | 3 + .../responses/mcp/chat_completions_handler.py | 101 +- .../mcp/chat_stream_round_billing.py | 281 +++++ .../responses/mcp/parked_stream_dispatch.py | 72 ++ litellm/responses/mcp/stream_round_billing.py | 58 +- .../integration/_support/hook_count_probe.py | 18 + .../integration/mcp/test_mcp_llm_endpoints.py | 261 +++- .../proxy/test_common_request_processing.py | 28 + .../mcp/test_chat_completions_handler.py | 1050 ++++++++++++++++- .../mcp/test_parked_stream_dispatch.py | 17 + 13 files changed, 1894 insertions(+), 97 deletions(-) create mode 100644 litellm/responses/mcp/chat_stream_round_billing.py create mode 100644 litellm/responses/mcp/parked_stream_dispatch.py create mode 100644 tests/unit/responses/mcp/test_parked_stream_dispatch.py diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4dd33fa5565..8bfb0b60693 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -568,6 +568,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: @@ -2566,6 +2571,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 682b2a2f538..098a404ea14 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2486,6 +2486,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": + 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 self.chunks + else None + ) + usage: Final[object] = getattr(partial, "usage", None) + if partial is None or 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 @@ -2495,23 +2514,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 a79d8fe29db..48eb58adaf2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -140,6 +140,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 in_flight_round_biller, 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_api import RoutingContext @@ -535,6 +536,16 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return await _dispatch_disconnect_billing_event(logging_obj, partial_response) +async def _bill_in_flight_round_on_disconnect(request_data: Mapping[str, object]) -> None: + bill_in_flight_round: Final = in_flight_round_biller(request_data.get("litellm_logging_obj")) + if bill_in_flight_round is None: + return + try: + await bill_in_flight_round() + 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( @@ -2083,6 +2094,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 @@ -4151,11 +4169,21 @@ class ProxyBaseLLMRequestProcessing: _captured_user_api_key_dict: Final = user_api_key_dict _captured_logging_obj: Final = logging_obj - async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None: + async def _on_deferred_stream_complete( + 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, + captured_logging_obj=_captured_logging_obj if logged_on is None else logged_on, assembled_response=assembled_response, cache_hit=cache_hit, ) @@ -4180,14 +4208,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", @@ -4524,6 +4546,7 @@ class ProxyBaseLLMRequestProcessing: client_disconnected, ) if recorded_client_disconnect: + await _bill_in_flight_round_on_disconnect(_REQUEST_FIELDS.validate_python(request_data)) 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/proxy/utils.py b/litellm/proxy/utils.py index 74e16628bc8..ddaa8d2e69d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,6 +91,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( ) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction +from litellm.responses.mcp.parked_stream_dispatch import held_round_spend from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, ModelInfo, Usage @@ -4943,6 +4944,8 @@ class ProxyLogging: if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS): return False match _args: + case (ModelResponse(), *_) if (held := held_round_spend(logging_obj)) is not None: + logging_obj.record_billed_rounds_for_failure(held) case (ModelResponse() as assembled, *_): logging_obj.record_assembled_response_for_failure(assembled) case (BilledRoundsDispatch() as billed,): diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index cd072825ef5..1e07b768f62 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -1,15 +1,18 @@ """Helpers for handling MCP-aware `/chat/completions` requests.""" +import asyncio import logging from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +import anyio from pydantic import TypeAdapter from typing_extensions import TypedDict, Unpack from litellm._internal_context import internal_sub_call from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.chat_stream_round_billing import ChatStreamRoundBilling, engaged_round_billing from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) @@ -19,7 +22,7 @@ from litellm.responses.mcp.round_billing import ( first_round_billed_on_failure, summed_cost_breakdown, ) -from litellm.types.utils import Message, ModelResponse +from litellm.types.utils import Message, ModelResponse, Usage from litellm.utils import CustomStreamWrapper if TYPE_CHECKING: @@ -258,10 +261,20 @@ async def acompletion_with_mcp( self.follow_up_stream = None self.follow_up_iterator = None self.follow_up_exhausted = False + self._round_billing: ChatStreamRoundBilling | None = engaged_round_billing(initial_stream.logging_obj) + self._unsent_usage_chunk: ModelResponseStream | None = None def __aiter__(self): return self + def __del__(self) -> None: + self.release_abandoned() + + def release_abandoned(self) -> None: + round_billing: Final[object] = getattr(self, "_round_billing", None) + if isinstance(round_billing, ChatStreamRoundBilling): + round_billing.release_abandoned() + def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_list_tools to the first chunk.""" from litellm.types.utils import ( @@ -321,18 +334,54 @@ async def acompletion_with_mcp( return chunk - async def _drain_inner_stream(self): + async def _drain_inner_stream(self) -> tuple[ModelResponseStream, ...]: try: - while True: - await self._stream_iterator.__anext__() - except StopAsyncIteration: - pass + return tuple([chunk async for chunk in initial_stream]) except Exception: logging.getLogger("LiteLLM").exception( "Error draining inner MCP stream after final chunk; spend logging may be incomplete" ) + return () - async def __anext__(self): + async def _complete_first_round(self) -> ModelResponseStream | None: + drained: Final = await self._drain_inner_stream() + if self._round_billing is not None: + self._round_billing.complete_first_round() + usage_chunks: Final = tuple( + chunk for chunk in drained if isinstance(getattr(chunk, "usage", None), Usage) + ) + return usage_chunks[-1] if usage_chunks else None + + async def _follow_up_or_release_first_round(self, usage_chunk: ModelResponseStream | None) -> None: + if self.tool_results and self.complete_response: + await self._prepare_follow_up_call() + if self.follow_up_stream is not None: + return + self._unsent_usage_chunk = usage_chunk + if self._round_billing is not None: + self._round_billing.finish() + + async def __anext__(self) -> ModelResponseStream: + round_billing: Final = self._round_billing + if round_billing is None: + return await self._next_chunk() + with round_billing.armed_while_reading(): + try: + return await self._next_chunk() + except StopAsyncIteration: + usage_chunk: Final = round_billing.finish() + if usage_chunk is None: + raise + return usage_chunk + except asyncio.CancelledError: + with anyio.CancelScope(shield=True): + await round_billing.bill_in_flight_round() + raise + except Exception: + round_billing.fail() + raise + + async def _next_chunk(self) -> ModelResponseStream: # Phase 1: Collect and yield initial stream chunks if not self.stream_exhausted: # Get the iterator from the stream wrapper @@ -362,29 +411,22 @@ async def acompletion_with_mcp( if is_final: self.stream_exhausted = True + usage_chunk: Final = await self._complete_first_round() await self._process_tool_calls() chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk) - if self.tool_results and self.complete_response: - await self._prepare_follow_up_call() - # Drain inner stream so CustomStreamWrapper fires its - # end-of-stream handler (dispatch_success_handlers → - # _ProxyDBLogger → LiteLLM_SpendLogs). The CSW may - # yield one usage chunk before raising StopAsyncIteration. - await self._drain_inner_stream() + await self._follow_up_or_release_first_round(usage_chunk) return chunk except StopAsyncIteration: self.stream_exhausted = True + drained_usage_chunk: Final = await self._complete_first_round() # Process tool calls after stream is exhausted await self._process_tool_calls() # If we have chunks, yield the final one with metadata if self.collected_chunks: final_chunk = self.collected_chunks[-1] final_chunk = self._add_mcp_tool_metadata_to_final_chunk(final_chunk) - # If we have tool results, prepare follow-up call - if self.tool_results and self.complete_response: - await self._prepare_follow_up_call() - await self._drain_inner_stream() + await self._follow_up_or_release_first_round(drained_usage_chunk) return final_chunk # Phase 2: Yield follow-up stream chunks if available @@ -400,7 +442,6 @@ async def acompletion_with_mcp( from litellm._logging import verbose_logger verbose_logger.debug("Follow-up chunk yielded: %s", chunk) - return chunk except StopAsyncIteration: self.follow_up_exhausted = True from litellm._logging import verbose_logger @@ -408,6 +449,10 @@ async def acompletion_with_mcp( verbose_logger.debug("Follow-up stream exhausted") # After follow-up stream is exhausted, check if we need to raise StopAsyncIteration raise StopAsyncIteration + client_chunk: Final = ( + chunk if self._round_billing is None else self._round_billing.client_chunk(chunk) + ) + return await self._next_chunk() if client_chunk is None else client_chunk # If we're here and follow_up_stream is None but we expected it, log a warning if ( @@ -420,7 +465,11 @@ async def acompletion_with_mcp( verbose_logger.warning("Follow-up stream was not created despite having tool results") - raise StopAsyncIteration + unsent_usage_chunk: Final = self._unsent_usage_chunk + if unsent_usage_chunk is None: + raise StopAsyncIteration + self._unsent_usage_chunk = None + return unsent_usage_chunk async def _process_tool_calls(self): """Process tool calls after streaming completes.""" @@ -488,10 +537,16 @@ async def acompletion_with_mcp( # This ensures the patch works correctly in tests import litellm + earlier_rounds: Final = None if self._round_billing is None else self._round_billing.bank_first_round() follow_up_response: Final = await litellm.acompletion(**follow_up_call_args) # Ensure follow-up response is a CustomStreamWrapper if isinstance(follow_up_response, CustomStreamWrapper): + if earlier_rounds is not None: + follow_up_response.bill_failure_after_earlier_rounds(earlier_rounds) + follow_up_response.messages = follow_up_messages + if self._round_billing is not None: + self._round_billing.follow_up_started(follow_up_response.partial_response) self.follow_up_stream = follow_up_response from litellm._logging import verbose_logger @@ -548,11 +603,13 @@ async def acompletion_with_mcp( def __aiter__(self): return self._custom_iterator + async def aclose(self) -> None: + self._custom_iterator.release_abandoned() + await super().aclose() + def __iter__(self): # For synchronous iteration, create a sync wrapper if self._sync_iterator is None: - import asyncio - try: self._sync_loop = asyncio.get_event_loop() except RuntimeError: diff --git a/litellm/responses/mcp/chat_stream_round_billing.py b/litellm/responses/mcp/chat_stream_round_billing.py new file mode 100644 index 00000000000..370de663606 --- /dev/null +++ b/litellm/responses/mcp/chat_stream_round_billing.py @@ -0,0 +1,281 @@ +import asyncio +from collections.abc import Callable, Generator +from contextlib import contextmanager +from dataclasses import dataclass, replace +from types import MappingProxyType +from typing import Final, TypeAlias + +from litellm.litellm_core_utils.asyncify import asyncify, 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, + HELD_ROUND_SPEND, + IN_FLIGHT_ROUND, + 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, ...] +PartialAnswer: TypeAlias = Callable[[], object] + + +@dataclass(frozen=True, slots=True) +class _AssembledRound: + response: ModelResponse + usage: Usage + cache_hit: bool + parked: Parked + + +@dataclass(frozen=True, slots=True) +class _PricedRound: + response: ModelResponse + usage: Usage + price: UncachedPrice + cache_hit: bool + + @property + def spend(self) -> BilledRoundsSpend: + return BilledRoundsSpend(self.usage, self.price.cost or 0.0, self.price.cost_breakdown) + + +@dataclass(frozen=True, slots=True) +class _HeldUsageChunk: + chunk: ModelResponseStream + usage: Usage + + +@dataclass(frozen=True, slots=True) +class _FirstRound: + priced: _PricedRound + parked: Parked + logged_alone: Parked + banked: LiteLLMLoggingObj | None = None + usage_chunk: _HeldUsageChunk | None = None + + def client_usage_chunk(self, billed: ModelResponse | None) -> ModelResponseStream | None: + if self.usage_chunk is None: + return None + 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({"id": self.priced.response.id, "usage": usage}) + ) + + +def _cost(response: ModelResponse) -> object: + hidden_params: Final = get_hidden_params(response) + return None if hidden_params is None else hidden_params.get("response_cost") + + +def _assembled_round(parked: Parked | None) -> _AssembledRound | None: + match parked: + case (ModelResponse() as response, cache_hit) if isinstance(usage := getattr(response, "usage", None), Usage): + return _AssembledRound(response, usage, cache_hit is True, parked) + case _: + return None + + +def engaged_round_billing(logging_obj: object) -> "ChatStreamRoundBilling | None": + return ChatStreamRoundBilling(logging_obj) if isinstance(logging_obj, LiteLLMLoggingObj) else None + + +class ChatStreamRoundBilling: + def __init__(self, logging_obj: LiteLLMLoggingObj) -> None: + self._logging_obj: Final = logging_obj + self._first: _FirstRound | None = None + self._partial_follow_up: PartialAnswer | 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 + if armed: + setattr(self._logging_obj, DISPATCH_RELEASE, self._log_success) + self._reading_loop = asyncio.get_running_loop() + try: + yield + finally: + if armed and self._peek_parked() is None: + self._disarm() + + def complete_first_round(self) -> None: + assembled: Final = _assembled_round(self._peek_parked()) + if assembled is not None: + self._first = _FirstRound(self._priced(assembled), assembled.parked, assembled.parked) + setattr(self._logging_obj, HELD_ROUND_SPEND, self._held_round_one_spend) + + def bank_first_round(self) -> BilledRoundsSpend | None: + self._bank_apart_from_the_follow_up() + return None if self._first is None else self._first.priced.spend + + def client_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None: + usage: Final[object] = getattr(chunk, "usage", None) + if self._first is None or not isinstance(usage, Usage): + return chunk + self._first = replace(self._first, usage_chunk=_HeldUsageChunk(chunk, usage)) + return None + + def finish(self) -> ModelResponseStream | None: + first: Final = self._first + billed: Final = self._release_rounds() + return None if first is None else first.client_usage_chunk(billed) + + 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 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.guardrail_scoped_copy())) + + def follow_up_started(self, partial_follow_up: PartialAnswer) -> None: + self._partial_follow_up = partial_follow_up + setattr(self._logging_obj, IN_FLIGHT_ROUND, self.bill_in_flight_round) + + async def bill_in_flight_round(self) -> None: + partial_follow_up: Final = self._partial_follow_up + if partial_follow_up is None or self._first_round_held_alone() is None: + return + final_round: Final = _assembled_round((await asyncify(partial_follow_up)(), False)) + first: Final = self._first_round_held_alone() + if final_round is None or first is None: + 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() or self._idle_reading_loop() is not None: + self.finish() + return + setattr(self._logging_obj, PARKED_DISPATCH, None) + self._first = None + self._disarm() + + def _bank_apart_from_the_follow_up(self) -> None: + first: Final = self._first + if first is None or first.banked is not None or self._peek_parked() is not first.parked: + return + banked: Final = self._logging_obj.attempt_scoped_copy() + if not self._owns_release(): + self._first = replace(first, banked=banked) + 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) + + def _release_rounds(self) -> ModelResponse | None: + parked: Final = self._take_parked() + first: Final = self._first + self._first = None + if parked is None: + self._disarm() + return None + if first is None or parked is first.logged_alone: + self._release(parked) + return None + final_round: Final = _assembled_round(parked) + if final_round is None: + self._release(first.logged_alone) + return None + billed, cache_hit = self._billed_for_both_rounds(first.priced, self._priced(final_round)) + self._release((billed, cache_hit)) + return billed + + def _billed_for_both_rounds(self, first: _PricedRound, final: _PricedRound) -> tuple[ModelResponse, bool]: + first_breakdown: Final = first.price.cost_breakdown + final_breakdown: Final = final.price.cost_breakdown + rounds: Final = billed_for_every_round(first.response, first_breakdown, final.response, final_breakdown) + usage: Final[object] = getattr(rounds, "usage", None) + charged_usage: Final = ( + usage.model_copy(update=MappingProxyType({"cost": _cost(rounds)})) if isinstance(usage, Usage) else usage + ) + billed: Final = rounds.model_copy(update=MappingProxyType({"id": first.response.id, "usage": charged_usage})) + cache_hit: Final = first.cache_hit and final.cache_hit + set_hidden_params(billed, {**(get_hidden_params(rounds) or {}), "cache_hit": cache_hit}) + self._logging_obj.cost_breakdown = summed_cost_breakdown( + first.response, first_breakdown, final.response, final_breakdown + ) + return billed, cache_hit + + 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) + priced: Final = assembled.response.model_copy() + set_hidden_params( + priced, + {**(get_hidden_params(assembled.response) or {}), "response_cost": price.cost, "cache_hit": served}, + ) + return _PricedRound(priced, assembled.usage, price, served) + + def _peek_parked(self) -> Parked | 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) + return parked + + def _first_round_held_alone(self) -> _FirstRound | None: + first: Final = self._first + return first if first is not None and self._holds_round_one_alone(first) else None + + def _held_round_one_spend(self) -> BilledRoundsSpend | None: + first: Final = self._first_round_held_alone() + return None if first is None else first.priced.spend + + 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) + 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) + return + self._log_through_the_release(parked) + + def _log_through_the_release(self, parked: Parked) -> 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) + + 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..792363f0bab --- /dev/null +++ b/litellm/responses/mcp/parked_stream_dispatch.py @@ -0,0 +1,72 @@ +from collections.abc import Coroutine +from dataclasses import dataclass +from typing import Final, Protocol + +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +PARKED_DISPATCH: Final = "_deferred_stream_complete_args" +DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" +IN_FLIGHT_ROUND: Final = "_bill_in_flight_round_on_disconnect" +HELD_ROUND_SPEND: Final = "_held_round_spend_for_failure" +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]: ... + + +class InFlightRoundBiller(Protocol): + def __call__(self) -> Coroutine[object, object, None]: ... + + +class HeldRoundSpend(Protocol): + def __call__(self) -> BilledRoundsSpend | 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 in_flight_round_biller(logging_obj: object) -> InFlightRoundBiller | None: + biller: Final[InFlightRoundBiller | None] = getattr(logging_obj, IN_FLIGHT_ROUND, None) + return biller if isinstance(logging_obj, LiteLLMLoggingObj) and callable(biller) else None + + +def held_round_spend(logging_obj: object) -> BilledRoundsSpend | None: + held: Final[HeldRoundSpend | None] = getattr(logging_obj, HELD_ROUND_SPEND, None) + return held() if isinstance(logging_obj, LiteLLMLoggingObj) and callable(held) else None + + +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/_support/hook_count_probe.py b/tests/integration/_support/hook_count_probe.py index 0d4056e8fb9..3cd72f1a6ba 100644 --- a/tests/integration/_support/hook_count_probe.py +++ b/tests/integration/_support/hook_count_probe.py @@ -75,6 +75,24 @@ class BlockingPostCallGuardrail(CustomGuardrail): raise HTTPException(status_code=400, detail="blocked by the hook-count guardrail") +class UsageRecordingPostCallGuardrail(CustomGuardrail): + """Records the token usage of the response it is handed in the request's guardrail information, so the spend row + shows which response the guardrail saw.""" + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: object, response: object + ) -> None: + usage: Final = getattr(response, "usage", None) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={ + "prompt_tokens": getattr(usage, "prompt_tokens", None), + "completion_tokens": getattr(usage, "completion_tokens", None), + }, + request_data=data, + guardrail_status="success", + ) + + probe: Final = HookCountProbe() sync_only_probe: Final = SyncOnlyHookCountProbe() diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 4b967802893..f6743a8751a 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -46,9 +46,10 @@ from pydantic import JsonValue, TypeAdapter pytestmark: Final = [pytest.mark.xdist_group("mcp_llm_private_gateway"), pytest.mark.shares_worker_state] -@pytest.fixture(scope="session") -def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: - directory: Final = tmp_path_factory.mktemp("mcp-llm-gateway") +@contextmanager +def _private_gateway( + directory: Path, guardrails: Sequence[Mapping[str, JsonValue]] = () +) -> Iterator[tuple[Gateway, Mapping[str, str]]]: stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) config: Final = directory / "stored-responses.yaml" (directory / f"{STREAM_BLOCK_MODULE}.py").write_text(STREAM_BLOCK_SOURCE) @@ -68,7 +69,8 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat { "guardrail_name": STREAM_BLOCK, "litellm_params": {"guardrail": f"{STREAM_BLOCK_MODULE}.BlockGuard", "mode": "post_call"}, - } + }, + *guardrails, ], } ) @@ -81,15 +83,61 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat yield private, services -@pytest.fixture -def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: - private, services = _llm_gateway +@contextmanager +def _serving(served: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + private, services = served with pytest.MonkeyPatch.context() as environment: for name, value in services.items(): environment.setenv(name, value) yield private +@pytest.fixture(scope="session") +def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: + with _private_gateway(tmp_path_factory.mktemp("mcp-llm-gateway")) as served: + yield served + + +@pytest.fixture +def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + with _serving(_llm_gateway) as private: + yield private + + +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]]]: + round_usage: Final[Mapping[str, JsonValue]] = { + "guardrail_name": ROUND_USAGE_GUARDRAIL, + "litellm_params": { + "guardrail": "integration._support.hook_count_probe.UsageRecordingPostCallGuardrail", + "mode": "post_call", + "default_on": True, + }, + } + 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 + + +@pytest.fixture +def guarded_gateway(_guarded_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + with _serving(_guarded_llm_gateway) as private: + yield private + + Surface = Literal["chat", "responses", "messages", "messages_bridge", "responses_bridge"] SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge") ADD: Final = {"a": 2, "b": 3} @@ -1032,6 +1080,184 @@ def test_nonstream_failed_follow_up_bills_round_one_on_the_failure_row( ) +def _streamed_chunks(response: httpx.Response) -> tuple[dict[str, object], ...]: + return tuple( + json.loads(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def test_streaming_chat_auto_execute_logs_the_answer_and_bills_both_rounds(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert response.status_code == 200, response.text + chunks: Final = _streamed_chunks(response) + text: Final = "".join( + str(chunk["choices"][0]["delta"].get("content") or "") for chunk in chunks if chunk.get("choices") + ) + assert text == ANSWER, response.text + usages: Final = tuple(chunk["usage"] for chunk in chunks if chunk.get("usage")) + assert usages, response.text + rows: Final = _spend_rows(response.headers["x-litellm-call-id"], "acompletion", "success") + assert len(rows) == 1, rows + row: Final = rows[0] + assert row["request_id"] == chunks[0]["id"], (row, chunks[0]) + assert (row["prompt_tokens"], row["completion_tokens"]) == ( + FIRST_ROUND[0] + FOLLOW_UP_ROUND[0], + FIRST_ROUND[1] + FOLLOW_UP_ROUND[1], + ), row + assert float(str(row["spend"])) == pytest.approx(_price(FIRST_ROUND, FOLLOW_UP_ROUND)), row + final_usage: Final = usages[-1] + assert (final_usage["prompt_tokens"], final_usage["completion_tokens"], final_usage["total_tokens"]) == ( + row["prompt_tokens"], + row["completion_tokens"], + sum((*FIRST_ROUND, *FOLLOW_UP_ROUND)), + ), usages + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_bills_round_one_once(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert FOLLOW_UP_OUTAGE in response.text, response.text + call_id: Final = response.headers["x-litellm-call-id"] + succeeded: Final = _spend_rows(call_id, "acompletion", "success") + assert len(succeeded) == 1, succeeded + assert (succeeded[0]["prompt_tokens"], succeeded[0]["completion_tokens"]) == FIRST_ROUND, succeeded + assert float(str(succeeded[0]["spend"])) == pytest.approx(_price(FIRST_ROUND)), succeeded + failed: Final = _spend_rows(call_id, "acompletion", "failure") + assert [float(str(row["spend"])) for row in failed] == [0.0], failed + + +def _guardrail_information(metadata: JsonValue) -> tuple[JsonValue, ...]: + entries: Final = object_value(metadata).get("guardrail_information") + return tuple(entries) if isinstance(entries, list) else () + + +def _round_usage_seen_by_the_guardrail(call_id: str) -> list[JsonValue]: + rows: Final = read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s AND status=%s', + (call_id, "acompletion", "success"), + ) + entries: Final = tuple(chain.from_iterable(_guardrail_information(row["metadata"]) for row in rows)) + return [ + entry.get("guardrail_response") + for entry in entries + if isinstance(entry, dict) and entry.get("guardrail_name") == ROUND_USAGE_GUARDRAIL + ] + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_runs_the_post_call_guardrail_on_round_one( + 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}) + assert FOLLOW_UP_OUTAGE in response.text, response.text + call_id: Final = response.headers["x-litellm-call-id"] + succeeded: Final = _spend_rows(call_id, "acompletion", "success") + assert len(succeeded) == 1, succeeded + assert (succeeded[0]["prompt_tokens"], succeeded[0]["completion_tokens"]) == FIRST_ROUND, succeeded + assert _round_usage_seen_by_the_guardrail(call_id) == [ + {"prompt_tokens": FIRST_ROUND[0], "completion_tokens": FIRST_ROUND[1]} + ] + + +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 +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send( + _granted_key(rig), + [AUTO], + prompt=f"{DIRECT} {rig.alias}", + stream=True, + stream_options={"include_usage": include_usage}, + ) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + chunks: Final = _streamed_chunks(response) + text: Final = "".join( + str(chunk["choices"][0]["delta"].get("content") or "") for chunk in chunks if chunk.get("choices") + ) + assert text == ANSWER, response.text + rows: Final = _spend_rows(response.headers["x-litellm-call-id"], "acompletion", "success") + assert len(rows) == 1 and (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == FIRST_ROUND, rows + usage_at: Final = tuple(index for index, chunk in enumerate(chunks) if chunk.get("usage")) + if not include_usage: + assert usage_at == (), response.text + return + assert usage_at == (len(chunks) - 1,), response.text + finished_at: Final = tuple( + index + for index, chunk in enumerate(chunks) + if chunk.get("choices") and chunk["choices"][0].get("finish_reason") is not None + ) + assert finished_at and finished_at[-1] < usage_at[0], response.text + usage: Final = chunks[-1]["usage"] + assert (usage["prompt_tokens"], usage["completion_tokens"], usage["total_tokens"]) == ( + *FIRST_ROUND, + sum(FIRST_ROUND), + ), usage + + NO_CACHE: Final = {"no-cache": True} @@ -1269,6 +1495,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 2d08609c93b..2db8941d304 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -83,6 +83,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( team_member_budget_entity_id, ) from litellm.proxy.utils import ProxyLogging +from litellm.responses.mcp.parked_stream_dispatch import IN_FLIGHT_ROUND from litellm.router import Router from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.utils import ModelResponseListIterator @@ -9051,6 +9052,33 @@ 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() + + async def cannot_price_the_in_flight_round() -> None: + raise RuntimeError("the in-flight round could not be priced") + + setattr(response.logging_obj, IN_FLIGHT_ROUND, cannot_price_the_in_flight_round) + stream: Final = types.SimpleNamespace(chunks=[], messages=None, aclose=closed) + 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 ec5968c2a1a..c0d32742e8a 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1,17 +1,25 @@ import asyncio +import gc import json +import queue import sys +import threading import types -from collections.abc import Mapping, Sequence +import warnings +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager from datetime import datetime from types import MappingProxyType from typing import Final, cast -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import respx from httpx import Response from mcp.types import Tool as MCPTool +from starlette.responses import StreamingResponse +from starlette.types import Message import litellm from litellm.integrations.custom_logger import CustomLogger @@ -25,7 +33,8 @@ from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.utils import ModelResponse, StandardLoggingPayload +from litellm.types.utils import ModelResponse, ModelResponseStream, StandardLoggingPayload, Usage +from litellm.utils import CustomStreamWrapper @pytest.mark.asyncio @@ -1556,6 +1565,7 @@ _CHAT_ROUND_TWO: Final = MappingProxyType( ), } ) +_NEVER_LOGGED_SECONDS: Final = 60 _CONTEXT_WINDOW_ERROR: Final = MappingProxyType( { "error": MappingProxyType( @@ -1575,16 +1585,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: @@ -1600,18 +1625,20 @@ async def _settled() -> None: await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) -async def _auto_execute_chat_echo( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], call_id: str -) -> tuple[_ChatRoundRecorder, ModelResponse | Exception]: +def _serve_chat_rounds( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response] +) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") litellm.in_memory_llm_clients_cache.flush_cache() respx_mock.post("https://mcp-chat-rounds.example.com/v1/chat/completions").mock(side_effect=list(rounds)) - recorder: Final = _ChatRoundRecorder() - logging_obj: Final = Logging( + + +def _chat_rounds_logging_obj(recorder: _ChatRoundRecorder, call_id: str, stream: bool) -> Logging: + return Logging( model="openai/gpt-5.6", messages=[{"role": "user", "content": "Call echo with PING"}], - stream=False, + stream=stream, call_type="acompletion", start_time=datetime(2026, 1, 1), litellm_call_id=call_id, @@ -1619,6 +1646,10 @@ async def _auto_execute_chat_echo( dynamic_async_success_callbacks=[recorder], dynamic_async_failure_callbacks=[recorder], ) + + +@contextmanager +def _echo_tool_served() -> Iterator[None]: echo_tool: Final = MCPTool.model_validate( {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, by_name=False, @@ -1635,6 +1666,16 @@ async def _auto_execute_chat_echo( AsyncMock(return_value=[{"tool_call_id": "call_echo", "result": "PING", "name": "echo"}]), ), ): + yield + + +async def _auto_execute_chat_echo( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], call_id: str +) -> tuple[_ChatRoundRecorder, ModelResponse | Exception]: + _serve_chat_rounds(respx_mock, monkeypatch, rounds) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, call_id, stream=False) + with _echo_tool_served(): try: response: Final = await litellm.acompletion( model="openai/gpt-5.6", @@ -1712,3 +1753,994 @@ async def test_non_stream_auto_execute_bills_the_paid_first_round_when_the_follo assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71, 19) assert recorder.failed[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE)) assert recorder.failed[0]["response_cost"] > 0 + + +_FIXED_MARGIN: Final = 0.001 + + +def _chat_stream_events(body: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + completed: Final = json.loads(json.dumps(body, default=dict)) + choice: Final = completed["choices"][0] + message: Final = choice["message"] + calls: Final = message.get("tool_calls") or [] + delta: Final = ( + {**message, "tool_calls": [{**call, "index": index} for index, call in enumerate(calls)]} if calls else message + ) + head: Final = {"id": completed["id"], "object": "chat.completion.chunk", "created": 1, "model": completed["model"]} + return ( + {**head, "choices": [{"index": 0, "delta": delta, "finish_reason": None}]}, + {**head, "choices": [{"index": 0, "delta": {}, "finish_reason": choice["finish_reason"]}]}, + {**head, "choices": [], "usage": completed["usage"]}, + ) + + +def _sse(events: Sequence[Mapping[str, object]]) -> bytes: + return b"".join(f"data: {json.dumps(event)}\n\n".encode() for event in events) + + +def _streamed(body: Mapping[str, object]) -> Response: + return Response( + 200, + content=_sse(_chat_stream_events(body)) + b"data: [DONE]\n\n", + headers={"content-type": "text/event-stream"}, + ) + + +class _CutOnceTheAnswerStarts(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _sse(_chat_stream_events(_CHAT_ROUND_TWO)[:1]) + raise httpx.ReadError("connection reset mid-stream") + + +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, 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"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + api_key="sk-mcp-chat-rounds", + api_base="https://mcp-chat-rounds.example.com/v1", + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + metadata=metadata, + ) + assert isinstance(stream, CustomStreamWrapper) + return stream + + +async def _read_all(stream: CustomStreamWrapper) -> tuple[ModelResponseStream, ...]: + return tuple([chunk async for chunk in stream]) + + +async def _read_until(stream: CustomStreamWrapper, content: str) -> None: + async for chunk in stream: + if chunk.choices and chunk.choices[0].delta.content == content: + return + pytest.fail(f"the stream never sent {content!r}") + + +def _streamed_usage(chunks: Sequence[ModelResponseStream]) -> Usage: + usages: Final = tuple(usage for chunk in chunks if isinstance(usage := getattr(chunk, "usage", None), Usage)) + assert len(usages) == 1, usages + return usages[0] + + +def _assert_one_stream_row_bills_both_rounds_once( + logged: Sequence[StandardLoggingPayload], chunks: Sequence[ModelResponseStream] +) -> None: + charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN + assert len(logged) == 1 + row: Final = logged[0] + assert row["id"] == chunks[0].id + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 32 + 64 + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert row["metadata"]["usage_object"]["cost"] == pytest.approx(row["response_cost"]) + usage: Final = _streamed_usage(chunks) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (71 + 149, 19 + 5, 90 + 154) + assert usage.cost == pytest.approx(row["response_cost"]) + + +def _prompt_tokens_the_follow_up_sent(respx_mock: respx.MockRouter) -> int: + follow_up_request: Final = json.loads(respx_mock.calls[1].request.content) + return litellm.token_counter(model="gpt-5.6", messages=follow_up_request["messages"]) + + +def _assert_one_row_bills_round_one(logged: Sequence[StandardLoggingPayload]) -> None: + assert len(logged) == 1 + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == (71, 19) + assert logged[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE)) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_read_to_the_end_logs_one_row_whose_usage_cost_counts_the_fixed_margin_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + 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-read-to-the-end", stream=True) + + with _echo_tool_served(): + chunks: Final = await _read_all(await _open_chat_echo_stream(logging_obj)) + await recorder.settled(logged=1) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once_when_the_proxy_releases( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_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-rounds-deferred", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + chunks: Final = await _read_all(stream) + await _settled() + logged_before_release: Final = recorder.logged + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await recorder.settled(logged=1) + + assert logged_before_release == () + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_bills_both_rounds_on_the_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.exceptions import GuardrailRaisedException + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_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-rounds-blocked", stream=True) + blocked: Final = GuardrailRaisedException(guardrail_name="output-policy", message="blocked") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + 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 recorder.settled(failed=1) + + charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71 + 149, 19 + 5) + assert recorder.failed[0]["response_cost"] == pytest.approx(charged_once) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_rejected_by_the_proxy_during_the_follow_up_bills_round_one_with_its_own_breakdown( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from fastapi import HTTPException + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_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-rejected-mid-follow-up", stream=True) + rejected: Final = HTTPException(status_code=400, detail="rejected by a stream hook") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + 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") + assert ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, rejected) + await logging_obj.async_failure_handler(rejected, "rejected by a stream hook") + await recorder.settled(failed=1) + + assert recorder.logged == () + assert len(recorder.failed) == 1 + failed: Final = recorder.failed[0] + breakdown: Final = failed["cost_breakdown"] + assert (failed["prompt_tokens"], failed["completion_tokens"]) == (71, 19) + assert failed["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE)) + assert breakdown is not None + assert breakdown["total_cost"] == pytest.approx(failed["response_cost"]) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_answered_without_a_tool_call_logs_its_one_round_once_the_answer_finishes( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_TWO),)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-one-round", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + async for chunk in stream: + if chunk.choices and chunk.choices[0].finish_reason is not None: + break + await recorder.settled(logged=1) + logged_once_the_answer_finished: Final = recorder.logged + after_the_answer: Final = await _read_all(stream) + await _settled() + + assert len(logged_once_the_answer_finished) == 1 + assert recorder.logged == logged_once_the_answer_finished + assert recorder.logged[0]["id"] == _CHAT_ROUND_TWO["id"] + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (149, 5) + assert after_the_answer[-1:] and isinstance(after_the_answer[-1].usage, Usage) + usage: Final = _streamed_usage(after_the_answer) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + recorder.logged[0]["prompt_tokens"], + recorder.logged[0]["completion_tokens"], + recorder.logged[0]["total_tokens"], + ) + assert recorder.logged[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_TWO)) + assert recorder.failed == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("closes", (True, False), ids=("aclose", "dropped")) +async def test_stream_auto_execute_bills_the_finished_first_round_when_the_caller_stops_reading_the_follow_up( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, closes: bool +): + _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-caller-stops", stream=True) + + async def stop_reading_once_the_follow_up_answers() -> None: + stream: Final = await _open_chat_echo_stream(logging_obj) + await _read_until(stream, "PING") + if closes: + await stream.aclose() + + with _echo_tool_served(): + await stop_reading_once_the_follow_up_answers() + 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( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging, _failure_fields_to_lift + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _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", stream=True) + request_data: Final = MappingProxyType({"litellm_logging_obj": logging_obj}) + + 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 + ) + with pytest.raises(litellm.ServiceUnavailableError) as broken: + 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 recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert len(recorder.failed) == 1 + failed: Final = recorder.failed[0] + unmargined: Final = ( + litellm.completion_cost( + completion_response=ModelResponse( + model="gpt-5.6", + usage=Usage( + prompt_tokens=failed["prompt_tokens"], + completion_tokens=failed["completion_tokens"], + total_tokens=failed["total_tokens"], + ), + ), + model="gpt-5.6", + custom_llm_provider="openai", + ) + - _FIXED_MARGIN + ) + assert failed["prompt_tokens"] == _prompt_tokens_the_follow_up_sent(respx_mock) + assert failed["completion_tokens"] > 0 + assert failed["response_cost"] == pytest.approx(unmargined) + request_failure: Final = _failure_fields_to_lift(request_data) + assert cast(StandardLoggingPayload, request_failure["standard_logging_object"])["id"] != recorder.logged[0]["id"] + assert request_failure["response_cost"] == pytest.approx(unmargined) + + +@pytest.mark.asyncio +@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-proxy-stream-hook", stream=True) + + with _echo_tool_served(): + 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_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_follow_up_streams_nothing_still_bills_the_first_round_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + nothing_streamed: Final = Response(200, content=b"data: [DONE]\n\n", headers={"content-type": "text/event-stream"}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), nothing_streamed)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-empty-follow-up", stream=True) + + with _echo_tool_served(): + await _read_all(await _open_chat_echo_stream(logging_obj)) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + + +_ANSWER_WITH_UNREADABLE_AUDIO: Final = MappingProxyType( + { + **_CHAT_ROUND_TWO, + "choices": ( + MappingProxyType( + { + "index": 0, + "finish_reason": "stop", + "message": MappingProxyType( + { + "role": "assistant", + "content": "PING", + "audio": MappingProxyType({"id": "audio_echo", "data": "not base64", "transcript": "PING"}), + } + ), + } + ), + ), + } +) + + +@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(_ANSWER_WITH_UNREADABLE_AUDIO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-unassembled", stream=True) + + with _echo_tool_served(): + chunks: Final = await _read_all(await _open_chat_echo_stream(logging_obj)) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert any(chunk.choices and chunk.choices[0].delta.content == "PING" for chunk in chunks) + usage: Final = _streamed_usage(chunks) + 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]) +@pytest.mark.parametrize("answer", [_CHAT_ROUND_TWO, _ANSWER_WITH_UNREADABLE_AUDIO], ids=("assembled", "unassembled")) +async def test_stream_auto_execute_sends_its_usage_chunk_under_the_id_of_the_row_it_billed( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + deferred_for_a_post_call_guardrail: bool, + answer: Mapping[str, object], +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(answer))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-usage-chunk-id", 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_all(stream) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await recorder.settled(logged=1) + + usage_chunks: Final = tuple(chunk for chunk in chunks if isinstance(getattr(chunk, "usage", None), Usage)) + answer_chunks: Final = tuple(chunk for chunk in chunks if chunk.choices and chunk.choices[0].delta.content == "PING") + assert len(recorder.logged) == 1 + assert [chunk.id for chunk in answer_chunks] == [_CHAT_ROUND_TWO["id"]] + assert len(usage_chunks) == 1 + assert usage_chunks[0].id == chunks[0].id == recorder.logged[0]["id"] == _CHAT_ROUND_ONE["id"] + + +def _assert_one_row_bills_round_one_and_the_partial_answer( + recorder: _ChatRoundRecorder, respx_mock: respx.MockRouter +) -> None: + 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 == _prompt_tokens_the_follow_up_sent(respx_mock) + assert 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 == () + + +@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_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +def _chat_rounds_router() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "mcp-chat-rounds", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-mcp-chat-rounds", + "api_base": "https://mcp-chat-rounds.example.com/v1", + }, + } + ] + ) + + +class _StallsOnceTheAnswerStarts(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _sse(_chat_stream_events(_CHAT_ROUND_TWO)[:1]) + await asyncio.Event().wait() + + +def _streams_content(body: object, content: str) -> bool: + lines: Final = body.decode().splitlines() if isinstance(body, bytes) else () + events: Final = tuple(json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: {")) + return any(event["choices"] and event["choices"][0]["delta"].get("content") == content for event in events) + + +async def _serve_until_the_client_hangs_up_after(response: StreamingResponse, content: str) -> None: + read: Final = asyncio.Event() + + async def receive() -> Message: + await read.wait() + return {"type": "http.disconnect"} + + async def send(message: Message) -> None: + if message["type"] == "http.response.body" and _streams_content(message.get("body"), content): + read.set() + + uvicorn_scope: Final = {"type": "http", "asgi": {"version": "3.0", "spec_version": "2.3"}, "headers": []} + await response(uvicorn_scope, receive, send) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_routed_through_the_proxy_bills_the_partial_answer_when_the_client_hangs_up( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, create_response + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds( + respx_mock, + monkeypatch, + ( + _streamed(_CHAT_ROUND_ONE), + Response(200, stream=_StallsOnceTheAnswerStarts(), headers={"content-type": "text/event-stream"}), + ), + ) + router: Final = _chat_rounds_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-routed-hang-up", stream=True) + + with _echo_tool_served(): + stream: Final = await router.acompletion( + model="mcp-chat-rounds", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=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 + ) + served: Final = await create_response( + proxy_server.async_data_generator( + response=stream, + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj, "model": "mcp-chat-rounds", "metadata": {}}, + ), + "text/event-stream", + {}, + ) + assert isinstance(served, StreamingResponse) + await _serve_until_the_client_hangs_up_after(served, "PING") + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +async def _stop_reading_once_the_proxy_sent(served: AsyncIterator[object], content: str) -> None: + async for frame in served: + if _streams_content(frame, content): + return + pytest.fail(f"the proxy never sent {content!r}") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_routed_through_the_proxy_bills_the_partial_answer_when_hung_up_between_reads( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + 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))) + router: Final = _chat_rounds_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-routed-hang-up-between-reads", stream=True) + + with _echo_tool_served(): + stream: Final = await router.acompletion( + model="mcp-chat-rounds", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=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 + ) + served: Final = proxy_server.async_data_generator( + response=stream, + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj, "model": "mcp-chat-rounds", "metadata": {}}, + ) + await _stop_reading_once_the_proxy_sent(served, "PING") + await served.aclose() + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +_LONGEST_WAIT_FOR_THE_LOOP_TO_MOVE_SECONDS: Final = 10 + + +class _CountsTheToolResultWhileTheLoopRuns: + def __init__(self, loop: asyncio.AbstractEventLoop, count: Callable[..., int]) -> None: + self._loop: Final = loop + self._count: Final = count + self._counting: Final = asyncio.Event() + self._loop_moved: Final = threading.Event() + self.loop_moved_while_counting: tuple[bool, ...] = () + + def __call__(self, model: str | None = None, messages: list[object] | None = None, **kwargs: object) -> int: + if any(isinstance(message, Mapping) and message.get("role") == "tool" for message in messages or ()): + self._loop.call_soon_threadsafe(self._counting.set) + moved: Final = self._loop_moved.wait(timeout=_LONGEST_WAIT_FOR_THE_LOOP_TO_MOVE_SECONDS) + self.loop_moved_while_counting = (*self.loop_moved_while_counting, moved) + return self._count(model=model, messages=messages, **kwargs) + + async def keep_the_loop_moving(self) -> None: + await self._counting.wait() + self._loop_moved.set() + + +@contextmanager +def _partial_count_watched(monkeypatch: pytest.MonkeyPatch) -> Iterator[_CountsTheToolResultWhileTheLoopRuns]: + from litellm.litellm_core_utils import streaming_chunk_builder_utils + + counter: Final = _CountsTheToolResultWhileTheLoopRuns( + asyncio.get_running_loop(), streaming_chunk_builder_utils.token_counter + ) + monkeypatch.setattr(streaming_chunk_builder_utils, "token_counter", counter) + other_task: Final = asyncio.create_task(counter.keep_the_loop_moving()) + try: + yield counter + finally: + other_task.cancel() + + +class _StallsAfterTheAnswerStarts(httpx.AsyncByteStream): + def __init__(self) -> None: + self.stalled: Final = asyncio.Event() + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _sse(_chat_stream_events(_CHAT_ROUND_TWO)[:1]) + self.stalled.set() + await asyncio.Event().wait() + + +@pytest.mark.asyncio +async def test_stream_auto_execute_cancelled_mid_answer_counts_the_partial_off_the_loop_and_bills_it_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + answer: Final = _StallsAfterTheAnswerStarts() + _serve_chat_rounds( + respx_mock, + monkeypatch, + (_streamed(_CHAT_ROUND_ONE), Response(200, stream=answer, headers={"content-type": "text/event-stream"})), + ) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-cancelled-off-the-loop", stream=True) + + with _echo_tool_served(), _partial_count_watched(monkeypatch) as counter: + stream: Final = await _open_chat_echo_stream(logging_obj) + reading: Final = asyncio.create_task(_read_all(stream)) + await asyncio.wait_for(answer.stalled.wait(), timeout=10) + reading.cancel() + with pytest.raises(asyncio.CancelledError): + await reading + await stream.aclose() + await recorder.settled(logged=1) + + assert counter.loop_moved_while_counting == (True,) + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_routed_through_the_proxy_and_closed_between_reads_counts_the_partial_off_the_loop( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + 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))) + router: Final = _chat_rounds_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-routed-closed-off-the-loop", stream=True) + + with _echo_tool_served(), _partial_count_watched(monkeypatch) as counter: + stream: Final = await router.acompletion( + model="mcp-chat-rounds", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + ) + served: Final = proxy_server.async_data_generator( + response=stream, + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj, "model": "mcp-chat-rounds", "metadata": {}}, + ) + await _stop_reading_once_the_proxy_sent(served, "PING") + await served.aclose() + await recorder.settled(logged=1) + + assert counter.loop_moved_while_counting == (True,) + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +class _StallsBeforeTheAnswerStarts(httpx.AsyncByteStream): + def __init__(self) -> None: + self.opened: Final = asyncio.Event() + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b": the answer has not started\n\n" + self.opened.set() + await asyncio.Event().wait() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_before_the_answer_starts_bills_round_one_alone_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 + + answer: Final = _StallsBeforeTheAnswerStarts() + _serve_chat_rounds( + respx_mock, + monkeypatch, + (_streamed(_CHAT_ROUND_ONE), Response(200, stream=answer, headers={"content-type": "text/event-stream"})), + ) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-before-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 + ) + reading: Final = asyncio.create_task(_read_all(stream)) + await asyncio.wait_for(answer.opened.wait(), timeout=10) + reading.cancel() + with pytest.raises(asyncio.CancelledError): + await reading + 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_row_bills_round_one(recorder.logged) + 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 = 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") + + +def _names_chat_stream_billing(text: str) -> bool: + return any(name in text for name in _CHAT_STREAM_BILLING_NAMES) + + +class _UnraisableRecorder: + def __init__(self) -> None: + self.raised: tuple[str, ...] = () + + def __call__(self, unraisable: "sys.UnraisableHookArgs") -> None: + self.raised = (*self.raised, f"{unraisable.object!r}: {unraisable.exc_value!r}") + + +def test_stream_auto_execute_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() + 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", stream=True), loop + ) + 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 [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))] == [] diff --git a/tests/unit/responses/mcp/test_parked_stream_dispatch.py b/tests/unit/responses/mcp/test_parked_stream_dispatch.py new file mode 100644 index 00000000000..dd380c582c8 --- /dev/null +++ b/tests/unit/responses/mcp/test_parked_stream_dispatch.py @@ -0,0 +1,17 @@ +from typing import Final + +import pytest + +from litellm.responses.mcp.parked_stream_dispatch import parked_success +from litellm.types.utils import ModelResponse + +_ASSEMBLED: Final = ModelResponse(id="chatcmpl-parked", model="gpt-5.6") + + +@pytest.mark.parametrize( + "parked", + [(), (_ASSEMBLED,), (_ASSEMBLED, False, "not a logger"), (_ASSEMBLED, False, None, None)], + ids=("empty", "assembled-alone", "third-part-is-not-a-logger", "four-parts"), +) +def test_a_parked_dispatch_of_an_unknown_shape_is_never_dispatched(parked: tuple[object, ...]): + assert parked_success(parked) is None