From 1802fe7228a5076b09941f1b9d998508512c0e38 Mon Sep 17 00:00:00 2001 From: Priyansh Nandwana Date: Thu, 27 Aug 2026 16:29:18 +0530 Subject: [PATCH 1/3] fix(proxy): gate streaming cost injection on stream_options.include_usage include_cost_in_streaming_usage is a process-wide flag, so every implementing path injected usage.cost for every caller on every route, whether or not the caller asked for usage. Turning it on also disabled the streaming fast path for all traffic, including requests that never carry a usage dict. Injection now consults the caller's stream_options.include_usage alongside the global flag. An explicit include_usage: false opts out on any protocol. Anthropic Messages, Vertex rawPredict and Gemini generateContent have no such field for a caller to set, so injection stays always-on there, now documented rather than accidental. The fast path is resolved per request, so streams that will never be injected into keep it. Fixes #38348 --- litellm/proxy/common_request_processing.py | 65 +++++++++- .../pass_through_endpoints/architecture.md | 2 +- .../streaming_handler.py | 15 +-- .../test_streaming_handler_interrupt.py | 112 ++++++++++++++++- .../proxy/test_common_request_processing.py | 114 ++++++++++++++++++ 5 files changed, 292 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 64b0c6c1967..2abc893f787 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3868,6 +3868,7 @@ class ProxyBaseLLMRequestProcessing: request: Request | None = None, flush_tail: Callable[[], bytes] | None = None, seal_open_frame: Callable[[bytes], str] | None = None, + protocol_supports_stream_options: bool = True, ) -> AsyncGenerator[str, None]: """ Shared streaming data generator: runs proxy iterator hook, per-chunk hook, @@ -3883,6 +3884,10 @@ class ProxyBaseLLMRequestProcessing: relays raw upstream bytes, so an upstream that hangs up mid-frame leaves the client inside an open frame, where an error frame would be swallowed or misparsed instead of raised. + + ``protocol_supports_stream_options`` says whether this route's protocol gives + callers a ``stream_options.include_usage`` to opt in with, which gates cost + injection; see ``_should_inject_cost_for_request``. """ verbose_proxy_logger.debug("inside generator") # Resolve per-stream (not per-chunk) whether the heavy per-chunk path @@ -3893,7 +3898,10 @@ class ProxyBaseLLMRequestProcessing: # await, response-string materialization, and cost-injection call are # pure overhead on the streaming hot path (the default config). caps: Final = ProxyLogging._callback_capabilities() - cost_injection_enabled: Final = bool(getattr(litellm, "include_cost_in_streaming_usage", False)) + cost_injection_enabled: Final = ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + request_data, + protocol_supports_stream_options=protocol_supports_stream_options, + ) fast_path = not caps.has_streaming_chunk_override and not caps.has_guardrail and not cost_injection_enabled debug_enabled: Final = verbose_proxy_logger.isEnabledFor(logging.DEBUG) stream_completed = False @@ -3935,7 +3943,10 @@ class ProxyBaseLLMRequestProcessing: model_name = request_data.get("model", "") chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( - chunk, model_name, request_data.get("litellm_logging_obj") + chunk, + model_name, + request_data.get("litellm_logging_obj"), + enabled=cost_injection_enabled, ) # Set before the yield: an async generator suspends at the yield, @@ -4017,6 +4028,7 @@ class ProxyBaseLLMRequestProcessing: proxy_logging_obj: ProxyLogging, request: Request | None = None, restamp_model: str | None = None, + protocol_supports_stream_options: bool = False, ) -> AsyncGenerator[str, None]: """ Anthropic /messages streaming data generator, which requires SSE events. @@ -4044,23 +4056,59 @@ class ProxyBaseLLMRequestProcessing: request=request, flush_tail=None if restamper is None else restamper.flush, seal_open_frame=seal_open_sse_frame, + protocol_supports_stream_options=protocol_supports_stream_options, ) + @staticmethod + def _should_inject_cost_for_request( + request_data: Mapping[str, Any] | None, + *, + protocol_supports_stream_options: bool = True, + ) -> bool: + """ + Whether this request's streamed usage events should carry ``usage.cost``. + + ``litellm.include_cost_in_streaming_usage`` is process-wide, so on its own it + injects for every caller on every route. OpenAI-protocol callers opt into usage + reporting per request via ``stream_options.include_usage``, so that opt-in gates + injection too. An explicit ``include_usage: false`` opts out on any protocol. + Anthropic Messages, Vertex ``rawPredict`` and Gemini ``generateContent`` have no + such field for a caller to set, so injection stays always-on there. + """ + if not getattr(litellm, "include_cost_in_streaming_usage", False): + return False + stream_options: Final = request_data.get("stream_options") if isinstance(request_data, Mapping) else None + if isinstance(stream_options, Mapping): + return bool(stream_options.get("include_usage", False)) + return not protocol_supports_stream_options + @overload @staticmethod def _process_chunk_with_cost_injection( - chunk: bytes, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + chunk: bytes, + model_name: str, + litellm_logging_obj: LiteLLMLoggingObj | None = None, + *, + enabled: bool | None = None, ) -> bytes: ... @overload @staticmethod def _process_chunk_with_cost_injection( - chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + chunk: object, + model_name: str, + litellm_logging_obj: LiteLLMLoggingObj | None = None, + *, + enabled: bool | None = None, ) -> object: ... @staticmethod def _process_chunk_with_cost_injection( - chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + chunk: object, + model_name: str, + litellm_logging_obj: LiteLLMLoggingObj | None = None, + *, + enabled: bool | None = None, ) -> object: """ Process a streaming chunk and inject cost information if enabled. @@ -4069,11 +4117,16 @@ class ProxyBaseLLMRequestProcessing: chunk: The streaming chunk (dict, str, bytes, or bytearray) model_name: Model name for cost calculation litellm_logging_obj: The call's logging object, used for pricing + enabled: Per-stream decision from ``_should_inject_cost_for_request``. + Falls back to the global flag alone when not passed. Returns: The processed chunk with cost information injected if applicable """ - if not getattr(litellm, "include_cost_in_streaming_usage", False): + injection_enabled: Final = ( + enabled if enabled is not None else bool(getattr(litellm, "include_cost_in_streaming_usage", False)) + ) + if not injection_enabled: return chunk try: diff --git a/litellm/proxy/pass_through_endpoints/architecture.md b/litellm/proxy/pass_through_endpoints/architecture.md index f7dd8077ab5..c47f6f9054e 100644 --- a/litellm/proxy/pass_through_endpoints/architecture.md +++ b/litellm/proxy/pass_through_endpoints/architecture.md @@ -62,7 +62,7 @@ sequenceDiagram | Streaming chunk collection | Collect chunks async for logging after stream completes | | Multipart form handling | Reconstruct multipart/form-data requests for file uploads | | Guardrails (opt-in) | Run content filtering when explicitly configured | -| Cost injection | Inject cost into streaming chunks when `include_cost_in_streaming_usage` enabled | +| Cost injection | Inject cost into streaming chunks when `include_cost_in_streaming_usage` enabled. OpenAI-protocol callers must also opt in via `stream_options.include_usage`; Anthropic and Vertex `rawPredict` have no such field, so injection is always-on there unless the caller sends an explicit `stream_options.include_usage: false` | ## What Does NOT Change diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 8f0f87e6e69..02e4691324c 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -6,7 +6,6 @@ from typing import Final, Protocol import httpx -import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks @@ -172,12 +171,11 @@ class PassThroughStreamingHandler: litellm_logging_obj=litellm_logging_obj, ) - # Resolve once per stream rather than re-reading the global + - # re-branching on every chunk. ``include_cost_in_streaming_usage`` is - # set at config load and stable for the process, matching how the - # proxy-level streaming fast path resolves it. cost_injection_active: Final = ( - bool(getattr(litellm, "include_cost_in_streaming_usage", False)) + ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + request_body, + protocol_supports_stream_options=endpoint_type == EndpointType.OPENAI, + ) and bool(model_name) and ( endpoint_type in (EndpointType.ANTHROPIC, EndpointType.OPENAI) @@ -220,7 +218,10 @@ class PassThroughStreamingHandler: complete_frames, pending = split_complete_sse_frames(pending + chunk) if complete_frames: yield ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( - complete_frames, resolved_model_name, litellm_logging_obj + complete_frames, + resolved_model_name, + litellm_logging_obj, + enabled=True, ) if pending: yield pending diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 00a3606bbf5..2d6b523ac46 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -398,12 +398,20 @@ def _openai_passthrough_stream_chunks(): ] -async def _collect_openai_passthrough_chunks(chunks, endpoint_type): +async def _collect_openai_passthrough_chunks(chunks, endpoint_type, request_body=None): + # Default to the opt-in body a real caller must send for the OpenAI protocol + # to emit a usage frame at all -- cost injection is gated on that opt-in. + if request_body is None: + request_body = { + "model": "gpt-4o-mini", + "stream": True, + "stream_options": {"include_usage": True}, + } response = _make_streaming_response(chunks) received = [] async for chunk in PassThroughStreamingHandler.chunk_processor( response=response, - request_body={"model": "gpt-4o-mini", "stream": True}, + request_body=request_body, litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=endpoint_type, start_time=datetime.now(), @@ -481,6 +489,106 @@ async def test_chunk_processor_streams_crlf_delimited_frames_live_and_injects_co assert final_payload["usage"]["cost"] > 0 +@pytest.mark.asyncio +async def test_chunk_processor_skips_injection_when_openai_caller_did_not_opt_in(monkeypatch): + """Regression: issue #38348 -- ``include_cost_in_streaming_usage`` is a process-wide + flag, but OpenAI-protocol callers opt into usage reporting per request via + ``stream_options.include_usage``. A caller that never asked for usage must not have + ``usage.cost`` injected into its stream just because the flag is on proxy-wide.""" + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + chunks = _openai_passthrough_stream_chunks() + + received = await _collect_openai_passthrough_chunks( + chunks, + EndpointType.OPENAI, + request_body={"model": "gpt-4o-mini", "stream": True}, + ) + + assert received == chunks + + +@pytest.mark.asyncio +async def test_chunk_processor_respects_explicit_include_usage_false(monkeypatch): + """Regression: issue #38348 -- an explicit ``include_usage: false`` is a caller + opting out, and must be honoured even with the global flag on.""" + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + chunks = _openai_passthrough_stream_chunks() + + received = await _collect_openai_passthrough_chunks( + chunks, + EndpointType.OPENAI, + request_body={ + "model": "gpt-4o-mini", + "stream": True, + "stream_options": {"include_usage": False}, + }, + ) + + assert received == chunks + + +@pytest.mark.asyncio +async def test_chunk_processor_anthropic_injects_without_stream_options(monkeypatch): + """The Anthropic Messages protocol has no ``stream_options`` for a caller to opt in + with, so injection stays always-on there while the flag is set -- issue #38348 asks + for that behaviour to be explicit rather than accidental.""" + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + frame = ( + b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' + b'"usage":{"input_tokens":11,"output_tokens":4}}\n\n' + ) + response = _make_streaming_response([frame]) + + received = [] + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-haiku-4-5", "stream": True}, + litellm_logging_obj=MagicMock(), + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + route_streaming_logging=AsyncMock(), + ): + received.append(chunk) + await asyncio.sleep(0) + + payload = json.loads(b"".join(received).decode("utf-8").split("data:", 1)[1].strip()) + assert payload["usage"]["cost"] > 0 + + +@pytest.mark.asyncio +async def test_chunk_processor_anthropic_respects_explicit_opt_out(monkeypatch): + """Even on Anthropic, a caller that explicitly sends ``include_usage: false`` opts + out of cost injection.""" + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + frame = ( + b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' + b'"usage":{"input_tokens":11,"output_tokens":4}}\n\n' + ) + response = _make_streaming_response([frame]) + + received = [] + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={ + "model": "claude-haiku-4-5", + "stream": True, + "stream_options": {"include_usage": False}, + }, + litellm_logging_obj=MagicMock(), + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + route_streaming_logging=AsyncMock(), + ): + received.append(chunk) + await asyncio.sleep(0) + + assert received == [frame] + + @pytest.mark.asyncio async def test_chunk_processor_flag_off_leaves_openai_passthrough_stream_byte_identical(monkeypatch): monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index c17f41a8b8f..520a1faa962 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -10226,3 +10226,117 @@ class TestStreamingContainerOwnershipRecordedBeforeDone: assert tuple(chunk for chunk, _ in observed) == self.CHUNKS assert tuple(count for _, count in observed) == (0, 0, 0, 0) recorder.assert_awaited_once() + + +class TestShouldInjectCostForRequest: + """Issue #38348: ``include_cost_in_streaming_usage`` is a process-wide flag, so on its + own it injects ``usage.cost`` for every caller on every route. Injection must also + consult the caller's per-request ``stream_options.include_usage`` opt-in.""" + + def test_global_flag_off_never_injects(self, monkeypatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + assert ( + ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + {"stream_options": {"include_usage": True}} + ) + is False + ) + + def test_openai_protocol_requires_caller_opt_in(self, monkeypatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + assert ( + ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + {"model": "gpt-4o-mini", "stream": True}, + protocol_supports_stream_options=True, + ) + is False + ) + + def test_openai_protocol_opted_in_injects(self, monkeypatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + assert ( + ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + {"stream_options": {"include_usage": True}}, + protocol_supports_stream_options=True, + ) + is True + ) + + @pytest.mark.parametrize("protocol_supports_stream_options", [True, False]) + def test_explicit_opt_out_is_honoured_on_every_protocol(self, monkeypatch, protocol_supports_stream_options): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + assert ( + ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + {"stream_options": {"include_usage": False}}, + protocol_supports_stream_options=protocol_supports_stream_options, + ) + is False + ) + + def test_protocol_without_stream_options_stays_always_on(self, monkeypatch): + """Anthropic Messages / Vertex rawPredict / Gemini give a caller no way to opt + in, so the flag remains always-on there.""" + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + assert ( + ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + {"model": "claude-haiku-4-5", "stream": True}, + protocol_supports_stream_options=False, + ) + is True + ) + + def test_missing_request_data_falls_back_to_protocol_default(self, monkeypatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + assert ProxyBaseLLMRequestProcessing._should_inject_cost_for_request(None) is False + assert ( + ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + None, protocol_supports_stream_options=False + ) + is True + ) + + def test_malformed_stream_options_falls_back_to_protocol_default(self, monkeypatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + assert ( + ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + {"stream_options": "include_usage"}, + protocol_supports_stream_options=False, + ) + is True + ) + + +class TestProcessChunkCostInjectionGate: + """``_process_chunk_with_cost_injection`` takes the per-stream decision from + ``_should_inject_cost_for_request`` and falls back to the global flag when the + caller does not pass one.""" + + @staticmethod + def _usage_chunk(): + return { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "choices": [], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + + def test_enabled_false_leaves_chunk_untouched(self, monkeypatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + chunk = self._usage_chunk() + assert ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + chunk, "gpt-4o-mini", None, enabled=False + ) is chunk + + def test_enabled_true_injects_even_with_global_flag_off(self, monkeypatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + result = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + self._usage_chunk(), "gpt-4o-mini", None, enabled=True + ) + assert result["usage"]["cost"] > 0 + + def test_omitted_enabled_falls_back_to_global_flag(self, monkeypatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + result = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + self._usage_chunk(), "gpt-4o-mini" + ) + assert result["usage"]["cost"] > 0 From 9869fd61438ab69ba345b8088efd79d2d210425d Mon Sep 17 00:00:00 2001 From: Priyansh Nandwana Date: Thu, 27 Aug 2026 16:48:11 +0530 Subject: [PATCH 2/3] fix(proxy): make the streaming cost-injection gate public The gate is called from the passthrough streaming handler as well as from common_request_processing, and a private cross-module call pushes basedpyright's reportPrivateUsage over its budget. It is a shared decision helper, so make it public rather than suppressing the rule. The two budget-reservation slow-path tests used include_cost_in_streaming_usage alone to force fast_path off. The gate now also needs the caller's stream_options.include_usage opt-in for that, so they pass one. --- litellm/proxy/common_request_processing.py | 8 ++++---- .../streaming_handler.py | 2 +- .../proxy/test_budget_reservation.py | 15 +++++++++------ .../proxy/test_common_request_processing.py | 18 +++++++++--------- 4 files changed, 23 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2abc893f787..153ac7f3d83 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3887,7 +3887,7 @@ class ProxyBaseLLMRequestProcessing: ``protocol_supports_stream_options`` says whether this route's protocol gives callers a ``stream_options.include_usage`` to opt in with, which gates cost - injection; see ``_should_inject_cost_for_request``. + injection; see ``should_inject_cost_for_request``. """ verbose_proxy_logger.debug("inside generator") # Resolve per-stream (not per-chunk) whether the heavy per-chunk path @@ -3898,7 +3898,7 @@ class ProxyBaseLLMRequestProcessing: # await, response-string materialization, and cost-injection call are # pure overhead on the streaming hot path (the default config). caps: Final = ProxyLogging._callback_capabilities() - cost_injection_enabled: Final = ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + cost_injection_enabled: Final = ProxyBaseLLMRequestProcessing.should_inject_cost_for_request( request_data, protocol_supports_stream_options=protocol_supports_stream_options, ) @@ -4060,7 +4060,7 @@ class ProxyBaseLLMRequestProcessing: ) @staticmethod - def _should_inject_cost_for_request( + def should_inject_cost_for_request( request_data: Mapping[str, Any] | None, *, protocol_supports_stream_options: bool = True, @@ -4117,7 +4117,7 @@ class ProxyBaseLLMRequestProcessing: chunk: The streaming chunk (dict, str, bytes, or bytearray) model_name: Model name for cost calculation litellm_logging_obj: The call's logging object, used for pricing - enabled: Per-stream decision from ``_should_inject_cost_for_request``. + enabled: Per-stream decision from ``should_inject_cost_for_request``. Falls back to the global flag alone when not passed. Returns: diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 02e4691324c..bebd92a96d0 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -172,7 +172,7 @@ class PassThroughStreamingHandler: ) cost_injection_active: Final = ( - ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + ProxyBaseLLMRequestProcessing.should_inject_cost_for_request( request_body, protocol_supports_stream_options=endpoint_type == EndpointType.OPENAI, ) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index b3913079bb2..a9c95d078e6 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -78,11 +78,12 @@ def spend_counter_state(): ps.prisma_client = original_prisma_client -def _request_body() -> dict: +def _request_body(*, include_usage: bool = False) -> dict: return { "model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 10, + **({"stream_options": {"include_usage": True}} if include_usage else {}), } @@ -3115,14 +3116,15 @@ async def test_streaming_cancel_in_slow_path_before_yield_refunds(spend_counter_ generator = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=MagicMock(), user_api_key_dict=valid_token, - request_data=_request_body(), + request_data=_request_body(include_usage=True), proxy_logging_obj=streaming_logging_obj, serialize_chunk=lambda chunk: chunk, serialize_error=lambda exc: str(exc), ) received = [] - # include_cost_in_streaming_usage forces fast_path off, so the hook above runs + # include_cost_in_streaming_usage plus the caller's stream_options.include_usage + # opt-in forces fast_path off, so the hook above runs with patch.object(litellm, "include_cost_in_streaming_usage", True, create=True): async def _drain(): async for chunk in generator: @@ -3191,15 +3193,16 @@ async def test_streaming_slow_path_processes_and_yields_chunk(spend_counter_stat generator = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=MagicMock(), user_api_key_dict=valid_token, - request_data=_request_body(), + request_data=_request_body(include_usage=True), proxy_logging_obj=streaming_logging_obj, serialize_chunk=lambda chunk: chunk, serialize_error=lambda exc: str(exc), ) received = [] - # include_cost_in_streaming_usage forces the slow path so the per-chunk hook, - # content accumulation, and cost-injection branch all run to a successful yield + # include_cost_in_streaming_usage plus the caller's stream_options.include_usage + # opt-in forces the slow path, so the per-chunk hook, content accumulation, and + # cost-injection branch all run to a successful yield with patch.object(litellm, "include_cost_in_streaming_usage", True, create=True): async for chunk in generator: received.append(chunk) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 520a1faa962..65fe75f16b6 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -10236,7 +10236,7 @@ class TestShouldInjectCostForRequest: def test_global_flag_off_never_injects(self, monkeypatch): monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) assert ( - ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + ProxyBaseLLMRequestProcessing.should_inject_cost_for_request( {"stream_options": {"include_usage": True}} ) is False @@ -10245,7 +10245,7 @@ class TestShouldInjectCostForRequest: def test_openai_protocol_requires_caller_opt_in(self, monkeypatch): monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) assert ( - ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + ProxyBaseLLMRequestProcessing.should_inject_cost_for_request( {"model": "gpt-4o-mini", "stream": True}, protocol_supports_stream_options=True, ) @@ -10255,7 +10255,7 @@ class TestShouldInjectCostForRequest: def test_openai_protocol_opted_in_injects(self, monkeypatch): monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) assert ( - ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + ProxyBaseLLMRequestProcessing.should_inject_cost_for_request( {"stream_options": {"include_usage": True}}, protocol_supports_stream_options=True, ) @@ -10266,7 +10266,7 @@ class TestShouldInjectCostForRequest: def test_explicit_opt_out_is_honoured_on_every_protocol(self, monkeypatch, protocol_supports_stream_options): monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) assert ( - ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + ProxyBaseLLMRequestProcessing.should_inject_cost_for_request( {"stream_options": {"include_usage": False}}, protocol_supports_stream_options=protocol_supports_stream_options, ) @@ -10278,7 +10278,7 @@ class TestShouldInjectCostForRequest: in, so the flag remains always-on there.""" monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) assert ( - ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + ProxyBaseLLMRequestProcessing.should_inject_cost_for_request( {"model": "claude-haiku-4-5", "stream": True}, protocol_supports_stream_options=False, ) @@ -10287,9 +10287,9 @@ class TestShouldInjectCostForRequest: def test_missing_request_data_falls_back_to_protocol_default(self, monkeypatch): monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) - assert ProxyBaseLLMRequestProcessing._should_inject_cost_for_request(None) is False + assert ProxyBaseLLMRequestProcessing.should_inject_cost_for_request(None) is False assert ( - ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + ProxyBaseLLMRequestProcessing.should_inject_cost_for_request( None, protocol_supports_stream_options=False ) is True @@ -10298,7 +10298,7 @@ class TestShouldInjectCostForRequest: def test_malformed_stream_options_falls_back_to_protocol_default(self, monkeypatch): monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) assert ( - ProxyBaseLLMRequestProcessing._should_inject_cost_for_request( + ProxyBaseLLMRequestProcessing.should_inject_cost_for_request( {"stream_options": "include_usage"}, protocol_supports_stream_options=False, ) @@ -10308,7 +10308,7 @@ class TestShouldInjectCostForRequest: class TestProcessChunkCostInjectionGate: """``_process_chunk_with_cost_injection`` takes the per-stream decision from - ``_should_inject_cost_for_request`` and falls back to the global flag when the + ``should_inject_cost_for_request`` and falls back to the global flag when the caller does not pass one.""" @staticmethod From bcadce93671992b35817b9c095ea72c80e0df011 Mon Sep 17 00:00:00 2001 From: Priyansh Nandwana Date: Mon, 14 Sep 2026 11:26:06 +0530 Subject: [PATCH 3/3] test(proxy): drop the restating docstrings and the request_body rebinding The test names already say what each case covers, and the collect helper now builds the opted-in body inline instead of reassigning its parameter. --- .../test_streaming_handler_interrupt.py | 25 ++++--------------- .../proxy/test_common_request_processing.py | 10 +------- 2 files changed, 6 insertions(+), 29 deletions(-) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 2d6b523ac46..c347d438a4a 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -398,20 +398,16 @@ def _openai_passthrough_stream_chunks(): ] +def _openai_opted_in_body(): + return {"model": "gpt-4o-mini", "stream": True, "stream_options": {"include_usage": True}} + + async def _collect_openai_passthrough_chunks(chunks, endpoint_type, request_body=None): - # Default to the opt-in body a real caller must send for the OpenAI protocol - # to emit a usage frame at all -- cost injection is gated on that opt-in. - if request_body is None: - request_body = { - "model": "gpt-4o-mini", - "stream": True, - "stream_options": {"include_usage": True}, - } response = _make_streaming_response(chunks) received = [] async for chunk in PassThroughStreamingHandler.chunk_processor( response=response, - request_body=request_body, + request_body=_openai_opted_in_body() if request_body is None else request_body, litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=endpoint_type, start_time=datetime.now(), @@ -491,10 +487,6 @@ async def test_chunk_processor_streams_crlf_delimited_frames_live_and_injects_co @pytest.mark.asyncio async def test_chunk_processor_skips_injection_when_openai_caller_did_not_opt_in(monkeypatch): - """Regression: issue #38348 -- ``include_cost_in_streaming_usage`` is a process-wide - flag, but OpenAI-protocol callers opt into usage reporting per request via - ``stream_options.include_usage``. A caller that never asked for usage must not have - ``usage.cost`` injected into its stream just because the flag is on proxy-wide.""" monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) chunks = _openai_passthrough_stream_chunks() @@ -509,8 +501,6 @@ async def test_chunk_processor_skips_injection_when_openai_caller_did_not_opt_in @pytest.mark.asyncio async def test_chunk_processor_respects_explicit_include_usage_false(monkeypatch): - """Regression: issue #38348 -- an explicit ``include_usage: false`` is a caller - opting out, and must be honoured even with the global flag on.""" monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) chunks = _openai_passthrough_stream_chunks() @@ -529,9 +519,6 @@ async def test_chunk_processor_respects_explicit_include_usage_false(monkeypatch @pytest.mark.asyncio async def test_chunk_processor_anthropic_injects_without_stream_options(monkeypatch): - """The Anthropic Messages protocol has no ``stream_options`` for a caller to opt in - with, so injection stays always-on there while the flag is set -- issue #38348 asks - for that behaviour to be explicit rather than accidental.""" monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) frame = ( b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' @@ -559,8 +546,6 @@ async def test_chunk_processor_anthropic_injects_without_stream_options(monkeypa @pytest.mark.asyncio async def test_chunk_processor_anthropic_respects_explicit_opt_out(monkeypatch): - """Even on Anthropic, a caller that explicitly sends ``include_usage: false`` opts - out of cost injection.""" monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) frame = ( b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 65fe75f16b6..047fa297b5c 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -10229,9 +10229,7 @@ class TestStreamingContainerOwnershipRecordedBeforeDone: class TestShouldInjectCostForRequest: - """Issue #38348: ``include_cost_in_streaming_usage`` is a process-wide flag, so on its - own it injects ``usage.cost`` for every caller on every route. Injection must also - consult the caller's per-request ``stream_options.include_usage`` opt-in.""" + """Issue #38348: cost injection honours the caller's stream_options.include_usage.""" def test_global_flag_off_never_injects(self, monkeypatch): monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) @@ -10274,8 +10272,6 @@ class TestShouldInjectCostForRequest: ) def test_protocol_without_stream_options_stays_always_on(self, monkeypatch): - """Anthropic Messages / Vertex rawPredict / Gemini give a caller no way to opt - in, so the flag remains always-on there.""" monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) assert ( ProxyBaseLLMRequestProcessing.should_inject_cost_for_request( @@ -10307,10 +10303,6 @@ class TestShouldInjectCostForRequest: class TestProcessChunkCostInjectionGate: - """``_process_chunk_with_cost_injection`` takes the per-stream decision from - ``should_inject_cost_for_request`` and falls back to the global flag when the - caller does not pass one.""" - @staticmethod def _usage_chunk(): return {