From b6ee181ec9bd3ad7837a31b473788d9cafbe6b6a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 06:01:06 -0700 Subject: [PATCH 01/10] 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 --- .../responses/mcp/chat_completions_handler.py | 61 ++- .../mcp/chat_stream_round_billing.py | 210 ++++++++++ .../integration/mcp/test_mcp_llm_endpoints.py | 51 +++ .../mcp/test_chat_completions_handler.py | 384 +++++++++++++++++- 4 files changed, 684 insertions(+), 22 deletions(-) create mode 100644 litellm/responses/mcp/chat_stream_round_billing.py diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index cd072825ef5..c917144a7c1 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -10,6 +10,7 @@ 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, ) @@ -258,10 +259,19 @@ 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) 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 ( @@ -332,7 +342,32 @@ async def acompletion_with_mcp( "Error draining inner MCP stream after final chunk; spend logging may be incomplete" ) - async def __anext__(self): + async def _complete_first_round(self) -> None: + await self._drain_inner_stream() + if self._round_billing is not None: + self._round_billing.complete_first_round() + + async def _follow_up_or_release_first_round(self) -> None: + if self.tool_results and self.complete_response: + await self._prepare_follow_up_call() + if self.follow_up_stream is None and 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: + round_billing.finish() + 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 +397,22 @@ async def acompletion_with_mcp( if is_final: self.stream_exhausted = True + 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() return chunk except StopAsyncIteration: self.stream_exhausted = True + 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() return final_chunk # Phase 2: Yield follow-up stream chunks if available @@ -400,7 +428,7 @@ async def acompletion_with_mcp( from litellm._logging import verbose_logger verbose_logger.debug("Follow-up chunk yielded: %s", chunk) - return chunk + return chunk if self._round_billing is None else self._round_billing.client_chunk(chunk) except StopAsyncIteration: self.follow_up_exhausted = True from litellm._logging import verbose_logger @@ -488,10 +516,13 @@ 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) self.follow_up_stream = follow_up_response from litellm._logging import verbose_logger @@ -548,6 +579,10 @@ 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: 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..7274652890e --- /dev/null +++ b/litellm/responses/mcp/chat_stream_round_billing.py @@ -0,0 +1,210 @@ +import asyncio +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass, replace +from types import MappingProxyType +from typing import Final, Protocol, TypeAlias + +from litellm.litellm_core_utils.asyncify import is_event_loop_running +from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.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, ...] + +_PARKED_DISPATCH: Final = "_deferred_stream_complete_args" +_DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" +_SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) + + +class _DispatchesSuccessHandlers(Protocol): + async def dispatch_success_handlers( + self, + result: object = None, + start_time: object = None, + end_time: object = None, + cache_hit: object = None, + prefer_async_handlers: bool = False, + ) -> None: ... + + +def _success_dispatcher(logger: LiteLLMLoggingObj) -> _DispatchesSuccessHandlers: + return logger + + +@dataclass(frozen=True, slots=True) +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 _FirstRound: + priced: _PricedRound + parked: Parked + logged_alone: Parked + banked: LiteLLMLoggingObj | None = None + + +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 + + @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) + 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) + + 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: + usage: Final[object] = getattr(chunk, "usage", None) + if self._first is None or not isinstance(usage, Usage): + return chunk + return chunk.model_copy(update=MappingProxyType({"usage": merged_chat_usage(self._first.priced.usage, usage)})) + + def finish(self) -> None: + parked: Final = self._take_parked() + first: Final = self._first + self._first = None + if parked is None: + self._disarm() + return + if first is None or parked is first.logged_alone: + self._release(parked) + return + final_round: Final = _assembled_round(parked) + if final_round is None: + self._release(first.logged_alone) + return + self._release(self._billed_for_both_rounds(first.priced, self._priced(final_round))) + + def fail(self) -> None: + self._bank_apart_from_the_follow_up() + first: Final = self._first + if first is None or first.banked is None or self._peek_parked() is not first.parked: + self.finish() + return + self._first = None + setattr(self._logging_obj, _PARKED_DISPATCH, None) + self._log_apart((*first.parked, first.banked)) + + def release_abandoned(self) -> None: + if not self._owns_release(): + return + if is_event_loop_running(): + 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 _billed_for_both_rounds(self, first: _PricedRound, final: _PricedRound) -> Parked: + first_breakdown: Final = first.price.cost_breakdown + final_breakdown: Final = final.price.cost_breakdown + billed: Final = billed_for_every_round(first.response, first_breakdown, final.response, final_breakdown) + cache_hit: Final = first.cache_hit and final.cache_hit + set_hidden_params(billed, {**(get_hidden_params(billed) 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 _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 + setattr(self._logging_obj, _DISPATCH_RELEASE, None) + self._log_apart(parked) + + def _log_apart(self, parked: Parked) -> None: + loop: Final = asyncio.get_running_loop() + spawn_logging_task(loop, self._log_success(*parked), task_name="Chat completions MCP stream success log") + + def _disarm(self) -> None: + if self._owns_release(): + setattr(self._logging_obj, _DISPATCH_RELEASE, None) + + async def _log_success(self, assembled: object, cache_hit: object, banked: object = None) -> None: + logger: Final = banked if isinstance(banked, LiteLLMLoggingObj) else self._logging_obj + await _success_dispatcher(logger).dispatch_success_handlers( + assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True + ) diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 490a8f831bc..2925154b9ad 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -950,6 +950,57 @@ 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[-1]["id"], (row, chunks[-1]) + 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 _tool_rows(call_id: str, expected: int) -> list[dict[str, JsonValue]]: return eventually( lambda: read_rows( diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index ec5968c2a1a..2833ea0c962 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1,13 +1,18 @@ import asyncio +import gc import json +import queue import sys 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 @@ -25,7 +30,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 @@ -1600,18 +1606,28 @@ 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]: +async def _settled_once(logged: Callable[[], bool]) -> None: + for _ in range(200): + if logged(): + break + await asyncio.sleep(0.01) + await _settled() + + +def _serve_chat_rounds( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response] +) -> None: 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 +1635,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 +1655,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 +1742,339 @@ 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) -> 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, + ) + 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[-1].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) + usage: Final = _streamed_usage(chunks) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (71 + 149, 19 + 5, 90 + 154) + + +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_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 _settled_once(lambda: len(recorder.logged) > 0) + + 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 _settled_once(lambda: len(recorder.failed) > 0) + + 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_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 _settled_once(lambda: len(recorder.logged) > 0) + logged_once_the_answer_finished: Final = recorder.logged + 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 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 _settled_once(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_round_one(recorder.logged) + + +@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 _settled_once(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + _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["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 +async def test_stream_auto_execute_leaves_no_deferred_log_armed_for_the_proxy_to_orphan_after_its_first_chunk( + 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-first-chunk-read-ahead", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + rest: Final = stream.__aiter__() + first_chunk: Final = await rest.__anext__() + armed_after_the_first_chunk: Final = getattr(logging_obj, "_on_deferred_stream_complete", None) + chunks: Final = (first_chunk, *[chunk async for chunk in rest]) + await _settled_once(lambda: len(recorder.logged) > 0) + + assert armed_after_the_first_chunk is None + _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 _settled_once(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_round_one(recorder.logged) + + +async def _stream_read_until_the_follow_up_answers( + logging_obj: Logging, opened: "queue.SimpleQueue[CustomStreamWrapper]" +) -> None: + stream: Final = await _open_chat_echo_stream(logging_obj) + await _read_until(stream, "PING") + opened.put(stream) + + +_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_dropped_after_its_event_loop_closed_discards_its_parked_log_cleanly( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + unraisable: Final = _UnraisableRecorder() + opened: Final = queue.SimpleQueue[CustomStreamWrapper]() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + + with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + asyncio.run( + _stream_read_until_the_follow_up_answers( + _chat_rounds_logging_obj(recorder, "mcp-chat-stream-loop-closed", stream=True), opened + ) + ) + opened.get_nowait() + gc.collect() + + assert recorder.logged == () + assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == [] From ccf2aa78d0020aabdc4b2b2c633611298e9766a5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 07:13:07 -0700 Subject: [PATCH 02/10] 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 --- .../responses/mcp/chat_completions_handler.py | 40 ++++++++++++------- .../integration/mcp/test_mcp_llm_endpoints.py | 40 +++++++++++++++++++ .../mcp/test_chat_completions_handler.py | 9 ++++- 3 files changed, 73 insertions(+), 16 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index c917144a7c1..f35089f2e1f 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -20,7 +20,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: @@ -260,6 +260,7 @@ async def acompletion_with_mcp( 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 @@ -331,26 +332,31 @@ 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 _complete_first_round(self) -> None: - await self._drain_inner_stream() + 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) -> 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 None and self._round_billing is not None: + 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: @@ -397,22 +403,22 @@ async def acompletion_with_mcp( if is_final: self.stream_exhausted = True - await self._complete_first_round() + usage_chunk: Final = await self._complete_first_round() await self._process_tool_calls() chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk) - await self._follow_up_or_release_first_round() + await self._follow_up_or_release_first_round(usage_chunk) return chunk except StopAsyncIteration: self.stream_exhausted = True - await self._complete_first_round() + 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) - await self._follow_up_or_release_first_round() + await self._follow_up_or_release_first_round(drained_usage_chunk) return final_chunk # Phase 2: Yield follow-up stream chunks if available @@ -448,7 +454,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.""" diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 2925154b9ad..490dd3cfa0a 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -1001,6 +1001,46 @@ def test_streaming_chat_auto_execute_whose_follow_up_fails_bills_round_one_once( assert [float(str(row["spend"])) for row in failed] == [0.0], failed +@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 + + def _tool_rows(call_id: str, expected: int) -> list[dict[str, JsonValue]]: return eventually( lambda: read_rows( diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 2833ea0c962..ce681ef7e6f 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1911,13 +1911,20 @@ async def test_stream_auto_execute_answered_without_a_tool_call_logs_its_one_rou break await _settled_once(lambda: len(recorder.logged) > 0) logged_once_the_answer_finished: Final = recorder.logged - await _read_all(stream) + 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 == () From 0ecbfa376b658e06403177fac8b7d36db04abb38 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:12:00 -0700 Subject: [PATCH 03/10] 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 --- litellm/proxy/common_request_processing.py | 6 +- .../responses/mcp/chat_completions_handler.py | 11 ++- .../mcp/chat_stream_round_billing.py | 82 ++++++++++++------ .../integration/_support/hook_count_probe.py | 18 ++++ .../integration/mcp/test_mcp_llm_endpoints.py | 85 +++++++++++++++++-- .../mcp/test_chat_completions_handler.py | 3 +- 6 files changed, 166 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f5d6417fb79..e011f01c162 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4053,11 +4053,13 @@ 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_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, ) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index f35089f2e1f..90438af1487 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -367,8 +367,10 @@ async def acompletion_with_mcp( try: return await self._next_chunk() except StopAsyncIteration: - round_billing.finish() - raise + usage_chunk: Final = round_billing.finish() + if usage_chunk is None: + raise + return usage_chunk except Exception: round_billing.fail() raise @@ -434,7 +436,6 @@ async def acompletion_with_mcp( from litellm._logging import verbose_logger verbose_logger.debug("Follow-up chunk yielded: %s", chunk) - return chunk if self._round_billing is None else self._round_billing.client_chunk(chunk) except StopAsyncIteration: self.follow_up_exhausted = True from litellm._logging import verbose_logger @@ -442,6 +443,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 ( diff --git a/litellm/responses/mcp/chat_stream_round_billing.py b/litellm/responses/mcp/chat_stream_round_billing.py index 7274652890e..4df79ca13ba 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, Coroutine, Generator from contextlib import contextmanager from dataclasses import dataclass, replace from types import MappingProxyType @@ -14,6 +14,7 @@ from litellm.responses.streaming_iterator import spawn_logging_task from litellm.types.utils import ModelResponse, ModelResponseStream, Usage Parked: TypeAlias = tuple[object, ...] +_Release: TypeAlias = Callable[..., Coroutine[object, object, None]] _PARKED_DISPATCH: Final = "_deferred_stream_complete_args" _DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" @@ -55,12 +56,31 @@ class _PricedRound: 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 + merged: Final = merged_chat_usage(self.priced.usage, self.usage_chunk.usage) + usage: Final = merged if billed is None else merged.model_copy(update=MappingProxyType({"cost": _cost(billed)})) + return self.usage_chunk.chunk.model_copy(update=MappingProxyType({"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: @@ -100,27 +120,17 @@ class ChatStreamRoundBilling: 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: + 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 - return chunk.model_copy(update=MappingProxyType({"usage": merged_chat_usage(self._first.priced.usage, usage)})) + self._first = replace(self._first, usage_chunk=_HeldUsageChunk(chunk, usage)) + return None - def finish(self) -> None: - parked: Final = self._take_parked() + def finish(self) -> ModelResponseStream | None: first: Final = self._first - self._first = None - if parked is None: - self._disarm() - return - if first is None or parked is first.logged_alone: - self._release(parked) - return - final_round: Final = _assembled_round(parked) - if final_round is None: - self._release(first.logged_alone) - return - self._release(self._billed_for_both_rounds(first.priced, self._priced(final_round))) + 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() @@ -130,7 +140,7 @@ class ChatStreamRoundBilling: return self._first = None setattr(self._logging_obj, _PARKED_DISPATCH, None) - self._log_apart((*first.parked, first.banked)) + self._log_through_the_release((*first.parked, first.banked)) def release_abandoned(self) -> None: if not self._owns_release(): @@ -154,12 +164,31 @@ class ChatStreamRoundBilling: self._first = replace(first, logged_alone=logged_alone, banked=banked) setattr(self._logging_obj, _PARKED_DISPATCH, logged_alone) - def _billed_for_both_rounds(self, first: _PricedRound, final: _PricedRound) -> Parked: + 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 - billed: Final = billed_for_every_round(first.response, first_breakdown, final.response, final_breakdown) + rounds: Final = billed_for_every_round(first.response, first_breakdown, final.response, final_breakdown) + billed: Final = rounds.model_copy(update=MappingProxyType({"id": first.response.id})) cache_hit: Final = first.cache_hit and final.cache_hit - set_hidden_params(billed, {**(get_hidden_params(billed) or {}), "cache_hit": 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 ) @@ -192,12 +221,13 @@ class ChatStreamRoundBilling: 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) - self._log_apart(parked) + self._log_through_the_release(parked) - def _log_apart(self, parked: Parked) -> None: - loop: Final = asyncio.get_running_loop() - spawn_logging_task(loop, self._log_success(*parked), task_name="Chat completions MCP stream success log") + def _log_through_the_release(self, parked: Parked) -> None: + armed: Final[_Release | None] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + setattr(self._logging_obj, _DISPATCH_RELEASE, None) + logged: Final = (self._log_success if armed is None else armed)(*parked) + spawn_logging_task(asyncio.get_running_loop(), logged, task_name="Chat completions MCP stream success log") def _disarm(self) -> None: if self._owns_release(): 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 d0fdc6935c7..c2f8845f0d5 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -45,9 +45,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" config.write_text( @@ -59,6 +60,7 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat **object_value(stock["litellm_settings"]), "cost_margin_config": {"hosted_vllm": {"fixed_amount": FIXED_MARGIN}}, }, + **({"guardrails": list(guardrails)} if guardrails else {}), } ) ) @@ -70,15 +72,50 @@ 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" + + +@pytest.fixture(scope="module") +def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: + guardrail: Final[Mapping[str, JsonValue]] = { + "guardrail_name": ROUND_USAGE_GUARDRAIL, + "litellm_params": { + "guardrail": "integration._support.hook_count_probe.UsageRecordingPostCallGuardrail", + "mode": "post_call", + "default_on": True, + }, + } + with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), (guardrail,)) 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"] SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge") ADD: Final = {"a": 2, "b": 3} @@ -973,7 +1010,7 @@ def test_streaming_chat_auto_execute_logs_the_answer_and_bills_both_rounds(gatew 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[-1]["id"], (row, chunks[-1]) + 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], @@ -1001,6 +1038,40 @@ def test_streaming_chat_auto_execute_whose_follow_up_fails_bills_round_one_once( 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]} + ] + + @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 diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index ce681ef7e6f..23ff0727098 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1823,13 +1823,14 @@ def _assert_one_stream_row_bills_both_rounds_once( 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[-1].id + 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) 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 _assert_one_row_bills_round_one(logged: Sequence[StandardLoggingPayload]) -> None: From c32692bd9b95361406f4dbcb4c607b8f71287b26 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:55:40 -0700 Subject: [PATCH 04/10] 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 --- litellm/litellm_core_utils/litellm_logging.py | 22 ++ .../litellm_core_utils/streaming_handler.py | 35 +- litellm/proxy/common_request_processing.py | 43 ++- .../responses/mcp/chat_completions_handler.py | 15 + .../mcp/chat_stream_round_billing.py | 107 +++--- .../responses/mcp/parked_stream_dispatch.py | 52 +++ litellm/responses/mcp/stream_round_billing.py | 58 ++-- .../integration/mcp/test_mcp_llm_endpoints.py | 89 ++++- .../proxy/test_common_request_processing.py | 28 ++ .../mcp/test_chat_completions_handler.py | 320 +++++++++++++++--- 10 files changed, 624 insertions(+), 145 deletions(-) create mode 100644 litellm/responses/mcp/parked_stream_dispatch.py diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8b66c27b930..82d2aaac81f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -564,6 +564,11 @@ def _resolve_mantle_region_for_cost( _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"}) _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) _REDACTED_RESULT: Final = TypeAdapter(object) +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" + + +def _without_guardrail_information(metadata: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in metadata.items() if key != _GUARDRAIL_INFORMATION_KEY} def _provider_response_id(source: object) -> str | None: @@ -2494,6 +2499,23 @@ class Logging(LiteLLMLoggingBaseClass): attempt.model_call_details = dict(self.model_call_details) return attempt + def guardrail_scoped_copy(self) -> "Logging": + request_params: Final = _LITELLM_PARAMS.validate_python( + self.litellm_params or {} # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped + ) + litellm_params: Final = { + **request_params, + **{ + key: _without_guardrail_information(_METADATA_MAPPING.validate_python(value)) + for key, value in request_params.items() + if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict) + }, + } + scoped: Final = copy.copy(self) + scoped.litellm_params = litellm_params + scoped.model_call_details = {**self.model_call_details, "litellm_params": litellm_params} + return scoped + def dispatch_attempt_failure_handlers( self, exception: Exception, traceback_exception: str ) -> Coroutine[None, None, None]: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 2348d8e57d6..ede5523eea4 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2458,6 +2458,25 @@ class CustomStreamWrapper: def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None: self._earlier_rounds_spend = earlier + def partial_response(self) -> "ModelResponse | TextCompletionResponse | None": + if not self.chunks: + return None + partial: Final = litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages if isinstance(self.messages, list) else None, + logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, + ) + if partial is None: + return None + usage: Final[object] = getattr(partial, "usage", None) + if not isinstance(usage, Usage): + return None + if self.model: + partial.model = self.model + backfill_missing_cache_usage_fields(usage) + return partial + def _record_partial_usage_for_failure(self) -> None: """ A stream that breaks mid-flight still billed the provider for the chunks @@ -2467,23 +2486,15 @@ class CustomStreamWrapper: moves it into terminal failure usage before opening the next attempt, so a recovered attempt logs its own usage. """ - if self.logging_obj is None or not self.chunks: + if self.logging_obj is None: return try: - partial_response: Final = litellm.stream_chunk_builder( - chunks=self.chunks, - messages=self.messages if isinstance(self.messages, list) else None, - logging_obj=self.logging_obj, - count_prompt_tokens=self.count_prompt_tokens, - ) + partial_response: Final = self.partial_response() if partial_response is None: return - usage: Final = cast(Usage | None, getattr(partial_response, "usage", None)) - if usage is None: + usage: Final[object] = getattr(partial_response, "usage", None) + if not isinstance(usage, Usage): return - if self.model: - partial_response.model = self.model - backfill_missing_cache_usage_fields(usage) partial_cost: Final = self.logging_obj.response_cost_calculator(result=partial_response) or 0.0 if self._earlier_rounds_spend is None: self.logging_obj.model_call_details["combined_usage_object"] = usage diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c22d5a9857d..a36884b90dd 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -137,6 +137,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, ) +from litellm.responses.mcp.parked_stream_dispatch import parked_success from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task from litellm.router import Router from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX @@ -527,6 +528,20 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return await _dispatch_disconnect_billing_event(logging_obj, partial_response) +@runtime_checkable +class _BillsInFlightRoundOnDisconnect(Protocol): + def bill_in_flight_round_on_disconnect(self) -> None: ... + + +def _bill_in_flight_round_on_disconnect(response: object) -> None: + if not isinstance(response, _BillsInFlightRoundOnDisconnect): + return + try: + response.bill_in_flight_round_on_disconnect() + except Exception as e: # noqa: BLE001 # disconnect billing is best-effort; never break stream teardown + verbose_proxy_logger.debug("Failed to bill the in-flight round after client disconnect: %s", e) + + async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, partial_response: object) -> bool: try: await logging_obj.dispatch_success_handlers( @@ -2052,6 +2067,13 @@ class _ClientStreamFailureReport: ) +def _metadata_logged_on(logged_on: LiteLLMLoggingObj) -> Mapping[str, object]: + logged_params: Final = object_mapping(logged_on.model_call_details.get("litellm_params")) or MappingProxyType({}) + return MappingProxyType( + {key: logged_params[key] for key in ("metadata", "litellm_metadata") if key in logged_params} + ) + + class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict, served_model: str | None = None): self.data = data @@ -4065,7 +4087,15 @@ class ProxyBaseLLMRequestProcessing: assembled_response: object, cache_hit: object, logged_on: LiteLLMLoggingObj | None = None ) -> None: await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( - captured_data=_captured_data, + captured_data=_captured_data + if logged_on is None + else { + **_captured_data, + **{ + key: value for key, value in _metadata_logged_on(logged_on).items() if key in _captured_data + }, + "litellm_logging_obj": logged_on, + }, captured_user_api_key_dict=_captured_user_api_key_dict, captured_logging_obj=_captured_logging_obj if logged_on is None else logged_on, assembled_response=assembled_response, @@ -4088,14 +4118,8 @@ class ProxyBaseLLMRequestProcessing: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine): GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) - case (assembled_response, cache_hit): - await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers( - assembled_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - prefer_async_handlers=True, - ) + case _ if (success := parked_success(args)) is not None: + await success.dispatch(_captured_native_logging_obj) case _: verbose_proxy_logger.error( "Deferred stream logging dropped: unexpected stored args shape %s", @@ -4423,6 +4447,7 @@ class ProxyBaseLLMRequestProcessing: client_disconnected, ) if recorded_client_disconnect: + _bill_in_flight_round_on_disconnect(response) deferred_stream_logging_armed: Final = _deferred_stream_logging_is_armed(request_data) ProxyLogging.fire_deferred_stream_logging(request_data) # A disconnect-time success event (the deferred-guardrail flush diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 90438af1487..06308a147fd 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -273,6 +273,18 @@ async def acompletion_with_mcp( if isinstance(round_billing, ChatStreamRoundBilling): round_billing.release_abandoned() + def bill_in_flight_round_on_disconnect(self) -> None: + follow_up: Final[object] = self.follow_up_stream + if ( + self._round_billing is None + or self.follow_up_exhausted + or not isinstance(follow_up, CustomStreamWrapper) + ): + return + partial: Final = follow_up.partial_response() + if isinstance(partial, ModelResponse): + self._round_billing.bill_partial_follow_up(partial) + def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_list_tools to the first chunk.""" from litellm.types.utils import ( @@ -598,6 +610,9 @@ async def acompletion_with_mcp( self._custom_iterator.release_abandoned() await super().aclose() + def bill_in_flight_round_on_disconnect(self) -> None: + self._custom_iterator.bill_in_flight_round_on_disconnect() + def __iter__(self): # For synchronous iteration, create a sync wrapper if self._sync_iterator is None: diff --git a/litellm/responses/mcp/chat_stream_round_billing.py b/litellm/responses/mcp/chat_stream_round_billing.py index 4df79ca13ba..3b0de7a0fde 100644 --- a/litellm/responses/mcp/chat_stream_round_billing.py +++ b/litellm/responses/mcp/chat_stream_round_billing.py @@ -1,39 +1,26 @@ import asyncio -from collections.abc import Callable, Coroutine, Generator +from collections.abc import Generator from contextlib import contextmanager from dataclasses import dataclass, replace from types import MappingProxyType -from typing import Final, Protocol, TypeAlias +from typing import Final, TypeAlias from litellm.litellm_core_utils.asyncify import is_event_loop_running from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.parked_stream_dispatch import ( + DISPATCH_RELEASE, + PARKED_DISPATCH, + SERVED_FROM_CACHE, + DeferredStreamRelease, + parked_success, +) from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown from litellm.responses.streaming_iterator import spawn_logging_task from litellm.types.utils import ModelResponse, ModelResponseStream, Usage Parked: TypeAlias = tuple[object, ...] -_Release: TypeAlias = Callable[..., Coroutine[object, object, None]] - -_PARKED_DISPATCH: Final = "_deferred_stream_complete_args" -_DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" -_SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) - - -class _DispatchesSuccessHandlers(Protocol): - async def dispatch_success_handlers( - self, - result: object = None, - start_time: object = None, - end_time: object = None, - cache_hit: object = None, - prefer_async_handlers: bool = False, - ) -> None: ... - - -def _success_dispatcher(logger: LiteLLMLoggingObj) -> _DispatchesSuccessHandlers: - return logger @dataclass(frozen=True, slots=True) @@ -73,8 +60,13 @@ class _FirstRound: def client_usage_chunk(self, billed: ModelResponse | None) -> ModelResponseStream | None: if self.usage_chunk is None: return None - merged: Final = merged_chat_usage(self.priced.usage, self.usage_chunk.usage) - usage: Final = merged if billed is None else merged.model_copy(update=MappingProxyType({"cost": _cost(billed)})) + usage: Final = ( + self.priced.usage.model_copy(update=MappingProxyType({"cost": self.priced.price.cost})) + if billed is None + else merged_chat_usage(self.priced.usage, self.usage_chunk.usage).model_copy( + update=MappingProxyType({"cost": _cost(billed)}) + ) + ) return self.usage_chunk.chunk.model_copy(update=MappingProxyType({"usage": usage})) @@ -99,12 +91,14 @@ class ChatStreamRoundBilling: def __init__(self, logging_obj: LiteLLMLoggingObj) -> None: self._logging_obj: Final = logging_obj self._first: _FirstRound | None = None + self._reading_loop: asyncio.AbstractEventLoop | None = None @contextmanager def armed_while_reading(self) -> Generator[None]: - armed: Final = getattr(self._logging_obj, _DISPATCH_RELEASE, None) is None + armed: Final = getattr(self._logging_obj, DISPATCH_RELEASE, None) is None if armed: - setattr(self._logging_obj, _DISPATCH_RELEASE, self._log_success) + setattr(self._logging_obj, DISPATCH_RELEASE, self._log_success) + self._reading_loop = asyncio.get_running_loop() try: yield finally: @@ -135,20 +129,30 @@ class ChatStreamRoundBilling: def fail(self) -> None: self._bank_apart_from_the_follow_up() first: Final = self._first - if first is None or first.banked is None or self._peek_parked() is not first.parked: + if first is None or first.banked is None or not self._holds_round_one_alone(first): self.finish() return self._first = None - setattr(self._logging_obj, _PARKED_DISPATCH, None) - self._log_through_the_release((*first.parked, first.banked)) + setattr(self._logging_obj, PARKED_DISPATCH, None) + self._log_through_the_release((*first.parked, first.banked.guardrail_scoped_copy())) + + def bill_partial_follow_up(self, partial: ModelResponse) -> None: + first: Final = self._first + final_round: Final = _assembled_round((partial, False)) + if first is None or final_round is None or not self._holds_round_one_alone(first): + return + self._first = None + setattr( + self._logging_obj, PARKED_DISPATCH, self._billed_for_both_rounds(first.priced, self._priced(final_round)) + ) def release_abandoned(self) -> None: if not self._owns_release(): return - if is_event_loop_running(): + if is_event_loop_running() or self._idle_reading_loop() is not None: self.finish() return - setattr(self._logging_obj, _PARKED_DISPATCH, None) + setattr(self._logging_obj, PARKED_DISPATCH, None) self._first = None self._disarm() @@ -162,7 +166,7 @@ class ChatStreamRoundBilling: return logged_alone: Final = (*first.parked, banked) self._first = replace(first, logged_alone=logged_alone, banked=banked) - setattr(self._logging_obj, _PARKED_DISPATCH, logged_alone) + setattr(self._logging_obj, PARKED_DISPATCH, logged_alone) def _release_rounds(self) -> ModelResponse | None: parked: Final = self._take_parked() @@ -196,7 +200,7 @@ class ChatStreamRoundBilling: def _priced(self, assembled: _AssembledRound) -> _PricedRound: served: Final = assembled.cache_hit or served_from_cache(assembled.response) - price: Final = _SERVED_FROM_CACHE if served else self._logging_obj.uncached_price(assembled.response) + price: Final = SERVED_FROM_CACHE if served else self._logging_obj.uncached_price(assembled.response) priced: Final = assembled.response.model_copy() set_hidden_params( priced, @@ -205,36 +209,47 @@ class ChatStreamRoundBilling: return _PricedRound(priced, assembled.usage, price, served) def _peek_parked(self) -> Parked | None: - parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None) + parked: Final[Parked | None] = getattr(self._logging_obj, PARKED_DISPATCH, None) return parked if isinstance(parked, tuple) else None def _take_parked(self) -> Parked | None: parked: Final = self._peek_parked() - setattr(self._logging_obj, _PARKED_DISPATCH, None) + setattr(self._logging_obj, PARKED_DISPATCH, None) return parked + def _holds_round_one_alone(self, first: _FirstRound) -> bool: + parked: Final = self._peek_parked() + return parked is first.parked or parked is first.logged_alone + def _owns_release(self) -> bool: - release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + release: Final[object] = getattr(self._logging_obj, DISPATCH_RELEASE, None) return release == self._log_success def _release(self, parked: Parked) -> None: - if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release(): - setattr(self._logging_obj, _PARKED_DISPATCH, parked) + if getattr(self._logging_obj, DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, PARKED_DISPATCH, parked) return self._log_through_the_release(parked) def _log_through_the_release(self, parked: Parked) -> None: - armed: Final[_Release | None] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) - setattr(self._logging_obj, _DISPATCH_RELEASE, None) + armed: Final[DeferredStreamRelease | None] = getattr(self._logging_obj, DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) logged: Final = (self._log_success if armed is None else armed)(*parked) + idle_reading_loop: Final = None if is_event_loop_running() else self._idle_reading_loop() + if idle_reading_loop is not None: + idle_reading_loop.run_until_complete(logged) + return spawn_logging_task(asyncio.get_running_loop(), logged, task_name="Chat completions MCP stream success log") + def _idle_reading_loop(self) -> asyncio.AbstractEventLoop | None: + loop: Final = self._reading_loop + return loop if loop is not None and not loop.is_closed() and not loop.is_running() else None + def _disarm(self) -> None: if self._owns_release(): - setattr(self._logging_obj, _DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) - async def _log_success(self, assembled: object, cache_hit: object, banked: object = None) -> None: - logger: Final = banked if isinstance(banked, LiteLLMLoggingObj) else self._logging_obj - await _success_dispatcher(logger).dispatch_success_handlers( - assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True - ) + async def _log_success(self, *parked: object) -> None: + success: Final = parked_success(parked) + if success is not None: + await success.dispatch(self._logging_obj) diff --git a/litellm/responses/mcp/parked_stream_dispatch.py b/litellm/responses/mcp/parked_stream_dispatch.py new file mode 100644 index 00000000000..99adbbbf886 --- /dev/null +++ b/litellm/responses/mcp/parked_stream_dispatch.py @@ -0,0 +1,52 @@ +from collections.abc import Coroutine +from dataclasses import dataclass +from typing import Final, Protocol + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import UncachedPrice + +PARKED_DISPATCH: Final = "_deferred_stream_complete_args" +DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" +SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) + + +class DispatchesSuccessHandlers(Protocol): + async def dispatch_success_handlers( + self, + result: object = None, + start_time: object = None, + end_time: object = None, + cache_hit: object = None, + prefer_async_handlers: bool = False, + ) -> None: ... + + +class DeferredStreamRelease(Protocol): + def __call__(self, *parked: object) -> Coroutine[object, object, None]: ... + + +def success_dispatcher(logger: LiteLLMLoggingObj) -> DispatchesSuccessHandlers: + return logger + + +@dataclass(frozen=True, slots=True) +class ParkedSuccess: + assembled: object + cache_hit: object + logged_on: LiteLLMLoggingObj | None = None + + async def dispatch(self, request_logger: LiteLLMLoggingObj) -> None: + logger: Final = request_logger if self.logged_on is None else self.logged_on + await success_dispatcher(logger).dispatch_success_handlers( + self.assembled, cache_hit=self.cache_hit, start_time=None, end_time=None, prefer_async_handlers=True + ) + + +def parked_success(parked: tuple[object, ...]) -> ParkedSuccess | None: + match parked: + case (assembled, cache_hit): + return ParkedSuccess(assembled, cache_hit) + case (assembled, cache_hit, LiteLLMLoggingObj() as logged_on): + return ParkedSuccess(assembled, cache_hit, logged_on) + case _: + return None diff --git a/litellm/responses/mcp/stream_round_billing.py b/litellm/responses/mcp/stream_round_billing.py index 443147230bf..5cb256a8ac0 100644 --- a/litellm/responses/mcp/stream_round_billing.py +++ b/litellm/responses/mcp/stream_round_billing.py @@ -2,11 +2,18 @@ import asyncio from collections.abc import Callable, Coroutine, Sequence from dataclasses import dataclass from datetime import datetime -from typing import Final, Protocol, TypeAlias +from typing import Final, TypeAlias from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.parked_stream_dispatch import ( + DISPATCH_RELEASE, + PARKED_DISPATCH, + SERVED_FROM_CACHE, + parked_success, + success_dispatcher, +) from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, spawn_logging_task from litellm.responses.utils import ResponseAPILoggingUtils @@ -16,25 +23,6 @@ from litellm.types.utils import CostBreakdown, ModelResponse, Usage Parked: TypeAlias = tuple[object, ...] RestoresHeaders: TypeAlias = Callable[[ResponsesAPIStreamingResponse], None] -_PARKED_DISPATCH: Final = "_deferred_stream_complete_args" -_DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" -_SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) - - -class _DispatchesSuccessHandlers(Protocol): - async def dispatch_success_handlers( - self, - result: object = None, - start_time: object = None, - end_time: object = None, - cache_hit: object = None, - prefer_async_handlers: bool = False, - ) -> None: ... - - -def _success_dispatcher(logger: LiteLLMLoggingObj) -> _DispatchesSuccessHandlers: - return logger - def _event_loop_running() -> bool: try: @@ -198,8 +186,8 @@ class StreamedRoundBilling: self._billed: _BilledRounds | None = None self._banked: Parked | None = None self._banked_row: _RoundRow | None = None - if getattr(logging_obj, _DISPATCH_RELEASE, None) is None: - setattr(logging_obj, _DISPATCH_RELEASE, self._run_parked) + if getattr(logging_obj, DISPATCH_RELEASE, None) is None: + setattr(logging_obj, DISPATCH_RELEASE, self._run_parked) def complete_round( self, @@ -286,7 +274,7 @@ class StreamedRoundBilling: def _price(self, response: ResponsesAPIResponse, round_parked: Parked | None) -> UncachedPrice: if served_from_cache(response): - return _SERVED_FROM_CACHE + return SERVED_FROM_CACHE match round_parked: case (ModelResponse() as assembled, bool()): return self._logging_obj.uncached_price(assembled) @@ -297,7 +285,7 @@ class StreamedRoundBilling: self._banked = parked self._banked_row = row if parked is not None: - setattr(self._logging_obj, _PARKED_DISPATCH, parked) + setattr(self._logging_obj, PARKED_DISPATCH, parked) def _bank_apart_from_later_rounds(self, billed: _BilledRounds) -> None: row: Final = self._banked_row @@ -320,7 +308,7 @@ class StreamedRoundBilling: @staticmethod async def _log_success(success: _SuccessLog) -> None: success.carry_bridged_usage() - await _success_dispatcher(success.logger).dispatch_success_handlers( + await success_dispatcher(success.logger).dispatch_success_handlers( success.row.logged, start_time=success.logger.start_time, end_time=success.row.end_time, @@ -329,10 +317,10 @@ class StreamedRoundBilling: ) def _take_parked(self) -> Parked | None: - parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None) + parked: Final[Parked | None] = getattr(self._logging_obj, PARKED_DISPATCH, None) if not isinstance(parked, tuple): return None - setattr(self._logging_obj, _PARKED_DISPATCH, None) + setattr(self._logging_obj, PARKED_DISPATCH, None) return parked async def _run_parked(self, *parked: object) -> None: @@ -341,29 +329,27 @@ class StreamedRoundBilling: await dispatch case (dispatch,) if asyncio.iscoroutine(dispatch): await dispatch - case (assembled, cache_hit): - await _success_dispatcher(self._logging_obj).dispatch_success_handlers( - assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True - ) + case _ if (success := parked_success(parked)) is not None: + await success.dispatch(self._logging_obj) case _: return def _owns_release(self) -> bool: - release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + release: Final[object] = getattr(self._logging_obj, DISPATCH_RELEASE, None) return release == self._run_parked def _release(self, parked: Parked) -> None: - if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release(): - setattr(self._logging_obj, _PARKED_DISPATCH, parked) + if getattr(self._logging_obj, DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, PARKED_DISPATCH, parked) return - setattr(self._logging_obj, _DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) spawn_logging_task( asyncio.get_running_loop(), self._run_parked(*parked), task_name="Responses MCP stream success logging" ) def _disarm(self) -> None: if self._owns_release(): - setattr(self._logging_obj, _DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) @staticmethod def _close(parked: Parked | None) -> None: diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 09dfe020cf5..d54e623f6a1 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -104,11 +104,12 @@ def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway ROUND_USAGE_GUARDRAIL: Final = "mcp-llm-round-usage" +CONTENT_FILTER: Final = f"mcp-llm-content-filter-{uuid.uuid4().hex[:8]}" @pytest.fixture(scope="module") def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: - guardrail: Final[Mapping[str, JsonValue]] = { + round_usage: Final[Mapping[str, JsonValue]] = { "guardrail_name": ROUND_USAGE_GUARDRAIL, "litellm_params": { "guardrail": "integration._support.hook_count_probe.UsageRecordingPostCallGuardrail", @@ -116,7 +117,17 @@ def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[t "default_on": True, }, } - with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), (guardrail,)) as served: + content_filter: Final[Mapping[str, JsonValue]] = { + "guardrail_name": CONTENT_FILTER, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "default_on": False, + "blocked_words": [{"keyword": "forbidden-keeper", "action": "BLOCK"}], + }, + } + guardrails: Final = (round_usage, content_filter) + with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), guardrails) as served: yield served @@ -1153,6 +1164,59 @@ def test_streaming_chat_auto_execute_whose_follow_up_fails_runs_the_post_call_gu ] +def _chat_rows(call_id: str, statuses: Sequence[str]) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + 'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s', + (call_id, "acompletion"), + ), + lambda found: _statuses(found) == sorted(statuses), + seconds=70, + ) + + +def _entries_of(guardrail_name: str, row: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + return tuple( + entry + for entry in _guardrail_information(row["metadata"]) + if isinstance(entry, dict) and entry.get("guardrail_name") == guardrail_name + ) + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_counts_its_post_call_content_filter_once( + guarded_gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(guarded_gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send( + _granted_key(rig), + [AUTO], + stream=True, + stream_options={"include_usage": True}, + guardrails=[CONTENT_FILTER], + ) + assert FOLLOW_UP_OUTAGE in response.text, response.text + rows: Final = _chat_rows(response.headers["x-litellm-call-id"], ("success", "failure")) + filtered: Final = tuple( + (row["status"], entry.get("guardrail_status")) for row in rows for entry in _entries_of(CONTENT_FILTER, row) + ) + assert filtered == (("failure", "guardrail_failed_to_respond"),), rows + (guardrail_id,) = frozenset( + str(entry.get("guardrail_id") or CONTENT_FILTER) + for row in rows + for entry in _entries_of(CONTENT_FILTER, row) + ) + counted: Final = eventually( + lambda: read_rows( + 'SELECT requests_evaluated, flagged_count FROM "LiteLLM_DailyGuardrailMetrics" WHERE guardrail_id=%s', + (guardrail_id,), + ), + lambda found: bool(found), + seconds=70, + ) + assert counted == [{"requests_evaluated": 1, "flagged_count": 1}], counted + + @pytest.mark.parametrize("include_usage", (True, False), ids=("include-usage", "no-usage")) def test_streaming_chat_auto_execute_answered_in_one_round_ends_with_the_usage_it_billed( gateway: Gateway, include_usage: bool @@ -1430,6 +1494,27 @@ def test_streaming_responses_follow_up_broken_mid_stream_bills_round_one_once_an assert float(str(failure["spend"])) == pytest.approx(_price(failed_round)), rows +def test_streaming_chat_follow_up_broken_mid_answer_bills_round_one_and_its_partial_answer_with_one_fixed_margin( + gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_breaks=True) + with _rig(gateway, "chat", upstream="messages_bridge", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert "the sum" in response.text, response.text + assert ANSWER not in response.text, response.text + assert len(rig.upstream_tools()) == 2 + round_one: Final = _price(FIRST_ROUND) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(round_one)) + assert _statuses(rows) == ["failure", "success"], rows + (success,) = (row for row in rows if row["status"] == "success") + (failure,) = (row for row in rows if row["status"] == "failure") + assert (success["prompt_tokens"], success["completion_tokens"]) == FIRST_ROUND, rows + assert float(str(success["spend"])) == pytest.approx(round_one), rows + partial_answer: Final = (int(str(failure["prompt_tokens"])), int(str(failure["completion_tokens"]))) + assert min(partial_answer) > 0, rows + assert float(str(failure["spend"])) == pytest.approx(_price(partial_answer)), rows + + @pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_once( gateway: Gateway, upstream: StreamedUpstream diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index fdcbcbfcd05..4731ea89284 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -9056,6 +9056,34 @@ class TestStreamingClientDisconnectBilling: proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + @pytest.mark.asyncio + async def test_disconnect_whose_in_flight_round_billing_raises_still_closes_the_stream_and_frees_the_slot(self): + import types + + response = await self._start_partial_stream() + closed: Final = AsyncMock() + + def cannot_price_the_in_flight_round() -> None: + raise RuntimeError("the in-flight round could not be priced") + + stream: Final = types.SimpleNamespace( + chunks=[], messages=None, aclose=closed, bill_in_flight_round_on_disconnect=cannot_price_the_in_flight_round + ) + proxy_logging_obj: Final = types.SimpleNamespace(arelease_max_parallel_requests_on_disconnect=AsyncMock()) + + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=stream, + stream_completed=False, + client_disconnected=True, + user_api_key_dict=MagicMock(), + proxy_logging_obj=proxy_logging_obj, + ) + + closed.assert_awaited_once() + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + async def _bill_and_collect_success_event(self, prepare=None, request_data=None): recorder = _RecordingSuccessLogger() original_callbacks = litellm.callbacks diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 23ff0727098..45a6978e952 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1562,6 +1562,7 @@ _CHAT_ROUND_TWO: Final = MappingProxyType( ), } ) +_NEVER_LOGGED_SECONDS: Final = 60 _CONTEXT_WINDOW_ERROR: Final = MappingProxyType( { "error": MappingProxyType( @@ -1581,16 +1582,31 @@ class _ChatRoundRecorder(CustomLogger): super().__init__() self.logged: tuple[StandardLoggingPayload, ...] = () self.failed: tuple[StandardLoggingPayload, ...] = () + self._recorded: Final = asyncio.Condition() async def async_log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.logged = (*self.logged, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + await self._announce() async def async_log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.failed = (*self.failed, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + await self._announce() + + async def _announce(self) -> None: + async with self._recorded: + self._recorded.notify_all() + + async def settled(self, logged: int = 0, failed: int = 0) -> None: + async with self._recorded: + await asyncio.wait_for( + self._recorded.wait_for(lambda: len(self.logged) >= logged and len(self.failed) >= failed), + timeout=_NEVER_LOGGED_SECONDS, + ) + await _settled() def _chat_round_cost(body: Mapping[str, object]) -> float: @@ -1606,14 +1622,6 @@ async def _settled() -> None: await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) -async def _settled_once(logged: Callable[[], bool]) -> None: - for _ in range(200): - if logged(): - break - await asyncio.sleep(0.01) - await _settled() - - def _serve_chat_rounds( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response] ) -> None: @@ -1785,7 +1793,9 @@ def _cut_once_the_answer_starts() -> Response: return Response(200, stream=_CutOnceTheAnswerStarts(), headers={"content-type": "text/event-stream"}) -async def _open_chat_echo_stream(logging_obj: Logging) -> CustomStreamWrapper: +async def _open_chat_echo_stream( + logging_obj: Logging, metadata: dict[str, object] | None = None +) -> CustomStreamWrapper: stream: Final = await litellm.acompletion( model="openai/gpt-5.6", messages=[{"role": "user", "content": "Call echo with PING"}], @@ -1795,6 +1805,7 @@ async def _open_chat_echo_stream(logging_obj: Logging) -> CustomStreamWrapper: stream=True, stream_options={"include_usage": True}, litellm_logging_obj=logging_obj, + metadata=metadata, ) assert isinstance(stream, CustomStreamWrapper) return stream @@ -1860,7 +1871,7 @@ async def test_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_ await _settled() logged_before_release: Final = recorder.logged ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - await _settled_once(lambda: len(recorder.logged) > 0) + await recorder.settled(logged=1) assert logged_before_release == () _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) @@ -1888,7 +1899,7 @@ async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_bills_both_r await _read_all(stream) assert ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked) await logging_obj.async_failure_handler(blocked, "blocked by output-policy") - await _settled_once(lambda: len(recorder.failed) > 0) + await recorder.settled(failed=1) charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN assert recorder.logged == () @@ -1910,7 +1921,7 @@ async def test_stream_auto_execute_answered_without_a_tool_call_logs_its_one_rou async for chunk in stream: if chunk.choices and chunk.choices[0].finish_reason is not None: break - await _settled_once(lambda: len(recorder.logged) > 0) + await recorder.settled(logged=1) logged_once_the_answer_finished: Final = recorder.logged after_the_answer: Final = await _read_all(stream) await _settled() @@ -1947,11 +1958,30 @@ async def test_stream_auto_execute_bills_the_finished_first_round_when_the_calle with _echo_tool_served(): await stop_reading_once_the_follow_up_answers() - await _settled_once(lambda: len(recorder.logged) > 0) + await recorder.settled(logged=1) _assert_one_row_bills_round_one(recorder.logged) +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_first_round_breaks_mid_answer_bills_only_its_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_cut_once_the_answer_starts(),)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-first-round-breaks", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + with pytest.raises(litellm.ServiceUnavailableError): + await _read_all(stream) + await recorder.settled(failed=1) + + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert recorder.failed[0]["completion_tokens"] > 0 + + @pytest.mark.asyncio @pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( @@ -1976,7 +2006,7 @@ async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_r await _read_all(stream) if not ProxyLogging._discard_deferred_stream_logging_for_failure(dict(request_data), broken.value): ProxyLogging.fire_deferred_stream_logging(dict(request_data)) - await _settled_once(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + await recorder.settled(logged=1, failed=1) _assert_one_row_bills_round_one(recorder.logged) assert len(recorder.failed) == 1 @@ -2004,23 +2034,67 @@ async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_r @pytest.mark.asyncio -async def test_stream_auto_execute_leaves_no_deferred_log_armed_for_the_proxy_to_orphan_after_its_first_chunk( +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_whose_follow_up_breaks_reports_the_request_guardrail_result_on_the_failure_row_only( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-breaks-guarded", stream=True) + metadata: Final[dict[str, object]] = {} + request_data: Final = {"metadata": metadata, "litellm_logging_obj": logging_obj} + content_filter: Final = CustomGuardrail(guardrail_name="request-content-filter", event_hook="post_call") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj, metadata) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError) as broken: + await _read_all(stream) + content_filter.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=str(broken.value), + request_data=request_data, + guardrail_status="guardrail_failed_to_respond", + ) + if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, broken.value): + ProxyLogging.fire_deferred_stream_logging(request_data) + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert recorder.logged[0]["guardrail_information"] is None + assert [entry["guardrail_name"] for entry in recorder.failed[0]["guardrail_information"] or ()] == [ + "request-content-filter" + ] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_served_through_the_proxy_stream_hook_logs_both_rounds_once_when_it_ends( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) recorder: Final = _ChatRoundRecorder() - logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-first-chunk-read-ahead", stream=True) + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-proxy-stream-hook", stream=True) with _echo_tool_served(): - stream: Final = await _open_chat_echo_stream(logging_obj) - rest: Final = stream.__aiter__() - first_chunk: Final = await rest.__anext__() - armed_after_the_first_chunk: Final = getattr(logging_obj, "_on_deferred_stream_complete", None) - chunks: Final = (first_chunk, *[chunk async for chunk in rest]) - await _settled_once(lambda: len(recorder.logged) > 0) + served: Final = ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=await _open_chat_echo_stream(logging_obj), + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj}, + ) + chunks: Final = tuple([chunk async for chunk in served]) + await recorder.settled(logged=1) - assert armed_after_the_first_chunk is None _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) @@ -2035,17 +2109,152 @@ async def test_stream_auto_execute_whose_follow_up_streams_nothing_still_bills_t with _echo_tool_served(): await _read_all(await _open_chat_echo_stream(logging_obj)) - await _settled_once(lambda: len(recorder.logged) > 0) + await recorder.settled(logged=1) _assert_one_row_bills_round_one(recorder.logged) -async def _stream_read_until_the_follow_up_answers( - logging_obj: Logging, opened: "queue.SimpleQueue[CustomStreamWrapper]" +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_follow_up_cannot_be_assembled_ends_with_the_usage_of_the_round_it_billed( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-unassembled", stream=True) + + def cannot_assemble(*args: object, **kwargs: object) -> ModelResponse: + raise ValueError("the follow-up's chunks cannot be assembled") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + await _read_until(stream, "PING") + monkeypatch.setattr(litellm, "stream_chunk_builder", cannot_assemble) + rest: Final = await _read_all(stream) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + usage: Final = _streamed_usage(rest) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (71, 19, 90) + assert usage.cost == pytest.approx(recorder.logged[0]["response_cost"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_during_the_follow_up_bills_round_one_and_the_partial_answer_on_one_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-mid-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _read_until(stream, "PING") + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + assert row["id"] == _CHAT_ROUND_ONE["id"] + partial_prompt, partial_completion = (row["prompt_tokens"] - 71, row["completion_tokens"] - 19) + assert partial_prompt > 0 and partial_completion > 0 + partial_answer: Final = litellm.completion_cost( + completion_response=ModelResponse( + model="gpt-5.6", + usage=Usage( + prompt_tokens=partial_prompt, + completion_tokens=partial_completion, + total_tokens=partial_prompt + partial_completion, + ), + ), + model="gpt-5.6", + custom_llm_provider="openai", + ) + assert row["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE) + partial_answer - _FIXED_MARGIN) + assert recorder.failed == () + + +async def _read_through_the_usage_chunk( + chunks: AsyncIterator[ModelResponseStream], read: tuple[ModelResponseStream, ...] = () +) -> tuple[ModelResponseStream, ...]: + chunk: Final = await anext(chunks) + if isinstance(getattr(chunk, "usage", None), Usage): + return (*read, chunk) + return await _read_through_the_usage_chunk(chunks, (*read, chunk)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_after_its_whole_answer_bills_both_rounds_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-after-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + chunks: Final = await _read_through_the_usage_chunk(aiter(stream)) + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + assert recorder.failed == () + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bridged_to_messages_bills_round_one_when_its_follow_up_breaks( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-bridged-to-messages", stream=True) + + async def bridged(stream: CustomStreamWrapper) -> AsyncIterator[ModelResponseStream]: + async for chunk in stream: + yield chunk + + with _echo_tool_served(): + relayed: Final = bridged(await _open_chat_echo_stream(logging_obj)) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=relayed, route_type="anthropic_messages", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError): + _ = [chunk async for chunk in relayed] + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + + +def _read_synchronously_until_the_follow_up_answers( + logging_obj: Logging, loop: asyncio.AbstractEventLoop, before_dropping: Callable[[], None] = lambda: None ) -> None: - stream: Final = await _open_chat_echo_stream(logging_obj) - await _read_until(stream, "PING") - opened.put(stream) + stream: Final = loop.run_until_complete(_open_chat_echo_stream(logging_obj)) + for chunk in stream: + if chunk.choices and chunk.choices[0].delta.content == "PING": + before_dropping() + return + pytest.fail("the stream never sent 'PING'") _CHAT_STREAM_BILLING_NAMES: Final = ("MCPStreamingIterator", "ChatStreamRoundBilling", "_log_success") @@ -2063,26 +2272,57 @@ class _UnraisableRecorder: self.raised = (*self.raised, f"{unraisable.object!r}: {unraisable.exc_value!r}") -def test_stream_auto_execute_dropped_after_its_event_loop_closed_discards_its_parked_log_cleanly( +def test_stream_auto_execute_read_synchronously_and_dropped_mid_answer_bills_the_finished_first_round( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) recorder: Final = _ChatRoundRecorder() unraisable: Final = _UnraisableRecorder() - opened: Final = queue.SimpleQueue[CustomStreamWrapper]() + loop: Final = asyncio.new_event_loop() gc.collect() monkeypatch.setattr(sys, "unraisablehook", unraisable) + asyncio.set_event_loop(loop) - with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - asyncio.run( - _stream_read_until_the_follow_up_answers( - _chat_rounds_logging_obj(recorder, "mcp-chat-stream-loop-closed", stream=True), opened + try: + with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _read_synchronously_until_the_follow_up_answers( + _chat_rounds_logging_obj(recorder, "mcp-chat-stream-sync-dropped", stream=True), loop ) - ) - opened.get_nowait() - gc.collect() + gc.collect() + loop.run_until_complete(recorder.settled(logged=1)) + finally: + asyncio.set_event_loop(None) + loop.close() + + _assert_one_row_bills_round_one(recorder.logged) + assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == [] + + +def test_stream_auto_execute_read_synchronously_and_dropped_after_its_loop_closed_drops_without_errors( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + unraisable: Final = _UnraisableRecorder() + loop: Final = asyncio.new_event_loop() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + asyncio.set_event_loop(loop) + + try: + with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _read_synchronously_until_the_follow_up_answers( + _chat_rounds_logging_obj(recorder, "mcp-chat-stream-sync-dropped-after-close", stream=True), + loop, + before_dropping=loop.close, + ) + gc.collect() + finally: + asyncio.set_event_loop(None) + loop.close() - assert recorder.logged == () assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == [] assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == [] From 301ce2ce2c7cc3d1acd3a0a047b35d7288a3fad8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 22:45:07 -0700 Subject: [PATCH 05/10] 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 --- .../mcp/test_chat_completions_handler.py | 35 ++++++++++++++----- 1 file changed, 26 insertions(+), 9 deletions(-) diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 45a6978e952..473cb6bc083 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -2118,22 +2118,39 @@ async def test_stream_auto_execute_whose_follow_up_streams_nothing_still_bills_t async def test_stream_auto_execute_whose_follow_up_cannot_be_assembled_ends_with_the_usage_of_the_round_it_billed( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): - _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + 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))) recorder: Final = _ChatRoundRecorder() logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-unassembled", stream=True) - def cannot_assemble(*args: object, **kwargs: object) -> ModelResponse: - raise ValueError("the follow-up's chunks cannot be assembled") - with _echo_tool_served(): - stream: Final = await _open_chat_echo_stream(logging_obj) - await _read_until(stream, "PING") - monkeypatch.setattr(litellm, "stream_chunk_builder", cannot_assemble) - rest: Final = await _read_all(stream) + 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) - usage: Final = _streamed_usage(rest) + 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"]) From 7da4c7fcdb0a5a61bb8ad08e96f4d4657120d3e5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 05:14:50 -0700 Subject: [PATCH 06/10] fix(mcp): count the fixed margin once in the combined stream row's usage cost --- .../responses/mcp/chat_stream_round_billing.py | 6 +++++- .../mcp/test_chat_completions_handler.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/litellm/responses/mcp/chat_stream_round_billing.py b/litellm/responses/mcp/chat_stream_round_billing.py index 3b0de7a0fde..ad0ae8fa6cf 100644 --- a/litellm/responses/mcp/chat_stream_round_billing.py +++ b/litellm/responses/mcp/chat_stream_round_billing.py @@ -190,7 +190,11 @@ class ChatStreamRoundBilling: 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) - billed: Final = rounds.model_copy(update=MappingProxyType({"id": first.response.id})) + 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( diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 473cb6bc083..2444d481ce7 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1839,6 +1839,7 @@ def _assert_one_stream_row_bills_both_rounds_once( 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"]) @@ -1850,6 +1851,22 @@ def _assert_one_row_bills_round_one(logged: Sequence[StandardLoggingPayload]) -> 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 69d30186bf2c6332cc4cc01794623adcbce5b376 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:49:04 -0700 Subject: [PATCH 07/10] 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 --- .../responses/mcp/chat_completions_handler.py | 6 +- .../mcp/test_chat_completions_handler.py | 130 +++++++++++++++--- 2 files changed, 117 insertions(+), 19 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 06308a147fd..ec68797d51f 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -1,5 +1,6 @@ """Helpers for handling MCP-aware `/chat/completions` requests.""" +import asyncio import logging from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast @@ -383,6 +384,9 @@ async def acompletion_with_mcp( if usage_chunk is None: raise return usage_chunk + except asyncio.CancelledError: + self.bill_in_flight_round_on_disconnect() + raise except Exception: round_billing.fail() raise @@ -616,8 +620,6 @@ async def acompletion_with_mcp( 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/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 2444d481ce7..44cdfeac650 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -17,6 +17,8 @@ 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 @@ -2172,6 +2174,28 @@ async def test_stream_auto_execute_whose_follow_up_cannot_be_assembled_ends_with assert usage.cost == pytest.approx(recorder.logged[0]["response_cost"]) +def _assert_one_row_bills_round_one_and_the_partial_answer(recorder: _ChatRoundRecorder) -> 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 > 0 and partial_completion > 0 + partial_answer: Final = litellm.completion_cost( + completion_response=ModelResponse( + model="gpt-5.6", + usage=Usage( + prompt_tokens=partial_prompt, + completion_tokens=partial_completion, + total_tokens=partial_prompt + partial_completion, + ), + ), + model="gpt-5.6", + custom_llm_provider="openai", + ) + assert row["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE) + partial_answer - _FIXED_MARGIN) + assert recorder.failed == () + + @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( @@ -2196,25 +2220,97 @@ async def test_stream_auto_execute_hung_up_during_the_follow_up_bills_round_one_ ) await recorder.settled(logged=1) - assert len(recorder.logged) == 1 - row: Final = recorder.logged[0] - assert row["id"] == _CHAT_ROUND_ONE["id"] - partial_prompt, partial_completion = (row["prompt_tokens"] - 71, row["completion_tokens"] - 19) - assert partial_prompt > 0 and partial_completion > 0 - partial_answer: Final = litellm.completion_cost( - completion_response=ModelResponse( - model="gpt-5.6", - usage=Usage( - prompt_tokens=partial_prompt, - completion_tokens=partial_completion, - total_tokens=partial_prompt + partial_completion, - ), + _assert_one_row_bills_round_one_and_the_partial_answer(recorder) + + +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"}), ), - 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 == () + 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", + }, + } + ] + ) + 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) async def _read_through_the_usage_chunk( From c0a2b3ff63ceeaa66b02f7f483f20d2036adbeb4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 09:19:55 -0700 Subject: [PATCH 08/10] 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 --- .../responses/mcp/chat_completions_handler.py | 1 + .../mcp/test_chat_completions_handler.py | 17 +++++++++++++---- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index ec68797d51f..2e05639e000 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -554,6 +554,7 @@ async def acompletion_with_mcp( 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 self.follow_up_stream = follow_up_response from litellm._logging import verbose_logger diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 44cdfeac650..1daa4718f83 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1847,6 +1847,11 @@ def _assert_one_stream_row_bills_both_rounds_once( 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) @@ -2045,6 +2050,7 @@ async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_r ) - _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) @@ -2174,12 +2180,15 @@ async def test_stream_auto_execute_whose_follow_up_cannot_be_assembled_ends_with assert usage.cost == pytest.approx(recorder.logged[0]["response_cost"]) -def _assert_one_row_bills_round_one_and_the_partial_answer(recorder: _ChatRoundRecorder) -> None: +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 > 0 and partial_completion > 0 + 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", @@ -2220,7 +2229,7 @@ async def test_stream_auto_execute_hung_up_during_the_follow_up_bills_round_one_ ) await recorder.settled(logged=1) - _assert_one_row_bills_round_one_and_the_partial_answer(recorder) + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) class _StallsOnceTheAnswerStarts(httpx.AsyncByteStream): @@ -2310,7 +2319,7 @@ async def test_stream_auto_execute_routed_through_the_proxy_bills_the_partial_an 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) + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) async def _read_through_the_usage_chunk( 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 09/10] 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, ...]: 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 10/10] 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