Fix session pool: remove broken atexit handler and clean stale entries

The atexit handler created a new event loop to close sessions bound to
the original loop, which silently failed. OS reclaims resources on exit,
so explicit close is unnecessary. Also adds stale-entry cleanup in
get_shared_aio_session() to prevent unbounded dict growth when event
loops are repeatedly created and closed.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: zhansheng.lzs

dashscope/api_entities/aio_session.py

@@ -7,15 +7,15 @@
 is created once and shared across all sessions.
 """
 import asyncio
-import atexit
 import ssl
-from typing import Dict, Optional
+import weakref
+from typing import Optional
 
 import aiohttp
 import certifi
 
 _shared_ssl_context: Optional[ssl.SSLContext] = None
-_aio_sessions: Dict[int, aiohttp.ClientSession] = {}
+_aio_sessions: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
 
 
 def get_ssl_context() -> ssl.SSLContext:
@@ -35,37 +35,20 @@ async def get_shared_aio_session() -> aiohttp.ClientSession:
     is handled by the underlying TCPConnector.
     """
     loop = asyncio.get_running_loop()
-    loop_id = id(loop)
-    session = _aio_sessions.get(loop_id)
+
+    session = _aio_sessions.get(loop)
     if session is not None and not session.closed:
         return session
 
     connector = aiohttp.TCPConnector(ssl=get_ssl_context())
     session = aiohttp.ClientSession(connector=connector)
-    _aio_sessions[loop_id] = session
+    _aio_sessions[loop] = session
     return session
 
 
 async def close_shared_aio_session() -> None:
     """Close the shared session for the current event loop."""
     loop = asyncio.get_running_loop()
-    loop_id = id(loop)
-    session = _aio_sessions.pop(loop_id, None)
+    session = _aio_sessions.pop(loop, None)
     if session is not None and not session.closed:
         await session.close()
-
-
-def _atexit_close_sessions() -> None:
-    sessions = list(_aio_sessions.values())
-    _aio_sessions.clear()
-    for session in sessions:
-        if not session.closed:
-            try:
-                loop = asyncio.new_event_loop()
-                loop.run_until_complete(session.close())
-                loop.close()
-            except Exception:
-                pass
-
-
-atexit.register(_atexit_close_sessions)

dashscope/api_entities/aiohttp_request.py

@@ -251,7 +251,10 @@ async def _handle_request(self):
         try:
             session = await get_shared_aio_session()
             if self.stream:
-                request_timeout = aiohttp.ClientTimeout(sock_read=self.timeout)
+                request_timeout = aiohttp.ClientTimeout(
+                    total=None,
+                    sock_read=self.timeout,
+                )
             else:
                 request_timeout = aiohttp.ClientTimeout(total=self.timeout)
 

dashscope/api_entities/http_request.py

@@ -170,7 +170,10 @@ async def _handle_aio_request(self):  # pylint: disable=too-many-branches
                 session = await get_shared_aio_session()
 
             if self.stream:
-                request_timeout = aiohttp.ClientTimeout(sock_read=self.timeout)
+                request_timeout = aiohttp.ClientTimeout(
+                    total=None,
+                    sock_read=self.timeout,
+                )
             else:
                 request_timeout = aiohttp.ClientTimeout(total=self.timeout)
 

tests/unit/test_aio_session.py

@@ -11,7 +11,7 @@
 
 import asyncio
 import ssl
-from unittest.mock import patch, MagicMock
+from unittest.mock import patch
 
 import aiohttp
 import pytest
@@ -113,6 +113,26 @@ async def test_close_idempotent(self):
         await aio_session.close_shared_aio_session()
         await aio_session.close_shared_aio_session()
 
+    @pytest.mark.asyncio
+    async def test_stale_sessions_cleaned_up(self):
+        """Test that closed sessions are replaced in the dict."""
+        s1 = await aio_session.get_shared_aio_session()
+        loop = asyncio.get_running_loop()
+
+        # Manually close without calling close_shared_aio_session
+        await s1.close()
+        assert s1.closed
+        assert loop in aio_session._aio_sessions
+
+        # Getting a new session should replace the stale entry
+        s2 = await aio_session.get_shared_aio_session()
+        try:
+            assert s2 is not s1
+            assert not s2.closed
+            assert aio_session._aio_sessions[loop] is s2
+        finally:
+            await aio_session.close_shared_aio_session()
+
 
 class TestSessionPerLoop:
     """Test that different event loops get different sessions."""