fix: Retain the consumers and transport registry when recreating the ClientFactory in remote_a2a_agent.py

PiperOrigin-RevId: 807762203

Author: google-genai-bot

src/google/adk/agents/remote_a2a_agent.py

@@ -188,12 +188,16 @@ async def _ensure_httpx_client(self) -> httpx.AsyncClient:
       )
       self._httpx_client_needs_cleanup = True
       if self._a2a_client_factory:
+        registry = self._a2a_client_factory._registry
         self._a2a_client_factory = A2AClientFactory(
             config=dataclasses.replace(
                 self._a2a_client_factory._config,
                 httpx_client=self._httpx_client,
-            )
+            ),
+            consumers=self._a2a_client_factory._consumers,
         )
+        for label, generator in registry.items():
+          self._a2a_client_factory.register(label, generator)
     if not self._a2a_client_factory:
       client_config = A2AClientConfig(
           httpx_client=self._httpx_client,

tests/unittests/agents/test_remote_a2a_agent.py

@@ -34,6 +34,7 @@
 # Import dependencies with version checking
 try:
   from a2a.client.client import ClientConfig
+  from a2a.client.client import Consumer
   from a2a.client.client_factory import ClientFactory
   from a2a.types import AgentCapabilities
   from a2a.types import AgentCard
@@ -246,6 +247,50 @@ async def test_ensure_factory_reuses_existing_client(self):
     assert client == existing_client
     assert agent._httpx_client_needs_cleanup is False
 
+  @pytest.mark.asyncio
+  async def test_ensure_httpx_client_updates_factory_with_new_client(self):
+    """Test that _ensure_httpx_client updates factory with new client."""
+    agent = RemoteA2aAgent(
+        name="test_agent",
+        agent_card=create_test_agent_card(),
+        a2a_client_factory=ClientFactory(
+            ClientConfig(httpx_client=None),
+        ),
+    )
+    assert agent._a2a_client_factory._config.httpx_client is None
+
+    client = await agent._ensure_httpx_client()
+
+    assert client is not None
+    assert agent._httpx_client == client
+    assert agent._httpx_client_needs_cleanup is True
+    assert agent._a2a_client_factory._config.httpx_client == client
+
+  @pytest.mark.asyncio
+  async def test_ensure_httpx_client_reregisters_transports_with_new_client(
+      self,
+  ):
+    """Test that _ensure_httpx_client registers transports with new client."""
+    factory = ClientFactory(
+        ClientConfig(httpx_client=None),
+    )
+    factory.register("transport_label", lambda: "test")
+    agent = RemoteA2aAgent(
+        name="test_agent",
+        agent_card=create_test_agent_card(),
+        a2a_client_factory=factory,
+    )
+    assert agent._a2a_client_factory._config.httpx_client is None
+    assert "transport_label" in agent._a2a_client_factory._registry
+
+    client = await agent._ensure_httpx_client()
+
+    assert client is not None
+    assert agent._httpx_client == client
+    assert agent._httpx_client_needs_cleanup is True
+    assert agent._a2a_client_factory._config.httpx_client == client
+    assert "transport_label" in agent._a2a_client_factory._registry
+
   @pytest.mark.asyncio
   async def test_resolve_agent_card_from_url_success(self):
     """Test successful agent card resolution from URL."""