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] 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 == ()