diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index ba3d19ef458..8682f7c486c 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 @@ -103,30 +104,99 @@ def get_team_id_from_data(data: dict) -> Optional[str]: return None -def add_shared_session_to_data(data: dict) -> 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). + + 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() + 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: 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 + # 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 + + # 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 + 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 - 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( @@ -195,7 +265,7 @@ async def route_request( """ 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..29bd9a491b7 --- /dev/null +++ b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py @@ -0,0 +1,182 @@ +""" +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.""" + 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 + + 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.""" + 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, + 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_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 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 + + +@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) + return new_session + + 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=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