diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 3172d3667e1..87355acca4c 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/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 diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 5cec763bb5d..e4e348b199f 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/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, ) diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py index 5515e6ce815..48e4b87a222 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py +++ b/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