fix(aiohttp_transport): close an owned aiohttp session when its transport is collected (#48650)

Backport of #48650 to lts/1.108.x.
Cherry-picked from d6edd37810 (main).

litellm/llms/custom_httpx/aiohttp_transport.py

@@ -289,6 +289,13 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
         cls._background_close_tasks.add(task)
         task.add_done_callback(cls._on_close_task_done)
 
+    def __del__(self) -> None:
+        session: Final = getattr(self, "client", None)
+        if not getattr(self, "_owns_session", False) or not isinstance(session, ClientSession):
+            return
+        with contextlib.suppress(Exception):
+            self._close_recycled_session(session)
+
     def get_valid_client_session(self) -> ClientSession:
         """
         Helper to get a valid ClientSession for the current event loop.

tests/integration/sdk/test_aiohttp_session_rebuild_wire.py

@@ -11,11 +11,30 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from typing import Final
 
 import pytest
-from pydantic import JsonValue, TypeAdapter
+from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
 
 CONFIGURED_KEEPALIVE_SECONDS: Final = 1
 IDLE_SECONDS: Final = 2
+EVICTED_CLIENTS: Final = 30
 RESPONSES: Final = TypeAdapter(list[dict[str, JsonValue]])
+CHAT_COMPLETION: Final = {
+    "id": "chatcmpl-scripted",
+    "object": "chat.completion",
+    "created": 1700000000,
+    "model": "gpt-5.4-nano",
+    "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
+    "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+}
+
+
+class EvictionReport(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    cache_size: int
+    cached_clients: int
+    established_connections: int
+    unclosed_messages: tuple[str, ...]
+
 
 REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent(
     """
@@ -38,6 +57,47 @@ REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent(
 )
 
 
+EVICTED_CLIENTS_EXCHANGE: Final = textwrap.dedent(
+    """
+    import asyncio, gc, json, sys
+    import psutil
+    import litellm
+    from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport
+
+    unclosed = []
+
+    def established(port: int) -> int:
+        connections = psutil.Process().net_connections("tcp")
+        return sum(1 for c in connections if c.status == "ESTABLISHED" and c.raddr and c.raddr.port == port)
+
+    async def one_call(base_url: str, index: int) -> None:
+        await litellm.acompletion(
+            model="openai/gpt-5.4-nano",
+            api_key=f"sk-distinct-{index}",
+            api_base=base_url,
+            messages=[{"role": "user", "content": "hi"}],
+        )
+
+    async def main(base_url: str, port: int, evicted: int) -> None:
+        asyncio.get_running_loop().set_exception_handler(lambda loop, ctx: unclosed.append(ctx["message"]))
+        cache = litellm.in_memory_llm_clients_cache
+        for index in range(cache.max_size_in_memory + evicted):
+            await one_call(base_url, index)
+        gc.collect()
+        await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
+        await asyncio.sleep(0)
+        print(json.dumps({
+            "cache_size": cache.max_size_in_memory,
+            "cached_clients": len(cache.cache_dict),
+            "established_connections": established(port),
+            "unclosed_messages": unclosed,
+        }))
+
+    asyncio.run(main(sys.argv[1], int(sys.argv[2]), int(sys.argv[3])))
+    """
+)
+
+
 class _ConnectionCountingPeer(ThreadingHTTPServer):
     daemon_threads = True
 
@@ -62,7 +122,10 @@ class _ConnectionHandler(BaseHTTPRequestHandler):
 
     def do_POST(self) -> None:
         self.rfile.read(int(self.headers["Content-Length"]))
-        body: Final = json.dumps({"connection": self.connection_number}).encode()
+        payload: Final = (
+            CHAT_COMPLETION if self.path.endswith("/chat/completions") else {"connection": self.connection_number}
+        )
+        body: Final = json.dumps(payload).encode()
         self.send_response(200)
         self.send_header("Content-Type", "application/json")
         self.send_header("Content-Length", str(len(body)))
@@ -107,3 +170,32 @@ def test_rebuilt_shared_session_drops_idle_connection_after_configured_keepalive
 ) -> None:
     observed: Final = _rebuilt_session_exchange(connection_counting_peer)
     assert observed == [{"connection": 1}, {"connection": 2}], observed
+
+
+def _evicted_clients_exchange(base_url: str) -> EvictionReport:
+    completed: Final = subprocess.run(
+        [
+            sys.executable,
+            "-P",
+            "-c",
+            EVICTED_CLIENTS_EXCHANGE,
+            base_url,
+            base_url.rsplit(":", 1)[1],
+            str(EVICTED_CLIENTS),
+        ],
+        capture_output=True,
+        text=True,
+        timeout=300,
+        check=False,
+    )
+    assert completed.returncode == 0, completed.stderr
+    return EvictionReport.model_validate_json(completed.stdout)
+
+
+def test_sdk_clients_evicted_from_a_full_client_cache_leave_no_open_session_or_connection(
+    connection_counting_peer: str,
+) -> None:
+    report: Final = _evicted_clients_exchange(connection_counting_peer)
+    assert report.cached_clients == report.cache_size, report
+    assert report.unclosed_messages == (), report
+    assert report.established_connections <= report.cache_size, report

tests/unit/caching/test_llm_caching_handler.py

@@ -11,11 +11,14 @@ See: https://github.com/BerriAI/litellm/pull/22247
 import asyncio, gc, httpx, importlib, litellm, os, threading, time, weakref
 import warnings
 
+import aiohttp
 import pytest
+from openai import AsyncOpenAI
 
 
 from litellm.caching.evicted_client_closer import EvictedClientCloser
 from litellm.caching.llm_caching_handler import LLMClientCache
+from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.llms.custom_httpx.http_handler import(
@@ -204,6 +207,36 @@ async def test_evicted_litellm_owned_client_is_closed_once_the_grace_window_elap
     assert client.closed is True
 
 
+@pytest.mark.asyncio
+async def test_an_evicted_owned_client_collected_inside_its_grace_window_still_closes_its_aiohttp_session():
+    """
+    The closer holds an evicted client weakly for the whole grace window. An SDK
+    process that exits, or a proxy that evicts and never calls the cache again,
+    collects the client first, so only its own finalizer can still close the
+    aiohttp session (LIT-9767).
+    """
+    clock = _FakeClock()
+    cache = LLMClientCache(
+        max_size_in_memory=2,
+        evicted_client_closer=EvictedClientCloser(grace_seconds=900.0, clock=clock),
+    )
+    session = aiohttp.ClientSession()
+    client = AsyncOpenAI(
+        api_key="sk-evicted",
+        http_client=httpx.AsyncClient(transport=LiteLLMAiohttpTransport(client=session)),
+    )
+    cache.set_cache("client-key", client, litellm_owned_client=True, ttl=600)
+
+    cache.ttl_dict = {key: 0 for key in cache.ttl_dict}
+    cache.expiration_heap = [(0, key) for _, key in cache.expiration_heap]
+    cache.evict_cache()
+    del client
+    gc.collect()
+    await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
+
+    assert session.closed, "the evicted client was collected inside the grace window and left its session open"
+
+
 @pytest.mark.asyncio
 async def test_evicted_caller_supplied_client_is_never_closed():
     """litellm does not own a client the caller passed in, so it must stay open."""

tests/unit/llms/custom_httpx/test_aiohttp_transport.py

@@ -1,8 +1,10 @@
 import aiohttp as aiohttp_aiohttp_handler, asyncio, importlib, litellm
 import concurrent.futures
+import gc
 import socket
 import sys
 import struct
+import threading
 from typing import Final
 
 import aiohttp
@@ -1310,3 +1312,74 @@ async def test_reset_after_request_is_not_replayed(monkeypatch: pytest.MonkeyPat
             assert len(received_requests) == 1
     finally:
         await handler.close()
+
+
+async def _new_session() -> aiohttp.ClientSession:
+    return aiohttp.ClientSession()
+
+
+async def _until_closed(session: aiohttp.ClientSession) -> None:
+    while not session.closed:
+        await asyncio.sleep(0)
+
+
+@pytest.mark.asyncio
+async def test_a_collected_transport_closes_the_session_it_owns():
+    """An SDK client evicted from the cache is collected long before the closer's grace window ends."""
+    session = aiohttp.ClientSession()
+    transport = LiteLLMAiohttpTransport(client=session)
+
+    del transport
+    gc.collect()
+    await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
+
+    assert session.closed, "the owned session outlived the transport that was collected before closing it"
+
+
+@pytest.mark.asyncio
+async def test_a_collected_transport_leaves_a_shared_session_open():
+    session = aiohttp.ClientSession()
+    try:
+        transport = LiteLLMAiohttpTransport(client=session, owns_session=False)
+
+        del transport
+        gc.collect()
+        await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
+
+        assert not session.closed, "a transport closed a session it does not own"
+    finally:
+        await session.close()
+
+
+def test_a_transport_collected_after_its_loop_closed_still_closes_its_session(caplog: pytest.LogCaptureFixture):
+    """asyncio.run() closes the loop the session lives on before the evicted client is collected."""
+    loop = asyncio.new_event_loop()
+    session = loop.run_until_complete(_new_session())
+    loop.close()
+    transport = LiteLLMAiohttpTransport(client=session)
+
+    del transport
+    gc.collect()
+
+    assert session.closed, "the owned session stayed open after its transport was collected on a closed loop"
+    del session
+    gc.collect()
+    assert "Unclosed client session" not in caplog.text, caplog.text
+
+
+def test_a_transport_collected_while_its_loop_runs_elsewhere_closes_its_session_on_that_loop():
+    loop = asyncio.new_event_loop()
+    thread = threading.Thread(target=loop.run_forever, daemon=True)
+    thread.start()
+    session = asyncio.run_coroutine_threadsafe(_new_session(), loop).result(timeout=10)
+    transport = LiteLLMAiohttpTransport(client=session)
+    try:
+        del transport
+        gc.collect()
+        asyncio.run_coroutine_threadsafe(_until_closed(session), loop).result(timeout=10)
+
+        assert session.closed, "the owned session stayed open after its transport was collected from another thread"
+    finally:
+        loop.call_soon_threadsafe(loop.stop)
+        thread.join(timeout=10)
+        loop.close()