fix(http): retry vector store search only on HTTP/2 GOAWAY and keep aiohttp pool and proxy behavior

Scope the one retry to a connection Google closes with an HTTP/2 GOAWAY and
log it. When HTTP/2 replaces the aiohttp transport, keep aiohttp's idle
keepalive and connection cap, and only then skip env proxies under
DISABLE_AIOHTTP_TRUST_ENV, so the global LITELLM_HTTP2 toggle keeps HTTPS_PROXY

litellm/llms/custom_httpx/http_handler.py

@@ -18,6 +18,7 @@ import certifi
 import httpx
 from aiohttp import ClientSession, DummyCookieJar, TCPConnector
 from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport
+from httpx._config import DEFAULT_LIMITS
 from httpx._types import CertTypes, RequestFiles
 from httpx._utils import get_environment_proxies
 
@@ -101,6 +102,13 @@ def _aiohttp_env_proxies_disabled() -> bool:
     return litellm.disable_aiohttp_trust_env or str_to_bool(os.getenv("DISABLE_AIOHTTP_TRUST_ENV", "False")) is True
 
 
+_AIOHTTP_EQUIVALENT_LIMITS: Final = httpx.Limits(
+    max_connections=AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT > 0 else None,
+    max_keepalive_connections=None,
+    keepalive_expiry=AIOHTTP_KEEPALIVE_TIMEOUT,
+)
+
+
 def _environment_proxy_mounts(
     build_proxy_transport: Callable[[str], _HttpxTransportT],
 ) -> Mapping[str, _HttpxTransportT | None]:
@@ -705,6 +713,8 @@ class AsyncHTTPHandler:
                 trust_env=False,
             )
         use_http2: Final = _resolve_http2(self.http2, ssl_verify)
+        replaces_aiohttp: Final = use_http2 and AsyncHTTPHandler.should_use_aiohttp_transport()
+        limits: Final = _AIOHTTP_EQUIVALENT_LIMITS if replaces_aiohttp else DEFAULT_LIMITS
         # Get unified SSL configuration
         ssl_config: Final = get_ssl_configuration(ssl_verify, http2=use_http2)
 
@@ -722,17 +732,21 @@ class AsyncHTTPHandler:
             shared_session=shared_session,
             http2=use_http2,
             cert=cert,
+            limits=limits,
         )
 
         # Get default headers (User-Agent, overridable via LITELLM_USER_AGENT)
         default_headers: Final = get_default_headers()
-        trust_env: Final = not (use_http2 and _aiohttp_env_proxies_disabled())
+        trust_env: Final = not (replaces_aiohttp and _aiohttp_env_proxies_disabled())
 
         return httpx.AsyncClient(
             transport=transport,
-            mounts=AsyncHTTPHandler.create_httpx_proxy_mounts(transport, verify=ssl_config, cert=cert, http2=use_http2)
+            mounts=AsyncHTTPHandler.create_httpx_proxy_mounts(
+                transport, verify=ssl_config, cert=cert, http2=use_http2, limits=limits
+            )
             if trust_env
             else None,
+            limits=limits,
             event_hooks=event_hooks,
             timeout=timeout,
             verify=ssl_config,
@@ -1233,6 +1247,7 @@ class AsyncHTTPHandler:
         shared_session: Optional["ClientSession"] = None,
         http2: bool = False,
         cert: CertTypes | None = None,
+        limits: httpx.Limits = DEFAULT_LIMITS,
     ) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | None:
         """
         - Creates a transport for httpx.AsyncClient
@@ -1265,6 +1280,7 @@ class AsyncHTTPHandler:
             http2=http2,
             verify=ssl_context if ssl_context is not None else ssl_verify is not False,
             cert=cert,
+            limits=limits,
         )
 
     _create_async_transport = create_async_transport
@@ -1416,6 +1432,7 @@ class AsyncHTTPHandler:
         http2: bool = False,
         verify: ssl.SSLContext | bool = True,
         cert: CertTypes | None = None,
+        limits: httpx.Limits = DEFAULT_LIMITS,
     ) -> AsyncHTTPTransport | None:
         """
         Creates an AsyncHTTPTransport
@@ -1425,7 +1442,11 @@ class AsyncHTTPHandler:
         """
         if litellm.force_ipv4:
             return AsyncHTTPTransport(
-                local_address=_IPV4_LOCAL_ADDRESS, verify=verify, cert=cert, http2=http2 or http2_enabled()
+                local_address=_IPV4_LOCAL_ADDRESS,
+                verify=verify,
+                cert=cert,
+                http2=http2 or http2_enabled(),
+                limits=limits,
             )
         else:
             return None
@@ -1436,12 +1457,13 @@ class AsyncHTTPHandler:
         verify: VerifyTypes,
         cert: CertTypes | None,
         http2: bool = False,
+        limits: httpx.Limits = DEFAULT_LIMITS,
     ) -> Mapping[str, AsyncHTTPTransport | None] | None:
         if not isinstance(transport, AsyncHTTPTransport):
             return None
         return _environment_proxy_mounts(
             lambda proxy_url: AsyncHTTPTransport(
-                proxy=proxy_url, verify=verify, cert=cert, http2=http2 or http2_enabled()
+                proxy=proxy_url, verify=verify, cert=cert, http2=http2 or http2_enabled(), limits=limits
             )
         )
 

litellm/llms/custom_httpx/llm_http_handler.py

@@ -25,6 +25,7 @@ from typing import (
 from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
 
 import httpx
+from h2.events import ConnectionTerminated
 from httpx import USE_CLIENT_DEFAULT
 from httpx._types import FileContent
 from openai.types.file_deleted import FileDeleted
@@ -340,17 +341,31 @@ class _SignedBodyKwargs(TypedDict, total=False):
     json: ReadOnly[dict[str, object]]
 
 
-async def _asend_idempotent(send: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response:
+_GOAWAY_RETRY_MESSAGE: Final = "Vector store search connection was closed by an HTTP/2 GOAWAY, retrying once"
+
+
+def _closed_by_http2_goaway(error: httpx.RemoteProtocolError) -> bool:
+    cause: Final = error.__cause__
+    return cause is not None and any(isinstance(arg, ConnectionTerminated) for arg in cause.args)
+
+
+async def _asend_with_goaway_retry(send: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response:
     try:
         return await send()
-    except httpx.RemoteProtocolError:
+    except httpx.RemoteProtocolError as error:
+        if not _closed_by_http2_goaway(error):
+            raise
+        verbose_logger.warning(_GOAWAY_RETRY_MESSAGE)
         return await send()
 
 
-def _send_idempotent(send: Callable[[], httpx.Response]) -> httpx.Response:
+def _send_with_goaway_retry(send: Callable[[], httpx.Response]) -> httpx.Response:
     try:
         return send()
-    except httpx.RemoteProtocolError:
+    except httpx.RemoteProtocolError as error:
+        if not _closed_by_http2_goaway(error):
+            raise
+        verbose_logger.warning(_GOAWAY_RETRY_MESSAGE)
         return send()
 
 
@@ -10405,7 +10420,7 @@ class BaseLLMHTTPHandler:
         request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body
 
         try:
-            response: Final = await _asend_idempotent(
+            response: Final = await _asend_with_goaway_retry(
                 partial(async_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout)
             )
         except httpx.TimeoutException:
@@ -10544,7 +10559,7 @@ class BaseLLMHTTPHandler:
         request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body
 
         try:
-            response: Final = _send_idempotent(
+            response: Final = _send_with_goaway_retry(
                 partial(sync_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout)
             )
         except httpx.TimeoutException:

tests/integration/sdk/test_http2_wire.py

@@ -4,18 +4,19 @@ import asyncio
 import contextlib
 import datetime
 import ipaddress
-import itertools
 import json
+import queue
 import socket
 import ssl
 import threading
 import time
-from collections.abc import Generator, Iterator
+from collections.abc import Callable, Generator, Iterator
 from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Final
 
+import httpx
 import pytest
 from cryptography import x509
 from cryptography.hazmat.primitives import hashes, serialization
@@ -24,6 +25,7 @@ from cryptography.x509.oid import NameOID
 from h2.config import H2Configuration
 from h2.connection import H2Connection
 from h2.events import StreamEnded
+from httpx._config import DEFAULT_LIMITS
 from hypercorn.asyncio import (
     serve,  # pyright: ignore[reportUnknownVariableType]  # hypercorn's serve signature passes through untyped worker hooks
 )
@@ -356,29 +358,50 @@ def test_vertex_search_leaves_a_caller_ssl_context_on_http1(
     assert _alpn_offered_by(async_context, http2_tls_peer) != "h2"
 
 
+def _route_https_through_a_dead_proxy(monkeypatch: pytest.MonkeyPatch) -> None:
+    for name in ("https_proxy", "HTTP_PROXY", "http_proxy", "ALL_PROXY", "all_proxy", "NO_PROXY", "no_proxy"):
+        monkeypatch.delenv(name, raising=False)
+    monkeypatch.delenv("DISABLE_AIOHTTP_TRANSPORT", raising=False)
+    monkeypatch.setenv("HTTPS_PROXY", f"http://127.0.0.1:{allocate_port()}")
+    monkeypatch.setenv("DISABLE_AIOHTTP_TRUST_ENV", "True")
+
+
+async def _async_peer_http_version(peer: TlsPeer, http2: bool) -> str:
+    from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+
+    handler: Final = AsyncHTTPHandler(ssl_verify=str(peer.cert_file), http2=http2)
+    try:
+        response: Final = await handler.client.post(f"{peer.url}/echo", json={"ping": "pong"})
+        return response.http_version
+    finally:
+        await handler.close()
+
+
 def test_http2_client_skips_env_proxies_when_aiohttp_trust_env_is_disabled(
     monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
 ) -> None:
     import litellm
-    from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 
     _set_http2(monkeypatch, False)
-    for name in ("https_proxy", "HTTP_PROXY", "http_proxy", "ALL_PROXY", "all_proxy", "NO_PROXY", "no_proxy"):
-        monkeypatch.delenv(name, raising=False)
-    monkeypatch.setenv("HTTPS_PROXY", f"http://127.0.0.1:{allocate_port()}")
-    monkeypatch.setenv("DISABLE_AIOHTTP_TRUST_ENV", "True")
-
-    async def _peer_http_version() -> str:
-        handler: Final = AsyncHTTPHandler(ssl_verify=str(http2_tls_peer.cert_file), http2=True)
-        try:
-            response: Final = await handler.client.post(f"{http2_tls_peer.url}/echo", json={"ping": "pong"})
-            return response.http_version
-        finally:
-            await handler.close()
+    _route_https_through_a_dead_proxy(monkeypatch)
 
     for force_ipv4 in (False, True):
         monkeypatch.setattr(litellm, "force_ipv4", force_ipv4)
-        assert asyncio.run(_peer_http_version()) == "HTTP/2"
+        assert asyncio.run(_async_peer_http_version(http2_tls_peer, http2=True)) == "HTTP/2"
+
+
+def test_global_http2_toggle_keeps_env_proxies_when_aiohttp_trust_env_is_disabled(
+    monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
+) -> None:
+    import litellm
+
+    _set_http2(monkeypatch, True)
+    _route_https_through_a_dead_proxy(monkeypatch)
+
+    for force_ipv4 in (False, True):
+        monkeypatch.setattr(litellm, "force_ipv4", force_ipv4)
+        with pytest.raises(httpx.ConnectError):
+            asyncio.run(_async_peer_http_version(http2_tls_peer, http2=False))
 
 
 def _answer_or_drain(connection: H2Connection, stream_id: int, drain: bool) -> None:
@@ -386,7 +409,7 @@ def _answer_or_drain(connection: H2Connection, stream_id: int, drain: bool) -> N
         connection.close_connection(last_stream_id=stream_id)
         return
     connection.send_headers(stream_id, [(":status", "200"), ("content-type", "application/json")])
-    connection.send_data(stream_id, json.dumps({"results": [{"id": "after-goaway"}]}).encode(), end_stream=True)
+    connection.send_data(stream_id, json.dumps({"results": [{"id": "answered"}]}).encode(), end_stream=True)
 
 
 def _reply_to_ended_streams(connection: H2Connection, tls: ssl.SSLSocket, data: bytes, drain: bool) -> bool:
@@ -407,28 +430,55 @@ def _serve_h2(context: ssl.SSLContext, raw: socket.socket, drain: bool) -> None:
                 return
 
 
+def _read_http1_request(tls: ssl.SSLSocket) -> bytes:
+    with tls.makefile("rb") as stream:
+        head: Final = tuple(iter(stream.readline, b"\r\n"))
+        length: Final = next(
+            (int(line.split(b":", 1)[1]) for line in head if line.lower().startswith(b"content-length:")), 0
+        )
+        return b"".join(head) + stream.read(length)
+
+
+def _hang_up_after_request(context: ssl.SSLContext, raw: socket.socket, connection_number: int) -> None:
+    with contextlib.suppress(OSError), raw, context.wrap_socket(raw, server_side=True) as tls:
+        _ = _read_http1_request(tls)
+
+
+def _drain_every_other_connection(context: ssl.SSLContext, raw: socket.socket, connection_number: int) -> None:
+    _serve_h2(context, raw, drain=connection_number % 2 == 1)
+
+
+def _answer_every_connection(context: ssl.SSLContext, raw: socket.socket, connection_number: int) -> None:
+    _serve_h2(context, raw, drain=False)
+
+
 @contextmanager
-def _goaway_peer(peer: TlsPeer) -> Generator[str]:
+def _raw_tls_peer(
+    peer: TlsPeer, alpn: str, serve_connection: Callable[[ssl.SSLContext, socket.socket, int], None]
+) -> Generator[tuple[str, queue.SimpleQueue[int]]]:
     context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
     context.load_cert_chain(peer.cert_file, peer.key_file)
-    context.set_alpn_protocols(["h2"])
+    context.set_alpn_protocols([alpn])
     stop: Final = threading.Event()
-    connections: Final = itertools.count(1)
+    accepted: Final = queue.SimpleQueue[int]()
 
     with socket.create_server(("127.0.0.1", 0)) as listener:
         listener.settimeout(0.1)
 
+        def _accept_one() -> None:
+            raw: Final = listener.accept()[0]
+            accepted.put(accepted.qsize() + 1)
+            threading.Thread(target=serve_connection, args=(context, raw, accepted.qsize()), daemon=True).start()
+
         def _accept() -> None:
             while not stop.is_set():
                 with contextlib.suppress(TimeoutError):
-                    threading.Thread(
-                        target=_serve_h2, args=(context, listener.accept()[0], next(connections) % 2 == 1), daemon=True
-                    ).start()
+                    _accept_one()
 
         acceptor: Final = threading.Thread(target=_accept, daemon=True)
         acceptor.start()
         try:
-            yield f"https://127.0.0.1:{listener.getsockname()[1]}"
+            yield f"https://127.0.0.1:{listener.getsockname()[1]}", accepted
         finally:
             stop.set()
             acceptor.join(timeout=5)
@@ -439,6 +489,46 @@ def test_vertex_search_retries_a_connection_the_server_drains_with_goaway(
 ) -> None:
     _trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
 
-    with _goaway_peer(http2_tls_peer) as search_url:
-        assert _vertex_search_file_ids(http2_tls_peer, search_url=search_url) == ["after-goaway"]
-        assert asyncio.run(_avertex_search_file_ids(http2_tls_peer, search_url=search_url)) == ["after-goaway"]
+    with _raw_tls_peer(http2_tls_peer, "h2", _drain_every_other_connection) as (search_url, _):
+        assert _vertex_search_file_ids(http2_tls_peer, search_url=search_url) == ["answered"]
+        assert asyncio.run(_avertex_search_file_ids(http2_tls_peer, search_url=search_url)) == ["answered"]
+
+
+def test_vertex_search_sends_once_when_an_http1_server_hangs_up(
+    monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
+) -> None:
+    import litellm
+
+    _trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
+    monkeypatch.setenv("LITELLM_HTTP2", "False")
+    monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
+
+    with _raw_tls_peer(http2_tls_peer, "http/1.1", _hang_up_after_request) as (search_url, accepted):
+        with pytest.raises(litellm.InternalServerError, match="Server disconnected"):
+            _vertex_search_file_ids(http2_tls_peer, search_url=search_url)
+        assert accepted.qsize() == 1
+        with pytest.raises(litellm.InternalServerError, match="Server disconnected"):
+            asyncio.run(_avertex_search_file_ids(http2_tls_peer, search_url=search_url))
+        assert accepted.qsize() == 2
+
+
+@pytest.mark.parametrize("force_ipv4", [False, True])
+def test_vertex_search_keeps_its_http2_connection_across_an_idle_gap(
+    monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, force_ipv4: bool
+) -> None:
+    import litellm
+    from litellm.constants import AIOHTTP_KEEPALIVE_TIMEOUT
+
+    idle_gap: Final = DEFAULT_LIMITS.keepalive_expiry + 1
+    assert idle_gap < AIOHTTP_KEEPALIVE_TIMEOUT
+    _trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
+    monkeypatch.setattr(litellm, "force_ipv4", force_ipv4)
+
+    async def _search_twice(search_url: str) -> list[str | None]:
+        first: Final = await _avertex_search_file_ids(http2_tls_peer, search_url=search_url)
+        await asyncio.sleep(idle_gap)
+        return [*first, *await _avertex_search_file_ids(http2_tls_peer, search_url=search_url)]
+
+    with _raw_tls_peer(http2_tls_peer, "h2", _answer_every_connection) as (search_url, accepted):
+        assert asyncio.run(_search_twice(search_url)) == ["answered", "answered"]
+        assert accepted.qsize() == 1