fix(aiohttp): wire SO_KEEPALIVE socket factory into aiohttp_transport session recreation

litellm/llms/custom_httpx/aiohttp_transport.py

@@ -160,13 +160,33 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
         client: Union[ClientSession, Callable[[], ClientSession]],
         ssl_verify: Optional[Union[bool, ssl.SSLContext]] = None,
         owns_session: bool = True,
+        session_factory: Callable[[], ClientSession] | None = None,
     ):
         self.client = client
         self._ssl_verify = ssl_verify  # Store for per-request SSL override
         super().__init__(client=client, owns_session=owns_session)
-        # Store the client factory for recreating sessions when needed
-        if callable(client):
-            self._client_factory = client
+        # Store the client factory for recreating sessions when needed. An
+        # explicit session_factory wins so a concrete `client` (e.g. a shared
+        # session) still recreates with keep-alive wiring instead of a bare
+        # ClientSession.
+        self._client_factory: Callable[[], ClientSession] | None = (
+            session_factory if session_factory is not None else client if callable(client) else None
+        )
+
+    def _new_session(self) -> ClientSession:
+        """
+        Build a fresh ClientSession, preferring the stored factory (which
+        carries the SO_KEEPALIVE socket factory) and falling back to a
+        keep-alive-aware default rather than a bare aiohttp.ClientSession().
+        """
+        if self._client_factory is not None:
+            return self._client_factory()
+
+        from litellm.llms.custom_httpx.http_handler import (
+            build_default_aiohttp_client_session,
+        )
+
+        return build_default_aiohttp_client_session()
 
     def _get_valid_client_session(self) -> ClientSession:
         """
@@ -179,20 +199,13 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
 
         # If we don't have a client or it's not a ClientSession, create one
         if not isinstance(self.client, ClientSession):
-            if hasattr(self, "_client_factory") and callable(self._client_factory):
-                self.client = self._client_factory()
-            else:
-                self.client = ClientSession()
+            self.client = self._new_session()
             # Don't return yet - check if the newly created session is valid
 
         # Check if the session itself is closed
         if self.client.closed:
             verbose_logger.debug("Session is closed, creating new session")
-            # Create a new session
-            if hasattr(self, "_client_factory") and callable(self._client_factory):
-                self.client = self._client_factory()
-            else:
-                self.client = ClientSession()
+            self.client = self._new_session()
             return self.client
 
         # Check if the existing session is still valid for the current event loop
@@ -215,17 +228,11 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
                     verbose_logger.debug(f"Error closing old session: {e}")
 
                 # Create a new session in the current event loop
-                if hasattr(self, "_client_factory") and callable(self._client_factory):
-                    self.client = self._client_factory()
-                else:
-                    self.client = ClientSession()
+                self.client = self._new_session()
 
         except (RuntimeError, AttributeError):
             # If we can't check the loop or session is invalid, recreate it
-            if hasattr(self, "_client_factory") and callable(self._client_factory):
-                self.client = self._client_factory()
-            else:
-                self.client = ClientSession()
+            self.client = self._new_session()
 
         return self.client
 
@@ -320,10 +327,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
             if "Session is closed" in str(e):
                 verbose_logger.debug(f"Session closed during request, retrying with new session: {e}")
                 # Force creation of a new session
-                if hasattr(self, "_client_factory") and callable(self._client_factory):
-                    self.client = self._client_factory()
-                else:
-                    self.client = ClientSession()
+                self.client = self._new_session()
                 client_session = self.client
 
                 # Retry the request with the new session

litellm/llms/custom_httpx/http_handler.py

@@ -105,6 +105,42 @@ def _build_aiohttp_keepalive_socket_factory() -> Optional[Callable[[Tuple[Any, .
     return factory
 
 
+def _build_aiohttp_transport_connector_kwargs(
+    base_kwargs: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """
+    Build the TCPConnector kwargs shared by every aiohttp session litellm
+    creates, so the SO_KEEPALIVE socket factory and pool tuning are applied
+    consistently across the AsyncHTTPHandler path and the OpenAI/Azure
+    aiohttp_transport fallbacks. socket_factory is only included when
+    keep-alive is enabled and aiohttp is new enough to accept it.
+    """
+    socket_factory = _build_aiohttp_keepalive_socket_factory()
+    return {
+        "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT,
+        "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE,
+        **(base_kwargs or {}),
+        **({"enable_cleanup_closed": True} if AIOHTTP_NEEDS_CLEANUP_CLOSED else {}),
+        **({"limit": AIOHTTP_CONNECTOR_LIMIT} if AIOHTTP_CONNECTOR_LIMIT > 0 else {}),
+        **({"limit_per_host": AIOHTTP_CONNECTOR_LIMIT_PER_HOST} if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0 else {}),
+        **({"socket_factory": socket_factory} if socket_factory is not None else {}),
+    }
+
+
+def build_default_aiohttp_client_session(trust_env: bool = False) -> ClientSession:
+    """
+    Create a ClientSession backed by a keep-alive-aware TCPConnector.
+
+    Used as the fallback session builder on the aiohttp_transport code path
+    (OpenAI/Azure) so recreated sessions still emit TCP keep-alive probes
+    instead of silently dropping back to a bare aiohttp.ClientSession().
+    """
+    return ClientSession(
+        connector=TCPConnector(**_build_aiohttp_transport_connector_kwargs()),
+        trust_env=trust_env,
+    )
+
+
 def get_default_headers() -> dict:
     """
     Get default headers for HTTP requests.
@@ -1013,6 +1049,14 @@ class AsyncHTTPHandler:
 
         verbose_logger.debug("Creating AiohttpTransport...")
 
+        transport_connector_kwargs = _build_aiohttp_transport_connector_kwargs(connector_kwargs)
+
+        def session_factory() -> ClientSession:
+            return ClientSession(
+                connector=TCPConnector(**transport_connector_kwargs),
+                trust_env=trust_env,
+            )
+
         # Use shared session if provided and valid
         if shared_session is not None and not shared_session.closed:
             verbose_logger.debug(f"SHARED SESSION: Reusing existing ClientSession (ID: {id(shared_session)})")
@@ -1020,32 +1064,13 @@ class AsyncHTTPHandler:
                 client=shared_session,
                 ssl_verify=ssl_for_transport,
                 owns_session=False,
+                session_factory=session_factory,
             )
 
         # Create new session only if none provided or existing one is invalid
         verbose_logger.debug("NEW SESSION: Creating new ClientSession (no shared session provided)")
-        transport_connector_kwargs = {
-            "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT,
-            "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE,
-            **connector_kwargs,
-        }
-        if AIOHTTP_NEEDS_CLEANUP_CLOSED:
-            transport_connector_kwargs["enable_cleanup_closed"] = True
-        if AIOHTTP_CONNECTOR_LIMIT > 0:
-            transport_connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT
-        if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0:
-            transport_connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST
-        # Returns None when SO_KEEPALIVE is disabled or aiohttp is too old to
-        # accept socket_factory — version detection lives inside the builder.
-        socket_factory = _build_aiohttp_keepalive_socket_factory()
-        if socket_factory is not None:
-            transport_connector_kwargs["socket_factory"] = socket_factory
-
         return LiteLLMAiohttpTransport(
-            client=lambda: ClientSession(
-                connector=TCPConnector(**transport_connector_kwargs),
-                trust_env=trust_env,
-            ),
+            client=session_factory,
             ssl_verify=ssl_for_transport,
         )
 

tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py

@@ -1,6 +1,8 @@
 import socket
 from unittest.mock import MagicMock, patch
 
+import aiohttp
+
 
 def _invoke_connector_factory(http_handler_module):
     """
@@ -159,3 +161,121 @@ def test_socket_factory_uses_tcp_keepalive_when_keepidle_unavailable(monkeypatch
         setsockopt_calls[(socket.IPPROTO_TCP, fake_socket_module.TCP_KEEPALIVE)] == 60
     )
     assert (socket.IPPROTO_TCP, getattr(socket, "TCP_KEEPIDLE", -1)) not in setsockopt_calls
+
+
+def _last_socket_factory(mock_tcp_connector):
+    """socket_factory kwarg of the most recent TCPConnector(...) construction."""
+    return mock_tcp_connector.call_args.kwargs.get("socket_factory")
+
+
+def test_shared_session_transport_recreates_with_socket_factory(monkeypatch):
+    """
+    Regression for issue #33567: when the aiohttp transport is built around a
+    caller-provided shared ClientSession, it must still recreate stale sessions
+    (closed / cross-loop) through a keep-alive-aware factory. Before the fix the
+    shared-session transport stored no factory, so _get_valid_client_session
+    fell back to a bare aiohttp.ClientSession() with no socket_factory and
+    AIOHTTP_SO_KEEPALIVE had zero effect on the Azure/OpenAI path.
+    """
+    from litellm.llms.custom_httpx import http_handler as http_handler_module
+
+    monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
+    monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
+
+    shared_session = MagicMock(name="shared_session", spec=aiohttp.ClientSession)
+    shared_session.closed = False
+
+    connector_mock = MagicMock(name="connector")
+    session_mock = MagicMock(name="session")
+
+    with patch.object(
+        http_handler_module, "TCPConnector", return_value=connector_mock
+    ) as mock_tcp_connector:
+        with patch.object(
+            http_handler_module, "ClientSession", return_value=session_mock
+        ):
+            transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(
+                shared_session=shared_session
+            )
+            # Transport reuses the shared session as the live client...
+            assert transport.client is shared_session
+            # ...but a keep-alive-aware factory backs every recreation.
+            assert transport._client_factory is not None
+            transport._new_session()
+
+    assert mock_tcp_connector.call_count >= 1
+    assert callable(_last_socket_factory(mock_tcp_connector))
+
+
+def test_new_session_falls_back_to_keepalive_default(monkeypatch):
+    """
+    A transport constructed with a concrete client and no factory must still
+    rebuild through the keep-alive-aware default rather than a bare session.
+    """
+    from litellm.llms.custom_httpx import http_handler as http_handler_module
+    from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport
+
+    monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
+    monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
+
+    # A non-callable client stands in for a concrete ClientSession so no
+    # factory is inferred from it.
+    concrete_session = object()
+    transport = LiteLLMAiohttpTransport(client=concrete_session, session_factory=None)
+    assert transport._client_factory is None
+
+    connector_mock = MagicMock(name="connector")
+    session_mock = MagicMock(name="session")
+
+    with patch.object(
+        http_handler_module, "TCPConnector", return_value=connector_mock
+    ) as mock_tcp_connector:
+        with patch.object(
+            http_handler_module, "ClientSession", return_value=session_mock
+        ):
+            transport._new_session()
+
+    assert mock_tcp_connector.call_count >= 1
+    assert callable(_last_socket_factory(mock_tcp_connector))
+
+
+def test_build_default_session_includes_socket_factory(monkeypatch):
+    from litellm.llms.custom_httpx import http_handler as http_handler_module
+
+    monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
+    monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
+
+    connector_mock = MagicMock(name="connector")
+    session_mock = MagicMock(name="session")
+
+    with patch.object(
+        http_handler_module, "TCPConnector", return_value=connector_mock
+    ) as mock_tcp_connector:
+        with patch.object(
+            http_handler_module, "ClientSession", return_value=session_mock
+        ):
+            http_handler_module.build_default_aiohttp_client_session()
+
+    assert mock_tcp_connector.call_count >= 1
+    assert callable(_last_socket_factory(mock_tcp_connector))
+
+
+def test_build_default_session_omits_socket_factory_when_disabled(monkeypatch):
+    from litellm.llms.custom_httpx import http_handler as http_handler_module
+
+    monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", False)
+    monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
+
+    connector_mock = MagicMock(name="connector")
+    session_mock = MagicMock(name="session")
+
+    with patch.object(
+        http_handler_module, "TCPConnector", return_value=connector_mock
+    ) as mock_tcp_connector:
+        with patch.object(
+            http_handler_module, "ClientSession", return_value=session_mock
+        ):
+            http_handler_module.build_default_aiohttp_client_session()
+
+    assert mock_tcp_connector.call_count >= 1
+    assert "socket_factory" not in mock_tcp_connector.call_args.kwargs