From 391e2c33f35ddb0906dc16b9be23eddfd6174604 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Mon, 7 Sep 2026 00:46:32 -0700 Subject: [PATCH] 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. --- litellm/proxy/common_request_processing.py | 128 ++++-- .../pass_through_endpoints.py | 2 +- litellm/router.py | 59 ++- .../proxy/test_common_request_processing.py | 231 +++++++++++ tests/test_litellm/test_router.py | 364 ++++++++++++++++++ 5 files changed, 743 insertions(+), 41 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5e6c9b34332..0695a8ec9d3 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,40 @@ 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]: + """The response headers to publish once the first chunk has been buffered.""" + 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 +932,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 +959,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 +988,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 +1004,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 +1026,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 +1551,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 +2159,46 @@ 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: dict) -> 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 + return str(model_id) if model_id else "" + + @staticmethod + def _stream_response_headers( + *, + hidden_params: Mapping[str, object], + request_data: dict, + 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=ProxyBaseLLMRequestProcessing._get_model_id_from_response(hidden_params, request_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=request_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 +2467,38 @@ 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, + request_data=self.data, + 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. + + `create_response` buffers the first chunk before it commits headers, so a + fallback that took over by then has already repointed `response`. + """ + if not getattr(response, "fallback_headers_adopted", False): + return custom_headers + return self._stream_response_headers( + hidden_params=get_hidden_params_dict(response), + request_data=self.data, + 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 +2636,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 +3088,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 3f450661946..43cc6585582 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -612,6 +612,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 @@ -2611,7 +2643,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__( @@ -2619,6 +2651,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) @@ -2696,6 +2729,11 @@ 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) + _adopting_wrapper = wrapper_ref() + if _adopting_wrapper is not None: + _adopting_wrapper.adopt_fallback_response_headers( + fallback_response, prepared_fallback_hidden_params + ) async for fallback_item in fallback_response: Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if ( @@ -2738,7 +2776,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( @@ -3167,13 +3209,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"): @@ -3230,6 +3273,11 @@ class Router: if hasattr(fallback_response, "__iter__"): prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) + _adopting_wrapper = wrapper_ref() + if _adopting_wrapper is not None: + _adopting_wrapper.adopt_fallback_response_headers( + fallback_response, prepared_fallback_hidden_params + ) for fallback_item in fallback_response: Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if ( @@ -3268,7 +3316,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..018b0afed91 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1809,6 +1809,147 @@ 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. + """ + refreshed_at = [] + + async def mock_generator(): + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + async def refresh_headers(): + refreshed_at.append("called") + return {"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 refreshed_at == ["called"] + 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. + """ + produced = {"first_chunk": False} + + async def mock_generator(): + produced["first_chunk"] = True + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + async def refresh_headers(): + served = "fallback-deployment" if produced["first_chunk"] 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 +8155,93 @@ 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" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 7c044310e14..a8dc161f958 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2386,6 +2386,370 @@ 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_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 + + +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.