fix(http_handler): defer finalizer client close while requests are in flight

litellm/caching/evicted_client_closer.py

@@ -206,6 +206,36 @@ def schedule(self, client: object) -> None:
             )
         )
 
+    def defer_if_busy(self, client: object) -> bool:
+        """Queue a deferred close for a client with a request in flight; no-op when idle.
+
+        For a finalized handler's client the sole-referrer refcount check has already
+        proven nothing else holds the client object, but a request in flight references
+        only the pooled connection, so it is invisible to that check. A busy client is
+        queued and closed by ``reap`` once idle and out of grace; returns True when the
+        close was deferred (the caller must NOT close the client itself), False when the
+        client is idle and the caller may close it directly.
+        """
+        if _close_function(client) is None:
+            return False
+        if not _has_connection_in_flight(client):
+            return False
+        self.mark_owned(client)
+        self.schedule(client)
+        return True
+
+    def close_or_defer(self, client: object) -> None:
+        """Close an unreferenced client now if idle, else queue it for a deferred close.
+
+        Same in-flight rule as ``defer_if_busy``; an idle client is closed immediately,
+        preserving the reclamation the finalizer used to do.
+        """
+        if _close_function(client) is None:
+            return
+        if self.defer_if_busy(client):
+            return
+        self._close(client)
+
     def reap(self) -> None:
         """Close every queued client that is due, idle, and closable from here.
 

litellm/llms/custom_httpx/http_handler.py

@@ -1102,13 +1102,34 @@ def __del__(self) -> None:
                 # here is the cross-loop path the transport refuses.
                 self._dispose_wrapped_aiohttp_session()
                 return
-            task: Final = loop.create_task(self._client.aclose())
+            task: Final = loop.create_task(self._finalizer_close_client(self._client))
             cls: Final = type(self)
             cls._finalizer_close_tasks.add(task)
             task.add_done_callback(cls._on_finalizer_close_done)
         except Exception:
             pass
 
+    @staticmethod
+    async def _finalizer_close_client(client: httpx.AsyncClient) -> None:
+        """Close a finalized handler's client, unless a request is still in flight.
+
+        The sole-referrer refcount guard in ``__del__`` proves nothing else holds the
+        client *object*, but a request in flight references only the pooled connection,
+        so it is invisible to that check: closing here would tear the pool down under
+        every live SSE stream (a cache-evicted handler is finalized the moment the cache
+        drops it — one batch of mid-turn stream deaths per handler-cache TTL per
+        process). A busy client is handed to the evicted-client closer, which closes it
+        once it reports no connection in flight and a grace window has passed; an idle
+        one is closed here so the finalizer keeps its reclamation. The in-flight check
+        runs inside this task, not in ``__del__``, so no closer lock is taken in GC
+        context.
+        """
+        from litellm.caching.evicted_client_closer import default_evicted_client_closer
+
+        if default_evicted_client_closer.defer_if_busy(client):
+            return
+        await client.aclose()
+
     @staticmethod
     def _create_async_transport(
         ssl_context: ssl.SSLContext | None = None,

tests/test_litellm/caching/test_evicted_client_closer.py

@@ -409,3 +409,57 @@ def test_a_reap_looks_at_what_is_due_rather_than_at_the_whole_queue():
         f"{CountingDeadline.comparisons} deadline comparisons for {evictions} evictions; "
         "a reap is walking the whole queue"
     )
+
+
+@pytest.mark.asyncio
+async def test_close_or_defer_closes_an_idle_client_immediately():
+    clock = FakeClock()
+    closer = make_closer(clock)
+    client = AsyncClient()
+
+    closer.close_or_defer(client)
+    await asyncio.sleep(0.05)
+
+    assert client.closed is True
+    assert closer.pending_count == 0
+
+
+@pytest.mark.asyncio
+async def test_close_or_defer_defers_a_client_with_a_request_on_the_wire():
+    server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0)
+    port = server.sockets[0].getsockname()[1]
+    clock = FakeClock()
+    closer = make_closer(clock)
+    client = httpx.AsyncClient()
+
+    async with asyncio.timeout(30):
+        async with client.stream("GET", f"http://127.0.0.1:{port}/") as response:
+            body_iter = response.aiter_raw()
+            await body_iter.__anext__()
+
+            closer.close_or_defer(client)
+
+            assert not client.is_closed
+            assert closer.pending_count == 1
+
+            remainder = b"".join([chunk async for chunk in body_iter])
+            assert b"hello" in remainder
+
+        clock.advance(61.0)
+        closer.reap()
+        await asyncio.sleep(0.05)
+
+    assert client.is_closed
+    assert closer.pending_count == 0
+    # No wait_closed(): on Python >= 3.12.1 it waits for every client
+    # transport, and a pooled keepalive connection would park it forever.
+    server.close()
+
+
+def test_close_or_defer_ignores_a_value_without_a_close_function():
+    clock = FakeClock()
+    closer = make_closer(clock)
+
+    closer.close_or_defer(object())
+
+    assert closer.pending_count == 0

tests/test_litellm/llms/custom_httpx/test_http_handler.py

@@ -1675,3 +1675,50 @@ async def aclose(self):
     finally:
         await handler.close()
     assert closed.is_set()
+
+
+async def _slow_chunked_upstream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+    """Serves a chunked body in two installments, so a stream is on the wire while the handler dies."""
+    await reader.read(4096)
+    writer.write(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n")
+    writer.write(b"5\r\nfirst\r\n")
+    await writer.drain()
+    await asyncio.sleep(0.4)
+    writer.write(b"4\r\nlast\r\n0\r\n\r\n")
+    await writer.drain()
+
+
+@pytest.mark.asyncio
+async def test_collected_handler_never_kills_a_stream_in_flight(monkeypatch):
+    """
+    Regression: a cache-evicted (hence collected) handler's finalizer used to close the
+    owned client while its pool still served live SSE streams, killing every one of them
+    mid-turn once per handler-cache TTL. The finalizer must defer to the evicted-client
+    closer while a connection is in flight, so the stream reads to completion.
+    """
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setattr(litellm, "force_ipv4", False)
+
+    server = await asyncio.start_server(_slow_chunked_upstream, "127.0.0.1", 0)
+    port = server.sockets[0].getsockname()[1]
+
+    handler = AsyncHTTPHandler()
+    async with asyncio.timeout(30):
+        request = handler.client.build_request("GET", f"http://127.0.0.1:{port}/")
+        response = await handler.client.send(request, stream=True)
+        body_iter = response.aiter_raw()
+        first = await body_iter.__anext__()
+        assert b"first" in first
+
+        del handler
+        gc.collect()
+        await asyncio.sleep(0.1)
+
+        remainder = b"".join([chunk async for chunk in body_iter])
+        assert b"last" in remainder
+
+        await response.aclose()
+    # No wait_closed(): the surviving client's pooled keepalive connection is
+    # the point of this test, and on Python >= 3.12.1 wait_closed() waits for
+    # every client transport, parking the suite forever.
+    server.close()