From 23614f5f5440bc8076007b1682403d4757d17106 Mon Sep 17 00:00:00 2001 From: Simon Sorg Date: Tue, 3 Nov 2026 16:19:03 +0100 Subject: [PATCH 1/2] fix(http): preserve retry client lifetime for streamed responses --- litellm/llms/custom_httpx/http_handler.py | 151 +++++++++++------- .../llms/custom_httpx/test_http_handler.py | 137 ++++++++++++++++ 2 files changed, 228 insertions(+), 60 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 3a6fa5f51da..4a5a9b78353 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,29 @@ def _anchor_handler_to(response: httpx.Response, handler: object) -> None: weakref.finalize(response, _drop_streaming_anchor, handler) +class _SingleUseClientStream(httpx.AsyncByteStream): + def __init__(self, stream: httpx.AsyncByteStream, client: httpx.AsyncClient) -> None: + self._stream = stream + self._client = client + self._closed = False + + 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 self._closed: + return + self._closed = True + try: + await self._stream.aclose() + finally: + await self._client.aclose() + + def blocked_cookie_jar() -> CookieJar: """A jar that stores no response cookie and sends none, for httpx clients. @@ -971,20 +994,17 @@ class AsyncHTTPHandler: except httpx.ConnectError: # Retry the request with a new session if there is a connection error new_client: Final = self.create_client(timeout=timeout, event_hooks=self.event_hooks) - try: - return await self.single_connection_request( - url=url, - method="POST", - client=new_client, - data=data, - json=json, - params=params, - headers=headers, - stream=stream, - content=content, - ) - finally: - await new_client.aclose() + return await self.single_connection_request( + url=url, + method="POST", + client=new_client, + data=data, + json=json, + params=params, + headers=headers, + stream=stream, + content=content, + ) except httpx.TimeoutException as e: _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: @@ -1046,9 +1066,18 @@ class AsyncHTTPHandler: retried.raise_for_status() except httpx.HTTPStatusError as retried_error: await _raise_masked_async_error(retried_error, stream) - return retried - finally: + except BaseException: await new_client.aclose() + raise + else: + if stream and not retried.is_closed: + assert isinstance(retried.stream, httpx.AsyncByteStream), ( + "An async retry must return an async stream" + ) + retried.stream = _SingleUseClientStream(retried.stream, new_client) + else: + await new_client.aclose() + return retried except httpx.TimeoutException as e: _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: @@ -1091,20 +1120,17 @@ class AsyncHTTPHandler: except httpx.ConnectError: # Retry the request with a new session if there is a connection error new_client: Final = self.create_client(timeout=timeout, event_hooks=self.event_hooks) - try: - return await self.single_connection_request( - url=url, - method="PATCH", - client=new_client, - data=data, - json=json, - params=params, - headers=headers, - stream=stream, - content=content, - ) - finally: - await new_client.aclose() + return await self.single_connection_request( + url=url, + method="PATCH", + client=new_client, + data=data, + json=json, + params=params, + headers=headers, + stream=stream, + content=content, + ) except httpx.TimeoutException as e: _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: @@ -1148,20 +1174,17 @@ class AsyncHTTPHandler: except (httpx.RemoteProtocolError, httpx.ConnectError): # Retry the request with a new session if there is a connection error new_client: Final = self.create_client(timeout=timeout, event_hooks=self.event_hooks) - try: - return await self.single_connection_request( - url=url, - method="DELETE", - client=new_client, - data=data, - json=json, - params=params, - headers=headers, - stream=stream, - content=content, - ) - finally: - await new_client.aclose() + return await self.single_connection_request( + url=url, + method="DELETE", + client=new_client, + data=data, + json=json, + params=params, + headers=headers, + stream=stream, + content=content, + ) except httpx.HTTPStatusError as e: await _raise_masked_async_error(e, stream) except Exception as e: @@ -1184,21 +1207,29 @@ class AsyncHTTPHandler: Used for retrying connection client errors. """ - # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) - request_data, request_content = _prepare_request_data_and_content(data, content) - - req: Final = client.build_request( - method, - url, - data=request_data, - json=json, - params=params, - headers=headers, - content=request_content, - ) - response: Final = await client.send(req, stream=stream) - response.raise_for_status() - return response + try: + request_data, request_content = _prepare_request_data_and_content(data, content) + req: Final = client.build_request( + method, + url, + data=request_data, + json=json, + params=params, + headers=headers, + content=request_content, + ) + response: Final = await client.send(req, stream=stream) + response.raise_for_status() + except BaseException: + await client.aclose() + raise + else: + if stream and not response.is_closed: + assert isinstance(response.stream, httpx.AsyncByteStream), "An async retry must return an async stream" + response.stream = _SingleUseClientStream(response.stream, client) + else: + await client.aclose() + return response # 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): From deb53fe784bccf51fb9a65833c819e4284ca819f Mon Sep 17 00:00:00 2001 From: Simon Sorg Date: Thu, 5 Nov 2026 19:38:50 +0100 Subject: [PATCH 2/2] refactor(http): shrink retry client stream ownership fix Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe --- litellm/llms/custom_httpx/http_handler.py | 150 +++++++++++----------- 1 file changed, 76 insertions(+), 74 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 4a5a9b78353..bd66530c51f 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -256,11 +256,10 @@ def _anchor_handler_to(response: httpx.Response, handler: object) -> None: weakref.finalize(response, _drop_streaming_anchor, handler) -class _SingleUseClientStream(httpx.AsyncByteStream): +class _RetryClientStream(httpx.AsyncByteStream): def __init__(self, stream: httpx.AsyncByteStream, client: httpx.AsyncClient) -> None: self._stream = stream self._client = client - self._closed = False async def __aiter__(self) -> AsyncIterator[bytes]: try: @@ -270,13 +269,20 @@ class _SingleUseClientStream(httpx.AsyncByteStream): await self.aclose() async def aclose(self) -> None: - if self._closed: - return - self._closed = True - try: - await self._stream.aclose() - finally: - await self._client.aclose() + 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: @@ -994,17 +1000,21 @@ class AsyncHTTPHandler: except httpx.ConnectError: # Retry the request with a new session if there is a connection error new_client: Final = self.create_client(timeout=timeout, event_hooks=self.event_hooks) - return await self.single_connection_request( - url=url, - method="POST", - client=new_client, - data=data, - json=json, - params=params, - headers=headers, - stream=stream, - content=content, - ) + try: + return await self.single_connection_request( + url=url, + method="POST", + client=new_client, + data=data, + json=json, + params=params, + headers=headers, + stream=stream, + content=content, + ) + except BaseException: + await new_client.aclose() + raise except httpx.TimeoutException as e: _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: @@ -1066,18 +1076,10 @@ class AsyncHTTPHandler: retried.raise_for_status() except httpx.HTTPStatusError as retried_error: await _raise_masked_async_error(retried_error, stream) + return await _hand_retry_client_to(retried, new_client) except BaseException: await new_client.aclose() raise - else: - if stream and not retried.is_closed: - assert isinstance(retried.stream, httpx.AsyncByteStream), ( - "An async retry must return an async stream" - ) - retried.stream = _SingleUseClientStream(retried.stream, new_client) - else: - await new_client.aclose() - return retried except httpx.TimeoutException as e: _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: @@ -1120,17 +1122,21 @@ class AsyncHTTPHandler: except httpx.ConnectError: # Retry the request with a new session if there is a connection error new_client: Final = self.create_client(timeout=timeout, event_hooks=self.event_hooks) - return await self.single_connection_request( - url=url, - method="PATCH", - client=new_client, - data=data, - json=json, - params=params, - headers=headers, - stream=stream, - content=content, - ) + try: + return await self.single_connection_request( + url=url, + method="PATCH", + client=new_client, + data=data, + json=json, + params=params, + headers=headers, + stream=stream, + content=content, + ) + except BaseException: + await new_client.aclose() + raise except httpx.TimeoutException as e: _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: @@ -1174,17 +1180,21 @@ class AsyncHTTPHandler: except (httpx.RemoteProtocolError, httpx.ConnectError): # Retry the request with a new session if there is a connection error new_client: Final = self.create_client(timeout=timeout, event_hooks=self.event_hooks) - return await self.single_connection_request( - url=url, - method="DELETE", - client=new_client, - data=data, - json=json, - params=params, - headers=headers, - stream=stream, - content=content, - ) + try: + return await self.single_connection_request( + url=url, + method="DELETE", + client=new_client, + data=data, + json=json, + params=params, + headers=headers, + stream=stream, + content=content, + ) + except BaseException: + await new_client.aclose() + raise except httpx.HTTPStatusError as e: await _raise_masked_async_error(e, stream) except Exception as e: @@ -1207,29 +1217,21 @@ class AsyncHTTPHandler: Used for retrying connection client errors. """ - try: - request_data, request_content = _prepare_request_data_and_content(data, content) - req: Final = client.build_request( - method, - url, - data=request_data, - json=json, - params=params, - headers=headers, - content=request_content, - ) - response: Final = await client.send(req, stream=stream) - response.raise_for_status() - except BaseException: - await client.aclose() - raise - else: - if stream and not response.is_closed: - assert isinstance(response.stream, httpx.AsyncByteStream), "An async retry must return an async stream" - response.stream = _SingleUseClientStream(response.stream, client) - else: - await client.aclose() - return response + # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) + request_data, request_content = _prepare_request_data_and_content(data, content) + + req: Final = client.build_request( + method, + url, + data=request_data, + json=json, + params=params, + headers=headers, + content=request_content, + ) + response: Final = await client.send(req, stream=stream) + response.raise_for_status() + 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