From c68808e0669fbae4713d1b80896db937c8b147a2 Mon Sep 17 00:00:00 2001 From: AkilanChithra Date: Mon, 12 Oct 2026 13:28:38 -0500 Subject: [PATCH 1/5] fix(streaming): preserve requested continuous usage --- .../litellm_core_utils/streaming_handler.py | 22 +++++-- .../test_streaming_handler.py | 59 +++++++++++++++++++ 2 files changed, 76 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index a78e6f17cc9..a7e68aab2f1 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -421,6 +421,13 @@ class CustomStreamWrapper: 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 _should_forward_continuous_usage(self) -> bool: + return ( + self.send_stream_usage is True + and self.stream_options is not None + and self.stream_options.get("continuous_usage_stats", False) is True + ) + def check_is_function_call(self, logging_obj) -> bool: from litellm.litellm_core_utils.prompt_templates.common_utils import ( is_function_call, @@ -1415,6 +1422,11 @@ class CustomStreamWrapper: return _ProviderChunkParsed(response_obj) def chunk_creator(self, chunk: Any): + finish_reason_seen: Final = ( + self.sent_last_chunk is True + or self.received_finish_reason is not None + or self.intermittent_finish_reason is not None + ) if hasattr(chunk, "id"): self.response_id = chunk.id provider_response_model: Final = _provider_response_model(chunk) @@ -1508,7 +1520,9 @@ class CustomStreamWrapper: except Exception: model_response.choices[0].delta = Delta() else: - if self.send_stream_usage is True: + if self.send_stream_usage is True and ( + not self._should_forward_continuous_usage() or not finish_reason_seen + ): model_response.choices = [] return model_response self._record_usage_only_chunk(model_response=model_response) @@ -1843,7 +1857,7 @@ class CustomStreamWrapper: obj_dict = response.model_dump() - if "usage" in obj_dict: + if "usage" in obj_dict and not self._should_forward_continuous_usage(): del obj_dict["usage"] response_hidden_params_for_model = cast( # cast-ok: preserve dynamic mapping behavior @@ -2053,10 +2067,8 @@ class CustomStreamWrapper: # directly to avoid expensive model_copy() per chunk. self.chunks.append(processed_chunk.model_copy()) - # Strip usage from the outgoing chunk so it's not sent twice - # (once in the chunk, once in _hidden_params). obj_dict = processed_chunk.model_dump() - if "usage" in obj_dict: + if "usage" in obj_dict and not self._should_forward_continuous_usage(): del obj_dict["usage"] processed_chunk = self.model_response_creator( chunk=obj_dict, diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 01fae0958e6..ce853372696 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -422,6 +422,65 @@ def test_strip_sse_data_from_chunk(): assert CustomStreamWrapper.strip_sse_data_from_chunk(None) is None +def _continuous_usage_fixture_chunks() -> tuple[ModelResponseStream, ...]: + return ( + ModelResponseStream( + id="chatcmpl-continuous-usage", + created=1, + model="stream-usage-test", + choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="a"))], + usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + ), + ModelResponseStream( + id="chatcmpl-continuous-usage", + created=1, + model="stream-usage-test", + choices=[StreamingChoices(index=0, delta=Delta(content="b"), finish_reason="stop")], + usage=Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12), + ), + ModelResponseStream( + id="chatcmpl-continuous-usage", + created=1, + model="stream-usage-test", + choices=[], + usage=Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12), + ), + ) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_forwards_running_counts(logging_obj: Logging, sync_mode: bool) -> None: + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=_continuous_usage_fixture_chunks()), + model="stream-usage-test", + logging_obj=logging_obj, + custom_llm_provider="hosted_vllm", + stream_options={"include_usage": True, "continuous_usage_stats": True}, + ) + collected: Final = tuple(wrapper) if sync_mode else tuple([chunk async for chunk in wrapper]) + content_chunks: Final = tuple(chunk for chunk in collected if chunk.choices and chunk.choices[0].delta.content) + + assert tuple(chunk.choices[0].delta.content for chunk in content_chunks) == ("a", "b") + assert tuple(getattr(chunk, "usage", None) for chunk in content_chunks) == ( + Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12), + ) + expected_final_usage: Final = Usage( + prompt_tokens=10, + completion_tokens=2, + total_tokens=12, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=2), + ) + assert collected[-1].usage == expected_final_usage + assert tuple(getattr(chunk, "usage", None) for chunk in collected) == ( + Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12), + None, + expected_final_usage, + ) + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_streaming_handler_with_usage( From f9b9251f78a1fed5ce7865180d50161b51ebc434 Mon Sep 17 00:00:00 2001 From: AkilanChithra Date: Mon, 12 Oct 2026 13:41:11 -0500 Subject: [PATCH 2/5] test(streaming): cover continuous usage edge cases --- .../test_streaming_handler.py | 169 ++++++++++++++++++ 1 file changed, 169 insertions(+) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index ce853372696..9cdabb5e980 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -481,6 +481,175 @@ async def test_continuous_usage_forwards_running_counts(logging_obj: Logging, sy ) +def _continuous_usage_stream( + logging_obj: Logging, + chunks: tuple[ModelResponseStream, ...], + stream_options: dict[str, bool] | None, +) -> CustomStreamWrapper: + return CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=chunks), + model="stream-usage-test", + logging_obj=logging_obj, + custom_llm_provider="hosted_vllm", + stream_options=stream_options, + ) + + +async def _continuous_usage_collected(wrapper: CustomStreamWrapper, sync_mode: bool) -> tuple[ModelResponseStream, ...]: + return tuple(wrapper) if sync_mode else tuple([chunk async for chunk in wrapper]) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.parametrize( + ("stream_options", "include_final"), + ( + (None, False), + ({"include_usage": True}, True), + ({"include_usage": True, "continuous_usage_stats": False}, True), + ({"include_usage": False, "continuous_usage_stats": True}, False), + ({"continuous_usage_stats": True}, False), + ), +) +@pytest.mark.asyncio +async def test_continuous_usage_default_contract( + logging_obj: Logging, sync_mode: bool, stream_options: dict[str, bool] | None, include_final: bool +) -> None: + wrapper: Final = _continuous_usage_stream(logging_obj, _continuous_usage_fixture_chunks(), stream_options) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + content: Final = tuple(chunk for chunk in collected if chunk.choices and chunk.choices[0].delta.content) + assert tuple(getattr(chunk, "usage", None) for chunk in content) == (None, None) + assert getattr(collected[-1], "usage", None) == ( + Usage( + prompt_tokens=10, + completion_tokens=2, + total_tokens=12, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=2), + ) + if include_final + else None + ) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_preserves_zero_counts(logging_obj: Logging, sync_mode: bool) -> None: + chunks: Final = _continuous_usage_fixture_chunks() + zero: Final = chunks[0].model_copy( + update={ + "choices": [StreamingChoices(index=0, delta=Delta(role="assistant", content=""))], + "usage": Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10), + } + ) + wrapper: Final = _continuous_usage_stream( + logging_obj, (zero, *chunks[1:]), {"include_usage": True, "continuous_usage_stats": True} + ) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + assert getattr(collected[0], "usage", None) == Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_does_not_invent_intermediate_counts(logging_obj: Logging, sync_mode: bool) -> None: + chunks: Final = tuple(chunk.model_copy(update={"usage": None}) for chunk in _continuous_usage_fixture_chunks()[:2]) + wrapper: Final = _continuous_usage_stream( + logging_obj, chunks, {"include_usage": True, "continuous_usage_stats": True} + ) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + content: Final = tuple(chunk for chunk in collected if chunk.choices and chunk.choices[0].delta.content) + assert tuple(getattr(chunk, "usage", None) for chunk in content) == (None, None) + assert collected[-1].usage is not None + assert collected[-1].usage.total_tokens == ( + collected[-1].usage.prompt_tokens + collected[-1].usage.completion_tokens + ) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_only_chunk_is_visible(logging_obj: Logging, sync_mode: bool) -> None: + chunks: Final = _continuous_usage_fixture_chunks() + progress: Final = chunks[-1].model_copy( + update={"usage": Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11)} + ) + wrapper: Final = _continuous_usage_stream( + logging_obj, (chunks[0], progress, *chunks[1:]), {"include_usage": True, "continuous_usage_stats": True} + ) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + assert collected[1].choices == [StreamingChoices(index=0, delta=Delta())] + assert getattr(collected[1], "usage", None) == Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11) + assert len(collected) == 5 + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_keeps_final_trailer_contract(logging_obj: Logging, sync_mode: bool) -> None: + chunks: Final = _continuous_usage_fixture_chunks() + late: Final = chunks[-1].model_copy(update={"usage": Usage(prompt_tokens=10, completion_tokens=3, total_tokens=13)}) + wrapper: Final = _continuous_usage_stream( + logging_obj, (*chunks[:2], late), {"include_usage": True, "continuous_usage_stats": True} + ) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + assert len(collected) == 4 + assert collected[-2].choices[0].finish_reason == "stop" + assert collected[-1].usage == Usage( + prompt_tokens=10, + completion_tokens=3, + total_tokens=13, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=3), + ) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_preserves_details_and_choices(logging_obj: Logging, sync_mode: bool) -> None: + chunks: Final = _continuous_usage_fixture_chunks() + detailed: Final = chunks[0].model_copy( + update={ + "choices": [ + StreamingChoices(index=0, delta=Delta(role="assistant", content="a")), + StreamingChoices(index=1, delta=Delta(content="other")), + ], + "usage": Usage( + prompt_tokens=10, + completion_tokens=1, + total_tokens=11, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=1), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=4), + ), + } + ) + default_wrapper: Final = _continuous_usage_stream(logging_obj, (detailed, *chunks[1:]), {"include_usage": True}) + default_first: Final = next(default_wrapper) if sync_mode else await anext(default_wrapper) + await default_wrapper.aclose() + wrapper: Final = _continuous_usage_stream( + logging_obj, (detailed, *chunks[1:]), {"include_usage": True, "continuous_usage_stats": True} + ) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + assert collected[0].choices == [ + StreamingChoices(index=0, delta=Delta(role="assistant", content="a")), + StreamingChoices(index=1, delta=Delta(content="other")), + ] + assert collected[0].choices == default_first.choices + assert getattr(collected[0], "usage", None) == Usage( + prompt_tokens=10, + completion_tokens=1, + total_tokens=11, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=1), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=4), + ) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_is_available_before_completion(logging_obj: Logging, sync_mode: bool) -> None: + wrapper: Final = _continuous_usage_stream( + logging_obj, _continuous_usage_fixture_chunks(), {"include_usage": True, "continuous_usage_stats": True} + ) + first: Final = next(wrapper) if sync_mode else await anext(wrapper) + assert first.choices[0].delta.content == "a" + assert getattr(first, "usage", None) == Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11) + await wrapper.aclose() + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_streaming_handler_with_usage( From 8130a0572ef678e77d41f461b9491b8ddbabbe65 Mon Sep 17 00:00:00 2001 From: AkilanChithra Date: Mon, 12 Oct 2026 13:53:33 -0500 Subject: [PATCH 3/5] fix(streaming): retain usage updates with empty choices --- .../litellm_core_utils/streaming_handler.py | 2 ++ .../test_streaming_handler.py | 31 +++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index a7e68aab2f1..2974b7b6bbb 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1082,6 +1082,8 @@ class CustomStreamWrapper: return self._handle_special_delta_content(model_response) else: if hasattr(model_response, "usage"): + if self._should_forward_continuous_usage() and model_response.usage is not None: + return model_response self.chunks.append(model_response) return diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 9cdabb5e980..af392debc9e 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -579,6 +579,37 @@ async def test_continuous_usage_only_chunk_is_visible(logging_obj: Logging, sync assert len(collected) == 5 +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.parametrize("content", (None, "")) +@pytest.mark.parametrize("completion_tokens", (0, 1)) +@pytest.mark.parametrize("continuous", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_empty_choice_is_visible( + logging_obj: Logging, sync_mode: bool, content: str | None, completion_tokens: int, continuous: bool +) -> None: + chunks: Final = _continuous_usage_fixture_chunks() + progress: Final = chunks[0].model_copy( + update={ + "choices": [StreamingChoices(index=0, delta=Delta(content=content))], + "usage": Usage(prompt_tokens=10, completion_tokens=completion_tokens, total_tokens=10 + completion_tokens), + } + ) + wrapper: Final = _continuous_usage_stream( + logging_obj, (chunks[0], progress, *chunks[1:]), {"include_usage": True, "continuous_usage_stats": continuous} + ) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + assert len(collected) == (5 if continuous else 4) + if continuous: + assert collected[1].choices == [StreamingChoices(index=0, delta=Delta(content=content))] + assert getattr(collected[1], "usage", None) == Usage( + prompt_tokens=10, completion_tokens=completion_tokens, total_tokens=10 + completion_tokens + ) + else: + assert collected[1].choices[0].delta.content == "b" + assert getattr(collected[1], "usage", None) is None + assert collected[-1].usage.total_tokens == 12 + + @pytest.mark.parametrize("sync_mode", (True, False)) @pytest.mark.asyncio async def test_continuous_usage_keeps_final_trailer_contract(logging_obj: Logging, sync_mode: bool) -> None: From d460487333df8b304173fb48245606692a9d8945 Mon Sep 17 00:00:00 2001 From: AkilanChithra Date: Mon, 12 Oct 2026 14:01:55 -0500 Subject: [PATCH 4/5] refactor(streaming): type usage stream options --- litellm/litellm_core_utils/streaming_handler.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 2974b7b6bbb..6f87fe234b4 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -221,7 +221,7 @@ class CustomStreamWrapper: model, logging_obj: LiteLLMLoggingObject, custom_llm_provider: str | None = None, - stream_options=None, + stream_options: Mapping[str, object] | None = None, make_call: Callable | None = None, _response_headers: dict | httpx.Headers | None = None, count_prompt_tokens: Callable[[], int] | None = None, @@ -280,7 +280,9 @@ class CustomStreamWrapper: self.response_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: Mapping[str, object] | None = 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 @@ -418,7 +420,7 @@ class CustomStreamWrapper: ) self._restore_consumer_correlation_context() - def check_send_stream_usage(self, stream_options: dict | None): + def check_send_stream_usage(self, stream_options: Mapping[str, object] | None) -> bool: return stream_options is not None and stream_options.get("include_usage", False) is True def _should_forward_continuous_usage(self) -> bool: @@ -1082,7 +1084,7 @@ class CustomStreamWrapper: return self._handle_special_delta_content(model_response) else: if hasattr(model_response, "usage"): - if self._should_forward_continuous_usage() and model_response.usage is not None: + if self._should_forward_continuous_usage() and getattr(model_response, "usage", None) is not None: return model_response self.chunks.append(model_response) return From 7575d53009f66220a21f163bf32eff454fc378d3 Mon Sep 17 00:00:00 2001 From: AkilanChithra Date: Mon, 12 Oct 2026 14:38:56 -0500 Subject: [PATCH 5/5] refactor(streaming): keep usage option types local --- litellm/litellm_core_utils/streaming_handler.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6f87fe234b4..e17d11a0a6c 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -221,7 +221,7 @@ class CustomStreamWrapper: model, logging_obj: LiteLLMLoggingObject, custom_llm_provider: str | None = None, - stream_options: Mapping[str, object] | None = None, + stream_options=None, make_call: Callable | None = None, _response_headers: dict | httpx.Headers | None = None, count_prompt_tokens: Callable[[], int] | None = None, @@ -280,9 +280,7 @@ class CustomStreamWrapper: self.response_id: str | None = None self.logging_loop = None self.rules = Rules() - self.stream_options: Mapping[str, object] | None = stream_options or getattr( - logging_obj, "stream_options", None - ) + self.stream_options = 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 @@ -420,14 +418,15 @@ class CustomStreamWrapper: ) self._restore_consumer_correlation_context() - def check_send_stream_usage(self, stream_options: Mapping[str, object] | None) -> bool: + 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 _should_forward_continuous_usage(self) -> bool: + options: Final[Mapping[str, object] | None] = getattr(self, "stream_options", None) return ( self.send_stream_usage is True - and self.stream_options is not None - and self.stream_options.get("continuous_usage_stats", False) is True + and options is not None + and options.get("continuous_usage_stats", False) is True ) def check_is_function_call(self, logging_obj) -> bool: