From 3a17303fdfe71956dfb59f6010d8a92be2011c52 Mon Sep 17 00:00:00 2001 From: Priyansh Nandwana Date: Thu, 27 Aug 2026 16:29:18 +0530 Subject: [PATCH 1/2] 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 | 64 +++++++++- .../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, 291 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f25fa46197e..8d41b6fb1ea 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3585,6 +3585,7 @@ class ProxyBaseLLMRequestProcessing: serialize_error: StreamErrorSerializer, request: Request | None = None, flush_tail: Callable[[], bytes] | None = None, + protocol_supports_stream_options: bool = True, ) -> AsyncGenerator[str, None]: """ Shared streaming data generator: runs proxy iterator hook, per-chunk hook, @@ -3594,6 +3595,9 @@ class ProxyBaseLLMRequestProcessing: ``flush_tail`` runs once after the upstream iterator completes cleanly and its non-empty result is yielded, so a serializer that buffers bytes across chunks can emit anything still held at end of stream. + ``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 @@ -3604,7 +3608,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 @@ -3645,7 +3652,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, @@ -3723,6 +3733,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 and Google /generateContent streaming data generator require SSE events. @@ -3748,23 +3759,59 @@ class ProxyBaseLLMRequestProcessing: ), request=request, flush_tail=None if restamper is None else restamper.flush, + 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. @@ -3773,11 +3820,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 4be0235adbb..238fa89352f 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.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -162,12 +161,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) @@ -199,7 +197,10 @@ class PassThroughStreamingHandler: ) # rebind-ok: SSE frame reassembly buffer across transport chunks 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 ea6adc35b9a..9fe3b97e9eb 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 @@ -392,12 +392,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(), @@ -475,6 +483,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 f7fe6ad9d39..cbd04dda0de 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -7945,3 +7945,117 @@ class TestDetachedStreamFailureHook: await logging_obj._on_detached_stream_failure(failure) assert [call["original_exception"] for call in recorder.calls] == [failure] + + +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 884fefa6f927e175f8415b1454a991c333cffb30 Mon Sep 17 00:00:00 2001 From: Priyansh Nandwana Date: Thu, 27 Aug 2026 16:48:11 +0530 Subject: [PATCH 2/2] 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 8d41b6fb1ea..464d16b752d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3597,7 +3597,7 @@ class ProxyBaseLLMRequestProcessing: chunks can emit anything still held at end of stream. ``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 @@ -3608,7 +3608,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, ) @@ -3763,7 +3763,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, @@ -3820,7 +3820,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 238fa89352f..0b8febde728 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -162,7 +162,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 b8fb6170d34..fb379fd541c 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -73,11 +73,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 {}), } @@ -2717,14 +2718,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: @@ -2793,15 +2795,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 cbd04dda0de..ee8749a7c9b 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -7955,7 +7955,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 @@ -7964,7 +7964,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, ) @@ -7974,7 +7974,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, ) @@ -7985,7 +7985,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, ) @@ -7997,7 +7997,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, ) @@ -8006,9 +8006,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 @@ -8017,7 +8017,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, ) @@ -8027,7 +8027,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