Merge bcadce9367 into dab2deb5ed

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:

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
 

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

tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py

@@ -398,12 +398,16 @@ def _openai_passthrough_stream_chunks():
     ]
 
 
-async def _collect_openai_passthrough_chunks(chunks, endpoint_type):
+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):
     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=_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(),
@@ -481,6 +485,95 @@ 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):
+    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):
+    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):
+    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):
+    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)

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)

tests/test_litellm/proxy/test_common_request_processing.py

@@ -10226,3 +10226,109 @@ 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: 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)
+        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):
+        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:
+    @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