fix(responses): keep provider response headers in streaming logging callbacks

litellm/responses/streaming_iterator.py

@@ -17,6 +17,7 @@
 
 import litellm
 from litellm.constants import (
+    EMPTY_MAPPING,
     LITELLM_MAX_STREAMING_DURATION_SECONDS,
     STREAM_SSE_DONE_STRING,
 )
@@ -273,6 +274,9 @@ def __init__(
         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 @@ def _log_completed_response(self, *, is_async: bool) -> None:
             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 @@ def _log_completed_response(self, *, is_async: bool) -> None:
             )
         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"""
 

tests/test_litellm/responses/test_streaming_iterator.py

@@ -6,7 +6,7 @@
 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 == {}