From 7b66c970e9706a5b96c1987bc445ff941e51fb66 Mon Sep 17 00:00:00 2001 From: voidborne-d Date: Tue, 17 Mar 2026 03:11:58 +0000 Subject: [PATCH] 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