diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index cb37725d79cf..4bf36a0d5c63 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/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: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 2357960f716b..45a49def59d5 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/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 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3c10239f8689..05401488f272 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2084,12 +2084,18 @@ 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 @@ -2097,7 +2103,10 @@ async def async_anthropic_messages_handler( # 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, diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index b18af060a201..0b4187d1bcf3 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/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.