Merge e3fbb6583d into 44cd858697

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

@@ -2482,6 +2482,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":
+        if not self.chunks:
+            return 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 partial is None:
+            return None
+        usage: Final[object] = getattr(partial, "usage", None)
+        if 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
@@ -2491,23 +2510,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/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,275 @@
+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,
+    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)
+
+    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 _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,62 @@
+from collections.abc import Coroutine
+from dataclasses import dataclass
+from typing import Final, Protocol
+
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.litellm_logging import UncachedPrice
+
+PARKED_DISPATCH: Final = "_deferred_stream_complete_args"
+DISPATCH_RELEASE: Final = "_on_deferred_stream_complete"
+IN_FLIGHT_ROUND: Final = "_bill_in_flight_round_on_disconnect"
+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]: ...
+
+
+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 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

@@ -81,6 +81,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
@@ -9049,6 +9050,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