diff --git a/src/google/adk/a2a/agent/_remote_a2a_agent.py b/src/google/adk/a2a/agent/_remote_a2a_agent.py index 6dfa73d75d2..a10d71d745e 100644 --- a/src/google/adk/a2a/agent/_remote_a2a_agent.py +++ b/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 diff --git a/tests/unittests/a2a/agent/test_remote_a2a_agent.py b/tests/unittests/a2a/agent/test_remote_a2a_agent.py index 3085ed4b24c..89aa8a5f7e7 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,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()