diff --git a/litellm/llms/custom_httpx/aiohttp_connector.py b/litellm/llms/custom_httpx/aiohttp_connector.py new file mode 100644 index 00000000000..2e8add7cfaf --- /dev/null +++ b/litellm/llms/custom_httpx/aiohttp_connector.py @@ -0,0 +1,40 @@ +from typing import TYPE_CHECKING + +from aiohttp import TCPConnector + +if TYPE_CHECKING: + from aiohttp.client_reqrep import ConnectionKey + from aiohttp.connector import Connection + from aiohttp.tracing import Trace + + +class HardenedTCPConnector(TCPConnector): + """ + TCPConnector that refuses to hand out keepalive connections which have been + flagged for closing while sitting idle in the pool. + + aiohttp's BaseConnector._get only checks that a pooled connection is still + connected and within the keepalive window; it does not re-check the + protocol's should_close flag. A connection is only pooled while should_close + is False, but the flag can flip to True afterwards - most notably via the + aiohttp 3.14.x regression (aio-libs/aiohttp#12953) where a stray sock_read + timer fires on an already-released connection, stamps a SocketTimeoutError on + the ResponseHandler and sets should_close without closing the transport. The + next request then reuses the poisoned connection and fails instantly with a + sub-millisecond "Connection timed out", spanning every provider that shares + the pool. + + Re-checking should_close at acquisition time drops such connections and keeps + pulling until a clean one (or none) is found, independent of the aiohttp + version in use. + """ + + async def _get(self, key: "ConnectionKey", traces: "list[Trace]") -> "Connection | None": + conn = await super()._get(key, traces) + if conn is None: + return None + protocol = conn.protocol + if protocol is not None and protocol.should_close: + conn.close() + return await self._get(key, traces) + return conn diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 5cec763bb5d..7edda674002 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -45,6 +45,7 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) +from litellm.llms.custom_httpx.aiohttp_connector import HardenedTCPConnector from litellm.types.llms.custom_http import * if TYPE_CHECKING: @@ -1043,7 +1044,7 @@ class AsyncHTTPHandler: return LiteLLMAiohttpTransport( client=lambda: ClientSession( - connector=TCPConnector(**transport_connector_kwargs), + connector=HardenedTCPConnector(**transport_connector_kwargs), trust_env=trust_env, ), ssl_verify=ssl_for_transport, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3b640ee54fd..74e2daf70bc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -809,8 +809,9 @@ async def proxy_shutdown_event(): async def _initialize_shared_aiohttp_session(): """Initialize shared aiohttp session for connection reuse with connection limits.""" try: - from aiohttp import ClientSession, TCPConnector + from aiohttp import ClientSession + from litellm.llms.custom_httpx.aiohttp_connector import HardenedTCPConnector from litellm.llms.custom_httpx.http_handler import ( _build_aiohttp_keepalive_socket_factory, ) @@ -829,7 +830,7 @@ async def _initialize_shared_aiohttp_session(): if socket_factory is not None: connector_kwargs["socket_factory"] = socket_factory - connector = TCPConnector(**connector_kwargs) + connector = HardenedTCPConnector(**connector_kwargs) session = ClientSession(connector=connector) verbose_proxy_logger.info( diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py index f279acfd60c..26a1187a939 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py @@ -9,7 +9,7 @@ def test_create_aiohttp_transport_sets_enable_cleanup_closed_when_needed(monkeyp monkeypatch.setattr(http_handler_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", True) with patch.object( - http_handler_module, "TCPConnector", return_value=connector_mock + http_handler_module, "HardenedTCPConnector", return_value=connector_mock ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock @@ -32,7 +32,7 @@ def test_create_aiohttp_transport_omits_enable_cleanup_closed_when_not_needed( monkeypatch.setattr(http_handler_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", False) with patch.object( - http_handler_module, "TCPConnector", return_value=connector_mock + http_handler_module, "HardenedTCPConnector", return_value=connector_mock ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_connector.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_connector.py new file mode 100644 index 00000000000..759e4added4 --- /dev/null +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_connector.py @@ -0,0 +1,134 @@ +import collections +from time import monotonic + +import pytest +from aiohttp.client_proto import ResponseHandler +from aiohttp.client_reqrep import ConnectionKey + +from litellm.llms.custom_httpx.aiohttp_connector import HardenedTCPConnector + + +class _FakeTransport: + """Minimal transport stand-in: stays 'open' until close()/abort().""" + + def __init__(self): + self._closing = False + + def is_closing(self): + return self._closing + + def close(self): + self._closing = True + + def abort(self): + self._closing = True + + def get_extra_info(self, name, default=None): + return default + + +def _make_key(host="api.openai.com"): + return ConnectionKey( + host=host, + port=443, + is_ssl=True, + ssl=True, + proxy=None, + proxy_auth=None, + proxy_headers_hash=None, + ) + + +def _make_pooled_proto(loop): + proto = ResponseHandler(loop=loop) + proto.transport = _FakeTransport() + return proto + + +def _pool(connector, key, proto): + connector._conns[key] = collections.deque([(proto, monotonic())]) + + +@pytest.mark.asyncio +async def test_get_drops_connection_poisoned_after_pooling(): + """ + Regression for aiohttp 3.14.x connection-pool poisoning (litellm #33820). + + A connection is only pooled while should_close is False. If a stray + sock_read timer later fires on the idle pooled connection it stamps a + timeout exception and flips should_close True *without* closing the + transport. Vanilla aiohttp _get would still hand it out (it only checks + is_connected + keepalive window). HardenedTCPConnector must refuse it. + """ + import asyncio + + loop = asyncio.get_running_loop() + connector = HardenedTCPConnector(keepalive_timeout=120) + try: + key = _make_key() + proto = _make_pooled_proto(loop) + + assert proto.should_close is False + _pool(connector, key, proto) + + # Emulate the post-release sock_read timer firing on the pooled conn. + proto._on_read_timeout() + assert proto.should_close is True + assert proto.is_connected() is True # transport was NOT closed + + conn = await connector._get(key, []) + assert conn is None, "poisoned connection must not be reused" + assert proto.is_connected() is False, "poisoned connection must be closed" + finally: + connector._conns.clear() + + +@pytest.mark.asyncio +async def test_get_reuses_healthy_connection(): + """A clean pooled connection is still reused - the fix must not break keepalive.""" + import asyncio + + loop = asyncio.get_running_loop() + connector = HardenedTCPConnector(keepalive_timeout=120) + try: + key = _make_key() + proto = _make_pooled_proto(loop) + _pool(connector, key, proto) + + conn = await connector._get(key, []) + assert conn is not None + assert conn.protocol is proto + finally: + connector._conns.clear() + + +@pytest.mark.asyncio +async def test_get_skips_poisoned_and_returns_next_healthy(): + """With a poisoned then a healthy conn queued, _get skips past the poison.""" + import asyncio + + loop = asyncio.get_running_loop() + connector = HardenedTCPConnector(keepalive_timeout=120) + try: + key = _make_key() + poisoned = _make_pooled_proto(loop) + healthy = _make_pooled_proto(loop) + connector._conns[key] = collections.deque( + [(poisoned, monotonic()), (healthy, monotonic())] + ) + + poisoned._on_read_timeout() + assert poisoned.should_close is True + + conn = await connector._get(key, []) + assert conn is not None + assert conn.protocol is healthy + assert poisoned.is_connected() is False + finally: + connector._conns.clear() + + +def test_http_handler_uses_hardened_connector(): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + assert http_handler_module.HardenedTCPConnector is HardenedTCPConnector 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..944cdcd2faf 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 @@ -27,7 +27,7 @@ def test_socket_factory_omitted_when_disabled(monkeypatch): session_mock = MagicMock(name="session") with patch.object( - http_handler_module, "TCPConnector", return_value=connector_mock + http_handler_module, "HardenedTCPConnector", return_value=connector_mock ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock @@ -48,7 +48,7 @@ def test_socket_factory_attached_when_enabled(monkeypatch): session_mock = MagicMock(name="session") with patch.object( - http_handler_module, "TCPConnector", return_value=connector_mock + http_handler_module, "HardenedTCPConnector", return_value=connector_mock ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock @@ -70,7 +70,7 @@ def test_socket_factory_skipped_on_old_aiohttp(monkeypatch): session_mock = MagicMock(name="session") with patch.object( - http_handler_module, "TCPConnector", return_value=connector_mock + http_handler_module, "HardenedTCPConnector", return_value=connector_mock ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock