Merge pull request #26730 from yassinkortam/fix/http-handler-keepalive

fix: add optional TCP SO_KEEPALIVE support to aiohttp's TCPConnector
2026-04-29 10:10:59 -07:00 · 2026-04-29 10:10:59 -07:00 · 9b3cd5ca25
commit 9b3cd5ca25
parent ea275659ac 848b79acb5
4 changed files with 240 additions and 0 deletions
--- a/litellm/constants.py
+++ b/litellm/constants.py
@ -224,6 +224,16 @@ AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(
 )
 AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120))
 AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300))
+# TCP keep-alive (SO_KEEPALIVE) — opt-in. Required when running behind NAT/LBs
+# whose idle timeout is shorter than provider response timeouts (e.g. AWS NAT
+# Gateway: 350s vs OpenAI/Azure: 600s). Without this, the kernel sends nothing
+# during a long provider call and the NAT reaps the flow before the response
+# arrives. Enabling SO_KEEPALIVE makes the kernel emit TCP probes that reset
+# the NAT idle timer.
+AIOHTTP_SO_KEEPALIVE = os.getenv("AIOHTTP_SO_KEEPALIVE", "False").lower() == "true"
+AIOHTTP_TCP_KEEPIDLE = int(os.getenv("AIOHTTP_TCP_KEEPIDLE", 60))
+AIOHTTP_TCP_KEEPINTVL = int(os.getenv("AIOHTTP_TCP_KEEPINTVL", 30))
+AIOHTTP_TCP_KEEPCNT = int(os.getenv("AIOHTTP_TCP_KEEPCNT", 5))
 # enable_cleanup_closed is only needed for Python versions with the SSL leak bug
 # Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960)
 # Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78
--- a/litellm/llms/custom_httpx/http_handler.py
+++ b/litellm/llms/custom_httpx/http_handler.py
@ -1,5 +1,7 @@
 import asyncio
+import inspect
 import os
+import socket
 import ssl
 import sys
 import time
@ -29,6 +31,10 @@ from litellm.constants import (
    AIOHTTP_CONNECTOR_LIMIT_PER_HOST,
    AIOHTTP_KEEPALIVE_TIMEOUT,
    AIOHTTP_NEEDS_CLEANUP_CLOSED,
+    AIOHTTP_SO_KEEPALIVE,
+    AIOHTTP_TCP_KEEPCNT,
+    AIOHTTP_TCP_KEEPIDLE,
+    AIOHTTP_TCP_KEEPINTVL,
    AIOHTTP_TTL_DNS_CACHE,
    COMPLETION_HTTP_FALLBACK_SECONDS,
    DEFAULT_SSL_CIPHERS,
@ -54,6 +60,57 @@ except Exception:
    version = "0.0.0"


+# aiohttp 3.10+ exposes a `socket_factory` kwarg on TCPConnector. Older
+# versions don't — detect once and skip the keep-alive wiring there.
+# https://docs.aiohttp.org/en/stable/client_reference.html#aiohttp.TCPConnector
+_AIOHTTP_SUPPORTS_SOCKET_FACTORY = (
+    "socket_factory" in inspect.signature(TCPConnector.__init__).parameters
+)
+
+
+def _build_aiohttp_keepalive_socket_factory() -> (
+    Optional[Callable[[Tuple[Any, ...]], socket.socket]]
+):
+    """
+    Build a socket_factory that enables SO_KEEPALIVE on aiohttp TCP sockets.
+
+    Why: by default, aiohttp creates sockets without SO_KEEPALIVE, so the kernel
+    sends nothing during a long idle TCP connection. NAT/LB hops (e.g. AWS NAT
+    Gateway, 350s idle timeout) reap the flow well before slow provider
+    responses (OpenAI/Azure: up to 600s) arrive. Enabling SO_KEEPALIVE makes
+    the kernel emit TCP probes that reset the NAT idle timer.
+
+    Returns None when AIOHTTP_SO_KEEPALIVE is disabled or aiohttp is too old.
+    """
+    if not AIOHTTP_SO_KEEPALIVE or not _AIOHTTP_SUPPORTS_SOCKET_FACTORY:
+        return None
+
+    def factory(addr_info: Tuple[Any, ...]) -> socket.socket:
+        family, type_, proto = addr_info[0], addr_info[1], addr_info[2]
+        sock = socket.socket(family=family, type=type_, proto=proto)
+        sock.setblocking(False)
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+        # Linux: TCP_KEEPIDLE is idle-before-first-probe.
+        # macOS/Darwin: TCP_KEEPALIVE is the equivalent.
+        if hasattr(socket, "TCP_KEEPIDLE"):
+            sock.setsockopt(
+                socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, AIOHTTP_TCP_KEEPIDLE
+            )
+        elif hasattr(socket, "TCP_KEEPALIVE"):
+            sock.setsockopt(
+                socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, AIOHTTP_TCP_KEEPIDLE
+            )
+        if hasattr(socket, "TCP_KEEPINTVL"):
+            sock.setsockopt(
+                socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, AIOHTTP_TCP_KEEPINTVL
+            )
+        if hasattr(socket, "TCP_KEEPCNT"):
+            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, AIOHTTP_TCP_KEEPCNT)
+        return sock
+
+    return factory
+
+
 def get_default_headers() -> dict:
    """
    Get default headers for HTTP requests.
@ -935,6 +992,11 @@ class AsyncHTTPHandler:
            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(
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@ -742,6 +742,10 @@ async def _initialize_shared_aiohttp_session():
    try:
        from aiohttp import ClientSession, TCPConnector

+        from litellm.llms.custom_httpx.http_handler import (
+            _build_aiohttp_keepalive_socket_factory,
+        )
+
        connector_kwargs: Dict[str, Any] = {
            "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT,
            "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE,
@ -752,6 +756,9 @@ async def _initialize_shared_aiohttp_session():
            connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT
        if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0:
            connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST
+        socket_factory = _build_aiohttp_keepalive_socket_factory()
+        if socket_factory is not None:
+            connector_kwargs["socket_factory"] = socket_factory

        connector = TCPConnector(**connector_kwargs)
        session = ClientSession(connector=connector)
--- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py
+++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py
@ -0,0 +1,161 @@
+import socket
+from unittest.mock import MagicMock, patch
+
+
+def _invoke_connector_factory(http_handler_module):
+    """
+    Drive the lambda factory installed on the transport so TCPConnector is
+    actually constructed. _create_aiohttp_transport returns a transport whose
+    _client_factory is the lambda that builds (TCPConnector → ClientSession);
+    invoking it directly avoids relying on _get_valid_client_session's internal
+    branching to trigger connector construction.
+    """
+    transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(
+        shared_session=None
+    )
+    transport._client_factory()
+    return transport
+
+
+def test_socket_factory_omitted_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
+        ):
+            _invoke_connector_factory(http_handler_module)
+
+    assert mock_tcp_connector.call_count >= 1
+    assert "socket_factory" not in mock_tcp_connector.call_args.kwargs
+
+
+def test_socket_factory_attached_when_enabled(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
+        ):
+            _invoke_connector_factory(http_handler_module)
+
+    assert mock_tcp_connector.call_count >= 1
+    factory = mock_tcp_connector.call_args.kwargs.get("socket_factory")
+    assert callable(factory)
+
+
+def test_socket_factory_skipped_on_old_aiohttp(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", 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
+        ):
+            _invoke_connector_factory(http_handler_module)
+
+    assert mock_tcp_connector.call_count >= 1
+    assert "socket_factory" not in mock_tcp_connector.call_args.kwargs
+
+
+def test_socket_factory_sets_keepalive_options(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)
+    monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 45)
+    monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPINTVL", 15)
+    monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPCNT", 4)
+
+    factory = http_handler_module._build_aiohttp_keepalive_socket_factory()
+    assert factory is not None
+
+    addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0))
+
+    fake_sock = MagicMock(spec=socket.socket)
+    with patch("socket.socket", return_value=fake_sock) as sock_ctor:
+        returned = factory(addr_info)
+
+    sock_ctor.assert_called_once_with(
+        family=socket.AF_INET, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP
+    )
+    assert returned is fake_sock
+    fake_sock.setblocking.assert_called_once_with(False)
+
+    setsockopt_calls = {
+        (call.args[0], call.args[1]): call.args[2]
+        for call in fake_sock.setsockopt.call_args_list
+    }
+    assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1
+
+    if hasattr(socket, "TCP_KEEPIDLE"):
+        assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE)] == 45
+    elif hasattr(socket, "TCP_KEEPALIVE"):
+        assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE)] == 45
+    if hasattr(socket, "TCP_KEEPINTVL"):
+        assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL)] == 15
+    if hasattr(socket, "TCP_KEEPCNT"):
+        assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPCNT)] == 4
+
+
+def test_socket_factory_uses_tcp_keepalive_when_keepidle_unavailable(monkeypatch):
+    """
+    Cover the macOS/Darwin branch: when TCP_KEEPIDLE is missing but TCP_KEEPALIVE
+    is present, the factory should fall back to TCP_KEEPALIVE for the idle timer.
+    Linux CI runners always have TCP_KEEPIDLE, so we patch socket itself to
+    simulate the BSD-derived environment.
+    """
+    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)
+    monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 60)
+
+    factory = http_handler_module._build_aiohttp_keepalive_socket_factory()
+    assert factory is not None
+
+    fake_socket_module = MagicMock(spec=[])
+    fake_socket_module.SOL_SOCKET = socket.SOL_SOCKET
+    fake_socket_module.SO_KEEPALIVE = socket.SO_KEEPALIVE
+    fake_socket_module.IPPROTO_TCP = socket.IPPROTO_TCP
+    fake_socket_module.TCP_KEEPALIVE = getattr(socket, "TCP_KEEPALIVE", 0x10)
+    fake_sock = MagicMock(spec=socket.socket)
+    fake_socket_module.socket = MagicMock(return_value=fake_sock)
+
+    addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0))
+
+    with patch.object(http_handler_module, "socket", fake_socket_module):
+        factory(addr_info)
+
+    setsockopt_calls = {
+        (call.args[0], call.args[1]): call.args[2]
+        for call in fake_sock.setsockopt.call_args_list
+    }
+    assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1
+    assert (
+        setsockopt_calls[(socket.IPPROTO_TCP, fake_socket_module.TCP_KEEPALIVE)] == 60
+    )
+    assert (socket.IPPROTO_TCP, getattr(socket, "TCP_KEEPIDLE", -1)) not in setsockopt_calls