scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-06 07:48:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 112

fix(mcp): bill round 2's partial when a routed stream's client hangs up between reads

Browse source 
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
This commit is contained in:
mateo-berri 2026-10-29 15:14:32 -07:00
parent c0a2b3ff63
commit f23da840b5
6 changed files with 143 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
litellm/proxy/common_request_processing.py
View file

@ -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

19
litellm/responses/mcp/chat_completions_handler.py
View file

@ -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:

18
litellm/responses/mcp/chat_stream_round_billing.py
View file

@ -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(

8
litellm/responses/mcp/parked_stream_dispatch.py
View file

@ -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):

6
tests/unit/proxy/test_common_request_processing.py
View file

@ -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(

122
tests/unit/responses/mcp/test_chat_completions_handler.py
View file

@ -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, ...]: