From a770b437d53f4ddd03a4bf24c113268ac9d36dc4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 30 Jul 2026 17:01:16 -0700 Subject: [PATCH] fix(proxy): gate default stream usage injection on provider support and neutralize client-sent strip marker Bytez and OCI param maps raise on stream_options when drop_params is unset, so the default injection would have broken every streamed chat completion routed to them. Injection now only happens when every router deployment behind the requested model (wildcards and aliases included) declares stream_options in its supported OpenAI params; providers that do not declare it either reject the param or already stream usage natively, so skipping them keeps old behavior instead of erroring. _litellm_strip_stream_usage arriving in the client request body is now overwritten at ingress (and popped in the experimental queue endpoint), so a client can no longer suppress the usage chunk it explicitly requested by planting the internal marker. --- litellm/proxy/common_request_processing.py | 59 +++++++- litellm/proxy/proxy_server.py | 1 + .../proxy/test_common_request_processing.py | 140 +++++++++++++++++- 3 files changed, 191 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 750954b69f7..d25ef8a3039 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -5,6 +5,7 @@ import math import time import traceback from datetime import datetime +from functools import lru_cache from typing import ( TYPE_CHECKING, Any, @@ -39,6 +40,9 @@ from litellm.constants import ( ) from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer +from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, @@ -245,25 +249,64 @@ async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[Any]"]) -> None pass +@lru_cache(maxsize=512) +def _litellm_model_supports_stream_options(litellm_model: str) -> bool: + try: + supported_params = get_supported_openai_params(model=litellm_model) + except Exception: # noqa: BLE001 # unmapped or malformed model strings must disable injection, not fail the request + return False + return supported_params is not None and "stream_options" in supported_params + + +def _deployment_litellm_model(deployment: Mapping[str, object]) -> Optional[str]: + litellm_params = deployment.get("litellm_params") + if isinstance(litellm_params, Mapping): + litellm_model = litellm_params.get("model") + else: + litellm_model = getattr(litellm_params, "model", None) + return litellm_model if isinstance(litellm_model, str) else None + + +def _model_deployments_support_stream_options( + model: object, + llm_router: Optional[Router], +) -> bool: + if not isinstance(model, str): + return False + deployments = llm_router.get_model_list(model_name=model) if llm_router is not None else None + deployment_models = tuple( + litellm_model + for deployment in deployments or () + for litellm_model in (_deployment_litellm_model(deployment),) + if litellm_model is not None + ) + candidate_models = deployment_models if deployment_models else (model,) + return all(_litellm_model_supports_stream_options(m) for m in candidate_models) + + def _stream_usage_tracking_updates( data: Mapping[str, object], general_settings: Mapping[str, object], route_type: str, + supports_stream_options: Callable[[], bool], ) -> Mapping[str, object]: + scrub = {"_litellm_strip_stream_usage": False} if "_litellm_strip_stream_usage" in data else {} if data.get("stream", False) is not True: - return {} + return scrub always_include = general_settings.get("always_include_stream_usage") stream_options = data.get("stream_options") if always_include is True: if "stream_options" not in data: - return {"stream_options": {"include_usage": True}} + return {**scrub, "stream_options": {"include_usage": True}} if isinstance(stream_options, dict) and "include_usage" not in stream_options: - return {"stream_options": {**stream_options, "include_usage": True}} - return {} + return {**scrub, "stream_options": {**stream_options, "include_usage": True}} + return scrub if always_include is False or route_type != "acompletion": - return {} + return scrub if isinstance(stream_options, dict) and stream_options.get("include_usage") is True: - return {} + return scrub + if not supports_stream_options(): + return scrub merged_stream_options = {**stream_options} if isinstance(stream_options, dict) else {} return { "stream_options": {**merged_stream_options, "include_usage": True}, @@ -1264,6 +1307,10 @@ class ProxyBaseLLMRequestProcessing: data=self.data, general_settings=general_settings, route_type=route_type, + supports_stream_options=lambda: _model_deployments_support_stream_options( + model=self.data.get("model"), + llm_router=llm_router, + ), ) ) ### CALL HOOKS ### - modify/reject incoming data before calling the model diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ab55e885cda..a60ea2da019 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13500,6 +13500,7 @@ async def async_queue_request( data = {} try: data = await request.json() # type: ignore + data.pop("_litellm_strip_stream_usage", None) # Include original request and headers in the data data["proxy_server_request"] = { diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 24b2007217e..28fb97006d8 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1,7 +1,7 @@ import asyncio import copy import datetime -from typing import AsyncGenerator, Optional +from typing import AsyncGenerator, Callable, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -5113,10 +5113,22 @@ class TestStreamingClientDisconnectBilling: proxy_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() -def _apply_stream_usage_tracking(data: dict, general_settings: dict, route_type: str) -> None: +def _apply_stream_usage_tracking( + data: dict, + general_settings: dict, + route_type: str, + supports_stream_options: Callable[[], bool] = lambda: True, +) -> None: from litellm.proxy.common_request_processing import _stream_usage_tracking_updates - data.update(_stream_usage_tracking_updates(data=data, general_settings=general_settings, route_type=route_type)) + data.update( + _stream_usage_tracking_updates( + data=data, + general_settings=general_settings, + route_type=route_type, + supports_stream_options=supports_stream_options, + ) + ) class TestApplyStreamUsageTracking: @@ -5203,3 +5215,125 @@ class TestApplyStreamUsageTracking: assert "stream_options" not in data assert "_litellm_strip_stream_usage" not in data + + def test_default_skips_injection_when_provider_lacks_stream_options_support(self): + data = {"stream": True, "model": "bytez-model"} + + _apply_stream_usage_tracking( + data=data, + general_settings={}, + route_type="acompletion", + supports_stream_options=lambda: False, + ) + + assert "stream_options" not in data + assert "_litellm_strip_stream_usage" not in data + + def test_client_supplied_strip_marker_is_neutralized(self): + data = { + "stream": True, + "stream_options": {"include_usage": True}, + "_litellm_strip_stream_usage": True, + } + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion") + + assert data["_litellm_strip_stream_usage"] is False + assert data["stream_options"] == {"include_usage": True} + + def test_client_supplied_strip_marker_is_neutralized_with_flag_true(self): + data = { + "stream": True, + "stream_options": {"include_usage": True}, + "_litellm_strip_stream_usage": True, + } + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["_litellm_strip_stream_usage"] is False + + def test_client_supplied_strip_marker_is_neutralized_on_non_streaming_request(self): + data = {"_litellm_strip_stream_usage": True} + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion") + + assert data["_litellm_strip_stream_usage"] is False + + +class TestModelDeploymentsSupportStreamOptions: + def _support(self, model, llm_router=None) -> bool: + from litellm.proxy.common_request_processing import ( + _model_deployments_support_stream_options, + ) + + return _model_deployments_support_stream_options(model=model, llm_router=llm_router) + + def test_openai_compatible_deployment_supports_stream_options(self): + router = litellm.Router( + model_list=[ + { + "model_name": "azure-nano", + "litellm_params": { + "model": "azure/gpt-5.4-nano", + "api_key": "fake", + "api_base": "https://example.openai.azure.com", + }, + } + ] + ) + + assert self._support("azure-nano", router) is True + + def test_deployment_on_provider_rejecting_stream_options_is_not_injected(self): + router = litellm.Router( + model_list=[ + { + "model_name": "tiny", + "litellm_params": {"model": "bytez/openai-community/gpt2", "api_key": "fake"}, + } + ] + ) + + assert self._support("tiny", router) is False + + def test_mixed_provider_model_group_is_not_injected(self): + router = litellm.Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "oci/cohere.command-r-plus", "api_key": "fake"}, + }, + ] + ) + + assert self._support("mixed", router) is False + + def test_wildcard_route_resolves_provider_support(self): + router = litellm.Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "fake"}, + } + ] + ) + + assert self._support("openai/gpt-4o", router) is True + + def test_provider_prefixed_model_without_router_is_resolved_directly(self): + assert self._support("openai/gpt-4o", None) is True + assert self._support("bytez/openai-community/gpt2", None) is False + + def test_unmapped_model_name_is_not_injected(self): + assert self._support("some-unmapped-public-alias", None) is False + + def test_non_string_model_is_not_injected(self): + assert self._support(None, None) is False