diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py index 27adcf813..172dd0fd4 100644 --- a/google/genai/_api_client.py +++ b/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] diff --git a/google/genai/tests/client/test_client_initialization.py b/google/genai/tests/client/test_client_initialization.py index 659d1ca4b..d353db2d0 100644 --- a/google/genai/tests/client/test_client_initialization.py +++ b/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