import asyncio import gc import io import os import pathlib import ssl import threading import weakref from collections.abc import Callable, Mapping from typing import Final from unittest.mock import MagicMock, patch import certifi import httpx import pytest from aiohttp import ClientSession, TCPConnector import litellm from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport from litellm.llms.custom_httpx.http_handler import ( _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER, AsyncHTTPHandler, HTTPHandler, MaskedHTTPStatusError, _get_httpx_client, get_ssl_configuration, ) from litellm.types.llms.custom_http import VerifyTypes @pytest.mark.asyncio async def test_async_post_streaming_status_error_should_not_wait_forever_for_body( monkeypatch, ): """ Vertex Anthropic streamRawPredict can return a pre-stream 4xx where the streamed error body never terminates. The handler must still surface the status promptly instead of blocking the downstream client. """ class HangingErrorStream(httpx.AsyncByteStream): async def __aiter__(self): await asyncio.Event().wait() if False: yield b"" async def mock_handler(request: httpx.Request) -> httpx.Response: return httpx.Response( 400, request=request, headers={"content-type": "application/json"}, stream=HangingErrorStream(), ) monkeypatch.setattr( "litellm.llms.custom_httpx.http_handler._STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS", 0.01, ) litellm_handler = AsyncHTTPHandler() await litellm_handler.client.aclose() litellm_handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock_handler)) try: with pytest.raises(MaskedHTTPStatusError) as exc_info: await asyncio.wait_for( litellm_handler.post( "https://vertex.example/streamRawPredict", stream=True, ), timeout=0.2, ) assert exc_info.value.status_code == 400 assert exc_info.value.response.status_code == 400 finally: await litellm_handler.close() def test_sync_post_streaming_status_error_should_not_wait_forever_for_body( monkeypatch, ): """ Keep the sync streaming error path aligned with the async path so a non-terminating streamed error body cannot block a worker thread forever. """ class HangingSyncErrorStream(httpx.SyncByteStream): def __init__(self): self.closed_event = threading.Event() def __iter__(self): self.closed_event.wait() if False: yield b"" def close(self): self.closed_event.set() def mock_handler(request: httpx.Request) -> httpx.Response: return httpx.Response( 400, request=request, headers={"content-type": "application/json"}, stream=HangingSyncErrorStream(), ) monkeypatch.setattr( "litellm.llms.custom_httpx.http_handler._STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS", 0.01, ) litellm_handler = HTTPHandler() litellm_handler.client.close() litellm_handler.client = httpx.Client(transport=httpx.MockTransport(mock_handler)) try: with pytest.raises(MaskedHTTPStatusError) as exc_info: litellm_handler.post( "https://vertex.example/streamRawPredict", stream=True, ) assert exc_info.value.status_code == 400 assert exc_info.value.response.status_code == 400 finally: litellm_handler.close() @pytest.mark.asyncio async def test_ssl_security_level(monkeypatch): # Ensure aiohttp transport is enabled for this test monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) with patch.dict(os.environ, clear=True): # Set environment variable for SSL security level monkeypatch.setenv("SSL_SECURITY_LEVEL", "DEFAULT@SECLEVEL=1") # Create async client with SSL verification disabled to isolate SSL context testing client = AsyncHTTPHandler() try: # Get the transport (should be LiteLLMAiohttpTransport) transport = client.client._transport assert isinstance(transport, LiteLLMAiohttpTransport) # Get the aiohttp ClientSession client_session = transport._get_valid_client_session() # Get the connector from the session connector = client_session.connector assert isinstance(connector, TCPConnector) # Get the SSL context from the connector ssl_context = connector._ssl # Verify that the SSL context exists and has the correct cipher string assert isinstance(ssl_context, ssl.SSLContext) finally: await client.close() @pytest.mark.asyncio async def test_force_ipv4_transport(monkeypatch: pytest.MonkeyPatch): """Test transport creation with force_ipv4 enabled""" monkeypatch.setattr(litellm, "force_ipv4", True) monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) transport = AsyncHTTPHandler._create_async_transport() # Should get an AsyncHTTPTransport (no real HTTP call — avoids CI hangs) assert isinstance(transport, httpx.AsyncHTTPTransport) @pytest.mark.asyncio async def test_aiohttp_disabled_transport(monkeypatch: pytest.MonkeyPatch): """Test transport creation with aiohttp disabled""" monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setattr(litellm, "force_ipv4", False) transport = AsyncHTTPHandler._create_async_transport() # Should get None when both aiohttp is disabled and force_ipv4 is False assert transport is None @pytest.mark.asyncio async def test_ssl_verification_with_aiohttp_transport(monkeypatch: pytest.MonkeyPatch): """ Test aiohttp respects ssl_verify=False We validate that the ssl settings for a litellm transport match what a ssl verify=False aiohttp client would have. """ import aiohttp # Ensure aiohttp transport is enabled for this test monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) litellm_async_client = AsyncHTTPHandler(ssl_verify=False) try: transport = litellm_async_client.client._transport assert isinstance(transport, LiteLLMAiohttpTransport) transport_connector = transport._get_valid_client_session().connector assert isinstance(transport_connector, TCPConnector) aiohttp_session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False)) try: aiohttp_connector = aiohttp_session.connector assert isinstance(aiohttp_connector, aiohttp.TCPConnector) # assert both litellm transport and aiohttp session have ssl_verify=False assert transport_connector._ssl == aiohttp_connector._ssl finally: await aiohttp_session.close() finally: await litellm_async_client.close() @pytest.mark.asyncio async def test_ssl_verification_with_shared_session(monkeypatch: pytest.MonkeyPatch): """ Test that ssl_verify=False is respected even with shared sessions. This was a bug where shared sessions bypassed SSL configuration because _create_aiohttp_transport returned immediately without passing ssl_verify to the LiteLLMAiohttpTransport constructor. The fix stores ssl_verify in the transport and passes it per-request. """ import aiohttp # Ensure aiohttp transport is enabled for this test monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) shared_session = aiohttp.ClientSession() try: # Create transport with shared session and ssl_verify=False transport = AsyncHTTPHandler._create_aiohttp_transport( ssl_verify=False, shared_session=shared_session, ) # Verify the transport uses the shared session assert transport.client is shared_session # Verify the SSL setting is stored in the transport for per-request use assert transport._ssl_verify is False finally: await shared_session.close() @pytest.mark.asyncio async def test_ssl_context_with_shared_session(monkeypatch: pytest.MonkeyPatch): """ Test that ssl_context is respected even with shared sessions. """ import aiohttp # Ensure aiohttp transport is enabled for this test monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) custom_ssl_context = ssl.create_default_context() # Create a shared session shared_session = aiohttp.ClientSession() try: # Create transport with shared session and custom ssl_context transport = AsyncHTTPHandler._create_aiohttp_transport( ssl_context=custom_ssl_context, shared_session=shared_session, ) # Verify the transport uses the shared session assert transport.client is shared_session # Verify the SSL context is stored in the transport for per-request use assert transport._ssl_verify is custom_ssl_context finally: await shared_session.close() def test_get_ssl_configuration(): """Test that get_ssl_configuration() returns a proper SSL context with certifi CA bundle when no environment variables are set.""" from litellm.llms.custom_httpx.http_handler import _ssl_context_cache # Clear cache to ensure ssl.create_default_context is called _ssl_context_cache.clear() with patch.dict(os.environ, clear=True): with patch("ssl.create_default_context") as mock_create_context: # Mock the return value mock_ssl_context = MagicMock(spec=ssl.SSLContext) mock_ssl_context.set_ciphers = MagicMock() mock_ssl_context.minimum_version = ssl.TLSVersion.TLSv1_2 mock_create_context.return_value = mock_ssl_context # Call the static method result = get_ssl_configuration() # Verify ssl.create_default_context was called with certifi's CA file expected_ca_file = certifi.where() mock_create_context.assert_called_once_with(cafile=expected_ca_file) # Verify it returns the mocked SSL context assert result == mock_ssl_context def test_get_ssl_configuration_integration(): """Integration test that _get_ssl_context() returns a working SSL context""" # Call the static method without mocking ssl_context = get_ssl_configuration() # Verify it returns an SSLContext instance assert isinstance(ssl_context, ssl.SSLContext) # Verify it has basic SSL context properties assert ssl_context.protocol is not None assert ssl_context.verify_mode is not None # Session Reuse Tests class MockClientSession: """Mock ClientSession that is not callable""" def __init__(self): self.closed = False @pytest.mark.asyncio async def test_create_aiohttp_transport_with_shared_session(): """Test that _create_aiohttp_transport reuses shared session when provided""" from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler # Create a mock shared session that's not callable mock_session = MockClientSession() # Test with shared session transport = AsyncHTTPHandler._create_aiohttp_transport( shared_session=mock_session # type: ignore ) # Verify the transport uses the shared session directly assert transport.client is mock_session assert not callable(transport.client) # Should not be callable @pytest.mark.asyncio async def test_async_handler_with_shared_session(): """Test AsyncHTTPHandler initialization with shared session""" from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler # Create a mock shared session mock_session = MockClientSession() # Create handler with shared session handler = AsyncHTTPHandler(shared_session=mock_session) # type: ignore # Verify the handler was created successfully assert handler is not None assert handler.client is not None @pytest.mark.asyncio async def test_get_async_httpx_client_with_shared_session(): """Test get_async_httpx_client with shared session""" from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, AsyncHTTPHandler as AsyncHTTPHandlerReload, ) from litellm.types.utils import LlmProviders # Create a mock shared session mock_session = MockClientSession() # Test with shared session client = get_async_httpx_client( llm_provider=LlmProviders.ANTHROPIC, shared_session=mock_session, # type: ignore ) # Verify the client was created successfully assert client is not None # Import locally to avoid stale reference after module reload in conftest assert isinstance(client, AsyncHTTPHandlerReload) @pytest.mark.asyncio async def test_get_async_httpx_client_without_shared_session(): """Test get_async_httpx_client without shared session (backward compatibility)""" from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, AsyncHTTPHandler as AsyncHTTPHandlerReload, ) from litellm.types.utils import LlmProviders # Test without shared session client = get_async_httpx_client(llm_provider=LlmProviders.ANTHROPIC, shared_session=None) # Verify the client was created successfully assert client is not None # Import locally to avoid stale reference after module reload in conftest assert isinstance(client, AsyncHTTPHandlerReload) @pytest.mark.asyncio async def test_session_reuse_chain(): """Test that session is properly passed through the entire call chain""" from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler # Create a mock shared session mock_session = MockClientSession() # Test the entire chain transport = AsyncHTTPHandler._create_async_transport( shared_session=mock_session # type: ignore ) # Verify the transport was created assert transport is not None # Test AsyncHTTPHandler creation handler = AsyncHTTPHandler(shared_session=mock_session) # type: ignore assert handler is not None def test_shared_session_parameter_in_acompletion(): """Test that acompletion function accepts shared_session parameter""" import inspect from litellm.main import acompletion # Get the function signature sig = inspect.signature(acompletion) params = list(sig.parameters.keys()) # Verify shared_session parameter exists assert "shared_session" in params # Verify the parameter type annotation shared_session_param = sig.parameters["shared_session"] assert "ClientSession" in str(shared_session_param.annotation) def test_shared_session_parameter_in_completion(): """Test that completion function accepts shared_session parameter""" import inspect from litellm.main import completion # Get the function signature sig = inspect.signature(completion) params = list(sig.parameters.keys()) # Verify shared_session parameter exists assert "shared_session" in params # Verify the parameter type annotation shared_session_param = sig.parameters["shared_session"] assert "ClientSession" in str(shared_session_param.annotation) @pytest.mark.asyncio async def test_session_reuse_integration(): """Integration test for session reuse functionality""" from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, AsyncHTTPHandler as AsyncHTTPHandlerReload, ) from litellm.types.utils import LlmProviders # Create a mock session mock_session = MockClientSession() # Create two clients with the same session client1 = get_async_httpx_client( llm_provider=LlmProviders.ANTHROPIC, shared_session=mock_session, # type: ignore ) client2 = get_async_httpx_client( llm_provider=LlmProviders.OPENAI, shared_session=mock_session, # type: ignore ) # Both clients should be created successfully assert client1 is not None assert client2 is not None # Both should be AsyncHTTPHandler instances # Import locally to avoid stale reference after module reload in conftest assert isinstance(client1, AsyncHTTPHandlerReload) assert isinstance(client2, AsyncHTTPHandlerReload) # Clean up await client1.close() await client2.close() @pytest.mark.parametrize( "env_curve,litellm_curve,expected_curve,should_call", [ # env_curve: SSL_ECDH_CURVE env var | litellm_curve: litellm.ssl_ecdh_curve variable # expected_curve: curve that should be set | should_call: whether set_ecdh_curve() should be called # Valid configurations ("X25519", None, "X25519", True), # Env var only ("prime256v1", None, "prime256v1", True), # Different valid curve (None, "secp384r1", "secp384r1", True), # litellm variable only ("X25519", "secp521r1", "X25519", True), # Env var takes precedence # Empty/None configurations - should skip ("", None, None, False), # Empty string - skip configuration (None, None, None, False), # None value - skip configuration ], ) def test_ssl_ecdh_curve(env_curve, litellm_curve, expected_curve, should_call, monkeypatch): """Test SSL ECDH curve configuration with valid curves and precedence""" from litellm.llms.custom_httpx.http_handler import _ssl_context_cache # Clear cache to ensure fresh SSL context creation _ssl_context_cache.clear() with patch.dict(os.environ, clear=True): if env_curve: monkeypatch.setenv("SSL_ECDH_CURVE", env_curve) monkeypatch.setattr(litellm, "ssl_ecdh_curve", litellm_curve) # Create a real SSL context and patch set_ecdh_curve on it # We need a real SSLContext instance (not a MagicMock) because _create_ssl_context # calls methods like set_ciphers() and minimum_version that require a real context. # We patch set_ecdh_curve specifically to verify it's called with the correct curve. real_ssl_context = ssl.create_default_context() with patch("ssl.create_default_context", return_value=real_ssl_context): with patch.object(real_ssl_context, "set_ecdh_curve") as mock_set_curve: ssl_context = get_ssl_configuration() if should_call: mock_set_curve.assert_called_once_with(expected_curve) else: mock_set_curve.assert_not_called() assert isinstance(ssl_context, ssl.SSLContext) def test_default_user_agent_is_litellm_version(monkeypatch): from litellm._version import version from litellm.llms.custom_httpx.http_handler import get_default_headers monkeypatch.delenv("LITELLM_USER_AGENT", raising=False) assert get_default_headers()["User-Agent"] == f"litellm/{version}" def test_user_agent_can_be_overridden_via_env_var(monkeypatch): from litellm.llms.custom_httpx.http_handler import get_default_headers monkeypatch.setenv("LITELLM_USER_AGENT", "Claude Code") assert get_default_headers()["User-Agent"] == "Claude Code" def test_user_agent_env_var_can_be_empty_string(monkeypatch): from litellm.llms.custom_httpx.http_handler import get_default_headers monkeypatch.setenv("LITELLM_USER_AGENT", "") assert get_default_headers()["User-Agent"] == "" def test_user_agent_override_is_not_appended_to_default(monkeypatch): from litellm.llms.custom_httpx.http_handler import HTTPHandler monkeypatch.delenv("LITELLM_USER_AGENT", raising=False) handler = HTTPHandler() try: req = handler.client.build_request( "GET", "https://example.com", headers={"user-agent": "Claude Code"}, ) assert req.headers.get_list("User-Agent") == ["Claude Code"] finally: handler.close() def test_sync_http_handler_uses_env_user_agent(monkeypatch): from litellm.llms.custom_httpx.http_handler import HTTPHandler monkeypatch.setenv("LITELLM_USER_AGENT", "Claude Code") handler = HTTPHandler() try: req = handler.client.build_request("GET", "https://example.com") assert req.headers.get("User-Agent") == "Claude Code" finally: handler.close() @pytest.mark.asyncio async def test_async_http_handler_uses_env_user_agent(monkeypatch): from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler monkeypatch.setenv("LITELLM_USER_AGENT", "Claude Code") handler = AsyncHTTPHandler() try: req = handler.client.build_request("GET", "https://example.com") assert req.headers.get("User-Agent") == "Claude Code" finally: await handler.close() @pytest.mark.asyncio async def test_httpx_handler_uses_env_user_agent(monkeypatch): from litellm.llms.custom_httpx.httpx_handler import HTTPHandler monkeypatch.setenv("LITELLM_USER_AGENT", "Claude Code") handler = HTTPHandler() try: req = handler.client.build_request("GET", "https://example.com") assert req.headers.get("User-Agent") == "Claude Code" finally: await handler.close() def test_get_httpx_client_applies_float_timeout_without_mocking_handler(): """ Exercise real _get_httpx_client + HTTPHandler: params={'timeout': x} must reach httpx.Client(timeout=...). Uses an uncommon timeout value to avoid colliding with other cached clients in-process. """ timeout = 3847.291 handler = _get_httpx_client(params={"timeout": timeout}) try: assert isinstance(handler, HTTPHandler) assert handler.client.timeout == httpx.Timeout(timeout) finally: handler.close() def test_get_httpx_client_applies_httpx_timeout_object_without_mocking_handler(): t = httpx.Timeout(40.0, connect=5.0) handler = _get_httpx_client(params={"timeout": t}) try: assert handler.client.timeout == t finally: handler.close() 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.""" captured = {} def mock_handler(request: httpx.Request) -> httpx.Response: captured["timeout"] = request.extensions.get("timeout") return httpx.Response(200, request=request, json={"ok": True}) handler = HTTPHandler() handler.client.close() handler.client = httpx.Client( transport=httpx.MockTransport(mock_handler), timeout=httpx.Timeout(5.0), ) try: handler.get("https://example.com/poll", timeout=99.0) assert captured["timeout"] == { "connect": 99.0, "read": 99.0, "write": 99.0, "pool": 99.0, } finally: handler.close() @pytest.mark.asyncio async def test_async_get_forwards_per_request_timeout(): captured = {} async def mock_handler(request: httpx.Request) -> httpx.Response: captured["timeout"] = request.extensions.get("timeout") return httpx.Response(200, request=request, json={"ok": True}) handler = AsyncHTTPHandler() await handler.client.aclose() handler.client = httpx.AsyncClient( transport=httpx.MockTransport(mock_handler), timeout=httpx.Timeout(5.0), ) try: await handler.get("https://example.com/poll", timeout=99.0) assert captured["timeout"] == { "connect": 99.0, "read": 99.0, "write": 99.0, "pool": 99.0, } finally: await handler.close() class TestDefaultCachedClientTimeoutHonorsRequestTimeout: """Cached default httpx clients must fall back to an explicit litellm.request_timeout. Regression for LIT-2369: get_async_httpx_client / _get_httpx_client hardcoded a 600s default and never consulted litellm.request_timeout, so provider calls with no per-model timeout (e.g. Bedrock) hung for 600s. """ def test_default_when_request_timeout_unset(self, monkeypatch: pytest.MonkeyPatch): from litellm.llms.custom_httpx.http_handler import ( _DEFAULT_TIMEOUT, _default_cached_client_timeout, ) monkeypatch.setattr(litellm, "request_timeout", litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS) monkeypatch.setattr(litellm, "request_timeout_explicitly_set", False) assert _default_cached_client_timeout() is _DEFAULT_TIMEOUT def test_uses_explicit_request_timeout(self, monkeypatch: pytest.MonkeyPatch): from litellm.llms.custom_httpx.http_handler import ( _default_cached_client_timeout, ) monkeypatch.setattr(litellm, "request_timeout", 300) monkeypatch.setattr(litellm, "request_timeout_explicitly_set", True) resolved = _default_cached_client_timeout() assert resolved.read == 300.0 assert resolved.connect == 5.0 def test_cached_async_client_built_with_explicit_request_timeout(self, monkeypatch: pytest.MonkeyPatch): from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.utils import LlmProviders monkeypatch.setattr(litellm, "request_timeout", 300) monkeypatch.setattr(litellm, "request_timeout_explicitly_set", True) litellm.in_memory_llm_clients_cache = LLMClientCache() client = get_async_httpx_client(llm_provider=LlmProviders.BEDROCK) assert client.timeout.read == 300.0 async def _read_http_request(reader: asyncio.StreamReader) -> None: raw = b"" while b"\r\n\r\n" not in raw: chunk = await reader.read(1024) if not chunk: return raw += chunk head, _, body = raw.partition(b"\r\n\r\n") content_length = next( (int(line.split(b":", 1)[1]) for line in head.split(b"\r\n") if line.lower().startswith(b"content-length")), 0, ) while len(body) < content_length: body += await reader.read(content_length - len(body)) @pytest.mark.asyncio async def test_init_held_async_handler_survives_external_client_close(): handler = AsyncHTTPHandler(timeout=42.5) held_client = handler.client await held_client.aclose() assert held_client.is_closed async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: await _read_http_request(reader) writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") await writer.drain() writer.close() server = await asyncio.start_server(respond, "127.0.0.1", 0) port = server.sockets[0].getsockname()[1] try: response = await handler.post(f"http://127.0.0.1:{port}/v1/compress", json={"messages": []}) finally: server.close() await server.wait_closed() assert response.status_code == 200 assert handler.client is not held_client assert handler.client.timeout == httpx.Timeout(42.5) await handler.close() @pytest.mark.asyncio async def test_init_held_async_handler_survives_evicted_client_close(): from litellm.caching.evicted_client_closer import EvictedClientCloser from litellm.caching.llm_caching_handler import LLMClientCache cache = LLMClientCache(evicted_client_closer=EvictedClientCloser(grace_seconds=0)) handler = AsyncHTTPHandler(timeout=42.5) held_client = handler.client cache.set_cache("init-held-handler", handler, litellm_owned_client=True, ttl=0) await asyncio.sleep(0.02) assert cache.get_cache("init-held-handler") is None await asyncio.sleep(0.05) assert held_client.is_closed async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: await _read_http_request(reader) writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") await writer.drain() writer.close() server = await asyncio.start_server(respond, "127.0.0.1", 0) port = server.sockets[0].getsockname()[1] try: response = await handler.post(f"http://127.0.0.1:{port}/v1/compress", json={"messages": []}) finally: server.close() await server.wait_closed() assert response.status_code == 200 assert handler.client is not held_client assert handler.client.timeout == httpx.Timeout(42.5) await handler.close() def test_init_held_sync_handler_recreates_closed_client(): from http.server import BaseHTTPRequestHandler, HTTPServer class OkRequestHandler(BaseHTTPRequestHandler): def do_GET(self): self.send_response(200) self.send_header("Content-Length", "2") self.end_headers() self.wfile.write(b"ok") def log_message(self, format, *args): pass handler = HTTPHandler(timeout=7) held_client = handler.client held_client.close() server = HTTPServer(("127.0.0.1", 0), OkRequestHandler) thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() try: response = handler.get(f"http://127.0.0.1:{server.server_port}/") finally: server.shutdown() server.server_close() thread.join(timeout=5) assert response.status_code == 200 assert handler.client is not held_client assert handler.client.timeout == httpx.Timeout(7) handler.close() def test_caller_supplied_sync_client_is_not_replaced_when_closed(): supplied = httpx.Client() handler = HTTPHandler(client=supplied) supplied.close() assert handler.client is supplied @pytest.mark.asyncio async def test_assigned_async_client_is_not_replaced(): handler = AsyncHTTPHandler() await handler.client.aclose() replacement = MagicMock() handler.client = replacement assert handler.client is replacement def test_concurrent_sync_heal_creates_exactly_one_replacement(): class GatedHealHandler(HTTPHandler): def __init__(self): self.heal_started = threading.Event() self.release_heal = threading.Event() self.heal_calls = 0 super().__init__(timeout=7) def create_client(self) -> httpx.Client: if hasattr(self, "_client"): self.heal_calls += 1 self.heal_started.set() assert self.release_heal.wait(timeout=5) return super().create_client() handler = GatedHealHandler() handler.client.close() seen = [] def grab_client(): seen.append(handler.client) first = threading.Thread(target=grab_client) second = threading.Thread(target=grab_client) first.start() assert handler.heal_started.wait(timeout=5) second.start() second.join(timeout=0.3) handler.release_heal.set() first.join(timeout=5) second.join(timeout=5) assert handler.heal_calls == 1 assert seen[0] is seen[1] assert not seen[0].is_closed handler.close() @pytest.fixture def fresh_llm_client_cache(): from litellm.caching.llm_caching_handler import LLMClientCache previous = getattr(litellm, "in_memory_llm_clients_cache", None) litellm.in_memory_llm_clients_cache = LLMClientCache() try: yield finally: litellm.in_memory_llm_clients_cache = previous def test_sole_referrer_handler_may_close_but_a_sharing_one_may_not(): from litellm.llms.custom_httpx.http_handler import _handler_may_close_client assert _handler_may_close_client(_CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER, owns_client=True) assert not _handler_may_close_client(_CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER + 1, owns_client=True) assert not _handler_may_close_client(_CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER, owns_client=False) @pytest.fixture def keepalive_server(): from http.server import BaseHTTPRequestHandler, HTTPServer from socketserver import ThreadingMixIn class OkRequestHandler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" def do_GET(self): self.send_response(200) self.send_header("Content-Length", "2") self.end_headers() self.wfile.write(b"ok") def log_message(self, format, *args): pass class ThreadedServer(ThreadingMixIn, HTTPServer): daemon_threads = True server = ThreadedServer(("127.0.0.1", 0), OkRequestHandler) thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() try: yield f"http://127.0.0.1:{server.server_port}/" finally: server.shutdown() server.server_close() thread.join(timeout=5) def test_exclusively_owned_sync_client_pool_is_closed_when_handler_is_collected(keepalive_server): handler = HTTPHandler() pool = handler._client._transport._pool client_ref = weakref.ref(handler._client) handler.get(keepalive_server) assert pool._connections, "setup failed: no pooled connection to release" del handler gc.collect() assert client_ref() is None assert pool._connections == [] @pytest.mark.asyncio async def test_exclusively_owned_async_client_pool_is_closed_when_handler_is_collected(keepalive_server, monkeypatch): monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setattr(litellm, "force_ipv4", False) handler = AsyncHTTPHandler() pool = handler.client._transport._pool await handler.get(keepalive_server) assert pool._connections, "setup failed: no pooled connection to release" del handler gc.collect() await asyncio.sleep(0.25) assert pool._connections == [] @pytest.mark.asyncio async def test_handed_out_async_client_pool_survives_handler_collection(keepalive_server, monkeypatch): monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setattr(litellm, "force_ipv4", False) handler = AsyncHTTPHandler() consumer_client = handler.client pool = consumer_client._transport._pool await handler.get(keepalive_server) assert pool._connections, "setup failed: no pooled connection to observe" del handler gc.collect() await asyncio.sleep(0.25) assert pool._connections != [] assert not consumer_client.is_closed await consumer_client.aclose() def test_handed_out_sync_client_pool_survives_handler_collection(keepalive_server): handler = HTTPHandler() consumer_client = handler.client pool = consumer_client._transport._pool handler.get(keepalive_server) assert pool._connections, "setup failed: no pooled connection to observe" del handler gc.collect() assert pool._connections != [] assert not consumer_client.is_closed consumer_client.close() def _mock_transport() -> httpx.MockTransport: """Answers anything with a short body, left unread when the caller asked to stream.""" def respond(request: httpx.Request) -> httpx.Response: return httpx.Response(200, request=request, content=b"ab") return httpx.MockTransport(respond) RELEASED_TOO_EARLY = "the handler was released while its response could still read" NEVER_RELEASED = "the handler outlived the response that was holding it" # Every method that can hand back a body the caller has not read yet, which is # every one that passes stream= down to send(). Parametrized so a method added # later is covered here rather than being the one that forgets to anchor. ASYNC_STREAMING_SENDS = ["post", "delete"] SYNC_STREAMING_SENDS = ["post", "patch", "put", "delete"] @pytest.mark.asyncio @pytest.mark.parametrize("method", ASYNC_STREAMING_SENDS) async def test_a_streaming_response_holds_its_handler_until_it_is_released(method): """The finalizer must not run while a body this handler issued can still arrive. ``_handler_may_close_client`` cannot see that body: it holds the connection it reads from and never the client. Anchoring the handler to the response is what withholds the close, and releasing the anchor is what still delivers one. """ handler = AsyncHTTPHandler() handler.client._transport = _mock_transport() ref = weakref.ref(handler) response = await getattr(handler, method)("https://example.invalid/stream", stream=True) del handler gc.collect() assert ref() is not None, RELEASED_TOO_EARLY assert await response.aread() == b"ab" del response gc.collect() assert ref() is None, NEVER_RELEASED @pytest.mark.parametrize("method", SYNC_STREAMING_SENDS) def test_a_sync_streaming_response_holds_its_handler_until_it_is_released(method): """The sync finalizer closes inline, so the same anchor has to hold it off.""" handler = HTTPHandler() handler.client._transport = _mock_transport() ref = weakref.ref(handler) response = getattr(handler, method)("https://example.invalid/stream", stream=True) del handler gc.collect() assert ref() is not None, RELEASED_TOO_EARLY assert response.read() == b"ab" del response gc.collect() assert ref() is None, NEVER_RELEASED @pytest.mark.asyncio async def test_a_fully_read_response_does_not_hold_its_handler(): """A non-streaming response is complete when ``post`` returns, so it anchors nothing. Otherwise every client close would wait on whatever the caller does next with a response it has already read. """ handler = AsyncHTTPHandler() handler.client._transport = _mock_transport() ref = weakref.ref(handler) response = await handler.post("https://example.invalid/whole") assert response.content == b"ab" del handler gc.collect() assert ref() is None, "a fully-read response pinned its handler" def test_sync_close_leaves_caller_supplied_client_open(): supplied = httpx.Client() handler = HTTPHandler(client=supplied) handler.close() assert not supplied.is_closed supplied.close() @pytest.mark.asyncio async def test_async_aexit_closes_an_owned_client_but_not_an_assigned_one(): owning = AsyncHTTPHandler() owned = owning.client await owning.__aexit__() assert owned.is_closed borrowing = AsyncHTTPHandler() original = borrowing.client assigned = httpx.AsyncClient() borrowing.client = assigned await borrowing.__aexit__() assert not assigned.is_closed await assigned.aclose() await original.aclose() @pytest.mark.asyncio async def test_async_close_leaves_assigned_client_open(): handler = AsyncHTTPHandler() owned = handler.client assigned = httpx.AsyncClient() handler.client = assigned await handler.close() assert not assigned.is_closed await assigned.aclose() await owned.aclose() def test_client_handed_out_by_sync_cache_survives_eviction_and_collection(fresh_llm_client_cache): from litellm.caching.llm_caching_handler import LLMClientCache handler = _get_httpx_client() consumer_client = handler.client handler_ref = weakref.ref(handler) assert not consumer_client.is_closed assert litellm.in_memory_llm_clients_cache.get_cache("httpx_client") is handler litellm.in_memory_llm_clients_cache = LLMClientCache() del handler gc.collect() assert litellm.in_memory_llm_clients_cache.get_cache("httpx_client") is None assert handler_ref() is None assert not consumer_client.is_closed consumer_client.close() @pytest.mark.asyncio async def test_client_handed_out_by_async_cache_survives_eviction_and_collection(fresh_llm_client_cache): from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.utils import LlmProviders handler = get_async_httpx_client(llm_provider=LlmProviders.OPENAI) consumer_client = handler.client handler_ref = weakref.ref(handler) assert not consumer_client.is_closed litellm.in_memory_llm_clients_cache = LLMClientCache() del handler gc.collect() await asyncio.sleep(0.1) assert handler_ref() is None assert not consumer_client.is_closed await consumer_client.aclose() _SET_COOKIE = "SESSION=upstream-a-secret; Path=/" def _cookie_recorder(): """A transport that hands out a Set-Cookie once, and records what comes back.""" seen = [] def handler(request: httpx.Request) -> httpx.Response: seen.append(request.headers.get("cookie")) if request.url.path == "/set": return httpx.Response(200, headers={"set-cookie": _SET_COOKIE}) return httpx.Response(200) return handler, seen @pytest.mark.asyncio async def test_async_client_never_replays_one_upstreams_cookie_to_another(): """LiteLLM's async clients are pooled and shared by every caller, so a cookie one upstream sets would be attached to every later request on a matching domain, reaching a different tenant's upstream. The client must persist no response cookie.""" handler, seen = _cookie_recorder() http_handler = AsyncHTTPHandler() client = http_handler.client client._transport = httpx.MockTransport(handler) await client.get("https://upstream-a.example.com/set") await client.get("https://upstream-b.example.com/rpc") await client.aclose() assert dict(client.cookies) == {}, "the shared client stored an upstream's cookie" assert seen == [None, None] def test_sync_client_never_replays_one_upstreams_cookie_to_another(): """Same invariant on the sync client, which is pooled the same way.""" handler, seen = _cookie_recorder() http_handler = HTTPHandler() client = http_handler.client client._transport = httpx.MockTransport(handler) client.get("https://upstream-a.example.com/set") client.get("https://upstream-b.example.com/rpc") client.close() assert dict(client.cookies) == {} assert seen == [None, None] @pytest.mark.asyncio async def test_aiohttp_session_never_replays_one_upstreams_cookie_to_another(): """The httpx jar is not the only one. AiohttpTransport is litellm's default transport and the aiohttp ClientSession keeps its own cookie jar, which httpx-level assertions cannot see, so blocking only the httpx jar leaves the leak intact on the real path. aiohttp's default jar refuses cookies for IP hosts, so this drives a hostname. An IP-addressed check passes whether or not the session jar is blocked.""" from aiohttp import DummyCookieJar from yarl import URL http_handler = AsyncHTTPHandler(timeout=61.0) transport = http_handler.client._transport assert isinstance(transport, LiteLLMAiohttpTransport), "aiohttp is no longer the default transport" session = transport.client() if callable(transport.client) else transport.client jar = session.cookie_jar assert isinstance(jar, DummyCookieJar) jar.update_cookies({"SESSION": "upstream-a-secret"}, URL("https://upstream-a.example.com")) assert len(jar) == 0 assert dict(jar.filter_cookies(URL("https://upstream-a.example.com"))) == {} await session.close() def _mint_session_on_dead_loop(handler: AsyncHTTPHandler) -> ClientSession: """Create the transport's real ClientSession on a loop that then closes. This is the lifecycle of every client minted for a short-lived event loop (the loop-id-keyed LLM client cache creates one handler per loop): the session outlives its loop and can only ever be disposed loop-lessly. """ transport = handler.client._transport assert isinstance(transport, LiteLLMAiohttpTransport) loop = asyncio.new_event_loop() async def _create() -> ClientSession: return transport._get_valid_client_session() session = loop.run_until_complete(_create()) loop.close() return session def test_finalizer_without_running_loop_closes_dead_loop_session(): """A handler finalized with no running event loop must still dispose its aiohttp session. The async close can never run in that context; without the synchronous fallback the session and its connector are abandoned to GC and emit "Unclosed client session" / "Unclosed connector" warnings.""" handler = AsyncHTTPHandler(timeout=61.0) session = _mint_session_on_dead_loop(handler) assert not session.closed del handler gc.collect() assert session.closed @pytest.mark.asyncio async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref(): """With a running loop, finalization schedules an async close and must keep a strong reference to the task until it completes — a bare create_task() result may be collected before it runs, leaving the session unclosed.""" handler = AsyncHTTPHandler(timeout=61.0) transport = handler.client._transport assert isinstance(transport, LiteLLMAiohttpTransport) session = transport._get_valid_client_session() assert not session.closed del transport baseline_tasks = set(AsyncHTTPHandler._finalizer_close_tasks) del handler gc.collect() scheduled = AsyncHTTPHandler._finalizer_close_tasks - baseline_tasks assert len(scheduled) == 1 await asyncio.gather(*scheduled) assert session.closed assert not (AsyncHTTPHandler._finalizer_close_tasks & scheduled) @pytest.mark.asyncio async def test_sync_close_helper_respects_session_ownership(): """The loop-less fallback closes only sessions the transport owns; a shared session (e.g. the proxy's) must never be closed by a handler.""" owned_handler = AsyncHTTPHandler(timeout=61.0) owned_transport = owned_handler.client._transport assert isinstance(owned_transport, LiteLLMAiohttpTransport) owned_session = owned_transport._get_valid_client_session() baseline = set(LiteLLMAiohttpTransport._background_close_tasks) owned_handler._dispose_wrapped_aiohttp_session() scheduled = LiteLLMAiohttpTransport._background_close_tasks - baseline await asyncio.gather(*scheduled) assert owned_session.closed shared_session = ClientSession() shared_handler = AsyncHTTPHandler(timeout=61.0, shared_session=shared_session) shared_transport = shared_handler.client._transport assert isinstance(shared_transport, LiteLLMAiohttpTransport) assert shared_transport._owns_session is False shared_handler._dispose_wrapped_aiohttp_session() assert not shared_session.closed await shared_session.close() await shared_handler.close() await owned_handler.close() @pytest.mark.asyncio async def test_finalizer_close_done_consumes_exception(): """A failing finalizer close must have its exception retrieved by the done callback, or asyncio emits "Task exception was never retrieved" at GC — the same log noise the finalizer path exists to eliminate.""" async def failing_close() -> None: raise RuntimeError("close failed") task = asyncio.get_running_loop().create_task(failing_close()) AsyncHTTPHandler._finalizer_close_tasks.add(task) await asyncio.sleep(0) AsyncHTTPHandler._on_finalizer_close_done(task) assert task not in AsyncHTTPHandler._finalizer_close_tasks cancelled = asyncio.get_running_loop().create_task(asyncio.sleep(30)) cancelled.cancel() await asyncio.sleep(0) AsyncHTTPHandler._on_finalizer_close_done(cancelled) @pytest.mark.asyncio async def test_finalizer_on_live_loop_disposes_foreign_loop_session_without_scheduling(): """GC on a live loop (e.g. the app's) of a handler whose session belongs to another, dead loop must not schedule aclose() here — that is the cross-loop path the transport refuses — and must still dispose the session.""" handler = AsyncHTTPHandler(timeout=61.0) session = await asyncio.to_thread(_mint_session_on_dead_loop, handler) assert not session.closed baseline_tasks = set(AsyncHTTPHandler._finalizer_close_tasks) del handler gc.collect() assert AsyncHTTPHandler._finalizer_close_tasks == baseline_tasks assert session.closed class _RetryClientHandler(AsyncHTTPHandler): def __init__(self, first: httpx.AsyncClient, retry: httpx.AsyncClient) -> None: self._retry_client: Final = retry super().__init__() self.client = first def create_client( self, timeout: float | httpx.Timeout | None = None, event_hooks: Mapping[str, list[Callable[..., object]]] | None = None, ssl_verify: VerifyTypes | None = None, shared_session: ClientSession | None = None, ) -> httpx.AsyncClient: return self._retry_client @pytest.mark.asyncio @pytest.mark.parametrize("method", ["post", "put", "patch", "delete"]) async def test_connection_error_retry_forwards_content(method: str): captured: list[bytes] = [] # mutable-ok: async closure capture buffer async def raise_connection_error(request: httpx.Request) -> httpx.Response: raise httpx.RemoteProtocolError("connection dropped", request=request) async def capture_and_succeed(request: httpx.Request) -> httpx.Response: captured.append(request.content) return httpx.Response(200, request=request) first: Final = httpx.AsyncClient(transport=httpx.MockTransport(raise_connection_error)) retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(capture_and_succeed)) async with first, retry: handler: Final = _RetryClientHandler(first=first, retry=retry) body = b'{"post": ["run1"]}' await getattr(handler, method)("https://api.example.com/runs/batch", content=body) assert captured == [body], "the retried request must carry the same content= body" await handler.close() @pytest.fixture def forward_proxy_server(): """Plain HTTP forward proxy that records the absolute URIs it is asked to fetch.""" from http.server import BaseHTTPRequestHandler, HTTPServer from socketserver import ThreadingMixIn seen_uris: list[str] = [] class RecordingProxyHandler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" def do_GET(self): seen_uris.append(self.path) self.send_response(200) self.send_header("Content-Length", "9") self.end_headers() self.wfile.write(b"via-proxy") def log_message(self, format, *args): pass class ThreadedServer(ThreadingMixIn, HTTPServer): daemon_threads = True server = ThreadedServer(("127.0.0.1", 0), RecordingProxyHandler) thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() try: yield f"http://127.0.0.1:{server.server_port}", seen_uris finally: server.shutdown() server.server_close() thread.join(timeout=5) # `.invalid` never resolves (RFC 6761), so the only way this request can succeed is through the proxy _PROXY_ONLY_UPSTREAM_URL = "http://upstream.invalid/v1/models" @pytest.mark.asyncio @pytest.mark.parametrize("disable_aiohttp_transport", [True, False]) @pytest.mark.parametrize("force_ipv4", [True, False]) async def test_async_handler_honours_proxy_env_for_every_transport( forward_proxy_server, monkeypatch: pytest.MonkeyPatch, disable_aiohttp_transport: bool, force_ipv4: bool ): proxy_url, seen_uris = forward_proxy_server monkeypatch.setenv("HTTP_PROXY", proxy_url) monkeypatch.delenv("NO_PROXY", raising=False) monkeypatch.delenv("no_proxy", raising=False) monkeypatch.setattr(litellm, "disable_aiohttp_transport", disable_aiohttp_transport) monkeypatch.setattr(litellm, "force_ipv4", force_ipv4) handler = AsyncHTTPHandler() try: response = await handler.get(_PROXY_ONLY_UPSTREAM_URL) finally: await handler.close() assert response.text == "via-proxy" assert seen_uris == [_PROXY_ONLY_UPSTREAM_URL] @pytest.mark.parametrize("force_ipv4", [True, False]) def test_sync_handler_honours_proxy_env(forward_proxy_server, monkeypatch: pytest.MonkeyPatch, force_ipv4: bool): proxy_url, seen_uris = forward_proxy_server monkeypatch.setenv("HTTP_PROXY", proxy_url) monkeypatch.delenv("NO_PROXY", raising=False) monkeypatch.delenv("no_proxy", raising=False) monkeypatch.setattr(litellm, "force_ipv4", force_ipv4) handler = HTTPHandler() try: response = handler.get(_PROXY_ONLY_UPSTREAM_URL) finally: handler.close() assert response.text == "via-proxy" assert seen_uris == [_PROXY_ONLY_UPSTREAM_URL] @pytest.mark.asyncio async def test_force_ipv4_httpx_transport_honours_no_proxy(keepalive_server, monkeypatch: pytest.MonkeyPatch): """NO_PROXY hosts must still go direct when the proxy mounts are supplied by litellm instead of httpx.""" monkeypatch.setenv("HTTP_PROXY", "http://proxy.invalid:3128") monkeypatch.setenv("NO_PROXY", "127.0.0.1") monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setattr(litellm, "force_ipv4", True) handler = AsyncHTTPHandler() try: response = await handler.get(keepalive_server) finally: await handler.close() assert response.text == "ok" @pytest.fixture def private_ca_tls_upstream(tmp_path: pathlib.Path): """HTTPS server behind a CONNECT proxy, both on localhost; the server's cert is signed by a test-only CA.""" import datetime import select import socket from http.server import BaseHTTPRequestHandler, HTTPServer from socketserver import ThreadingMixIn from cryptography import x509 from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.x509.oid import NameOID key = rsa.generate_private_key(public_exponent=65537, key_size=2048) name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "upstream.invalid")]) now = datetime.datetime.now(datetime.timezone.utc) cert = ( x509.CertificateBuilder() .subject_name(name) .issuer_name(name) .public_key(key.public_key()) .serial_number(x509.random_serial_number()) .not_valid_before(now - datetime.timedelta(minutes=1)) .not_valid_after(now + datetime.timedelta(hours=1)) .add_extension(x509.SubjectAlternativeName([x509.DNSName("upstream.invalid")]), critical=False) .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True) .sign(key, hashes.SHA256()) ) ca_pem = tmp_path / "ca.pem" ca_pem.write_bytes(cert.public_bytes(serialization.Encoding.PEM)) key_pem = tmp_path / "key.pem" key_pem.write_bytes( key.private_bytes( serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption() ) ) class OkTlsHandler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" def do_GET(self): self.send_response(200) self.send_header("Content-Length", "6") self.end_headers() self.wfile.write(b"ok-tls") def log_message(self, format, *args): pass class ThreadedServer(ThreadingMixIn, HTTPServer): daemon_threads = True tls_server = ThreadedServer(("127.0.0.1", 0), OkTlsHandler) server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) server_ctx.load_cert_chain(str(ca_pem), str(key_pem)) tls_server.socket = server_ctx.wrap_socket(tls_server.socket, server_side=True) tls_port = tls_server.server_port class ConnectProxyHandler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" def do_CONNECT(self): upstream = socket.create_connection(("127.0.0.1", tls_port)) self.send_response(200, "Connection established") self.end_headers() sockets = [self.connection, upstream] while True: readable, _, _ = select.select(sockets, [], [], 5) if not readable: break for src in readable: data = src.recv(65536) if not data: upstream.close() return (upstream if src is self.connection else self.connection).sendall(data) def log_message(self, format, *args): pass proxy_server = ThreadedServer(("127.0.0.1", 0), ConnectProxyHandler) threads = [ threading.Thread(target=tls_server.serve_forever, daemon=True), threading.Thread(target=proxy_server.serve_forever, daemon=True), ] for thread in threads: thread.start() try: yield f"http://127.0.0.1:{proxy_server.server_port}", str(ca_pem) finally: for server in (proxy_server, tls_server): server.shutdown() server.server_close() for thread in threads: thread.join(timeout=5) @pytest.mark.asyncio async def test_force_ipv4_https_proxy_mount_uses_handler_ca_bundle( private_ca_tls_upstream, monkeypatch: pytest.MonkeyPatch ): proxy_url, ca_pem = private_ca_tls_upstream monkeypatch.setenv("HTTPS_PROXY", proxy_url) monkeypatch.delenv("NO_PROXY", raising=False) monkeypatch.delenv("no_proxy", raising=False) monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setattr(litellm, "force_ipv4", True) handler = AsyncHTTPHandler(ssl_verify=ca_pem) try: response = await handler.get("https://upstream.invalid/v1/models") finally: await handler.close() assert response.text == "ok-tls" def test_sync_force_ipv4_https_proxy_mount_uses_handler_ca_bundle( private_ca_tls_upstream, monkeypatch: pytest.MonkeyPatch ): proxy_url, ca_pem = private_ca_tls_upstream monkeypatch.setenv("HTTPS_PROXY", proxy_url) monkeypatch.delenv("NO_PROXY", raising=False) monkeypatch.delenv("no_proxy", raising=False) monkeypatch.setattr(litellm, "force_ipv4", True) handler = HTTPHandler(ssl_verify=ca_pem) try: response = handler.get("https://upstream.invalid/v1/models") finally: handler.close() assert response.text == "ok-tls" @pytest.mark.asyncio async def test_put_can_refuse_to_follow_a_redirect(): """The client follows redirects by default; a caller uploading to a URL it did not choose must be able to opt out.""" hops: list[str] = [] # mutable-ok: the fake transport records the paths it was asked for async def mock_handler(request: httpx.Request) -> httpx.Response: hops.append(request.url.path) if request.url.path == "/first": return httpx.Response(302, request=request, headers={"location": "/second"}) return httpx.Response(200, request=request) handler = AsyncHTTPHandler() await handler.client.aclose() handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock_handler), follow_redirects=True) try: followed = await handler.put("https://uploads.example/first", data=b"x") assert followed.status_code == 200 assert hops == ["/first", "/second"] hops.clear() with pytest.raises(MaskedHTTPStatusError) as refused: await handler.put("https://uploads.example/first", data=b"x", follow_redirects=False) assert refused.value.status_code == 302 assert hops == ["/first"] finally: await handler.close() @pytest.mark.asyncio async def test_a_retried_put_stays_a_put_and_still_refuses_redirects(): """ The connection-error retry used to resend as POST through a client that follows redirects. Storage answers a POST to a presigned PUT url with 403 or 405, so the batch looked permanently rejected, and the redirect refusal the caller asked for was silently lost. """ attempts: list[tuple[str, str]] = [] # mutable-ok: the fake transports record what they were asked for async def refusing_transport(request: httpx.Request) -> httpx.Response: attempts.append((request.method, request.url.path)) raise httpx.ConnectError("connection reset", request=request) async def retry_transport(request: httpx.Request) -> httpx.Response: attempts.append((request.method, request.url.path)) if request.url.path == "/first": return httpx.Response(302, request=request, headers={"location": "/second"}) return httpx.Response(200, request=request) class HandlerWithFakeRetryClient(AsyncHTTPHandler): def create_client(self, *args, **kwargs) -> httpx.AsyncClient: return httpx.AsyncClient(transport=httpx.MockTransport(retry_transport), follow_redirects=True) handler = HandlerWithFakeRetryClient() await handler.client.aclose() handler.client = httpx.AsyncClient(transport=httpx.MockTransport(refusing_transport)) try: with pytest.raises(MaskedHTTPStatusError) as refused: await handler.put("https://uploads.example/first", data=b"x", follow_redirects=False) assert refused.value.status_code == 302 assert attempts == [("PUT", "/first"), ("PUT", "/first")] finally: await handler.client.aclose() @pytest.mark.asyncio @pytest.mark.parametrize("target", ["https://example.com/final.json?next=1", "https://other.example/final.json?next=1"]) async def test_bounded_get_preserves_sdk_redirect_auth_and_query_handling(respx_mock, monkeypatch, target): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") respx_mock.get("https://example.com/spec.json?original=1").respond(302, headers={"location": target}) destination = respx_mock.get(target).respond(200, json={"paths": {}}) handler = AsyncHTTPHandler() try: response = await handler.get( "https://example.com/spec.json?original=1", max_response_bytes=100, follow_redirects=True, headers={"Authorization": "Bearer sentinel", "Accept-Encoding": "gzip"}, timeout=2.0, ) finally: await handler.close() assert response.json() == {"paths": {}} request = destination.calls[0].request assert request.headers.get("authorization") == (None if "other.example" in target else "Bearer sentinel") assert request.headers["accept-encoding"] == "identity" assert str(request.url) == target assert request.extensions["timeout"]["read"] == 2.0 @pytest.mark.asyncio async def test_bounded_get_stops_redirect_loops(respx_mock, monkeypatch): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") route = respx_mock.get("https://example.com/spec.json").respond(302, headers={"location": "/spec.json"}) handler = AsyncHTTPHandler() try: with pytest.raises(ValueError, match="Too many redirects"): await handler.get("https://example.com/spec.json", max_response_bytes=100, follow_redirects=True) finally: await handler.close() assert route.call_count == 11 @pytest.mark.asyncio async def test_bounded_get_closes_stream_on_cancellation(respx_mock, monkeypatch): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") started = asyncio.Event() closed = asyncio.Event() class SlowStream(httpx.AsyncByteStream): async def __aiter__(self): yield b"x" started.set() await asyncio.Event().wait() async def aclose(self): closed.set() respx_mock.get("https://example.com/slow.json").respond(200, stream=SlowStream()) handler = AsyncHTTPHandler() try: task = asyncio.create_task(handler.get("https://example.com/slow.json", max_response_bytes=100)) await asyncio.wait_for(started.wait(), timeout=1) task.cancel() with pytest.raises(asyncio.CancelledError): await task finally: await handler.close() assert closed.is_set() @pytest.mark.asyncio async def test_http2_flag_bypasses_aiohttp_transport(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) monkeypatch.setattr(litellm, "force_ipv4", False) monkeypatch.delenv("LITELLM_HTTP2", raising=False) monkeypatch.delenv("DISABLE_AIOHTTP_TRANSPORT", raising=False) monkeypatch.setattr(litellm, "http2", True) assert AsyncHTTPHandler._should_use_aiohttp_transport() is False assert AsyncHTTPHandler._create_async_transport() is None 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 @pytest.mark.asyncio async def test_http2_disabled_by_default(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(litellm, "http2", False) monkeypatch.delenv("LITELLM_HTTP2", raising=False) monkeypatch.delenv("DISABLE_AIOHTTP_TRANSPORT", raising=False) monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) assert AsyncHTTPHandler._should_use_aiohttp_transport() is True