mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(aiohttp): respect ssl_verify with shared sessions (#20349)
* fix(aiohttp): respect ssl_verify with shared sessions * fix(aiohttp): resolve mypy error for ssl parameter type Pass ssl kwarg conditionally to aiohttp request() only when explicitly configured, since None is not a valid value for the ssl parameter (expected SSLContext | bool | Fingerprint).
This commit is contained in:
parent
1afe3032fd
commit
3bbc25a3f0
3 changed files with 114 additions and 2 deletions
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import os
|
||||
import ssl
|
||||
import typing
|
||||
import urllib.request
|
||||
from typing import Callable, Dict, Optional, Union
|
||||
|
|
@ -139,8 +140,13 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
|
|||
Credit to: https://github.com/karpetrosyan/httpx-aiohttp for this implementation
|
||||
"""
|
||||
|
||||
def __init__(self, client: Union[ClientSession, Callable[[], ClientSession]]):
|
||||
def __init__(
|
||||
self,
|
||||
client: Union[ClientSession, Callable[[], ClientSession]],
|
||||
ssl_verify: Optional[Union[bool, ssl.SSLContext]] = None,
|
||||
):
|
||||
self.client = client
|
||||
self._ssl_verify = ssl_verify # Store for per-request SSL override
|
||||
super().__init__(client=client)
|
||||
# Store the client factory for recreating sessions when needed
|
||||
if callable(client):
|
||||
|
|
@ -214,6 +220,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
|
|||
timeout: dict,
|
||||
proxy: Optional[str],
|
||||
sni_hostname: Optional[str],
|
||||
ssl_verify: Optional[Union[bool, ssl.SSLContext]] = None,
|
||||
) -> ClientResponse:
|
||||
"""
|
||||
Helper function to make an aiohttp request with the given parameters.
|
||||
|
|
@ -224,6 +231,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
|
|||
timeout: Timeout settings dict with 'connect', 'read', 'pool' keys
|
||||
proxy: Optional proxy URL
|
||||
sni_hostname: Optional SNI hostname for SSL
|
||||
ssl_verify: Optional SSL verification setting (False to disable, SSLContext for custom)
|
||||
|
||||
Returns:
|
||||
ClientResponse from aiohttp
|
||||
|
|
@ -237,6 +245,13 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
|
|||
data = request.stream # type: ignore
|
||||
request.headers.pop("transfer-encoding", None) # handled by aiohttp
|
||||
|
||||
# Only pass ssl kwarg when explicitly configured, to avoid
|
||||
# overriding the session/connector defaults with None (which is
|
||||
# not a valid value for aiohttp's ssl parameter).
|
||||
ssl_kwargs: Dict[str, Union[bool, ssl.SSLContext]] = {}
|
||||
if ssl_verify is not None:
|
||||
ssl_kwargs["ssl"] = ssl_verify
|
||||
|
||||
response = await client_session.request(
|
||||
method=request.method,
|
||||
url=YarlURL(str(request.url), encoded=True),
|
||||
|
|
@ -251,6 +266,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
|
|||
),
|
||||
proxy=proxy,
|
||||
server_hostname=sni_hostname,
|
||||
**ssl_kwargs,
|
||||
).__aenter__()
|
||||
|
||||
return response
|
||||
|
|
@ -268,6 +284,9 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
|
|||
# Resolve proxy settings from environment variables
|
||||
proxy = await self._get_proxy_settings(request)
|
||||
|
||||
# Use stored SSL configuration for per-request override
|
||||
ssl_config = self._ssl_verify
|
||||
|
||||
try:
|
||||
with map_aiohttp_exceptions():
|
||||
response = await self._make_aiohttp_request(
|
||||
|
|
@ -276,6 +295,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
|
|||
timeout=timeout,
|
||||
proxy=proxy,
|
||||
sni_hostname=sni_hostname,
|
||||
ssl_verify=ssl_config,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
# Handle the case where session was closed between our check and actual use
|
||||
|
|
@ -296,6 +316,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
|
|||
timeout=timeout,
|
||||
proxy=proxy,
|
||||
sni_hostname=sni_hostname,
|
||||
ssl_verify=ssl_config,
|
||||
)
|
||||
else:
|
||||
# Re-raise if it's a different RuntimeError
|
||||
|
|
|
|||
|
|
@ -846,6 +846,16 @@ class AsyncHTTPHandler:
|
|||
if str_to_bool(os.getenv("AIOHTTP_TRUST_ENV", "False")) is True:
|
||||
trust_env = True
|
||||
|
||||
#########################################################
|
||||
# Determine SSL config to pass to transport for per-request override
|
||||
# This ensures ssl_verify works even with shared sessions
|
||||
#########################################################
|
||||
ssl_for_transport: Optional[Union[bool, ssl.SSLContext]] = None
|
||||
if ssl_context is not None:
|
||||
ssl_for_transport = ssl_context
|
||||
elif ssl_verify is False:
|
||||
ssl_for_transport = False
|
||||
|
||||
verbose_logger.debug("Creating AiohttpTransport...")
|
||||
|
||||
# Use shared session if provided and valid
|
||||
|
|
@ -853,7 +863,10 @@ class AsyncHTTPHandler:
|
|||
verbose_logger.debug(
|
||||
f"SHARED SESSION: Reusing existing ClientSession (ID: {id(shared_session)})"
|
||||
)
|
||||
return LiteLLMAiohttpTransport(client=shared_session)
|
||||
return LiteLLMAiohttpTransport(
|
||||
client=shared_session,
|
||||
ssl_verify=ssl_for_transport,
|
||||
)
|
||||
|
||||
# Create new session only if none provided or existing one is invalid
|
||||
verbose_logger.debug(
|
||||
|
|
@ -877,6 +890,7 @@ class AsyncHTTPHandler:
|
|||
connector=TCPConnector(**transport_connector_kwargs),
|
||||
trust_env=trust_env,
|
||||
),
|
||||
ssl_verify=ssl_for_transport,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -140,6 +140,83 @@ async def test_ssl_verification_with_aiohttp_transport():
|
|||
litellm.disable_aiohttp_transport = original_disable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ssl_verification_with_shared_session():
|
||||
"""
|
||||
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
|
||||
original_disable = litellm.disable_aiohttp_transport
|
||||
litellm.disable_aiohttp_transport = False
|
||||
|
||||
try:
|
||||
# Create a shared session (simulating what happens in production)
|
||||
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()
|
||||
finally:
|
||||
# Restore original setting
|
||||
litellm.disable_aiohttp_transport = original_disable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ssl_context_with_shared_session():
|
||||
"""
|
||||
Test that ssl_context is respected even with shared sessions.
|
||||
"""
|
||||
import aiohttp
|
||||
|
||||
# Ensure aiohttp transport is enabled for this test
|
||||
original_disable = litellm.disable_aiohttp_transport
|
||||
litellm.disable_aiohttp_transport = False
|
||||
|
||||
try:
|
||||
# Create a custom SSL context
|
||||
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()
|
||||
finally:
|
||||
# Restore original setting
|
||||
litellm.disable_aiohttp_transport = original_disable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_transport_trust_env_setting(monkeypatch):
|
||||
"""Test that trust_env setting is properly configured in aiohttp transport"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue