scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-10 08:28:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

fix(mcp): count the fixed margin once in the combined stream row's usage cost

Browse source
This commit is contained in:
mateo-berri 2026-10-29 05:14:50 -07:00
parent e4db619d6f
commit 7da4c7fcdb
2 changed files with 22 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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