scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 178

fix(router): emit x-litellm-overhead-duration-ms header for streaming requests (#22027)

Browse source 
* fix(router): preserve _hidden_params in FallbackStreamWrapper so x-litellm-overhead-duration-ms is emitted for streaming requests

* test(router): add regression test for FallbackStreamWrapper _hidden_params preservation
This commit is contained in:
Ishaan Jaff 2026-02-24 11:56:16 -08:00 • committed by Sameer Kankute
parent e0818b1696
commit ba3f30f04d
3 changed files with 269 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/router.py
View file

@ -1579,6 +1579,9 @@ class Router:
logging_obj=model_response.logging_obj,
)
self._async_generator = async_generator
# Preserve hidden params (including litellm_overhead_time_ms) from original response
if hasattr(model_response, "_hidden_params"):
self._hidden_params = model_response._hidden_params.copy()
def __aiter__(self):
return self
@ -7002,9 +7005,9 @@ class Router:
raise ValueError("Deployment not found")
## GET BASE MODEL
base_model = deployment.get("model_info", {}).get("base_model", None)
base_model = (deployment.get("model_info") or {}).get("base_model", None)
if base_model is None:
base_model = deployment.get("litellm_params", {}).get("base_model", None)
base_model = (deployment.get("litellm_params") or {}).get("base_model", None)
model = base_model
@ -7019,7 +7022,7 @@ class Router:
raise ValueError(
f"Deployment missing valid litellm_params. "
f"Got: {type(litellm_params_data).__name__}, "
f"deployment_id: {deployment.get('model_info', {}).get('id', 'unknown')}"
f"deployment_id: {(deployment.get('model_info') or {}).get('id', 'unknown')}"
)
_model, custom_llm_provider, _, _ = litellm.get_llm_provider(
model=litellm_params.model,
@ -7039,10 +7042,10 @@ class Router:
if potential_models is not None:
for potential_model in potential_models:
try:
if potential_model.get("model_info", {}).get(
if (potential_model.get("model_info") or {}).get(
"id"
) == deployment.get("model_info", {}).get("id"):
model = potential_model.get("litellm_params", {}).get(
) == (deployment.get("model_info") or {}).get("id"):
model = (potential_model.get("litellm_params") or {}).get(
"model"
)
break
@ -7063,9 +7066,10 @@ class Router:
model_info = litellm.get_model_info(model=model_info_name)
## CHECK USER SET MODEL INFO
user_model_info = deployment.get("model_info", {})
user_model_info = deployment.get("model_info") or {}
model_info.update(user_model_info)
if model_info is not None:
model_info.update(user_model_info)
return model_info

201
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -1,4 +1,6 @@
import copy
import datetime
from typing import AsyncGenerator
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -1348,3 +1350,202 @@ class TestOverrideOpenAIResponseModel:
# Verify the model was not changed
assert response_obj.model == fallback_model
class TestStreamingOverheadHeader:
"""
Tests that x-litellm-overhead-duration-ms is emitted in streaming responses.
Regression tests for: streaming requests not including overhead header.
"""
def test_get_custom_headers_includes_overhead_when_set(self):
"""
get_custom_headers() returns x-litellm-overhead-duration-ms
when litellm_overhead_time_ms is in hidden_params.
"""
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key_dict.tpm_limit = None
mock_user_api_key_dict.rpm_limit = None
mock_user_api_key_dict.max_budget = None
mock_user_api_key_dict.spend = 0.0
mock_user_api_key_dict.allowed_model_region = None
hidden_params = {
"litellm_overhead_time_ms": 42.5,
"_response_ms": 500.0,
"model_id": "test-model-id",
"api_base": "https://api.openai.com",
}
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=mock_user_api_key_dict,
call_id="test-call-id",
model_id="test-model-id",
cache_key="",
api_base="https://api.openai.com",
version="1.0.0",
response_cost=0.001,
model_region="",
hidden_params=hidden_params,
)
assert "x-litellm-overhead-duration-ms" in headers
assert headers["x-litellm-overhead-duration-ms"] == "42.5"
def test_get_custom_headers_omits_overhead_when_none(self):
"""
get_custom_headers() omits x-litellm-overhead-duration-ms
when litellm_overhead_time_ms is not in hidden_params.
"""
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key_dict.tpm_limit = None
mock_user_api_key_dict.rpm_limit = None
mock_user_api_key_dict.max_budget = None
mock_user_api_key_dict.spend = 0.0
mock_user_api_key_dict.allowed_model_region = None
hidden_params = {
"_response_ms": 500.0,
"model_id": "test-model-id",
}
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=mock_user_api_key_dict,
call_id="test-call-id",
model_id="test-model-id",
cache_key="",
api_base="https://api.openai.com",
version="1.0.0",
response_cost=0.001,
model_region="",
hidden_params=hidden_params,
)
# Should be absent (None gets filtered by exclude_values)
assert "x-litellm-overhead-duration-ms" not in headers
def test_update_response_metadata_sets_overhead_on_stream_wrapper(self):
"""
update_response_metadata() sets litellm_overhead_time_ms on
a streaming response's _hidden_params when llm_api_duration_ms is available.
"""
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
update_response_metadata,
)
# Mock the logging object with llm_api_duration_ms set
mock_logging_obj = MagicMock()
mock_logging_obj.model_call_details = {
"llm_api_duration_ms": 200.0,
"litellm_params": {},
}
mock_logging_obj.caching_details = None
mock_logging_obj.callback_duration_ms = None
mock_logging_obj.litellm_call_id = "test-call-id"
mock_logging_obj._response_cost_calculator = MagicMock(return_value=0.001)
# Simulate a streaming result object with _hidden_params (like CustomStreamWrapper)
stream_result = MagicMock()
stream_result._hidden_params = {
"model_id": "test-model-id",
"api_base": "https://api.openai.com",
"additional_headers": {},
}
start_time = datetime.datetime.now() - datetime.timedelta(milliseconds=300)
end_time = datetime.datetime.now()
update_response_metadata(
result=stream_result,
logging_obj=mock_logging_obj,
model="gpt-4o",
kwargs={},
start_time=start_time,
end_time=end_time,
)
assert "litellm_overhead_time_ms" in stream_result._hidden_params
overhead = stream_result._hidden_params["litellm_overhead_time_ms"]
assert overhead is not None
assert isinstance(overhead, float)
# overhead = total_response_ms (~300ms) - llm_api_duration_ms (200ms) = ~100ms
assert overhead > 0
@pytest.mark.asyncio
async def test_streaming_response_includes_overhead_header(self):
"""
StreamingResponse returned by create_response() includes
x-litellm-overhead-duration-ms in its headers.
"""
async def mock_generator() -> AsyncGenerator[str, None]:
yield 'data: {"id":"chatcmpl-test","choices":[{"delta":{"content":"hi"}}]}\n\n'
yield "data: [DONE]\n\n"
headers = {
"x-litellm-overhead-duration-ms": "42.5",
"x-litellm-call-id": "test-call-id",
"x-litellm-model-id": "test-model-id",
}
response = await create_response(
generator=mock_generator(),
media_type="text/event-stream",
headers=headers,
)
assert isinstance(response, StreamingResponse)
assert response.headers.get("x-litellm-overhead-duration-ms") == "42.5"
def test_streaming_overhead_header_in_custom_headers_from_stream_hidden_params(
self,
):
"""
Verifies that when get_custom_headers() is called with a streaming
response's hidden_params (containing litellm_overhead_time_ms),
the x-litellm-overhead-duration-ms header is correctly populated.
This tests the critical path: update_response_metadata sets the value
→ get_custom_headers reads it → StreamingResponse header is set.
"""
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key_dict.tpm_limit = None
mock_user_api_key_dict.rpm_limit = None
mock_user_api_key_dict.max_budget = None
mock_user_api_key_dict.spend = 0.0
mock_user_api_key_dict.allowed_model_region = None
# This is what CustomStreamWrapper._hidden_params looks like after
# update_response_metadata() has been called on it
hidden_params = {
"model_id": "openai-gpt4o-deployment",
"api_base": "https://api.openai.com",
"additional_headers": {},
"litellm_overhead_time_ms": 55.3, # set by update_response_metadata
"_response_ms": 280.0,
"litellm_call_id": "test-call-id",
"response_cost": 0.002,
"cache_key": None,
"fastest_response_batch_completion": None,
"callback_duration_ms": None,
}
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=mock_user_api_key_dict,
call_id="test-call-id",
model_id=hidden_params.get("model_id"),
cache_key=hidden_params.get("cache_key") or "",
api_base=hidden_params.get("api_base") or "",
version="1.0.0",
response_cost=hidden_params.get("response_cost"),
model_region="",
hidden_params=hidden_params,
)
# The overhead header must be present and correct
assert "x-litellm-overhead-duration-ms" in custom_headers, (
"x-litellm-overhead-duration-ms header must be emitted during streaming. "
"It was missing — this is the streaming overhead header regression."
)
assert custom_headers["x-litellm-overhead-duration-ms"] == "55.3"

57
tests/test_litellm/test_router.py
View file

@ -1297,6 +1297,61 @@ async def test_acompletion_streaming_iterator_edge_cases():
print("✓ Edge case tests passed!")
@pytest.mark.asyncio
async def test_acompletion_streaming_iterator_preserves_hidden_params():
"""
Regression test: FallbackStreamWrapper must copy _hidden_params from the
original CustomStreamWrapper so that x-litellm-overhead-duration-ms (and
other hidden params) are present in the proxy response headers for streaming.
"""
from unittest.mock import MagicMock
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
}
],
)
# Simulate a CustomStreamWrapper that already has timing metadata set by
# update_response_metadata (litellm_overhead_time_ms, _response_ms, etc.)
mock_response = MagicMock()
mock_response.model = "gpt-4"
mock_response.custom_llm_provider = "openai"
mock_response.logging_obj = MagicMock()
mock_response._hidden_params = {
"litellm_overhead_time_ms": 12.34,
"_response_ms": 500.0,
"litellm_call_id": "test-call-id",
"api_base": "https://api.openai.com",
"additional_headers": {},
}
# Make the mock iterable (yields nothing — we only care about hidden_params)
async def _empty():
return
yield # make it an async generator
mock_response.__aiter__ = lambda self: _empty().__aiter__()
result = await router._acompletion_streaming_iterator(
model_response=mock_response,
messages=[{"role": "user", "content": "hi"}],
initial_kwargs={"model": "gpt-4", "stream": True},
)
# The returned FallbackStreamWrapper must carry the original _hidden_params
assert hasattr(result, "_hidden_params"), "result must have _hidden_params"
assert result._hidden_params.get("litellm_overhead_time_ms") == 12.34, (
"litellm_overhead_time_ms must be preserved — "
"this is what drives x-litellm-overhead-duration-ms in streaming responses"
)
assert result._hidden_params.get("litellm_call_id") == "test-call-id"
assert result._hidden_params.get("_response_ms") == 500.0
@pytest.mark.asyncio
async def test_async_function_with_fallbacks_common_utils():
"""Test the async_function_with_fallbacks_common_utils method"""
@ -1858,7 +1913,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name():
litellm_credential_name to actual credential values (for UI-created models).
"""
from litellm.types.utils import CredentialItem
# Setup credential list with a test credential
litellm.credential_list = [
CredentialItem(