fix(aiohttp): only set enable_cleanup_closed when required (#21897)

* fix(aiohttp): only set enable_cleanup_closed when required

* add tests

litellm/llms/custom_httpx/http_handler.py

@@ -28,6 +28,7 @@ from litellm.constants import (
     AIOHTTP_CONNECTOR_LIMIT,
     AIOHTTP_CONNECTOR_LIMIT_PER_HOST,
     AIOHTTP_KEEPALIVE_TIMEOUT,
+    AIOHTTP_NEEDS_CLEANUP_CLOSED,
     AIOHTTP_TTL_DNS_CACHE,
     DEFAULT_SSL_CIPHERS,
 )
@@ -876,9 +877,10 @@ class AsyncHTTPHandler:
         transport_connector_kwargs = {
             "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT,
             "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE,
-            "enable_cleanup_closed": True,
             **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:

litellm/proxy/proxy_server.py

@@ -714,8 +714,9 @@ async def _initialize_shared_aiohttp_session():
         connector_kwargs = {
             "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT,
             "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE,
-            "enable_cleanup_closed": True,
         }
+        if AIOHTTP_NEEDS_CLEANUP_CLOSED:
+            connector_kwargs["enable_cleanup_closed"] = True
         if AIOHTTP_CONNECTOR_LIMIT > 0:
             connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT
         if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0:

tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py

@@ -0,0 +1,31 @@
+from unittest.mock import MagicMock, patch
+
+
+def test_create_aiohttp_transport_sets_enable_cleanup_closed_when_needed(monkeypatch):
+    from litellm.llms.custom_httpx import http_handler as http_handler_module
+
+    connector_mock = MagicMock(name="connector")
+    session_mock = MagicMock(name="session")
+    monkeypatch.setattr(http_handler_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", True)
+
+    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=None)
+            transport._get_valid_client_session()
+
+    assert mock_tcp_connector.call_args.kwargs["enable_cleanup_closed"] is True
+
+
+def test_create_aiohttp_transport_omits_enable_cleanup_closed_when_not_needed(monkeypatch):
+    from litellm.llms.custom_httpx import http_handler as http_handler_module
+
+    connector_mock = MagicMock(name="connector")
+    session_mock = MagicMock(name="session")
+    monkeypatch.setattr(http_handler_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", False)
+
+    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=None)
+            transport._get_valid_client_session()
+
+    assert "enable_cleanup_closed" not in mock_tcp_connector.call_args.kwargs

tests/test_litellm/proxy/test_aiohttp_cleanup_closed.py

@@ -0,0 +1,34 @@
+import asyncio
+from unittest.mock import MagicMock, patch
+
+
+def test_initialize_shared_aiohttp_session_sets_enable_cleanup_closed_when_needed(
+    monkeypatch,
+):
+    from litellm.proxy import proxy_server as proxy_server_module
+
+    connector_mock = MagicMock(name="connector")
+    session_mock = MagicMock(name="session")
+    monkeypatch.setattr(proxy_server_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", True)
+
+    with patch("aiohttp.TCPConnector", return_value=connector_mock) as mock_tcp_connector:
+        with patch("aiohttp.ClientSession", return_value=session_mock):
+            asyncio.run(proxy_server_module._initialize_shared_aiohttp_session())
+
+    assert mock_tcp_connector.call_args.kwargs["enable_cleanup_closed"] is True
+
+
+def test_initialize_shared_aiohttp_session_omits_enable_cleanup_closed_when_not_needed(
+    monkeypatch,
+):
+    from litellm.proxy import proxy_server as proxy_server_module
+
+    connector_mock = MagicMock(name="connector")
+    session_mock = MagicMock(name="session")
+    monkeypatch.setattr(proxy_server_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", False)
+
+    with patch("aiohttp.TCPConnector", return_value=connector_mock) as mock_tcp_connector:
+        with patch("aiohttp.ClientSession", return_value=session_mock):
+            asyncio.run(proxy_server_module._initialize_shared_aiohttp_session())
+
+    assert "enable_cleanup_closed" not in mock_tcp_connector.call_args.kwargs