Merge deb53fe784 into 85bc58085e

litellm/llms/custom_httpx/http_handler.py

@@ -8,7 +8,7 @@ import sys
 import threading
 import time
 import weakref
-from collections.abc import AsyncIterable, Callable, Iterable, Mapping
+from collections.abc import AsyncIterable, AsyncIterator, Callable, Iterable, Mapping
 from http.cookiejar import CookieJar, DefaultCookiePolicy
 from io import BytesIO
 from types import MappingProxyType
@@ -256,6 +256,35 @@ def _anchor_handler_to(response: httpx.Response, handler: object) -> None:
     weakref.finalize(response, _drop_streaming_anchor, handler)
 
 
+class _RetryClientStream(httpx.AsyncByteStream):
+    def __init__(self, stream: httpx.AsyncByteStream, client: httpx.AsyncClient) -> None:
+        self._stream = stream
+        self._client = client
+
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        try:
+            async for chunk in self._stream:
+                yield chunk
+        finally:
+            await self.aclose()
+
+    async def aclose(self) -> None:
+        if not self._client.is_closed:
+            try:
+                await self._stream.aclose()
+            finally:
+                await self._client.aclose()
+
+
+async def _hand_retry_client_to(response: httpx.Response, client: httpx.AsyncClient) -> httpx.Response:
+    if response.is_closed:
+        await client.aclose()
+    else:
+        assert isinstance(response.stream, httpx.AsyncByteStream), "An async retry must return an async stream"
+        response.stream = _RetryClientStream(response.stream, client)
+    return response
+
+
 def blocked_cookie_jar() -> CookieJar:
     """A jar that stores no response cookie and sends none, for httpx clients.
 
@@ -983,8 +1012,9 @@ class AsyncHTTPHandler:
                     stream=stream,
                     content=content,
                 )
-            finally:
+            except BaseException:
                 await new_client.aclose()
+                raise
         except httpx.TimeoutException as e:
             _raise_timeout(e, start_time)
         except httpx.HTTPStatusError as e:
@@ -1046,9 +1076,10 @@ class AsyncHTTPHandler:
                     retried.raise_for_status()
                 except httpx.HTTPStatusError as retried_error:
                     await _raise_masked_async_error(retried_error, stream)
-                return retried
-            finally:
+                return await _hand_retry_client_to(retried, new_client)
+            except BaseException:
                 await new_client.aclose()
+                raise
         except httpx.TimeoutException as e:
             _raise_timeout(e, start_time)
         except httpx.HTTPStatusError as e:
@@ -1103,8 +1134,9 @@ class AsyncHTTPHandler:
                     stream=stream,
                     content=content,
                 )
-            finally:
+            except BaseException:
                 await new_client.aclose()
+                raise
         except httpx.TimeoutException as e:
             _raise_timeout(e, start_time)
         except httpx.HTTPStatusError as e:
@@ -1160,8 +1192,9 @@ class AsyncHTTPHandler:
                     stream=stream,
                     content=content,
                 )
-            finally:
+            except BaseException:
                 await new_client.aclose()
+                raise
         except httpx.HTTPStatusError as e:
             await _raise_masked_async_error(e, stream)
         except Exception as e:
@@ -1198,7 +1231,7 @@ class AsyncHTTPHandler:
         )
         response: Final = await client.send(req, stream=stream)
         response.raise_for_status()
-        return response
+        return await _hand_retry_client_to(response, client)
 
     # Strong references to finalizer-scheduled client-close tasks. A bare
     # create_task() result may be garbage-collected before it runs, leaving

tests/unit/llms/custom_httpx/test_http_handler.py

@@ -1528,6 +1528,143 @@ class _RetryClientHandler(AsyncHTTPHandler):
         return self._retry_client
 
 
+class _ClientBoundStream(httpx.AsyncByteStream):
+    def __init__(self, client: httpx.AsyncClient, error: Exception | None = None) -> None:
+        self.client: Final = client
+        self.error: Final = error
+        self.close_count = 0
+
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        if self.client.is_closed:
+            raise httpx.ReadError("connection closed before reading")
+        yield b"data: first\n\n"
+        if self.error is not None:
+            raise self.error
+        yield b"data: [DONE]\n\n"
+
+    async def aclose(self) -> None:
+        self.close_count += 1
+
+
+class _PendingClientBoundStream(_ClientBoundStream):
+    def __init__(self, client: httpx.AsyncClient) -> None:
+        super().__init__(client)
+        self.started: Final = asyncio.Event()
+        self.release: Final = asyncio.Event()
+
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        self.started.set()
+        await self.release.wait()
+        yield b"data: [DONE]\n\n"
+
+
+async def _fail_connection(request: httpx.Request) -> httpx.Response:
+    raise httpx.ConnectError("first connection failed", request=request)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["post", "put", "patch", "delete"])
+@pytest.mark.parametrize("streaming", [True, False])
+async def test_retry_body_survives_until_consumed_and_releases_client(method: str, streaming: bool) -> None:
+    retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, stream=body)))
+    body: Final = _ClientBoundStream(retry)
+    async with httpx.AsyncClient(transport=httpx.MockTransport(_fail_connection)) as first:
+        handler: Final = _RetryClientHandler(first, retry)
+        try:
+            response: Final = await getattr(handler, method)("https://upstream.test/stream", stream=streaming)
+            assert retry.is_closed is not streaming
+            assert await response.aread() == b"data: first\n\ndata: [DONE]\n\n"
+            assert retry.is_closed
+            assert body.close_count == 1
+            await response.aclose()
+            assert body.close_count == 1
+        finally:
+            await retry.aclose()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["post", "put", "patch", "delete"])
+async def test_retry_stream_explicit_close_releases_client(method: str) -> None:
+    retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, stream=body)))
+    body: Final = _ClientBoundStream(retry)
+    async with httpx.AsyncClient(transport=httpx.MockTransport(_fail_connection)) as first:
+        handler: Final = _RetryClientHandler(first, retry)
+        try:
+            response: Final = await getattr(handler, method)("https://upstream.test/stream", stream=True)
+            assert not retry.is_closed
+            await response.aclose()
+            await response.aclose()
+            assert retry.is_closed
+            assert body.close_count == 1
+        finally:
+            await retry.aclose()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["post", "put", "patch", "delete"])
+async def test_retry_stream_read_error_releases_client_and_preserves_error(method: str) -> None:
+    failure: Final = httpx.ReadError("upstream truncated the body")
+    retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, stream=body)))
+    body: Final = _ClientBoundStream(retry, failure)
+    async with httpx.AsyncClient(transport=httpx.MockTransport(_fail_connection)) as first:
+        handler: Final = _RetryClientHandler(first, retry)
+        try:
+            response: Final = await getattr(handler, method)("https://upstream.test/stream", stream=True)
+            with pytest.raises(httpx.ReadError) as raised:
+                await response.aread()
+            assert raised.value is failure
+            assert retry.is_closed
+            assert body.close_count == 1
+            await response.aclose()
+            assert body.close_count == 1
+        finally:
+            await retry.aclose()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["post", "put", "patch", "delete"])
+async def test_failed_retry_releases_client(method: str) -> None:
+    retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(_fail_connection))
+    async with httpx.AsyncClient(transport=httpx.MockTransport(_fail_connection)) as first:
+        handler: Final = _RetryClientHandler(first, retry)
+        with pytest.raises(httpx.ConnectError):
+            await getattr(handler, method)("https://upstream.test/stream", stream=True)
+        assert retry.is_closed
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["post", "put", "patch", "delete"])
+async def test_buffered_retry_response_releases_client_even_when_streaming_requested(method: str) -> None:
+    retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, content=b"done")))
+    async with httpx.AsyncClient(transport=httpx.MockTransport(_fail_connection)) as first:
+        handler: Final = _RetryClientHandler(first, retry)
+        response: Final = await getattr(handler, method)("https://upstream.test/stream", stream=True)
+        assert await response.aread() == b"done"
+        assert retry.is_closed
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["post", "put", "patch", "delete"])
+async def test_retry_stream_cancellation_releases_client(method: str) -> None:
+    retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, stream=body)))
+    body: Final = _PendingClientBoundStream(retry)
+    async with httpx.AsyncClient(transport=httpx.MockTransport(_fail_connection)) as first:
+        handler: Final = _RetryClientHandler(first, retry)
+        try:
+            response: Final = await getattr(handler, method)("https://upstream.test/stream", stream=True)
+            reading: Final = asyncio.create_task(response.aread())
+            await body.started.wait()
+            reading.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await reading
+            assert retry.is_closed
+            assert body.close_count == 1
+            await response.aclose()
+            assert body.close_count == 1
+        finally:
+            await retry.aclose()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("method", ["post", "put", "patch", "delete"])
 async def test_connection_error_retry_forwards_content(method: str):