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

src/google/adk/a2a/agent/_remote_a2a_agent.py

@@ -667,6 +667,21 @@ 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
+    # Concurrent first requests across sessions would otherwise each fetch
+    # the card and build a client.
+    self.lock = asyncio.Lock()
+
+
 @a2a_experimental
 class RemoteA2aAgent(BaseAgent):
   """Agent that communicates with a remote A2A agent via A2A client.
@@ -704,6 +719,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 +867,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 +879,8 @@ 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._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
@@ -1245,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)
@@ -2293,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

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,105 @@ 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
+
+  @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()