fix(aiohttp_transport): keep the owned session open while a response stream is still being read (LIT-10237) (#48920)

Backport of #48920 to lts/1.108.x.
Cherry-picked from 4a70d80aef (main).

litellm/llms/custom_httpx/aiohttp_transport.py

@@ -6,6 +6,7 @@ import ssl
 import sys
 import typing
 import urllib.request
+import weakref
 from collections.abc import Callable, Generator
 from typing import ClassVar, Final
 
@@ -109,6 +110,24 @@ def map_aiohttp_exceptions() -> Generator[None, None, None]:
         raise mapped_exc(message) from exc
 
 
+def _release_transport_anchor(_transport: object) -> None:
+    """Release a transport anchored to a response stream. See ``_anchor_transport_to``."""
+
+
+def _anchor_transport_to(stream: "AiohttpResponseStream", transport: object) -> None:
+    """Keep the transport alive for as long as the stream can still read.
+
+    The transport finalizer closes the session it owns, and nothing on the path from a
+    live response back to that session references the transport, so a client dropped
+    mid-stream would close the session under the body. ``weakref.finalize`` holds the
+    transport in its own registry rather than on the stream, which keeps the transport
+    out of the response's reference cycle (an exception carrying the response, for one)
+    so it is finalized by refcount once the anchor drops and can still schedule an async
+    close, instead of being collected in the same cyclic pass as the session it closes.
+    """
+    weakref.finalize(stream, _release_transport_anchor, transport)
+
+
 class AiohttpResponseStream(httpx.AsyncByteStream):
     CHUNK_SIZE = 1024 * 16
 
@@ -463,10 +482,12 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
                 # Re-raise if it's a different RuntimeError
                 raise
 
+        stream: Final = AiohttpResponseStream(response)
+        _anchor_transport_to(stream, self)
         return httpx.Response(
             status_code=response.status,
             headers=response.headers,
-            stream=AiohttpResponseStream(response),
+            stream=stream,
             request=request,
         )
 

tests/unit/llms/custom_httpx/test_aiohttp_transport.py

@@ -5,6 +5,7 @@ import socket
 import sys
 import struct
 import threading
+import weakref
 from typing import Final
 
 import aiohttp
@@ -1383,3 +1384,57 @@ def test_a_transport_collected_while_its_loop_runs_elsewhere_closes_its_session_
         loop.call_soon_threadsafe(loop.stop)
         thread.join(timeout=10)
         loop.close()
+
+
+@pytest.mark.asyncio
+async def test_a_stream_keeps_reading_after_the_client_that_opened_it_is_collected():
+    from aiohttp import web
+
+    release: Final = asyncio.Event()
+    body: Final = tuple(f"chunk{i}\n".encode() for i in range(4))
+
+    async def handler(request: web.Request) -> web.StreamResponse:
+        response: Final = web.StreamResponse()
+        await response.prepare(request)
+        await response.write(body[0])
+        await release.wait()
+        for chunk in body[1:]:
+            await response.write(chunk)
+        await response.write_eof()
+        return response
+
+    app: Final = web.Application()
+    app.router.add_get("/stream", handler)
+    runner: Final = web.AppRunner(app)
+    await runner.setup()
+    sock: Final = socket.socket()
+    sock.bind(("127.0.0.1", 0))
+    port: Final[int] = sock.getsockname()[1]
+    await web.SockSite(runner, sock).start()
+
+    transport: Final = LiteLLMAiohttpTransport(client=aiohttp.ClientSession)
+    client: Final = httpx.AsyncClient(transport=transport)
+    transport_ref: Final = weakref.ref(transport)
+    try:
+        response: Final = await client.send(client.build_request("GET", f"http://127.0.0.1:{port}/stream"), stream=True)
+        session: Final = transport.client
+        assert isinstance(session, aiohttp.ClientSession)
+        chunks: Final = response.aiter_bytes()
+        first: Final = await chunks.__anext__()
+
+        del client, transport
+        gc.collect()
+        await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
+        release.set()
+
+        rest: Final = [chunk async for chunk in chunks]
+        assert first + b"".join(rest) == b"".join(body), (first, rest)
+        assert not session.closed, "the session was closed under a stream that was still being read"
+
+        del response, chunks
+        gc.collect()
+        await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
+        assert transport_ref() is None, "the finished stream kept the transport alive"
+        assert session.closed, "the owned session stayed open after its last stream finished"
+    finally:
+        await runner.cleanup()