From e04e5d71138ea1356aacea1ea22ef060855dd596 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 7 Sep 2026 12:59:57 -0700 Subject: [PATCH] fix(router): keep provider response headers on streaming chat completions (#40091) * fix(router): keep provider response headers on streaming chat completions The Router re-wraps a deployment's CustomStreamWrapper in FallbackStreamWrapper (and its sync twin) so a mid-stream failure can fail over. Neither wrapper forwarded `_response_headers`, so every streaming chat completion handed the proxy's callbacks and its response-header builder a wrapper with no provider headers, and a successful mid-stream fallback still published the failed deployment's identity, `x-request-id` and rate limit counters. Forward `_response_headers` into both wrappers, repoint the wrapper at the deployment that served the stream once a fallback takes over, and rebuild the proxy's response headers from that deployment while `create_response` still has the first chunk buffered. * fix(router): follow a nested fallback to the deployment that served the stream A fallback the router picks is itself a fallback-aware wrapper, and it only repoints at its own fallback once it yields, so reading its hidden params at selection time named a deployment that produced no output. Re-read them when the first fallback item arrives, which is still before the proxy commits response headers. Also addresses review feedback: the streaming header builder reads self.data instead of taking a coarse request_data parameter, and the new test recorder local is Final. * test(router): cover the fallback header adoption helper directly The router_code_coverage gate wants every router.py function named in a router test, and this also pins the weak-reference behavior: a wrapper collected mid-stream must not break the generator still draining it. * refactor(proxy): take a read-only mapping for the model-id lookup _get_model_id_from_response only reads its request payload, so a Mapping says what it needs and the two metadata hops are narrowed instead of assumed to be dicts. * test: drop mutable recorder locals and routine comments from the new tests An AsyncMock await_count and an asyncio.Event say the same thing as a list and a dict that the test mutates. * chore(router): justify the two rebinds in the fallback loops Both are the one-shot re-read that follows a nested fallback, so they get the repo's rebind-ok note like the rest of the file. --- litellm/proxy/common_request_processing.py | 126 ++-- .../pass_through_endpoints.py | 2 +- litellm/router.py | 83 ++- .../proxy/test_common_request_processing.py | 247 +++++++- tests/test_litellm/test_router.py | 574 ++++++++++++++++++ 5 files changed, 985 insertions(+), 47 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5e6c9b34332..9720e4b1cf8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -756,7 +756,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: @@ -888,25 +888,39 @@ def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]: return f"data: {json.dumps({'error': error_obj})}\n\n", "data: [DONE]\n\n" +def _sse_stream_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + """`headers` plus the two that stop reverse proxies from buffering SSE (issue #28384).""" + return MappingProxyType({**headers, **_TTFT_KEEPALIVE_HEADERS}) + + +async def _resolve_stream_headers( + headers: Mapping[str, str], refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None +) -> Mapping[str, str]: + if refresh_headers is None: + return headers + try: + return await refresh_headers() + except Exception as e: # noqa: BLE001 # a stream whose first chunk is already paid for must not fail over its headers + verbose_proxy_logger.exception("Error refreshing streaming response headers: %s", e) + return headers + + 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. + + ``refresh_headers`` is consulted once the first chunk has been buffered, for + callers whose headers can only be known then. """ - # 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", - } first_chunk_value: str | None = None final_status_code = default_status_code @@ -917,6 +931,7 @@ async def create_response( # Now get the first chunk from the actual generator first_chunk_value = await _buffer_first_chunk_honoring_disconnect(generator, request) + resolved_headers: Final = await _resolve_stream_headers(headers, refresh_headers) if first_chunk_value is not None: try: @@ -943,7 +958,7 @@ async def create_response( return JSONResponse( status_code=final_status_code, content={"error": error_dict}, - headers=headers, + headers=resolved_headers, ) except Exception as e: verbose_proxy_logger.debug("Error parsing first chunk value: %s", e) @@ -972,7 +987,7 @@ async def create_response( return StreamingResponse( empty_gen(), media_type=media_type, - headers=streaming_headers, + headers=_sse_stream_headers(await _resolve_stream_headers(headers, refresh_headers)), status_code=default_status_code, ) except Exception as e: @@ -988,7 +1003,7 @@ async def create_response( return StreamingResponse( error_gen_message(), media_type=media_type, - headers=streaming_headers, + headers=_sse_stream_headers(await _resolve_stream_headers(headers, refresh_headers)), status_code=error_status, ) @@ -1010,7 +1025,7 @@ async def create_response( return _UpstreamClosingStreamingResponse( combined_generator(), media_type=media_type, - headers=streaming_headers, + headers=_sse_stream_headers(resolved_headers), status_code=final_status_code, upstream_generator=generator, ) @@ -1535,7 +1550,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _merge_passthrough_streaming_headers( response_headers: httpx.Headers | dict | None, - custom_headers: dict, + custom_headers: Mapping[str, str], ) -> dict: """ Merge upstream passthrough headers with proxy/custom headers. @@ -2143,14 +2158,45 @@ class ProxyBaseLLMRequestProcessing: return fallback_model_group @staticmethod - def _get_model_id_from_response(hidden_params: dict, data: dict) -> str: + def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str: """Extract model_id from hidden_params with fallback to litellm_metadata.""" model_id = hidden_params.get("model_id", None) or "" if not model_id: - litellm_metadata: Final = data.get("litellm_metadata", {}) or {} - model_info: Final = litellm_metadata.get("model_info", {}) or {} - model_id = model_info.get("id", "") or "" - return model_id + litellm_metadata: Final = data.get("litellm_metadata") + model_info: Final = litellm_metadata.get("model_info") if isinstance(litellm_metadata, Mapping) else None + model_id = (model_info.get("id") or "") if isinstance(model_info, Mapping) else "" + return str(model_id) if model_id else "" + + def _stream_response_headers( + self, + *, + hidden_params: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + logging_obj: LiteLLMLoggingObj, + version: str | None, + callback_headers: Mapping[str, str], + ) -> Mapping[str, str]: + """The streaming response headers describing `hidden_params`' deployment.""" + return MappingProxyType( + { + **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(hidden_params, self.data), + cache_key=hidden_params.get("cache_key") or "", + api_base=hidden_params.get("api_base") or "", + version=version, + response_cost=hidden_params.get("response_cost") or "", + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + fastest_response_batch_completion=hidden_params.get("fastest_response_batch_completion"), + request_data=self.data, + hidden_params=hidden_params, + litellm_logging_obj=logging_obj, + **(hidden_params.get("additional_headers") or MappingProxyType({})), + ), + **callback_headers, + } + ) @staticmethod def _get_deployment_model_name( @@ -2419,31 +2465,32 @@ 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, - ) - # Call response headers hook for streaming success - callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + stream_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=dict(request.headers), ) - if callback_headers: - custom_headers.update(callback_headers) + custom_headers: Final = self._stream_response_headers( + hidden_params=hidden_params, + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, + version=version, + callback_headers=stream_callback_headers or MappingProxyType({}), + ) + + async def refresh_stream_headers() -> Mapping[str, str]: + """`custom_headers` rebuilt for whichever deployment served the stream.""" + if not getattr(response, "fallback_headers_adopted", False): + return custom_headers + return self._stream_response_headers( + hidden_params=get_hidden_params_dict(response), + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, + version=version, + callback_headers=stream_callback_headers or MappingProxyType({}), + ) # Preserve the original client-requested model (pre-alias mapping) for downstream # streaming generators. Pre-call processing can rewrite `self.data["model"]` for @@ -2581,6 +2628,7 @@ class ProxyBaseLLMRequestProcessing: media_type="text/event-stream", headers=custom_headers, request=request, + refresh_headers=refresh_stream_headers, ) ### CALL HOOKS ### - modify outgoing data @@ -3032,7 +3080,7 @@ class ProxyBaseLLMRequestProcessing: response: Any, proxy_logging_obj: "ProxyLogging", user_api_key_dict: "UserAPIKeyAuth", - custom_headers: dict, + custom_headers: Mapping[str, str], request_headers: dict[str, str], ) -> Response | None: if not self._has_post_call_guardrails_for_passthrough(): diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3cb9acc6110..f1f823e59b5 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -322,7 +322,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): def get_response_headers( headers: httpx.Headers, litellm_call_id: str | None = None, - custom_headers: dict | None = None, + custom_headers: Mapping[str, str] | None = None, ) -> dict: # Exclude headers that uvicorn writes itself (server, date) and # encoding/length headers that don't survive re-serialization. diff --git a/litellm/router.py b/litellm/router.py index c7a254cbc02..95cabfad4bd 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -616,6 +616,38 @@ RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( RETRY_BREADCRUMB_LIMIT: Final = 4 +class FallbackAwareStreamWrapper(CustomStreamWrapper): + """Base for the Router's chat-completion stream wrappers, which are built around the + attempt the Router picked first and have to repoint themselves when a fallback takes over.""" + + fallback_headers_adopted: bool = False + + def adopt_fallback_response_headers( + self, + fallback_response: object, + prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]], + ) -> None: + """Repoint this wrapper at the deployment that served the stream. + + Replaces rather than merges, so the failed attempt's `x-request-id`, rate limit + counters, `model_id` and `api_base` cannot reach the proxy's response headers or + its callbacks. + """ + self._response_headers = getattr(fallback_response, "_response_headers", None) + fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params + if fallback_hidden_params: + self._hidden_params = { # mutable-ok: the rest of litellm writes into _hidden_params + **fallback_hidden_params, + # dict() because add_retry_fallback_headers mutates additional_headers in place + "additional_headers": dict(fallback_headers), # mutable-ok: see above + } + self._base_hidden_params = { # mutable-ok: CustomStreamWrapper keeps this snapshot as a dict + **self._hidden_params, + "response_cost": None, + } + self.fallback_headers_adopted = True + + class Router: model_names: set = set() cache_responses: bool | None = False @@ -2576,6 +2608,18 @@ class Router: return fallback_hidden_params, {} return fallback_hidden_params, cast("dict[str, object]", fallback_headers) + @staticmethod + def _adopt_fallback_response_headers( + wrapper_ref: "weakref.ref[FallbackAwareStreamWrapper]", + fallback_response: object, + ) -> tuple[dict[str, object], dict[str, object]]: + """Repoint the wrapper at `fallback_response`, returning its prepared hidden params.""" + prepared: Final = Router._prepare_fallback_hidden_params(fallback_response) + adopting_wrapper: Final = wrapper_ref() + if adopting_wrapper is not None: + adopting_wrapper.adopt_fallback_response_headers(fallback_response, prepared) + return prepared + @staticmethod def _apply_fallback_hidden_params_to_item( fallback_item: object, @@ -2615,7 +2659,7 @@ class Router: held_slot: Final = deployment_slot if deployment_slot is not None else contextlib.AsyncExitStack() - class FallbackStreamWrapper(CustomStreamWrapper): + class FallbackStreamWrapper(FallbackAwareStreamWrapper): def __init__(self, async_generator: AsyncGenerator): # Copy attributes from the original model_response super().__init__( @@ -2623,6 +2667,7 @@ class Router: model=model_response.model, custom_llm_provider=model_response.custom_llm_provider, logging_obj=model_response.logging_obj, + _response_headers=getattr(model_response, "_response_headers", None), ) self._async_generator = async_generator inner_chunks: Final[object] = getattr(model_response, "chunks", None) @@ -2699,8 +2744,17 @@ 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) + prepared_fallback_hidden_params = Router._adopt_fallback_response_headers( + wrapper_ref, fallback_response + ) + fallback_headers_are_settled = False async for fallback_item in fallback_response: + if not fallback_headers_are_settled: + fallback_headers_are_settled = True # rebind-ok: one-shot latch + # a fallback that failed over again only repoints itself once it yields + prepared_fallback_hidden_params = ( # rebind-ok: re-read once the fallback yields + Router._adopt_fallback_response_headers(wrapper_ref, fallback_response) + ) Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if ( fallback_item @@ -2742,7 +2796,11 @@ class Router: e, ) - return FallbackStreamWrapper(stream_with_fallbacks()) + wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks()) + # weak, so the generator closing over it does not keep the wrapper out of + # refcount teardown and delay the `finally` that releases the deployment slot + wrapper_ref: Final = weakref.ref(wrapped_response) + return wrapped_response @staticmethod def _extract_partial_responses_usage( @@ -3171,13 +3229,14 @@ class Router: """ from litellm.exceptions import MidStreamFallbackError - class SyncFallbackStreamWrapper(CustomStreamWrapper): + class SyncFallbackStreamWrapper(FallbackAwareStreamWrapper): def __init__(self, sync_generator: Generator): super().__init__( completion_stream=sync_generator, model=model_response.model, custom_llm_provider=model_response.custom_llm_provider, logging_obj=model_response.logging_obj, + _response_headers=getattr(model_response, "_response_headers", None), ) self._sync_generator = sync_generator if hasattr(model_response, "_hidden_params"): @@ -3233,8 +3292,17 @@ class Router: ) if hasattr(fallback_response, "__iter__"): - prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) + prepared_fallback_hidden_params = Router._adopt_fallback_response_headers( + wrapper_ref, fallback_response + ) + fallback_headers_are_settled = False for fallback_item in fallback_response: + if not fallback_headers_are_settled: + fallback_headers_are_settled = True # rebind-ok: one-shot latch + # a fallback that failed over again only repoints itself once it yields + prepared_fallback_hidden_params = ( # rebind-ok: re-read once the fallback yields + Router._adopt_fallback_response_headers(wrapper_ref, fallback_response) + ) Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if ( fallback_item @@ -3272,7 +3340,10 @@ class Router: close_err, ) - return SyncFallbackStreamWrapper(stream_with_fallbacks()) + wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks()) + # weak, for the same reason as the async twin + wrapper_ref: Final = weakref.ref(wrapped_response) + return wrapped_response async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs): """ diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 96d9b0c5a26..6acd9d7258e 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2,7 +2,7 @@ import asyncio import copy import datetime import json -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import AsyncGenerator, Callable, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -1809,6 +1809,146 @@ class TestCommonRequestProcessingHelpers: response = await create_response(mock_generator(), "text/event-stream", custom_headers) assert response.headers["x-custom-header"] == "TestValue" + async def test_create_streaming_response_refresh_headers_after_first_chunk(self): + """LIT-6767: headers a caller can only resolve once the first chunk exists. + + A pre-first-chunk fallback replaces the deployment while the response + headers are still uncommitted, so ``refresh_headers`` is consulted after + the first chunk is buffered and its result wins. + """ + + async def mock_generator(): + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + refresh_headers: Final = AsyncMock( + return_value={"x-litellm-model-id": "fallback-deployment", "llm_provider-x-request-id": "req-FALLBACK"} + ) + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment", "llm_provider-x-request-id": "req-FAILED"}, + refresh_headers=refresh_headers, + ) + assert isinstance(response, StreamingResponse) + assert refresh_headers.await_count == 1 + assert response.headers["x-litellm-model-id"] == "fallback-deployment" + assert response.headers["llm_provider-x-request-id"] == "req-FALLBACK" + # the buffering headers are still applied on top of the refreshed set + assert response.headers["x-accel-buffering"] == "no" + assert response.headers["cache-control"] == "no-cache" + + async def test_create_streaming_response_refreshes_only_after_the_first_chunk(self): + """LIT-6767: the refresh has to be consulted after the generator produced a chunk. + + A pre-first-chunk fallback only repoints the response while that first chunk is + being produced, so a refresh consulted any earlier still describes the attempt + that failed and the headers go out wrong. + """ + first_chunk_produced: Final = asyncio.Event() + + async def mock_generator(): + first_chunk_produced.set() + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + async def refresh_headers(): + served = "fallback-deployment" if first_chunk_produced.is_set() else "failed-deployment" + return {"x-litellm-model-id": served} + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + refresh_headers=refresh_headers, + ) + assert response.headers["x-litellm-model-id"] == "fallback-deployment" + + async def test_create_streaming_response_empty_stream_uses_refreshed_headers(self): + """LIT-6767: a fallback that served nothing still gets to name itself. + + The empty-generator branch returns its own StreamingResponse, so it needs the + refreshed headers too or the client is told the failed deployment answered. + """ + + async def mock_generator(): + return + yield # make it an async generator + + async def refresh_headers(): + return {"x-litellm-model-id": "fallback-deployment"} + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + refresh_headers=refresh_headers, + ) + assert isinstance(response, StreamingResponse) + assert response.headers["x-litellm-model-id"] == "fallback-deployment" + assert response.headers["x-accel-buffering"] == "no" + + async def test_create_streaming_response_without_refresh_headers_is_unchanged(self): + """LIT-6767: the default keeps the caller-supplied headers verbatim.""" + + async def mock_generator(): + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + ) + assert response.headers["x-litellm-model-id"] == "failed-deployment" + + async def test_create_streaming_response_refresh_headers_failure_keeps_stream(self): + """LIT-6767: the first chunk is already paid for, so a failing refresh + falls back to the caller's headers instead of erroring the stream.""" + + async def mock_generator(): + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + async def refresh_headers(): + raise RuntimeError("boom") + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + refresh_headers=refresh_headers, + ) + assert isinstance(response, StreamingResponse) + assert response.status_code == status.HTTP_200_OK + assert response.headers["x-litellm-model-id"] == "failed-deployment" + assert await self.consume_stream(response) == [ + 'data: {"content": "data"}\n\n', + "data: [DONE]\n\n", + ] + + async def test_create_response_first_chunk_error_uses_refreshed_headers(self): + """LIT-6767: the JSON error response built from a bad first chunk carries + the refreshed headers too, so it cannot describe a deployment that no + longer served the request.""" + + async def mock_generator(): + yield 'data: {"error": {"code": 403, "message": "forbidden"}}\n\n' + yield "data: [DONE]\n\n" + + async def refresh_headers(): + return {"x-litellm-model-id": "fallback-deployment"} + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + refresh_headers=refresh_headers, + ) + assert isinstance(response, JSONResponse) + assert response.headers["x-litellm-model-id"] == "fallback-deployment" + async def test_create_streaming_response_disables_proxy_buffering(self): """Regression for #28384: every StreamingResponse create_response returns must carry the headers that stop nginx/ingress/Envoy from buffering the @@ -8014,3 +8154,108 @@ class TestDetachedStreamFailureHook: await logging_obj._on_detached_stream_failure(failure) assert [call["original_exception"] for call in recorder.calls] == [failure] + + +class TestStreamingResponseHeadersFollowFallback: + """LIT-6767: the streaming branch has to publish the deployment that served the stream.""" + + @staticmethod + def _fallback_adopting_stream(): + class _Stream: + def __init__(self) -> None: + self._hidden_params = { + "model_id": "failed-deployment", + "api_base": "http://127.0.0.1:20769/v1", + "additional_headers": {"llm_provider-stale-marker": "failed-deployment"}, + } + self.fallback_headers_adopted = False + + def adopt(self) -> None: + self._hidden_params = { + "model_id": "served-deployment", + "api_base": "https://api.openai.com", + "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, + } + self.fallback_headers_adopted = True + + return _Stream() + + @pytest.mark.asyncio + async def test_streaming_headers_name_the_deployment_that_served(self, monkeypatch): + """A pre-first-chunk fallback repoints the stream while the headers are still + uncommitted, so the published headers must describe the fallback, not the attempt + the Router picked first.""" + stream = self._fallback_adopting_stream() + + def select_data_generator(**kwargs): + async def generator(): + stream.adopt() + yield 'data: {"choices": [{"delta": {"content": "OK"}}]}\n\n' + yield "data: [DONE]\n\n" + + return generator() + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "lit-6767-call" + logging_obj._defer_async_logging = False + logging_obj._on_deferred_stream_complete = None + logging_obj.cost_breakdown = None + + processor = ProxyBaseLLMRequestProcessing( + data={"model": "oa-midfail", "stream": True, "litellm_logging_obj": logging_obj} + ) + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_success_hook = AsyncMock( + side_effect=lambda data, user_api_key_dict, response: response + ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-header": "kept"} + ) + + async def fake_route_request(**kwargs): + async def call(): + return stream + + return call() + + monkeypatch.setattr( + litellm.proxy.common_request_processing, "route_request", fake_route_request + ) + + result = await processor.base_process_llm_request( + request=Request(scope={"type": "http", "headers": []}), + fastapi_response=Response(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + select_data_generator=select_data_generator, + is_streaming_request=True, + skip_pre_call_logic=True, + ) + + assert isinstance(result, StreamingResponse) + assert result.headers["x-litellm-model-id"] == "served-deployment" + assert result.headers["x-litellm-model-api-base"] == "https://api.openai.com" + assert result.headers["llm_provider-x-request-id"] == "req-SERVED" + assert "llm_provider-stale-marker" not in result.headers + assert result.headers["x-callback-header"] == "kept" + + +class TestPassthroughHeadersAcceptImmutableMappings: + """LIT-6767: the streaming branch now hands the passthrough helpers an immutable mapping.""" + + def test_merge_passthrough_streaming_headers_accepts_a_read_only_mapping(self): + merged = ProxyBaseLLMRequestProcessing._merge_passthrough_streaming_headers( + response_headers=httpx.Headers({"content-type": "text/event-stream", "transfer-encoding": "chunked"}), + custom_headers=MappingProxyType({"x-litellm-model-id": "served-deployment"}), + ) + + assert merged["x-litellm-model-id"] == "served-deployment" + assert merged["content-type"] == "text/event-stream" + # the excluded hop-by-hop header is still dropped + assert "transfer-encoding" not in merged diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index e9b77076448..eed34c79a06 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2387,6 +2387,580 @@ async def test_acompletion_streaming_iterator_preserves_hidden_params(): assert result._hidden_params.get("_response_ms") == 500.0 +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_preserves_response_headers(): + """LIT-6767: the returned wrapper must carry the provider's raw response headers. + + Proxy callbacks read ``_response_headers`` off the object the router hands + back. The wrapper used to be built without it, so every streaming chat + completion reported zero raw provider headers while the non-streaming path + reported the full set. + """ + from unittest.mock import MagicMock + + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + async def _empty(): + return + yield # make it an async generator + + provider_headers = { + "x-request-id": "req-provider-123", + "x-ratelimit-remaining-requests": "42", + # a provider must never be able to spoof an internal header + "x-litellm-model-id": "spoofed", + } + source = CustomStreamWrapper( + completion_stream=_empty(), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers=provider_headers, + ) + + result = await router._acompletion_streaming_iterator( + model_response=source, + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + + assert result._response_headers == provider_headers + additional_headers = result._hidden_params["additional_headers"] + assert additional_headers["llm_provider-x-request-id"] == "req-provider-123" + assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42" + # internal-header protection survives: the provider value is namespaced, never promoted + assert additional_headers["llm_provider-x-litellm-model-id"] == "spoofed" + assert "x-litellm-model-id" not in additional_headers + + +def test_completion_streaming_iterator_preserves_response_headers(): + """LIT-6767, sync counterpart of the async header-preservation test.""" + from unittest.mock import MagicMock + + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + provider_headers = {"x-request-id": "req-provider-sync", "openai-organization": "org-real"} + source = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers=provider_headers, + ) + + result = router._completion_streaming_iterator( + model_response=source, + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + + assert result._response_headers == provider_headers + assert result._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req-provider-sync" + + +def test_adopt_fallback_response_headers_replaces_rather_than_merges(): + """LIT-6767: direct unit for FallbackAwareStreamWrapper.adopt_fallback_response_headers. + + Values from the failed attempt must not survive, so the wrapper replaces both + ``_response_headers`` and ``_hidden_params`` instead of merging them. + """ + from unittest.mock import MagicMock + + from litellm.router import FallbackAwareStreamWrapper, Router + + wrapper = FallbackAwareStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers={"x-request-id": "req-FAILED"}, + ) + wrapper._hidden_params = { + "model_id": "failed-deployment", + "only_on_failed_attempt": "stale", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + } + + fallback = MagicMock() + fallback._response_headers = {"x-request-id": "req-FALLBACK"} + fallback._hidden_params = { + "model_id": "fallback-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, + } + + wrapper.adopt_fallback_response_headers( + fallback, Router._prepare_fallback_hidden_params(fallback) + ) + + assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} + assert wrapper._hidden_params["model_id"] == "fallback-deployment" + assert "only_on_failed_attempt" not in wrapper._hidden_params + assert wrapper._hidden_params is not fallback._hidden_params + # the snapshot CustomStreamWrapper caches at init has to follow, or a chunk built + # from it would still be stamped with the deployment that failed + assert wrapper._base_hidden_params["model_id"] == "fallback-deployment" + # the nested header dict is copied too, so a later mutation on the fallback + # response cannot reach headers the proxy has already published + assert wrapper._hidden_params["additional_headers"] is not fallback._hidden_params["additional_headers"] + fallback._hidden_params["additional_headers"]["llm_provider-x-request-id"] = "req-MUTATED" + assert wrapper._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"} + + +def test_adopt_fallback_response_headers_survives_a_collected_wrapper(): + """LIT-6767: adoption still returns the fallback's params once the wrapper is gone.""" + import weakref + from unittest.mock import MagicMock + + from litellm.router import FallbackAwareStreamWrapper, Router + + fallback: Final = MagicMock() + fallback._response_headers = {"x-request-id": "req-FALLBACK"} + fallback._hidden_params = { + "model_id": "fallback-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, + } + + wrapper = FallbackAwareStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + ) + live_ref: Final = weakref.ref(wrapper) + prepared: Final = Router._adopt_fallback_response_headers(live_ref, fallback) + assert prepared == (fallback._hidden_params, fallback._hidden_params["additional_headers"]) + assert wrapper.fallback_headers_adopted is True + assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} + + dead_ref: Final = weakref.ref(wrapper) + del wrapper + assert dead_ref() is None + assert Router._adopt_fallback_response_headers(dead_ref, fallback) == prepared + + +def test_adopt_fallback_response_headers_drops_headers_the_fallback_cannot_replace(): + """LIT-6767: a fallback that carries no raw provider headers publishes none. + + Keeping the failed attempt's raw headers would hand the client and the callbacks a + provider ``x-request-id`` for a request that deployment never served, which is the + leak this fix exists to close. + """ + from unittest.mock import MagicMock + + from litellm.router import FallbackAwareStreamWrapper, Router + + wrapper = FallbackAwareStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers={"x-request-id": "req-FAILED"}, + ) + wrapper._hidden_params = {"model_id": "failed-deployment", "additional_headers": {}} + + fallback = MagicMock() + fallback._response_headers = None + fallback._hidden_params = {"model_id": "fallback-deployment"} + + wrapper.adopt_fallback_response_headers( + fallback, Router._prepare_fallback_hidden_params(fallback) + ) + + assert wrapper._response_headers is None + assert wrapper._hidden_params["model_id"] == "fallback-deployment" + assert wrapper.fallback_headers_adopted is True + + +def test_adopt_fallback_response_headers_keeps_identity_when_fallback_has_none(): + """A fallback response carrying no hidden params keeps the identity headers. + + Publishing no ``x-litellm-*`` header at all for a request the fallback served is + worse than keeping what is there, so only the raw provider headers are dropped. + """ + from unittest.mock import MagicMock + + from litellm.router import FallbackAwareStreamWrapper, Router + + wrapper = FallbackAwareStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers={"x-request-id": "req-FAILED"}, + ) + hidden_params_before = wrapper._hidden_params + + fallback = object() + wrapper.adopt_fallback_response_headers( + fallback, Router._prepare_fallback_hidden_params(fallback) + ) + + assert wrapper._response_headers is None + assert wrapper._hidden_params is hidden_params_before + assert wrapper.fallback_headers_adopted is True + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_adopts_fallback_response_headers(): + """LIT-6767: after a successful pre-first-chunk fallback, the wrapper must + describe the deployment that served the stream, with no value left over + from the attempt that failed.""" + from unittest.mock import MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + failed_error = MidStreamFallbackError( + message="upstream died before the first chunk", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=True, + ) + + class FailedStream: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + self._response_headers = {"x-request-id": "req-FAILED"} + self._hidden_params = { + "model_id": "failed-deployment", + "api_base": "https://failed.example", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + "only_on_failed_attempt": "stale", + } + + def __aiter__(self): + return self + + async def __anext__(self): + raise failed_error + + class FallbackStream: + def __init__(self): + self._response_headers = {"x-request-id": "req-FALLBACK"} + self._hidden_params = { + "model_id": "fallback-deployment", + "api_base": "https://fallback.example", + "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, + } + self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._chunks) + except StopIteration: + raise StopAsyncIteration from None + + fallback_stream = FallbackStream() + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=fallback_stream, + ): + result = await router._acompletion_streaming_iterator( + model_response=FailedStream(), + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + # the failed attempt is what the wrapper is built from + assert result._response_headers == {"x-request-id": "req-FAILED"} + # the very first chunk the fallback produces must already be published under + # the fallback's identity: the proxy commits response headers once that chunk + # is buffered, so adopting any later is adopting too late + await result.__anext__() + assert result._response_headers == {"x-request-id": "req-FALLBACK"} + assert result._hidden_params["model_id"] == "fallback-deployment" + async for _ in result: + pass + + assert result._response_headers == {"x-request-id": "req-FALLBACK"} + assert result._hidden_params["model_id"] == "fallback-deployment" + assert result._hidden_params["api_base"] == "https://fallback.example" + assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"} + # stale values are removed, not merged over + assert "only_on_failed_attempt" not in result._hidden_params + # and the wrapper holds its own copy, so later fallback mutations cannot leak in + assert result._hidden_params is not fallback_stream._hidden_params + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback(): + """LIT-6767: a fallback that itself fails over before its first chunk. + + The selected fallback still describes its own failed attempt at selection time, so + the wrapper has to re-read it once a chunk exists or it publishes a deployment that + produced no output. + """ + from unittest.mock import MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + failed_error: Final = MidStreamFallbackError( + message="upstream died before the first chunk", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=True, + ) + + class FailedStream: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + self._response_headers = {"x-request-id": "req-FAILED"} + self._hidden_params = { + "model_id": "failed-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + } + + def __aiter__(self): + return self + + async def __anext__(self): + raise failed_error + + class NestedFallbackStream: + """A fallback that repoints itself at a third deployment as it yields.""" + + def __init__(self): + self._response_headers = {"x-request-id": "req-MIDDLE"} + self._hidden_params = { + "model_id": "middle-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"}, + } + self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + chunk = next(self._chunks) + except StopIteration: + raise StopAsyncIteration from None + self._response_headers = {"x-request-id": "req-SERVED"} + self._hidden_params = { + "model_id": "served-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, + } + return chunk + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=NestedFallbackStream(), + ): + result = await router._acompletion_streaming_iterator( + model_response=FailedStream(), + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + first_chunk: Final = await result.__anext__() + # the proxy commits response headers once this chunk is buffered + assert result._response_headers == {"x-request-id": "req-SERVED"} + assert result._hidden_params["model_id"] == "served-deployment" + assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-SERVED"} + # and the chunk itself carries the same deployment + assert first_chunk._hidden_params["model_id"] == "served-deployment" + async for _ in result: + pass + + assert result._response_headers == {"x-request-id": "req-SERVED"} + assert result._hidden_params["model_id"] == "served-deployment" + + +def test_completion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback(): + """LIT-6767, sync counterpart of the nested-fallback adoption test.""" + from unittest.mock import MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + failed_error: Final = MidStreamFallbackError( + message="upstream died before the first chunk", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=True, + ) + + class FailedStream: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + self._response_headers = {"x-request-id": "req-FAILED"} + self._hidden_params = { + "model_id": "failed-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + } + + def __iter__(self): + return self + + def __next__(self): + raise failed_error + + class NestedFallbackStream: + """A fallback that repoints itself at a third deployment as it yields.""" + + def __init__(self): + self._response_headers = {"x-request-id": "req-MIDDLE"} + self._hidden_params = { + "model_id": "middle-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"}, + } + self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) + + def __iter__(self): + return self + + def __next__(self): + chunk = next(self._chunks) + self._response_headers = {"x-request-id": "req-SERVED"} + self._hidden_params = { + "model_id": "served-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, + } + return chunk + + with patch.object(router, "function_with_fallbacks", return_value=NestedFallbackStream()): + result = router._completion_streaming_iterator( + model_response=FailedStream(), + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + first_chunk: Final = next(result) + assert result._response_headers == {"x-request-id": "req-SERVED"} + assert result._hidden_params["model_id"] == "served-deployment" + assert first_chunk._hidden_params["model_id"] == "served-deployment" + for _ in result: + pass + + assert result._response_headers == {"x-request-id": "req-SERVED"} + assert result._hidden_params["model_id"] == "served-deployment" + + +def test_completion_streaming_iterator_adopts_fallback_response_headers(): + """LIT-6767, sync counterpart of the fallback-adoption test.""" + from unittest.mock import MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + failed_error = MidStreamFallbackError( + message="upstream died before the first chunk", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=True, + ) + + class FailedStream: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + self._response_headers = {"x-request-id": "req-FAILED"} + self._hidden_params = { + "model_id": "failed-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + "only_on_failed_attempt": "stale", + } + + def __iter__(self): + return self + + def __next__(self): + raise failed_error + + class FallbackStream: + def __init__(self): + self._response_headers = {"x-request-id": "req-FALLBACK"} + self._hidden_params = { + "model_id": "fallback-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, + } + + def __iter__(self): + return iter([]) + + with patch.object(router, "function_with_fallbacks", return_value=FallbackStream()): + result = router._completion_streaming_iterator( + model_response=FailedStream(), + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + assert result._response_headers == {"x-request-id": "req-FAILED"} + for _ in result: + pass + + assert result._response_headers == {"x-request-id": "req-FALLBACK"} + assert result._hidden_params["model_id"] == "fallback-deployment" + assert "only_on_failed_attempt" not in result._hidden_params + + def test_completion_streaming_iterator_fallback_on_429(): """Sync streaming: MidStreamFallbackError (429 pre-first-chunk) triggers fallback.