From 7ee36c2a3ab74f01543033cc5a32fbaba0863c26 Mon Sep 17 00:00:00 2001 From: shin-bot-litellm Date: Fri, 6 Feb 2026 23:56:16 -0800 Subject: [PATCH] fix(http_handler): bypass cache when shared_session is provided for aiohttp tracing (#20630) * Add http support to custom code guardrails + Unified guardrails for MCP + Agent guardrail support (#20619) * fix: fix styling * fix(custom_code_guardrail.py): add http support for custom code guardrails allows users to call external guardrails on litellm with minimal code changes (no custom handlers) Test guardrail integrations more easily * feat(a2a/): add guardrails for agent interactions allows the same guardrails for llm's to be applied to agents as well * fix(a2a/): support passing guardrails to a2a from the UI * style(code-editor): allow editing custom code guardrails on ui + add examples of pre/post calls for custom code guardrails * feat(mcp/): support custom code guardrails for mcp calls allows custom code guardrails to work on mcp input * feat(chatui.tsx): support guardrails on mcp tool calls on playground * fix(mypy): resolve missing return statements and type casting issues (#20618) * fix(mypy): resolve missing return statements and type casting issues * fix(pangea): use elif to prevent UnboundLocalError and handle None messages Address Greptile review feedback: - Make branches mutually exclusive using elif to prevent input_messages from being overwritten - Handle case where data.get('messages') returns None to avoid passing invalid payload to Pangea API --------- Co-authored-by: Shin * [Feat] MCP Gateway - Allow setting MCP Servers as Private/Public available on Internet (#20607) * update MCPAuthenticatedUser * add available_on_public_internet for MCPs * update claude.md * init IPAddressUtils * init available_on_public_internet * add on REST endpoints * filter with IP * TestIsInternalIp * _extract_mcp_headers_from_request * init get_mcp_client_ip * _get_general_settings * allowed_server_ids * address PR comments * get_mcp_server_by_name fix * fix server * fix review comments * get_public_mcp_servers * address _get_allowed_mcp_servers * fixing user_id * [Feat] IP-Based Access Control for MCP Servers (#20620) * update MCPAuthenticatedUser * add available_on_public_internet for MCPs * update claude.md * init IPAddressUtils * init available_on_public_internet * add on REST endpoints * filter with IP * TestIsInternalIp * _extract_mcp_headers_from_request * init get_mcp_client_ip * _get_general_settings * allowed_server_ids * address PR comments * get_mcp_server_by_name fix * fix server * fix review comments * get_public_mcp_servers * address _get_allowed_mcp_servers * test fix * fix linting * inint ui types * add ui for managing MCP private/public * add ui * fixes * add to schema * add types * fix endpoint * add endpoint * update manager * test mcp * dont use external party for ip address * Add OpenAI/Azure release test suite with HTTP client lifecycle regression detection (#20622) * docs (#20626) * docs * fix(mypy): resolve type checking errors in 5 files (#20627) - a2a_protocol/exception_mapping_utils.py: Fix type ignore comment for None assignment - caching/redis_cache.py: Add type ignore for async ping return type - caching/redis_cluster_cache.py: Add type ignore for async ping return type - llms/deprecated_providers/palm.py: Add type ignore for palm.generate_text - proxy/auth/handle_jwt.py: Add type ignore for jwt.decode options argument All changes add appropriate type: ignore comments to handle library typing inconsistencies. * fix(test): update deprecated gemini embedding model (#20621) Replace text-embedding-004 with gemini-embedding-001. The old model was deprecated and returns 404: 'models/text-embedding-004 is not found for API version v1beta' Co-authored-by: Shin * ui new buil * 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 --------- Co-authored-by: Krish Dholakia Co-authored-by: Shin Co-authored-by: Ishaan Jaff Co-authored-by: yuneng-jiang Co-authored-by: Alexsander Hamir Co-authored-by: shin-bot-litellm --- litellm/llms/custom_httpx/http_handler.py | 23 +++++- .../llms/custom_httpx/test_http_handler.py | 79 +++++++++++++++++++ 2 files changed, 100 insertions(+), 2 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 95f411c397..5cf6efe5ba 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1206,7 +1206,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(): @@ -1233,12 +1254,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( diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index c249bd9970..b0011fd8f7 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -480,6 +480,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"""