diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 3a6fa5f51da..bd66530c51f 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/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 diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 3e583bf8dff..75865572052 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/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):