fix(mcp): hold round one's log for the guardrail, bill the chunk cost once, and keep round one's id

When the follow-up call of a streamed chat MCP auto-execute fails, round one's
success log now goes through the proxy's deferred guardrail release instead of
straight to the handlers, so a post-call guardrail runs on it and the row
carries it. Without a proxy release it logs directly as before

The last usage chunk now takes its cost from the combined row's billed total,
so a fixed cost margin is counted once in the chunk, as it is on the row. The
follow-up's usage chunk is held until the follow-up ends to make that possible

The combined row now keeps round one's id, the stream's first chunk id, which
matches how a single-round stream is logged

litellm/proxy/common_request_processing.py

@@ -4053,11 +4053,13 @@ class ProxyBaseLLMRequestProcessing:
             _captured_user_api_key_dict: Final = user_api_key_dict
             _captured_logging_obj: Final = logging_obj
 
-            async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None:
+            async def _on_deferred_stream_complete(
+                assembled_response: object, cache_hit: object, logged_on: LiteLLMLoggingObj | None = None
+            ) -> None:
                 await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
                     captured_data=_captured_data,
                     captured_user_api_key_dict=_captured_user_api_key_dict,
-                    captured_logging_obj=_captured_logging_obj,
+                    captured_logging_obj=_captured_logging_obj if logged_on is None else logged_on,
                     assembled_response=assembled_response,
                     cache_hit=cache_hit,
                 )

litellm/responses/mcp/chat_completions_handler.py

@@ -367,8 +367,10 @@ async def acompletion_with_mcp(
                     try:
                         return await self._next_chunk()
                     except StopAsyncIteration:
-                        round_billing.finish()
-                        raise
+                        usage_chunk: Final = round_billing.finish()
+                        if usage_chunk is None:
+                            raise
+                        return usage_chunk
                     except Exception:
                         round_billing.fail()
                         raise
@@ -434,7 +436,6 @@ async def acompletion_with_mcp(
                         from litellm._logging import verbose_logger
 
                         verbose_logger.debug("Follow-up chunk yielded: %s", chunk)
-                        return chunk if self._round_billing is None else self._round_billing.client_chunk(chunk)
                     except StopAsyncIteration:
                         self.follow_up_exhausted = True
                         from litellm._logging import verbose_logger
@@ -442,6 +443,10 @@ async def acompletion_with_mcp(
                         verbose_logger.debug("Follow-up stream exhausted")
                         # After follow-up stream is exhausted, check if we need to raise StopAsyncIteration
                         raise StopAsyncIteration
+                    client_chunk: Final = (
+                        chunk if self._round_billing is None else self._round_billing.client_chunk(chunk)
+                    )
+                    return await self._next_chunk() if client_chunk is None else client_chunk
 
                 # If we're here and follow_up_stream is None but we expected it, log a warning
                 if (

litellm/responses/mcp/chat_stream_round_billing.py

@@ -1,5 +1,5 @@
 import asyncio
-from collections.abc import Generator
+from collections.abc import Callable, Coroutine, Generator
 from contextlib import contextmanager
 from dataclasses import dataclass, replace
 from types import MappingProxyType
@@ -14,6 +14,7 @@ from litellm.responses.streaming_iterator import spawn_logging_task
 from litellm.types.utils import ModelResponse, ModelResponseStream, Usage
 
 Parked: TypeAlias = tuple[object, ...]
+_Release: TypeAlias = Callable[..., Coroutine[object, object, None]]
 
 _PARKED_DISPATCH: Final = "_deferred_stream_complete_args"
 _DISPATCH_RELEASE: Final = "_on_deferred_stream_complete"
@@ -55,12 +56,31 @@ class _PricedRound:
         return BilledRoundsSpend(self.usage, self.price.cost or 0.0, self.price.cost_breakdown)
 
 
+@dataclass(frozen=True, slots=True)
+class _HeldUsageChunk:
+    chunk: ModelResponseStream
+    usage: Usage
+
+
 @dataclass(frozen=True, slots=True)
 class _FirstRound:
     priced: _PricedRound
     parked: Parked
     logged_alone: Parked
     banked: LiteLLMLoggingObj | None = None
+    usage_chunk: _HeldUsageChunk | None = None
+
+    def client_usage_chunk(self, billed: ModelResponse | None) -> ModelResponseStream | None:
+        if self.usage_chunk is None:
+            return None
+        merged: Final = merged_chat_usage(self.priced.usage, self.usage_chunk.usage)
+        usage: Final = merged if billed is None else merged.model_copy(update=MappingProxyType({"cost": _cost(billed)}))
+        return self.usage_chunk.chunk.model_copy(update=MappingProxyType({"usage": usage}))
+
+
+def _cost(response: ModelResponse) -> object:
+    hidden_params: Final = get_hidden_params(response)
+    return None if hidden_params is None else hidden_params.get("response_cost")
 
 
 def _assembled_round(parked: Parked | None) -> _AssembledRound | None:
@@ -100,27 +120,17 @@ class ChatStreamRoundBilling:
         self._bank_apart_from_the_follow_up()
         return None if self._first is None else self._first.priced.spend
 
-    def client_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream:
+    def client_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None:
         usage: Final[object] = getattr(chunk, "usage", None)
         if self._first is None or not isinstance(usage, Usage):
             return chunk
-        return chunk.model_copy(update=MappingProxyType({"usage": merged_chat_usage(self._first.priced.usage, usage)}))
+        self._first = replace(self._first, usage_chunk=_HeldUsageChunk(chunk, usage))
+        return None
 
-    def finish(self) -> None:
-        parked: Final = self._take_parked()
+    def finish(self) -> ModelResponseStream | None:
         first: Final = self._first
-        self._first = None
-        if parked is None:
-            self._disarm()
-            return
-        if first is None or parked is first.logged_alone:
-            self._release(parked)
-            return
-        final_round: Final = _assembled_round(parked)
-        if final_round is None:
-            self._release(first.logged_alone)
-            return
-        self._release(self._billed_for_both_rounds(first.priced, self._priced(final_round)))
+        billed: Final = self._release_rounds()
+        return None if first is None else first.client_usage_chunk(billed)
 
     def fail(self) -> None:
         self._bank_apart_from_the_follow_up()
@@ -130,7 +140,7 @@ class ChatStreamRoundBilling:
             return
         self._first = None
         setattr(self._logging_obj, _PARKED_DISPATCH, None)
-        self._log_apart((*first.parked, first.banked))
+        self._log_through_the_release((*first.parked, first.banked))
 
     def release_abandoned(self) -> None:
         if not self._owns_release():
@@ -154,12 +164,31 @@ class ChatStreamRoundBilling:
         self._first = replace(first, logged_alone=logged_alone, banked=banked)
         setattr(self._logging_obj, _PARKED_DISPATCH, logged_alone)
 
-    def _billed_for_both_rounds(self, first: _PricedRound, final: _PricedRound) -> Parked:
+    def _release_rounds(self) -> ModelResponse | None:
+        parked: Final = self._take_parked()
+        first: Final = self._first
+        self._first = None
+        if parked is None:
+            self._disarm()
+            return None
+        if first is None or parked is first.logged_alone:
+            self._release(parked)
+            return None
+        final_round: Final = _assembled_round(parked)
+        if final_round is None:
+            self._release(first.logged_alone)
+            return None
+        billed, cache_hit = self._billed_for_both_rounds(first.priced, self._priced(final_round))
+        self._release((billed, cache_hit))
+        return billed
+
+    def _billed_for_both_rounds(self, first: _PricedRound, final: _PricedRound) -> tuple[ModelResponse, bool]:
         first_breakdown: Final = first.price.cost_breakdown
         final_breakdown: Final = final.price.cost_breakdown
-        billed: Final = billed_for_every_round(first.response, first_breakdown, final.response, final_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}))
         cache_hit: Final = first.cache_hit and final.cache_hit
-        set_hidden_params(billed, {**(get_hidden_params(billed) or {}), "cache_hit": cache_hit})
+        set_hidden_params(billed, {**(get_hidden_params(rounds) or {}), "cache_hit": cache_hit})
         self._logging_obj.cost_breakdown = summed_cost_breakdown(
             first.response, first_breakdown, final.response, final_breakdown
         )
@@ -192,12 +221,13 @@ class ChatStreamRoundBilling:
         if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release():
             setattr(self._logging_obj, _PARKED_DISPATCH, parked)
             return
-        setattr(self._logging_obj, _DISPATCH_RELEASE, None)
-        self._log_apart(parked)
+        self._log_through_the_release(parked)
 
-    def _log_apart(self, parked: Parked) -> None:
-        loop: Final = asyncio.get_running_loop()
-        spawn_logging_task(loop, self._log_success(*parked), task_name="Chat completions MCP stream success log")
+    def _log_through_the_release(self, parked: Parked) -> None:
+        armed: Final[_Release | None] = getattr(self._logging_obj, _DISPATCH_RELEASE, None)
+        setattr(self._logging_obj, _DISPATCH_RELEASE, None)
+        logged: Final = (self._log_success if armed is None else armed)(*parked)
+        spawn_logging_task(asyncio.get_running_loop(), logged, task_name="Chat completions MCP stream success log")
 
     def _disarm(self) -> None:
         if self._owns_release():

tests/integration/_support/hook_count_probe.py

@@ -75,6 +75,24 @@ class BlockingPostCallGuardrail(CustomGuardrail):
         raise HTTPException(status_code=400, detail="blocked by the hook-count guardrail")
 
 
+class UsageRecordingPostCallGuardrail(CustomGuardrail):
+    """Records the token usage of the response it is handed in the request's guardrail information, so the spend row
+    shows which response the guardrail saw."""
+
+    async def async_post_call_success_hook(
+        self, data: dict[str, object], user_api_key_dict: object, response: object
+    ) -> None:
+        usage: Final = getattr(response, "usage", None)
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_json_response={
+                "prompt_tokens": getattr(usage, "prompt_tokens", None),
+                "completion_tokens": getattr(usage, "completion_tokens", None),
+            },
+            request_data=data,
+            guardrail_status="success",
+        )
+
+
 probe: Final = HookCountProbe()
 sync_only_probe: Final = SyncOnlyHookCountProbe()
 

tests/integration/mcp/test_mcp_llm_endpoints.py

@@ -45,9 +45,10 @@ from pydantic import JsonValue, TypeAdapter
 pytestmark: Final = [pytest.mark.xdist_group("mcp_llm_private_gateway"), pytest.mark.shares_worker_state]
 
 
-@pytest.fixture(scope="session")
-def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]:
-    directory: Final = tmp_path_factory.mktemp("mcp-llm-gateway")
+@contextmanager
+def _private_gateway(
+    directory: Path, guardrails: Sequence[Mapping[str, JsonValue]] = ()
+) -> Iterator[tuple[Gateway, Mapping[str, str]]]:
     stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
     config: Final = directory / "stored-responses.yaml"
     config.write_text(
@@ -59,6 +60,7 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat
                     **object_value(stock["litellm_settings"]),
                     "cost_margin_config": {"hosted_vllm": {"fixed_amount": FIXED_MARGIN}},
                 },
+                **({"guardrails": list(guardrails)} if guardrails else {}),
             }
         )
     )
@@ -70,15 +72,50 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat
         yield private, services
 
 
-@pytest.fixture
-def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]:
-    private, services = _llm_gateway
+@contextmanager
+def _serving(served: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]:
+    private, services = served
     with pytest.MonkeyPatch.context() as environment:
         for name, value in services.items():
             environment.setenv(name, value)
         yield private
 
 
+@pytest.fixture(scope="session")
+def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]:
+    with _private_gateway(tmp_path_factory.mktemp("mcp-llm-gateway")) as served:
+        yield served
+
+
+@pytest.fixture
+def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]:
+    with _serving(_llm_gateway) as private:
+        yield private
+
+
+ROUND_USAGE_GUARDRAIL: Final = "mcp-llm-round-usage"
+
+
+@pytest.fixture(scope="module")
+def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]:
+    guardrail: Final[Mapping[str, JsonValue]] = {
+        "guardrail_name": ROUND_USAGE_GUARDRAIL,
+        "litellm_params": {
+            "guardrail": "integration._support.hook_count_probe.UsageRecordingPostCallGuardrail",
+            "mode": "post_call",
+            "default_on": True,
+        },
+    }
+    with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), (guardrail,)) as served:
+        yield served
+
+
+@pytest.fixture
+def guarded_gateway(_guarded_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]:
+    with _serving(_guarded_llm_gateway) as private:
+        yield private
+
+
 Surface = Literal["chat", "responses", "messages", "messages_bridge"]
 SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge")
 ADD: Final = {"a": 2, "b": 3}
@@ -973,7 +1010,7 @@ def test_streaming_chat_auto_execute_logs_the_answer_and_bills_both_rounds(gatew
         rows: Final = _spend_rows(response.headers["x-litellm-call-id"], "acompletion", "success")
         assert len(rows) == 1, rows
         row: Final = rows[0]
-        assert row["request_id"] == chunks[-1]["id"], (row, chunks[-1])
+        assert row["request_id"] == chunks[0]["id"], (row, chunks[0])
         assert (row["prompt_tokens"], row["completion_tokens"]) == (
             FIRST_ROUND[0] + FOLLOW_UP_ROUND[0],
             FIRST_ROUND[1] + FOLLOW_UP_ROUND[1],
@@ -1001,6 +1038,40 @@ 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
 
 
+def _guardrail_information(metadata: JsonValue) -> tuple[JsonValue, ...]:
+    entries: Final = object_value(metadata).get("guardrail_information")
+    return tuple(entries) if isinstance(entries, list) else ()
+
+
+def _round_usage_seen_by_the_guardrail(call_id: str) -> list[JsonValue]:
+    rows: Final = read_rows(
+        'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s AND status=%s',
+        (call_id, "acompletion", "success"),
+    )
+    entries: Final = tuple(chain.from_iterable(_guardrail_information(row["metadata"]) for row in rows))
+    return [
+        entry.get("guardrail_response")
+        for entry in entries
+        if isinstance(entry, dict) and entry.get("guardrail_name") == ROUND_USAGE_GUARDRAIL
+    ]
+
+
+def test_streaming_chat_auto_execute_whose_follow_up_fails_runs_the_post_call_guardrail_on_round_one(
+    guarded_gateway: Gateway,
+) -> None:
+    billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True)
+    with _rig(guarded_gateway, "chat", billing=billing, prices=PRICES) as rig:
+        response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True})
+        assert FOLLOW_UP_OUTAGE in response.text, response.text
+        call_id: Final = response.headers["x-litellm-call-id"]
+        succeeded: Final = _spend_rows(call_id, "acompletion", "success")
+        assert len(succeeded) == 1, succeeded
+        assert (succeeded[0]["prompt_tokens"], succeeded[0]["completion_tokens"]) == FIRST_ROUND, succeeded
+        assert _round_usage_seen_by_the_guardrail(call_id) == [
+            {"prompt_tokens": FIRST_ROUND[0], "completion_tokens": FIRST_ROUND[1]}
+        ]
+
+
 @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

tests/unit/responses/mcp/test_chat_completions_handler.py

@@ -1823,13 +1823,14 @@ def _assert_one_stream_row_bills_both_rounds_once(
     charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN
     assert len(logged) == 1
     row: Final = logged[0]
-    assert row["id"] == chunks[-1].id
+    assert row["id"] == chunks[0].id
     assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5)
     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)
     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"])
 
 
 def _assert_one_row_bills_round_one(logged: Sequence[StandardLoggingPayload]) -> None: