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