fix: reap aiohttp sessions and auth locks belonging to closed event loops

google/genai/_api_client.py

@@ -953,6 +953,20 @@ def _is_session_closed(self, session: Any) -> bool:
       return bool(session._auth_request._closed)
     return False
 
+  @staticmethod
+  def _discard_closed_loop_entries(entries: dict[Any, Any]) -> None:
+    """Drops entries of a loop-keyed cache whose event loop has been closed.
+
+    Servers that run every request on a fresh loop (ADK's sync `Runner.run()`,
+    used by Agent Engine, calls `asyncio.run()` per request) would otherwise
+    grow these caches without bound, pinning one aiohttp session, connector and
+    set of sockets per request served. A session on a closed loop cannot be
+    awaited shut, but dropping the last reference to it lets the garbage
+    collector release the connector and its sockets.
+    """
+    for loop in [loop for loop in entries if loop.is_closed()]:
+      del entries[loop]
+
   @property
   def _aiohttp_session(
       self,
@@ -976,6 +990,7 @@ async def _get_aiohttp_session(
     loop = asyncio.get_running_loop()
 
     with self._sync_auth_lock:
+      self._discard_closed_loop_entries(self._aiohttp_sessions)
       session = self._aiohttp_sessions.get(loop)
       if session is not None and self._is_session_closed(session):
         session = None
@@ -1319,6 +1334,7 @@ async def _get_async_auth_lock(self) -> asyncio.Lock:
     """
     loop = asyncio.get_running_loop()
     with self._sync_auth_lock:
+      self._discard_closed_loop_entries(self._async_auth_locks)
       if loop not in self._async_auth_locks:
         self._async_auth_locks[loop] = asyncio.Lock()
       return self._async_auth_locks[loop]

google/genai/tests/client/test_client_initialization.py

@@ -22,6 +22,7 @@
 import os
 import ssl
 import sys
+import threading
 from unittest import mock
 
 import certifi
@@ -1676,8 +1677,6 @@ async def mock_async_operation(op_id: int):
 @pytest.mark.asyncio
 async def test_get_async_auth_lock_creation_lock_lifecycle():
   """Tests the creation lock lifecycle and cleanup."""
-  import threading
-
   client = Client(
       vertexai=True, project="fake_project_id", location="fake-location"
   )
@@ -1947,3 +1946,83 @@ async def test_async_mtls_uses_refreshable_credentials(monkeypatch):
     assert passed_creds.valid == True
     mock_creds.expired = True
     assert passed_creds.valid == False
+
+
+def _run_on_fresh_loops(coro_fn, count):
+  """Runs coro_fn() count times, each on its own thread and its own loop.
+
+  This is the shape ADK's sync `Runner.run()` gives Agent Engine: every request
+  gets a new thread and a new `asyncio.run()` loop that is closed on the way
+  out, while the genai client itself is a long-lived singleton.
+
+  Args:
+    coro_fn: Zero-argument callable returning the coroutine to run.
+    count: How many loops to run it on, one after another.
+  """
+  for _ in range(count):
+    thread = threading.Thread(target=lambda: asyncio.run(coro_fn()))
+    thread.start()
+    thread.join()
+
+
+@requires_aiohttp
+def test_aiohttp_sessions_not_retained_for_closed_event_loops():
+  """Sessions belonging to finished loops must not accumulate. b/496663148."""
+  client = Client(
+      vertexai=True, project="fake_project_id", location="fake-location"
+  )
+  api_client.has_aiohttp = True
+  base_client = client._api_client
+
+  _run_on_fresh_loops(base_client._get_aiohttp_session, 10)
+
+  # Reaping happens on access, so the newest loop's entry survives until the
+  # next call. What matters is that the cache stays bounded by the number of
+  # live loops rather than growing once per request served.
+  assert len(base_client._aiohttp_sessions) <= 1
+
+
+def test_async_auth_locks_not_retained_for_closed_event_loops():
+  """Auth locks belonging to finished loops must not accumulate. b/496663148."""
+  client = Client(
+      vertexai=True, project="fake_project_id", location="fake-location"
+  )
+  base_client = client._api_client
+
+  _run_on_fresh_loops(base_client._get_async_auth_lock, 10)
+
+  assert len(base_client._async_auth_locks) <= 1
+
+
+@requires_aiohttp
+def test_aiohttp_session_kept_per_live_event_loop():
+  """Each live loop keeps its own session; reaping must not steal it.
+
+  Reusing one loop's session on another is what raised `RuntimeError: ... got
+  Future ... attached to a different loop` in b/496663148.
+  """
+  client = Client(
+      vertexai=True, project="fake_project_id", location="fake-location"
+  )
+  api_client.has_aiohttp = True
+  base_client = client._api_client
+  sessions = []
+  barrier = threading.Barrier(3)
+
+  def hold_loop_open():
+    async def run():
+      session = await base_client._get_aiohttp_session()
+      sessions.append(session)
+      # Keep this loop alive until every thread has its own session.
+      await asyncio.get_running_loop().run_in_executor(None, barrier.wait)
+      assert await base_client._get_aiohttp_session() is session
+
+    asyncio.run(run())
+
+  threads = [threading.Thread(target=hold_loop_open) for _ in range(3)]
+  for thread in threads:
+    thread.start()
+  for thread in threads:
+    thread.join()
+
+  assert len({id(session) for session in sessions}) == 3