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.
This commit is contained in:
mateo-berri 2026-07-30 17:01:16 -07:00
parent 43efd02d35
commit a770b437d5
3 changed files with 191 additions and 9 deletions

View file

@ -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

View file

@ -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"] = {

View file

@ -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