fix(proxy): bill partial Responses API streams when the client disconnects (#47611)

* fix(proxy): bill partial Responses API streams when the client disconnects

Streamed /v1/responses requests and the default-route OpenAI and Azure /v1/messages requests behind them logged a 499 with zero spend when the client went away mid-stream, while /v1/chat/completions already billed the chunks it had sent. The Responses API iterator now keeps the response.created payload and can hand the disconnect cleanup a response.incomplete event with counted usage and the served service_tier, and the messages-route byte stream and router wrappers delegate to it.

* fix(proxy): satisfy lint gates for partial Responses API disconnect billing

Make the disconnect helper public so the proxy, router and messages adapter can import it without reportPrivateUsage, type the SSE byte stream as an AsyncGenerator so aclose is known, read litellm_call_id off the logging object, and keep the new tests on Final locals with the recorder injected through dynamic_async_success_callbacks instead of litellm.callbacks.

* fix(router): bill the primary stream's tokens after a mid-stream hop and chat-bridge Responses streams on disconnect

The router's fallback wrapper now folds the primary Responses stream's relayed tokens into the fallback's disconnect event the way _combine_responses_fallback_usage does for a natural end of stream, and LiteLLMCompletionStreamingIterator (chat-only providers served through /v1/responses) builds its disconnect event from the chat completion chunks it collected. stamp_responses_usage_cost is now public so the bridge module can price that event.

* Potential fix for pull request finding 'CodeQL / Log Injection'

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>

litellm/llms/anthropic/pass_through/responses_adapters/streaming_iterator.py

@@ -2,7 +2,7 @@
 ## Translates OpenAI call to Anthropic `/v1/messages` format
 import asyncio
 from collections import deque
-from collections.abc import AsyncIterator, Iterator, Mapping
+from collections.abc import AsyncGenerator, Iterator, Mapping
 from typing import TYPE_CHECKING, Any, Final
 
 from pydantic import ConfigDict, field_validator
@@ -21,7 +21,10 @@ from litellm.llms.anthropic.pass_through.messages.utils import (
 )
 from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults
 from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame
-from litellm.responses.streaming_iterator import stream_error_status_and_message
+from litellm.responses.streaming_iterator import (
+    partial_responses_incomplete_event_on_disconnect,
+    stream_error_status_and_message,
+)
 from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
 from litellm.types.llms.base import LiteLLMBaseModel
 
@@ -33,6 +36,7 @@ from .transformation import (
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
     from litellm.llms.anthropic.pass_through.safeguards import SafeguardsEvaluator
+    from litellm.types.llms.openai import ResponseIncompleteEvent
 
 
 class _UpstreamFailure(LiteLLMBaseModel):
@@ -488,8 +492,38 @@ class AnthropicResponsesStreamWrapper:
 
         raise StopAsyncIteration
 
-    async def async_anthropic_sse_wrapper(self) -> AsyncIterator[bytes]:
-        """Yield SSE-encoded bytes for each Anthropic event chunk."""
+    def async_anthropic_sse_wrapper(self) -> "AnthropicResponsesSSEByteStream":
+        """SSE-encoded bytes for each Anthropic event chunk."""
+        return AnthropicResponsesSSEByteStream(self)
+
+    async def _sse_frames(self) -> AsyncGenerator[bytes, None]:
         safeguard_results: Final = StreamedSafeguardResults(self.safeguards_evaluator)
         async for chunk in self:
             yield anthropic_sse_frame(await safeguard_results.observe(chunk))
+
+    def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
+        return partial_responses_incomplete_event_on_disconnect(self.responses_stream)
+
+
+class AnthropicResponsesSSEByteStream:
+    """
+    The byte stream the proxy relays for a default-route /v1/messages request.
+    Keeps the Responses API iterator reachable so a client disconnect can bill
+    the partial output the way native /v1/responses streams do.
+    """
+
+    def __init__(self, wrapper: AnthropicResponsesStreamWrapper) -> None:
+        self._wrapper: Final = wrapper
+        self._frames: Final = wrapper._sse_frames()  # pyright: ignore[reportPrivateUsage]  # the wrapper's own byte stream
+
+    def __aiter__(self) -> "AnthropicResponsesSSEByteStream":
+        return self
+
+    async def __anext__(self) -> bytes:
+        return await self._frames.__anext__()
+
+    async def aclose(self) -> None:
+        await self._frames.aclose()
+
+    def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
+        return self._wrapper.partial_response_on_disconnect()

litellm/proxy/common_request_processing.py

@@ -128,6 +128,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.streaming_iterator import partial_responses_incomplete_event_on_disconnect
 from litellm.router import Router
 from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
 from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
@@ -455,6 +456,13 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
         # A natural end-of-stream success event already fired and released the
         # slot; do not bill again, and let the caller skip the slot release.
         return True
+    partial_responses_event: Final = partial_responses_incomplete_event_on_disconnect(response)
+    if partial_responses_event is not None:
+        verbose_proxy_logger.debug(
+            "Billing partial streamed Responses API spend after client disconnect, litellm_call_id=%s",
+            sanitize_for_log(logging_obj.litellm_call_id),
+        )
+        return await _dispatch_disconnect_billing_event(logging_obj, partial_responses_event)
     chunks: Final[object] = getattr(response, "chunks", None)
     if not isinstance(chunks, list) or not chunks:
         return False
@@ -484,6 +492,10 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
     partial_usage: Final = getattr(partial_response, "usage", None)
     if isinstance(partial_usage, Usage):
         backfill_missing_cache_usage_fields(partial_usage)
+    return await _dispatch_disconnect_billing_event(logging_obj, partial_response)
+
+
+async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, partial_response: object) -> bool:
     try:
         await logging_obj.dispatch_success_handlers(
             partial_response,

litellm/responses/litellm_completion_transformation/streaming_iterator.py

@@ -22,7 +22,7 @@ from litellm.responses.litellm_completion_transformation.tool_search import (
 from litellm.responses.litellm_completion_transformation.transformation import (
     LiteLLMCompletionResponsesConfig,
 )
-from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
+from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, stamp_responses_usage_cost
 from litellm.responses.utils import ResponsesAPIRequestUtils
 from litellm.types.llms.openai import (
     PART_UNION_TYPES,
@@ -1363,6 +1363,43 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
             return False
         return not any(getattr(part, "text", None) for part in getattr(item, "content", None) or ())
 
+    def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None:
+        """
+        The response.incomplete event to bill when the client went away mid-stream.
+        Built from the chat completion chunks relayed so far, so the usage is what
+        stream_chunk_builder counts for them.
+        """
+        if self.completed_response is not None or not self.collected_chat_completion_chunks:
+            return None
+        built: Final = stream_chunk_builder(
+            chunks=self.collected_chat_completion_chunks,
+            messages=LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(  # pyright: ignore[reportUnknownMemberType]  # the transformer's signature is partially untyped
+                input=self.request_input,
+                responses_api_request=dict(self.responses_api_request),
+            ),
+            logging_obj=self.litellm_logging_obj,
+        )
+        if not isinstance(built, ModelResponse):
+            return None
+        responses_api_response: Final = (
+            LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(  # pyright: ignore[reportUnknownMemberType]  # the transformer's signature is partially untyped
+                request_input=self.request_input,
+                chat_completion_response=built,
+                responses_api_request=self.responses_api_request,
+            )
+        )
+        partial: Final = self._with_encoded_response_id(
+            responses_api_response.model_copy(
+                update={
+                    "id": self._cached_response_id or responses_api_response.id,
+                    "status": "incomplete",
+                    "output": list(self._output_with_streamed_item_ids(responses_api_response)),
+                }
+            )
+        )
+        stamp_responses_usage_cost(partial, self.litellm_logging_obj)
+        return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial)
+
     def _emit_terminal_response_event(
         self, litellm_model_response: ModelResponse
     ) -> ResponseCompletedEvent | ResponseIncompleteEvent | None:

litellm/responses/streaming_iterator.py

@@ -58,10 +58,12 @@ from litellm.types.integrations.custom_logger import converted_stream_requested
 from litellm.types.llms.openai import (
     PART_UNION_TYPES,
     ResponseAPIUsage,
+    ResponseIncompleteEvent,
     ResponsesAPIResponse,
     ResponsesAPIStreamEvents,
     ResponsesAPIStreamingResponse,
 )
+from litellm.types.responses.main import GenericResponseOutputItem, OutputText
 from litellm.types.utils import CallTypes
 from litellm.utils import async_post_call_success_deployment_hook
 
@@ -385,6 +387,7 @@ class BaseResponsesAPIStreamingIterator:
         self._output_started = False
         self._generated_content = ""
         self._generated_tool_arguments = ""
+        self._created_response: ResponsesAPIResponse | None = None
         self._completed_response_cached = False
         self._completed_response_logged = False
         self._completed_response_cache_hit: bool | None = None
@@ -544,6 +547,10 @@ class BaseResponsesAPIStreamingIterator:
                 # Store the completed response (also for incomplete/failed so logging still fires)
                 _chunk_type: Final = getattr(openai_responses_api_chunk, "type", None)
                 openai_types = get_openai_response_types()
+                if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED:
+                    _created: Final[object] = getattr(openai_responses_api_chunk, "response", None)
+                    if isinstance(_created, ResponsesAPIResponse):
+                        self._created_response = _created
                 if openai_responses_api_chunk and _chunk_type in (
                     openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
                     openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
@@ -575,7 +582,7 @@ class BaseResponsesAPIStreamingIterator:
                         else openai_responses_api_chunk.model_copy(update={"response": _billed_response})
                     )
                     self.completed_response = _terminal_chunk
-                    _stamp_responses_usage_cost(_billed_response, self.logging_obj)
+                    stamp_responses_usage_cost(_billed_response, self.logging_obj)
 
                     if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:
                         self._handle_logging_failed_response()
@@ -1030,6 +1037,34 @@ class BaseResponsesAPIStreamingIterator:
         if pending:
             await asyncio.wait(pending)
 
+    def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None:
+        """
+        The response.incomplete event to bill when the client went away mid-stream,
+        the same shape the success handler assembles for a natural end of stream.
+        Its response is the response.created payload (it carries the served
+        service_tier) plus the text and tool arguments relayed so far, with counted
+        usage since the provider never sent any. None when nothing was relayed yet
+        or the stream already reached its terminal event and billed itself.
+        """
+        created: Final = self._created_response
+        generated: Final = self._generated_content + self._generated_tool_arguments
+        if self.completed_response is not None or created is None or not generated:
+            return None
+        usage: Final = _estimate_usage_safely(
+            self.model or "", self.request_data.get("input"), self.request_data, generated
+        )
+        if usage is None:
+            return None
+        partial: Final = created.model_copy(
+            update={
+                "status": "incomplete",
+                "usage": usage,
+                "output": _partial_output_items(self._generated_content),
+            }
+        )
+        stamp_responses_usage_cost(partial, self.logging_obj)
+        return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial)
+
     def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None:
         self._yielded_first_chunk = True
         if event.type not in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES:
@@ -1561,6 +1596,37 @@ def _add_text_like_part_events(
         )
 
 
+@runtime_checkable
+class _BillsPartialStreamOnDisconnect(Protocol):
+    def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: ...
+
+
+def partial_responses_incomplete_event_on_disconnect(stream: object) -> ResponseIncompleteEvent | None:
+    """Ask the stream behind a proxy or router wrapper for the partial response to bill, if it can build one."""
+    if not isinstance(stream, _BillsPartialStreamOnDisconnect):
+        return None
+    try:
+        partial: Final = stream.partial_response_on_disconnect()
+    except Exception as e:  # noqa: BLE001  # disconnect billing is best-effort; never break stream teardown
+        verbose_logger.debug("Failed to build partial Responses API response for disconnect billing: %s", e)
+        return None
+    return partial if isinstance(partial, ResponseIncompleteEvent) else None
+
+
+def _partial_output_items(generated_text: str) -> Sequence[GenericResponseOutputItem]:
+    if not generated_text:
+        return []
+    return [
+        GenericResponseOutputItem(
+            type="message",
+            id=f"msg_{uuid.uuid4()}",
+            status="incomplete",
+            role="assistant",
+            content=[OutputText(type="output_text", text=generated_text, annotations=[])],
+        )
+    ]
+
+
 def _billed_terminal_response(
     response_obj: object, estimate: Callable[[], ResponseAPIUsage | None] | None
 ) -> ResponsesAPIResponse | None:
@@ -1667,7 +1733,7 @@ def _estimate_usage_safely(
         return None
 
 
-def _stamp_responses_usage_cost(
+def stamp_responses_usage_cost(
     response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
 ) -> None:
     if response_obj is None or logging_obj is None:
@@ -1717,7 +1783,7 @@ def build_synthetic_response_events(
     chunk_size: int,
 ) -> list[ResponsesAPIStreamingResponse]:
     openai_types: Final = get_openai_response_types()
-    _stamp_responses_usage_cost(transformed, logging_obj)
+    stamp_responses_usage_cost(transformed, logging_obj)
 
     events: Final[list[ResponsesAPIStreamingResponse]] = [
         _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed),

litellm/router.py

@@ -478,6 +478,7 @@ if TYPE_CHECKING:
     from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject
     from litellm.types.llms.openai import (
         ResponseAPIUsage,
+        ResponseIncompleteEvent,
         ResponseInputParam,
         ResponsesAPIResponse,
     )
@@ -983,6 +984,11 @@ class FallbackAwareAnthropicMessagesStream:
         self._source_iterator = fallback_response
         self.fallback_headers_adopted = True
 
+    def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
+        from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
+
+        return partial_responses_incomplete_event_on_disconnect(self._source_iterator)
+
     def follow_source_attribution(self) -> None:
         """
         A retry's or a fallback's stream carries a wrapper of its own, so a hop it makes before its
@@ -3911,6 +3917,7 @@ class Router:
                 from datetime import datetime
 
                 self._async_generator = async_generator
+                self._billing_source: object = source_iterator
                 # Mirror every attribute BaseResponsesAPIStreamingIterator.__init__
                 # would have set. The wrapper bypasses super().__init__ (it has no
                 # httpx.Response of its own and no provider config to drive), so
@@ -3959,10 +3966,29 @@ class Router:
 
             def adopt_fallback_headers(self, fallback_response: object) -> tuple[dict[str, object], dict[str, object]]:
                 prepared: Final = Router._prepare_fallback_hidden_params(fallback_response)
+                self._billing_source = fallback_response
                 self._hidden_params = {**prepared[0], "additional_headers": prepared[1]}
                 self.fallback_headers_adopted = True
                 return prepared
 
+            def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
+                """
+                The event to bill on a client disconnect. After a mid-stream hop the primary
+                stream's relayed tokens are folded into the fallback's event, the way
+                _combine_responses_fallback_usage does for a natural end of stream.
+                """
+                from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
+
+                current: Final = partial_responses_incomplete_event_on_disconnect(self._billing_source)
+                if self._billing_source is source_iterator:
+                    return current
+                primary: Final = partial_responses_incomplete_event_on_disconnect(source_iterator)
+                if current is None or primary is None:
+                    return current or primary
+                if primary.response.usage is not None:
+                    Router._combine_responses_fallback_usage(current, primary.response.usage)
+                return current
+
             def __aiter__(self):
                 return self
 

tests/unit/proxy/test_common_request_processing.py

@@ -9011,6 +9011,118 @@ def _apply_stream_usage_tracking(
     )
 
 
+    async def _start_partial_responses_stream(self, recorder: _RecordingSuccessLogger, *, relay_text: bool = True):
+        from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
+        from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
+
+        mocked: Final = await litellm.aresponses(
+            model="gpt-4o-mini",
+            input="tell me a story",
+            mock_response="The codename is AZURE-FALCON-42 and the story is long.",
+            stream=True,
+            api_key="test-key",
+        )
+        logging_obj: Final = mocked.logging_obj
+        logging_obj.dynamic_async_success_callbacks = [recorder]
+        stream: Final = ResponsesAPIStreamingIterator(
+            response=httpx.Response(200),
+            model="gpt-4o-mini",
+            responses_api_provider_config=OpenAIResponsesAPIConfig(),
+            logging_obj=logging_obj,
+            request_data={"model": "gpt-4o-mini", "input": "tell me a story"},
+            custom_llm_provider="openai",
+        )
+        created: Final = {
+            "type": "response.created",
+            "sequence_number": 0,
+            "response": {
+                "id": "resp_partial",
+                "object": "response",
+                "created_at": 1700000000.0,
+                "model": "gpt-4o-mini",
+                "status": "in_progress",
+                "output": [],
+                "parallel_tool_calls": True,
+                "tool_choice": "auto",
+                "tools": [],
+                "service_tier": "priority",
+            },
+        }
+        delta: Final = {
+            "type": "response.output_text.delta",
+            "sequence_number": 1,
+            "item_id": "msg_partial",
+            "output_index": 0,
+            "content_index": 0,
+            "delta": "The codename is AZURE-FALCON-42",
+        }
+        for event in (created, delta) if relay_text else (created,):
+            stream._process_chunk(json.dumps(event))
+        return stream
+
+    async def _bill_and_wait(self, recorder: _RecordingSuccessLogger, streamed: object, logging_obj: object) -> bool:
+        billed: Final = await _bill_partial_streamed_spend_on_disconnect({"litellm_logging_obj": logging_obj}, streamed)
+        for _ in range(50):
+            if recorder.success_events:
+                break
+            await asyncio.sleep(0.1)
+        await asyncio.sleep(0.5)
+        return billed
+
+    @pytest.mark.asyncio
+    async def test_disconnect_bills_partial_native_responses_stream_with_counted_usage(self):
+        from litellm.types.llms.openai import ResponsesAPIResponse
+
+        recorder: Final = _RecordingSuccessLogger()
+        stream: Final = await self._start_partial_responses_stream(recorder)
+        billed: Final = await self._bill_and_wait(recorder, stream, stream.logging_obj)
+
+        assert billed is True
+        assert len(recorder.success_events) == 1
+        partial_response: Final = recorder.success_events[0]["response_obj"]
+        assert isinstance(partial_response, ResponsesAPIResponse)
+        assert partial_response.status == "incomplete"
+        assert getattr(partial_response, "service_tier") == "priority"
+        assert partial_response.usage.output_tokens == litellm.token_counter(
+            model="gpt-4o-mini", text="The codename is AZURE-FALCON-42"
+        )
+        assert partial_response.usage.input_tokens > 0
+
+    @pytest.mark.asyncio
+    async def test_disconnect_bills_partial_default_route_messages_stream_through_the_responses_adapter(self):
+        from litellm.llms.anthropic.pass_through.responses_adapters.streaming_iterator import (
+            AnthropicResponsesStreamWrapper,
+        )
+        from litellm.router import FallbackAwareAnthropicMessagesStream
+        from litellm.types.llms.openai import ResponsesAPIResponse
+
+        async def _sse_frames() -> AsyncGenerator[bytes, None]:
+            yield b"event: message_start\n\n"
+
+        recorder: Final = _RecordingSuccessLogger()
+        stream: Final = await self._start_partial_responses_stream(recorder)
+        wrapper: Final = AnthropicResponsesStreamWrapper(responses_stream=stream, model="gpt-4o-mini")
+        streamed: Final = FallbackAwareAnthropicMessagesStream(_sse_frames(), wrapper.async_anthropic_sse_wrapper())
+        billed: Final = await self._bill_and_wait(recorder, streamed, stream.logging_obj)
+
+        assert billed is True
+        assert len(recorder.success_events) == 1
+        partial_response: Final = recorder.success_events[0]["response_obj"]
+        assert isinstance(partial_response, ResponsesAPIResponse)
+        assert partial_response.usage.output_tokens == litellm.token_counter(
+            model="gpt-4o-mini", text="The codename is AZURE-FALCON-42"
+        )
+
+    @pytest.mark.asyncio
+    async def test_disconnect_before_any_responses_output_bills_nothing(self):
+        recorder: Final = _RecordingSuccessLogger()
+        stream: Final = await self._start_partial_responses_stream(recorder, relay_text=False)
+        billed: Final = await self._bill_and_wait(recorder, stream, stream.logging_obj)
+
+        assert billed is False
+        assert recorder.success_events == []
+
+
 class TestApplyStreamUsageTracking:
     @pytest.mark.parametrize(
         ("general_settings", "requested_stream_options", "expected_stream_options"),

tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py

@@ -6566,3 +6566,40 @@ def test_a_continued_chain_leaves_the_turns_metadata_alone_without_request_corre
     assert request["litellm_trace_id"] == "chain-session"
     assert "litellm_metadata" not in request
     assert request["metadata"] == {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"}
+
+
+@pytest.mark.asyncio
+async def test_bridge_partial_response_on_disconnect_bills_the_relayed_chat_chunks():
+    from litellm.responses.litellm_completion_transformation.streaming_iterator import (
+        LiteLLMCompletionStreamingIterator,
+    )
+
+    chat_stream: Final = await litellm.acompletion(
+        model="gpt-4o-mini",
+        messages=[{"role": "user", "content": "tell me a story"}],
+        mock_response="The codename is AZURE-FALCON-42 and the story is long.",
+        stream=True,
+        api_key="test-key",
+    )
+    iterator: Final = LiteLLMCompletionStreamingIterator(
+        model="gpt-4o-mini",
+        litellm_custom_stream_wrapper=chat_stream,
+        request_input="tell me a story",
+        responses_api_request={},
+        custom_llm_provider="openai",
+    )
+    assert iterator.partial_response_on_disconnect() is None, "nothing relayed yet, nothing to bill"
+    events: Final = [await iterator.__anext__() for _ in range(6)]
+    relayed_text: Final = "".join(
+        getattr(event, "delta", "") for event in events if event.type == "response.output_text.delta"
+    )
+    assert relayed_text, "the bridge should have relayed some text by now"
+
+    incomplete_event: Final = iterator.partial_response_on_disconnect()
+
+    assert incomplete_event is not None
+    assert incomplete_event.type == "response.incomplete"
+    assert incomplete_event.response.status == "incomplete"
+    assert incomplete_event.response.usage is not None
+    assert incomplete_event.response.usage.output_tokens == litellm.token_counter(model="gpt-4o-mini", text=relayed_text)
+    assert incomplete_event.response.usage.input_tokens > 0

tests/unit/responses/test_responses_streaming_iterator.py

@@ -7,7 +7,10 @@ crashed with exit 139 whenever any CustomLogger was registered).
 """
 
 import asyncio
+import json
 import time
+from collections.abc import Mapping
+from typing import Final
 
 import httpx
 import pytest
@@ -17,6 +20,7 @@ from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module
 from litellm.responses import streaming_iterator as responses_streaming_iterator_module
 from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
+from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
 from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
 from litellm.types.llms.openai import ResponsesAPIResponse
 
@@ -163,3 +167,99 @@ async def test_sync_callbacks_run_only_after_async_handler_completes(recording_e
     submit_times = recording_executor.submit_times_for(logging_obj)
     assert len(submit_times) == 1
     assert submit_times[0] >= recorder.async_hook_finished
+
+
+_CREATED_RESPONSE_WITH_PRIORITY_TIER: Final[Mapping[str, object]] = {
+    "id": "resp_partial",
+    "object": "response",
+    "created_at": 1700000000.0,
+    "model": "gpt-5.4-nano",
+    "status": "in_progress",
+    "output": [],
+    "parallel_tool_calls": True,
+    "tool_choice": "auto",
+    "tools": [],
+    "service_tier": "priority",
+}
+_CREATED_WITH_PRIORITY_TIER: Final[Mapping[str, object]] = {
+    "type": "response.created",
+    "sequence_number": 0,
+    "response": _CREATED_RESPONSE_WITH_PRIORITY_TIER,
+}
+
+
+def _text_delta(text: str, sequence_number: int) -> Mapping[str, object]:
+    return {
+        "type": "response.output_text.delta",
+        "sequence_number": sequence_number,
+        "item_id": "msg_partial",
+        "output_index": 0,
+        "content_index": 0,
+        "delta": text,
+    }
+
+
+def _completed_after(text: str) -> Mapping[str, object]:
+    return {
+        "type": "response.completed",
+        "sequence_number": 9,
+        "response": {
+            **_CREATED_RESPONSE_WITH_PRIORITY_TIER,
+            "status": "completed",
+            "usage": {"input_tokens": 5, "output_tokens": 7, "total_tokens": 12},
+        },
+    }
+
+
+def _live_iterator(logging_obj: LitellmLogging) -> ResponsesAPIStreamingIterator:
+    return ResponsesAPIStreamingIterator(
+        response=httpx.Response(200),
+        model="gpt-5.4-nano",
+        responses_api_provider_config=OpenAIResponsesAPIConfig(),
+        logging_obj=logging_obj,
+        request_data={"model": "gpt-5.4-nano", "input": "tell me a story"},
+        custom_llm_provider="openai",
+    )
+
+
+def test_partial_response_on_disconnect_bills_the_relayed_text_at_the_served_service_tier():
+    iterator: Final = _live_iterator(_make_logging_obj())
+    for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time ", 1), _text_delta("a fox slept.", 2)):
+        iterator._process_chunk(json.dumps(event))
+
+    incomplete_event: Final = iterator.partial_response_on_disconnect()
+
+    assert incomplete_event is not None
+    assert incomplete_event.type == "response.incomplete"
+    partial: Final = incomplete_event.response
+    assert partial.status == "incomplete"
+    assert getattr(partial, "service_tier") == "priority"
+    assert partial.output[0].content[0].text == "Once upon a time a fox slept."
+    assert partial.usage.output_tokens == litellm.token_counter(
+        model="gpt-5.4-nano", text="Once upon a time a fox slept."
+    )
+    assert partial.usage.input_tokens == litellm.token_counter(
+        model="gpt-5.4-nano", messages=[{"role": "user", "content": "tell me a story"}]
+    )
+    assert partial.usage.total_tokens == partial.usage.input_tokens + partial.usage.output_tokens
+
+
+@pytest.mark.parametrize(
+    "events",
+    [
+        (),
+        (_CREATED_WITH_PRIORITY_TIER,),
+        (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1), _completed_after("Once upon a time")),
+    ],
+    ids=["nothing_relayed", "created_only", "already_completed"],
+)
+@pytest.mark.asyncio
+async def test_partial_response_on_disconnect_is_none_when_there_is_nothing_to_bill_or_the_stream_billed_itself(
+    events: tuple[Mapping[str, object], ...],
+):
+    iterator: Final = _live_iterator(_make_logging_obj())
+    for event in events:
+        iterator._process_chunk(json.dumps(event))
+
+    assert iterator.partial_response_on_disconnect() is None
+    await iterator._await_pending_logging()

tests/unit/responses/test_streaming_iterator.py

@@ -839,26 +839,26 @@ def _responses_api_response_with_usage() -> ResponsesAPIResponse:
 
 
 def test_stamp_responses_usage_cost_stamps_computed_cost():
-    from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
+    from litellm.responses.streaming_iterator import stamp_responses_usage_cost
 
     response = _responses_api_response_with_usage()
     logging_obj = Mock(spec=LiteLLMLoggingObj)
     logging_obj.response_cost_calculator.return_value = 0.000704
 
-    _stamp_responses_usage_cost(response, logging_obj)
+    stamp_responses_usage_cost(response, logging_obj)
 
     assert getattr(response.usage, "cost", None) == pytest.approx(0.000704)
     logging_obj.response_cost_calculator.assert_called_once_with(result=response)
 
 
 def test_stamp_responses_usage_cost_keeps_provider_reported_cost():
-    from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
+    from litellm.responses.streaming_iterator import stamp_responses_usage_cost
 
     response = _responses_api_response_with_usage()
     setattr(response.usage, "cost", 0.5)
     logging_obj = Mock(spec=LiteLLMLoggingObj)
 
-    _stamp_responses_usage_cost(response, logging_obj)
+    stamp_responses_usage_cost(response, logging_obj)
 
     assert getattr(response.usage, "cost", None) == pytest.approx(0.5)
     logging_obj.response_cost_calculator.assert_not_called()
@@ -878,7 +878,7 @@ def _unvalidated_response_with_dict_usage(usage: dict) -> ResponsesAPIResponse:
 
 
 def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage():
-    from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
+    from litellm.responses.streaming_iterator import stamp_responses_usage_cost
     response = _unvalidated_response_with_dict_usage(
         {
             "input_tokens": 29,
@@ -890,7 +890,7 @@ def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage():
     )
     logging_obj = Mock(spec=LiteLLMLoggingObj)
 
-    _stamp_responses_usage_cost(response, logging_obj)
+    stamp_responses_usage_cost(response, logging_obj)
 
     assert isinstance(response.usage, ResponseAPIUsage)
     assert response.usage.cost == pytest.approx(3e-05)
@@ -899,12 +899,12 @@ def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage():
 
 
 def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost():
-    from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
+    from litellm.responses.streaming_iterator import stamp_responses_usage_cost
     response = _unvalidated_response_with_dict_usage({"input_tokens": 29, "output_tokens": 120, "total_tokens": 149})
     logging_obj = Mock(spec=LiteLLMLoggingObj)
     logging_obj.response_cost_calculator.return_value = 0.000704
 
-    _stamp_responses_usage_cost(response, logging_obj)
+    stamp_responses_usage_cost(response, logging_obj)
 
     assert isinstance(response.usage, ResponseAPIUsage)
     assert response.usage.cost == pytest.approx(0.000704)
@@ -912,13 +912,13 @@ def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost():
 
 
 def test_stamp_responses_usage_cost_survives_calculator_failure():
-    from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
+    from litellm.responses.streaming_iterator import stamp_responses_usage_cost
 
     response = _responses_api_response_with_usage()
     logging_obj = Mock(spec=LiteLLMLoggingObj)
     logging_obj.response_cost_calculator.side_effect = RuntimeError("cost map unavailable")
 
-    _stamp_responses_usage_cost(response, logging_obj)
+    stamp_responses_usage_cost(response, logging_obj)
 
     assert getattr(response.usage, "cost", None) is None
 

tests/unit/test_router/test_router.py

@@ -29668,3 +29668,68 @@ def test_pre_call_checks_ignores_an_image_response_format_string(monkeypatch: py
     )
 
     assert [d["model_info"]["id"] for d in filtered] == ["only"]
+
+
+@pytest.mark.asyncio
+async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_streams_tokens():
+    """After a mid-stream hop, a disconnect bills the primary's relayed text plus the fallback's, not the fallback alone."""
+    import json
+
+    from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
+    from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
+
+    def _relayed_iterator(text: str) -> ResponsesAPIStreamingIterator:
+        logging_obj: Final = MagicMock(spec=LiteLLMLogging)
+        logging_obj.model_call_details = {"litellm_params": {}}
+        logging_obj.completion_start_time = None
+        iterator: Final = ResponsesAPIStreamingIterator(
+            response=httpx.Response(200),
+            model="gpt-5.4",
+            responses_api_provider_config=OpenAIResponsesAPIConfig(),
+            logging_obj=logging_obj,
+            request_data={"model": "gpt-5.4", "input": "Hi"},
+            custom_llm_provider="openai",
+        )
+        created: Final = {
+            "id": "resp_1",
+            "object": "response",
+            "created_at": 1700000000.0,
+            "model": "gpt-5.4",
+            "status": "in_progress",
+            "output": [],
+            "parallel_tool_calls": True,
+            "tool_choice": "auto",
+            "tools": [],
+        }
+        iterator._process_chunk(json.dumps({"type": "response.created", "sequence_number": 0, "response": created}))
+        iterator._process_chunk(
+            json.dumps(
+                {
+                    "type": "response.output_text.delta",
+                    "sequence_number": 1,
+                    "item_id": "msg_1",
+                    "output_index": 0,
+                    "content_index": 0,
+                    "delta": text,
+                }
+            )
+        )
+        return iterator
+
+    router: Final = Router(
+        model_list=[{"model_name": "primary", "litellm_params": {"model": "openai/gpt-5.4", "api_key": "k"}}]
+    )
+    primary: Final = _relayed_iterator("Once upon a time there was a lighthouse keeper")
+    fallback: Final = _relayed_iterator("who lit the lamp")
+    wrapped: Final = await router._aresponses_streaming_iterator(
+        response=primary,
+        initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": AsyncMock()},
+    )
+    primary_tokens: Final = primary.partial_response_on_disconnect().response.usage.output_tokens
+    fallback_tokens: Final = fallback.partial_response_on_disconnect().response.usage.output_tokens
+
+    assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens
+
+    wrapped.adopt_fallback_headers(fallback)
+
+    assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens + fallback_tokens