Merge pull request #23808 from voidborne-d/fix/shared-aiohttp-session-auto-recovery

fix: auto-recover shared aiohttp session when closed
(cherry picked from commit cec3e9e7d4)

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 []

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