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