From 4b07387bec91989934e9fa3592395d37c6282b7d Mon Sep 17 00:00:00 2001 From: Ishan Katoch <98383932+ishan-1010@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:17:16 +0530 Subject: [PATCH 1/2] fix(a2a): resolve a remote agent's card and client once across invocations 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 #7629 --- src/google/adk/a2a/agent/_remote_a2a_agent.py | 94 ++++++++++++++++++- .../a2a/agent/test_remote_a2a_agent.py | 62 ++++++++++++ 2 files changed, 153 insertions(+), 3 deletions(-) diff --git a/src/google/adk/a2a/agent/_remote_a2a_agent.py b/src/google/adk/a2a/agent/_remote_a2a_agent.py index 6dfa73d75d2..ab261466e27 100644 --- a/src/google/adk/a2a/agent/_remote_a2a_agent.py +++ b/src/google/adk/a2a/agent/_remote_a2a_agent.py @@ -667,6 +667,18 @@ async def _before_request( ) +class _SharedResolution: + """The resolved card and clients of a RemoteA2aAgent and all its clones.""" + + def __init__(self) -> None: + self.agent_card: Optional[AgentCard] = None + self.a2a_client: Optional[A2AClient] = None + self.a2a_client_factory: Optional[A2AClientFactory] = None + self.httpx_client: Optional[httpx.AsyncClient] = None + self.httpx_client_needs_cleanup = False + self.is_resolved = False + + @a2a_experimental class RemoteA2aAgent(BaseAgent): """Agent that communicates with a remote A2A agent via A2A client. @@ -704,6 +716,83 @@ class RemoteA2aAgent(BaseAgent): ``transfer_to_agent`` target. """ + # The agent tree is cloned for every invocation (BaseAgent.clone() is a + # shallow model_copy), so the resolved card and clients live on an object + # the clones share: the first turn resolves them, and cleanup() on this + # instance closes what any clone opened. Deleting one of these attributes + # resets it, which unittest.mock.patch.object does on exit. + @property + def _agent_card(self) -> Optional[AgentCard]: + return self._resolution.agent_card + + @_agent_card.setter + def _agent_card(self, value: Optional[AgentCard]) -> None: + self._resolution.agent_card = value + + @_agent_card.deleter + def _agent_card(self) -> None: + self._resolution.agent_card = None + + @property + def _a2a_client(self) -> Optional[A2AClient]: + return self._resolution.a2a_client + + @_a2a_client.setter + def _a2a_client(self, value: Optional[A2AClient]) -> None: + self._resolution.a2a_client = value + + @_a2a_client.deleter + def _a2a_client(self) -> None: + self._resolution.a2a_client = None + + @property + def _a2a_client_factory(self) -> Optional[A2AClientFactory]: + return self._resolution.a2a_client_factory + + @_a2a_client_factory.setter + def _a2a_client_factory(self, value: Optional[A2AClientFactory]) -> None: + self._resolution.a2a_client_factory = value + + @_a2a_client_factory.deleter + def _a2a_client_factory(self) -> None: + self._resolution.a2a_client_factory = None + + @property + def _httpx_client(self) -> Optional[httpx.AsyncClient]: + return self._resolution.httpx_client + + @_httpx_client.setter + def _httpx_client(self, value: Optional[httpx.AsyncClient]) -> None: + self._resolution.httpx_client = value + + @_httpx_client.deleter + def _httpx_client(self) -> None: + self._resolution.httpx_client = None + + @property + def _httpx_client_needs_cleanup(self) -> bool: + return self._resolution.httpx_client_needs_cleanup + + @_httpx_client_needs_cleanup.setter + def _httpx_client_needs_cleanup(self, value: bool) -> None: + self._resolution.httpx_client_needs_cleanup = value + + @_httpx_client_needs_cleanup.deleter + def _httpx_client_needs_cleanup(self) -> None: + self._resolution.httpx_client_needs_cleanup = False + + @property + def _is_resolved(self) -> bool: + return self._resolution.is_resolved + + @_is_resolved.setter + def _is_resolved(self, value: bool) -> None: + self._resolution.is_resolved = value + + @_is_resolved.deleter + def _is_resolved(self) -> None: + self._resolution.is_resolved = False + def __init__( self, name: str, @@ -775,9 +864,8 @@ def __init__( if agent_card is None: raise ValueError("agent_card cannot be None") - self._agent_card: Optional[AgentCard] = None + self._resolution = _SharedResolution() self._agent_card_source: Optional[str] = None - self._a2a_client: Optional[A2AClient] = None # This is stored to support backward compatible usage of class. # In future, the client is expected to be present in the factory. self._httpx_client = httpx_client @@ -788,7 +876,7 @@ def __init__( self._is_resolved = False self._genai_part_converter = genai_part_converter self._a2a_part_converter = a2a_part_converter - self._a2a_client_factory: Optional[A2AClientFactory] = a2a_client_factory + self._a2a_client_factory = a2a_client_factory self._a2a_request_meta_provider = a2a_request_meta_provider self._full_history_when_stateless_param = full_history_when_stateless self._context_builder = context_builder diff --git a/tests/unittests/a2a/agent/test_remote_a2a_agent.py b/tests/unittests/a2a/agent/test_remote_a2a_agent.py index 3085ed4b24c..1c751fee2ed 100644 --- a/tests/unittests/a2a/agent/test_remote_a2a_agent.py +++ b/tests/unittests/a2a/agent/test_remote_a2a_agent.py @@ -65,6 +65,7 @@ from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_ERROR_RESULT from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_SUCCESS_RESULT from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME +from google.adk.agents.llm_agent import LlmAgent from google.adk.agents.remote_a2a_agent import A2A_METADATA_PREFIX from google.adk.agents.remote_a2a_agent import AgentCardResolutionError from google.adk.agents.remote_a2a_agent import RemoteA2aAgent @@ -9069,3 +9070,64 @@ async def override_metadata(ctx, response, event): assert turn[-1].custom_metadata["a2a:task_id"] == "caller-task" assert turn[-1].custom_metadata["a2a:context_id"] == "caller-context" assert turn[-1].custom_metadata["a2a:response"]["id"] == "currency-task" + + +class TestRemoteA2aAgentClonesShareResolution: + """The agent tree is cloned per invocation, so clones share resolution.""" + + _CARD_URL = "https://example.com/.well-known/agent-card.json" + + @pytest.mark.asyncio + async def test_card_is_fetched_once_across_turns(self): + from tests.unittests import testing_utils + + remote = RemoteA2aAgent(name="remote", agent_card=self._CARD_URL) + root = LlmAgent( + name="root", + model=testing_utils.MockModel.create(responses=["ok", "ok", "ok"]), + sub_agents=[remote], + ) + runner = testing_utils.InMemoryRunner(root_agent=root) + resolve = AsyncMock(return_value=create_test_agent_card()) + + with patch.object(RemoteA2aAgent, "_resolve_agent_card", resolve): + for _ in range(3): + await runner.run_async(testing_utils.get_user_content("hi")) + + assert resolve.await_count == 1 + assert remote._agent_card is not None + await remote.cleanup() + + @pytest.mark.asyncio + async def test_clones_reuse_one_card_and_client(self): + agent = RemoteA2aAgent(name="remote", agent_card=self._CARD_URL) + resolve = AsyncMock(return_value=create_test_agent_card()) + + with patch.object(RemoteA2aAgent, "_resolve_agent_card", resolve): + first = await agent.clone()._ensure_resolved() + second = await agent.clone()._ensure_resolved() + + try: + assert first is second + assert resolve.await_count == 1 + assert agent._agent_card is not None + assert agent._httpx_client is not None + finally: + await agent.cleanup() + + @pytest.mark.asyncio + async def test_cleanup_closes_client_created_by_clone(self): + agent = RemoteA2aAgent(name="remote", agent_card=self._CARD_URL) + clone = agent.clone() + with patch.object( + RemoteA2aAgent, + "_resolve_agent_card", + AsyncMock(return_value=create_test_agent_card()), + ): + await clone._ensure_resolved() + client = clone._httpx_client + + await agent.cleanup() + + assert client.is_closed + assert clone._httpx_client is None From 7836316898acc6c29238326030cc5b8ff7ea8d1a Mon Sep 17 00:00:00 2001 From: Ishan Katoch <98383932+ishan-1010@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:49:40 +0530 Subject: [PATCH 2/2] fix(a2a): reset the shared client on cleanup and resolve it under a lock 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. --- src/google/adk/a2a/agent/_remote_a2a_agent.py | 78 +++++++++++-------- .../a2a/agent/test_remote_a2a_agent.py | 41 ++++++++++ 2 files changed, 86 insertions(+), 33 deletions(-) diff --git a/src/google/adk/a2a/agent/_remote_a2a_agent.py b/src/google/adk/a2a/agent/_remote_a2a_agent.py index ab261466e27..a10d71d745e 100644 --- a/src/google/adk/a2a/agent/_remote_a2a_agent.py +++ b/src/google/adk/a2a/agent/_remote_a2a_agent.py @@ -677,6 +677,9 @@ def __init__(self) -> None: self.httpx_client: Optional[httpx.AsyncClient] = None self.httpx_client_needs_cleanup = False self.is_resolved = False + # Concurrent first requests across sessions would otherwise each fetch + # the card and build a client. + self.lock = asyncio.Lock() @a2a_experimental @@ -877,6 +880,7 @@ def __init__( self._genai_part_converter = genai_part_converter self._a2a_part_converter = a2a_part_converter self._a2a_client_factory = a2a_client_factory + self._configured_a2a_client_factory = a2a_client_factory self._a2a_request_meta_provider = a2a_request_meta_provider self._full_history_when_stateless_param = full_history_when_stateless self._context_builder = context_builder @@ -1333,44 +1337,47 @@ async def _ensure_resolved( return client # Shared (cached) resolution path. - if not self._agent_card: - - # Resolve agent card if needed - agent_card = await self._resolve_agent_card(ctx) - - # Validate agent card - await self._validate_agent_card(agent_card) + async with self._resolution.lock: + if self._is_resolved and self._a2a_client: + return self._a2a_client + if not self._agent_card: + + # Resolve agent card if needed + agent_card = await self._resolve_agent_card(ctx) + + # Validate agent card + await self._validate_agent_card(agent_card) + + # Store the card only once it has validated. A rejected card left on + # the instance reads as already resolved, so the next call skips the + # check and talks to the origin that card named. + self._agent_card = agent_card + + # A public card may already have been resolved for transfer selection. + # Preserve the existing behavior of adopting its description when the + # remote agent itself is initialized. + if not self.description and self._agent_card.description: + self.description = _adopted_card_description( + self._agent_card.description, self._agent_card_source + ) - # Store the card only once it has validated. A rejected card left on - # the instance reads as already resolved, so the next call skips the - # check and talks to the origin that card named. - self._agent_card = agent_card - - # A public card may already have been resolved for transfer selection. - # Preserve the existing behavior of adopting its description when the - # remote agent itself is initialized. - if not self.description and self._agent_card.description: - self.description = _adopted_card_description( - self._agent_card.description, self._agent_card_source - ) + # Initialize A2A client + if not self._a2a_client: + await self._ensure_httpx_client() + if not self._a2a_client_factory: + raise AgentCardResolutionError( + f"Failed to initialize A2A client factory for {self.name}" + ) + self._a2a_client = self._a2a_client_factory.create(self._agent_card) - # Initialize A2A client - if not self._a2a_client: - await self._ensure_httpx_client() - if not self._a2a_client_factory: + if not self._a2a_client: raise AgentCardResolutionError( - f"Failed to initialize A2A client factory for {self.name}" + f"Failed to initialize A2A client for {self.name}" ) - self._a2a_client = self._a2a_client_factory.create(self._agent_card) - if not self._a2a_client: - raise AgentCardResolutionError( - f"Failed to initialize A2A client for {self.name}" - ) - - self._is_resolved = True - logger.info("Successfully resolved remote A2A agent: %s", self.name) - return self._a2a_client + self._is_resolved = True + logger.info("Successfully resolved remote A2A agent: %s", self.name) + return self._a2a_client except Exception as e: logger.error("Failed to resolve remote A2A agent %s: %s", self.name, e) @@ -2381,3 +2388,8 @@ async def cleanup(self) -> None: ) finally: self._httpx_client = None + # The A2A client and the factory were bound to the closed client, and + # clones share them, so resolve again on the next request. + self._a2a_client = None + self._a2a_client_factory = self._configured_a2a_client_factory + self._is_resolved = False diff --git a/tests/unittests/a2a/agent/test_remote_a2a_agent.py b/tests/unittests/a2a/agent/test_remote_a2a_agent.py index 1c751fee2ed..89aa8a5f7e7 100644 --- a/tests/unittests/a2a/agent/test_remote_a2a_agent.py +++ b/tests/unittests/a2a/agent/test_remote_a2a_agent.py @@ -9131,3 +9131,44 @@ async def test_cleanup_closes_client_created_by_clone(self): assert client.is_closed assert clone._httpx_client is None + + @pytest.mark.asyncio + async def test_resolve_after_cleanup_builds_a_new_client(self): + agent = RemoteA2aAgent(name="remote", agent_card=self._CARD_URL) + clone = agent.clone() + with patch.object( + RemoteA2aAgent, + "_resolve_agent_card", + AsyncMock(return_value=create_test_agent_card()), + ): + before = await clone._ensure_resolved() + closed_client = clone._httpx_client + await agent.cleanup() + after = await clone._ensure_resolved() + + try: + assert after is not before + assert clone._httpx_client is not closed_client + assert not clone._httpx_client.is_closed + finally: + await agent.cleanup() + + @pytest.mark.asyncio + async def test_concurrent_first_resolution_fetches_the_card_once(self): + agent = RemoteA2aAgent(name="remote", agent_card=self._CARD_URL) + + async def slow_resolve(*args, **kwargs): + await asyncio.sleep(0.01) + return create_test_agent_card() + + resolve = AsyncMock(side_effect=slow_resolve) + with patch.object(RemoteA2aAgent, "_resolve_agent_card", resolve): + clients = await asyncio.gather( + *(agent.clone()._ensure_resolved() for _ in range(5)) + ) + + try: + assert resolve.await_count == 1 + assert all(client is clients[0] for client in clients) + finally: + await agent.cleanup()