From 5972227b6d2bbe01d43e373ec5931c7cc6821658 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Mon, 7 Sep 2026 22:07:12 -0700 Subject: [PATCH] fix(http_handler): keep a handler alive while a response it issued is still reading _handler_may_close_client withholds the finalizer's close from a client that someone else references. A streaming response is not one of those referrers: it holds the connection it reads from, never the client, so the refcount says "sole referrer" for exactly the client that is busiest, and __del__ tears the pool down mid-body. Both handlers, both transports, and it is the shape #24929 reports: a provider's streaming call returns the response and drops the handler, and get_async_httpx_client lets a cached handler go after an hour. A streaming send now anchors the handler to the response it returns, so the handler is finalized once the caller is done with the body and the ordinary close applies with nothing in flight. The anchor is a weakref.finalize, which holds the handler in its own registry rather than on the response. That keeps the handler out of the response's reference cycle, so it is finalized by refcount and can still schedule an async close, rather than inside a cyclic collection that reaps its aiohttp session in the same pass; and a handler serving several streams is released only once all of them are done, because each anchor holds it separately. Only a streaming send anchors. A non-streaming response has been read in full by the time the method returns, so pinning the handler to it would delay every client close behind whatever the caller does with the response. The alternative was to ask the connection pool whether a request was in flight, reusing EvictedClientCloser's check. It reads client._transport, so it answers "idle" for any client with a proxy configured, where httpx routes through client._mounts; and having found the client busy it can only poll, which never terminates for a response the caller abandons unread, since httpx leaves that connection checked out. The response's own lifetime is the condition both were approximating. --- litellm/llms/custom_httpx/http_handler.py | 40 +++ tests/local_testing/conftest.py | 3 + .../test_handler_gc_does_not_close_client.py | 302 ++++++++++-------- 3 files changed, 210 insertions(+), 135 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index e1f0fc9e7d3..612d5997ec5 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -7,6 +7,7 @@ import ssl import sys import threading import time +import weakref from collections.abc import AsyncIterable, Callable, Iterable, Mapping from http.cookiejar import CookieJar, DefaultCookiePolicy from types import MappingProxyType @@ -178,6 +179,33 @@ def _handler_may_close_client(client_refcount: int, owns_client: bool) -> bool: return owns_client and client_refcount <= _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER +def _drop_streaming_anchor(_handler: object) -> None: + """Release a handler anchored to a streaming response. See ``_anchor_handler_to``. + + The work is the reference held until this point, so there is nothing to do here. + """ + + +def _anchor_handler_to(response: httpx.Response, handler: object) -> None: + """Keep the handler alive for as long as a streaming response can still read. + + A body still arriving reads through the handler's connection pool, and closing + the client tears that pool down. The refcount ``_handler_may_close_client`` + reads cannot see that body: the reference graph runs response -> stream -> + connection and stops there, so a client carrying one looks exactly like an + unreferenced client, and the finalizer closes it mid-body. + + ``weakref.finalize`` holds the handler in its own registry rather than on the + response, which matters twice. The handler stays out of the response's + reference cycle, so it is finalized by refcount once the anchor drops and can + still schedule an async close, instead of being finalized inside a cyclic + collection that reaps its aiohttp session in the same pass. And a handler + serving several streams collects only once every one of them is done, because + each anchor holds it separately. + """ + weakref.finalize(response, _drop_streaming_anchor, handler) + + def blocked_cookie_jar() -> CookieJar: """A jar that stores no response cookie and sends none, for httpx clients. @@ -704,6 +732,8 @@ class AsyncHTTPHandler: content=request_content, ) response: Final = await self.client.send(req, stream=stream) + if stream: + _anchor_handler_to(response, self) response.raise_for_status() return response except (httpx.RemoteProtocolError, httpx.ConnectError): @@ -898,6 +928,8 @@ class AsyncHTTPHandler: content=request_content, ) response: Final = await self.client.send(req, stream=stream) + if stream: + _anchor_handler_to(response, self) response.raise_for_status() return response except (httpx.RemoteProtocolError, httpx.ConnectError): @@ -1362,6 +1394,8 @@ class HTTPHandler: content=request_content, ) response: Final = self.client.send(req, stream=stream) + if stream: + _anchor_handler_to(response, self) response.raise_for_status() return response except httpx.TimeoutException: @@ -1412,6 +1446,8 @@ class HTTPHandler: content=request_content, ) response: Final = self.client.send(req, stream=stream) + if stream: + _anchor_handler_to(response, self) response.raise_for_status() return response except httpx.TimeoutException: @@ -1462,6 +1498,8 @@ class HTTPHandler: content=request_content, ) response: Final = self.client.send(req, stream=stream) + if stream: + _anchor_handler_to(response, self) return response except httpx.TimeoutException: raise litellm.Timeout( @@ -1511,6 +1549,8 @@ class HTTPHandler: content=request_content, ) response: Final = self.client.send(req, stream=stream) + if stream: + _anchor_handler_to(response, self) response.raise_for_status() return response except httpx.TimeoutException: diff --git a/tests/local_testing/conftest.py b/tests/local_testing/conftest.py index 5535a62bb81..228457f4d55 100644 --- a/tests/local_testing/conftest.py +++ b/tests/local_testing/conftest.py @@ -75,6 +75,9 @@ _VCR_INCOMPATIBLE_FILES = frozenset( "test_router_caching.py", # Hits the local fake OpenAI endpoint on 127.0.0.1; nothing to record. "test_fake_openai_endpoint.py", + # Needs the real connection pool a collected handler tears down; vcrpy + # patches the transport that pool lives in. + "test_handler_gc_does_not_close_client.py", } ) diff --git a/tests/local_testing/test_handler_gc_does_not_close_client.py b/tests/local_testing/test_handler_gc_does_not_close_client.py index 3874405a261..1a6ab1b1827 100644 --- a/tests/local_testing/test_handler_gc_does_not_close_client.py +++ b/tests/local_testing/test_handler_gc_does_not_close_client.py @@ -1,36 +1,39 @@ """ -Garbage-collecting an HTTP handler must not close the httpx client it holds. +Collecting an HTTP handler must not abort a response that is still on the wire. -``HTTPHandler`` and ``AsyncHTTPHandler`` used to close their client from -``__del__``. Closing an httpx client tears down the connection pool that every -in-flight response is streaming through, and it permanently invalidates the -client for future requests -- including for callers who only ever borrowed -``handler.client``. Since nothing in a response's reference graph points back -at the handler, and since litellm caches handlers behind a one-hour TTL, the -handler routinely became collectable while its client was still in use. +``HTTPHandler`` and ``AsyncHTTPHandler`` close their client from ``__del__``. +Closing a client tears down the connection pool, which aborts every response +still streaming through it. ``_handler_may_close_client`` already withholds the +close from a client someone else holds, but a streaming response holds the +connection it is reading from and never the client, so the refcount it reads +says "sole referrer" for exactly the client that is busiest. The handler is +routinely collectable at that moment: a provider's streaming call returns the +response and drops the handler, and ``get_async_httpx_client`` caches handlers +behind a one-hour TTL and then lets them go. -``LLMClientCache`` documents the invariant this broke: evicted clients "may -still be in use by in-flight requests", so they are left to normal garbage -collection rather than closed eagerly. A finalizer that closes on collection -defeats exactly that. +The fix anchors the handler to the streaming response, so these tests turn on +*when* the handler is collected rather than on whether it is: pinned while the +body can still arrive, released once the caller is done with the response. -Each test below is one shape the finalizers broke; all of them fail if either -``__del__`` comes back. Async cases run on both transports, because litellm -defaults to aiohttp and only uses httpcore when aiohttp is disabled. +Nothing here re-tests the shapes ``_handler_may_close_client`` covers -- a +borrowed ``handler.client``, a caller-supplied client, an evicted-but-held +client. Those are pinned in ``tests/test_litellm/llms/custom_httpx/ +test_http_handler.py``. What is uncovered there is the in-flight response, so no +test here may keep the client in a local: that inflates the very refcount under +test, and the test then passes on a broken handler. They hold weak references +instead, which the refcount does not count. These live here rather than under ``tests/test_litellm/`` because they need a real connection pool: a mocked transport goes on yielding chunks after its client is closed, so the very teardown under test is what a mock cannot -reproduce. The server is a hermetic, credential-free ``ThreadingHTTPServer`` -on an ephemeral loopback port, and needs no network access beyond it. +reproduce. The server is a hermetic, credential-free ``ThreadingHTTPServer`` on +an ephemeral loopback port, and needs no network access beyond it. Related: https://github.com/BerriAI/litellm/issues/24929 """ import asyncio import gc -import os -import sys import threading import time import weakref @@ -39,8 +42,6 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import httpx import pytest -sys.path.insert(0, os.path.abspath("../..")) - import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import ( @@ -54,15 +55,12 @@ FRAME_COUNT = 6 # Generous: the server emits all frames in ~0.3s. A client whose pool was torn # down mid-stream can stall silently instead of raising, so reads are bounded. READ_TIMEOUT_SECONDS = 15.0 +RELEASE_TIMEOUT_SECONDS = 3.0 BOTH_TRANSPORTS = pytest.mark.parametrize("disable_aiohttp_transport", [False, True], ids=["aiohttp", "httpcore"]) -# Every test here rests on the handler actually being collected at the ``del``. -# If something ever pins it, the test would pass while guarding nothing, so each -# one checks the premise. The check comes last: a reintroduced finalizer also -# fails it, by resurrecting the handler into the task it creates for ``close()``, -# and the transport error is the more useful thing to see first. -HANDLER_NOT_COLLECTED = "handler was not collected; this test no longer exercises the finalizer path" +STILL_PINNED = "the handler was released while its response could still read" +NOT_RELEASED = "the handler outlived the response that was holding it" class _ChunkedSSEServer: @@ -129,93 +127,44 @@ async def _read_frames(response: httpx.Response) -> int: return b"".join(chunks).count(b"data: frame-") +async def _wait_until(is_done, failure: str) -> None: + deadline = time.monotonic() + RELEASE_TIMEOUT_SECONDS + while time.monotonic() < deadline: + if is_done(): + return + await asyncio.sleep(0.05) + pytest.fail(failure) + + @pytest.mark.asyncio @BOTH_TRANSPORTS async def test_async_stream_survives_handler_collection(monkeypatch, disable_aiohttp_transport): - """A response being streamed keeps working after its handler is collected.""" + """A response still streaming keeps working after its handler goes out of scope. + + The caller holds the response and nothing else, which is what a provider's + streaming path is left with once ``post(..., stream=True)`` has returned. + """ _select_transport(monkeypatch, disable_aiohttp_transport) with _ChunkedSSEServer() as server: handler = AsyncHTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) - client = handler.client - try: - response = await client.send(client.build_request("GET", server.url), stream=True) + response = await handler.post(server.url, stream=True) - # The handler loses its last reference while the body is still streaming. - ref = weakref.ref(handler) - del handler - gc.collect() + ref = weakref.ref(handler) + del handler + gc.collect() + await asyncio.sleep(0) # let any close the finalizer scheduled run - frames = await asyncio.wait_for(_read_frames(response), timeout=READ_TIMEOUT_SECONDS) - assert frames == FRAME_COUNT - assert client.is_closed is False - assert ref() is None, HANDLER_NOT_COLLECTED - finally: - await client.aclose() + assert ref() is not None, STILL_PINNED + assert await asyncio.wait_for(_read_frames(response), timeout=READ_TIMEOUT_SECONDS) == FRAME_COUNT - -@pytest.mark.asyncio -@BOTH_TRANSPORTS -async def test_borrowed_async_client_outlives_its_handler(monkeypatch, disable_aiohttp_transport): - """A caller that keeps only ``handler.client`` can still send requests once the handler is gone. - - This is the shape at litellm/a2a_protocol/main.py (``httpx_client = - _async_handler.client``, handed to the a2a SDK) and at - litellm/proxy/pass_through_endpoints/pass_through_endpoints.py (``async_client - = async_client_obj.client``). Both take the handler from - ``get_async_httpx_client``, so the cache pins it for - ``_DEFAULT_TTL_FOR_HTTPX_CLIENTS`` (one hour) and then lets it go on eviction, - at which point it is collected while the borrowed client is still serving a - longer-lived consumer: ``create_a2a_client`` hands its client to the a2a SDK - and documents it as "create client once, reuse for multiple requests". - """ - _select_transport(monkeypatch, disable_aiohttp_transport) - - with _ChunkedSSEServer(frame_count=1, frame_delay=0.0) as server: - handler = AsyncHTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) - client = handler.client - try: - ref = weakref.ref(handler) - del handler - gc.collect() - # A finalizer would close the client from a task, so let the loop turn. - await asyncio.sleep(0.05) - - assert client.is_closed is False - response = await client.get(server.url) - assert response.status_code == 200 - assert ref() is None, HANDLER_NOT_COLLECTED - finally: - await client.aclose() - - -def test_sync_handler_collection_does_not_close_a_caller_owned_client(monkeypatch): - """A throwaway handler wrapped around someone else's client must not close it. - - litellm/llms/azure/azure.py builds ``HTTPHandler(client=litellm.client_session)`` - for a single image generation and drops it. With a finalizer, that one call - left the user's shared session closed for the rest of the process. - """ - monkeypatch.setattr(litellm, "force_ipv4", False) - - with _ChunkedSSEServer(frame_count=1, frame_delay=0.0) as server: - caller_client = httpx.Client(timeout=httpx.Timeout(10.0, connect=5.0)) - monkeypatch.setattr(litellm, "client_session", caller_client) - try: - handler = HTTPHandler(client=litellm.client_session) - ref = weakref.ref(handler) - del handler - gc.collect() - - assert caller_client.is_closed is False - assert caller_client.get(server.url).status_code == 200 - assert ref() is None, HANDLER_NOT_COLLECTED - finally: - caller_client.close() + del response + gc.collect() + assert ref() is None, NOT_RELEASED def test_sync_stream_survives_handler_collection(monkeypatch): - """A sync response being streamed keeps working after its handler is collected. + """The sync handler closes inline from its finalizer, so a stream must hold it off. litellm/main.py builds a sync handler only for non-streaming calls, commented "Keep this here, otherwise, the httpx.client closes and streaming is @@ -225,26 +174,115 @@ def test_sync_stream_survives_handler_collection(monkeypatch): with _ChunkedSSEServer() as server: handler = HTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) - client = handler.client - try: - response = client.send(client.build_request("GET", server.url), stream=True) + response = handler.post(server.url, stream=True) - # The handler loses its last reference while the body is still streaming. - ref = weakref.ref(handler) - del handler + ref = weakref.ref(handler) + del handler + gc.collect() + assert ref() is not None, STILL_PINNED + + # Joined before counting, as in ``_read_frames``. + chunks = [] + for chunk in response.iter_bytes(): + chunks.append(chunk) gc.collect() + assert b"".join(chunks).count(b"data: frame-") == FRAME_COUNT - # Joined before counting, as in ``_read_frames``. - chunks = [] - for chunk in response.iter_bytes(): - chunks.append(chunk) - gc.collect() + del response + gc.collect() + assert ref() is None, NOT_RELEASED - assert b"".join(chunks).count(b"data: frame-") == FRAME_COUNT - assert client.is_closed is False - assert ref() is None, HANDLER_NOT_COLLECTED - finally: - client.close() + +@pytest.mark.asyncio +@BOTH_TRANSPORTS +async def test_an_abandoned_stream_still_releases_its_handler(monkeypatch, disable_aiohttp_transport): + """A caller that drops a stream unread must not pin the handler for good. + + Tying the handler to the response's own lifetime is what bounds this. No + deadline, and no poll of the connection's state, can tell an abandoned body + from one the upstream is merely slow to finish: httpx leaves the connection + checked out until the response is read or closed, and a legitimate stream is + bounded only by how long the upstream keeps sending. + """ + _select_transport(monkeypatch, disable_aiohttp_transport) + + with _ChunkedSSEServer() as server: + handler = AsyncHTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) + client_ref = weakref.ref(handler.client) + response = await handler.post(server.url, stream=True) + + ref = weakref.ref(handler) + del handler, response + gc.collect() + + assert ref() is None, NOT_RELEASED + await _wait_until( + lambda: client_ref() is None or client_ref().is_closed, + "the client outlived the abandoned stream without being closed", + ) + + +@pytest.mark.asyncio +@BOTH_TRANSPORTS +async def test_the_pool_is_released_once_the_stream_it_carried_ends(monkeypatch, disable_aiohttp_transport): + """Holding the finalizer off must defer the close, not drop it. + + Otherwise a collected handler leaks its pool for every streaming request it + was carrying, and on aiohttp warns "Unclosed client session" when the + collector eventually takes it. The pool and the session are children of the + client, so keeping one here does not inflate the refcount the finalizer + reads, the way keeping the client would. + """ + _select_transport(monkeypatch, disable_aiohttp_transport) + + with _ChunkedSSEServer() as server: + handler = AsyncHTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) + transport = handler.client._transport + if disable_aiohttp_transport: + pool = transport._pool + + def is_released() -> bool: + return pool.connections == [] + else: + session = transport._get_valid_client_session() + + def is_released() -> bool: + return session.closed + + response = await handler.post(server.url, stream=True) + + del handler, transport + gc.collect() + assert not is_released(), "the pool was torn down while it was still carrying a body" + + assert await asyncio.wait_for(_read_frames(response), timeout=READ_TIMEOUT_SECONDS) == FRAME_COUNT + del response + gc.collect() + + await _wait_until(is_released, "the pool outlived the stream it carried, unclosed") + + +@pytest.mark.asyncio +@BOTH_TRANSPORTS +async def test_a_non_streaming_response_does_not_pin_its_handler(monkeypatch, disable_aiohttp_transport): + """Only a body that can still arrive holds the handler. + + A non-streaming response has been read in full by the time ``post`` returns, + so pinning the handler to it would delay every client close behind whatever + the caller goes on to do with the response. + """ + _select_transport(monkeypatch, disable_aiohttp_transport) + + with _ChunkedSSEServer(frame_count=1, frame_delay=0.0) as server: + handler = AsyncHTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) + response = await handler.post(server.url) + assert response.status_code == 200 + + ref = weakref.ref(handler) + del handler + gc.collect() + + assert ref() is None, "a fully-read response pinned its handler" @pytest.mark.asyncio @@ -254,30 +292,24 @@ async def test_cached_handler_eviction_does_not_abort_an_in_flight_stream(monkey ``get_async_httpx_client`` caches handlers for an hour. When that TTL expires the cache drops the only reference to a handler whose client is - still streaming -- the production shape of #24929, and the case - ``LLMClientCache`` documents as "may still be in use by in-flight requests". + still streaming -- the production shape of #24929. """ _select_transport(monkeypatch, disable_aiohttp_transport) monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) with _ChunkedSSEServer() as server: handler = get_async_httpx_client(llm_provider=LlmProviders.OPENAI) - client = handler.client - try: - response = await client.send(client.build_request("GET", server.url), stream=True) + response = await handler.post(server.url, stream=True) - # An hour passes: the TTL expires and the cache lets the handler go. - ref = weakref.ref(handler) - litellm.in_memory_llm_clients_cache.flush_cache() - del handler - gc.collect() + # An hour passes: the TTL expires and the cache lets the handler go. + ref = weakref.ref(handler) + litellm.in_memory_llm_clients_cache.flush_cache() + del handler + gc.collect() - frames = await asyncio.wait_for(_read_frames(response), timeout=READ_TIMEOUT_SECONDS) - assert frames == FRAME_COUNT - assert client.is_closed is False + assert ref() is not None, STILL_PINNED + assert await asyncio.wait_for(_read_frames(response), timeout=READ_TIMEOUT_SECONDS) == FRAME_COUNT - # And the evicted client is still usable for the next request. - assert (await client.get(server.url)).status_code == 200 - assert ref() is None, HANDLER_NOT_COLLECTED - finally: - await client.aclose() + del response + gc.collect() + assert ref() is None, NOT_RELEASED