fix(mcp): count the follow-up's own messages in round 2's partial prompt

The follow-up stream shares the request's logger, so it took the client's
original messages as its prompt. When round 2 ended early, from a hang-up or a
broken upstream, its partial usage counted only those messages and left out the
assistant tool call and the tool results the follow-up actually sent. The
follow-up stream now counts the messages the follow-up call sent

litellm/responses/mcp/chat_completions_handler.py

@@ -554,6 +554,7 @@ async def acompletion_with_mcp(
                 if isinstance(follow_up_response, CustomStreamWrapper):
                     if earlier_rounds is not None:
                         follow_up_response.bill_failure_after_earlier_rounds(earlier_rounds)
+                    follow_up_response.messages = follow_up_messages
                     self.follow_up_stream = follow_up_response
                     from litellm._logging import verbose_logger
 

tests/unit/responses/mcp/test_chat_completions_handler.py

@@ -1847,6 +1847,11 @@ def _assert_one_stream_row_bills_both_rounds_once(
     assert usage.cost == pytest.approx(row["response_cost"])
 
 
+def _prompt_tokens_the_follow_up_sent(respx_mock: respx.MockRouter) -> int:
+    follow_up_request: Final = json.loads(respx_mock.calls[1].request.content)
+    return litellm.token_counter(model="gpt-5.6", messages=follow_up_request["messages"])
+
+
 def _assert_one_row_bills_round_one(logged: Sequence[StandardLoggingPayload]) -> None:
     assert len(logged) == 1
     assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == (71, 19)
@@ -2045,6 +2050,7 @@ async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_r
         )
         - _FIXED_MARGIN
     )
+    assert failed["prompt_tokens"] == _prompt_tokens_the_follow_up_sent(respx_mock)
     assert failed["completion_tokens"] > 0
     assert failed["response_cost"] == pytest.approx(unmargined)
     request_failure: Final = _failure_fields_to_lift(request_data)
@@ -2174,12 +2180,15 @@ async def test_stream_auto_execute_whose_follow_up_cannot_be_assembled_ends_with
     assert usage.cost == pytest.approx(recorder.logged[0]["response_cost"])
 
 
-def _assert_one_row_bills_round_one_and_the_partial_answer(recorder: _ChatRoundRecorder) -> None:
+def _assert_one_row_bills_round_one_and_the_partial_answer(
+    recorder: _ChatRoundRecorder, respx_mock: respx.MockRouter
+) -> None:
     assert len(recorder.logged) == 1
     row: Final = recorder.logged[0]
     assert row["id"] == _CHAT_ROUND_ONE["id"]
     partial_prompt, partial_completion = (row["prompt_tokens"] - 71, row["completion_tokens"] - 19)
-    assert partial_prompt > 0 and partial_completion > 0
+    assert partial_prompt == _prompt_tokens_the_follow_up_sent(respx_mock)
+    assert partial_completion > 0
     partial_answer: Final = litellm.completion_cost(
         completion_response=ModelResponse(
             model="gpt-5.6",
@@ -2220,7 +2229,7 @@ async def test_stream_auto_execute_hung_up_during_the_follow_up_bills_round_one_
         )
     await recorder.settled(logged=1)
 
-    _assert_one_row_bills_round_one_and_the_partial_answer(recorder)
+    _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock)
 
 
 class _StallsOnceTheAnswerStarts(httpx.AsyncByteStream):
@@ -2310,7 +2319,7 @@ async def test_stream_auto_execute_routed_through_the_proxy_bills_the_partial_an
         await _serve_until_the_client_hangs_up_after(served, "PING")
     await recorder.settled(logged=1)
 
-    _assert_one_row_bills_round_one_and_the_partial_answer(recorder)
+    _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock)
 
 
 async def _read_through_the_usage_chunk(