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"]