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.
This commit is contained in:
Maximilian Roos 2026-09-07 22:07:12 -07:00
parent aa759ef750
commit 5972227b6d
No known key found for this signature in database
GPG key ID: 4744C397627707AE
3 changed files with 210 additions and 135 deletions

View file

@ -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:

View file

@ -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",
}
)

View file

@ -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