diff --git a/src/google/adk/models/apigee_llm.py b/src/google/adk/models/apigee_llm.py index a4b57ce688..3866272712 100644 --- a/src/google/adk/models/apigee_llm.py +++ b/src/google/adk/models/apigee_llm.py @@ -20,6 +20,7 @@ import collections.abc import enum from functools import cached_property +from functools import partial import json import logging import os @@ -28,6 +29,7 @@ from typing import Generator from typing import Optional from typing import TYPE_CHECKING +import weakref from google.adk import version as adk_version from google.genai import types @@ -402,6 +404,11 @@ def _validate_model_string(model: str) -> bool: return False +# Keeps a strong reference to fire-and-forget cleanup tasks so they are not +# garbage collected before they finish (see CompletionsHTTPClient._cleanup_client). +_CLEANUP_TASKS: set[asyncio.Task] = set() + + class CompletionsHTTPClient: """A generic HTTP client for completions, compatible with OpenAI API.""" @@ -427,17 +434,29 @@ def _client(self) -> httpx.AsyncClient: timeout=None, follow_redirects=True, ) - atexit.register(self._cleanup_client, client) + # Register with a weakref.proxy so the atexit registry does not keep the + # client (and its connection pool) alive for the whole process. Keep the + # bound callback per instance so close()/aclose() can unregister just this + # client's handler without affecting other CompletionsHTTPClient instances. + self._atexit_callback = partial(self._cleanup_client, weakref.proxy(client)) + atexit.register(self._atexit_callback) return client @staticmethod def _cleanup_client(client: httpx.AsyncClient) -> None: """Cleans up the httpx client.""" - if client.is_closed: + try: + if client.is_closed: + return + except ReferenceError: + # The client was already garbage collected via its weakref.proxy. return try: loop = asyncio.get_running_loop() - loop.create_task(client.aclose()) + task = loop.create_task(client.aclose()) + # Retain a reference so the task is not garbage collected while pending. + _CLEANUP_TASKS.add(task) + task.add_done_callback(_CLEANUP_TASKS.discard) except RuntimeError: try: # This fails if asyncio.run is already called in main and is closing. @@ -449,10 +468,12 @@ def close(self) -> None: if '_client' not in self.__dict__: return self._cleanup_client(self._client) + atexit.unregister(self._atexit_callback) async def aclose(self) -> None: if '_client' not in self.__dict__: return + atexit.unregister(self._atexit_callback) if self._client.is_closed: return await self._client.aclose() diff --git a/tests/unittests/models/test_completions_http_client.py b/tests/unittests/models/test_completions_http_client.py index 5022862cfe..1e04a18188 100644 --- a/tests/unittests/models/test_completions_http_client.py +++ b/tests/unittests/models/test_completions_http_client.py @@ -12,9 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import atexit +import gc import json from unittest import mock from unittest.mock import AsyncMock +import weakref from google.adk.models.apigee_llm import ChatCompletionsResponseHandler from google.adk.models.apigee_llm import CompletionsHTTPClient @@ -829,3 +832,93 @@ def test_process_chunk_with_refusal_streaming(): final_response.content.parts[0].text == 'Hello\n[[REFUSAL]]: I refuse to answer' ) + + +def test_client_creation_registers_atexit_cleanup(): + local_client = CompletionsHTTPClient(base_url='https://localhost') + with mock.patch( + 'google.adk.models.apigee_llm.atexit.register' + ) as mock_register: + _ = local_client._client + mock_register.assert_called_once_with(local_client._atexit_callback) + + +def test_close_unregisters_atexit_cleanup(): + local_client = CompletionsHTTPClient(base_url='https://localhost') + _ = local_client._client + callback = local_client._atexit_callback + with mock.patch( + 'google.adk.models.apigee_llm.atexit.unregister' + ) as mock_unregister: + local_client.close() + mock_unregister.assert_called_once_with(callback) + + +@pytest.mark.asyncio +async def test_aclose_unregisters_atexit_cleanup(): + local_client = CompletionsHTTPClient(base_url='https://localhost') + _ = local_client._client + callback = local_client._atexit_callback + with mock.patch( + 'google.adk.models.apigee_llm.atexit.unregister' + ) as mock_unregister: + await local_client.aclose() + mock_unregister.assert_called_once_with(callback) + assert local_client._client.is_closed + + +def test_close_without_client_does_not_touch_atexit(): + local_client = CompletionsHTTPClient(base_url='https://localhost') + with mock.patch( + 'google.adk.models.apigee_llm.atexit.unregister' + ) as mock_unregister: + local_client.close() + mock_unregister.assert_not_called() + + +def test_atexit_registration_does_not_pin_client(): + # The atexit handler is registered against a weakref.proxy, so the registry + # must not keep the underlying httpx client alive after it is closed and its + # owner is dropped. Regression test for the leaked AsyncClient / connection + # pool that grew unbounded in long-running deployments. + local_client = CompletionsHTTPClient(base_url='https://localhost') + httpx_client = local_client._client + client_ref = weakref.ref(httpx_client) + + local_client.close() + del httpx_client + del local_client + gc.collect() + + assert client_ref() is None + + +def test_cleanup_client_tolerates_dead_weakref_proxy(): + local_client = CompletionsHTTPClient(base_url='https://localhost') + httpx_client = local_client._client + proxy = weakref.proxy(httpx_client) + del httpx_client + del local_client + gc.collect() + + # The proxy now points at a collected object; cleanup must not raise. + CompletionsHTTPClient._cleanup_client(proxy) + + +@pytest.mark.asyncio +async def test_cleanup_client_retains_pending_task(): + from google.adk.models import apigee_llm + + local_client = CompletionsHTTPClient(base_url='https://localhost') + httpx_client = local_client._client + + apigee_llm._CLEANUP_TASKS.clear() + CompletionsHTTPClient._cleanup_client(httpx_client) + + # While the aclose() task is pending it must be held in the module-level set + # so it is not garbage collected before it completes. + assert len(apigee_llm._CLEANUP_TASKS) == 1 + task = next(iter(apigee_llm._CLEANUP_TASKS)) + await task + assert apigee_llm._CLEANUP_TASKS == set() + assert httpx_client.is_closed