From deb53fe784bccf51fb9a65833c819e4284ca819f Mon Sep 17 00:00:00 2001 From: Simon Sorg Date: Thu, 5 Nov 2026 19:38:50 +0100 Subject: [PATCH] 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