fix(models): stop CompletionsHTTPClient leaking its httpx client via atexit

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