diff --git a/src/anthropic/lib/middleware/_fallbacks.py b/src/anthropic/lib/middleware/_fallbacks.py index 15edbc232..cb6426392 100644 --- a/src/anthropic/lib/middleware/_fallbacks.py +++ b/src/anthropic/lib/middleware/_fallbacks.py @@ -431,7 +431,9 @@ def _splice_fallback_stream( pin=pin, ) return APIResponse( - raw=_spliced_http_response(response.http_response, _FrameByteStream(frames)), + raw=_spliced_http_response( + response.http_response, _FrameByteStream(frames, response=response.http_response) + ), cast_to=response._cast_to, client=response._client, stream=True, @@ -461,7 +463,9 @@ def _splice_fallback_stream_async( pin=pin, ) return AsyncAPIResponse( - raw=_spliced_http_response(response.http_response, _AsyncFrameByteStream(frames)), + raw=_spliced_http_response( + response.http_response, _AsyncFrameByteStream(frames, response=response.http_response) + ), cast_to=response._cast_to, client=response._client, stream=True, @@ -1692,8 +1696,9 @@ def _spliced_http_response( class _FrameByteStream(httpx2.SyncByteStream): - def __init__(self, frames: Generator[bytes, None, None]) -> None: + def __init__(self, frames: Generator[bytes, None, None], *, response: httpx2.Response) -> None: self._frames = frames + self._response = response @override def __iter__(self) -> Iterator[bytes]: @@ -1701,12 +1706,17 @@ def __iter__(self) -> Iterator[bytes]: @override def close(self) -> None: - self._frames.close() + try: + self._frames.close() + finally: + # An unread generator never enters its finally block. + self._response.close() class _AsyncFrameByteStream(httpx2.AsyncByteStream): - def __init__(self, frames: AsyncGenerator[bytes, None]) -> None: + def __init__(self, frames: AsyncGenerator[bytes, None], *, response: httpx2.Response) -> None: self._frames = frames + self._response = response @override def __aiter__(self) -> AsyncIterator[bytes]: @@ -1714,4 +1724,8 @@ def __aiter__(self) -> AsyncIterator[bytes]: @override async def aclose(self) -> None: - await self._frames.aclose() + try: + await self._frames.aclose() + finally: + # An unread generator never enters its finally block. + await self._response.aclose() diff --git a/tests/lib/test_refusal_fallback_streaming.py b/tests/lib/test_refusal_fallback_streaming.py index e130240df..8de9261d9 100644 --- a/tests/lib/test_refusal_fallback_streaming.py +++ b/tests/lib/test_refusal_fallback_streaming.py @@ -6,6 +6,7 @@ import logging from typing import Any, List, Protocol, cast from pathlib import Path +from unittest.mock import Mock from typing_extensions import Literal import httpx2 @@ -1303,6 +1304,36 @@ def test_entry_overrides_apply_to_the_hop_request(self, respx_mock: MockRouter) class TestCancellation: + def test_closing_an_unread_stream_closes_the_underlying_response(self) -> None: + response = httpx2.Response( + 200, headers={"content-type": "text/event-stream"}, stream=httpx2.ByteStream(STREAM_A.encode()) + ) + handler = Mock(return_value=response) + with make_sync_client( + http_client=httpx2.Client(transport=httpx2.MockTransport(handler)), + middleware=[BetaRefusalFallbackMiddleware(FALLBACKS)], + ) as client: + with create_stream(client): + assert not response.is_closed + + assert response.is_closed + handler.assert_called_once() + + async def test_closing_an_unread_async_stream_closes_the_underlying_response(self) -> None: + response = httpx2.Response( + 200, headers={"content-type": "text/event-stream"}, stream=httpx2.ByteStream(STREAM_A.encode()) + ) + handler = Mock(return_value=response) + async with make_async_client( + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), + middleware=[BetaRefusalFallbackMiddleware(FALLBACKS)], + ) as client: + async with await create_stream_async(client): + assert not response.is_closed + + assert response.is_closed + handler.assert_called_once() + @pytest.mark.respx(base_url=base_url) def test_closing_the_stream_mid_passthrough_tears_down_without_a_fallback_request( self, respx_mock: MockRouter