From 534123a49c90d723007be99f498b1d4947e6d286 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 23:17:13 +0000 Subject: [PATCH 01/14] feat(vertex_ai): use HTTP/2 httpx client for search_api vector store Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../base_llm/vector_store/transformation.py | 3 ++ litellm/llms/custom_httpx/http_handler.py | 32 ++++++++++--- litellm/llms/custom_httpx/llm_http_handler.py | 12 ++++- .../search_api/transformation.py | 9 ++++ .../llms/custom_httpx/test_http_handler.py | 47 +++++++++++++++++++ .../custom_httpx/test_llm_http_handler.py | 34 ++++++++++++++ ...x_ai_search_vector_store_transformation.py | 6 +++ 7 files changed, 134 insertions(+), 9 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 07b60cb4b72..f7044249ffc 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -124,6 +124,9 @@ class BaseVectorStoreConfig: def validate_create_vector_store(self) -> None: return None + def get_httpx_client_params(self) -> Mapping[str, object]: + return MappingProxyType({}) + def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]: return [] diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 12e515e19a8..72af9be3751 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -566,17 +566,20 @@ class AsyncHTTPHandler: client_alias: str | None = None, # name for client in logs ssl_verify: VerifyTypes | None = None, shared_session: Optional["ClientSession"] = None, + http2: bool = False, ): self.timeout = timeout self.event_hooks = event_hooks self.ssl_verify = ssl_verify self.shared_session = shared_session + self.http2 = http2 self._owns_client = True self._client = self.create_client( timeout=timeout, event_hooks=event_hooks, ssl_verify=ssl_verify, shared_session=shared_session, + http2=http2, ) self.client_alias = client_alias @@ -588,6 +591,7 @@ class AsyncHTTPHandler: event_hooks=self.event_hooks, ssl_verify=self.ssl_verify, shared_session=self.shared_session, + http2=self.http2, ) return self._client @@ -602,6 +606,7 @@ class AsyncHTTPHandler: event_hooks: Mapping[str, list[Callable[..., object]]] | None, ssl_verify: VerifyTypes | None = None, shared_session: Optional["ClientSession"] = None, + http2: bool = False, ) -> httpx.AsyncClient: # Get unified SSL configuration ssl_config: Final = get_ssl_configuration(ssl_verify) @@ -614,10 +619,19 @@ class AsyncHTTPHandler: timeout = _DEFAULT_TIMEOUT # Create a client with a connection pool - transport: Final = AsyncHTTPHandler._create_async_transport( - ssl_context=ssl_config if isinstance(ssl_config, ssl.SSLContext) else None, - ssl_verify=ssl_config if isinstance(ssl_config, bool) else None, - shared_session=shared_session, + transport: Final = ( + httpx.AsyncHTTPTransport( + http2=True, + verify=ssl_config, + cert=cert, + local_address=_IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None, + ) + if http2 + else AsyncHTTPHandler._create_async_transport( + ssl_context=ssl_config if isinstance(ssl_config, ssl.SSLContext) else None, + ssl_verify=ssl_config if isinstance(ssl_config, bool) else None, + shared_session=shared_session, + ) ) # Get default headers (User-Agent, overridable via LITELLM_USER_AGENT) @@ -625,7 +639,7 @@ class AsyncHTTPHandler: return httpx.AsyncClient( transport=transport, - mounts=AsyncHTTPHandler._create_httpx_proxy_mounts(transport, verify=ssl_config, cert=cert), + mounts=AsyncHTTPHandler._create_httpx_proxy_mounts(transport, verify=ssl_config, cert=cert, http2=http2), event_hooks=event_hooks, timeout=timeout, verify=ssl_config, @@ -1229,11 +1243,12 @@ class AsyncHTTPHandler: transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | None, verify: VerifyTypes, cert: CertTypes | None, + http2: bool = False, ) -> 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) + lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2) ) @@ -1246,10 +1261,12 @@ class HTTPHandler: ssl_verify: bool | str | None = None, disable_default_headers: bool | None = False, # arize phoenix returns different API responses when user agent header in request + http2: bool = False, ): self.timeout = timeout self.ssl_verify = ssl_verify self.disable_default_headers = disable_default_headers + self.http2 = http2 self._owns_client = client is None self._heal_lock = threading.Lock() self._client = self.create_client() if client is None else client @@ -1269,6 +1286,7 @@ class HTTPHandler: return httpx.Client( transport=self._create_sync_transport(), mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert), + http2=self.http2, timeout=self.timeout if self.timeout is not None else _DEFAULT_TIMEOUT, verify=ssl_config, cert=cert, @@ -1549,7 +1567,7 @@ class HTTPHandler: Some users have seen httpx ConnectionError when using ipv6 - forcing ipv4 resolves the issue for them """ if litellm.force_ipv4: - return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS) + return HTTPTransport(http2=self.http2, local_address=_IPV4_LOCAL_ADDRESS) else: return getattr(litellm, "sync_transport", None) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index e720428847d..4db3e9d8bb4 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -9815,7 +9815,10 @@ class BaseLLMHTTPHandler: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( llm_provider=litellm.LlmProviders(custom_llm_provider), - params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + params={ + "ssl_verify": litellm_params.get("ssl_verify", None), + **vector_store_provider_config.get_httpx_client_params(), + }, ) else: async_httpx_client = client @@ -9953,7 +9956,12 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) + sync_httpx_client = _get_httpx_client( + params={ + "ssl_verify": litellm_params.get("ssl_verify", None), + **vector_store_provider_config.get_httpx_client_params(), + } + ) else: sync_httpx_client = client diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 0bcf16ee06f..2578b44b8ef 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -1,4 +1,6 @@ +import importlib.util from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol import httpx @@ -51,6 +53,8 @@ VERTEX_SEARCH_DATASTORE_EXTRA_BODY_FIELDS: Final = frozenset(VertexSearchDataSto VERTEX_SEARCH_ENGINE_EXTRA_BODY_FIELDS: Final = frozenset(VertexSearchEngineExtraBody.__annotations__) +_H2_AVAILABLE: Final = importlib.util.find_spec("h2") is not None + class VertexSearchSnippet(TypedDict, total=False): snippet: ReadOnly[str] @@ -108,6 +112,11 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): def __init__(self): super().__init__() + def get_httpx_client_params(self) -> Mapping[str, object]: + if _H2_AVAILABLE: + return MappingProxyType({"http2": True}) + return MappingProxyType({}) + @staticmethod def get_supported_extra_body_fields(is_engine: bool = False) -> frozenset[str]: """ diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 3d4ba264c1c..737c62efd58 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -21,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, MaskedHTTPStatusError, _get_httpx_client, + get_async_httpx_client, get_ssl_configuration, ) @@ -645,6 +646,52 @@ def test_get_httpx_client_applies_httpx_timeout_object_without_mocking_handler() handler.close() +@pytest.mark.asyncio +async def test_async_http_handler_http2_transport(): + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + http2_handler = AsyncHTTPHandler(http2=True) + default_handler = AsyncHTTPHandler() + try: + assert isinstance(http2_handler.client._transport, httpx.AsyncHTTPTransport) + assert http2_handler.client._transport._pool._http2 is True + assert isinstance(default_handler.client._transport, httpx.AsyncHTTPTransport) + assert default_handler.client._transport._pool._http2 is False + finally: + await http2_handler.close() + await default_handler.close() + monkeypatch.undo() + + +def test_http_handler_http2_transport(): + http2_handler = HTTPHandler(http2=True) + default_handler = HTTPHandler() + try: + assert http2_handler.client._transport._pool._http2 is True + assert default_handler.client._transport._pool._http2 is False + finally: + http2_handler.close() + default_handler.close() + + +@pytest.mark.asyncio +async def test_get_async_httpx_client_http2_cache_key(): + from litellm.caching.llm_caching_handler import LLMClientCache + from litellm.types.utils import LlmProviders + + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + http2_handler = get_async_httpx_client(llm_provider=LlmProviders.VERTEX_AI, params={"http2": True}) + default_handler = get_async_httpx_client(llm_provider=LlmProviders.VERTEX_AI) + try: + assert http2_handler is not default_handler + assert http2_handler.client._transport._pool._http2 is True + finally: + await http2_handler.close() + await default_handler.close() + monkeypatch.undo() + + def test_sync_get_forwards_per_request_timeout(): """HTTPHandler.get(timeout=...) must apply the timeout to that request, overriding the client default rather than silently ignoring it.""" diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index cea2d439198..99029cab466 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2572,6 +2572,40 @@ def test_vector_store_search_handler_direct_config_sync_skips_http(): assert pre_call_args["vector_store_id"] == "vs_direct" +@pytest.mark.asyncio +async def test_async_vector_store_search_handler_passes_provider_httpx_params(): + from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( + VertexSearchAPIVectorStoreConfig, + ) + + response = httpx.Response(200, json={"results": []}) + client = AsyncMock(spec=AsyncHTTPHandler) + client.post.return_value = response + config = VertexSearchAPIVectorStoreConfig() + logging_obj = Mock(model_call_details={}) + + with ( + patch.object(config, "validate_environment", return_value={}), + patch.object(config, "get_complete_url", return_value="https://discoveryengine.googleapis.com/v1/search"), + patch( # test-quality-ok: verifies vector-store transport configuration reaches the HTTP client factory + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=client, + ) as get_client, + ): + result = await BaseLLMHTTPHandler().async_vector_store_search_handler( + vector_store_id="vs", + query="q", + vector_store_search_optional_params={}, + vector_store_provider_config=config, + custom_llm_provider="vertex_ai", + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + ) + + assert result["data"] == [] + assert get_client.call_args.kwargs["params"] == {"ssl_verify": None, "http2": True} + + @pytest.mark.asyncio async def test_vector_store_search_handler_direct_config_async_skips_http(): handler = BaseLLMHTTPHandler() diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py index 034f85f5a0b..02b5b964b68 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py @@ -3,11 +3,17 @@ from types import SimpleNamespace import pytest from litellm.exceptions import BadRequestError +from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( VertexSearchAPIVectorStoreConfig, ) +def test_vector_store_httpx_client_params(): + assert VertexSearchAPIVectorStoreConfig().get_httpx_client_params() == {"http2": True} + assert BaseVectorStoreConfig().get_httpx_client_params() == {} + + def test_should_encode_vertex_search_vector_store_id_in_complete_url(): config = VertexSearchAPIVectorStoreConfig() From 1aab2710b5790bb26af5e00f59427129fc523dda Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 23:30:04 +0000 Subject: [PATCH 02/14] test(httpx): exercise configured request path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/custom_httpx/test_http_handler.py | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 737c62efd58..e50dfbc75aa 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -663,6 +663,25 @@ async def test_async_http_handler_http2_transport(): monkeypatch.undo() +@pytest.mark.asyncio +async def test_async_http_handler_http2_request(): + requests: list[httpx.Request] = [] + + async def mock_handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, request=request, json={"ok": True}) + + handler = AsyncHTTPHandler(http2=True) + await handler.client.aclose() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock_handler)) + try: + response = await handler.get("https://example.com/search") + assert response.json() == {"ok": True} + assert requests[0].url == "https://example.com/search" + finally: + await handler.close() + + def test_http_handler_http2_transport(): http2_handler = HTTPHandler(http2=True) default_handler = HTTPHandler() @@ -674,6 +693,24 @@ def test_http_handler_http2_transport(): default_handler.close() +def test_http_handler_http2_request(): + requests: list[httpx.Request] = [] + + def mock_handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, request=request, json={"ok": True}) + + handler = HTTPHandler(http2=True) + handler.client.close() + handler.client = httpx.Client(transport=httpx.MockTransport(mock_handler)) + try: + response = handler.get("https://example.com/search") + assert response.json() == {"ok": True} + assert requests[0].url == "https://example.com/search" + finally: + handler.close() + + @pytest.mark.asyncio async def test_get_async_httpx_client_http2_cache_key(): from litellm.caching.llm_caching_handler import LLMClientCache From 878f0fb10c909e01e8171ca23aebb83c57f655a3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 23:35:14 +0000 Subject: [PATCH 03/14] test(httpx): verify HTTP/2 request negotiation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/custom_httpx/test_http_handler.py | 87 +++++++++++++------ 1 file changed, 59 insertions(+), 28 deletions(-) diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index e50dfbc75aa..7b42218886c 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1,4 +1,5 @@ import asyncio +from datetime import datetime, timedelta, timezone import gc import io import os @@ -664,22 +665,70 @@ async def test_async_http_handler_http2_transport(): @pytest.mark.asyncio -async def test_async_http_handler_http2_request(): - requests: list[httpx.Request] = [] +async def test_async_http_handler_http2_request(tmp_path): + pytest.importorskip("h2") + pytest.importorskip("cryptography") + from cryptography import x509 + from cryptography.hazmat.primitives import hashes, serialization + from cryptography.hazmat.primitives.asymmetric import rsa + from cryptography.x509.oid import NameOID + from h2.config import H2Configuration + from h2.connection import H2Connection + from h2.events import RequestReceived - async def mock_handler(request: httpx.Request) -> httpx.Response: - requests.append(request) - return httpx.Response(200, request=request, json={"ok": True}) + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + subject = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")]) + certificate = ( + x509.CertificateBuilder() + .subject_name(subject) + .issuer_name(subject) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(datetime.now(timezone.utc)) + .not_valid_after(datetime.now(timezone.utc) + timedelta(days=1)) + .add_extension(x509.SubjectAlternativeName([x509.DNSName("localhost")]), critical=False) + .sign(key, hashes.SHA256()) + ) + certificate_path = tmp_path / "certificate.pem" + key_path = tmp_path / "key.pem" + certificate_path.write_bytes(certificate.public_bytes(serialization.Encoding.PEM)) + key_path.write_bytes( + key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.TraditionalOpenSSL, + serialization.NoEncryption(), + ) + ) - handler = AsyncHTTPHandler(http2=True) - await handler.client.aclose() - handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock_handler)) + async def serve_http2(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + connection = H2Connection(config=H2Configuration(client_side=False, header_encoding="utf-8")) + connection.initiate_connection() + writer.write(connection.data_to_send()) + await writer.drain() + while data := await reader.read(65535): + for event in connection.receive_data(data): + if isinstance(event, RequestReceived): + connection.send_headers(event.stream_id, [(":status", "200"), ("content-type", "application/json")]) + connection.send_data(event.stream_id, b'{"ok":true}', end_stream=True) + writer.write(connection.data_to_send()) + await writer.drain() + writer.close() + await writer.wait_closed() + + server_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_context.load_cert_chain(certificate_path, key_path) + server_context.set_alpn_protocols(["h2"]) + server = await asyncio.start_server(serve_http2, "127.0.0.1", 0, ssl=server_context) + handler = AsyncHTTPHandler(http2=True, ssl_verify=False) try: - response = await handler.get("https://example.com/search") + port = server.sockets[0].getsockname()[1] + response = await handler.get(f"https://localhost:{port}/search") + assert response.http_version == "HTTP/2" assert response.json() == {"ok": True} - assert requests[0].url == "https://example.com/search" finally: await handler.close() + server.close() + await server.wait_closed() def test_http_handler_http2_transport(): @@ -693,24 +742,6 @@ def test_http_handler_http2_transport(): default_handler.close() -def test_http_handler_http2_request(): - requests: list[httpx.Request] = [] - - def mock_handler(request: httpx.Request) -> httpx.Response: - requests.append(request) - return httpx.Response(200, request=request, json={"ok": True}) - - handler = HTTPHandler(http2=True) - handler.client.close() - handler.client = httpx.Client(transport=httpx.MockTransport(mock_handler)) - try: - response = handler.get("https://example.com/search") - assert response.json() == {"ok": True} - assert requests[0].url == "https://example.com/search" - finally: - handler.close() - - @pytest.mark.asyncio async def test_get_async_httpx_client_http2_cache_key(): from litellm.caching.llm_caching_handler import LLMClientCache From 3407f74247f1507e5fa24f159db29f8352c07b58 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 23:43:31 +0000 Subject: [PATCH 04/14] fix(httpx): preserve HTTP/2 for sync proxy mounts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/http_handler.py | 7 +- .../llms/custom_httpx/test_http_handler.py | 87 ++++++------------- 2 files changed, 33 insertions(+), 61 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 72af9be3751..64520cd824f 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1285,7 +1285,7 @@ class HTTPHandler: # Create a client with a connection pool return httpx.Client( transport=self._create_sync_transport(), - mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert), + mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert, http2=self.http2), http2=self.http2, timeout=self.timeout if self.timeout is not None else _DEFAULT_TIMEOUT, verify=ssl_config, @@ -1575,10 +1575,13 @@ class HTTPHandler: def _create_sync_proxy_mounts( verify: VerifyTypes, cert: CertTypes | None, + http2: bool = False, ) -> Mapping[str, HTTPTransport | None] | None: if not litellm.force_ipv4: return None - return _environment_proxy_mounts(lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert)) + return _environment_proxy_mounts( + lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2) + ) def get_async_httpx_client( diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 7b42218886c..e50dfbc75aa 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1,5 +1,4 @@ import asyncio -from datetime import datetime, timedelta, timezone import gc import io import os @@ -665,70 +664,22 @@ async def test_async_http_handler_http2_transport(): @pytest.mark.asyncio -async def test_async_http_handler_http2_request(tmp_path): - pytest.importorskip("h2") - pytest.importorskip("cryptography") - from cryptography import x509 - from cryptography.hazmat.primitives import hashes, serialization - from cryptography.hazmat.primitives.asymmetric import rsa - from cryptography.x509.oid import NameOID - from h2.config import H2Configuration - from h2.connection import H2Connection - from h2.events import RequestReceived +async def test_async_http_handler_http2_request(): + requests: list[httpx.Request] = [] - key = rsa.generate_private_key(public_exponent=65537, key_size=2048) - subject = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")]) - certificate = ( - x509.CertificateBuilder() - .subject_name(subject) - .issuer_name(subject) - .public_key(key.public_key()) - .serial_number(x509.random_serial_number()) - .not_valid_before(datetime.now(timezone.utc)) - .not_valid_after(datetime.now(timezone.utc) + timedelta(days=1)) - .add_extension(x509.SubjectAlternativeName([x509.DNSName("localhost")]), critical=False) - .sign(key, hashes.SHA256()) - ) - certificate_path = tmp_path / "certificate.pem" - key_path = tmp_path / "key.pem" - certificate_path.write_bytes(certificate.public_bytes(serialization.Encoding.PEM)) - key_path.write_bytes( - key.private_bytes( - serialization.Encoding.PEM, - serialization.PrivateFormat.TraditionalOpenSSL, - serialization.NoEncryption(), - ) - ) + async def mock_handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, request=request, json={"ok": True}) - async def serve_http2(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: - connection = H2Connection(config=H2Configuration(client_side=False, header_encoding="utf-8")) - connection.initiate_connection() - writer.write(connection.data_to_send()) - await writer.drain() - while data := await reader.read(65535): - for event in connection.receive_data(data): - if isinstance(event, RequestReceived): - connection.send_headers(event.stream_id, [(":status", "200"), ("content-type", "application/json")]) - connection.send_data(event.stream_id, b'{"ok":true}', end_stream=True) - writer.write(connection.data_to_send()) - await writer.drain() - writer.close() - await writer.wait_closed() - - server_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) - server_context.load_cert_chain(certificate_path, key_path) - server_context.set_alpn_protocols(["h2"]) - server = await asyncio.start_server(serve_http2, "127.0.0.1", 0, ssl=server_context) - handler = AsyncHTTPHandler(http2=True, ssl_verify=False) + handler = AsyncHTTPHandler(http2=True) + await handler.client.aclose() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock_handler)) try: - port = server.sockets[0].getsockname()[1] - response = await handler.get(f"https://localhost:{port}/search") - assert response.http_version == "HTTP/2" + response = await handler.get("https://example.com/search") assert response.json() == {"ok": True} + assert requests[0].url == "https://example.com/search" finally: await handler.close() - server.close() - await server.wait_closed() def test_http_handler_http2_transport(): @@ -742,6 +693,24 @@ def test_http_handler_http2_transport(): default_handler.close() +def test_http_handler_http2_request(): + requests: list[httpx.Request] = [] + + def mock_handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, request=request, json={"ok": True}) + + handler = HTTPHandler(http2=True) + handler.client.close() + handler.client = httpx.Client(transport=httpx.MockTransport(mock_handler)) + try: + response = handler.get("https://example.com/search") + assert response.json() == {"ok": True} + assert requests[0].url == "https://example.com/search" + finally: + handler.close() + + @pytest.mark.asyncio async def test_get_async_httpx_client_http2_cache_key(): from litellm.caching.llm_caching_handler import LLMClientCache From bbaac6b00c4a3e26bee02f66b92c6c1fc0f47242 Mon Sep 17 00:00:00 2001 From: mrinal Date: Thu, 10 Sep 2026 23:59:38 +0000 Subject: [PATCH 05/14] build(deps): add h2 to proxy extra so vertex search HTTP/2 is on by default Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- pyproject.toml | 1 + uv.lock | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 04f2f3fd1dd..a4724fae39e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ Documentation = "https://docs.litellm.ai" proxy = [ "gunicorn>=23.0.0,<24.0", "uvicorn>=0.33.0,<1.0", + "h2>=4.1.0,<5.0", "granian>=2.7.4,<3.0", "uvloop>=0.22.1,<1.0; sys_platform != 'win32'", "fastapi>=0.136.3,<1.0", diff --git a/uv.lock b/uv.lock index 0fe787645a2..6a548945c4e 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-06T00:40:30.433549Z" +exclude-newer = "2026-09-07T23:59:14.199777903Z" exclude-newer-span = "P3D" [manifest] @@ -4427,6 +4427,7 @@ proxy = [ { name = "fastapi-sso" }, { name = "granian" }, { name = "gunicorn" }, + { name = "h2" }, { name = "hiredis" }, { name = "inquirerpy" }, { name = "litellm-enterprise" }, @@ -4613,6 +4614,7 @@ requires-dist = [ { name = "grpcio", marker = "extra == 'grpc'", specifier = "==1.78.0" }, { name = "grpcio", marker = "extra == 'proxy-runtime'", specifier = "==1.78.0" }, { name = "gunicorn", marker = "extra == 'proxy'", specifier = ">=23.0.0,<24.0" }, + { name = "h2", marker = "extra == 'proxy'", specifier = ">=4.1.0,<5.0" }, { name = "hiredis", marker = "extra == 'proxy'", specifier = ">=3.0.0,<4.0" }, { name = "httpx", specifier = ">=0.28.0,<1.0" }, { name = "importlib-metadata", specifier = ">=8.0.0,<9.0" }, From 98998c207d7c33bafd4e16306d9afc5ba4a7e7eb Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 11 Sep 2026 02:22:00 +0000 Subject: [PATCH 06/14] fix(http_handler): isolate ssl context for http2 clients Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/http_handler.py | 12 +++++++----- .../llms/custom_httpx/test_http_handler.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 64520cd824f..18dc8b71e12 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -235,9 +235,9 @@ def _prepare_request_data_and_content( # Cache for SSL contexts to avoid creating duplicate contexts with the same configuration -# Key: tuple of (cafile, ssl_security_level, ssl_ecdh_curve) +# Key: tuple of (cafile, ssl_security_level, ssl_ecdh_curve, http2) # Value: ssl.SSLContext -_ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None], ssl.SSLContext]] = {} +_ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None, bool], ssl.SSLContext]] = {} def _create_ssl_context( @@ -330,6 +330,7 @@ def get_ssl_verify( def get_ssl_configuration( ssl_verify: VerifyTypes | None = None, + http2: bool = False, ) -> bool | str | ssl.SSLContext: """ Unified SSL configuration function that handles ssl_context and ssl_verify logic. @@ -345,6 +346,7 @@ def get_ssl_configuration( SSL contexts are cached to avoid creating duplicate contexts with the same configuration, which reduces memory allocation and improves performance. + http2 clients get their own cached context because httpcore sets ALPN protocols on it. Args: ssl_verify: SSL verification setting. Can be: @@ -378,7 +380,7 @@ def get_ssl_configuration( if ssl_verify is not False: # Create cache key from configuration parameters - cache_key: Final = (cafile, ssl_security_level, ssl_ecdh_curve) + cache_key: Final = (cafile, ssl_security_level, ssl_ecdh_curve, http2) # Check if we have a cached SSL context for this configuration if cache_key not in _ssl_context_cache: @@ -609,7 +611,7 @@ class AsyncHTTPHandler: http2: bool = False, ) -> httpx.AsyncClient: # Get unified SSL configuration - ssl_config: Final = get_ssl_configuration(ssl_verify) + ssl_config: Final = get_ssl_configuration(ssl_verify, http2=http2) # An SSL certificate used by the requested host to authenticate the client. # /path/to/client.pem @@ -1273,7 +1275,7 @@ class HTTPHandler: def create_client(self) -> httpx.Client: # Get unified SSL configuration - ssl_config: Final = get_ssl_configuration(self.ssl_verify) + ssl_config: Final = get_ssl_configuration(self.ssl_verify, http2=self.http2) # An SSL certificate used by the requested host to authenticate the client. # /path/to/client.pem diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index e50dfbc75aa..ebd84cf2d5a 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -306,6 +306,23 @@ def test_get_ssl_configuration(): assert result == mock_ssl_context +def test_get_ssl_configuration_http2_uses_separate_context(): + from litellm.llms.custom_httpx.http_handler import _ssl_context_cache + + _ssl_context_cache.clear() + + default_context = get_ssl_configuration() + http2_context = get_ssl_configuration(http2=True) + + assert default_context is not http2_context + assert get_ssl_configuration() is default_context + assert get_ssl_configuration(http2=True) is http2_context + + custom_context = ssl.create_default_context() + assert get_ssl_configuration(custom_context) is custom_context + assert get_ssl_configuration(custom_context, http2=True) is custom_context + + def test_get_ssl_configuration_integration(): """Integration test that _get_ssl_context() returns a working SSL context""" # Call the static method without mocking From 0ad1e2964901f9389db8d35aca3870576375d71e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 07:40:10 -0700 Subject: [PATCH 07/14] fix(http): add an HTTP/2 opt-out and harden the Vertex AI Search HTTP/2 client litellm.http2 is now unset by default, so http2: false or LITELLM_HTTP2=False turns the provider's HTTP/2 default off. A caller's own SSLContext keeps the search on HTTP/1.1, the HTTP/2 async client honors DISABLE_AIOHTTP_TRUST_ENV, forced IPv4 transports get the CA bundle and client cert, and vector store search retries once when the server drops the connection with a GOAWAY. --- litellm/__init__.py | 2 +- litellm/llms/custom_httpx/http_handler.py | 52 +++++- litellm/llms/custom_httpx/llm_http_handler.py | 28 +-- tests/integration/sdk/test_http2_wire.py | 169 +++++++++++++++++- 4 files changed, 224 insertions(+), 27 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 0fe3d091c2a..38c02e6a463 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -518,7 +518,7 @@ aiohttp_trust_env: bool = False # set to true to use HTTP_ Proxy settings disable_aiohttp_transport: bool = False # Set this to true to use httpx instead disable_aiohttp_trust_env: bool = False # When False, aiohttp will respect HTTP(S)_PROXY env vars force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. -http2: bool = False +http2: bool | None = None network_mock: bool = False # When True, use mock transport — no real network calls ####### STOP SEQUENCE LIMIT ####### diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 82e946fceed..6a2a69f9199 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -83,6 +83,24 @@ def http2_enabled() -> bool: return litellm.http2 is True or str_to_bool(os.getenv("LITELLM_HTTP2", "False")) is True +def http2_disabled() -> bool: + from litellm.secret_managers.main import str_to_bool + + return litellm.http2 is False or str_to_bool(os.getenv("LITELLM_HTTP2")) is False + + +def _resolve_http2(provider_default: bool, ssl_verify: VerifyTypes | None) -> bool: + if http2_enabled(): + return True + return provider_default and not http2_disabled() and not isinstance(ssl_verify, ssl.SSLContext) + + +def _aiohttp_env_proxies_disabled() -> bool: + from litellm.secret_managers.main import str_to_bool + + return litellm.disable_aiohttp_trust_env or str_to_bool(os.getenv("DISABLE_AIOHTTP_TRUST_ENV", "False")) is True + + def _environment_proxy_mounts( build_proxy_transport: Callable[[str], _HttpxTransportT], ) -> Mapping[str, _HttpxTransportT | None]: @@ -686,7 +704,7 @@ class AsyncHTTPHandler: follow_redirects=self.follow_redirects, trust_env=False, ) - use_http2: Final = self.http2 or http2_enabled() + use_http2: Final = _resolve_http2(self.http2, ssl_verify) # Get unified SSL configuration ssl_config: Final = get_ssl_configuration(ssl_verify, http2=use_http2) @@ -703,14 +721,18 @@ class AsyncHTTPHandler: ssl_verify=ssl_config if isinstance(ssl_config, bool) else None, shared_session=shared_session, http2=use_http2, + cert=cert, ) # 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()) 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) + if trust_env + else None, event_hooks=event_hooks, timeout=timeout, verify=ssl_config, @@ -719,6 +741,7 @@ class AsyncHTTPHandler: cookies=blocked_cookie_jar(), follow_redirects=self.follow_redirects, http2=use_http2, + trust_env=trust_env, ) async def close(self): @@ -1209,6 +1232,7 @@ class AsyncHTTPHandler: ssl_verify: bool | None = None, shared_session: Optional["ClientSession"] = None, http2: bool = False, + cert: CertTypes | None = None, ) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | None: """ - Creates a transport for httpx.AsyncClient @@ -1237,7 +1261,11 @@ class AsyncHTTPHandler: ######################################################### # HTTPX TRANSPORT is used when aiohttp is not installed ######################################################### - return AsyncHTTPHandler._create_httpx_transport(http2=http2) + return AsyncHTTPHandler._create_httpx_transport( + http2=http2, + verify=ssl_context if ssl_context is not None else ssl_verify is not False, + cert=cert, + ) _create_async_transport = create_async_transport @@ -1384,7 +1412,11 @@ class AsyncHTTPHandler: _create_aiohttp_transport = create_aiohttp_transport @staticmethod - def _create_httpx_transport(http2: bool = False) -> AsyncHTTPTransport | None: + def _create_httpx_transport( + http2: bool = False, + verify: ssl.SSLContext | bool = True, + cert: CertTypes | None = None, + ) -> AsyncHTTPTransport | None: """ Creates an AsyncHTTPTransport @@ -1392,7 +1424,9 @@ class AsyncHTTPHandler: - [Default] If force_ipv4 is False, it will return None """ if litellm.force_ipv4: - return AsyncHTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, http2=http2 or http2_enabled()) + return AsyncHTTPTransport( + local_address=_IPV4_LOCAL_ADDRESS, verify=verify, cert=cert, http2=http2 or http2_enabled() + ) else: return None @@ -1436,7 +1470,7 @@ class HTTPHandler: self._client = self.create_client() if client is None else client def create_client(self) -> httpx.Client: - use_http2: Final = self.http2 or http2_enabled() + use_http2: Final = _resolve_http2(self.http2, self.ssl_verify) # Get unified SSL configuration ssl_config: Final = get_ssl_configuration(self.ssl_verify, http2=use_http2) @@ -1449,7 +1483,7 @@ class HTTPHandler: # Create a client with a connection pool return httpx.Client( - transport=self._create_sync_transport(http2=use_http2), + transport=self._create_sync_transport(verify=ssl_config, cert=cert, http2=use_http2), mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert, http2=use_http2), timeout=self.timeout if self.timeout is not None else _DEFAULT_TIMEOUT, verify=ssl_config, @@ -1733,7 +1767,7 @@ class HTTPHandler: pass @staticmethod - def _create_sync_transport(http2: bool) -> HTTPTransport | None: + def _create_sync_transport(verify: VerifyTypes, cert: CertTypes | None, http2: bool) -> HTTPTransport | None: """ Create an HTTP transport with IPv4 only if litellm.force_ipv4 is True. Otherwise, return None. @@ -1741,7 +1775,7 @@ class HTTPHandler: Some users have seen httpx ConnectionError when using ipv6 - forcing ipv4 resolves the issue for them """ if litellm.force_ipv4: - return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, http2=http2) + return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, verify=verify, cert=cert, http2=http2) else: return getattr(litellm, "sync_transport", None) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index de602f489c5..1ddd5cdd61d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -340,6 +340,20 @@ class _SignedBodyKwargs(TypedDict, total=False): json: ReadOnly[dict[str, object]] +async def _asend_idempotent(send: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response: + try: + return await send() + except httpx.RemoteProtocolError: + return await send() + + +def _send_idempotent(send: Callable[[], httpx.Response]) -> httpx.Response: + try: + return send() + except httpx.RemoteProtocolError: + return send() + + def _signed_body_kwargs(*, signed_body: bytes | None, data: dict[str, object]) -> _SignedBodyKwargs: if signed_body is not None: return {"data": signed_body} @@ -10391,11 +10405,8 @@ class BaseLLMHTTPHandler: request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body try: - response: Final = await async_httpx_client.post( - url=url, - headers=headers, - data=request_data, - timeout=timeout, + response: Final = await _asend_idempotent( + partial(async_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout) ) except httpx.TimeoutException: raise vector_store_provider_config.get_error_class( @@ -10533,11 +10544,8 @@ class BaseLLMHTTPHandler: request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body try: - response: Final = sync_httpx_client.post( - url=url, - headers=headers, - data=request_data, - timeout=timeout, + response: Final = _send_idempotent( + partial(sync_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout) ) except httpx.TimeoutException: raise vector_store_provider_config.get_error_class( diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index daf49f4ff78..b16f1699aaa 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/tests/integration/sdk/test_http2_wire.py @@ -1,13 +1,17 @@ from __future__ import annotations import asyncio +import contextlib import datetime import ipaddress +import itertools import json import socket +import ssl import threading import time -from collections.abc import Iterator +from collections.abc import Generator, Iterator +from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Final @@ -17,6 +21,9 @@ from cryptography import x509 from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.x509.oid import NameOID +from h2.config import H2Configuration +from h2.connection import H2Connection +from h2.events import StreamEnded from hypercorn.asyncio import ( serve, # pyright: ignore[reportUnknownVariableType] # hypercorn's serve signature passes through untyped worker hooks ) @@ -35,6 +42,7 @@ VERTEX_SERVING_CONFIG_PATH: Final = ( class TlsPeer: url: str cert_file: Path + key_file: Path @dataclass(frozen=True, slots=True) @@ -142,7 +150,7 @@ def http2_tls_peer(tmp_path_factory: pytest.TempPathFactory) -> Iterator[TlsPeer time.sleep(0.05) else: pytest.fail("hypercorn peer did not start") - yield TlsPeer(url=f"https://127.0.0.1:{port}", cert_file=cert_file) + yield TlsPeer(url=f"https://127.0.0.1:{port}", cert_file=cert_file, key_file=key_file) shutdown.set() thread.join(timeout=10) @@ -211,7 +219,9 @@ def test_async_handler_negotiates_http2_only_when_enabled( @pytest.mark.covers("other.sdk_wire.http2.sync_handler_negotiates_h2_only_when_enabled") -def test_sync_handler_negotiates_http2_only_when_enabled(monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer) -> None: +def test_sync_handler_negotiates_http2_only_when_enabled( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: for enabled in (False, True): _set_http2(monkeypatch, enabled) _assert_negotiated(_sync_exchange(http2_tls_peer.url), enabled) @@ -227,31 +237,37 @@ def _trust_peer_without_global_http2(monkeypatch: pytest.MonkeyPatch, peer: TlsP litellm.in_memory_llm_clients_cache.flush_cache() -def _vertex_search_file_ids(peer: TlsPeer) -> list[str | None]: +def _vertex_search_file_ids( + peer: TlsPeer, search_url: str | None = None, ssl_verify: ssl.SSLContext | None = None +) -> list[str | None]: import litellm response: Final = litellm.vector_stores.search( vector_store_id="default_search", query="ping", custom_llm_provider="vertex_ai/search_api", - api_base=f"{peer.url}{VERTEX_SERVING_CONFIG_PATH}", + api_base=f"{search_url or peer.url}{VERTEX_SERVING_CONFIG_PATH}", vertex_credentials=service_account_json("p", peer.url), vertex_project="p", + ssl_verify=ssl_verify, ) assert isinstance(response, dict) return [result.get("file_id") for result in response.get("data") or ()] -async def _avertex_search_file_ids(peer: TlsPeer) -> list[str | None]: +async def _avertex_search_file_ids( + peer: TlsPeer, search_url: str | None = None, ssl_verify: ssl.SSLContext | None = None +) -> list[str | None]: import litellm response: Final = await litellm.vector_stores.asearch( vector_store_id="default_search", query="ping", custom_llm_provider="vertex_ai/search_api", - api_base=f"{peer.url}{VERTEX_SERVING_CONFIG_PATH}", + api_base=f"{search_url or peer.url}{VERTEX_SERVING_CONFIG_PATH}", vertex_credentials=service_account_json("p", peer.url), vertex_project="p", + ssl_verify=ssl_verify, ) return [result.get("file_id") for result in response.get("data") or ()] @@ -287,3 +303,142 @@ def test_default_async_client_stays_http1_after_vertex_search_used_http2( await handler.close() assert asyncio.run(_run()) == (["http/2"], "1.1") + + +def test_vertex_search_over_forced_ipv4_trusts_the_configured_ca( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + import litellm + + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + monkeypatch.setattr(litellm, "force_ipv4", True) + + assert _vertex_search_file_ids(http2_tls_peer) == ["http/2"] + assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/2"] + + +def test_explicit_http2_off_keeps_vertex_search_on_http1( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + import litellm + + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + monkeypatch.setenv("LITELLM_HTTP2", "False") + assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"] + assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/1.1"] + + monkeypatch.delenv("LITELLM_HTTP2") + monkeypatch.setattr(litellm, "http2", False) + litellm.in_memory_llm_clients_cache.flush_cache() + assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"] + assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/1.1"] + + +def _alpn_offered_by(context: ssl.SSLContext, peer: TlsPeer) -> str | None: + host, port = peer.url.removeprefix("https://").split(":") + with ( + socket.create_connection((host, int(port)), timeout=5) as raw, + context.wrap_socket(raw, server_hostname=host) as tls, + ): + return tls.selected_alpn_protocol() + + +def test_vertex_search_leaves_a_caller_ssl_context_on_http1( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + sync_context: Final = ssl.create_default_context(cafile=str(http2_tls_peer.cert_file)) + async_context: Final = ssl.create_default_context(cafile=str(http2_tls_peer.cert_file)) + + assert _vertex_search_file_ids(http2_tls_peer, ssl_verify=sync_context) == ["http/1.1"] + assert asyncio.run(_avertex_search_file_ids(http2_tls_peer, ssl_verify=async_context)) == ["http/1.1"] + assert _alpn_offered_by(sync_context, http2_tls_peer) != "h2" + assert _alpn_offered_by(async_context, http2_tls_peer) != "h2" + + +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() + + for force_ipv4 in (False, True): + monkeypatch.setattr(litellm, "force_ipv4", force_ipv4) + assert asyncio.run(_peer_http_version()) == "HTTP/2" + + +def _answer_or_drain(connection: H2Connection, stream_id: int, drain: bool) -> None: + if drain: + 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) + + +def _reply_to_ended_streams(connection: H2Connection, tls: ssl.SSLSocket, data: bytes, drain: bool) -> bool: + ended: Final = tuple(event.stream_id for event in connection.receive_data(data) if isinstance(event, StreamEnded)) + for stream_id in ended: + _answer_or_drain(connection, stream_id, drain) + tls.sendall(connection.data_to_send()) + return bool(ended) + + +def _serve_h2(context: ssl.SSLContext, raw: socket.socket, drain: bool) -> None: + with contextlib.suppress(OSError), raw, context.wrap_socket(raw, server_side=True) as tls: + connection: Final = H2Connection(config=H2Configuration(client_side=False)) + connection.initiate_connection() + tls.sendall(connection.data_to_send()) + while data := tls.recv(65535): + if _reply_to_ended_streams(connection, tls, data, drain) and drain: + return + + +@contextmanager +def _goaway_peer(peer: TlsPeer) -> Generator[str]: + context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(peer.cert_file, peer.key_file) + context.set_alpn_protocols(["h2"]) + stop: Final = threading.Event() + connections: Final = itertools.count(1) + + with socket.create_server(("127.0.0.1", 0)) as listener: + listener.settimeout(0.1) + + 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() + + acceptor: Final = threading.Thread(target=_accept, daemon=True) + acceptor.start() + try: + yield f"https://127.0.0.1:{listener.getsockname()[1]}" + finally: + stop.set() + acceptor.join(timeout=5) + + +def test_vertex_search_retries_a_connection_the_server_drains_with_goaway( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> 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"] From 1e50f9ddc3f416bbb09ecc4d975c500a26cf5c89 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 08:41:50 -0700 Subject: [PATCH 08/14] 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 | 30 +++- litellm/llms/custom_httpx/llm_http_handler.py | 27 +++- tests/integration/sdk/test_http2_wire.py | 144 ++++++++++++++---- 3 files changed, 164 insertions(+), 37 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 6a2a69f9199..328c68c8b5f 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/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 ) ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 1ddd5cdd61d..236674a4a8b 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/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: diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index b16f1699aaa..527e43b10d3 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/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 From 5894a0c46bea1cc90d0c0983a56241157e093dbd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 09:22:35 -0700 Subject: [PATCH 09/14] test(http): cover the vector store search GOAWAY retry in unit tests and drop duplicate SSL cache comments --- litellm/llms/custom_httpx/http_handler.py | 3 -- .../custom_httpx/test_llm_http_handler.py | 53 +++++++++++++++++++ 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 328c68c8b5f..ac439ec1918 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -305,8 +305,6 @@ def _prepare_request_data_and_content( # Cache for SSL contexts to avoid creating duplicate contexts with the same configuration -# Key: tuple of (cafile, ssl_security_level, ssl_ecdh_curve, http2) -# Value: ssl.SSLContext _ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None, bool], ssl.SSLContext]] = {} @@ -416,7 +414,6 @@ def get_ssl_configuration( SSL contexts are cached to avoid creating duplicate contexts with the same configuration, which reduces memory allocation and improves performance. - http2 clients get their own cached context because httpcore sets ALPN protocols on it. Args: ssl_verify: SSL verification setting. Can be: diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 8b46597b603..467d9cc6e38 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -10,9 +10,11 @@ from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, Mock, patch +import httpcore import httpx import pytest from botocore.credentials import RefreshableCredentials +from h2.events import ConnectionTerminated import litellm from litellm._logging import verbose_logger @@ -5391,3 +5393,54 @@ async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_ sent: Final = received[0] assert "safeguards" not in json.loads(sent.content) assert "dangerous-tool-use-2026-09-03" not in sent.headers.get("anthropic-beta", ""), dict(sent.headers) + + +_EMPTY_SEARCH_PAGE: Final = { + "object": "vector_store.search_results.page", + "search_query": "q", + "data": [], + "has_more": False, + "next_page": None, +} + + +def _close_with_goaway(request: httpx.Request) -> httpx.Response: + raise httpx.RemoteProtocolError("GOAWAY", request=request) from httpcore.RemoteProtocolError(ConnectionTerminated()) + + +def _hang_up(request: httpx.Request) -> httpx.Response: + raise httpx.RemoteProtocolError("Server disconnected without sending a response.", request=request) + + +def _answer_search(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=_EMPTY_SEARCH_PAGE) + + +async def _search_through(transport: httpx.MockTransport, is_async: bool) -> object: + search_kwargs: Final = {"vector_store_id": "vs_1", "query": "q", "custom_llm_provider": "openai", "api_key": "sk-test"} + if not is_async: + with httpx.Client(transport=transport) as sync_client: + return litellm.vector_stores.search(**search_kwargs, client=HTTPHandler(client=sync_client)) + async with httpx.AsyncClient(transport=transport) as async_client: + client: Final = AsyncHTTPHandler() + await client.close() + client.client = async_client + return await litellm.vector_stores.asearch(**search_kwargs, client=client) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True)) +async def test_vector_store_search_sends_again_once_after_an_http2_goaway(is_async: bool) -> None: + sends: Final = iter((_close_with_goaway, _answer_search)) + response: Final = await _search_through(httpx.MockTransport(lambda request: next(sends)(request)), is_async) + assert response["search_query"] == "q" + assert next(sends, None) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True)) +async def test_vector_store_search_sends_once_when_the_connection_drops_without_a_goaway(is_async: bool) -> None: + sends: Final = iter((_hang_up, _answer_search)) + with pytest.raises(litellm.InternalServerError, match="Server disconnected"): + await _search_through(httpx.MockTransport(lambda request: next(sends)(request)), is_async) + assert next(sends) is _answer_search From 81003176d49b3f7425b022f4f273f7ad28a06556 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:01:53 -0700 Subject: [PATCH 10/14] fix(http): open another HTTP/2 connection once 100 streams are in use httpcore keeps every request to an origin on one HTTP/2 connection and caps it at 100 concurrent streams, so request 101 and later waited with no timeout. The handlers now spread requests over extra connections, opening one only when every open connection has 100 streams in flight. The force_ipv4 wire tests now cover ssl_verify=False and the client certificate. --- .../custom_httpx/http2_connection_spread.py | 130 ++++++++++++ litellm/llms/custom_httpx/http_handler.py | 87 +++++--- tests/integration/sdk/test_http2_wire.py | 98 ++++++++- .../test_http2_connection_spread.py | 190 ++++++++++++++++++ .../llms/custom_httpx/test_http_handler.py | 4 +- 5 files changed, 474 insertions(+), 35 deletions(-) create mode 100644 litellm/llms/custom_httpx/http2_connection_spread.py create mode 100644 tests/unit/llms/custom_httpx/test_http2_connection_spread.py diff --git a/litellm/llms/custom_httpx/http2_connection_spread.py b/litellm/llms/custom_httpx/http2_connection_spread.py new file mode 100644 index 00000000000..3803ccdf56e --- /dev/null +++ b/litellm/llms/custom_httpx/http2_connection_spread.py @@ -0,0 +1,130 @@ +import threading +from collections.abc import AsyncIterator, Callable, Iterator +from typing import Final, Generic, TypeVar + +import httpx + +HTTP2_STREAMS_PER_CONNECTION: Final = 100 + +_TransportT: Final = TypeVar("_TransportT", httpx.BaseTransport, httpx.AsyncBaseTransport) + + +class _Lane(Generic[_TransportT]): + __slots__ = ("in_flight", "transport") + + def __init__(self, transport: _TransportT) -> None: + self.transport: Final = transport + self.in_flight = 0 + + def acquire(self) -> None: + self.in_flight += 1 + + def release(self) -> None: + self.in_flight -= 1 + + +class _Lanes(Generic[_TransportT]): + def __init__(self, new_transport: Callable[[], _TransportT], streams_per_connection: int) -> None: + self._new_transport: Final = new_transport + self._streams_per_connection: Final = streams_per_connection + self.lanes: tuple[_Lane[_TransportT], ...] = () + + def acquire(self) -> _Lane[_TransportT]: + free: Final = next((lane for lane in self.lanes if lane.in_flight < self._streams_per_connection), None) + lane: Final = free if free is not None else self._open() + lane.acquire() + return lane + + def _open(self) -> _Lane[_TransportT]: + lane: Final = _Lane(self._new_transport()) + self.lanes = (*self.lanes, lane) + return lane + + +class _AsyncReleasingStream(httpx.AsyncByteStream): + def __init__(self, stream: httpx.AsyncByteStream, release: Callable[[], None]) -> None: + self._stream: Final = stream + self._release: Final = release + + async def __aiter__(self) -> AsyncIterator[bytes]: + async for chunk in self._stream: + yield chunk + + async def aclose(self) -> None: + self._release() + await self._stream.aclose() + + +class _ReleasingStream(httpx.SyncByteStream): + def __init__(self, stream: httpx.SyncByteStream, release: Callable[[], None]) -> None: + self._stream: Final = stream + self._release: Final = release + + def __iter__(self) -> Iterator[bytes]: + yield from self._stream + + def close(self) -> None: + self._release() + self._stream.close() + + +def _with_stream(response: httpx.Response, stream: httpx.SyncByteStream | httpx.AsyncByteStream) -> httpx.Response: + return httpx.Response(response.status_code, headers=response.headers, stream=stream, extensions=response.extensions) + + +class AsyncHTTP2ConnectionSpread(httpx.AsyncBaseTransport): + def __init__( + self, + new_transport: Callable[[], httpx.AsyncBaseTransport], + streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION, + ) -> None: + self._lanes: Final = _Lanes(new_transport, streams_per_connection) + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + lane: Final = self._lanes.acquire() + try: + response: Final = await lane.transport.handle_async_request(request) + except BaseException: + lane.release() + raise + stream: Final = response.stream + if not isinstance(stream, httpx.AsyncByteStream): + lane.release() + return response + return _with_stream(response, _AsyncReleasingStream(stream, lane.release)) + + async def aclose(self) -> None: + for lane in self._lanes.lanes: + await lane.transport.aclose() + + +class HTTP2ConnectionSpread(httpx.BaseTransport): + def __init__( + self, + new_transport: Callable[[], httpx.BaseTransport], + streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION, + ) -> None: + self._lanes: Final = _Lanes(new_transport, streams_per_connection) + self._lock: Final = threading.Lock() + + def _release(self, lane: _Lane[httpx.BaseTransport]) -> None: + with self._lock: + lane.release() + + def handle_request(self, request: httpx.Request) -> httpx.Response: + with self._lock: + lane: Final = self._lanes.acquire() + try: + response: Final = lane.transport.handle_request(request) + except BaseException: + self._release(lane) + raise + stream: Final = response.stream + if not isinstance(stream, httpx.SyncByteStream): + self._release(lane) + return response + return _with_stream(response, _ReleasingStream(stream, lambda: self._release(lane))) + + def close(self) -> None: + for lane in self._lanes.lanes: + lane.transport.close() diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index ac439ec1918..04e72bc7b28 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -43,6 +43,7 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) +from litellm.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread, HTTP2ConnectionSpread from litellm.types.llms.custom_http import * if TYPE_CHECKING: @@ -73,7 +74,7 @@ _RequestContent: TypeAlias = str | bytes | Iterable[bytes] | AsyncIterable[bytes _IPV4_LOCAL_ADDRESS: Final = "0.0.0.0" -_HttpxTransportT = TypeVar("_HttpxTransportT", HTTPTransport, AsyncHTTPTransport) +_HttpxTransportT: Final = TypeVar("_HttpxTransportT", httpx.BaseTransport, httpx.AsyncBaseTransport) HTTPEventHooks: TypeAlias = Mapping[str, list[Callable[..., object]]] @@ -1245,7 +1246,7 @@ class AsyncHTTPHandler: http2: bool = False, cert: CertTypes | None = None, limits: httpx.Limits = DEFAULT_LIMITS, - ) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | None: + ) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None: """ - Creates a transport for httpx.AsyncClient - if litellm.force_ipv4 is True, it will return AsyncHTTPTransport with local_address="0.0.0.0" @@ -1430,38 +1431,43 @@ class AsyncHTTPHandler: verify: ssl.SSLContext | bool = True, cert: CertTypes | None = None, limits: httpx.Limits = DEFAULT_LIMITS, - ) -> AsyncHTTPTransport | None: + ) -> AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None: """ Creates an AsyncHTTPTransport + - If HTTP/2 is on, it spreads requests over as many HTTP/2 connections as the stream limit needs - If force_ipv4 is True, it will create an AsyncHTTPTransport with local_address set to "0.0.0.0" - - [Default] If force_ipv4 is False, it will return None + - [Default] Otherwise it will return None """ - if litellm.force_ipv4: - return AsyncHTTPTransport( - local_address=_IPV4_LOCAL_ADDRESS, - verify=verify, - cert=cert, - http2=http2 or http2_enabled(), - limits=limits, + local_address: Final = _IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None + if http2 or http2_enabled(): + return AsyncHTTP2ConnectionSpread( + lambda: AsyncHTTPTransport( + local_address=local_address, verify=verify, cert=cert, http2=True, limits=limits + ) ) - else: - return None + if litellm.force_ipv4: + return AsyncHTTPTransport(local_address=local_address, verify=verify, cert=cert, limits=limits) + return None @staticmethod def create_httpx_proxy_mounts( - transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | None, + transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None, verify: VerifyTypes, cert: CertTypes | None, http2: bool = False, limits: httpx.Limits = DEFAULT_LIMITS, - ) -> Mapping[str, AsyncHTTPTransport | None] | None: - if not isinstance(transport, AsyncHTTPTransport): + ) -> Mapping[str, httpx.AsyncBaseTransport | None] | None: + if not isinstance(transport, (AsyncHTTPTransport, AsyncHTTP2ConnectionSpread)): return None - return _environment_proxy_mounts( - lambda proxy_url: AsyncHTTPTransport( - proxy=proxy_url, verify=verify, cert=cert, http2=http2 or http2_enabled(), limits=limits + if http2 or http2_enabled(): + return _environment_proxy_mounts( + lambda proxy_url: AsyncHTTP2ConnectionSpread( + lambda: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True, limits=limits) + ) ) + return _environment_proxy_mounts( + lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, limits=limits) ) _create_httpx_proxy_mounts = create_httpx_proxy_mounts @@ -1500,10 +1506,12 @@ class HTTPHandler: # Get default headers (User-Agent, overridable via LITELLM_USER_AGENT) default_headers: Final = get_default_headers() if not self.disable_default_headers else None + transport: Final = self._create_sync_transport(verify=ssl_config, cert=cert, http2=use_http2) + # Create a client with a connection pool return httpx.Client( - transport=self._create_sync_transport(verify=ssl_config, cert=cert, http2=use_http2), - mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert, http2=use_http2), + transport=transport, + mounts=self._create_sync_proxy_mounts(transport, verify=ssl_config, cert=cert, http2=use_http2), timeout=self.timeout if self.timeout is not None else _DEFAULT_TIMEOUT, verify=ssl_config, cert=cert, @@ -1786,29 +1794,44 @@ class HTTPHandler: pass @staticmethod - def _create_sync_transport(verify: VerifyTypes, cert: CertTypes | None, http2: bool) -> HTTPTransport | None: + def _create_sync_transport( + verify: VerifyTypes, cert: CertTypes | None, http2: bool + ) -> HTTPTransport | HTTP2ConnectionSpread | None: """ - Create an HTTP transport with IPv4 only if litellm.force_ipv4 is True. - Otherwise, return None. + Create an HTTP transport with IPv4 only if litellm.force_ipv4 is True, and one that spreads requests over + as many HTTP/2 connections as the stream limit needs if HTTP/2 is on. + Otherwise, return litellm.sync_transport or None. Some users have seen httpx ConnectionError when using ipv6 - forcing ipv4 resolves the issue for them """ + custom_transport: Final = None if litellm.force_ipv4 else getattr(litellm, "sync_transport", None) + if custom_transport is not None: + return custom_transport + local_address: Final = _IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None + if http2: + return HTTP2ConnectionSpread( + lambda: HTTPTransport(local_address=local_address, verify=verify, cert=cert, http2=True) + ) if litellm.force_ipv4: - return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, verify=verify, cert=cert, http2=http2) - else: - return getattr(litellm, "sync_transport", None) + return HTTPTransport(local_address=local_address, verify=verify, cert=cert) + return None @staticmethod def _create_sync_proxy_mounts( + transport: httpx.BaseTransport | None, verify: VerifyTypes, cert: CertTypes | None, http2: bool, - ) -> Mapping[str, HTTPTransport | None] | None: - if not litellm.force_ipv4: + ) -> Mapping[str, httpx.BaseTransport | None] | None: + if transport is None or transport is getattr(litellm, "sync_transport", None): return None - return _environment_proxy_mounts( - lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2) - ) + if http2: + return _environment_proxy_mounts( + lambda proxy_url: HTTP2ConnectionSpread( + lambda: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True) + ) + ) + return _environment_proxy_mounts(lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert)) def get_async_httpx_client( diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index 527e43b10d3..4ff38b82837 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/tests/integration/sdk/test_http2_wire.py @@ -11,9 +11,11 @@ import ssl import threading import time from collections.abc import Callable, Generator, Iterator +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final import httpx @@ -35,6 +37,9 @@ from integration._support.ports import allocate_port from integration._support.vertex import service_account_json STREAM_CHUNKS: Final = 3 +HELD_SEARCHES: Final = 150 +HOLD_SECONDS: Final = 10.0 +_HOLD_ROOMS: Final = MappingProxyType({"hold-async": queue.SimpleQueue[int](), "hold-sync": queue.SimpleQueue[int]()}) VERTEX_SERVING_CONFIG_PATH: Final = ( "/v1/projects/p/locations/global/collections/default_collection/engines/e/servingConfigs/default_search" ) @@ -96,6 +101,13 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab if message["type"] == "http.request" and not message["more_body"]: break version: Final = scope["http_version"] + hold_room: Final = _HOLD_ROOMS.get(scope["path"].split("/")[1]) + if hold_room is not None and scope["path"].endswith(":search"): + client: Final = scope["client"] + await _send_json( + send, {"results": [{"id": await _hold_until_all_arrive(hold_room, client[1] if client else 0)}]} + ) + return if scope["path"] == "/_oauth/token": await _send_json(send, {"access_token": "scripted", "expires_in": 3600, "token_type": "Bearer"}) return @@ -119,6 +131,14 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab await _send_json(send, {"http_version": version}) +async def _hold_until_all_arrive(arrivals: queue.SimpleQueue[int], client_port: int) -> str: + arrivals.put(client_port) + deadline: Final = time.monotonic() + HOLD_SECONDS + while arrivals.qsize() < HELD_SEARCHES and time.monotonic() < deadline: + await asyncio.sleep(0.01) + return "held-together" if arrivals.qsize() >= HELD_SEARCHES else "released-alone" + + async def _send_json(send: ASGISendCallable, body: dict[str, object]) -> None: await send( HTTPResponseStartEvent(type="http.response.start", status=200, headers=[(b"content-type", b"application/json")]) @@ -454,11 +474,17 @@ def _answer_every_connection(context: ssl.SSLContext, raw: socket.socket, connec @contextmanager def _raw_tls_peer( - peer: TlsPeer, alpn: str, serve_connection: Callable[[ssl.SSLContext, socket.socket, int], None] + peer: TlsPeer, + alpn: str, + serve_connection: Callable[[ssl.SSLContext, socket.socket, int], None], + client_ca: Path | None = 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([alpn]) + if client_ca is not None: + context.verify_mode = ssl.CERT_REQUIRED + context.load_verify_locations(client_ca) stop: Final = threading.Event() accepted: Final = queue.SimpleQueue[int]() @@ -532,3 +558,73 @@ def test_vertex_search_keeps_its_http2_connection_across_an_idle_gap( 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 + + +def _held_burst(peer: TlsPeer, room: str) -> list[str]: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + + url: Final = f"{peer.url}/{room}/default_search:search" + if room == "hold-sync": + handler: Final = HTTPHandler(ssl_verify=str(peer.cert_file), http2=True) + try: + with ThreadPoolExecutor(max_workers=HELD_SEARCHES) as pool: + responses: Final = tuple(pool.map(lambda _: handler.post(url, json={}), range(HELD_SEARCHES))) + return [response.json()["results"][0]["id"] for response in responses] + finally: + handler.close() + + async def _burst() -> list[str]: + async_handler: Final = AsyncHTTPHandler(ssl_verify=str(peer.cert_file), http2=True) + try: + async_responses: Final = await asyncio.gather( + *(async_handler.post(url, json={}) for _ in range(HELD_SEARCHES)) + ) + return [response.json()["results"][0]["id"] for response in async_responses] + finally: + await async_handler.close() + + return asyncio.run(_burst()) + + +@pytest.mark.parametrize("room", ["hold-async", "hold-sync"]) +def test_http2_handlers_open_another_connection_once_the_streams_run_out( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, room: str +) -> None: + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + + search_ids: Final = _held_burst(http2_tls_peer, room) + arrivals: Final = _HOLD_ROOMS[room] + client_ports: Final = frozenset(arrivals.get_nowait() for _ in range(arrivals.qsize())) + + assert search_ids == ["held-together"] * HELD_SEARCHES + assert len(client_ports) == 2 + + +def test_forced_ipv4_handlers_honour_ssl_verify_false(monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer) -> None: + import litellm + + monkeypatch.setattr(litellm, "force_ipv4", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + for enabled in (False, True): + _set_http2(monkeypatch, enabled) + _assert_negotiated(_async_exchange(http2_tls_peer.url), enabled) + _assert_negotiated(_sync_exchange(http2_tls_peer.url), enabled) + + +def test_vertex_search_over_forced_ipv4_presents_the_client_certificate( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, tmp_path: Path +) -> None: + import litellm + + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + monkeypatch.setattr(litellm, "force_ipv4", True) + client_pem: Final = tmp_path / "client.pem" + client_pem.write_bytes(http2_tls_peer.cert_file.read_bytes() + http2_tls_peer.key_file.read_bytes()) + monkeypatch.setenv("SSL_CERTIFICATE", str(client_pem)) + + with _raw_tls_peer(http2_tls_peer, "h2", _answer_every_connection, client_ca=http2_tls_peer.cert_file) 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"] diff --git a/tests/unit/llms/custom_httpx/test_http2_connection_spread.py b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py new file mode 100644 index 00000000000..eddbc18f41a --- /dev/null +++ b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py @@ -0,0 +1,190 @@ +import asyncio +import threading +import time +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from typing import Final + +import httpx +import pytest + +from litellm.llms.custom_httpx.http2_connection_spread import ( + HTTP2_STREAMS_PER_CONNECTION, + AsyncHTTP2ConnectionSpread, + HTTP2ConnectionSpread, +) + +URL: Final = "https://search.example/v1/search" +BURST: Final = HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2 + + +class _GatedAsyncConnection(httpx.AsyncBaseTransport): + def __init__(self, gate: asyncio.Event) -> None: + self._gate: Final = gate + self.in_flight = 0 + self.peak = 0 + self.served = 0 + self.closed = False + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + await self._gate.wait() + self.in_flight -= 1 + self.served += 1 + return httpx.Response(200, stream=httpx.ByteStream(b"ok")) + + async def aclose(self) -> None: + self.closed = True + + +class _AsyncConnections: + def __init__(self, gate: asyncio.Event) -> None: + self._gate: Final = gate + self.opened: tuple[_GatedAsyncConnection, ...] = () + + def __call__(self) -> _GatedAsyncConnection: + connection: Final = _GatedAsyncConnection(self._gate) + self.opened = (*self.opened, connection) + return connection + + def in_flight(self) -> int: + return sum(connection.in_flight for connection in self.opened) + + +class _RefusedConnection(httpx.AsyncBaseTransport): + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("refused", request=request) + + +class _RefusedConnections: + def __init__(self) -> None: + self.opened = 0 + + def __call__(self) -> _RefusedConnection: + self.opened += 1 + return _RefusedConnection() + + +class _GatedConnection(httpx.BaseTransport): + def __init__(self, gate: threading.Event) -> None: + self._gate: Final = gate + self._lock: Final = threading.Lock() + self.in_flight = 0 + self.peak = 0 + self.served = 0 + + def handle_request(self, request: httpx.Request) -> httpx.Response: + with self._lock: + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + assert self._gate.wait(timeout=10) + with self._lock: + self.in_flight -= 1 + self.served += 1 + return httpx.Response(200, stream=httpx.ByteStream(b"ok")) + + +class _Connections: + def __init__(self, gate: threading.Event) -> None: + self._gate: Final = gate + self.opened: tuple[_GatedConnection, ...] = () + + def __call__(self) -> _GatedConnection: + connection: Final = _GatedConnection(self._gate) + self.opened = (*self.opened, connection) + return connection + + def in_flight(self) -> int: + return sum(connection.in_flight for connection in self.opened) + + +async def _until(condition: Callable[[], bool]) -> None: + for _ in range(1000): + if condition(): + return + await asyncio.sleep(0) + pytest.fail("condition never held") + + +def _wait_until(condition: Callable[[], bool]) -> None: + deadline: Final = time.monotonic() + 10 + while not condition(): + assert time.monotonic() < deadline, "condition never held" + time.sleep(0.01) + + +@pytest.mark.asyncio +async def test_async_spread_opens_another_connection_only_once_streams_run_out(): + gate: Final = asyncio.Event() + connections: Final = _AsyncConnections(gate) + async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections)) as client: + pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(BURST)) + await _until(lambda: connections.in_flight() == BURST) + assert [connection.peak for connection in connections.opened] == [ + HTTP2_STREAMS_PER_CONNECTION, + BURST - HTTP2_STREAMS_PER_CONNECTION, + ] + + gate.set() + assert {response.status_code for response in await asyncio.gather(*pending)} == {200} + assert (await client.get(URL)).text == "ok" + assert [connection.served for connection in connections.opened] == [ + HTTP2_STREAMS_PER_CONNECTION + 1, + BURST - HTTP2_STREAMS_PER_CONNECTION, + ] + assert all(connection.closed for connection in connections.opened) + + +@pytest.mark.asyncio +async def test_async_spread_holds_a_stream_until_its_response_is_closed(): + gate: Final = asyncio.Event() + gate.set() + connections: Final = _AsyncConnections(gate) + async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections)) as client: + open_responses: Final = tuple( + [ + await client.send(client.build_request("GET", URL), stream=True) + for _ in range(HTTP2_STREAMS_PER_CONNECTION) + ] + ) + await client.get(URL) + assert len(connections.opened) == 2 + + for response in open_responses: + await response.aclose() + await client.get(URL) + assert [connection.served for connection in connections.opened] == [HTTP2_STREAMS_PER_CONNECTION + 1, 1] + + +@pytest.mark.asyncio +async def test_async_spread_frees_the_stream_when_the_request_fails(): + connections: Final = _RefusedConnections() + async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections)) as client: + for _ in range(BURST): + with pytest.raises(httpx.ConnectError): + await client.get(URL) + assert connections.opened == 1 + + +def test_sync_spread_opens_another_connection_only_once_streams_run_out(): + gate: Final = threading.Event() + connections: Final = _Connections(gate) + with ( + httpx.Client(transport=HTTP2ConnectionSpread(connections)) as client, + ThreadPoolExecutor(max_workers=BURST) as pool, + ): + pending: Final = tuple(pool.submit(client.get, URL) for _ in range(BURST)) + _wait_until(lambda: connections.in_flight() == BURST) + assert [connection.peak for connection in connections.opened] == [ + HTTP2_STREAMS_PER_CONNECTION, + BURST - HTTP2_STREAMS_PER_CONNECTION, + ] + + gate.set() + assert {future.result().status_code for future in pending} == {200} + assert client.get(URL).text == "ok" + assert [connection.served for connection in connections.opened] == [ + HTTP2_STREAMS_PER_CONNECTION + 1, + BURST - HTTP2_STREAMS_PER_CONNECTION, + ] diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index bd4ed4c99cc..861b2bfa58d 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1917,12 +1917,12 @@ async def test_http2_flag_bypasses_aiohttp_transport(monkeypatch: pytest.MonkeyP monkeypatch.setattr(litellm, "http2", True) assert AsyncHTTPHandler.should_use_aiohttp_transport() is False - assert AsyncHTTPHandler.create_async_transport() is None + assert not isinstance(AsyncHTTPHandler.create_async_transport(), LiteLLMAiohttpTransport) monkeypatch.setattr(litellm, "http2", False) monkeypatch.setenv("LITELLM_HTTP2", "True") assert AsyncHTTPHandler.should_use_aiohttp_transport() is False - assert AsyncHTTPHandler.create_async_transport() is None + assert not isinstance(AsyncHTTPHandler.create_async_transport(), LiteLLMAiohttpTransport) @pytest.mark.asyncio From c9a8846278916b021d610cd1598c6332438dda5a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 12:01:53 -0700 Subject: [PATCH 11/14] fix(vertex_ai): reuse one Vertex search OAuth token across searches Every vector store search built a new search config, so the token cache it carries was thrown away and each search fetched a new Google OAuth token with a blocking call on the event loop. Under a burst of searches that blocking fetch held back every HTTP/2 response. The provider config is now shared, so searches with the same credentials reuse one token until it expires. --- .../search_api/transformation.py | 3 ++ litellm/utils.py | 4 +-- tests/integration/sdk/test_http2_wire.py | 33 ++++++++++++++++++- 3 files changed, 37 insertions(+), 3 deletions(-) diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 7b058bd0334..d7676d54451 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -461,3 +461,6 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): input_cost_per_query: Final = model_info.get("input_cost_per_query") or 0.0 return input_cost_per_query, 0.0 + + +vertex_search_api_vector_store_config: Final = VertexSearchAPIVectorStoreConfig() diff --git a/litellm/utils.py b/litellm/utils.py index 14ee35e571e..19c70b73e48 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10120,10 +10120,10 @@ class ProviderConfigManager: return VertexVectorStoreConfig() elif api_type == "search_api": from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( - VertexSearchAPIVectorStoreConfig, + vertex_search_api_vector_store_config, ) - return VertexSearchAPIVectorStoreConfig() + return vertex_search_api_vector_store_config elif litellm.LlmProviders.BEDROCK == provider: from litellm.llms.bedrock.vector_stores.transformation import ( BedrockVectorStoreConfig, diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index 4ff38b82837..793dd8d490b 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/tests/integration/sdk/test_http2_wire.py @@ -40,6 +40,7 @@ STREAM_CHUNKS: Final = 3 HELD_SEARCHES: Final = 150 HOLD_SECONDS: Final = 10.0 _HOLD_ROOMS: Final = MappingProxyType({"hold-async": queue.SimpleQueue[int](), "hold-sync": queue.SimpleQueue[int]()}) +_TOKEN_REQUESTS: Final = queue.SimpleQueue[str]() VERTEX_SERVING_CONFIG_PATH: Final = ( "/v1/projects/p/locations/global/collections/default_collection/engines/e/servingConfigs/default_search" ) @@ -108,7 +109,8 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab send, {"results": [{"id": await _hold_until_all_arrive(hold_room, client[1] if client else 0)}]} ) return - if scope["path"] == "/_oauth/token": + if scope["path"].endswith("/_oauth/token"): + _TOKEN_REQUESTS.put(scope["path"]) await _send_json(send, {"access_token": "scripted", "expires_in": 3600, "token_type": "Bearer"}) return if scope["path"].endswith(":search"): @@ -327,6 +329,35 @@ def test_default_async_client_stays_http1_after_vertex_search_used_http2( assert asyncio.run(_run()) == (["http/2"], "1.1") +def test_vertex_search_fetches_one_token_for_many_searches( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + import litellm + + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + credentials: Final = service_account_json("p", f"{http2_tls_peer.url}/token-reuse") + search: Final = MappingProxyType( + { + "vector_store_id": "default_search", + "query": "ping", + "custom_llm_provider": "vertex_ai/search_api", + "api_base": f"{http2_tls_peer.url}{VERTEX_SERVING_CONFIG_PATH}", + "vertex_credentials": credentials, + "vertex_project": "p", + } + ) + + async def _searches() -> None: + await litellm.vector_stores.asearch(**search) + await asyncio.gather(*(litellm.vector_stores.asearch(**search) for _ in range(5))) + + asyncio.run(_searches()) + litellm.vector_stores.search(**search) + + token_paths: Final = tuple(iter(lambda: None if _TOKEN_REQUESTS.empty() else _TOKEN_REQUESTS.get(), None)) + assert token_paths.count("/token-reuse/_oauth/token") == 1 + + def test_vertex_search_over_forced_ipv4_trusts_the_configured_ca( monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer ) -> None: From 10153423ce687e4bdf535843683660b0ec7c7801 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:24:00 -0700 Subject: [PATCH 12/14] fix(http): keep the HTTP/2 connection spread within the connection cap --- .../custom_httpx/http2_connection_spread.py | 18 +++++++++--- litellm/llms/custom_httpx/http_handler.py | 12 +++++--- tests/integration/sdk/test_http2_wire.py | 28 +++++++++++++++++++ .../test_http2_connection_spread.py | 14 ++++++++++ 4 files changed, 64 insertions(+), 8 deletions(-) diff --git a/litellm/llms/custom_httpx/http2_connection_spread.py b/litellm/llms/custom_httpx/http2_connection_spread.py index 3803ccdf56e..cabdd5011cd 100644 --- a/litellm/llms/custom_httpx/http2_connection_spread.py +++ b/litellm/llms/custom_httpx/http2_connection_spread.py @@ -24,17 +24,25 @@ class _Lane(Generic[_TransportT]): class _Lanes(Generic[_TransportT]): - def __init__(self, new_transport: Callable[[], _TransportT], streams_per_connection: int) -> None: + def __init__( + self, new_transport: Callable[[], _TransportT], streams_per_connection: int, max_connections: int | None + ) -> None: self._new_transport: Final = new_transport self._streams_per_connection: Final = streams_per_connection + self._max_connections: Final = max_connections self.lanes: tuple[_Lane[_TransportT], ...] = () def acquire(self) -> _Lane[_TransportT]: free: Final = next((lane for lane in self.lanes if lane.in_flight < self._streams_per_connection), None) - lane: Final = free if free is not None else self._open() + lane: Final = free if free is not None else self._open_or_share() lane.acquire() return lane + def _open_or_share(self) -> _Lane[_TransportT]: + if self._max_connections is not None and len(self.lanes) >= self._max_connections: + return min(self.lanes, key=lambda lane: lane.in_flight) + return self._open() + def _open(self) -> _Lane[_TransportT]: lane: Final = _Lane(self._new_transport()) self.lanes = (*self.lanes, lane) @@ -77,8 +85,9 @@ class AsyncHTTP2ConnectionSpread(httpx.AsyncBaseTransport): self, new_transport: Callable[[], httpx.AsyncBaseTransport], streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION, + max_connections: int | None = None, ) -> None: - self._lanes: Final = _Lanes(new_transport, streams_per_connection) + self._lanes: Final = _Lanes(new_transport, streams_per_connection, max_connections) async def handle_async_request(self, request: httpx.Request) -> httpx.Response: lane: Final = self._lanes.acquire() @@ -103,8 +112,9 @@ class HTTP2ConnectionSpread(httpx.BaseTransport): self, new_transport: Callable[[], httpx.BaseTransport], streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION, + max_connections: int | None = None, ) -> None: - self._lanes: Final = _Lanes(new_transport, streams_per_connection) + self._lanes: Final = _Lanes(new_transport, streams_per_connection, max_connections) self._lock: Final = threading.Lock() def _release(self, lane: _Lane[httpx.BaseTransport]) -> None: diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 04e72bc7b28..e83d9e66758 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1444,7 +1444,8 @@ class AsyncHTTPHandler: return AsyncHTTP2ConnectionSpread( lambda: AsyncHTTPTransport( local_address=local_address, verify=verify, cert=cert, http2=True, limits=limits - ) + ), + max_connections=limits.max_connections, ) if litellm.force_ipv4: return AsyncHTTPTransport(local_address=local_address, verify=verify, cert=cert, limits=limits) @@ -1463,7 +1464,8 @@ class AsyncHTTPHandler: if http2 or http2_enabled(): return _environment_proxy_mounts( lambda proxy_url: AsyncHTTP2ConnectionSpread( - lambda: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True, limits=limits) + lambda: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True, limits=limits), + max_connections=limits.max_connections, ) ) return _environment_proxy_mounts( @@ -1810,7 +1812,8 @@ class HTTPHandler: local_address: Final = _IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None if http2: return HTTP2ConnectionSpread( - lambda: HTTPTransport(local_address=local_address, verify=verify, cert=cert, http2=True) + lambda: HTTPTransport(local_address=local_address, verify=verify, cert=cert, http2=True), + max_connections=DEFAULT_LIMITS.max_connections, ) if litellm.force_ipv4: return HTTPTransport(local_address=local_address, verify=verify, cert=cert) @@ -1828,7 +1831,8 @@ class HTTPHandler: if http2: return _environment_proxy_mounts( lambda proxy_url: HTTP2ConnectionSpread( - lambda: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True) + lambda: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True), + max_connections=DEFAULT_LIMITS.max_connections, ) ) return _environment_proxy_mounts(lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert)) diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index 793dd8d490b..07b62a7d69e 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/tests/integration/sdk/test_http2_wire.py @@ -41,6 +41,7 @@ HELD_SEARCHES: Final = 150 HOLD_SECONDS: Final = 10.0 _HOLD_ROOMS: Final = MappingProxyType({"hold-async": queue.SimpleQueue[int](), "hold-sync": queue.SimpleQueue[int]()}) _TOKEN_REQUESTS: Final = queue.SimpleQueue[str]() +_COUNTED_CLIENT_PORTS: Final = queue.SimpleQueue[int]() VERTEX_SERVING_CONFIG_PATH: Final = ( "/v1/projects/p/locations/global/collections/default_collection/engines/e/servingConfigs/default_search" ) @@ -113,6 +114,11 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab _TOKEN_REQUESTS.put(scope["path"]) await _send_json(send, {"access_token": "scripted", "expires_in": 3600, "token_type": "Bearer"}) return + if scope["path"] == "/counted:search": + counted_client: Final = scope["client"] + _COUNTED_CLIENT_PORTS.put(counted_client[1] if counted_client else 0) + await _send_json(send, {"results": [{"id": "counted"}]}) + return if scope["path"].endswith(":search"): await _send_json(send, {"results": [{"id": f"http/{version}"}]}) return @@ -631,6 +637,28 @@ def test_http2_handlers_open_another_connection_once_the_streams_run_out( assert len(client_ports) == 2 +def test_http2_transport_opens_no_more_connections_than_its_cap(http2_tls_peer: TlsPeer) -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + async def _burst() -> frozenset[str]: + transport: Final = AsyncHTTPHandler.create_async_transport( + ssl_context=ssl.create_default_context(cafile=str(http2_tls_peer.cert_file)), + http2=True, + limits=httpx.Limits(max_connections=1), + ) + async with httpx.AsyncClient(transport=transport) as client: + responses: Final = await asyncio.gather( + *(client.post(f"{http2_tls_peer.url}/counted:search", json={}) for _ in range(HELD_SEARCHES)) + ) + return frozenset(response.json()["results"][0]["id"] for response in responses) + + assert asyncio.run(_burst()) == {"counted"} + client_ports: Final = frozenset( + iter(lambda: None if _COUNTED_CLIENT_PORTS.empty() else _COUNTED_CLIENT_PORTS.get(), None) + ) + assert len(client_ports) == 1 + + def test_forced_ipv4_handlers_honour_ssl_verify_false(monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer) -> None: import litellm diff --git a/tests/unit/llms/custom_httpx/test_http2_connection_spread.py b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py index eddbc18f41a..36d79e0260b 100644 --- a/tests/unit/llms/custom_httpx/test_http2_connection_spread.py +++ b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py @@ -136,6 +136,20 @@ async def test_async_spread_opens_another_connection_only_once_streams_run_out() assert all(connection.closed for connection in connections.opened) +@pytest.mark.asyncio +async def test_async_spread_shares_the_least_busy_connection_once_the_connection_cap_is_reached(): + gate: Final = asyncio.Event() + connections: Final = _AsyncConnections(gate) + burst: Final = 2 * HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2 + async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, max_connections=2)) as client: + pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(burst)) + await _until(lambda: connections.in_flight() == burst) + assert [connection.peak for connection in connections.opened] == [burst // 2, burst // 2] + + gate.set() + assert {response.status_code for response in await asyncio.gather(*pending)} == {200} + + @pytest.mark.asyncio async def test_async_spread_holds_a_stream_until_its_response_is_closed(): gate: Final = asyncio.Event() From 5e9a1820da4ce3c6f96f87d2fe2364d27d1e3da7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:55:52 -0700 Subject: [PATCH 13/14] fix(http): keep sync vertex search on HTTP/1.1, since httpcore's sync HTTP/2 is not thread-safe httpcore 1.0.9 hands two threads the same HTTP/2 stream id, so threaded sync searches failed with protocol errors. Only the async client, which the proxy uses, now defaults to HTTP/2, and the sync connection spread is removed. --- .../base_llm/vector_store/transformation.py | 2 +- .../custom_httpx/http2_connection_spread.py | 74 ++++-------------- litellm/llms/custom_httpx/http_handler.py | 37 +++------ litellm/llms/custom_httpx/llm_http_handler.py | 9 +-- .../search_api/transformation.py | 2 +- tests/integration/sdk/test_http2_wire.py | 77 ++++++++----------- .../test_http2_connection_spread.py | 67 ---------------- 7 files changed, 57 insertions(+), 211 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index d15c30a8eb5..513ffc342a7 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -124,7 +124,7 @@ class BaseVectorStoreConfig: def validate_create_vector_store(self) -> None: return None - def get_httpx_client_params(self) -> Mapping[str, object]: + def get_async_httpx_client_params(self) -> Mapping[str, object]: return MappingProxyType({}) def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]: diff --git a/litellm/llms/custom_httpx/http2_connection_spread.py b/litellm/llms/custom_httpx/http2_connection_spread.py index cabdd5011cd..1fd7fb2c4a8 100644 --- a/litellm/llms/custom_httpx/http2_connection_spread.py +++ b/litellm/llms/custom_httpx/http2_connection_spread.py @@ -1,18 +1,15 @@ -import threading -from collections.abc import AsyncIterator, Callable, Iterator -from typing import Final, Generic, TypeVar +from collections.abc import AsyncIterator, Callable +from typing import Final import httpx HTTP2_STREAMS_PER_CONNECTION: Final = 100 -_TransportT: Final = TypeVar("_TransportT", httpx.BaseTransport, httpx.AsyncBaseTransport) - -class _Lane(Generic[_TransportT]): +class _Lane: __slots__ = ("in_flight", "transport") - def __init__(self, transport: _TransportT) -> None: + def __init__(self, transport: httpx.AsyncBaseTransport) -> None: self.transport: Final = transport self.in_flight = 0 @@ -23,27 +20,30 @@ class _Lane(Generic[_TransportT]): self.in_flight -= 1 -class _Lanes(Generic[_TransportT]): +class _Lanes: def __init__( - self, new_transport: Callable[[], _TransportT], streams_per_connection: int, max_connections: int | None + self, + new_transport: Callable[[], httpx.AsyncBaseTransport], + streams_per_connection: int, + max_connections: int | None, ) -> None: self._new_transport: Final = new_transport self._streams_per_connection: Final = streams_per_connection self._max_connections: Final = max_connections - self.lanes: tuple[_Lane[_TransportT], ...] = () + self.lanes: tuple[_Lane, ...] = () - def acquire(self) -> _Lane[_TransportT]: + def acquire(self) -> _Lane: free: Final = next((lane for lane in self.lanes if lane.in_flight < self._streams_per_connection), None) lane: Final = free if free is not None else self._open_or_share() lane.acquire() return lane - def _open_or_share(self) -> _Lane[_TransportT]: + def _open_or_share(self) -> _Lane: if self._max_connections is not None and len(self.lanes) >= self._max_connections: return min(self.lanes, key=lambda lane: lane.in_flight) return self._open() - def _open(self) -> _Lane[_TransportT]: + def _open(self) -> _Lane: lane: Final = _Lane(self._new_transport()) self.lanes = (*self.lanes, lane) return lane @@ -63,20 +63,7 @@ class _AsyncReleasingStream(httpx.AsyncByteStream): await self._stream.aclose() -class _ReleasingStream(httpx.SyncByteStream): - def __init__(self, stream: httpx.SyncByteStream, release: Callable[[], None]) -> None: - self._stream: Final = stream - self._release: Final = release - - def __iter__(self) -> Iterator[bytes]: - yield from self._stream - - def close(self) -> None: - self._release() - self._stream.close() - - -def _with_stream(response: httpx.Response, stream: httpx.SyncByteStream | httpx.AsyncByteStream) -> httpx.Response: +def _with_stream(response: httpx.Response, stream: httpx.AsyncByteStream) -> httpx.Response: return httpx.Response(response.status_code, headers=response.headers, stream=stream, extensions=response.extensions) @@ -105,36 +92,3 @@ class AsyncHTTP2ConnectionSpread(httpx.AsyncBaseTransport): async def aclose(self) -> None: for lane in self._lanes.lanes: await lane.transport.aclose() - - -class HTTP2ConnectionSpread(httpx.BaseTransport): - def __init__( - self, - new_transport: Callable[[], httpx.BaseTransport], - streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION, - max_connections: int | None = None, - ) -> None: - self._lanes: Final = _Lanes(new_transport, streams_per_connection, max_connections) - self._lock: Final = threading.Lock() - - def _release(self, lane: _Lane[httpx.BaseTransport]) -> None: - with self._lock: - lane.release() - - def handle_request(self, request: httpx.Request) -> httpx.Response: - with self._lock: - lane: Final = self._lanes.acquire() - try: - response: Final = lane.transport.handle_request(request) - except BaseException: - self._release(lane) - raise - stream: Final = response.stream - if not isinstance(stream, httpx.SyncByteStream): - self._release(lane) - return response - return _with_stream(response, _ReleasingStream(stream, lambda: self._release(lane))) - - def close(self) -> None: - for lane in self._lanes.lanes: - lane.transport.close() diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index e83d9e66758..2ce6d1ea032 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -43,7 +43,7 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) -from litellm.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread, HTTP2ConnectionSpread +from litellm.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread from litellm.types.llms.custom_http import * if TYPE_CHECKING: @@ -1485,19 +1485,17 @@ class HTTPHandler: disable_default_headers: bool | None = False, # arize phoenix returns different API responses when user agent header in request follow_redirects: bool = True, - http2: bool = False, ): self.timeout = timeout self.ssl_verify = ssl_verify self.disable_default_headers = disable_default_headers self.follow_redirects = follow_redirects - self.http2 = http2 self._owns_client = client is None self._heal_lock = threading.Lock() self._client = self.create_client() if client is None else client def create_client(self) -> httpx.Client: - use_http2: Final = _resolve_http2(self.http2, self.ssl_verify) + use_http2: Final = http2_enabled() # Get unified SSL configuration ssl_config: Final = get_ssl_configuration(self.ssl_verify, http2=use_http2) @@ -1796,28 +1794,16 @@ class HTTPHandler: pass @staticmethod - def _create_sync_transport( - verify: VerifyTypes, cert: CertTypes | None, http2: bool - ) -> HTTPTransport | HTTP2ConnectionSpread | None: + def _create_sync_transport(verify: VerifyTypes, cert: CertTypes | None, http2: bool) -> HTTPTransport | None: """ - Create an HTTP transport with IPv4 only if litellm.force_ipv4 is True, and one that spreads requests over - as many HTTP/2 connections as the stream limit needs if HTTP/2 is on. + Create an HTTP transport with IPv4 only if litellm.force_ipv4 is True. Otherwise, return litellm.sync_transport or None. Some users have seen httpx ConnectionError when using ipv6 - forcing ipv4 resolves the issue for them """ - custom_transport: Final = None if litellm.force_ipv4 else getattr(litellm, "sync_transport", None) - if custom_transport is not None: - return custom_transport - local_address: Final = _IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None - if http2: - return HTTP2ConnectionSpread( - lambda: HTTPTransport(local_address=local_address, verify=verify, cert=cert, http2=True), - max_connections=DEFAULT_LIMITS.max_connections, - ) if litellm.force_ipv4: - return HTTPTransport(local_address=local_address, verify=verify, cert=cert) - return None + return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, verify=verify, cert=cert, http2=http2) + return getattr(litellm, "sync_transport", None) @staticmethod def _create_sync_proxy_mounts( @@ -1828,14 +1814,9 @@ class HTTPHandler: ) -> Mapping[str, httpx.BaseTransport | None] | None: if transport is None or transport is getattr(litellm, "sync_transport", None): return None - if http2: - return _environment_proxy_mounts( - lambda proxy_url: HTTP2ConnectionSpread( - lambda: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True), - max_connections=DEFAULT_LIMITS.max_connections, - ) - ) - return _environment_proxy_mounts(lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert)) + return _environment_proxy_mounts( + lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2) + ) def get_async_httpx_client( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 236674a4a8b..3895193d779 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -10349,7 +10349,7 @@ class BaseLLMHTTPHandler: llm_provider=litellm.LlmProviders(custom_llm_provider), params={ "ssl_verify": litellm_params.get("ssl_verify", None), - **vector_store_provider_config.get_httpx_client_params(), + **vector_store_provider_config.get_async_httpx_client_params(), }, ) else: @@ -10487,12 +10487,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = get_httpx_client( - params={ - "ssl_verify": litellm_params.get("ssl_verify", None), - **vector_store_provider_config.get_httpx_client_params(), - } - ) + sync_httpx_client = get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index d7676d54451..dfd34c706fd 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -234,7 +234,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): def __init__(self): super().__init__() - def get_httpx_client_params(self) -> Mapping[str, object]: + def get_async_httpx_client_params(self) -> Mapping[str, object]: return MappingProxyType({"http2": True}) @staticmethod diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index 07b62a7d69e..b8f2ae9b76a 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/tests/integration/sdk/test_http2_wire.py @@ -39,7 +39,7 @@ from integration._support.vertex import service_account_json STREAM_CHUNKS: Final = 3 HELD_SEARCHES: Final = 150 HOLD_SECONDS: Final = 10.0 -_HOLD_ROOMS: Final = MappingProxyType({"hold-async": queue.SimpleQueue[int](), "hold-sync": queue.SimpleQueue[int]()}) +_HELD_ARRIVALS: Final = queue.SimpleQueue[int]() _TOKEN_REQUESTS: Final = queue.SimpleQueue[str]() _COUNTED_CLIENT_PORTS: Final = queue.SimpleQueue[int]() VERTEX_SERVING_CONFIG_PATH: Final = ( @@ -103,12 +103,9 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab if message["type"] == "http.request" and not message["more_body"]: break version: Final = scope["http_version"] - hold_room: Final = _HOLD_ROOMS.get(scope["path"].split("/")[1]) - if hold_room is not None and scope["path"].endswith(":search"): + if scope["path"] == "/hold:search": client: Final = scope["client"] - await _send_json( - send, {"results": [{"id": await _hold_until_all_arrive(hold_room, client[1] if client else 0)}]} - ) + await _send_json(send, {"results": [{"id": await _hold_until_all_arrive(client[1] if client else 0)}]}) return if scope["path"].endswith("/_oauth/token"): _TOKEN_REQUESTS.put(scope["path"]) @@ -139,12 +136,12 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab await _send_json(send, {"http_version": version}) -async def _hold_until_all_arrive(arrivals: queue.SimpleQueue[int], client_port: int) -> str: - arrivals.put(client_port) +async def _hold_until_all_arrive(client_port: int) -> str: + _HELD_ARRIVALS.put(client_port) deadline: Final = time.monotonic() + HOLD_SECONDS - while arrivals.qsize() < HELD_SEARCHES and time.monotonic() < deadline: + while _HELD_ARRIVALS.qsize() < HELD_SEARCHES and time.monotonic() < deadline: await asyncio.sleep(0.01) - return "held-together" if arrivals.qsize() >= HELD_SEARCHES else "released-alone" + return "held-together" if _HELD_ARRIVALS.qsize() >= HELD_SEARCHES else "released-alone" async def _send_json(send: ASGISendCallable, body: dict[str, object]) -> None: @@ -302,11 +299,11 @@ async def _avertex_search_file_ids( return [result.get("file_id") for result in response.get("data") or ()] -def test_sync_vertex_search_negotiates_http2_without_global_flag( +def test_sync_vertex_search_stays_on_http1_without_global_flag( monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer ) -> None: _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) - assert _vertex_search_file_ids(http2_tls_peer) == ["http/2"] + assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"] def test_async_vertex_search_negotiates_http2_without_global_flag( @@ -372,7 +369,7 @@ def test_vertex_search_over_forced_ipv4_trusts_the_configured_ca( _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) monkeypatch.setattr(litellm, "force_ipv4", True) - assert _vertex_search_file_ids(http2_tls_peer) == ["http/2"] + assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"] assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/2"] @@ -553,8 +550,9 @@ def test_vertex_search_retries_a_connection_the_server_drains_with_goaway( _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) 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"] + _set_http2(monkeypatch, True) + assert _vertex_search_file_ids(http2_tls_peer, search_url=search_url) == ["answered"] def test_vertex_search_sends_once_when_an_http1_server_hangs_up( @@ -597,41 +595,25 @@ def test_vertex_search_keeps_its_http2_connection_across_an_idle_gap( assert accepted.qsize() == 1 -def _held_burst(peer: TlsPeer, room: str) -> list[str]: - from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler - - url: Final = f"{peer.url}/{room}/default_search:search" - if room == "hold-sync": - handler: Final = HTTPHandler(ssl_verify=str(peer.cert_file), http2=True) - try: - with ThreadPoolExecutor(max_workers=HELD_SEARCHES) as pool: - responses: Final = tuple(pool.map(lambda _: handler.post(url, json={}), range(HELD_SEARCHES))) - return [response.json()["results"][0]["id"] for response in responses] - finally: - handler.close() - - async def _burst() -> list[str]: - async_handler: Final = AsyncHTTPHandler(ssl_verify=str(peer.cert_file), http2=True) - try: - async_responses: Final = await asyncio.gather( - *(async_handler.post(url, json={}) for _ in range(HELD_SEARCHES)) - ) - return [response.json()["results"][0]["id"] for response in async_responses] - finally: - await async_handler.close() - - return asyncio.run(_burst()) - - -@pytest.mark.parametrize("room", ["hold-async", "hold-sync"]) -def test_http2_handlers_open_another_connection_once_the_streams_run_out( - monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, room: str +def test_async_http2_handler_opens_another_connection_once_the_streams_run_out( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer ) -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) - search_ids: Final = _held_burst(http2_tls_peer, room) - arrivals: Final = _HOLD_ROOMS[room] - client_ports: Final = frozenset(arrivals.get_nowait() for _ in range(arrivals.qsize())) + async def _burst() -> list[str]: + handler: Final = AsyncHTTPHandler(ssl_verify=str(http2_tls_peer.cert_file), http2=True) + try: + responses: Final = await asyncio.gather( + *(handler.post(f"{http2_tls_peer.url}/hold:search", json={}) for _ in range(HELD_SEARCHES)) + ) + return [response.json()["results"][0]["id"] for response in responses] + finally: + await handler.close() + + search_ids: Final = asyncio.run(_burst()) + client_ports: Final = frozenset(_HELD_ARRIVALS.get_nowait() for _ in range(_HELD_ARRIVALS.qsize())) assert search_ids == ["held-together"] * HELD_SEARCHES assert len(client_ports) == 2 @@ -685,5 +667,6 @@ def test_vertex_search_over_forced_ipv4_presents_the_client_certificate( 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"] + _set_http2(monkeypatch, True) + assert _vertex_search_file_ids(http2_tls_peer, search_url=search_url) == ["answered"] diff --git a/tests/unit/llms/custom_httpx/test_http2_connection_spread.py b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py index 36d79e0260b..d637631110f 100644 --- a/tests/unit/llms/custom_httpx/test_http2_connection_spread.py +++ b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py @@ -1,8 +1,5 @@ import asyncio -import threading -import time from collections.abc import Callable -from concurrent.futures import ThreadPoolExecutor from typing import Final import httpx @@ -11,7 +8,6 @@ import pytest from litellm.llms.custom_httpx.http2_connection_spread import ( HTTP2_STREAMS_PER_CONNECTION, AsyncHTTP2ConnectionSpread, - HTTP2ConnectionSpread, ) URL: Final = "https://search.example/v1/search" @@ -66,39 +62,6 @@ class _RefusedConnections: return _RefusedConnection() -class _GatedConnection(httpx.BaseTransport): - def __init__(self, gate: threading.Event) -> None: - self._gate: Final = gate - self._lock: Final = threading.Lock() - self.in_flight = 0 - self.peak = 0 - self.served = 0 - - def handle_request(self, request: httpx.Request) -> httpx.Response: - with self._lock: - self.in_flight += 1 - self.peak = max(self.peak, self.in_flight) - assert self._gate.wait(timeout=10) - with self._lock: - self.in_flight -= 1 - self.served += 1 - return httpx.Response(200, stream=httpx.ByteStream(b"ok")) - - -class _Connections: - def __init__(self, gate: threading.Event) -> None: - self._gate: Final = gate - self.opened: tuple[_GatedConnection, ...] = () - - def __call__(self) -> _GatedConnection: - connection: Final = _GatedConnection(self._gate) - self.opened = (*self.opened, connection) - return connection - - def in_flight(self) -> int: - return sum(connection.in_flight for connection in self.opened) - - async def _until(condition: Callable[[], bool]) -> None: for _ in range(1000): if condition(): @@ -107,13 +70,6 @@ async def _until(condition: Callable[[], bool]) -> None: pytest.fail("condition never held") -def _wait_until(condition: Callable[[], bool]) -> None: - deadline: Final = time.monotonic() + 10 - while not condition(): - assert time.monotonic() < deadline, "condition never held" - time.sleep(0.01) - - @pytest.mark.asyncio async def test_async_spread_opens_another_connection_only_once_streams_run_out(): gate: Final = asyncio.Event() @@ -179,26 +135,3 @@ async def test_async_spread_frees_the_stream_when_the_request_fails(): with pytest.raises(httpx.ConnectError): await client.get(URL) assert connections.opened == 1 - - -def test_sync_spread_opens_another_connection_only_once_streams_run_out(): - gate: Final = threading.Event() - connections: Final = _Connections(gate) - with ( - httpx.Client(transport=HTTP2ConnectionSpread(connections)) as client, - ThreadPoolExecutor(max_workers=BURST) as pool, - ): - pending: Final = tuple(pool.submit(client.get, URL) for _ in range(BURST)) - _wait_until(lambda: connections.in_flight() == BURST) - assert [connection.peak for connection in connections.opened] == [ - HTTP2_STREAMS_PER_CONNECTION, - BURST - HTTP2_STREAMS_PER_CONNECTION, - ] - - gate.set() - assert {future.result().status_code for future in pending} == {200} - assert client.get(URL).text == "ok" - assert [connection.served for connection in connections.opened] == [ - HTTP2_STREAMS_PER_CONNECTION + 1, - BURST - HTTP2_STREAMS_PER_CONNECTION, - ] From 40810f2a76631ed1c4975738d657b89344035537 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:34:58 -0700 Subject: [PATCH 14/14] fix(http): split the connection cap across HTTP/2 lanes Each lane had its own pool with the full connection cap, so an origin answering HTTP/1.1, or several origins on one client, could exceed max_connections once more than 100 requests were in flight. Lanes now open only while lanes times 100 streams stay under the cap, and each lane pool gets at most 100 of it. --- .../custom_httpx/http2_connection_spread.py | 27 +++++--- litellm/llms/custom_httpx/http_handler.py | 12 ++-- .../test_http2_connection_spread.py | 61 +++++++++++++++++-- 3 files changed, 81 insertions(+), 19 deletions(-) diff --git a/litellm/llms/custom_httpx/http2_connection_spread.py b/litellm/llms/custom_httpx/http2_connection_spread.py index 1fd7fb2c4a8..0fc1d8c38b9 100644 --- a/litellm/llms/custom_httpx/http2_connection_spread.py +++ b/litellm/llms/custom_httpx/http2_connection_spread.py @@ -23,13 +23,13 @@ class _Lane: class _Lanes: def __init__( self, - new_transport: Callable[[], httpx.AsyncBaseTransport], + new_transport: Callable[[httpx.Limits], httpx.AsyncBaseTransport], + limits: httpx.Limits, streams_per_connection: int, - max_connections: int | None, ) -> None: self._new_transport: Final = new_transport + self._limits: Final = limits self._streams_per_connection: Final = streams_per_connection - self._max_connections: Final = max_connections self.lanes: tuple[_Lane, ...] = () def acquire(self) -> _Lane: @@ -39,15 +39,26 @@ class _Lanes: return lane def _open_or_share(self) -> _Lane: - if self._max_connections is not None and len(self.lanes) >= self._max_connections: + cap: Final = self._limits.max_connections + if cap is not None and len(self.lanes) * self._streams_per_connection >= cap: return min(self.lanes, key=lambda lane: lane.in_flight) return self._open() def _open(self) -> _Lane: - lane: Final = _Lane(self._new_transport()) + lane: Final = _Lane(self._new_transport(self._next_lane_limits())) self.lanes = (*self.lanes, lane) return lane + def _next_lane_limits(self) -> httpx.Limits: + cap: Final = self._limits.max_connections + if cap is None: + return self._limits + return httpx.Limits( + max_connections=min(self._streams_per_connection, cap - len(self.lanes) * self._streams_per_connection), + max_keepalive_connections=self._limits.max_keepalive_connections, + keepalive_expiry=self._limits.keepalive_expiry, + ) + class _AsyncReleasingStream(httpx.AsyncByteStream): def __init__(self, stream: httpx.AsyncByteStream, release: Callable[[], None]) -> None: @@ -70,11 +81,11 @@ def _with_stream(response: httpx.Response, stream: httpx.AsyncByteStream) -> htt class AsyncHTTP2ConnectionSpread(httpx.AsyncBaseTransport): def __init__( self, - new_transport: Callable[[], httpx.AsyncBaseTransport], + new_transport: Callable[[httpx.Limits], httpx.AsyncBaseTransport], + limits: httpx.Limits, streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION, - max_connections: int | None = None, ) -> None: - self._lanes: Final = _Lanes(new_transport, streams_per_connection, max_connections) + self._lanes: Final = _Lanes(new_transport, limits, streams_per_connection) async def handle_async_request(self, request: httpx.Request) -> httpx.Response: lane: Final = self._lanes.acquire() diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 2ce6d1ea032..e00902d9047 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1442,10 +1442,10 @@ class AsyncHTTPHandler: local_address: Final = _IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None if http2 or http2_enabled(): return AsyncHTTP2ConnectionSpread( - lambda: AsyncHTTPTransport( - local_address=local_address, verify=verify, cert=cert, http2=True, limits=limits + lambda lane_limits: AsyncHTTPTransport( + local_address=local_address, verify=verify, cert=cert, http2=True, limits=lane_limits ), - max_connections=limits.max_connections, + limits, ) if litellm.force_ipv4: return AsyncHTTPTransport(local_address=local_address, verify=verify, cert=cert, limits=limits) @@ -1464,8 +1464,10 @@ class AsyncHTTPHandler: if http2 or http2_enabled(): return _environment_proxy_mounts( lambda proxy_url: AsyncHTTP2ConnectionSpread( - lambda: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True, limits=limits), - max_connections=limits.max_connections, + lambda lane_limits: AsyncHTTPTransport( + proxy=proxy_url, verify=verify, cert=cert, http2=True, limits=lane_limits + ), + limits, ) ) return _environment_proxy_mounts( diff --git a/tests/unit/llms/custom_httpx/test_http2_connection_spread.py b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py index d637631110f..a83c013c6a6 100644 --- a/tests/unit/llms/custom_httpx/test_http2_connection_spread.py +++ b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py @@ -12,6 +12,7 @@ from litellm.llms.custom_httpx.http2_connection_spread import ( URL: Final = "https://search.example/v1/search" BURST: Final = HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2 +UNCAPPED: Final = httpx.Limits(max_connections=None) class _GatedAsyncConnection(httpx.AsyncBaseTransport): @@ -39,7 +40,7 @@ class _AsyncConnections: self._gate: Final = gate self.opened: tuple[_GatedAsyncConnection, ...] = () - def __call__(self) -> _GatedAsyncConnection: + def __call__(self, _limits: httpx.Limits) -> _GatedAsyncConnection: connection: Final = _GatedAsyncConnection(self._gate) self.opened = (*self.opened, connection) return connection @@ -57,11 +58,39 @@ class _RefusedConnections: def __init__(self) -> None: self.opened = 0 - def __call__(self) -> _RefusedConnection: + def __call__(self, _limits: httpx.Limits) -> _RefusedConnection: self.opened += 1 return _RefusedConnection() +class _HTTP1Connections: + def __init__(self, gate: asyncio.Event) -> None: + self._gate: Final = gate + self.requested = 0 + self.open = 0 + + def __call__(self, limits: httpx.Limits) -> "_HTTP1Pool": + assert limits.max_connections is not None + return _HTTP1Pool(self, asyncio.Semaphore(limits.max_connections)) + + async def serve(self, pool_slots: asyncio.Semaphore) -> None: + self.requested += 1 + async with pool_slots: + self.open += 1 + await self._gate.wait() + self.open -= 1 + + +class _HTTP1Pool(httpx.AsyncBaseTransport): + def __init__(self, connections: _HTTP1Connections, slots: asyncio.Semaphore) -> None: + self._connections: Final = connections + self._slots: Final = slots + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + await self._connections.serve(self._slots) + return httpx.Response(200, stream=httpx.ByteStream(b"ok")) + + async def _until(condition: Callable[[], bool]) -> None: for _ in range(1000): if condition(): @@ -74,7 +103,7 @@ async def _until(condition: Callable[[], bool]) -> None: async def test_async_spread_opens_another_connection_only_once_streams_run_out(): gate: Final = asyncio.Event() connections: Final = _AsyncConnections(gate) - async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections)) as client: + async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, UNCAPPED)) as client: pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(BURST)) await _until(lambda: connections.in_flight() == BURST) assert [connection.peak for connection in connections.opened] == [ @@ -97,7 +126,10 @@ async def test_async_spread_shares_the_least_busy_connection_once_the_connection gate: Final = asyncio.Event() connections: Final = _AsyncConnections(gate) burst: Final = 2 * HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2 - async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, max_connections=2)) as client: + two_connections_of_streams: Final = httpx.Limits(max_connections=2 * HTTP2_STREAMS_PER_CONNECTION) + async with httpx.AsyncClient( + transport=AsyncHTTP2ConnectionSpread(connections, two_connections_of_streams) + ) as client: pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(burst)) await _until(lambda: connections.in_flight() == burst) assert [connection.peak for connection in connections.opened] == [burst // 2, burst // 2] @@ -106,12 +138,29 @@ async def test_async_spread_shares_the_least_busy_connection_once_the_connection assert {response.status_code for response in await asyncio.gather(*pending)} == {200} +@pytest.mark.asyncio +async def test_async_spread_keeps_http1_connections_within_the_connection_cap(): + gate: Final = asyncio.Event() + connections: Final = _HTTP1Connections(gate) + cap: Final = HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2 + burst: Final = 2 * cap + async with httpx.AsyncClient( + transport=AsyncHTTP2ConnectionSpread(connections, httpx.Limits(max_connections=cap)) + ) as client: + pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(burst)) + await _until(lambda: connections.requested == burst) + assert connections.open == cap + + gate.set() + assert {response.status_code for response in await asyncio.gather(*pending)} == {200} + + @pytest.mark.asyncio async def test_async_spread_holds_a_stream_until_its_response_is_closed(): gate: Final = asyncio.Event() gate.set() connections: Final = _AsyncConnections(gate) - async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections)) as client: + async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, UNCAPPED)) as client: open_responses: Final = tuple( [ await client.send(client.build_request("GET", URL), stream=True) @@ -130,7 +179,7 @@ async def test_async_spread_holds_a_stream_until_its_response_is_closed(): @pytest.mark.asyncio async def test_async_spread_frees_the_stream_when_the_request_fails(): connections: Final = _RefusedConnections() - async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections)) as client: + async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, UNCAPPED)) as client: for _ in range(BURST): with pytest.raises(httpx.ConnectError): await client.get(URL)