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>

Author: anxkhn

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()

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