Merge f633abb117 into 38b4d31447

litellm/litellm_core_utils/core_helpers.py

@@ -670,6 +670,17 @@ def _assign_attributes(target: object, values: Iterable[tuple[str, object]]) ->
         setattr(target, key, value)
 
 
+def preserve_streaming_choice_provider_extras(
+    model_response: "ModelResponseStream", original_chunk: "ModelResponseStream"
+) -> None:
+    from litellm.types.utils import STREAMING_CHOICE_OPENAI_FIELDS
+
+    for outgoing, upstream in zip(model_response.choices, original_chunk.choices):
+        for key, value in (upstream.__pydantic_extra__ or {}).items():
+            if key not in STREAMING_CHOICE_OPENAI_FIELDS:
+                setattr(outgoing, key, value)
+
+
 def safe_deep_copy(data):
     """
     Safe Deep Copy

litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py

@@ -25,6 +25,7 @@ from litellm.types.llms.openai import (
     OpenAIModerationResponse,
 )
 from litellm.types.utils import (
+    STREAMING_CHOICE_OPENAI_FIELDS,
     ChatCompletionDeltaToolCall,
     ChatCompletionMessageCustomToolCall,
     ChatCompletionMessageToolCall,
@@ -55,8 +56,10 @@ from litellm.types.utils import Logprobs as TextCompletionLogprobs
 from .get_headers import get_response_headers
 
 _MESSAGE_FIELDS: Final[frozenset] = frozenset(Message.model_fields.keys())
-_CHOICES_FIELDS: Final[frozenset] = frozenset(Choices.model_fields.keys())
+_CHOICES_FIELDS: Final[frozenset[str]] = frozenset(Choices.model_fields.keys())
+_CHOICES_EXPLICIT_KWARGS: Final[frozenset[str]] = _CHOICES_FIELDS | {"enhancements"}
 _MODEL_RESPONSE_FIELDS: Final[frozenset] = frozenset(ModelResponse.model_fields.keys()) | {"usage"}
+_CACHED_CHOICE_RESERVED_KEYS: Final[frozenset[str]] = STREAMING_CHOICE_OPENAI_FIELDS | {"message", "finish_details"}
 
 
 def _normalize_images_for_message(
@@ -145,6 +148,10 @@ def _split_assembled_content_for_replay(content: str | None) -> list[str]:
     return _REPLAY_CONTENT_SLICE_RE.findall(content)
 
 
+def _choice_provider_extras(choice: Mapping[str, object], reserved_keys: frozenset[str]) -> Mapping[str, object]:
+    return {key: value for key, value in choice.items() if key not in reserved_keys}
+
+
 def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None:
     # Rebuild the delta as content-only so every accumulate-able field (role,
     # tool_calls, reasoning_content, thinking_blocks, audio, images,
@@ -158,6 +165,12 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None:
         del choice.enhancements
 
 
+def _clear_cached_choice_provider_extras(choice: StreamingChoices) -> None:
+    for key in tuple(choice.__pydantic_extra__ or {}):
+        if key not in STREAMING_CHOICE_OPENAI_FIELDS:
+            delattr(choice, key)
+
+
 def _first_replay_delta(
     message: Mapping[str, object], reasoning_items: Sequence[ChatCompletionReasoningItem] | None
 ) -> Delta:
@@ -242,7 +255,13 @@ async def convert_to_streaming_response_async(
 
         logprobs = choice.get("logprobs", None)
 
-        choice = StreamingChoices(finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs)
+        choice = StreamingChoices(
+            finish_reason=finish_reason,
+            index=idx,
+            delta=delta,
+            logprobs=logprobs,
+            **_choice_provider_extras(_FIELD_MAP.validate_python(choice), _CACHED_CHOICE_RESERVED_KEYS),
+        )
         choice_list.append(choice)
 
     model_response_object.choices = choice_list
@@ -291,6 +310,8 @@ async def convert_to_streaming_response_async(
         slice_chunk.choices[0].delta.content = piece
         if i > 0:
             _clear_later_replay_slice_metadata(slice_chunk.choices[0])
+        if i != last_idx:
+            _clear_cached_choice_provider_extras(slice_chunk.choices[0])
         slice_chunk.choices[0].finish_reason = original_finish_reason if i == last_idx else None
         if i == last_idx and original_usage is not None:
             setattr(slice_chunk, "usage", original_usage)
@@ -332,6 +353,7 @@ def convert_to_streaming_response(
             delta=delta,
             logprobs=logprobs,
             enhancements=enhancements,
+            **_choice_provider_extras(_FIELD_MAP.validate_python(choice), _CACHED_CHOICE_RESERVED_KEYS),
         )
 
         choice_list.append(choice)
@@ -376,6 +398,8 @@ def convert_to_streaming_response(
         slice_chunk.choices[0].delta.content = piece
         if i > 0:
             _clear_later_replay_slice_metadata(slice_chunk.choices[0])
+        if i != last_idx:
+            _clear_cached_choice_provider_extras(slice_chunk.choices[0])
         slice_chunk.choices[0].finish_reason = original_finish_reason if i == last_idx else None
         if i == last_idx and original_usage is not None:
             setattr(slice_chunk, "usage", original_usage)
@@ -834,6 +858,9 @@ def convert_to_model_response_object(
                     logprobs=logprobs,
                     enhancements=enhancements,
                     provider_specific_fields=provider_specific_fields,
+                    **_choice_provider_extras(
+                        _FIELD_MAP.validate_python(provider_specific_fields), _CHOICES_EXPLICIT_KWARGS
+                    ),
                 )
                 if provider_sent_logprobs and logprobs is None:
                     choice.logprobs = None

litellm/litellm_core_utils/streaming_handler.py

@@ -17,14 +17,14 @@ import httpx
 from openai.types.chat import ChatCompletionChunk
 from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
 from openai.types.chat.chat_completion_chunk import ChoiceDelta
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel, TypeAdapter, ValidationError
 from typing_extensions import NotRequired, TypedDict
 
 import litellm
 from litellm import verbose_logger
 from litellm._uuid import uuid
 from litellm.litellm_core_utils.asyncify import asyncify
-from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost
+from litellm.litellm_core_utils.core_helpers import as_str_mapping, chunk_usage, provider_reported_usage_cost
 from litellm.litellm_core_utils.hidden_params import (
     HIDDEN_PARAMS_ATTR,
     get_hidden_params,
@@ -40,6 +40,7 @@ from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer
 from litellm.types.llms.base import LiteLLMBaseModel
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import (
+    STREAMING_CHOICE_OPENAI_FIELDS,
     CacheCreationTokenDetails,
     CompletionTokensDetailsWrapper,
     Delta,
@@ -70,6 +71,7 @@ TOOL_CALLS_ATTRIBUTE: Final = "tool_calls"
 FUNCTION_CALL_ATTRIBUTE: Final = "function_call"
 
 _SYNC_ITER_EXHAUSTED: Final = object()
+_STREAM_OPTIONS_VALUE: Final = TypeAdapter(object)
 
 _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__)
 _CompleteStreamingResponse = ModelResponse | TextCompletionResponse | ModelResponseStream | None
@@ -337,6 +339,7 @@ class CustomStreamWrapper:
         self.system_fingerprint: str | None = None
         self._provider_response_model: str | None = None
         self.received_finish_reason: str | None = None
+        self._deferred_finish_chunk: ModelResponseStream | None = None
         self.intermittent_finish_reason: str | None = None  # finish reasons that show up mid-stream
         self.special_tokens = [
             "<|assistant|>",
@@ -373,10 +376,18 @@ class CustomStreamWrapper:
         self.received_model_id: str | None = None
         self.logging_loop = None
         self.rules = Rules()
-        self.stream_options = stream_options or getattr(logging_obj, "stream_options", None)
+        self.stream_options: object = _STREAM_OPTIONS_VALUE.validate_python(
+            stream_options or getattr(logging_obj, "stream_options", None)
+        )
         self.messages = getattr(logging_obj, "messages", None)
         self.sent_stream_usage = False
         self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False
+        stream_options_fields: Final = as_str_mapping(self.stream_options)
+        self._continuous_usage_stats = (
+            isinstance(self.stream_options, dict)
+            and stream_options_fields is not None
+            and stream_options_fields.get("continuous_usage_stats") is True
+        ) and self.send_stream_usage
         self.caller_requested_stream_usage: Final = (
             self.check_send_stream_usage(getattr(logging_obj, "stream_options", None))
             and getattr(logging_obj, "stream_usage_injected", False) is not True
@@ -543,8 +554,11 @@ class CustomStreamWrapper:
                     )
         self._restore_consumer_correlation_context()
 
-    def check_send_stream_usage(self, stream_options: dict | None):
-        return stream_options is not None and stream_options.get("include_usage", False) is True
+    def check_send_stream_usage(self, stream_options: object) -> bool:
+        if stream_options is None:
+            return False
+        get_option: Final[Callable[[str, object], object] | None] = getattr(stream_options, "get", None)
+        return get_option is not None and get_option("include_usage", False) is True
 
     def check_is_function_call(self, logging_obj) -> bool:
         from litellm.litellm_core_utils.prompt_templates.common_utils import (
@@ -591,9 +605,7 @@ class CustomStreamWrapper:
 
         last_content: Final = self.chunks[-1].choices[0].delta.content
 
-        if (
-            not isinstance(last_content, str) or len(last_content) <= 2
-        ):  # ignore empty content - https://github.com/BerriAI/litellm/issues/5158#issuecomment-2287156946
+        if last_content is None or len(last_content) <= 2:
             self._repeated_messages_count = 1
             return
 
@@ -1129,6 +1141,7 @@ class CustomStreamWrapper:
         plain_delta_created: bool = False,
     ):
         from litellm.litellm_core_utils.core_helpers import (
+            preserve_streaming_choice_provider_extras,
             preserve_upstream_non_openai_attributes,
         )
 
@@ -1227,6 +1240,24 @@ class CustomStreamWrapper:
 
                 self._optional_combine_thinking_block_in_choices(model_response=model_response)
 
+                if self.received_finish_reason is not None:
+                    self._deferred_finish_chunk = ModelResponseStream(choices=model_response.choices).model_copy(
+                        deep=True
+                    )
+                    return model_response.model_copy(
+                        update={
+                            "choices": [
+                                StreamingChoices(
+                                    **{
+                                        key: value
+                                        for key, value in choice.model_dump().items()
+                                        if key in STREAMING_CHOICE_OPENAI_FIELDS
+                                    }
+                                )
+                                for choice in model_response.choices
+                            ]
+                        }
+                    )
                 return model_response
             else:
                 return
@@ -1261,6 +1292,10 @@ class CustomStreamWrapper:
                     model_response=model_response,
                     original_chunk=_original_chunk,
                 )
+                preserve_streaming_choice_provider_extras(
+                    model_response=model_response,
+                    original_chunk=_original_chunk,
+                )
 
             if _is_delta_empty:
                 model_response.choices[0].delta = Delta(content=None)  # ensure empty delta chunk returned
@@ -1694,7 +1729,7 @@ class CustomStreamWrapper:
                         model_response.choices[0].delta = plain_delta if plain_delta is not None else Delta()
                         plain_delta_created = plain_delta is not None  # rebind-ok: the first chunk builds its delta
                 else:
-                    if self.send_stream_usage is True:
+                    if self.send_stream_usage is True and not self._continuous_usage_stats:
                         model_response.choices = []
                         return model_response
                     self._record_usage_only_chunk(model_response=model_response)
@@ -1912,9 +1947,13 @@ class CustomStreamWrapper:
             self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
 
     def finish_reason_handler(self):
+        from litellm.litellm_core_utils.core_helpers import preserve_streaming_choice_provider_extras
+
         model_response: Final = self.model_response_creator(
             hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id}
         )
+        if self._deferred_finish_chunk is not None:
+            preserve_streaming_choice_provider_extras(model_response, self._deferred_finish_chunk)
         _finish_reason: Final = self.received_finish_reason or self.intermittent_finish_reason
         if _finish_reason is not None:
             model_response.choices[0].finish_reason = _finish_reason
@@ -2071,7 +2110,7 @@ class CustomStreamWrapper:
                     # ModelResponseStream declares `usage` as a field, so
                     # hasattr(response, "usage") is always True — must check
                     # `is not None` to avoid running this path on every chunk.
-                    if getattr(response, "usage", None) is not None:
+                    if getattr(response, "usage", None) is not None and not self._continuous_usage_stats:
                         usage_to_preserve = response.usage
                         if usage_to_preserve:
                             response_hidden_params_for_usage = cast(  # cast-ok: preserve dynamic mapping behavior
@@ -2173,6 +2212,8 @@ class CustomStreamWrapper:
                     self.chunks.append(processed_chunk)
                 if self.stream_options is None:
                     self._set_final_chunk_usage(processed_chunk)
+                elif self._continuous_usage_stats:
+                    processed_chunk.usage = calculate_total_usage(chunks=self.chunks)
                 ## LOGGING
                 executor.submit(
                     self.run_success_logging_and_cache_storage,
@@ -2259,7 +2300,7 @@ class CustomStreamWrapper:
                         hasattr(processed_chunk, "usage") and getattr(processed_chunk, "usage", None) is not None
                     )
 
-                    if _has_usage:
+                    if _has_usage and not self._continuous_usage_stats:
                         usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk)
                         if usage_free_chunk is None:
                             continue
@@ -2406,6 +2447,8 @@ class CustomStreamWrapper:
                 self.chunks.append(processed_chunk)
             if self.stream_options is None:
                 self._set_final_chunk_usage(processed_chunk)
+            elif self._continuous_usage_stats:
+                processed_chunk.usage = calculate_total_usage(chunks=self.chunks)
             # see sync __next__'s sibling branch: deliberately do NOT restore
             # here - this chunk is still this call's own data, and restoring
             # before returning it would corrupt the caller's own log

litellm/proxy/common_request_processing.py

@@ -59,6 +59,7 @@ from litellm.litellm_core_utils.bug_report import (
     strip_bug_report_notice,
 )
 from litellm.litellm_core_utils.core_helpers import (
+    as_str_mapping,
     get_or_create_metadata_bucket,
     independent_snapshot,
     is_expected_client_error,
@@ -591,6 +592,7 @@ def stream_usage_tracking_updates(
         return scrub
     always_include: Final = general_settings.get("always_include_stream_usage")
     stream_options: Final = data.get("stream_options")
+    stream_options_fields: Final = as_str_mapping(stream_options)
     if always_include is True:
         if "stream_options" not in data:
             return {**scrub, "stream_options": {"include_usage": True}, STREAM_USAGE_INJECTED_FIELD: True}
@@ -609,7 +611,11 @@ def stream_usage_tracking_updates(
         return scrub
     if not supports_stream_options():
         return scrub
-    merged_stream_options: Final = {**stream_options} if isinstance(stream_options, dict) else {}
+    merged_stream_options: Final = (
+        {key: value for key, value in stream_options_fields.items() if key != "continuous_usage_stats"}
+        if isinstance(stream_options, dict) and stream_options_fields is not None
+        else {}
+    )
     return {
         "stream_options": {**merged_stream_options, "include_usage": True},
         "_litellm_strip_stream_usage": True,

litellm/proxy/proxy_server.py

@@ -192,6 +192,7 @@ from litellm.router_utils.routing_groups import parse_routing_groups
 from litellm.types.caching import RedisPipelineIncrementOperation
 from litellm.types.utils import (
     PRICING_OVERRIDES_KEY,
+    STREAMING_CHOICE_OPENAI_FIELDS,
     CountContextManagement,
     ModelResponse,
     ModelResponseStream,
@@ -10056,6 +10057,13 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) -
     return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending
 
 
+def _streaming_choice_carries_provider_extras(choice: object) -> bool:
+    extra: Final[Mapping[str, object] | None] = getattr(choice, "__pydantic_extra__", None)
+    if not isinstance(extra, dict):
+        return False
+    return any(key not in STREAMING_CHOICE_OPENAI_FIELDS for key in extra)
+
+
 def _fast_serialize_simple_model_response_stream(
     chunk: ModelResponseStream,
 ) -> bytes | None:
@@ -10078,6 +10086,8 @@ def _fast_serialize_simple_model_response_stream(
     choice: Final = choices[0]
     if getattr(choice, "logprobs", None) is not None or getattr(choice, "enhancements", None) is not None:
         return None
+    if _streaming_choice_carries_provider_extras(chunk.choices[0]):
+        return None
 
     delta: Final[object] = getattr(choice, "delta", None)
     if delta is None:

litellm/types/utils.py

@@ -2098,6 +2098,9 @@ class UnreportedUsage(Usage):
     pass
 
 
+STREAMING_CHOICE_OPENAI_FIELDS: Final = frozenset({"index", "delta", "finish_reason", "logprobs", "enhancements"})
+
+
 class StreamingChoices(OpenAIObject):
     def __init__(
         self,

tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py

@@ -2968,3 +2968,64 @@ def test_convert_to_model_response_object_function_output():
     assert result.usage.completion_tokens_details == CompletionTokensDetailsWrapper(reasoning_tokens=0)
 
     assert result._hidden_params is not None
+
+
+def test_convert_to_model_response_object_keeps_provider_choice_level_extras_at_choice_level():
+    response_object = {
+        "id": "chatcmpl-vllm-extras",
+        "object": "chat.completion",
+        "created": 1,
+        "model": "HostedVLLM/Qwen2.5-0.5B-Instruct",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "Hi"},
+                "finish_reason": "stop",
+                "stop_reason": "eos",
+                "token_ids": [11, 22, 33],
+            }
+        ],
+        "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
+    }
+
+    result = convert_to_model_response_object(
+        model_response_object=ModelResponse(),
+        response_object=response_object,
+        response_type="completion",
+    )
+
+    choice: Final = result.choices[0]
+    assert getattr(choice, "stop_reason", None) == "eos"
+    assert getattr(choice, "token_ids", None) == [11, 22, 33]
+    assert choice.provider_specific_fields is not None
+    assert choice.provider_specific_fields["stop_reason"] == "eos"
+    assert choice.provider_specific_fields["token_ids"] == [11, 22, 33]
+    dumped: Final = json.loads(choice.model_dump_json())
+    assert dumped["stop_reason"] == "eos"
+    assert dumped["token_ids"] == [11, 22, 33]
+
+
+def test_convert_to_model_response_object_accepts_a_choice_enhancements_key():
+    response_object = {
+        "id": "chatcmpl-enhancements",
+        "object": "chat.completion",
+        "created": 1,
+        "model": "azure/gpt-4o",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "Hi"},
+                "finish_reason": "stop",
+                "enhancements": {"grounding": True},
+            }
+        ],
+        "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
+    }
+
+    result = convert_to_model_response_object(
+        model_response_object=ModelResponse(),
+        response_object=response_object,
+        response_type="completion",
+    )
+
+    assert result.choices[0].enhancements == {"grounding": True}

tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py

@@ -460,3 +460,18 @@ def test_sync_non_enumerated_delta_fields_only_on_first_chunk():
     assert getattr(chunks[0].choices[0].delta, "annotations", None) == annotations
     for c in chunks[1:]:
         assert getattr(c.choices[0].delta, "annotations", None) is None
+
+
+@pytest.mark.parametrize("replay", ["async", "sync"])
+@pytest.mark.asyncio
+async def test_a_replayed_stream_keeps_cached_choice_provider_extras_on_the_finish_chunk(replay):
+    cached = _async_payload()
+    cached["choices"][0]["stop_reason"] = "eos"
+    cached["choices"][0]["token_ids"] = [9]
+
+    chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached))
+
+    assert len(chunks) > 1
+    assert [getattr(c.choices[0], "stop_reason", None) for c in chunks] == [None] * (len(chunks) - 1) + ["eos"]
+    assert [getattr(c.choices[0], "token_ids", None) for c in chunks] == [None] * (len(chunks) - 1) + [[9]]
+    assert chunks[-1].choices[0].finish_reason == "stop"

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -15,6 +15,7 @@ import httpx
 import pytest
 from openai.types.chat import ChatCompletionChunk
 from pydantic import BaseModel, TypeAdapter
+
 import litellm
 from litellm import verbose_logger
 from litellm._logging import session_id_var, trace_id_var
@@ -7846,6 +7847,179 @@ async def test_leaving_an_async_with_block_closes_the_provider_stream():
     assert stream.close_calls == 1
 
 
+def _vllm_wrapper(
+    chunks: Sequence[ModelResponseStream | ChatCompletionChunk],
+    stream_options: dict | None = None,
+    custom_llm_provider: str = "hosted_vllm",
+) -> CustomStreamWrapper:
+    return CustomStreamWrapper(
+        completion_stream=ModelResponseListIterator(model_responses=chunks),
+        model="hosted_vllm/m",
+        custom_llm_provider=custom_llm_provider,
+        stream_options=stream_options,
+        logging_obj=Logging(
+            model="hosted_vllm/m",
+            messages=[{"role": "user", "content": "hi"}],
+            stream=True,
+            call_type="acompletion",
+            start_time=time.time(),
+            litellm_call_id="vllm-choice",
+            function_id="test",
+        ),
+    )
+
+
+def _vllm_chunk(choice: dict, usage: dict | None = None) -> ModelResponseStream:
+    body: Final = {"id": "c", "created": 1, "model": "hosted_vllm/m", "choices": [choice]}
+    if usage is not None:
+        body["usage"] = usage
+    return ModelResponseStream(**body)
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_vllm_choice_extras_survive_the_terminal_finish_chunk(sync_mode: bool):
+    chunks: Final = [
+        _vllm_chunk({"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}),
+        _vllm_chunk({"index": 0, "delta": {}, "finish_reason": "stop", "stop_reason": "eos", "token_ids": [9]}),
+    ]
+    wrapper: Final = _vllm_wrapper(chunks)
+
+    emitted: Final = [chunk async for chunk in wrapper] if not sync_mode else list(wrapper)
+
+    finish_choice: Final = emitted[-1].choices[0]
+    assert getattr(finish_choice, "stop_reason", None) == "eos"
+    assert getattr(finish_choice, "token_ids", None) == [9]
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.parametrize("source", ["sdk", "http_handler", "cache"])
+@pytest.mark.parametrize("content", ["", "last token"])
+@pytest.mark.parametrize(("stop_reason", "token_ids"), [("eos", (9,)), (0, ())])
+@pytest.mark.asyncio
+async def test_vllm_finish_choice_extras_survive_serialized_live_and_cached_streams(
+    sync_mode: bool, source: str, content: str, stop_reason: str | int, token_ids: tuple[int, ...]
+) -> None:
+    from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
+        convert_to_streaming_response,
+        convert_to_streaming_response_async,
+    )
+    from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
+    from litellm.proxy.proxy_server import _serialize_streaming_chunk
+
+    terminal_choice: Final = {
+        "index": 0,
+        "finish_reason": "stop",
+        "stop_reason": stop_reason,
+        "token_ids": list(token_ids),
+    }
+    cached_response: Final = {
+        "id": "c",
+        "created": 1,
+        "model": "hosted_vllm/m",
+        "choices": [{**terminal_choice, "message": {"role": "assistant", "content": content}}],
+    }
+    replay_chunks: Final = (
+        list(convert_to_streaming_response(cached_response))
+        if sync_mode
+        else [chunk async for chunk in convert_to_streaming_response_async(cached_response)]
+    )
+    parser: Final = OpenAIChatCompletionStreamingHandler(streaming_response=None, sync_stream=True)
+    live_body: Final = {
+        "id": "c",
+        "object": "chat.completion.chunk",
+        "created": 1,
+        "model": "hosted_vllm/m",
+        "choices": [{**terminal_choice, "delta": {"content": content}}],
+    }
+    live_chunk: Final = ChatCompletionChunk(**live_body) if source == "sdk" else parser.chunk_parser(live_body)
+    wrapper: Final = _vllm_wrapper(
+        replay_chunks if source == "cache" else [live_chunk],
+        custom_llm_provider="cached_response" if source == "cache" else "hosted_vllm",
+    )
+
+    emitted: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper]
+    serialized: Final = [json.loads(_serialize_streaming_chunk(chunk)) for chunk in emitted]
+    finish_choices: Final = [chunk["choices"][0] for chunk in serialized if chunk["choices"][0].get("finish_reason")]
+
+    assert len(finish_choices) == 1
+    assert finish_choices[0]["stop_reason"] == terminal_choice["stop_reason"]
+    assert finish_choices[0]["token_ids"] == terminal_choice["token_ids"]
+    assert not finish_choices[0]["delta"].get("content")
+    assert [
+        chunk["choices"][0]["token_ids"] for chunk in serialized if "token_ids" in chunk["choices"][0]
+    ] == [terminal_choice["token_ids"]]
+    assert "".join(chunk["choices"][0]["delta"].get("content") or "" for chunk in serialized) == content
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_continuous_usage_stats_keeps_usage_on_every_chunk(sync_mode: bool):
+    counts: Final = [
+        {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
+        {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
+        {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13},
+    ]
+    chunks: Final = [
+        _vllm_chunk({"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}, usage=counts[0]),
+        _vllm_chunk({"index": 0, "delta": {}, "finish_reason": "stop"}, usage=counts[1]),
+        ModelResponseStream(id="c", created=1, model="hosted_vllm/m", choices=[], usage=Usage(**counts[2])),
+    ]
+    wrapper: Final = _vllm_wrapper(chunks, stream_options={"include_usage": True, "continuous_usage_stats": True})
+
+    emitted: Final = [chunk async for chunk in wrapper] if not sync_mode else list(wrapper)
+
+    assert all(getattr(chunk, "usage", None) is not None for chunk in emitted)
+    assert [chunk.usage.total_tokens for chunk in emitted] == [11, 12, 13]
+    assert litellm.stream_chunk_builder(chunks=emitted).usage.total_tokens == 13
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_continuous_usage_stats_finishes_before_the_final_usage_chunk(sync_mode: bool):
+    content_usage: Final = {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}
+    final_usage: Final = {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13}
+    chunks: Final = [
+        _vllm_chunk({"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}, usage=content_usage),
+        _vllm_chunk({"index": 0, "delta": {"content": "!"}, "finish_reason": "length"}, usage=final_usage),
+        ModelResponseStream(id="c", created=1, model="hosted_vllm/m", choices=[], usage=Usage(**final_usage)),
+    ]
+    wrapper: Final = _vllm_wrapper(chunks, stream_options={"include_usage": True, "continuous_usage_stats": True})
+
+    emitted: Final = [chunk async for chunk in wrapper] if not sync_mode else list(wrapper)
+
+    finish_indexes: Final = [i for i, chunk in enumerate(emitted) if chunk.choices and chunk.choices[0].finish_reason]
+    assert len(finish_indexes) == 1
+    assert emitted[finish_indexes[0]].choices[0].finish_reason == "length"
+    assert finish_indexes[0] < len(emitted) - 1
+    terminal_chunk: Final = emitted[-1]
+    assert not (terminal_chunk.choices and terminal_chunk.choices[0].finish_reason)
+    assert [chunk.usage.total_tokens for chunk in emitted] == [
+        content_usage["total_tokens"],
+        *[final_usage["total_tokens"]] * 3,
+    ]
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_include_usage_without_continuous_usage_stats_still_emits_only_a_final_usage_chunk(sync_mode: bool):
+    counts: Final = [
+        {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
+        {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13},
+    ]
+    chunks: Final = [
+        _vllm_chunk({"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}, usage=counts[0]),
+        _vllm_chunk({"index": 0, "delta": {}, "finish_reason": "stop"}, usage=counts[1]),
+    ]
+    wrapper: Final = _vllm_wrapper(chunks, stream_options={"include_usage": True})
+
+    emitted: Final = [chunk async for chunk in wrapper] if not sync_mode else list(wrapper)
+
+    content_chunks: Final = [chunk for chunk in emitted[:-1] if chunk.choices and chunk.choices[0].delta.content]
+    assert content_chunks
+    assert all(getattr(chunk, "usage", None) is None for chunk in content_chunks)
+    assert getattr(emitted[-1], "usage", None) is not None
+
 def _final_chunk_with_dict_usage(choices) -> ModelResponseStream:
     """A stream chunk whose ``usage`` is a plain dict.
 

tests/unit/proxy/proxy_server/test_streaming_helpers.py

@@ -544,6 +544,19 @@ def test_fast_serialize_simple_model_response_stream_returns_bytes_payload():
     }
 
 
+def test_fast_serialize_falls_back_when_choice_carries_provider_extras():
+    chunk = _simple_chunk()
+    chunk.choices[0].stop_reason = "eos"
+    chunk.choices[0].token_ids = [1, 2, 3]
+
+    assert _fast_serialize_simple_model_response_stream(chunk) is None
+
+    serialized: Final = _serialize_streaming_chunk(chunk)
+    payload: Final = json.loads(serialized)
+    assert payload["choices"][0]["stop_reason"] == "eos"
+    assert payload["choices"][0]["token_ids"] == [1, 2, 3]
+
+
 def test_fast_serialize_simple_model_response_stream_with_usage_returns_none_invalid():
     """Fast path bails (returns None) when ``usage`` is populated — the slow
     path is required to preserve usage fields. Returning None here is the

tests/unit/proxy/test_common_request_processing.py

@@ -9371,6 +9371,23 @@ class TestApplyStreamUsageTracking:
         assert "stream_options" not in data
         assert "_litellm_strip_stream_usage" not in data
 
+    def test_client_continuous_usage_stats_is_not_injected_onto_the_caller(self):
+        data = {"stream": True, "stream_options": {"continuous_usage_stats": True}}
+
+        _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion")
+
+        assert data["stream_options"] == {"include_usage": True}
+        assert data["_litellm_strip_stream_usage"] is True
+
+    def test_client_requested_continuous_usage_stats_reaches_the_backend_intact(self):
+        requested: Final = {"include_usage": True, "continuous_usage_stats": True}
+        data = {"stream": True, "stream_options": dict(requested)}
+
+        _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion")
+
+        assert data["stream_options"] == requested
+        assert "_litellm_strip_stream_usage" not in data
+
     def test_client_supplied_strip_marker_is_neutralized(self):
         data = {
             "stream": True,