fix(anthropic_messages): forward provider response headers on streaming /v1/messages responses

litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py

@@ -193,6 +193,14 @@ async def __anext__(self) -> bytes:
 
         raise StopAsyncIteration
 
+    async def aclose(self) -> None:
+        from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
+            aclose_if_supported,
+        )
+
+        await aclose_if_supported(self._inner)
+        await aclose_if_supported(self._follow_up_iterator)
+
     async def _process_agentic_hooks(self) -> None:
         """Rebuild the Anthropic response from collected SSE bytes and call hooks."""
         if self._hook_processing_done:

litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py

@@ -1,8 +1,13 @@
 import asyncio
 import json
 from datetime import datetime
-from typing import Any, AsyncIterator, List, Union
+from typing import Any, AsyncIterator, List, Protocol, Union, runtime_checkable
 
+import httpx
+from pydantic import TypeAdapter
+from typing_extensions import TypedDict
+
+from litellm.litellm_core_utils.core_helpers import process_response_headers
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.proxy.pass_through_endpoints.success_handler import (
     PassThroughEndpointLogging,
@@ -13,6 +18,58 @@
 GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ = PassThroughEndpointLogging()
 
 
+class AnthropicMessagesStreamHiddenParams(TypedDict):
+    additional_headers: dict[str, str]
+
+
+@runtime_checkable
+class SupportsAclose(Protocol):
+    async def aclose(self) -> None: ...
+
+
+async def aclose_if_supported(stream: object) -> None:
+    if isinstance(stream, SupportsAclose):
+        await stream.aclose()
+
+
+_RESPONSE_HEADERS_ADAPTER: TypeAdapter[dict[str, str]] = TypeAdapter(dict[str, str])
+
+
+def anthropic_messages_stream_hidden_params(
+    response_headers: httpx.Headers,
+) -> AnthropicMessagesStreamHiddenParams:
+    return AnthropicMessagesStreamHiddenParams(
+        additional_headers=_RESPONSE_HEADERS_ADAPTER.validate_python(process_response_headers(response_headers))
+    )
+
+
+class AnthropicMessagesStreamingResponse:
+    """
+    Wraps the /v1/messages SSE byte stream so upstream provider response
+    headers (e.g. Bedrock's x-amzn-requestid / x-amzn-trace-id) survive as
+    ``_hidden_params["additional_headers"]``, which the proxy forwards to
+    clients as ``llm_provider-*`` response headers. Bare async generators
+    cannot carry attributes, so header context was previously dropped.
+    """
+
+    def __init__(
+        self,
+        completion_stream: AsyncIterator[bytes],
+        hidden_params: AnthropicMessagesStreamHiddenParams,
+    ) -> None:
+        self.completion_stream = completion_stream
+        self._hidden_params = hidden_params
+
+    def __aiter__(self) -> "AnthropicMessagesStreamingResponse":
+        return self
+
+    async def __anext__(self) -> bytes:
+        return await self.completion_stream.__anext__()
+
+    async def aclose(self) -> None:
+        await aclose_if_supported(self.completion_stream)
+
+
 class BaseAnthropicMessagesStreamingIterator:
     """
     Base class for Anthropic Messages streaming iterators that provides common logic

litellm/llms/custom_httpx/llm_http_handler.py

@@ -2084,20 +2084,29 @@ async def async_anthropic_messages_handler(
 
         initial_response: Union[AsyncIterator, AnthropicMessagesResponse]
         if stream:
+            from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
+                AnthropicMessagesStreamingResponse,
+                anthropic_messages_stream_hidden_params,
+            )
+
             completion_stream = anthropic_messages_provider_config.get_async_streaming_response_iterator(
                 model=model,
                 httpx_response=response,
                 request_body=request_body,
                 litellm_logging_obj=logging_obj,
             )
+            stream_hidden_params = anthropic_messages_stream_hidden_params(response.headers)
 
             if not self._has_agentic_completion_hook(logging_obj):
                 # No callback overrides async_should_run_agentic_loop, so the
                 # agentic wrapper's only effect would be buffering every chunk
                 # and rebuilding the response from SSE at end-of-stream to call
                 # hooks that all return (False, {}). Stream through directly and
                 # skip that per-chunk + end-of-stream overhead.
-                return completion_stream
+                return AnthropicMessagesStreamingResponse(
+                    completion_stream=completion_stream,
+                    hidden_params=stream_hidden_params,
+                )
 
             from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
                 AgenticAnthropicStreamingIterator,
@@ -2114,7 +2123,10 @@ async def async_anthropic_messages_handler(
                 custom_llm_provider=custom_llm_provider,
                 kwargs={**kwargs, "api_key": api_key} if api_key else kwargs,
             )
-            return initial_response
+            return AnthropicMessagesStreamingResponse(
+                completion_stream=initial_response,
+                hidden_params=stream_hidden_params,
+            )
         else:
             initial_response = anthropic_messages_provider_config.transform_anthropic_messages_response(
                 model=model,

tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py

@@ -435,6 +435,227 @@ def capture_validate(*args, **kwargs):
         assert captured_headers["X-Auth-Token"] == "token123"
 
 
+@pytest.mark.asyncio
+async def test_async_anthropic_messages_handler_streaming_forwards_provider_response_headers():
+    """
+    Regression test for LIT-3724 (issue 2): streaming /v1/messages responses
+    dropped the upstream provider's HTTP response headers, so Bedrock's
+    x-amzn-requestid / x-amzn-trace-id never reached clients even with
+    `return_response_headers: true`. The returned stream object must carry
+    them in `_hidden_params["additional_headers"]` (llm_provider-* prefixed),
+    which the proxy merges into the client-facing response headers.
+    """
+    from collections.abc import AsyncIterator as ABCAsyncIterator
+
+    from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
+        AnthropicMessagesConfig,
+    )
+
+    handler = BaseLLMHTTPHandler()
+
+    sse_body = (
+        b'event: message_start\ndata: {"type": "message_start"}\n\n'
+        b'event: message_stop\ndata: {"type": "message_stop"}\n\n'
+    )
+    upstream_response = httpx.Response(
+        200,
+        headers={
+            "x-amzn-requestid": "amzn-req-123",
+            "x-amzn-trace-id": "Root=1-abc-def",
+        },
+        content=sse_body,
+        request=httpx.Request("POST", "https://api.anthropic.com/v1/messages"),
+    )
+    mock_client = AsyncMock(spec=AsyncHTTPHandler)
+    mock_client.post = AsyncMock(return_value=upstream_response)
+
+    mock_logging_obj = Mock()
+    mock_logging_obj.model_call_details = {}
+
+    result = await handler.async_anthropic_messages_handler(
+        model="claude-sonnet-4-20250514",
+        messages=[{"role": "user", "content": "Hello"}],
+        anthropic_messages_provider_config=AnthropicMessagesConfig(),
+        anthropic_messages_optional_request_params={"max_tokens": 32},
+        custom_llm_provider="anthropic",
+        litellm_params=GenericLiteLLMParams(),
+        logging_obj=mock_logging_obj,
+        client=mock_client,
+        api_key="sk-test",
+        stream=True,
+        kwargs={},
+    )
+
+    assert isinstance(result, ABCAsyncIterator)
+
+    additional_headers = result._hidden_params["additional_headers"]
+    assert additional_headers["llm_provider-x-amzn-requestid"] == "amzn-req-123"
+    assert additional_headers["llm_provider-x-amzn-trace-id"] == "Root=1-abc-def"
+
+    collected = b"".join([chunk async for chunk in result])
+    assert b"message_start" in collected
+    assert b"message_stop" in collected
+
+
+@pytest.mark.asyncio
+async def test_async_anthropic_messages_handler_agentic_streaming_forwards_provider_response_headers():
+    """
+    Companion to the test above for the agentic branch: when a callback
+    overrides async_should_run_agentic_loop, the handler wraps
+    AgenticAnthropicStreamingIterator in AnthropicMessagesStreamingResponse.
+    That wrapping must still expose the provider headers and delegate
+    iteration through the two-phase agentic iterator unchanged.
+    """
+    from collections.abc import AsyncIterator as ABCAsyncIterator
+
+    from litellm.integrations.custom_logger import CustomLogger
+    from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
+        AgenticAnthropicStreamingIterator,
+    )
+    from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
+        AnthropicMessagesConfig,
+    )
+
+    class NoOpAgenticCallback(CustomLogger):
+        async def async_should_run_agentic_loop(
+            self,
+            response,
+            model,
+            messages,
+            tools,
+            stream,
+            custom_llm_provider,
+            kwargs,
+        ):
+            return False, {}
+
+    handler = BaseLLMHTTPHandler()
+
+    sse_body = (
+        b'event: message_start\ndata: {"type": "message_start"}\n\n'
+        b'event: message_stop\ndata: {"type": "message_stop"}\n\n'
+    )
+    upstream_response = httpx.Response(
+        200,
+        headers={"x-amzn-requestid": "amzn-req-456"},
+        content=sse_body,
+        request=httpx.Request("POST", "https://api.anthropic.com/v1/messages"),
+    )
+    mock_client = AsyncMock(spec=AsyncHTTPHandler)
+    mock_client.post = AsyncMock(return_value=upstream_response)
+
+    mock_logging_obj = Mock()
+    mock_logging_obj.model_call_details = {}
+    mock_logging_obj.dynamic_success_callbacks = [NoOpAgenticCallback()]
+
+    result = await handler.async_anthropic_messages_handler(
+        model="claude-sonnet-4-20250514",
+        messages=[{"role": "user", "content": "Hello"}],
+        anthropic_messages_provider_config=AnthropicMessagesConfig(),
+        anthropic_messages_optional_request_params={"max_tokens": 32},
+        custom_llm_provider="anthropic",
+        litellm_params=GenericLiteLLMParams(),
+        logging_obj=mock_logging_obj,
+        client=mock_client,
+        api_key="sk-test",
+        stream=True,
+        kwargs={},
+    )
+
+    assert isinstance(result, ABCAsyncIterator)
+    assert isinstance(result.completion_stream, AgenticAnthropicStreamingIterator)
+    assert result._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "amzn-req-456"
+
+    collected = b"".join([chunk async for chunk in result])
+    assert b"message_start" in collected
+    assert b"message_stop" in collected
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_streaming_response_aclose_closes_upstream_stream():
+    """
+    Regression test: the proxy's streaming cleanup calls aclose on the
+    handler's return value (see _finalize_streaming_generator_cleanup's
+    hasattr(response, "aclose") check). The wrapper must forward aclose to
+    the upstream stream so provider connections are released on client
+    disconnect instead of lingering until garbage collection.
+    """
+    from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
+        AnthropicMessagesStreamingResponse,
+    )
+
+    class UpstreamTracker:
+        def __init__(self):
+            self.closed = False
+
+    tracker = UpstreamTracker()
+
+    async def upstream():
+        try:
+            yield b'data: {"type": "message_start"}\n\n'
+            yield b'data: {"type": "message_stop"}\n\n'
+        finally:
+            tracker.closed = True
+
+    stream = AnthropicMessagesStreamingResponse(
+        completion_stream=upstream(),
+        hidden_params={"additional_headers": {}},
+    )
+
+    first_chunk = await stream.__anext__()
+    assert b"message_start" in first_chunk
+    assert tracker.closed is False
+
+    await stream.aclose()
+    assert tracker.closed is True
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_streaming_response_aclose_closes_agentic_upstream_stream():
+    from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
+        AgenticAnthropicStreamingIterator,
+    )
+    from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
+        AnthropicMessagesStreamingResponse,
+    )
+
+    class UpstreamTracker:
+        def __init__(self):
+            self.closed = False
+
+    tracker = UpstreamTracker()
+
+    async def upstream():
+        try:
+            yield b'data: {"type": "message_start"}\n\n'
+            yield b'data: {"type": "message_stop"}\n\n'
+        finally:
+            tracker.closed = True
+
+    agentic_iterator = AgenticAnthropicStreamingIterator(
+        completion_stream=upstream(),
+        http_handler=Mock(),
+        model="claude-sonnet-4-20250514",
+        messages=[{"role": "user", "content": "Hello"}],
+        anthropic_messages_provider_config=Mock(),
+        anthropic_messages_optional_request_params={},
+        logging_obj=Mock(),
+        custom_llm_provider="anthropic",
+        kwargs={},
+    )
+    stream = AnthropicMessagesStreamingResponse(
+        completion_stream=agentic_iterator,
+        hidden_params={"additional_headers": {}},
+    )
+
+    first_chunk = await stream.__anext__()
+    assert b"message_start" in first_chunk
+    assert tracker.closed is False
+
+    await stream.aclose()
+    assert tracker.closed is True
+
+
 @pytest.mark.asyncio
 async def test_async_anthropic_messages_handler_passes_litellm_metadata():
     """Ensure litellm_metadata from kwargs is forwarded via update_from_kwargs.