mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
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
This commit is contained in:
mateo-berri
parent
f23da840b5
commit
e3fbb6583d
7 changed files with 214 additions and 38 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
17
tests/unit/responses/mcp/test_parked_stream_dispatch.py
Normal file
17
tests/unit/responses/mcp/test_parked_stream_dispatch.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue