From f23da840b59c623960c45fa269afd42e89173fe3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 15:14:32 -0700 Subject: [PATCH] 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 --- litellm/proxy/common_request_processing.py | 16 +-- .../responses/mcp/chat_completions_handler.py | 19 +-- .../mcp/chat_stream_round_billing.py | 18 ++- .../responses/mcp/parked_stream_dispatch.py | 8 +- .../proxy/test_common_request_processing.py | 6 +- .../mcp/test_chat_completions_handler.py | 122 ++++++++++++++++-- 6 files changed, 143 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 73627c51466..7f7db2a010e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -137,7 +137,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, ) -from litellm.responses.mcp.parked_stream_dispatch import parked_success +from litellm.responses.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_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX @@ -528,16 +528,12 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return await _dispatch_disconnect_billing_event(logging_obj, partial_response) -@runtime_checkable -class _BillsInFlightRoundOnDisconnect(Protocol): - def bill_in_flight_round_on_disconnect(self) -> None: ... - - -def _bill_in_flight_round_on_disconnect(response: object) -> None: - if not isinstance(response, _BillsInFlightRoundOnDisconnect): +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: - response.bill_in_flight_round_on_disconnect() + 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) @@ -4446,7 +4442,7 @@ class ProxyBaseLLMRequestProcessing: client_disconnected, ) if recorded_client_disconnect: - _bill_in_flight_round_on_disconnect(response) + _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/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 2e05639e000..bfa2bc8a479 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -274,18 +274,6 @@ async def acompletion_with_mcp( if isinstance(round_billing, ChatStreamRoundBilling): round_billing.release_abandoned() - def bill_in_flight_round_on_disconnect(self) -> None: - follow_up: Final[object] = self.follow_up_stream - if ( - self._round_billing is None - or self.follow_up_exhausted - or not isinstance(follow_up, CustomStreamWrapper) - ): - return - partial: Final = follow_up.partial_response() - if isinstance(partial, ModelResponse): - self._round_billing.bill_partial_follow_up(partial) - def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_list_tools to the first chunk.""" from litellm.types.utils import ( @@ -385,7 +373,7 @@ async def acompletion_with_mcp( raise return usage_chunk except asyncio.CancelledError: - self.bill_in_flight_round_on_disconnect() + round_billing.bill_in_flight_round() raise except Exception: round_billing.fail() @@ -555,6 +543,8 @@ async def acompletion_with_mcp( 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 @@ -615,9 +605,6 @@ async def acompletion_with_mcp( self._custom_iterator.release_abandoned() await super().aclose() - def bill_in_flight_round_on_disconnect(self) -> None: - self._custom_iterator.bill_in_flight_round_on_disconnect() - def __iter__(self): # For synchronous iteration, create a sync wrapper if self._sync_iterator is None: diff --git a/litellm/responses/mcp/chat_stream_round_billing.py b/litellm/responses/mcp/chat_stream_round_billing.py index ad0ae8fa6cf..ac70377fa6e 100644 --- a/litellm/responses/mcp/chat_stream_round_billing.py +++ b/litellm/responses/mcp/chat_stream_round_billing.py @@ -1,5 +1,5 @@ import asyncio -from collections.abc import Generator +from collections.abc import Callable, Generator from contextlib import contextmanager from dataclasses import dataclass, replace from types import MappingProxyType @@ -11,6 +11,7 @@ from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, Uncach from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.responses.mcp.parked_stream_dispatch import ( DISPATCH_RELEASE, + IN_FLIGHT_ROUND, PARKED_DISPATCH, SERVED_FROM_CACHE, DeferredStreamRelease, @@ -21,6 +22,7 @@ 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) @@ -91,6 +93,7 @@ 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 @@ -136,10 +139,17 @@ class ChatStreamRoundBilling: setattr(self._logging_obj, PARKED_DISPATCH, None) self._log_through_the_release((*first.parked, first.banked.guardrail_scoped_copy())) - def bill_partial_follow_up(self, partial: ModelResponse) -> None: + 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) + + def bill_in_flight_round(self) -> None: first: Final = self._first - final_round: Final = _assembled_round((partial, False)) - if first is None or final_round is None or not self._holds_round_one_alone(first): + partial_follow_up: Final = self._partial_follow_up + if first is None or partial_follow_up is None or not self._holds_round_one_alone(first): + return + final_round: Final = _assembled_round((partial_follow_up(), False)) + if final_round is None: return self._first = None setattr( diff --git a/litellm/responses/mcp/parked_stream_dispatch.py b/litellm/responses/mcp/parked_stream_dispatch.py index 99adbbbf886..093587400f5 100644 --- a/litellm/responses/mcp/parked_stream_dispatch.py +++ b/litellm/responses/mcp/parked_stream_dispatch.py @@ -1,4 +1,4 @@ -from collections.abc import Coroutine +from collections.abc import Callable, Coroutine from dataclasses import dataclass from typing import Final, Protocol @@ -7,6 +7,7 @@ from litellm.litellm_core_utils.litellm_logging import UncachedPrice PARKED_DISPATCH: Final = "_deferred_stream_complete_args" DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" +IN_FLIGHT_ROUND: Final = "_bill_in_flight_round_on_disconnect" SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) @@ -42,6 +43,11 @@ class ParkedSuccess: ) +def in_flight_round_biller(logging_obj: object) -> Callable[[], object] | None: + biller: Final[object] = getattr(logging_obj, IN_FLIGHT_ROUND, None) + return biller if isinstance(logging_obj, LiteLLMLoggingObj) and callable(biller) else None + + def parked_success(parked: tuple[object, ...]) -> ParkedSuccess | None: match parked: case (assembled, cache_hit): diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 4731ea89284..a5d6a33877f 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -77,6 +77,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 @@ -9066,9 +9067,8 @@ class TestStreamingClientDisconnectBilling: def cannot_price_the_in_flight_round() -> None: raise RuntimeError("the in-flight round could not be priced") - stream: Final = types.SimpleNamespace( - chunks=[], messages=None, aclose=closed, bill_in_flight_round_on_disconnect=cannot_price_the_in_flight_round - ) + 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( diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 1daa4718f83..8af56157cb8 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -2232,6 +2232,21 @@ async def test_stream_auto_execute_hung_up_during_the_follow_up_bills_round_one_ _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]) @@ -2277,18 +2292,7 @@ async def test_stream_auto_execute_routed_through_the_proxy_bills_the_partial_an Response(200, stream=_StallsOnceTheAnswerStarts(), headers={"content-type": "text/event-stream"}), ), ) - router: Final = 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", - }, - } - ] - ) + 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) @@ -2322,6 +2326,100 @@ async def test_stream_auto_execute_routed_through_the_proxy_bills_the_partial_an _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) + + +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, ...]: