From c19d49d9192f26a7e15c9571d1c3a4e7ee45092a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:40:13 -0700 Subject: [PATCH] fix(responses): keep provider response headers in streaming logging callbacks (#38131) * fix(responses): keep provider response headers in streaming logging callbacks The responses streaming iterator captures the provider's HTTP response headers into its own _hidden_params, but never puts them on the completed response, and the model_validate(model_dump()) copy made for logging drops pydantic private attributes. Success callbacks and StandardLoggingPayload.hidden_params.additional_headers therefore saw an empty dict for streaming /v1/responses, so Azure's apim-request-id was unreadable from the callback payload. Restore the headers on the nested response of the logging copy, preferring any the provider transform already set (the fake_stream path) and falling back to the ones the iterator captured from the stream. Skipped when the copy fell back to the original event, so a serialization failure never leaves logging-only state on the caller's object. * fix: satisfy LIT002 mutable-collection gate in header restore --------- Co-authored-by: Yucheng Zhu --- litellm/responses/streaming_iterator.py | 40 +++++ .../responses/test_streaming_iterator.py | 161 +++++++++++++++++- 2 files changed, 200 insertions(+), 1 deletion(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7871c85220c..f271655f5e3 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -17,6 +17,7 @@ from typing_extensions import TypeIs import litellm from litellm.constants import ( + EMPTY_MAPPING, LITELLM_MAX_STREAMING_DURATION_SECONDS, STREAM_SSE_DONE_STRING, ) @@ -273,6 +274,9 @@ class BaseResponsesAPIStreamingIterator: self._hidden_params["additional_headers"] = process_response_headers( self.response.headers or {} ) # GUARANTEE OPENAI HEADERS IN RESPONSE + self._raw_response_headers: Mapping[str, str] = MappingProxyType( + dict(self.response.headers or {}) # mutable-ok: immediately frozen by MappingProxyType + ) def _check_max_streaming_duration(self) -> None: """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" @@ -446,6 +450,7 @@ class BaseResponsesAPIStreamingIterator: except Exception: # Fallback to original if serialization fails pass + self._restore_provider_response_headers(logging_response) end_time: Final = datetime.now() if is_async: @@ -480,6 +485,41 @@ class BaseResponsesAPIStreamingIterator: ) self._run_post_success_hooks(end_time=end_time) + def _restore_provider_response_headers(self, logging_response: object) -> None: + """Re-apply the provider's response headers to the copy handed to logging callbacks. + + ``model_validate(model_dump())`` above drops pydantic private attributes, so the + ``_hidden_params`` the provider transform set on the nested response are lost. Returns early + when that copy fell back to the original event, so logging-only state never lands on the + object the caller is iterating. + """ + if logging_response is self.completed_response: + return + target: Final[object] = getattr(logging_response, "response", None) + existing_hidden: Final[object] = getattr(target, "_hidden_params", None) + if not isinstance(existing_hidden, Mapping): + return + existing: Final[Mapping[str, object]] = existing_hidden + source_hidden: Final[object] = getattr( + getattr(self.completed_response, "response", None), "_hidden_params", None + ) + source: Final[Mapping[str, object]] = source_hidden if isinstance(source_hidden, Mapping) else EMPTY_MAPPING + processed: Final[object] = source.get("additional_headers") or self._hidden_params.get("additional_headers") + raw: Final[object] = source.get("headers") or self._raw_response_headers + headers: Final[Mapping[str, object]] = processed if isinstance(processed, Mapping) else EMPTY_MAPPING + raw_headers: Final[Mapping[str, object]] = raw if isinstance(raw, Mapping) else EMPTY_MAPPING + # rebuild by value and let existing keys win: sharing the source dicts would alias what the proxy + # splats into the client's HTTP headers, and copying non-header keys would carry response_cost + setattr( # noqa: B010 # target is typed object here, so a plain attribute store does not type check + target, + "_hidden_params", + { # mutable-ok: the cost calculator writes optional_params into _hidden_params + "additional_headers": {**headers}, # mutable-ok: fresh copy, logging callbacks may mutate it + "headers": {**raw_headers}, # mutable-ok: fresh copy, logging callbacks may mutate it + **existing, + }, + ) + def _handle_logging_completed_response(self): """Base implementation - should be overridden by subclasses""" diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 9edcaaef034..c226c0b4d09 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -6,7 +6,7 @@ completion_start_time = end_time.""" import json from datetime import datetime from typing import Optional -from unittest.mock import Mock +from unittest.mock import Mock, patch import httpx import pytest @@ -378,3 +378,162 @@ def test_stamp_responses_usage_cost_survives_calculator_failure(): _stamp_responses_usage_cost(response, logging_obj) assert getattr(response.usage, "cost", None) is None + + +def _capture_dispatch(logged: list): + """Record the object handed to the success handlers. + + ``Mock(spec=LiteLLMLoggingObj).dispatch_success_handlers`` is an AsyncMock whose side effect + only runs when the coroutine is awaited, so capture with a plain function instead. + """ + + async def _noop() -> None: + return None + + def _dispatch(result, **kwargs): + logged.append(result) + return _noop() + + return _dispatch + + +def _headers_config(*, transform_hidden_params: Optional[dict] = None) -> Mock: + """Config whose completed event carries a real ResponsesAPIResponse, so the logging copy + performs a genuine model_dump/model_validate round trip.""" + mock_config = Mock(spec=BaseResponsesAPIConfig) + + def _transform(model, parsed_chunk, logging_obj): + evt_type = parsed_chunk.get("type") + if evt_type != "response.completed": + stub = Mock() + stub.type = evt_type + return stub + response = ResponsesAPIResponse( + id="resp_headers", + created_at=1, + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + if transform_hidden_params is not None: + response._hidden_params.update(transform_hidden_params) + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=response, + ) + + mock_config.transform_streaming_response.side_effect = _transform + return mock_config + + +def _make_header_iterator( + *, + headers: dict, + config: Mock, + logging_obj: LiteLLMLoggingObj, +) -> ResponsesAPIStreamingIterator: + async def aiter_bytes(): + yield _sse_event({"type": "response.completed"}) + + mock_response = Mock() + mock_response.headers = headers + mock_response.aiter_bytes = aiter_bytes + + return ResponsesAPIStreamingIterator( + response=mock_response, + model="gpt-4o-mini", + responses_api_provider_config=config, + logging_obj=logging_obj, + litellm_metadata={}, + custom_llm_provider="azure", + ) + + +@pytest.mark.asyncio +async def test_streaming_logging_response_carries_provider_response_headers(): + """LIT-6055: the provider headers the iterator captured must reach the logged response, so + custom loggers can read Azure's apim-request-id from the callback payload.""" + logging_obj = _logging_obj_stub() + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator( + headers={"apim-request-id": "azure-correlation-1", "x-ms-region": "East US 2"}, + config=_headers_config(), + logging_obj=logging_obj, + ) + async for _ in iterator: + pass + + assert len(logged) == 1 + hidden_params = logged[0].response._hidden_params + assert hidden_params["additional_headers"]["llm_provider-apim-request-id"] == "azure-correlation-1" + assert hidden_params["additional_headers"]["llm_provider-x-ms-region"] == "East US 2" + assert hidden_params["headers"]["apim-request-id"] == "azure-correlation-1" + # the proxy builds the client's response headers from the iterator's own dict, so the logged + # response must hold copies rather than alias it + assert hidden_params["additional_headers"] is not iterator._hidden_params["additional_headers"] + assert hidden_params["headers"] is not iterator._raw_response_headers + + +@pytest.mark.asyncio +async def test_streaming_logging_copy_preserves_transform_hidden_params(): + """LIT-6055: model_validate(model_dump()) drops pydantic private attributes, so headers a + provider transform already set on the response (fake_stream) must be re-applied.""" + logging_obj = _logging_obj_stub() + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator( + headers={}, + config=_headers_config( + transform_hidden_params={ + "additional_headers": {"llm_provider-apim-request-id": "from-transform"}, + "headers": {"apim-request-id": "from-transform"}, + "response_cost": 0.5, + } + ), + logging_obj=logging_obj, + ) + async for _ in iterator: + pass + + assert len(logged) == 1 + hidden_params = logged[0].response._hidden_params + assert hidden_params["additional_headers"]["llm_provider-apim-request-id"] == "from-transform" + assert hidden_params["headers"]["apim-request-id"] == "from-transform" + assert iterator.completed_response is not logged[0] + # only the header keys travel: response_cost would short-circuit the cost calculator + assert "response_cost" not in hidden_params + + +@pytest.mark.asyncio +async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched(): + """LIT-6055: when the logging copy falls back to the original event, the header restore must + not stamp logging-only state onto the object the caller is iterating.""" + logging_obj = _logging_obj_stub() + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator( + headers={"apim-request-id": "azure-correlation-1"}, + config=_headers_config(), + logging_obj=logging_obj, + ) + async for _ in iterator: + pass + + assert len(logged) == 1 + iterator._completed_response_logged = False + logged.clear() + with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")): + iterator._log_completed_response(is_async=True) + + assert logged == [iterator.completed_response] + assert iterator.completed_response.response._hidden_params == {}