Skip to content

fix(a2a): resolve a remote agent's card and client once across invocations - #7672

Open
ishan-1010 wants to merge 2 commits into
google:mainfrom
ishan-1010:fix-remote-a2a-resolve-once
Open

ishan-1010 wants to merge 2 commits into
google:mainfrom
ishan-1010:fix-remote-a2a-resolve-once

Conversation

@ishan-1010

Copy link
Copy Markdown

Link to Issue or Description of Change

Problem:
The agent tree is cloned for every invocation and BaseAgent.clone() is a shallow model_copy, so RemoteA2aAgent cached the resolved card, httpx client and A2A client on a copy that gets thrown away. Every turn fetched /.well-known/agent-card.json again and opened a new httpx.AsyncClient, even on turns that never transfer, and cleanup() / runner.close() on the instance you passed in couldn't close them.

Solution:
Keep that state (_agent_card, _a2a_client, _a2a_client_factory, _httpx_client, _httpx_client_needs_cleanup, _is_resolved) on a small _SharedResolution object the agent and its clones share. The attributes stay as properties with the same names, so the rest of the class is unchanged. The per-invocation path for authenticated cards still builds its client per invocation and doesn't cache it.

Testing Plan

Unit Tests:

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

New TestRemoteA2aAgentClonesShareResolution, all three fail on main:

  • test_card_is_fetched_once_across_turns: 3 runner turns through a root LlmAgent, card resolved once (3 on main)
  • test_clones_reuse_one_card_and_client: two clones get the same A2A client
  • test_cleanup_closes_client_created_by_clone: cleanup() on the original closes the httpx client a clone opened
tests/unittests/a2a: 951 passed, 55 skipped
tests/unittests: the only failures are 50 in labs/antigravity, integrations/vmaas and test_release_dependencies, which fail the same way on main here (optional deps not installed)
pre-commit (isort, pyink, ruff, addlicense, compliance checks, codespell): passed
mypy on _remote_a2a_agent.py: no new errors (1 existing on main)

Manual End-to-End (E2E) Tests:

The repro script from #7629 (real uvicorn A2A server, fake BaseLlm, 5 turns that never transfer):

main (facd7dd)                          this PR
agent card GETs: 5                      agent card GETs: 1
httpx.AsyncClient created: 5            httpx.AsyncClient created: 1
card cached on the instance: False      card cached on the instance: True

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.

…tions

The agent tree is cloned for every invocation, and BaseAgent.clone() is a
shallow model_copy, so RemoteA2aAgent cached the resolved agent card, the
httpx client and the A2A client on a copy that was thrown away. Every turn
fetched the card again and opened a new httpx client, even when it never
transferred to the remote agent, and cleanup() on the original instance
could not close those clients.

Keep that state on an object the agent and its clones share.

Fixes google#7629

@tonydzi tonydzi left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mycroft here, Anton's synthetic AI cofounder. I read _ensure_resolved and cleanup on this branch (4b07387) against the new shared _SharedResolution.

Sharing the card/client across clones makes cleanup() (_remote_a2a_agent.py, ~L2370) act on state that live invocations also use, and cleanup() only half-resets it:

  1. It closes self._httpx_client and sets it to None, but leaves _a2a_client, _a2a_client_factory (which rebind_client_factory_httpx bound to that client around L1052) and _is_resolved untouched. The early return at ~L1311 (not per_invocation_card and self._is_resolved and self._a2a_client) therefore keeps handing out a client whose transport is already closed. Before this PR a clone had its own state, so closing one instance did not reach into another; now Runner.close() on the original instance (or any clone's cleanup()) can break the in-flight and the next invocations of every clone. test_cleanup_closes_client_created_by_clone asserts client.is_closed and clone._httpx_client is None, but not that a later _ensure_resolved() rebuilds a working client. Resetting _a2a_client, the factory binding and is_resolved in cleanup() (or dropping the _SharedResolution object) would close that.

  2. The shared path in _ensure_resolved has awaits between check and set (_resolve_agent_card, _ensure_httpx_client, then self._a2a_client = self._a2a_client_factory.create(...) after if not self._a2a_client) and no lock. _ensure_httpx_client already carries the comment "Another caller may have created the client during the await", so the race is known. With per-clone state it required concurrent calls on one instance; with the shared object, every concurrent first request across sessions on a server joins it (N card fetches, N clients, last write wins). An asyncio.Lock on the _SharedResolution around the shared branch would make the "fetched once" claim in test_card_is_fetched_once_across_turns hold under concurrency too, and that test runs turns sequentially, so it cannot show it.

I read the code paths only; I did not run them.

cleanup() closed the httpx client but kept the A2A client, the factory
bound to it and the resolved flag, so the next request (from any clone,
since they share this state) got a client whose transport was closed.
Reset them too, so the next request resolves again.

Concurrent first requests across sessions each fetched the card and
built a client. Run the shared resolution under a lock the clones share.
@ishan-1010

Copy link
Copy Markdown
Author

@tonydzi both fixed in 7836316.

  1. cleanup() now also drops the A2A client, restores the configured factory and clears the resolved flag (here), so the next request from any clone resolves again. test resolves, cleans up, and checks the next _ensure_resolved() returns a new client on an open httpx client.
  2. the shared path runs under an asyncio.Lock on _SharedResolution with a re-check after acquiring it (here). test runs 5 clones' first resolution concurrently and checks one card fetch and one client.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

RemoteA2aAgent fetches the agent card and creates a new httpx client on every turn

2 participants