From 7b66c970e9706a5b96c1987bc445ff941e51fb66 Mon Sep 17 00:00:00 2001 From: voidborne-d Date: Tue, 17 Mar 2026 03:11:58 +0000 Subject: [PATCH 1/7] fix: auto-recover shared aiohttp session when closed (#23806) When the shared aiohttp session closes (due to network interruption, idle timeout, or Redis failover side effects), the proxy permanently falls back to creating a new HTTPS connection per request, losing the benefit of connection pooling for the entire pod lifetime. Fix: make add_shared_session_to_data() async and recreate the session when it is found closed, restoring connection pooling automatically. Fixes #23806 --- litellm/proxy/route_llm_request.py | 29 ++++-- .../proxy/test_aiohttp_session_recovery.py | 99 +++++++++++++++++++ 2 files changed, 122 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/proxy/test_aiohttp_session_recovery.py diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index e5fc9fe76a4..265a311d393 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -123,23 +123,40 @@ def get_team_id_from_data(data: dict) -> Optional[str]: return None -def add_shared_session_to_data(data: dict) -> None: +async def add_shared_session_to_data(data: dict) -> None: """ Add shared aiohttp session for connection reuse (prevents cold starts). + If the session was closed (e.g. due to network interruption or idle timeout), + automatically recreates it so connection pooling is restored. Silently continues without session reuse if import fails or session is unavailable. Args: data: Dictionary to add the shared session to """ try: + import litellm.proxy.proxy_server as proxy_server from litellm._logging import verbose_proxy_logger - from litellm.proxy.proxy_server import shared_aiohttp_session - if shared_aiohttp_session is not None and not shared_aiohttp_session.closed: - data["shared_session"] = shared_aiohttp_session + session = proxy_server.shared_aiohttp_session + + if session is not None and not session.closed: + data["shared_session"] = session verbose_proxy_logger.info( - f"SESSION REUSE: Attached shared aiohttp session to request (ID: {id(shared_aiohttp_session)})" + f"SESSION REUSE: Attached shared aiohttp session to request (ID: {id(session)})" ) + elif session is not None and session.closed: + # Session was created at startup but has since closed — recreate it + verbose_proxy_logger.warning( + f"SESSION REUSE: Shared aiohttp session is closed (ID: {id(session)}), recreating..." + ) + new_session = await proxy_server._initialize_shared_aiohttp_session() + if new_session is not None: + proxy_server.shared_aiohttp_session = new_session + data["shared_session"] = new_session + else: + verbose_proxy_logger.info( + "SESSION REUSE: Failed to recreate shared session, continuing without session reuse" + ) else: verbose_proxy_logger.info( "SESSION REUSE: No shared session available for this request" @@ -248,7 +265,7 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin """ Common helper to route the request """ - add_shared_session_to_data(data) + await add_shared_session_to_data(data) team_id = get_team_id_from_data(data) router_model_names = llm_router.model_names if llm_router is not None else [] diff --git a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py new file mode 100644 index 00000000000..e87ca5a1632 --- /dev/null +++ b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py @@ -0,0 +1,99 @@ +""" +Tests for shared aiohttp session auto-recovery. + +When the shared session closes (e.g. network interruption, idle timeout), +add_shared_session_to_data should recreate it instead of permanently +falling back to per-request connections. + +Fixes: https://github.com/BerriAI/litellm/issues/23806 +""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +@pytest.mark.asyncio +async def test_add_shared_session_attaches_open_session(): + """When the shared session is open, it should be attached to data.""" + from litellm.proxy.route_llm_request import add_shared_session_to_data + + mock_session = MagicMock() + mock_session.closed = False + + with patch( + "litellm.proxy.proxy_server.shared_aiohttp_session", mock_session + ): + data = {} + await add_shared_session_to_data(data) + assert data["shared_session"] is mock_session + + +@pytest.mark.asyncio +async def test_add_shared_session_recreates_closed_session(): + """When the shared session is closed, it should be recreated.""" + from litellm.proxy import proxy_server as proxy_server_module + from litellm.proxy.route_llm_request import add_shared_session_to_data + + closed_session = MagicMock() + closed_session.closed = True + + new_session = MagicMock() + new_session.closed = False + + with patch.object( + proxy_server_module, + "shared_aiohttp_session", + closed_session, + ): + with patch.object( + proxy_server_module, + "_initialize_shared_aiohttp_session", + new_callable=AsyncMock, + return_value=new_session, + ) as mock_init: + data = {} + await add_shared_session_to_data(data) + + mock_init.assert_called_once() + assert data["shared_session"] is new_session + assert proxy_server_module.shared_aiohttp_session is new_session + + +@pytest.mark.asyncio +async def test_add_shared_session_handles_recreation_failure(): + """When recreation fails, data should not contain shared_session.""" + from litellm.proxy import proxy_server as proxy_server_module + from litellm.proxy.route_llm_request import add_shared_session_to_data + + closed_session = MagicMock() + closed_session.closed = True + + with patch.object( + proxy_server_module, + "shared_aiohttp_session", + closed_session, + ): + with patch.object( + proxy_server_module, + "_initialize_shared_aiohttp_session", + new_callable=AsyncMock, + return_value=None, + ): + data = {} + await add_shared_session_to_data(data) + assert "shared_session" not in data + + +@pytest.mark.asyncio +async def test_add_shared_session_no_session_available(): + """When no session was ever created, data should not contain shared_session.""" + from litellm.proxy.route_llm_request import add_shared_session_to_data + + with patch( + "litellm.proxy.proxy_server.shared_aiohttp_session", None + ): + data = {} + await add_shared_session_to_data(data) + assert "shared_session" not in data From ab4fda2eebbc6e1dbc772e93d4be8c18b58deb8c Mon Sep 17 00:00:00 2001 From: voidborne-d Date: Tue, 17 Mar 2026 08:08:44 +0000 Subject: [PATCH 2/7] fix: add asyncio.Lock to prevent session/connector leak on concurrent recreation When multiple requests detect a closed shared session simultaneously, they would each create a new aiohttp.ClientSession, leaking intermediate sessions and their TCP connectors. Added double-checked locking pattern with asyncio.Lock to ensure only one coroutine recreates the session. Added concurrent recreation test case. --- litellm/proxy/route_llm_request.py | 45 ++++++++++++---- .../proxy/test_aiohttp_session_recovery.py | 51 +++++++++++++++++++ 2 files changed, 85 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 265a311d393..376bf07f2b5 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,3 +1,4 @@ +import asyncio from typing import TYPE_CHECKING, Any, Literal, Optional from fastapi import HTTPException, status @@ -123,11 +124,24 @@ def get_team_id_from_data(data: dict) -> Optional[str]: return None +_shared_session_lock: Optional[asyncio.Lock] = None + + +def _get_shared_session_lock() -> asyncio.Lock: + """Lazily create the shared session lock (must be called within a running event loop).""" + global _shared_session_lock + if _shared_session_lock is None: + _shared_session_lock = asyncio.Lock() + return _shared_session_lock + + async def add_shared_session_to_data(data: dict) -> None: """ Add shared aiohttp session for connection reuse (prevents cold starts). If the session was closed (e.g. due to network interruption or idle timeout), automatically recreates it so connection pooling is restored. + Uses an asyncio.Lock to prevent race conditions where multiple concurrent + requests could each create a new session, leaking intermediate ones. Silently continues without session reuse if import fails or session is unavailable. Args: @@ -146,23 +160,32 @@ async def add_shared_session_to_data(data: dict) -> None: ) elif session is not None and session.closed: # Session was created at startup but has since closed — recreate it - verbose_proxy_logger.warning( - f"SESSION REUSE: Shared aiohttp session is closed (ID: {id(session)}), recreating..." - ) - new_session = await proxy_server._initialize_shared_aiohttp_session() - if new_session is not None: - proxy_server.shared_aiohttp_session = new_session - data["shared_session"] = new_session - else: - verbose_proxy_logger.info( - "SESSION REUSE: Failed to recreate shared session, continuing without session reuse" + # Use lock to prevent concurrent recreation (avoids session/connector leak) + lock = _get_shared_session_lock() + async with lock: + # Double-check under lock — another coroutine may have already recreated it + session = proxy_server.shared_aiohttp_session + if session is not None and not session.closed: + data["shared_session"] = session + return + + verbose_proxy_logger.warning( + f"SESSION REUSE: Shared aiohttp session is closed (ID: {id(session)}), recreating..." ) + new_session = await proxy_server._initialize_shared_aiohttp_session() + if new_session is not None: + proxy_server.shared_aiohttp_session = new_session + data["shared_session"] = new_session + else: + verbose_proxy_logger.info( + "SESSION REUSE: Failed to recreate shared session, continuing without session reuse" + ) else: verbose_proxy_logger.info( "SESSION REUSE: No shared session available for this request" ) except Exception: - # Silently continue without session reuse if import fails or session unavailable + # Silently continue without session reuse if import fails or session is unavailable pass diff --git a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py index e87ca5a1632..b71a61d2c29 100644 --- a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py +++ b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py @@ -97,3 +97,54 @@ async def test_add_shared_session_no_session_available(): data = {} await add_shared_session_to_data(data) assert "shared_session" not in data + + +@pytest.mark.asyncio +async def test_add_shared_session_concurrent_recreation_uses_lock(): + """When multiple coroutines detect a closed session concurrently, + only one should recreate it (double-checked locking via asyncio.Lock).""" + import litellm.proxy.route_llm_request as route_module + from litellm.proxy import proxy_server as proxy_server_module + from litellm.proxy.route_llm_request import add_shared_session_to_data + + # Reset the module-level lock so each test is isolated + route_module._shared_session_lock = None + + closed_session = MagicMock() + closed_session.closed = True + + new_session = MagicMock() + new_session.closed = False + + call_count = 0 + + async def mock_init(): + nonlocal call_count + call_count += 1 + # Simulate some async work + await asyncio.sleep(0.01) + proxy_server_module.shared_aiohttp_session = new_session + return new_session + + with patch.object( + proxy_server_module, + "shared_aiohttp_session", + closed_session, + ): + with patch.object( + proxy_server_module, + "_initialize_shared_aiohttp_session", + side_effect=mock_init, + ): + # Launch 5 concurrent calls + results = [{} for _ in range(5)] + await asyncio.gather( + *(add_shared_session_to_data(d) for d in results) + ) + + # Only 1 coroutine should have called _initialize (the rest see the + # re-checked session as open under the lock) + assert call_count == 1, f"Expected 1 init call, got {call_count}" + # All should have the new session + for d in results: + assert d.get("shared_session") is new_session From 9e09bbc1dfbccc551a8898ab3268a1014c7c81d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?d=20=F0=9F=94=B9?= <258577966+voidborne-d@users.noreply.github.com> Date: Tue, 17 Mar 2026 09:54:01 +0000 Subject: [PATCH 3/7] fix: reset _shared_session_lock in all tests for event loop isolation Address Greptile P1 review: tests that exercise the closed-session code path need to reset the module-level lock to avoid RuntimeError on Python < 3.10 when asyncio.Lock is reused across different event loops. --- tests/test_litellm/proxy/test_aiohttp_session_recovery.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py index b71a61d2c29..763744609e3 100644 --- a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py +++ b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py @@ -33,9 +33,13 @@ async def test_add_shared_session_attaches_open_session(): @pytest.mark.asyncio async def test_add_shared_session_recreates_closed_session(): """When the shared session is closed, it should be recreated.""" + import litellm.proxy.route_llm_request as route_module from litellm.proxy import proxy_server as proxy_server_module from litellm.proxy.route_llm_request import add_shared_session_to_data + # Reset the module-level lock so each test uses the current event loop + route_module._shared_session_lock = None + closed_session = MagicMock() closed_session.closed = True @@ -64,9 +68,13 @@ async def test_add_shared_session_recreates_closed_session(): @pytest.mark.asyncio async def test_add_shared_session_handles_recreation_failure(): """When recreation fails, data should not contain shared_session.""" + import litellm.proxy.route_llm_request as route_module from litellm.proxy import proxy_server as proxy_server_module from litellm.proxy.route_llm_request import add_shared_session_to_data + # Reset the module-level lock so each test uses the current event loop + route_module._shared_session_lock = None + closed_session = MagicMock() closed_session.closed = True From 32ecd241168c57fcc1de28c75b3cb72524af8ea2 Mon Sep 17 00:00:00 2001 From: d Date: Tue, 17 Mar 2026 13:09:26 +0000 Subject: [PATCH 4/7] fix: address P2 review feedback - exception handling and warning accuracy - Add try/except around _initialize_shared_aiohttp_session call to catch and log exceptions (instead of letting them bubble to outer handler) - Fix warning message when re-checked session is None (was incorrectly logging closed session ID on a None session) - Add debug logging to outer except handler instead of bare pass - Add test for _initialize_shared_aiohttp_session raising exception --- litellm/proxy/route_llm_request.py | 35 +++++++++++++++---- .../proxy/test_aiohttp_session_recovery.py | 30 ++++++++++++++++ 2 files changed, 59 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 376bf07f2b5..96e6f705cca 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -169,10 +169,23 @@ async def add_shared_session_to_data(data: dict) -> None: data["shared_session"] = session return - verbose_proxy_logger.warning( - f"SESSION REUSE: Shared aiohttp session is closed (ID: {id(session)}), recreating..." - ) - new_session = await proxy_server._initialize_shared_aiohttp_session() + # session could be None here (if another coroutine set it to None) + # or closed — either way we need to recreate + if session is not None: + verbose_proxy_logger.warning( + f"SESSION REUSE: Shared aiohttp session is closed (ID: {id(session)}), recreating..." + ) + else: + verbose_proxy_logger.warning( + "SESSION REUSE: Shared aiohttp session is None after re-check, recreating..." + ) + try: + new_session = await proxy_server._initialize_shared_aiohttp_session() + except Exception: + verbose_proxy_logger.exception( + "SESSION REUSE: Exception during shared session recreation" + ) + new_session = None if new_session is not None: proxy_server.shared_aiohttp_session = new_session data["shared_session"] = new_session @@ -185,8 +198,18 @@ async def add_shared_session_to_data(data: dict) -> None: "SESSION REUSE: No shared session available for this request" ) except Exception: - # Silently continue without session reuse if import fails or session is unavailable - pass + # Continue without session reuse — this outer handler covers import failures + # and other unexpected errors to avoid breaking the request path. + # Inner recovery logic has its own specific exception handling. + try: + from litellm._logging import verbose_proxy_logger + + verbose_proxy_logger.debug( + "SESSION REUSE: Unexpected error in session setup, continuing without reuse", + exc_info=True, + ) + except Exception: + pass async def route_request( # noqa: PLR0915 - Complex routing function, refactoring tracked separately diff --git a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py index 763744609e3..224a4aa1392 100644 --- a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py +++ b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py @@ -94,6 +94,36 @@ async def test_add_shared_session_handles_recreation_failure(): assert "shared_session" not in data +@pytest.mark.asyncio +async def test_add_shared_session_handles_recreation_exception(): + """When _initialize_shared_aiohttp_session raises, data should not contain shared_session.""" + import litellm.proxy.route_llm_request as route_module + from litellm.proxy import proxy_server as proxy_server_module + from litellm.proxy.route_llm_request import add_shared_session_to_data + + # Reset the module-level lock so each test uses the current event loop + route_module._shared_session_lock = None + + closed_session = MagicMock() + closed_session.closed = True + + with patch.object( + proxy_server_module, + "shared_aiohttp_session", + closed_session, + ): + with patch.object( + proxy_server_module, + "_initialize_shared_aiohttp_session", + new_callable=AsyncMock, + side_effect=RuntimeError("connection pool exhausted"), + ): + data = {} + await add_shared_session_to_data(data) + # Should gracefully handle exception — no shared_session attached + assert "shared_session" not in data + + @pytest.mark.asyncio async def test_add_shared_session_no_session_available(): """When no session was ever created, data should not contain shared_session.""" From ef22144854f0248163ff9250060ed49afff23bdc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?d=20=F0=9F=94=B9?= <258577966+voidborne-d@users.noreply.github.com> Date: Tue, 17 Mar 2026 18:07:15 +0000 Subject: [PATCH 5/7] address P2 feedback: add lock docstring warning, remove redundant mock write - Add WARNING docstring to _get_shared_session_lock() about not resetting the lock to None while coroutines may be in the recovery path - Remove redundant proxy_server_module.shared_aiohttp_session assignment in mock_init (add_shared_session_to_data overwrites it synchronously) --- litellm/proxy/route_llm_request.py | 7 ++++++- tests/test_litellm/proxy/test_aiohttp_session_recovery.py | 1 - 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 96e6f705cca..79e8f41972f 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -128,7 +128,12 @@ _shared_session_lock: Optional[asyncio.Lock] = None def _get_shared_session_lock() -> asyncio.Lock: - """Lazily create the shared session lock (must be called within a running event loop).""" + """Lazily create the shared session lock (must be called within a running event loop). + + WARNING: Do not reset _shared_session_lock to None while any coroutine may be + executing the session-recovery path; doing so breaks the double-checked locking + guarantee and can cause duplicate session creation. + """ global _shared_session_lock if _shared_session_lock is None: _shared_session_lock = asyncio.Lock() diff --git a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py index 224a4aa1392..a2b09527962 100644 --- a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py +++ b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py @@ -161,7 +161,6 @@ async def test_add_shared_session_concurrent_recreation_uses_lock(): call_count += 1 # Simulate some async work await asyncio.sleep(0.01) - proxy_server_module.shared_aiohttp_session = new_session return new_session with patch.object( From ca8f5cffa0051dfb9f607b4f4f7fe925312c9a18 Mon Sep 17 00:00:00 2001 From: voidborne-d Date: Tue, 17 Mar 2026 18:52:57 +0000 Subject: [PATCH 6/7] style: apply black formatting to fix CI lint check --- litellm/proxy/route_llm_request.py | 4 +++- .../proxy/test_aiohttp_session_recovery.py | 12 +++--------- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 79e8f41972f..f1590b16c24 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -185,7 +185,9 @@ async def add_shared_session_to_data(data: dict) -> None: "SESSION REUSE: Shared aiohttp session is None after re-check, recreating..." ) try: - new_session = await proxy_server._initialize_shared_aiohttp_session() + new_session = ( + await proxy_server._initialize_shared_aiohttp_session() + ) except Exception: verbose_proxy_logger.exception( "SESSION REUSE: Exception during shared session recreation" diff --git a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py index a2b09527962..f089d10425c 100644 --- a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py +++ b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py @@ -22,9 +22,7 @@ async def test_add_shared_session_attaches_open_session(): mock_session = MagicMock() mock_session.closed = False - with patch( - "litellm.proxy.proxy_server.shared_aiohttp_session", mock_session - ): + with patch("litellm.proxy.proxy_server.shared_aiohttp_session", mock_session): data = {} await add_shared_session_to_data(data) assert data["shared_session"] is mock_session @@ -129,9 +127,7 @@ async def test_add_shared_session_no_session_available(): """When no session was ever created, data should not contain shared_session.""" from litellm.proxy.route_llm_request import add_shared_session_to_data - with patch( - "litellm.proxy.proxy_server.shared_aiohttp_session", None - ): + with patch("litellm.proxy.proxy_server.shared_aiohttp_session", None): data = {} await add_shared_session_to_data(data) assert "shared_session" not in data @@ -175,9 +171,7 @@ async def test_add_shared_session_concurrent_recreation_uses_lock(): ): # Launch 5 concurrent calls results = [{} for _ in range(5)] - await asyncio.gather( - *(add_shared_session_to_data(d) for d in results) - ) + await asyncio.gather(*(add_shared_session_to_data(d) for d in results)) # Only 1 coroutine should have called _initialize (the rest see the # re-checked session as open under the lock) From 88f59e1465ae55097abc3520e1f6d271b5f39bbf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?d=20=F0=9F=94=B9?= <258577966+voidborne-d@users.noreply.github.com> Date: Wed, 18 Mar 2026 00:54:23 +0000 Subject: [PATCH 7/7] fix: use AsyncMock for concurrent test consistency MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address review feedback from greptile — use new_callable=AsyncMock on the concurrent test's patch.object to ensure the mock is properly typed as async, even though side_effect already handles the coroutine. --- tests/test_litellm/proxy/test_aiohttp_session_recovery.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py index f089d10425c..29bd9a491b7 100644 --- a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py +++ b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py @@ -167,6 +167,7 @@ async def test_add_shared_session_concurrent_recreation_uses_lock(): with patch.object( proxy_server_module, "_initialize_shared_aiohttp_session", + new_callable=AsyncMock, side_effect=mock_init, ): # Launch 5 concurrent calls