Revert "Revert "fix(http_handler): bypass cache when shared_session is provided for aiohttp tracing (#20630)""

This reverts commit 0b69c21eca.
This commit is contained in:
Ishaan Jaffer 2026-02-21 12:32:41 -08:00
parent f9398eef80
commit ba0d541b19

View file

@ -559,6 +559,85 @@ async def test_shared_session_each_call_gets_new_handler():
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"""