fix(router): refresh streaming fallback headers and add spend logger proof

This commit is contained in:
Yujong Lee 2026-09-02 12:53:51 -07:00
parent 8d8266d01d
commit eda63531cf
9 changed files with 438 additions and 39 deletions

View file

@ -57,7 +57,7 @@
"limit": 5601
},
"reportMissingTypeArgument": {
"limit": 15306
"limit": 15304
},
"reportMissingTypeStubs": {
"limit": 40
@ -99,7 +99,7 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44364
"limit": 44353
},
"reportUnknownLambdaType": {
"limit": 109
@ -108,10 +108,10 @@
"limit": 38350
},
"reportUnknownParameterType": {
"limit": 19626
"limit": 19624
},
"reportUnknownVariableType": {
"limit": 29890
"limit": 29883
},
"reportUnnecessaryCast": {
"limit": 111

View file

@ -0,0 +1,55 @@
# Store upstream response headers in spend logs
Copy `response_header_logger.py` next to your proxy configuration as `custom_callbacks.py`, then register the logger:
```yaml
model_list:
- model_name: openrouter-header-test
litellm_params:
model: openrouter/openai/gpt-5.6-luna
api_key: os.environ/OPENROUTER_API_KEY
litellm_settings:
callbacks: [custom_callbacks.proxy_handler_instance]
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
database_url: os.environ/DATABASE_URL
proxy_batch_write_at: 1
```
Start the proxy with this repository's code and your configuration:
```sh
uv run --no-sync litellm --config config.yaml --port 4013
```
The logger captures response headers before stream iteration, enriches successful spend logs through `async_logging_hook`, and recovers failure headers from the exception chain. A failure without recoverable headers retains the earlier capture. When a fallback response is captured, its headers replace the failed attempts headers. Request state stays in request metadata, so concurrent calls do not share headers on the logger instance
Send a request with existing custom metadata:
```sh
curl --fail-with-body -sS -N http://localhost:4013/v1/chat/completions \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" \
-H 'Content-Type: application/json' \
-d '{"model":"openrouter-header-test","messages":[{"role":"user","content":"Reply exactly OK"}],"stream":true,"max_tokens":32,"metadata":{"spend_logs_metadata":{"verification_case":"stream","existing_field":"keep-me"}}}'
```
After the background write, read the stored headers:
```sh
curl --fail-with-body -sS http://localhost:4013/spend/logs \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" |
jq '.[] | . as $row |
(.metadata | if type == "string" then fromjson else . end).spend_logs_metadata as $m |
{request_id: $row.request_id, case: $m.verification_case,
existing_field: $m.existing_field, headers: $m.upstream_response_headers}'
```
Headers appear in `metadata.spend_logs_metadata.upstream_response_headers`. Provider names receive the `llm_provider-` prefix; existing custom metadata and LiteLLM's internal-header protections are preserved
Set `stream` to false for the non-streaming control. To reproduce an upstream routing failure, add `"num_retries":0,"provider":{"only":["header-verification-nonexistent-provider"],"allow_fallbacks":false}` to the request. This returned HTTP 404 in verification
The logger was verified with OpenRouter chat completions, including simultaneous streams, a non-streaming control, and an upstream 404. The unit tests also cover a headerless error after header capture and both metadata keys used by the logging pipeline. Live commands, spend-log/database readbacks, and the exact tested revisions are in [PR #39376](https://github.com/BerriAI/litellm/pull/39376)
An earlier logger without early capture lost headers on a live streaming 429. That 429 did not recur after the logger update; retention for that error shape is covered by a unit test, not a second live 429 reproduction
Unmodified v1.98.0 still needs the shared HTTP-handler propagation fix as well as the router fix. The logger alone cannot recover headers dropped before the callback. These are headers returned by OpenRouter, which may differ from those of the provider behind it. Other providers, every retry/fallback attempt, and every failure type have not been live-tested

View file

@ -0,0 +1,86 @@
from collections.abc import Mapping
from typing import Final, cast
import httpx
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
def _as_dict(value: object) -> dict[str, object]:
return cast(dict[str, object], value) if isinstance(value, dict) else {}
def _child(parent: dict[str, object], key: str) -> dict[str, object]:
child: Final = _as_dict(parent.get(key))
parent[key] = child # rebind-ok: callback metadata must be updated in place
return child
def _headers(value: object) -> dict[str, str]:
if isinstance(value, httpx.Headers):
return dict(value)
return {name: header for name, header in _as_dict(value).items() if isinstance(header, str)}
def get_error_headers(error: BaseException | None, seen: tuple[int, ...] = ()) -> dict[str, str]:
if error is None or id(error) in seen:
return {}
response: Final[object] = getattr(error, "response", None)
headers: Final = _headers(getattr(response, "headers", None)) or _headers(getattr(error, "headers", None))
if headers:
return headers
visited: Final = (*seen, id(error))
return get_error_headers(error.__cause__, visited) or get_error_headers(error.__context__, visited)
def save_headers(metadata: dict[str, object], headers: Mapping[str, object]) -> None:
existing: Final = _as_dict(metadata.get("spend_logs_metadata"))
metadata["spend_logs_metadata"] = { # rebind-ok: callback metadata must be updated in place
**existing,
"upstream_response_headers": {
**_as_dict(existing.get("upstream_response_headers")),
**headers,
},
}
class ResponseHeaderLogger(CustomLogger):
async def async_post_call_response_headers_hook(
self,
data: dict[str, object],
user_api_key_dict: UserAPIKeyAuth,
response: object,
request_headers: dict[str, str] | None = None,
litellm_call_info: dict[str, object] | None = None,
) -> None:
headers: Final = _headers(getattr(response, "_response_headers", None))
spend_metadata: Final = _child(_child(data, "metadata"), "spend_logs_metadata")
spend_metadata["upstream_response_headers"] = {f"llm_provider-{name}": value for name, value in headers.items()}
async def async_logging_hook(
self, kwargs: dict[str, object], result: object, call_type: str
) -> tuple[dict[str, object], object]:
payload: Final = _as_dict(kwargs.get("standard_logging_object"))
hidden: Final = _as_dict(payload.get("hidden_params"))
headers: Final = _as_dict(hidden.get("additional_headers"))
params: Final = _child(kwargs, "litellm_params")
key: Final = "litellm_metadata" if params.get("litellm_metadata") else "metadata"
save_headers(_child(params, key), headers)
return kwargs, result
async def async_post_call_failure_hook(
self,
request_data: dict[str, object],
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
traceback_str: str | None = None,
) -> None:
headers: Final = get_error_headers(original_exception)
save_headers(
_child(request_data, "metadata"),
{f"llm_provider-{name}": value for name, value in headers.items()},
)
proxy_handler_instance: Final = ResponseHeaderLogger()

View file

@ -719,7 +719,7 @@ class _UpstreamClosingStreamingResponse(StreamingResponse):
content: AsyncGenerator[str, None],
*,
media_type: str | None = None,
headers: dict | None = None,
headers: Mapping[str, str] | None = None,
status_code: int = status.HTTP_200_OK,
upstream_generator: AsyncGenerator[str, None] | None = None,
) -> None:
@ -857,22 +857,23 @@ def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]:
async def create_response(
generator: AsyncGenerator[str, None],
media_type: str,
headers: dict,
headers: Mapping[str, str],
default_status_code: int = status.HTTP_200_OK,
request: Request | None = None,
refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None = None,
) -> StreamingResponse | JSONResponse:
"""
Create streaming response, checking if the first chunk is an error.
If the first chunk is an error, return a standard JSON error response.
Otherwise, return StreamingResponse and stream all content.
"""
# Tell buffering reverse proxies (nginx, ingress-nginx, Envoy) to flush SSE
# immediately instead of releasing the whole stream in one batch (issue #28384).
streaming_headers: Final = {
**headers,
"Cache-Control": "no-cache",
"X-Accel-Buffering": "no",
}
async def streaming_headers() -> Mapping[str, str]:
current_headers: Final = await refresh_headers() if refresh_headers is not None else headers
return MappingProxyType({**current_headers, "Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
first_chunk_value: str | None = None
final_status_code = default_status_code
@ -909,7 +910,7 @@ async def create_response(
return JSONResponse(
status_code=final_status_code,
content={"error": error_dict},
headers=headers,
headers=await refresh_headers() if refresh_headers is not None else headers,
)
except Exception as e:
verbose_proxy_logger.debug("Error parsing first chunk value: %s", e)
@ -938,7 +939,7 @@ async def create_response(
return StreamingResponse(
empty_gen(),
media_type=media_type,
headers=streaming_headers,
headers=await streaming_headers(),
status_code=default_status_code,
)
except Exception as e:
@ -954,7 +955,7 @@ async def create_response(
return StreamingResponse(
error_gen_message(),
media_type=media_type,
headers=streaming_headers,
headers=await streaming_headers(),
status_code=error_status,
)
@ -976,7 +977,7 @@ async def create_response(
return _UpstreamClosingStreamingResponse(
combined_generator(),
media_type=media_type,
headers=streaming_headers,
headers=await streaming_headers(),
status_code=final_status_code,
upstream_generator=generator,
)
@ -2379,28 +2380,46 @@ class ProxyBaseLLMRequestProcessing:
if self._is_streaming_request(
data=self.data, is_streaming_request=is_streaming_request
) or self._is_streaming_response(response): # use generate_responses to stream responses
custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=logging_obj.litellm_call_id,
model_id=model_id,
cache_key=cache_key,
api_base=api_base,
version=version,
response_cost=response_cost,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
fastest_response_batch_completion=fastest_response_batch_completion,
request_data=self.data,
hidden_params=hidden_params,
litellm_logging_obj=logging_obj,
**additional_headers,
)
def get_stream_headers() -> dict[str, str]:
current_hidden: Final = get_hidden_params_dict(response)
return ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=logging_obj.litellm_call_id,
model_id=self._get_model_id_from_response(current_hidden, self.data),
cache_key=current_hidden.get("cache_key") or "",
api_base=current_hidden.get("api_base") or "",
version=version,
response_cost=current_hidden.get("response_cost") or "",
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
fastest_response_batch_completion=current_hidden.get("fastest_response_batch_completion"),
request_data=self.data,
hidden_params=current_hidden,
litellm_logging_obj=logging_obj,
**(current_hidden.get("additional_headers") or MappingProxyType({})),
)
custom_headers: Final = get_stream_headers()
request_headers: Final = dict(request.headers)
async def refresh_stream_headers() -> Mapping[str, str]:
if get_hidden_params_dict(response) is hidden_params:
return custom_headers
refreshed: Final = get_stream_headers()
updated_callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook(
data=self.data,
user_api_key_dict=user_api_key_dict,
response=response,
request_headers=request_headers,
)
return MappingProxyType({**refreshed, **(updated_callback_headers or MappingProxyType({}))})
# Call response headers hook for streaming success
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
data=self.data,
user_api_key_dict=user_api_key_dict,
response=response,
request_headers=dict(request.headers),
request_headers=request_headers,
)
if callback_headers:
custom_headers.update(callback_headers)
@ -2508,6 +2527,8 @@ class ProxyBaseLLMRequestProcessing:
)
# Non-streaming response - fall through to normal response handling
elif select_data_generator:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
selected_data_generator = select_data_generator(
response=response,
user_api_key_dict=user_api_key_dict,
@ -2536,6 +2557,7 @@ class ProxyBaseLLMRequestProcessing:
media_type="text/event-stream",
headers=custom_headers,
request=request,
refresh_headers=refresh_stream_headers if isinstance(response, CustomStreamWrapper) else None,
)
### CALL HOOKS ### - modify outgoing data

View file

@ -2499,6 +2499,10 @@ class Router:
if hasattr(model_response, "_hidden_params"):
self._hidden_params = model_response._hidden_params.copy()
def update_headers(self, source: object) -> None:
self._response_headers = getattr(source, "_response_headers", None)
self._hidden_params = get_hidden_params_dict(source).copy()
def __aiter__(self):
return self
@ -2552,8 +2556,10 @@ class Router:
# If fallback returns a streaming response, iterate over it
if hasattr(fallback_response, "__aiter__"):
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
wrapped_response.update_headers(fallback_response)
async for fallback_item in fallback_response:
wrapped_response.update_headers(fallback_response)
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params)
if (
fallback_item
@ -2601,7 +2607,8 @@ class Router:
e,
)
return FallbackStreamWrapper(stream_with_fallbacks())
wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks())
return wrapped_response
@staticmethod
def _extract_partial_responses_usage(
@ -3043,6 +3050,10 @@ class Router:
if hasattr(model_response, "_hidden_params"):
self._hidden_params = model_response._hidden_params.copy()
def update_headers(self, source: object) -> None:
self._response_headers = getattr(source, "_response_headers", None)
self._hidden_params = get_hidden_params_dict(source).copy()
def __iter__(self):
return self
@ -3093,8 +3104,10 @@ class Router:
)
if hasattr(fallback_response, "__iter__"):
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
wrapped_response.update_headers(fallback_response)
for fallback_item in fallback_response:
wrapped_response.update_headers(fallback_response)
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params)
if (
fallback_item
@ -3132,7 +3145,8 @@ class Router:
close_err,
)
return SyncFallbackStreamWrapper(stream_with_fallbacks())
wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks())
return wrapped_response
async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs):
"""

View file

@ -565,3 +565,96 @@ def test_redis_cache_completion_stream():
# test_redis_cache_completion_stream()
@pytest.mark.asyncio
async def test_response_header_logger_preserves_headers_on_headerless_stream_failure():
from types import SimpleNamespace
from typing import Final
from cookbook.logging_observability.response_header_logger import ResponseHeaderLogger
from litellm.proxy._types import UserAPIKeyAuth
logger: Final = ResponseHeaderLogger()
auth: Final = UserAPIKeyAuth()
data: Final = {"metadata": {"spend_logs_metadata": {"existing_field": "keep"}}}
response: Final = SimpleNamespace(
_response_headers={
"x-generation-id": "stream-request",
"x-litellm-attempted-fallbacks": "spoofed",
}
)
await logger.async_post_call_response_headers_hook(data, auth, response)
await logger.async_post_call_failure_hook(data, RuntimeError("stream failed"), auth)
assert data["metadata"]["spend_logs_metadata"] == {
"existing_field": "keep",
"upstream_response_headers": {
"llm_provider-x-generation-id": "stream-request",
"llm_provider-x-litellm-attempted-fallbacks": "spoofed",
},
}
@pytest.mark.asyncio
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
async def test_response_header_logger_success_preserves_capture_and_metadata(metadata_key: str):
from types import SimpleNamespace
from typing import Final
from cookbook.logging_observability.response_header_logger import ResponseHeaderLogger
from litellm.proxy._types import UserAPIKeyAuth
logger: Final = ResponseHeaderLogger()
metadata: Final = {"spend_logs_metadata": {"existing_field": "keep"}}
await logger.async_post_call_response_headers_hook(
{"metadata": metadata},
UserAPIKeyAuth(),
SimpleNamespace(_response_headers={"x-generation-id": "stream-request"}),
)
kwargs: Final = {
"litellm_params": {metadata_key: metadata},
"standard_logging_object": {
"hidden_params": {
"additional_headers": {
"llm_provider-x-litellm-response-cost": 0.01,
}
}
},
}
await logger.async_logging_hook(kwargs, None, "acompletion")
assert metadata["spend_logs_metadata"] == {
"existing_field": "keep",
"upstream_response_headers": {
"llm_provider-x-generation-id": "stream-request",
"llm_provider-x-litellm-response-cost": 0.01,
},
}
@pytest.mark.asyncio
@pytest.mark.parametrize("fallback_headers", [None, {"x-request-id": "winner"}])
async def test_response_header_logger_replaces_headers_after_fallback(fallback_headers: dict[str, str] | None):
from types import SimpleNamespace
from typing import Final
from cookbook.logging_observability.response_header_logger import ResponseHeaderLogger
from litellm.proxy._types import UserAPIKeyAuth
logger: Final = ResponseHeaderLogger()
auth: Final = UserAPIKeyAuth()
data: Final = {"metadata": {"spend_logs_metadata": {"existing_field": "keep"}}}
await logger.async_post_call_response_headers_hook(
data,
auth,
SimpleNamespace(_response_headers={"x-request-id": "failed", "x-failed-only": "remove"}),
)
await logger.async_post_call_response_headers_hook(data, auth, SimpleNamespace(_response_headers=fallback_headers))
await logger.async_post_call_failure_hook(data, RuntimeError("stream failed"), auth)
assert data["metadata"]["spend_logs_metadata"] == {
"existing_field": "keep",
"upstream_response_headers": {
f"llm_provider-{name}": value for name, value in (fallback_headers or {}).items()
},
}

View file

@ -7665,3 +7665,35 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_
records = [r for r in caplog.records if "_handle_llm_api_exception(): Exception occured" in r.getMessage()]
assert len(records) == 1
assert (records[0].exc_info is not None) is expect_traceback
@pytest.mark.asyncio
@pytest.mark.parametrize("stream_kind", ["success", "empty", "error"])
async def test_create_response_refreshes_headers_before_committing(stream_kind: str):
from typing import Final
headers: Final = {"llm_provider-x-request-id": "failed", "llm_provider-x-failed-only": "remove"}
async def stream():
if stream_kind == "error":
yield 'data: {"error":{"message":"failed","code":429}}\n\n'
elif stream_kind == "success":
yield 'data: {"id":"winner","choices":[]}\n\n'
async def refresh_headers():
return {"llm_provider-x-request-id": "winner", "x-litellm-model-group": "fallback"}
response: Final = await create_response(
stream(),
"text/event-stream",
headers,
refresh_headers=refresh_headers,
)
assert response.headers["llm_provider-x-request-id"] == "winner"
assert response.headers["x-litellm-model-group"] == "fallback"
assert "llm_provider-x-failed-only" not in response.headers
if stream_kind == "error":
assert response.status_code == 429
else:
assert response.headers["cache-control"] == "no-cache"
assert response.headers["x-accel-buffering"] == "no"

View file

@ -2346,7 +2346,7 @@ async def test_completion_streaming_iterator_preserves_response_headers(
_response_headers=response_headers,
)
router: Final = Router(model_list=[])
upstream._hidden_params["additional_headers"]["x-litellm-model-group"] = "real-group"
await router.set_response_headers(response=upstream, model_group="real-group")
iterator_kwargs: Final = dict(
model_response=upstream,
messages=[{"role": "user", "content": "Hello"}],
@ -2373,6 +2373,103 @@ async def test_completion_streaming_iterator_preserves_response_headers(
assert StandardLoggingPayloadSetup.get_hidden_params(complete._hidden_params)["additional_headers"] == expected
@pytest.mark.asyncio
@pytest.mark.parametrize("is_async", [False, True])
@pytest.mark.parametrize("fallback_headers", [None, {"x-request-id": "winner", "x-litellm-model-group": "spoofed"}])
async def test_completion_streaming_iterator_replaces_failed_provider_headers(
is_async: bool, fallback_headers: dict[str, str] | None
) -> None:
from typing import Final
from unittest.mock import Mock
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
client: Final = AsyncMock(spec=AsyncHTTPHandler) if is_async else Mock(spec=HTTPHandler)
client.post.return_value = httpx.Response(
200,
content=(
b'data: {"id":"winner","object":"chat.completion.chunk","created":1,'
b'"model":"test-model","choices":[{"index":0,"delta":{"content":"Hello"},'
b'"finish_reason":null}]}\n\n'
b"data: [DONE]\n\n"
),
request=httpx.Request("POST", "https://provider.test/v1/chat/completions"),
)
request: Final = dict(
model="hosted_vllm/test-model",
messages=[{"role": "user", "content": "Hello"}],
api_base="https://provider.test/v1",
api_key="fake-key",
stream=True,
client=client,
)
original: Final = await litellm.acompletion(**request) if is_async else litellm.completion(**request)
fallback: Final = CustomStreamWrapper(
completion_stream=original.completion_stream,
model=original.model,
custom_llm_provider=original.custom_llm_provider,
logging_obj=original.logging_obj,
_response_headers=fallback_headers,
)
class FallbackRouter(Router):
async def async_function_with_fallbacks_common_utils(self, **kwargs):
return fallback
def function_with_fallbacks(self, **kwargs):
return fallback
class FailedStream:
model = "failed-model"
custom_llm_provider = "hosted_vllm"
logging_obj = original.logging_obj
chunks = []
_response_headers = {"x-request-id": "failed", "x-failed-only": "remove"}
_hidden_params = {
"additional_headers": {"llm_provider-x-request-id": "failed", "llm_provider-x-failed-only": "remove"}
}
def __iter__(self):
return self
def __next__(self):
raise MidStreamFallbackError(
message="limited",
model=self.model,
llm_provider=self.custom_llm_provider,
is_pre_first_chunk=True,
)
def __aiter__(self):
return self
async def __anext__(self):
return self.__next__()
router: Final = FallbackRouter(model_list=[])
await router.set_response_headers(response=fallback, model_group="fallback")
iterator_kwargs: Final = dict(
model_response=FailedStream(),
messages=[{"role": "user", "content": "Hello"}],
initial_kwargs={"model": "primary", "stream": True},
)
wrapped: Final = (
await router._acompletion_streaming_iterator(**iterator_kwargs)
if is_async
else router._completion_streaming_iterator(**iterator_kwargs)
)
chunks: Final = [chunk async for chunk in wrapped] if is_async else list(wrapped)
expected: Final = {
**{f"llm_provider-{name}": value for name, value in (fallback_headers or {}).items()},
"x-litellm-model-group": "fallback",
}
assert wrapped._response_headers == fallback_headers
assert wrapped._hidden_params["additional_headers"] == expected
assert all(chunk._hidden_params["additional_headers"] == expected for chunk in chunks)
assert litellm.stream_chunk_builder(chunks).choices[0].message.content == "Hello"
def test_completion_streaming_iterator_fallback_on_429():
"""Sync streaming: MidStreamFallbackError (429 pre-first-chunk) triggers fallback.

View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 22364
"limit": 22363
},
"LIT002": {
"limit": 26777
"limit": 26776
},
"LIT003": {
"limit": 269
@ -27,7 +27,7 @@
"limit": 0
},
"LIT010": {
"limit": 16507
"limit": 16505
},
"LIT011": {
"limit": 5535