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

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."""

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(

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