From 561039539edc5126bcec75375541bb27963e6283 Mon Sep 17 00:00:00 2001 From: Anas Khan <83116240+anxkhn@users.noreply.github.com> Date: Fri, 3 Jul 2026 14:59:12 +0530 Subject: [PATCH] fix(models): stop CompletionsHTTPClient leaking its httpx client via atexit CompletionsHTTPClient._client registered its per-instance httpx.AsyncClient with atexit.register(self._cleanup_client, client), leaving a permanent strong reference to the client and its connection pool in the process-global atexit registry. Nothing ever called atexit.unregister, so neither close(), aclose(), nor __del__ could release the client before interpreter exit. A long-running deployment that creates many CompletionsHTTPClient instances therefore grew the atexit registry and retained httpx clients without bound. Register the cleanup callback against weakref.proxy(client) so the registry no longer pins the client, mirroring the existing pattern in bigquery_agent_analytics_plugin, and unregister the per-instance callback in close()/aclose() once the client is explicitly closed. Tolerate a dead proxy in _cleanup_client, and retain the fire-and-forget aclose() task in a module-level set so it is not garbage collected while pending. Add regression tests covering the register/unregister pairing and asserting the client is releasable after close(). Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com> --- src/google/adk/models/apigee_llm.py | 27 +++++- .../models/test_completions_http_client.py | 93 +++++++++++++++++++ 2 files changed, 117 insertions(+), 3 deletions(-) diff --git a/src/google/adk/models/apigee_llm.py b/src/google/adk/models/apigee_llm.py index a4b57ce6885..38662727127 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 5022862cfea..1e04a181882 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