fix(mcp): bill both calls of a streamed chat completions MCP auto-execute (#48899)

* fix(mcp): bill both rounds of a streamed chat completions MCP auto-execute

A streamed /v1/chat/completions request with an MCP tool set to require_approval "never" runs two model calls on one logging object, and the final-stream dedupe kept only the first round's spend row. The final usage chunk also showed only the second round.

The MCP stream iterator now holds the first round's assembled response, prices it, and logs one success row that carries the usage and cost of both rounds, with the fixed margin charged once. The final usage chunk shows the same summed usage. A single-round stream still logs its one row when the answer finishes.

When the follow-up round fails, the first round is logged from a copy of the logging object taken before the follow-up opened, and the failure row carries only the follow-up's partial usage with no second fixed margin. With proxy post-call guardrails armed, only (assembled, cache_hit) tuples are parked for the proxy, and the deferred log is armed only while the iterator reads, so the proxy's orphaned-callback cleanup never sees it after the Router's first-chunk read

* 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

* 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

* fix(mcp): bill every chat MCP stream round once on failure, hang-up, and sync drop

When the follow-up call fails, round 1's success row is now logged on a copy of the
logging object that does not share the request's guardrail results, so a post-call
guardrail that reports on the failed request counts once instead of on both rows

A client hang-up during the follow-up now bills round 1 plus the follow-up's partial
answer on one success row with round 1's id and the fixed margin counted once, the
same way a plain stream's hang-up bills its partial answer

The last usage chunk now carries round 1's usage and cost when only round 1 is billed,
and a synchronous reader that drops the stream mid-answer bills round 1 on the loop it
read with instead of discarding it

The deferred release's logged_on argument is typed and every release closure accepts
it, and the parked dispatch constants and protocols both MCP billing paths use now live
in one module. The unit tests wait on a logging signal instead of a fixed delay and
check the billed row through the proxy stream hook instead of private callback state

* test(mcp): drive the unassembled follow-up through an upstream audio delta that is not base64

The test replaced litellm.stream_chunk_builder with a stub that raises. Round 2 now streams an audio delta whose data is not valid base64, so the real builder raises, the stream parks only best-effort usage, and the client's last usage chunk must carry round 1's tokens and cost

* fix(mcp): count the fixed margin once in the combined stream row's usage cost

* fix(mcp): bill round 2's partial answer when a client hangs up on a routed stream

On the proxy the Router wraps the MCP stream, and when a client hangs up during
round 2 the router's own cleanup closes that stream before the proxy's disconnect
cleanup runs. The close logged round 1 alone, so round 2's partial answer was
never billed. The MCP stream now bills round 1 plus round 2's partial answer as
soon as the hang-up's cancellation reaches its read, before anything closes it

* 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

* fix(mcp): bill round 2's partial when a routed stream's client hangs up between reads

On a proxy the stream the disconnect cleanup holds is the router's fallback wrapper, so the
cleanup never found the MCP stream's in-flight billing and released round 1 alone when the client
hung up while the proxy sat at a yield. The round billing now registers its in-flight biller on
the request's logging object, next to the parked release, and the cleanup looks it up there, so
any wrapper chain bills round 1 plus round 2's partial once

* fix(mcp): send the MCP stream usage chunk under the billed row id and count the partial off the loop

A streamed chat completion that runs a LiteLLM-hosted MCP tool bills both model calls on one row under round 1's
chatcmpl id, but the usage chunk it sent the client was round 2's held chunk and kept round 2's id, so a caller could
not join the usage chunk to its row. The usage chunk now carries round 1's id in every path that sends it, with or
without a deferred post-call guardrail, while its usage and cost and every content chunk's id stay as they were

When the client hangs up during round 2, assembling round 2's partial answer counts the follow-up messages, tool
result included, with the local tokenizer. That ran on the event loop and stalled every other request on the worker
for about a second on a large tool result. The assembly now runs on a worker thread through asyncify from both the
reader's cancellation and the proxy's disconnect cleanup, shielded from cancellation, and round 1 is checked again
after the await so a release that happened meanwhile is never billed twice

* test(mcp): wait on the loop with a hang guard instead of a 300 ms deadline

* fix(mcp): give a proxy-rejected follow-up's failure row round 1's own cost breakdown

* refactor(streaming): fold the partial response's empty guards into one check

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm/litellm_core_utils/litellm_logging.py

@@ -568,6 +568,11 @@ def _resolve_mantle_region_for_cost(
 _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"})
 _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object])
 _REDACTED_RESULT: Final = TypeAdapter(object)
+_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information"
+
+
+def _without_guardrail_information(metadata: Mapping[str, object]) -> Mapping[str, object]:
+    return {key: value for key, value in metadata.items() if key != _GUARDRAIL_INFORMATION_KEY}
 
 
 def _provider_response_id(source: object) -> str | None:
@@ -2566,6 +2571,23 @@ class Logging(LiteLLMLoggingBaseClass):
         attempt.model_call_details = dict(self.model_call_details)
         return attempt
 
+    def guardrail_scoped_copy(self) -> "Logging":
+        request_params: Final = _LITELLM_PARAMS.validate_python(
+            self.litellm_params or {}  # pyright: ignore[reportUnknownMemberType]  # Logging.litellm_params is untyped
+        )
+        litellm_params: Final = {
+            **request_params,
+            **{
+                key: _without_guardrail_information(_METADATA_MAPPING.validate_python(value))
+                for key, value in request_params.items()
+                if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict)
+            },
+        }
+        scoped: Final = copy.copy(self)
+        scoped.litellm_params = litellm_params
+        scoped.model_call_details = {**self.model_call_details, "litellm_params": litellm_params}
+        return scoped
+
     def dispatch_attempt_failure_handlers(
         self, exception: Exception, traceback_exception: str
     ) -> Coroutine[None, None, None]:

litellm/litellm_core_utils/streaming_handler.py

@@ -2486,6 +2486,25 @@ class CustomStreamWrapper:
     def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None:
         self._earlier_rounds_spend = earlier
 
+    def partial_response(self) -> "ModelResponse | TextCompletionResponse | None":
+        partial: Final = (
+            litellm.stream_chunk_builder(
+                chunks=self.chunks,
+                messages=self.messages if isinstance(self.messages, list) else None,
+                logging_obj=self.logging_obj,
+                count_prompt_tokens=self.count_prompt_tokens,
+            )
+            if self.chunks
+            else None
+        )
+        usage: Final[object] = getattr(partial, "usage", None)
+        if partial is None or not isinstance(usage, Usage):
+            return None
+        if self.model:
+            partial.model = self.model
+        backfill_missing_cache_usage_fields(usage)
+        return partial
+
     def _record_partial_usage_for_failure(self) -> None:
         """
         A stream that breaks mid-flight still billed the provider for the chunks
@@ -2495,23 +2514,15 @@ class CustomStreamWrapper:
         moves it into terminal failure usage before opening the next attempt,
         so a recovered attempt logs its own usage.
         """
-        if self.logging_obj is None or not self.chunks:
+        if self.logging_obj is None:
             return
         try:
-            partial_response: Final = litellm.stream_chunk_builder(
-                chunks=self.chunks,
-                messages=self.messages if isinstance(self.messages, list) else None,
-                logging_obj=self.logging_obj,
-                count_prompt_tokens=self.count_prompt_tokens,
-            )
+            partial_response: Final = self.partial_response()
             if partial_response is None:
                 return
-            usage: Final = cast(Usage | None, getattr(partial_response, "usage", None))
-            if usage is None:
+            usage: Final[object] = getattr(partial_response, "usage", None)
+            if not isinstance(usage, Usage):
                 return
-            if self.model:
-                partial_response.model = self.model
-            backfill_missing_cache_usage_fields(usage)
             partial_cost: Final = self.logging_obj.response_cost_calculator(result=partial_response) or 0.0
             if self._earlier_rounds_spend is None:
                 self.logging_obj.model_call_details["combined_usage_object"] = usage

litellm/proxy/common_request_processing.py

@@ -140,6 +140,7 @@ from litellm.proxy.utils import (  # noqa: F401  # legacy module exports
     _check_and_merge_model_level_guardrails,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     check_and_merge_model_level_guardrails,
 )
+from litellm.responses.mcp.parked_stream_dispatch import in_flight_round_biller, parked_success
 from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task
 from litellm.router import Router
 from litellm.router_api import RoutingContext
@@ -535,6 +536,16 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
     return await _dispatch_disconnect_billing_event(logging_obj, partial_response)
 
 
+async def _bill_in_flight_round_on_disconnect(request_data: Mapping[str, object]) -> None:
+    bill_in_flight_round: Final = in_flight_round_biller(request_data.get("litellm_logging_obj"))
+    if bill_in_flight_round is None:
+        return
+    try:
+        await bill_in_flight_round()
+    except Exception as e:  # noqa: BLE001  # disconnect billing is best-effort; never break stream teardown
+        verbose_proxy_logger.debug("Failed to bill the in-flight round after client disconnect: %s", e)
+
+
 async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, partial_response: object) -> bool:
     try:
         await logging_obj.dispatch_success_handlers(
@@ -2083,6 +2094,13 @@ class _ClientStreamFailureReport:
             )
 
 
+def _metadata_logged_on(logged_on: LiteLLMLoggingObj) -> Mapping[str, object]:
+    logged_params: Final = object_mapping(logged_on.model_call_details.get("litellm_params")) or MappingProxyType({})
+    return MappingProxyType(
+        {key: logged_params[key] for key in ("metadata", "litellm_metadata") if key in logged_params}
+    )
+
+
 class ProxyBaseLLMRequestProcessing:
     def __init__(self, data: dict, served_model: str | None = None):
         self.data = data
@@ -4151,11 +4169,21 @@ 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_data=_captured_data
+                    if logged_on is None
+                    else {
+                        **_captured_data,
+                        **{
+                            key: value for key, value in _metadata_logged_on(logged_on).items() if key in _captured_data
+                        },
+                        "litellm_logging_obj": logged_on,
+                    },
                     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,
                 )
@@ -4180,14 +4208,8 @@ class ProxyBaseLLMRequestProcessing:
                     GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine)
                 case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine):
                     GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine)
-                case (assembled_response, cache_hit):
-                    await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers(
-                        assembled_response,
-                        cache_hit=cache_hit,
-                        start_time=None,
-                        end_time=None,
-                        prefer_async_handlers=True,
-                    )
+                case _ if (success := parked_success(args)) is not None:
+                    await success.dispatch(_captured_native_logging_obj)
                 case _:
                     verbose_proxy_logger.error(
                         "Deferred stream logging dropped: unexpected stored args shape %s",
@@ -4524,6 +4546,7 @@ class ProxyBaseLLMRequestProcessing:
                     client_disconnected,
                 )
             if recorded_client_disconnect:
+                await _bill_in_flight_round_on_disconnect(_REQUEST_FIELDS.validate_python(request_data))
                 deferred_stream_logging_armed: Final = _deferred_stream_logging_is_armed(request_data)
                 ProxyLogging.fire_deferred_stream_logging(request_data)
                 # A disconnect-time success event (the deferred-guardrail flush

litellm/proxy/utils.py

@@ -91,6 +91,7 @@ from litellm.proxy.common_utils.openai_error_payload import (
 )
 from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
 from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction
+from litellm.responses.mcp.parked_stream_dispatch import held_round_spend
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.proxy.model_listing import ModelInfoResponse
 from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, ModelInfo, Usage
@@ -4943,6 +4944,8 @@ class ProxyLogging:
         if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS):
             return False
         match _args:
+            case (ModelResponse(), *_) if (held := held_round_spend(logging_obj)) is not None:
+                logging_obj.record_billed_rounds_for_failure(held)
             case (ModelResponse() as assembled, *_):
                 logging_obj.record_assembled_response_for_failure(assembled)
             case (BilledRoundsDispatch() as billed,):

litellm/responses/mcp/chat_completions_handler.py

@@ -1,15 +1,18 @@
 """Helpers for handling MCP-aware `/chat/completions` requests."""
 
+import asyncio
 import logging
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, cast
 
+import anyio
 from pydantic import TypeAdapter
 from typing_extensions import TypedDict, Unpack
 
 from litellm._internal_context import internal_sub_call
 from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.responses.mcp.chat_stream_round_billing import ChatStreamRoundBilling, engaged_round_billing
 from litellm.responses.mcp.litellm_proxy_mcp_handler import (
     LiteLLM_Proxy_MCP_Handler,
 )
@@ -19,7 +22,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:
@@ -258,10 +261,20 @@ async def acompletion_with_mcp(
                 self.follow_up_stream = None
                 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
 
+            def __del__(self) -> None:
+                self.release_abandoned()
+
+            def release_abandoned(self) -> None:
+                round_billing: Final[object] = getattr(self, "_round_billing", None)
+                if isinstance(round_billing, ChatStreamRoundBilling):
+                    round_billing.release_abandoned()
+
             def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream:
                 """Add mcp_list_tools to the first chunk."""
                 from litellm.types.utils import (
@@ -321,18 +334,54 @@ 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 __anext__(self):
+            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, 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 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:
+                round_billing: Final = self._round_billing
+                if round_billing is None:
+                    return await self._next_chunk()
+                with round_billing.armed_while_reading():
+                    try:
+                        return await self._next_chunk()
+                    except StopAsyncIteration:
+                        usage_chunk: Final = round_billing.finish()
+                        if usage_chunk is None:
+                            raise
+                        return usage_chunk
+                    except asyncio.CancelledError:
+                        with anyio.CancelScope(shield=True):
+                            await round_billing.bill_in_flight_round()
+                        raise
+                    except Exception:
+                        round_billing.fail()
+                        raise
+
+            async def _next_chunk(self) -> ModelResponseStream:
                 # Phase 1: Collect and yield initial stream chunks
                 if not self.stream_exhausted:
                     # Get the iterator from the stream wrapper
@@ -362,29 +411,22 @@ async def acompletion_with_mcp(
 
                         if is_final:
                             self.stream_exhausted = True
+                            usage_chunk: Final = await self._complete_first_round()
                             await self._process_tool_calls()
                             chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk)
-                            if self.tool_results and self.complete_response:
-                                await self._prepare_follow_up_call()
-                            # Drain inner stream so CustomStreamWrapper fires its
-                            # end-of-stream handler (dispatch_success_handlers →
-                            # _ProxyDBLogger → LiteLLM_SpendLogs).  The CSW may
-                            # yield one usage chunk before raising StopAsyncIteration.
-                            await self._drain_inner_stream()
+                            await self._follow_up_or_release_first_round(usage_chunk)
 
                         return chunk
                     except StopAsyncIteration:
                         self.stream_exhausted = True
+                        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)
-                            # If we have tool results, prepare follow-up call
-                            if self.tool_results and self.complete_response:
-                                await self._prepare_follow_up_call()
-                            await self._drain_inner_stream()
+                            await self._follow_up_or_release_first_round(drained_usage_chunk)
                             return final_chunk
 
                 # Phase 2: Yield follow-up stream chunks if available
@@ -400,7 +442,6 @@ async def acompletion_with_mcp(
                         from litellm._logging import verbose_logger
 
                         verbose_logger.debug("Follow-up chunk yielded: %s", chunk)
-                        return chunk
                     except StopAsyncIteration:
                         self.follow_up_exhausted = True
                         from litellm._logging import verbose_logger
@@ -408,6 +449,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 (
@@ -420,7 +465,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."""
@@ -488,10 +537,16 @@ async def acompletion_with_mcp(
                 # This ensures the patch works correctly in tests
                 import litellm
 
+                earlier_rounds: Final = None if self._round_billing is None else self._round_billing.bank_first_round()
                 follow_up_response: Final = await litellm.acompletion(**follow_up_call_args)
 
                 # Ensure follow-up response is a CustomStreamWrapper
                 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
+                    if self._round_billing is not None:
+                        self._round_billing.follow_up_started(follow_up_response.partial_response)
                     self.follow_up_stream = follow_up_response
                     from litellm._logging import verbose_logger
 
@@ -548,11 +603,13 @@ async def acompletion_with_mcp(
             def __aiter__(self):
                 return self._custom_iterator
 
+            async def aclose(self) -> None:
+                self._custom_iterator.release_abandoned()
+                await super().aclose()
+
             def __iter__(self):
                 # For synchronous iteration, create a sync wrapper
                 if self._sync_iterator is None:
-                    import asyncio
-
                     try:
                         self._sync_loop = asyncio.get_event_loop()
                     except RuntimeError:

litellm/responses/mcp/chat_stream_round_billing.py

@@ -0,0 +1,281 @@
+import asyncio
+from collections.abc import Callable, Generator
+from contextlib import contextmanager
+from dataclasses import dataclass, replace
+from types import MappingProxyType
+from typing import Final, TypeAlias
+
+from litellm.litellm_core_utils.asyncify import asyncify, is_event_loop_running
+from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params
+from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.responses.mcp.parked_stream_dispatch import (
+    DISPATCH_RELEASE,
+    HELD_ROUND_SPEND,
+    IN_FLIGHT_ROUND,
+    PARKED_DISPATCH,
+    SERVED_FROM_CACHE,
+    DeferredStreamRelease,
+    parked_success,
+)
+from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown
+from litellm.responses.streaming_iterator import spawn_logging_task
+from litellm.types.utils import ModelResponse, ModelResponseStream, Usage
+
+Parked: TypeAlias = tuple[object, ...]
+PartialAnswer: TypeAlias = Callable[[], object]
+
+
+@dataclass(frozen=True, slots=True)
+class _AssembledRound:
+    response: ModelResponse
+    usage: Usage
+    cache_hit: bool
+    parked: Parked
+
+
+@dataclass(frozen=True, slots=True)
+class _PricedRound:
+    response: ModelResponse
+    usage: Usage
+    price: UncachedPrice
+    cache_hit: bool
+
+    @property
+    def spend(self) -> BilledRoundsSpend:
+        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
+        usage: Final = (
+            self.priced.usage.model_copy(update=MappingProxyType({"cost": self.priced.price.cost}))
+            if billed is None
+            else merged_chat_usage(self.priced.usage, self.usage_chunk.usage).model_copy(
+                update=MappingProxyType({"cost": _cost(billed)})
+            )
+        )
+        return self.usage_chunk.chunk.model_copy(
+            update=MappingProxyType({"id": self.priced.response.id, "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:
+    match parked:
+        case (ModelResponse() as response, cache_hit) if isinstance(usage := getattr(response, "usage", None), Usage):
+            return _AssembledRound(response, usage, cache_hit is True, parked)
+        case _:
+            return None
+
+
+def engaged_round_billing(logging_obj: object) -> "ChatStreamRoundBilling | None":
+    return ChatStreamRoundBilling(logging_obj) if isinstance(logging_obj, LiteLLMLoggingObj) else None
+
+
+class ChatStreamRoundBilling:
+    def __init__(self, logging_obj: LiteLLMLoggingObj) -> None:
+        self._logging_obj: Final = logging_obj
+        self._first: _FirstRound | None = None
+        self._partial_follow_up: PartialAnswer | None = None
+        self._reading_loop: asyncio.AbstractEventLoop | None = None
+
+    @contextmanager
+    def armed_while_reading(self) -> Generator[None]:
+        armed: Final = getattr(self._logging_obj, DISPATCH_RELEASE, None) is None
+        if armed:
+            setattr(self._logging_obj, DISPATCH_RELEASE, self._log_success)
+        self._reading_loop = asyncio.get_running_loop()
+        try:
+            yield
+        finally:
+            if armed and self._peek_parked() is None:
+                self._disarm()
+
+    def complete_first_round(self) -> None:
+        assembled: Final = _assembled_round(self._peek_parked())
+        if assembled is not None:
+            self._first = _FirstRound(self._priced(assembled), assembled.parked, assembled.parked)
+            setattr(self._logging_obj, HELD_ROUND_SPEND, self._held_round_one_spend)
+
+    def bank_first_round(self) -> BilledRoundsSpend | None:
+        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 | None:
+        usage: Final[object] = getattr(chunk, "usage", None)
+        if self._first is None or not isinstance(usage, Usage):
+            return chunk
+        self._first = replace(self._first, usage_chunk=_HeldUsageChunk(chunk, usage))
+        return None
+
+    def finish(self) -> ModelResponseStream | None:
+        first: Final = self._first
+        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()
+        first: Final = self._first
+        if first is None or first.banked is None or not self._holds_round_one_alone(first):
+            self.finish()
+            return
+        self._first = None
+        setattr(self._logging_obj, PARKED_DISPATCH, None)
+        self._log_through_the_release((*first.parked, first.banked.guardrail_scoped_copy()))
+
+    def follow_up_started(self, partial_follow_up: PartialAnswer) -> None:
+        self._partial_follow_up = partial_follow_up
+        setattr(self._logging_obj, IN_FLIGHT_ROUND, self.bill_in_flight_round)
+
+    async def bill_in_flight_round(self) -> None:
+        partial_follow_up: Final = self._partial_follow_up
+        if partial_follow_up is None or self._first_round_held_alone() is None:
+            return
+        final_round: Final = _assembled_round((await asyncify(partial_follow_up)(), False))
+        first: Final = self._first_round_held_alone()
+        if final_round is None or first is None:
+            return
+        self._first = None
+        setattr(
+            self._logging_obj, PARKED_DISPATCH, self._billed_for_both_rounds(first.priced, self._priced(final_round))
+        )
+
+    def release_abandoned(self) -> None:
+        if not self._owns_release():
+            return
+        if is_event_loop_running() or self._idle_reading_loop() is not None:
+            self.finish()
+            return
+        setattr(self._logging_obj, PARKED_DISPATCH, None)
+        self._first = None
+        self._disarm()
+
+    def _bank_apart_from_the_follow_up(self) -> None:
+        first: Final = self._first
+        if first is None or first.banked is not None or self._peek_parked() is not first.parked:
+            return
+        banked: Final = self._logging_obj.attempt_scoped_copy()
+        if not self._owns_release():
+            self._first = replace(first, banked=banked)
+            return
+        logged_alone: Final = (*first.parked, banked)
+        self._first = replace(first, logged_alone=logged_alone, banked=banked)
+        setattr(self._logging_obj, PARKED_DISPATCH, logged_alone)
+
+    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
+        rounds: Final = billed_for_every_round(first.response, first_breakdown, final.response, final_breakdown)
+        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(
+            first.response, first_breakdown, final.response, final_breakdown
+        )
+        return billed, cache_hit
+
+    def _priced(self, assembled: _AssembledRound) -> _PricedRound:
+        served: Final = assembled.cache_hit or served_from_cache(assembled.response)
+        price: Final = SERVED_FROM_CACHE if served else self._logging_obj.uncached_price(assembled.response)
+        priced: Final = assembled.response.model_copy()
+        set_hidden_params(
+            priced,
+            {**(get_hidden_params(assembled.response) or {}), "response_cost": price.cost, "cache_hit": served},
+        )
+        return _PricedRound(priced, assembled.usage, price, served)
+
+    def _peek_parked(self) -> Parked | None:
+        parked: Final[Parked | None] = getattr(self._logging_obj, PARKED_DISPATCH, None)
+        return parked if isinstance(parked, tuple) else None
+
+    def _take_parked(self) -> Parked | None:
+        parked: Final = self._peek_parked()
+        setattr(self._logging_obj, PARKED_DISPATCH, None)
+        return parked
+
+    def _first_round_held_alone(self) -> _FirstRound | None:
+        first: Final = self._first
+        return first if first is not None and self._holds_round_one_alone(first) else None
+
+    def _held_round_one_spend(self) -> BilledRoundsSpend | None:
+        first: Final = self._first_round_held_alone()
+        return None if first is None else first.priced.spend
+
+    def _holds_round_one_alone(self, first: _FirstRound) -> bool:
+        parked: Final = self._peek_parked()
+        return parked is first.parked or parked is first.logged_alone
+
+    def _owns_release(self) -> bool:
+        release: Final[object] = getattr(self._logging_obj, DISPATCH_RELEASE, None)
+        return release == self._log_success
+
+    def _release(self, parked: Parked) -> None:
+        if getattr(self._logging_obj, DISPATCH_RELEASE, None) is not None and not self._owns_release():
+            setattr(self._logging_obj, PARKED_DISPATCH, parked)
+            return
+        self._log_through_the_release(parked)
+
+    def _log_through_the_release(self, parked: Parked) -> None:
+        armed: Final[DeferredStreamRelease | 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)
+        idle_reading_loop: Final = None if is_event_loop_running() else self._idle_reading_loop()
+        if idle_reading_loop is not None:
+            idle_reading_loop.run_until_complete(logged)
+            return
+        spawn_logging_task(asyncio.get_running_loop(), logged, task_name="Chat completions MCP stream success log")
+
+    def _idle_reading_loop(self) -> asyncio.AbstractEventLoop | None:
+        loop: Final = self._reading_loop
+        return loop if loop is not None and not loop.is_closed() and not loop.is_running() else None
+
+    def _disarm(self) -> None:
+        if self._owns_release():
+            setattr(self._logging_obj, DISPATCH_RELEASE, None)
+
+    async def _log_success(self, *parked: object) -> None:
+        success: Final = parked_success(parked)
+        if success is not None:
+            await success.dispatch(self._logging_obj)

litellm/responses/mcp/parked_stream_dispatch.py

@@ -0,0 +1,72 @@
+from collections.abc import Coroutine
+from dataclasses import dataclass
+from typing import Final, Protocol
+
+from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+
+PARKED_DISPATCH: Final = "_deferred_stream_complete_args"
+DISPATCH_RELEASE: Final = "_on_deferred_stream_complete"
+IN_FLIGHT_ROUND: Final = "_bill_in_flight_round_on_disconnect"
+HELD_ROUND_SPEND: Final = "_held_round_spend_for_failure"
+SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None)
+
+
+class DispatchesSuccessHandlers(Protocol):
+    async def dispatch_success_handlers(
+        self,
+        result: object = None,
+        start_time: object = None,
+        end_time: object = None,
+        cache_hit: object = None,
+        prefer_async_handlers: bool = False,
+    ) -> None: ...
+
+
+class DeferredStreamRelease(Protocol):
+    def __call__(self, *parked: object) -> Coroutine[object, object, None]: ...
+
+
+class InFlightRoundBiller(Protocol):
+    def __call__(self) -> Coroutine[object, object, None]: ...
+
+
+class HeldRoundSpend(Protocol):
+    def __call__(self) -> BilledRoundsSpend | None: ...
+
+
+def success_dispatcher(logger: LiteLLMLoggingObj) -> DispatchesSuccessHandlers:
+    return logger
+
+
+@dataclass(frozen=True, slots=True)
+class ParkedSuccess:
+    assembled: object
+    cache_hit: object
+    logged_on: LiteLLMLoggingObj | None = None
+
+    async def dispatch(self, request_logger: LiteLLMLoggingObj) -> None:
+        logger: Final = request_logger if self.logged_on is None else self.logged_on
+        await success_dispatcher(logger).dispatch_success_handlers(
+            self.assembled, cache_hit=self.cache_hit, start_time=None, end_time=None, prefer_async_handlers=True
+        )
+
+
+def in_flight_round_biller(logging_obj: object) -> InFlightRoundBiller | None:
+    biller: Final[InFlightRoundBiller | None] = getattr(logging_obj, IN_FLIGHT_ROUND, None)
+    return biller if isinstance(logging_obj, LiteLLMLoggingObj) and callable(biller) else None
+
+
+def held_round_spend(logging_obj: object) -> BilledRoundsSpend | None:
+    held: Final[HeldRoundSpend | None] = getattr(logging_obj, HELD_ROUND_SPEND, None)
+    return held() if isinstance(logging_obj, LiteLLMLoggingObj) and callable(held) else None
+
+
+def parked_success(parked: tuple[object, ...]) -> ParkedSuccess | None:
+    match parked:
+        case (assembled, cache_hit):
+            return ParkedSuccess(assembled, cache_hit)
+        case (assembled, cache_hit, LiteLLMLoggingObj() as logged_on):
+            return ParkedSuccess(assembled, cache_hit, logged_on)
+        case _:
+            return None

litellm/responses/mcp/stream_round_billing.py

@@ -2,11 +2,18 @@ import asyncio
 from collections.abc import Callable, Coroutine, Sequence
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Final, Protocol, TypeAlias
+from typing import Final, TypeAlias
 
 from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache
 from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.responses.mcp.parked_stream_dispatch import (
+    DISPATCH_RELEASE,
+    PARKED_DISPATCH,
+    SERVED_FROM_CACHE,
+    parked_success,
+    success_dispatcher,
+)
 from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown
 from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, spawn_logging_task
 from litellm.responses.utils import ResponseAPILoggingUtils
@@ -16,25 +23,6 @@ from litellm.types.utils import CostBreakdown, ModelResponse, Usage
 Parked: TypeAlias = tuple[object, ...]
 RestoresHeaders: TypeAlias = Callable[[ResponsesAPIStreamingResponse], None]
 
-_PARKED_DISPATCH: Final = "_deferred_stream_complete_args"
-_DISPATCH_RELEASE: Final = "_on_deferred_stream_complete"
-_SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None)
-
-
-class _DispatchesSuccessHandlers(Protocol):
-    async def dispatch_success_handlers(
-        self,
-        result: object = None,
-        start_time: object = None,
-        end_time: object = None,
-        cache_hit: object = None,
-        prefer_async_handlers: bool = False,
-    ) -> None: ...
-
-
-def _success_dispatcher(logger: LiteLLMLoggingObj) -> _DispatchesSuccessHandlers:
-    return logger
-
 
 def _event_loop_running() -> bool:
     try:
@@ -198,8 +186,8 @@ class StreamedRoundBilling:
         self._billed: _BilledRounds | None = None
         self._banked: Parked | None = None
         self._banked_row: _RoundRow | None = None
-        if getattr(logging_obj, _DISPATCH_RELEASE, None) is None:
-            setattr(logging_obj, _DISPATCH_RELEASE, self._run_parked)
+        if getattr(logging_obj, DISPATCH_RELEASE, None) is None:
+            setattr(logging_obj, DISPATCH_RELEASE, self._run_parked)
 
     def complete_round(
         self,
@@ -286,7 +274,7 @@ class StreamedRoundBilling:
 
     def _price(self, response: ResponsesAPIResponse, round_parked: Parked | None) -> UncachedPrice:
         if served_from_cache(response):
-            return _SERVED_FROM_CACHE
+            return SERVED_FROM_CACHE
         match round_parked:
             case (ModelResponse() as assembled, bool()):
                 return self._logging_obj.uncached_price(assembled)
@@ -297,7 +285,7 @@ class StreamedRoundBilling:
         self._banked = parked
         self._banked_row = row
         if parked is not None:
-            setattr(self._logging_obj, _PARKED_DISPATCH, parked)
+            setattr(self._logging_obj, PARKED_DISPATCH, parked)
 
     def _bank_apart_from_later_rounds(self, billed: _BilledRounds) -> None:
         row: Final = self._banked_row
@@ -320,7 +308,7 @@ class StreamedRoundBilling:
     @staticmethod
     async def _log_success(success: _SuccessLog) -> None:
         success.carry_bridged_usage()
-        await _success_dispatcher(success.logger).dispatch_success_handlers(
+        await success_dispatcher(success.logger).dispatch_success_handlers(
             success.row.logged,
             start_time=success.logger.start_time,
             end_time=success.row.end_time,
@@ -329,10 +317,10 @@ class StreamedRoundBilling:
         )
 
     def _take_parked(self) -> Parked | None:
-        parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None)
+        parked: Final[Parked | None] = getattr(self._logging_obj, PARKED_DISPATCH, None)
         if not isinstance(parked, tuple):
             return None
-        setattr(self._logging_obj, _PARKED_DISPATCH, None)
+        setattr(self._logging_obj, PARKED_DISPATCH, None)
         return parked
 
     async def _run_parked(self, *parked: object) -> None:
@@ -341,29 +329,27 @@ class StreamedRoundBilling:
                 await dispatch
             case (dispatch,) if asyncio.iscoroutine(dispatch):
                 await dispatch
-            case (assembled, cache_hit):
-                await _success_dispatcher(self._logging_obj).dispatch_success_handlers(
-                    assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True
-                )
+            case _ if (success := parked_success(parked)) is not None:
+                await success.dispatch(self._logging_obj)
             case _:
                 return
 
     def _owns_release(self) -> bool:
-        release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None)
+        release: Final[object] = getattr(self._logging_obj, DISPATCH_RELEASE, None)
         return release == self._run_parked
 
     def _release(self, parked: Parked) -> None:
-        if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release():
-            setattr(self._logging_obj, _PARKED_DISPATCH, parked)
+        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)
+        setattr(self._logging_obj, DISPATCH_RELEASE, None)
         spawn_logging_task(
             asyncio.get_running_loop(), self._run_parked(*parked), task_name="Responses MCP stream success logging"
         )
 
     def _disarm(self) -> None:
         if self._owns_release():
-            setattr(self._logging_obj, _DISPATCH_RELEASE, None)
+            setattr(self._logging_obj, DISPATCH_RELEASE, None)
 
     @staticmethod
     def _close(parked: Parked | None) -> None:

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

@@ -46,9 +46,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"
     (directory / f"{STREAM_BLOCK_MODULE}.py").write_text(STREAM_BLOCK_SOURCE)
@@ -68,7 +69,8 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat
                     {
                         "guardrail_name": STREAM_BLOCK,
                         "litellm_params": {"guardrail": f"{STREAM_BLOCK_MODULE}.BlockGuard", "mode": "post_call"},
-                    }
+                    },
+                    *guardrails,
                 ],
             }
         )
@@ -81,15 +83,61 @@ 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"
+CONTENT_FILTER: Final = f"mcp-llm-content-filter-{uuid.uuid4().hex[:8]}"
+
+
+@pytest.fixture(scope="module")
+def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]:
+    round_usage: Final[Mapping[str, JsonValue]] = {
+        "guardrail_name": ROUND_USAGE_GUARDRAIL,
+        "litellm_params": {
+            "guardrail": "integration._support.hook_count_probe.UsageRecordingPostCallGuardrail",
+            "mode": "post_call",
+            "default_on": True,
+        },
+    }
+    content_filter: Final[Mapping[str, JsonValue]] = {
+        "guardrail_name": CONTENT_FILTER,
+        "litellm_params": {
+            "guardrail": "litellm_content_filter",
+            "mode": "post_call",
+            "default_on": False,
+            "blocked_words": [{"keyword": "forbidden-keeper", "action": "BLOCK"}],
+        },
+    }
+    guardrails: Final = (round_usage, content_filter)
+    with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), guardrails) 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", "responses_bridge"]
 SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge")
 ADD: Final = {"a": 2, "b": 3}
@@ -1032,6 +1080,184 @@ def test_nonstream_failed_follow_up_bills_round_one_on_the_failure_row(
         )
 
 
+def _streamed_chunks(response: httpx.Response) -> tuple[dict[str, object], ...]:
+    return tuple(
+        json.loads(line.removeprefix("data: "))
+        for line in response.text.splitlines()
+        if line.startswith("data: ") and line != "data: [DONE]"
+    )
+
+
+def test_streaming_chat_auto_execute_logs_the_answer_and_bills_both_rounds(gateway: Gateway) -> 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], stream=True, stream_options={"include_usage": True})
+        assert response.status_code == 200, 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
+        usages: Final = tuple(chunk["usage"] for chunk in chunks if chunk.get("usage"))
+        assert usages, response.text
+        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[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],
+        ), row
+        assert float(str(row["spend"])) == pytest.approx(_price(FIRST_ROUND, FOLLOW_UP_ROUND)), row
+        final_usage: Final = usages[-1]
+        assert (final_usage["prompt_tokens"], final_usage["completion_tokens"], final_usage["total_tokens"]) == (
+            row["prompt_tokens"],
+            row["completion_tokens"],
+            sum((*FIRST_ROUND, *FOLLOW_UP_ROUND)),
+        ), usages
+
+
+def test_streaming_chat_auto_execute_whose_follow_up_fails_bills_round_one_once(gateway: Gateway) -> None:
+    billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True)
+    with _rig(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 float(str(succeeded[0]["spend"])) == pytest.approx(_price(FIRST_ROUND)), succeeded
+        failed: Final = _spend_rows(call_id, "acompletion", "failure")
+        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]}
+        ]
+
+
+def _chat_rows(call_id: str, statuses: Sequence[str]) -> list[dict[str, JsonValue]]:
+    return eventually(
+        lambda: read_rows(
+            'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s',
+            (call_id, "acompletion"),
+        ),
+        lambda found: _statuses(found) == sorted(statuses),
+        seconds=70,
+    )
+
+
+def _entries_of(guardrail_name: str, row: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]:
+    return tuple(
+        entry
+        for entry in _guardrail_information(row["metadata"])
+        if isinstance(entry, dict) and entry.get("guardrail_name") == guardrail_name
+    )
+
+
+def test_streaming_chat_auto_execute_whose_follow_up_fails_counts_its_post_call_content_filter_once(
+    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},
+            guardrails=[CONTENT_FILTER],
+        )
+        assert FOLLOW_UP_OUTAGE in response.text, response.text
+        rows: Final = _chat_rows(response.headers["x-litellm-call-id"], ("success", "failure"))
+        filtered: Final = tuple(
+            (row["status"], entry.get("guardrail_status")) for row in rows for entry in _entries_of(CONTENT_FILTER, row)
+        )
+        assert filtered == (("failure", "guardrail_failed_to_respond"),), rows
+        (guardrail_id,) = frozenset(
+            str(entry.get("guardrail_id") or CONTENT_FILTER)
+            for row in rows
+            for entry in _entries_of(CONTENT_FILTER, row)
+        )
+        counted: Final = eventually(
+            lambda: read_rows(
+                'SELECT requests_evaluated, flagged_count FROM "LiteLLM_DailyGuardrailMetrics" WHERE guardrail_id=%s',
+                (guardrail_id,),
+            ),
+            lambda found: bool(found),
+            seconds=70,
+        )
+        assert counted == [{"requests_evaluated": 1, "flagged_count": 1}], counted
+
+
+@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
+
+
 NO_CACHE: Final = {"no-cache": True}
 
 
@@ -1269,6 +1495,27 @@ def test_streaming_responses_follow_up_broken_mid_stream_bills_round_one_once_an
         assert float(str(failure["spend"])) == pytest.approx(_price(failed_round)), rows
 
 
+def test_streaming_chat_follow_up_broken_mid_answer_bills_round_one_and_its_partial_answer_with_one_fixed_margin(
+    gateway: Gateway,
+) -> None:
+    billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_breaks=True)
+    with _rig(gateway, "chat", upstream="messages_bridge", billing=billing, prices=PRICES) as rig:
+        response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True})
+        assert "the sum" in response.text, response.text
+        assert ANSWER not in response.text, response.text
+        assert len(rig.upstream_tools()) == 2
+        round_one: Final = _price(FIRST_ROUND) + FIXED_MARGIN
+        rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(round_one))
+        assert _statuses(rows) == ["failure", "success"], rows
+        (success,) = (row for row in rows if row["status"] == "success")
+        (failure,) = (row for row in rows if row["status"] == "failure")
+        assert (success["prompt_tokens"], success["completion_tokens"]) == FIRST_ROUND, rows
+        assert float(str(success["spend"])) == pytest.approx(round_one), rows
+        partial_answer: Final = (int(str(failure["prompt_tokens"])), int(str(failure["completion_tokens"])))
+        assert min(partial_answer) > 0, rows
+        assert float(str(failure["spend"])) == pytest.approx(_price(partial_answer)), rows
+
+
 @pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS)
 def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_once(
     gateway: Gateway, upstream: StreamedUpstream

tests/unit/proxy/test_common_request_processing.py

@@ -83,6 +83,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import (
     team_member_budget_entity_id,
 )
 from litellm.proxy.utils import ProxyLogging
+from litellm.responses.mcp.parked_stream_dispatch import IN_FLIGHT_ROUND
 from litellm.router import Router
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.utils import ModelResponseListIterator
@@ -9051,6 +9052,33 @@ class TestStreamingClientDisconnectBilling:
 
         proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once()
 
+    @pytest.mark.asyncio
+    async def test_disconnect_whose_in_flight_round_billing_raises_still_closes_the_stream_and_frees_the_slot(self):
+        import types
+
+        response = await self._start_partial_stream()
+        closed: Final = AsyncMock()
+
+        async def cannot_price_the_in_flight_round() -> None:
+            raise RuntimeError("the in-flight round could not be priced")
+
+        setattr(response.logging_obj, IN_FLIGHT_ROUND, cannot_price_the_in_flight_round)
+        stream: Final = types.SimpleNamespace(chunks=[], messages=None, aclose=closed)
+        proxy_logging_obj: Final = types.SimpleNamespace(arelease_max_parallel_requests_on_disconnect=AsyncMock())
+
+        await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup(
+            request=None,
+            request_data={"litellm_logging_obj": response.logging_obj},
+            response=stream,
+            stream_completed=False,
+            client_disconnected=True,
+            user_api_key_dict=MagicMock(),
+            proxy_logging_obj=proxy_logging_obj,
+        )
+
+        closed.assert_awaited_once()
+        proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once()
+
     async def _bill_and_collect_success_event(self, prepare=None, request_data=None):
         recorder = _RecordingSuccessLogger()
         original_callbacks = litellm.callbacks

tests/unit/responses/mcp/test_parked_stream_dispatch.py

@@ -0,0 +1,17 @@
+from typing import Final
+
+import pytest
+
+from litellm.responses.mcp.parked_stream_dispatch import parked_success
+from litellm.types.utils import ModelResponse
+
+_ASSEMBLED: Final = ModelResponse(id="chatcmpl-parked", model="gpt-5.6")
+
+
+@pytest.mark.parametrize(
+    "parked",
+    [(), (_ASSEMBLED,), (_ASSEMBLED, False, "not a logger"), (_ASSEMBLED, False, None, None)],
+    ids=("empty", "assembled-alone", "third-part-is-not-a-logger", "four-parts"),
+)
+def test_a_parked_dispatch_of_an_unknown_shape_is_never_dispatched(parked: tuple[object, ...]):
+    assert parked_success(parked) is None