From e3fbb6583daa4c1ac313cb943be1842b81c012e3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 18:50:44 -0700 Subject: [PATCH] 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 --- litellm/proxy/common_request_processing.py | 6 +- .../responses/mcp/chat_completions_handler.py | 4 +- .../mcp/chat_stream_round_billing.py | 20 +- .../responses/mcp/parked_stream_dispatch.py | 10 +- .../proxy/test_common_request_processing.py | 2 +- .../mcp/test_chat_completions_handler.py | 193 +++++++++++++++--- .../mcp/test_parked_stream_dispatch.py | 17 ++ 7 files changed, 214 insertions(+), 38 deletions(-) create mode 100644 tests/unit/responses/mcp/test_parked_stream_dispatch.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7f7db2a010e..b08ec55340f 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -528,12 +528,12 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return await _dispatch_disconnect_billing_event(logging_obj, partial_response) -def _bill_in_flight_round_on_disconnect(request_data: Mapping[str, object]) -> None: +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: - bill_in_flight_round() + 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) @@ -4442,7 +4442,7 @@ class ProxyBaseLLMRequestProcessing: client_disconnected, ) if recorded_client_disconnect: - _bill_in_flight_round_on_disconnect(_REQUEST_FIELDS.validate_python(request_data)) + 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/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index bfa2bc8a479..1e07b768f62 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -5,6 +5,7 @@ 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 @@ -373,7 +374,8 @@ async def acompletion_with_mcp( raise return usage_chunk except asyncio.CancelledError: - round_billing.bill_in_flight_round() + with anyio.CancelScope(shield=True): + await round_billing.bill_in_flight_round() raise except Exception: round_billing.fail() diff --git a/litellm/responses/mcp/chat_stream_round_billing.py b/litellm/responses/mcp/chat_stream_round_billing.py index ac70377fa6e..7623d42e295 100644 --- a/litellm/responses/mcp/chat_stream_round_billing.py +++ b/litellm/responses/mcp/chat_stream_round_billing.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, replace from types import MappingProxyType from typing import Final, TypeAlias -from litellm.litellm_core_utils.asyncify import is_event_loop_running +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 @@ -69,7 +69,9 @@ class _FirstRound: update=MappingProxyType({"cost": _cost(billed)}) ) ) - return self.usage_chunk.chunk.model_copy(update=MappingProxyType({"usage": usage})) + return self.usage_chunk.chunk.model_copy( + update=MappingProxyType({"id": self.priced.response.id, "usage": usage}) + ) def _cost(response: ModelResponse) -> object: @@ -143,13 +145,13 @@ class ChatStreamRoundBilling: 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 + async def bill_in_flight_round(self) -> None: 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): + if partial_follow_up is None or self._first_round_held_alone() is None: return - final_round: Final = _assembled_round((partial_follow_up(), False)) - if final_round is None: + 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( @@ -231,6 +233,10 @@ class ChatStreamRoundBilling: 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 _holds_round_one_alone(self, first: _FirstRound) -> bool: parked: Final = self._peek_parked() return parked is first.parked or parked is first.logged_alone diff --git a/litellm/responses/mcp/parked_stream_dispatch.py b/litellm/responses/mcp/parked_stream_dispatch.py index 093587400f5..85bbc2d902a 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 Callable, Coroutine +from collections.abc import Coroutine from dataclasses import dataclass from typing import Final, Protocol @@ -26,6 +26,10 @@ class DeferredStreamRelease(Protocol): def __call__(self, *parked: object) -> Coroutine[object, object, None]: ... +class InFlightRoundBiller(Protocol): + def __call__(self) -> Coroutine[object, object, None]: ... + + def success_dispatcher(logger: LiteLLMLoggingObj) -> DispatchesSuccessHandlers: return logger @@ -43,8 +47,8 @@ class ParkedSuccess: ) -def in_flight_round_biller(logging_obj: object) -> Callable[[], object] | None: - biller: Final[object] = getattr(logging_obj, IN_FLIGHT_ROUND, None) +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 diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index a5d6a33877f..28b2a0fef0b 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -9064,7 +9064,7 @@ class TestStreamingClientDisconnectBilling: response = await self._start_partial_stream() closed: Final = AsyncMock() - def cannot_price_the_in_flight_round() -> None: + 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) diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 8af56157cb8..10ed42063a0 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -3,6 +3,7 @@ import gc import json import queue import sys +import threading import types import warnings from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence @@ -2139,33 +2140,33 @@ async def test_stream_auto_execute_whose_follow_up_streams_nothing_still_bills_t _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 ): - 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"} - ), - } - ), - } - ), - ), - } - ) - _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(answer_with_unreadable_audio))) + _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) @@ -2180,6 +2181,40 @@ async def test_stream_auto_execute_whose_follow_up_cannot_be_assembled_ends_with 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: @@ -2374,6 +2409,118 @@ 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) +_PARTIAL_COUNT_BLOCKS_SECONDS: Final = 0.3 + + +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=_PARTIAL_COUNT_BLOCKS_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() 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