diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index df5b10b3bdc..2c5f455692c 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -1,10 +1,11 @@ import asyncio +import concurrent.futures import contextlib import os import ssl import typing import urllib.request -from typing import Any, Callable, Dict, Optional, Union +from typing import Any, Callable, ClassVar, Dict, Optional, Union import aiohttp import aiohttp.client_exceptions @@ -138,6 +139,11 @@ class LiteLLMAiohttpTransport(AiohttpTransport): Credit to: https://github.com/karpetrosyan/httpx-aiohttp for this implementation """ + # Strong references to scheduled session-close tasks. A bare + # asyncio.create_task() result may be garbage-collected before it runs, + # leaving the recycled session unclosed ("Unclosed client session"). + _background_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes + def __init__( self, client: Union[ClientSession, Callable[[], ClientSession]], @@ -164,6 +170,92 @@ class LiteLLMAiohttpTransport(AiohttpTransport): self._owns_session = True return session + @classmethod + def _on_close_task_done(cls, task: "asyncio.Task[None]") -> None: + cls._background_close_tasks.discard(task) + if task.cancelled(): + return + exc = task.exception() + if exc is not None: + verbose_logger.debug("Error closing recycled aiohttp session: %s", exc) + + @staticmethod + def _on_threadsafe_close_done(future: "concurrent.futures.Future[None]") -> None: + if future.cancelled(): + return + exc = future.exception() + if exc is not None: + verbose_logger.debug("Error closing recycled aiohttp session on its own loop: %s", exc) + + @staticmethod + def _mark_connector_closed(session: ClientSession) -> None: + """Synchronously dispose a session whose event loop is gone. + + An async close can no longer run on a closed loop. BaseConnector._close + is the same synchronous teardown aiohttp's own finalizer (__del__) + uses: it is guarded for closed loops, releases pooled connections, and + flips the flags that ClientSession.closed / BaseConnector.closed read - + so no "Unclosed client session" / "Unclosed connector" warnings reach + the event-loop exception handler at garbage collection. + """ + connector = getattr(session, "_connector", None) + close_sync = getattr(connector, "_close", None) + if not callable(close_sync): + return + try: + close_sync() + except (RuntimeError, AttributeError, OSError) as e: + verbose_logger.debug("Best-effort connector close failed: %s", e) + + def _close_recycled_session(self, session: ClientSession) -> None: + """Deterministically dispose a ClientSession this transport is replacing. + + Covers the three lifecycles a recycled session can be in: + - its loop is the current running loop: schedule an async close and keep + a strong reference to the task until it completes; + - its loop is still running elsewhere (e.g. another thread): hand the + close to that loop thread-safely; + - its loop is stopped or closed, or there is no running loop: fall + back to the synchronous finalizer-safe teardown. + """ + if session.closed: + return + + session_loop = getattr(session, "_loop", None) + try: + current_loop: Optional[asyncio.AbstractEventLoop] = asyncio.get_running_loop() + except RuntimeError: + current_loop = None + + if session_loop is not None and session_loop is not current_loop: + if not session_loop.is_closed() and session_loop.is_running(): + # The session's loop is running somewhere else (e.g. another + # thread): closing from here would touch that loop's internals + # unsafely; hand the close to its own loop. + try: + future = asyncio.run_coroutine_threadsafe(session.close(), session_loop) + except RuntimeError as e: # loop shut down between the checks + verbose_logger.debug("Threadsafe session close failed: %s", e) + self._mark_connector_closed(session) + else: + future.add_done_callback(self._on_threadsafe_close_done) + return + + # Foreign loop that is stopped or closed: an async close can no + # longer run there, and running it on the current loop would touch + # another loop's internals. Dispose synchronously instead. + self._mark_connector_closed(session) + return + + if current_loop is None: + self._mark_connector_closed(session) + return + + task = current_loop.create_task(session.close()) + cls = type(self) + cls._background_close_tasks.add(task) + task.add_done_callback(cls._on_close_task_done) + def _get_valid_client_session(self) -> ClientSession: """ Helper to get a valid ClientSession for the current event loop. @@ -193,21 +285,25 @@ class LiteLLMAiohttpTransport(AiohttpTransport): # Close old session to prevent leaks old_session = self.client try: - if self._owns_session and not old_session.closed: - try: - asyncio.create_task(old_session.close()) - except RuntimeError: - # Different event loop - can't schedule task, rely on GC - verbose_logger.debug("Old session from different loop, relying on GC") + if self._owns_session: + self._close_recycled_session(old_session) except Exception as e: verbose_logger.debug(f"Error closing old session: {e}") # Create a new session in the current event loop self.client = self._rebuild_session() - except (RuntimeError, AttributeError): - # If we can't check the loop or session is invalid, recreate it + except (RuntimeError, AttributeError) as e: + # If we can't check the loop or session is invalid, recreate it, + # but still dispose of the session being replaced. + old_session = self.client + if self._owns_session: + try: + self._close_recycled_session(old_session) + except (RuntimeError, AttributeError, OSError) as close_error: + verbose_logger.debug(f"Error closing old session: {close_error}") self.client = self._rebuild_session() + verbose_logger.debug(f"Error checking session loop, created new session: {e}") return self.client @@ -301,7 +397,14 @@ class LiteLLMAiohttpTransport(AiohttpTransport): # Handle the case where session was closed between our check and actual use if "Session is closed" in str(e): verbose_logger.debug(f"Session closed during request, retrying with new session: {e}") - # Force creation of a new session + # Dispose of the session that actually faulted. Do NOT read + # self.client here: a concurrent task may already have + # replaced it with a healthy session that must stay open. + # Guarded by isinstance: factory-injected sessions may be + # duck-typed test doubles without a close() coroutine. + # Read _owns_session before _rebuild_session() claims ownership. + if self._owns_session and isinstance(client_session, ClientSession): + self._close_recycled_session(client_session) self.client = self._rebuild_session() client_session = self.client diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py index 0550898c73d..b0b092a541f 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py @@ -1,4 +1,5 @@ import asyncio +import concurrent.futures import os import sys @@ -827,3 +828,305 @@ async def test_stale_loop_rebuild_does_not_close_unowned_session(): shared_session._loop = running_loop other_loop.close() await shared_session.close() + + +# --------------------------------------------------------------------------- +# Recycled-session leak tests (#24230) +# --------------------------------------------------------------------------- + + +async def _new_session() -> aiohttp.ClientSession: + return aiohttp.ClientSession() + + +def _make_session_on_dead_loop() -> aiohttp.ClientSession: + """Create a ClientSession bound to an event loop that is then closed. + + Runs in a worker thread: the caller may already be inside a running + event loop, where a nested run_until_complete is forbidden. + """ + import threading + + result: dict = {} + + def build() -> None: + loop = asyncio.new_event_loop() + try: + result["session"] = loop.run_until_complete(_new_session()) + finally: + loop.close() + + thread = threading.Thread(target=build) + thread.start() + thread.join(5) + return result["session"] + + +def _flaky_get_running_loop_factory(): + """get_running_loop stand-in that fails once, then delegates. + + Reproduces #24230: a transient loop-inspection failure sends + _get_valid_client_session into its (RuntimeError, AttributeError) + fallback branch. + """ + real_get_running_loop = asyncio.get_running_loop + calls = {"count": 0} + + def flaky(): + calls["count"] += 1 + if calls["count"] == 1: + raise RuntimeError("simulated loop inspection failure") + return real_get_running_loop() + + return flaky + + +@pytest.mark.asyncio +async def test_fallback_recreate_closes_previous_session(): + """ + Regression test for #24230: when loop inspection fails and the fallback + branch recreates the session, the replaced session must still be closed - + not silently abandoned to the garbage collector. + """ + from unittest.mock import patch + + old_session = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + transport.client = old_session + + with patch( + "litellm.llms.custom_httpx.aiohttp_transport.asyncio.get_running_loop", + side_effect=_flaky_get_running_loop_factory(), + ): + new_session = transport._get_valid_client_session() + + try: + assert new_session is not old_session + for _ in range(3): + await asyncio.sleep(0) + assert old_session.closed, "replaced session must be closed, not leaked" + finally: + await new_session.close() + if not old_session.closed: + await old_session.close() + + +@pytest.mark.asyncio +async def test_replaced_session_emits_no_unclosed_warnings(): + """ + Regression test for #24230: a session replaced by the fallback branch must + not surface "Unclosed client session" / "Unclosed connector" warnings when + the garbage collector finalizes it. + """ + import gc + import warnings as warnings_mod + from unittest.mock import patch + + old_session = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + transport.client = old_session + + with patch( + "litellm.llms.custom_httpx.aiohttp_transport.asyncio.get_running_loop", + side_effect=_flaky_get_running_loop_factory(), + ): + new_session = transport._get_valid_client_session() + + try: + for _ in range(3): + await asyncio.sleep(0) + + del old_session + with warnings_mod.catch_warnings(record=True) as caught: + warnings_mod.simplefilter("always") + gc.collect() + + unclosed = [ + str(w.message) + for w in caught + if "Unclosed client session" in str(w.message) or "Unclosed connector" in str(w.message) + ] + assert not unclosed, f"leaked session warnings: {unclosed}" + finally: + await new_session.close() + + +@pytest.mark.asyncio +async def test_dead_loop_session_closed_synchronously_on_recycle(): + """ + Regression test for #24230: a session whose event loop is already closed + cannot run an async close anywhere. Recycling it must dispose of it + deterministically, the session reads closed as soon as the recycle + returns, so no finalizer warning window remains. + """ + old_session = _make_session_on_dead_loop() + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + transport.client = old_session + + new_session = transport._get_valid_client_session() + + try: + assert new_session is not old_session + assert old_session.closed, "session from a closed loop must be disposed synchronously at recycle" + finally: + await new_session.close() + + +@pytest.mark.asyncio +async def test_close_task_strongly_referenced_until_done(): + """ + Regression test for #24230: scheduled session-close tasks must be strongly + referenced (and pruned on completion) so they cannot be garbage-collected + before they run. + """ + old_session = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + + transport._close_recycled_session(old_session) + + assert LiteLLMAiohttpTransport._background_close_tasks, "close task must be strongly referenced while pending" + for _ in range(5): + await asyncio.sleep(0) + assert old_session.closed + assert not LiteLLMAiohttpTransport._background_close_tasks, "completed close tasks must be pruned from the registry" + + +@pytest.mark.asyncio +async def test_session_from_other_running_loop_closed_threadsafe(): + """ + Regression test for #24230: a session that belongs to a loop still running + in another thread must be closed on its own loop (thread-safe), not driven + from the current loop. + """ + import threading + import time + + ready = threading.Event() + holder: dict = {} + + def worker() -> None: + loop = asyncio.new_event_loop() + holder["loop"] = loop + + async def make() -> None: + holder["session"] = aiohttp.ClientSession() + + loop.run_until_complete(make()) + ready.set() + loop.run_forever() + loop.close() + + thread = threading.Thread(target=worker, daemon=True) + thread.start() + assert ready.wait(5), "worker loop failed to start" + + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + transport.client = holder["session"] + + new_session = transport._get_valid_client_session() + + try: + deadline = time.monotonic() + 5 + while not holder["session"].closed and time.monotonic() < deadline: + await asyncio.sleep(0.01) + assert holder["session"].closed, "foreign-loop session was never closed" + finally: + holder["loop"].call_soon_threadsafe(holder["loop"].stop) + thread.join(5) + await new_session.close() + + +def test_threadsafe_close_done_callback_tolerates_cancelled_future(): + """ + Regression test for #24230 (review finding): when the foreign loop stops + before the handed-off close coroutine runs, asyncio cancels the + concurrent.futures.Future. The done-callback must return quietly instead + of letting future.exception() raise CancelledError (a BaseException that + escapes _invoke_callbacks and crashes the foreign loop's thread). + """ + future: "concurrent.futures.Future[None]" = concurrent.futures.Future() + future.cancel() + + LiteLLMAiohttpTransport._on_threadsafe_close_done(future) + + +@pytest.mark.asyncio +async def test_session_closed_retry_does_not_close_concurrent_replacement(): + """ + Regression test for #24230 (review finding): when the "Session is closed" + retry fires, the handler must dispose the session that actually faulted, + not self.client - a concurrent task may already have replaced self.client + with a healthy session, which must stay open. + """ + from unittest.mock import patch + + faulted_session = aiohttp.ClientSession() + healthy_replacement = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + transport.client = faulted_session + + calls = {"n": 0} + + async def fake_make_request(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + # simulate a concurrent task replacing the shared session between + # the failed await and the exception handler + transport.client = healthy_replacement + raise RuntimeError("Session is closed") + raise StopAsyncIteration("stop after retry dispatch") + + with patch.object(transport, "_make_aiohttp_request", side_effect=fake_make_request): + with pytest.raises(Exception): + await transport.handle_async_request(httpx.Request("GET", "http://example.com")) + + try: + assert not healthy_replacement.closed, "concurrent replacement session must not be closed by the retry handler" + for _ in range(3): + await asyncio.sleep(0) + assert faulted_session.closed, "the faulted session must be disposed" + finally: + await faulted_session.close() + await healthy_replacement.close() + new_session = transport.client + if isinstance(new_session, aiohttp.ClientSession): + await new_session.close() + + +@pytest.mark.asyncio +async def test_stopped_loop_session_disposed_synchronously_on_recycle(): + """ + Regression test for #24230 (review finding): a session whose loop is + stopped but not yet closed cannot safely run an async close on another + loop, and nothing will ever process a close handed to the stopped loop. + Recycling must dispose it synchronously, like the closed-loop case. + """ + import threading + + result: dict = {} + + def build() -> None: + loop = asyncio.new_event_loop() + + async def make() -> None: + result["session"] = aiohttp.ClientSession() + + loop.run_until_complete(make()) + result["loop"] = loop # stopped, deliberately NOT closed + + thread = threading.Thread(target=build) + thread.start() + thread.join(5) + + old_session = result["session"] + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + transport.client = old_session + + new_session = transport._get_valid_client_session() + + try: + assert new_session is not old_session + assert old_session.closed, "session from a stopped (not yet closed) loop must be disposed synchronously" + finally: + await new_session.close() + result["loop"].close()