Skip to content

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

Open
anxkhn wants to merge 2 commits into
google:mainfrom
anxkhn:fix/apigee-completions-client-atexit-leak
Open

fix(models): stop CompletionsHTTPClient leaking its httpx client via atexit#6379
anxkhn wants to merge 2 commits into
google:mainfrom
anxkhn:fix/apigee-completions-client-atexit-leak

Conversation

@anxkhn

@anxkhn anxkhn commented Jul 12, 2026

Copy link
Copy Markdown
Contributor

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

  • N/A (no existing issue). Described below following the bug report template, per
    CONTRIBUTING ("you may instead describe the bug or feature directly within the
    PR description, following the structure of our issue templates").

2. Or, if no issue exists, describe the change:

Problem:

CompletionsHTTPClient._client (the OpenAI-compatible HTTP client used by
ApigeeLlm, in src/google/adk/models/apigee_llm.py) registers its per-instance
httpx.AsyncClient for cleanup with atexit.register(self._cleanup_client, client) on first access. atexit holds its callbacks (and their arguments) for
the entire life of the process, so this keeps a strong reference to the
AsyncClient and its connection pool in the process-global atexit registry.
Nothing ever calls atexit.unregister, so close(), aclose(), and __del__
cannot actually release the client before interpreter exit. A long-running
deployment that constructs many CompletionsHTTPClient / ApigeeLlm instances
therefore grows the atexit registry and leaks one open connection pool per
instance, without bound.

There is also a smaller latent issue in _cleanup_client: it schedules
loop.create_task(client.aclose()) without keeping a reference to the resulting
task, so the task can be garbage collected before it finishes ("Task was
destroyed but it is pending").

Steps to Reproduce:

  1. pip install google-adk
  2. Run the snippet below (it exercises only the client's lifecycle; no Apigee
    credentials or network calls are needed):
import gc, weakref
from google.adk.models.apigee_llm import CompletionsHTTPClient

refs = []
for _ in range(5):
    c = CompletionsHTTPClient(base_url="https://localhost")
    httpx_client = c._client   # triggers the atexit registration
    refs.append(weakref.ref(httpx_client))
    c.close()                  # explicit close
    del httpx_client
    del c

gc.collect()
alive = sum(1 for r in refs if r() is not None)
print(f"{alive}/5 httpx clients still alive after close() + del + gc")

Observed Behavior (before this change):

5/5 httpx clients still alive after close() + del + gc

Every client (and its connection pool) is pinned by the atexit registry even
after an explicit close(), an owner drop, and a full GC.

Expected Behavior:

After close() (or once the owner is dropped and finalized), the client should
be releasable and its connection pool freed; the process should not accumulate
one retained AsyncClient per constructed instance.

Solution:

  • Register the cleanup callback against weakref.proxy(client) so the atexit
    registry no longer keeps the client alive. This mirrors the pattern already
    used in bigquery_agent_analytics_plugin.py
    (atexit.register(self._atexit_cleanup, weakref.proxy(batch_processor))).
  • Store the bound callback per instance
    (self._atexit_callback = partial(self._cleanup_client, weakref.proxy(client)))
    and unregister exactly that callback in close() / aclose(). A per-instance
    partial is required because _cleanup_client is a shared @staticmethod, and
    atexit.unregister(func) removes every registration of func; unregistering
    the unique partial removes only this instance's handler and leaves other
    instances untouched.
  • Tolerate a dead proxy in _cleanup_client (catch ReferenceError), matching
    the same in-repo precedent.
  • Retain the fire-and-forget aclose() task in a module-level set until it
    completes, so it is not garbage collected while pending.

Observed Behavior (after this change):

0/5 httpx clients still alive after close() + del + gc

The same snippet reports 0/5 alive. The del-only path (drop the owner without
calling close(), relying on __del__ + the non-pinning proxy) also reports
0/5.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Added tests/unittests/models/test_completions_http_client.py regression tests
covering: the register/unregister pairing on close() and aclose(), the no-op
when no client was ever created, that the client is releasable (weakref goes to
None) after close() + GC, that cleanup tolerates a dead weakref.proxy, and
that the pending aclose() task is retained until it finishes.

$ pytest tests/unittests/models/test_completions_http_client.py \
         tests/unittests/models/test_apigee_llm.py -q
82 passed

Manual End-to-End (E2E) Tests:

The reproduction snippet above is the manual check: 5/5 retained clients before
the change, 0/5 after. It needs no Apigee setup because it only drives the HTTP
client's construction/cleanup lifecycle.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

The correct non-pinning pattern already exists in this codebase
(bigquery_agent_analytics_plugin.py); this change brings
CompletionsHTTPClient in line with it.

…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>
@adk-bot adk-bot added the models [Component] This issue is related to model support label Jul 12, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

models [Component] This issue is related to model support

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants