mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(http_handler): bypass cache when shared_session is provided for aiohttp tracing
When users pass a shared_session with trace_configs to acompletion(), the get_async_httpx_client() function was ignoring it and returning a cached client without the user's tracing configuration. This fix bypasses the cache when shared_session is provided, ensuring the user's ClientSession (with its trace_configs, connector settings, etc.) is actually used for the request. Fixes #20174
This commit is contained in:
parent
51af66fdb2
commit
222508c9cc
2 changed files with 100 additions and 2 deletions
|
|
@ -1192,7 +1192,28 @@ def get_async_httpx_client(
|
|||
If not present, creates a new client
|
||||
|
||||
Caches the new client and returns it.
|
||||
|
||||
Note: When shared_session is provided, the cache is bypassed to ensure
|
||||
the user's session (with its trace_configs, connector settings, etc.)
|
||||
is used for the request.
|
||||
"""
|
||||
# When shared_session is provided, bypass cache and create a new handler
|
||||
# that uses the user's session directly. This preserves the user's
|
||||
# session configuration including trace_configs for aiohttp tracing.
|
||||
if shared_session is not None:
|
||||
verbose_logger.debug(
|
||||
f"shared_session provided (ID: {id(shared_session)}), bypassing client cache"
|
||||
)
|
||||
if params is not None:
|
||||
handler_params = {k: v for k, v in params.items() if k != "disable_aiohttp_transport"}
|
||||
handler_params["shared_session"] = shared_session
|
||||
return AsyncHTTPHandler(**handler_params)
|
||||
else:
|
||||
return AsyncHTTPHandler(
|
||||
timeout=httpx.Timeout(timeout=600.0, connect=5.0),
|
||||
shared_session=shared_session,
|
||||
)
|
||||
|
||||
_params_key_name = ""
|
||||
if params is not None:
|
||||
for key, value in params.items():
|
||||
|
|
@ -1219,12 +1240,10 @@ def get_async_httpx_client(
|
|||
if params is not None:
|
||||
# Filter out params that are only used for cache key, not for AsyncHTTPHandler.__init__
|
||||
handler_params = {k: v for k, v in params.items() if k != "disable_aiohttp_transport"}
|
||||
handler_params["shared_session"] = shared_session
|
||||
_new_client = AsyncHTTPHandler(**handler_params)
|
||||
else:
|
||||
_new_client = AsyncHTTPHandler(
|
||||
timeout=httpx.Timeout(timeout=600.0, connect=5.0),
|
||||
shared_session=shared_session,
|
||||
)
|
||||
|
||||
cache.set_cache(
|
||||
|
|
|
|||
|
|
@ -403,6 +403,85 @@ async def test_session_reuse_integration():
|
|||
await client2.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shared_session_bypasses_cache():
|
||||
"""
|
||||
Test that when shared_session is provided, the cache is bypassed.
|
||||
|
||||
This is critical for aiohttp tracing support - users need their custom
|
||||
ClientSession (with trace_configs) to be used, not a cached session.
|
||||
|
||||
Related: GitHub issue #20174
|
||||
"""
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
# First, get a cached client without shared_session
|
||||
cached_client = get_async_httpx_client(
|
||||
llm_provider=LlmProviders.ANTHROPIC,
|
||||
shared_session=None
|
||||
)
|
||||
|
||||
# Now create a mock shared session
|
||||
mock_session = MockClientSession()
|
||||
|
||||
# Get a client WITH shared_session - this should NOT return the cached client
|
||||
client_with_session = get_async_httpx_client(
|
||||
llm_provider=LlmProviders.ANTHROPIC, # Same provider!
|
||||
shared_session=mock_session # type: ignore
|
||||
)
|
||||
|
||||
# The clients should be DIFFERENT - cache should be bypassed when shared_session is provided
|
||||
assert client_with_session is not cached_client, \
|
||||
"Cache should be bypassed when shared_session is provided"
|
||||
|
||||
# Verify the shared_session handler is using our mock session
|
||||
# The transport should have our mock_session as its client
|
||||
transport = client_with_session.client._transport
|
||||
if hasattr(transport, 'client'):
|
||||
assert transport.client is mock_session, \
|
||||
"Handler should use the provided shared_session"
|
||||
|
||||
# Clean up
|
||||
await cached_client.close()
|
||||
await client_with_session.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shared_session_each_call_gets_new_handler():
|
||||
"""
|
||||
Test that each call with shared_session creates a new handler.
|
||||
|
||||
This ensures user sessions (with their trace_configs, etc.) are always
|
||||
used and not affected by caching.
|
||||
"""
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
# Create two different mock sessions
|
||||
mock_session1 = MockClientSession()
|
||||
mock_session2 = MockClientSession()
|
||||
|
||||
# Get clients with different sessions for the same provider
|
||||
client1 = get_async_httpx_client(
|
||||
llm_provider=LlmProviders.ANTHROPIC,
|
||||
shared_session=mock_session1 # type: ignore
|
||||
)
|
||||
|
||||
client2 = get_async_httpx_client(
|
||||
llm_provider=LlmProviders.ANTHROPIC, # Same provider
|
||||
shared_session=mock_session2 # type: ignore # Different session
|
||||
)
|
||||
|
||||
# Should be different clients, each using their own session
|
||||
assert client1 is not client2, \
|
||||
"Different shared_sessions should create different handlers"
|
||||
|
||||
# Clean up
|
||||
await client1.close()
|
||||
await client2.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_validation():
|
||||
"""Test that session validation works correctly"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue