fix(proxy): extend response headers hook to streaming, TTS, image gen, and pass-through (#24232)

* fix(proxy): extend response headers hook to streaming, TTS, image gen, and pass-through

* test: mock post_call_response_headers_hook in audio speech route tests

litellm/proxy/image_endpoints/endpoints.py

@@ -173,6 +173,16 @@ async def image_generation(
             )
         )
 
+        # Call response headers hook (matches base_process_llm_request behavior)
+        callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
+            data=data,
+            user_api_key_dict=user_api_key_dict,
+            response=response,
+            request_headers=dict(request.headers),
+        )
+        if callback_headers:
+            fastapi_response.headers.update(callback_headers)
+
         return response
     except Exception as e:
         await proxy_logging_obj.post_call_failure_hook(

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -1099,6 +1099,20 @@ async def pass_through_request(  # noqa: PLR0915
                     status_code=e.response.status_code, detail=await e.response.aread()
                 )
 
+            # Call response headers hook for streaming pass-through
+            _response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
+                headers=response.headers,
+                litellm_call_id=litellm_call_id,
+            )
+            callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
+                data=_parsed_body or {},
+                user_api_key_dict=user_api_key_dict,
+                response=response,
+                request_headers=dict(request.headers),
+            )
+            if callback_headers:
+                _response_headers.update(callback_headers)
+
             return StreamingResponse(
                 PassThroughStreamingHandler.chunk_processor(
                     response=response,
@@ -1109,10 +1123,7 @@ async def pass_through_request(  # noqa: PLR0915
                     passthrough_success_handler_obj=pass_through_endpoint_logging,
                     url_route=str(url),
                 ),
-                headers=HttpPassThroughEndpointHelpers.get_response_headers(
-                    headers=response.headers,
-                    litellm_call_id=litellm_call_id,
-                ),
+                headers=_response_headers,
                 status_code=response.status_code,
             )
 
@@ -1151,6 +1162,20 @@ async def pass_through_request(  # noqa: PLR0915
                     status_code=e.response.status_code, detail=await e.response.aread()
                 )
 
+            # Call response headers hook for detected streaming pass-through
+            _response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
+                headers=response.headers,
+                litellm_call_id=litellm_call_id,
+            )
+            callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
+                data=_parsed_body or {},
+                user_api_key_dict=user_api_key_dict,
+                response=response,
+                request_headers=dict(request.headers),
+            )
+            if callback_headers:
+                _response_headers.update(callback_headers)
+
             return StreamingResponse(
                 PassThroughStreamingHandler.chunk_processor(
                     response=response,
@@ -1161,10 +1186,7 @@ async def pass_through_request(  # noqa: PLR0915
                     passthrough_success_handler_obj=pass_through_endpoint_logging,
                     url_route=str(url),
                 ),
-                headers=HttpPassThroughEndpointHelpers.get_response_headers(
-                    headers=response.headers,
-                    litellm_call_id=litellm_call_id,
-                ),
+                headers=_response_headers,
                 status_code=response.status_code,
             )
 
@@ -1303,6 +1325,16 @@ async def pass_through_request(  # noqa: PLR0915
             api_base=str(url._uri_reference),
         )
 
+        # Call response headers hook
+        callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
+            data=_parsed_body or {},
+            user_api_key_dict=user_api_key_dict,
+            response=response,
+            request_headers=dict(request.headers),
+        )
+        if callback_headers:
+            custom_headers.update(callback_headers)
+
         response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
             headers=response.headers,
             custom_headers=custom_headers,

litellm/proxy/proxy_server.py

@@ -9191,6 +9191,16 @@ async def audio_speech(
             hidden_params=hidden_params,
         )
 
+        # Call response headers hook (matches audio_transcription behavior)
+        callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
+            data=data,
+            user_api_key_dict=user_api_key_dict,
+            response=response,
+            request_headers=dict(request.headers),
+        )
+        if callback_headers:
+            custom_headers.update(callback_headers)
+
         # Determine media type based on model type
         media_type = "audio/mpeg"  # Default for OpenAI TTS
         request_model = data.get("model", "")

litellm/proxy/utils.py

@@ -2496,7 +2496,8 @@ class ProxyLogging:
         )
 
         return {
-            "custom_llm_provider": hidden_params.get("custom_llm_provider"),
+            "custom_llm_provider": hidden_params.get("custom_llm_provider")
+            or getattr(response, "custom_llm_provider", None),
             "model_info": model_info,
             "api_base": hidden_params.get("api_base"),
             "model_id": hidden_params.get("model_id"),

tests/test_litellm/proxy/hooks/test_post_call_response_headers_hook.py

@@ -13,7 +13,6 @@ from unittest.mock import patch
 
 sys.path.insert(0, os.path.abspath("../../../.."))
 
-import litellm
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.proxy._types import UserAPIKeyAuth
 
@@ -336,3 +335,110 @@ async def test_litellm_call_info_backwards_compatible():
 
         assert result == {"x-test": "1"}
         assert injector.called is True
+
+
+# --- Tests for custom_llm_provider fallback (streaming response types) ---
+
+
+@pytest.mark.asyncio
+async def test_litellm_call_info_fallback_to_response_attribute():
+    """Test that _build_litellm_call_info falls back to response.custom_llm_provider
+    when _hidden_params doesn't contain it (streaming response types)."""
+    inspector = CallInfoInspectorLogger()
+
+    class MockStreamResponse:
+        """Mimics CustomStreamWrapper: custom_llm_provider as attribute,
+        _hidden_params without it."""
+
+        custom_llm_provider = "bedrock"
+        _hidden_params = {
+            "model_id": "model-xyz",
+            "api_base": "https://bedrock.us-east-1.amazonaws.com",
+        }
+
+    with patch("litellm.callbacks", [inspector]):
+        from litellm.proxy.utils import ProxyLogging
+        from litellm.caching.caching import DualCache
+
+        proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
+
+        await proxy_logging.post_call_response_headers_hook(
+            data={
+                "model": "claude-3",
+                "metadata": {"model_info": {"id": "model-xyz"}},
+            },
+            user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
+            response=MockStreamResponse(),
+        )
+
+        assert inspector.called is True
+        assert inspector.received_call_info is not None
+        assert inspector.received_call_info["custom_llm_provider"] == "bedrock"
+        assert (
+            inspector.received_call_info["api_base"]
+            == "https://bedrock.us-east-1.amazonaws.com"
+        )
+        assert inspector.received_call_info["model_id"] == "model-xyz"
+
+
+@pytest.mark.asyncio
+async def test_litellm_call_info_fallback_no_hidden_params():
+    """Test that _build_litellm_call_info works when response has no _hidden_params
+    at all (LiteLLMCompletionStreamingIterator case)."""
+    inspector = CallInfoInspectorLogger()
+
+    class MockIteratorResponse:
+        """Mimics LiteLLMCompletionStreamingIterator: custom_llm_provider as attribute,
+        no _hidden_params attribute at all."""
+
+        custom_llm_provider = "vertex_ai"
+
+    with patch("litellm.callbacks", [inspector]):
+        from litellm.proxy.utils import ProxyLogging
+        from litellm.caching.caching import DualCache
+
+        proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
+
+        await proxy_logging.post_call_response_headers_hook(
+            data={"model": "gemini-pro", "metadata": {}},
+            user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
+            response=MockIteratorResponse(),
+        )
+
+        assert inspector.called is True
+        assert inspector.received_call_info is not None
+        assert inspector.received_call_info["custom_llm_provider"] == "vertex_ai"
+        assert inspector.received_call_info["api_base"] is None
+        assert inspector.received_call_info["model_id"] is None
+
+
+@pytest.mark.asyncio
+async def test_litellm_call_info_hidden_params_takes_priority():
+    """Test that _hidden_params.custom_llm_provider takes priority over
+    the response attribute when both are present."""
+    inspector = CallInfoInspectorLogger()
+
+    class MockResponse:
+        custom_llm_provider = "attribute_value"
+        _hidden_params = {
+            "custom_llm_provider": "hidden_params_value",
+            "api_base": "https://example.com",
+            "model_id": "m1",
+        }
+
+    with patch("litellm.callbacks", [inspector]):
+        from litellm.proxy.utils import ProxyLogging
+        from litellm.caching.caching import DualCache
+
+        proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
+
+        await proxy_logging.post_call_response_headers_hook(
+            data={"model": "test", "metadata": {}},
+            user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
+            response=MockResponse(),
+        )
+
+        assert (
+            inspector.received_call_info["custom_llm_provider"]
+            == "hidden_params_value"
+        )

tests/test_litellm/proxy/image_endpoints/test_endpoints.py

@@ -43,11 +43,15 @@ async def test_image_generation_prompt_rerouting(monkeypatch):
     async def fake_post_call_success_hook(*, data, user_api_key_dict, response):
         return response
 
+    async def fake_post_call_response_headers_hook(**kwargs):
+        return {"x-callback-test": "value"}
+
     fake_proxy_logger = SimpleNamespace(
         pre_call_hook=fake_pre_call_hook,
         update_request_status=fake_update_request_status,
         post_call_failure_hook=fake_post_call_failure_hook,
         post_call_success_hook=fake_post_call_success_hook,
+        post_call_response_headers_hook=fake_post_call_response_headers_hook,
     )
 
     captured_route_request_data: Dict[str, Any] = {}
@@ -110,3 +114,4 @@ async def test_image_generation_prompt_rerouting(monkeypatch):
     assert pre_call_input["messages"][0]["content"] == "original prompt"
     assert captured_route_request_data["prompt"] == "sanitized prompt"
     assert "messages" not in captured_route_request_data
+    assert response.headers.get("x-callback-test") == "value"

tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py

@@ -1803,6 +1803,9 @@ class TestForwardHeaders:
             mock_logging_obj.pre_call_hook = AsyncMock(return_value=mock_request_body)
             mock_logging_obj.post_call_success_hook = AsyncMock()
             mock_logging_obj.post_call_failure_hook = AsyncMock()
+            mock_logging_obj.post_call_response_headers_hook = AsyncMock(
+                return_value={}
+            )
 
             # Call pass_through_request with forward_headers=True
             result = await pass_through_request(
@@ -1901,6 +1904,9 @@ class TestForwardHeaders:
             mock_logging_obj.pre_call_hook = AsyncMock(return_value=mock_request_body)
             mock_logging_obj.post_call_success_hook = AsyncMock()
             mock_logging_obj.post_call_failure_hook = AsyncMock()
+            mock_logging_obj.post_call_response_headers_hook = AsyncMock(
+                return_value={}
+            )
 
             # Call pass_through_request with forward_headers=False (default)
             result = await pass_through_request(

tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py

@@ -957,6 +957,9 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
                             return_value={"test": "data"}
                         )
                         mock_proxy_logging.post_call_failure_hook = AsyncMock()
+                        mock_proxy_logging.post_call_response_headers_hook = AsyncMock(
+                            return_value={"x-callback-test": "value"}
+                        )
 
                         # Setup mock for http response
                         mock_response = MagicMock()
@@ -1069,6 +1072,9 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream():
                     return_value={"model": "claude-3", "stream": True}
                 )
                 mock_proxy_logging.post_call_failure_hook = AsyncMock()
+                mock_proxy_logging.post_call_response_headers_hook = AsyncMock(
+                    return_value={"x-callback-test": "value"}
+                )
 
                 upstream_response = MagicMock()
                 upstream_response.status_code = 200
@@ -1134,6 +1140,9 @@ async def test_pass_through_request_sse_response_marks_logging_obj_as_stream():
                     return_value={"model": "claude-3"}
                 )
                 mock_proxy_logging.post_call_failure_hook = AsyncMock()
+                mock_proxy_logging.post_call_response_headers_hook = AsyncMock(
+                    return_value={"x-callback-test": "value"}
+                )
 
                 upstream_response = MagicMock()
                 upstream_response.status_code = 200
@@ -1975,6 +1984,9 @@ async def test_pass_through_request_query_params_forwarding():
                         mock_proxy_logging.pre_call_hook = AsyncMock(
                             return_value=test_body
                         )
+                        mock_proxy_logging.post_call_response_headers_hook = AsyncMock(
+                            return_value={"x-callback-test": "value"}
+                        )
 
                         # Setup mock for http response
                         mock_response = MagicMock()
@@ -2703,6 +2715,10 @@ async def test_pass_through_request_non_streaming_uses_content_for_state_raw_bod
             "litellm.proxy.proxy_server.proxy_logging_obj.pre_call_hook",
             new=AsyncMock(side_effect=_hook_mutates_body),
         ),
+        patch(
+            "litellm.proxy.proxy_server.proxy_logging_obj.post_call_response_headers_hook",
+            new=AsyncMock(return_value={}),
+        ),
         patch(
             "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler",
             new=AsyncMock(),
@@ -2763,6 +2779,10 @@ async def test_pass_through_request_streaming_uses_content_for_state_raw_body():
             "litellm.proxy.proxy_server.proxy_logging_obj.pre_call_hook",
             new=AsyncMock(side_effect=lambda **kw: kw["data"]),
         ),
+        patch(
+            "litellm.proxy.proxy_server.proxy_logging_obj.post_call_response_headers_hook",
+            new=AsyncMock(return_value={}),
+        ),
         patch(
             "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler",
             new=AsyncMock(),

tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py

@@ -130,6 +130,7 @@ class TestPassthroughPostCallGuardrails:
         mock_proxy_logging.post_call_success_hook = AsyncMock(
             return_value=_GEMINI_RESPONSE
         )
+        mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={})
 
         with _common_patches(mock_proxy_logging, mock_response):
             await pass_through_request(
@@ -155,6 +156,7 @@ class TestPassthroughPostCallGuardrails:
         mock_proxy_logging = MagicMock()
         mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
         mock_proxy_logging.post_call_success_hook = AsyncMock()
+        mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={})
 
         with _common_patches(mock_proxy_logging, mock_response):
             result = await pass_through_request(

tests/test_litellm/proxy/proxy_server/test_routes_audio.py

@@ -26,6 +26,7 @@ def patched_speech(monkeypatch):
         MagicMock(
             pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]),
             post_call_failure_hook=AsyncMock(),
+            post_call_response_headers_hook=AsyncMock(return_value={}),
             update_request_status=AsyncMock(),
         ),
     )
@@ -61,6 +62,7 @@ def patched_speech_error(monkeypatch):
         MagicMock(
             pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]),
             post_call_failure_hook=AsyncMock(),
+            post_call_response_headers_hook=AsyncMock(return_value={}),
             update_request_status=AsyncMock(),
         ),
     )

tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py

@@ -68,6 +68,7 @@ async def test_audio_speech_success_does_not_call_post_call_success_hook(
     mock_logging.post_call_failure_hook = mock_failure_hook
     mock_logging.pre_call_hook = mock_pre_call
     mock_logging.update_request_status = mock_update_status
+    mock_logging.post_call_response_headers_hook = AsyncMock(return_value={})
 
     async def _mock_route_request(*, data, route_type, llm_router, user_model):
         assert route_type == "aspeech"