fix: Fixes for initializing RemoteA2aAgent - passing in preferred transport, protocol version, and auth headers

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 896068855

Author: wukath

src/google/adk/integrations/agent_registry/agent_registry.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Generator
 from enum import Enum
 import logging
 import os
@@ -55,6 +56,12 @@
 
 AGENT_REGISTRY_BASE_URL = "https://agentregistry.googleapis.com/v1alpha"
 
+_TRANSPORT_MAPPING = {
+    "HTTP_JSON": A2ATransport.http_json,
+    "JSONRPC": A2ATransport.jsonrpc,
+    "GRPC": A2ATransport.grpc,
+}
+
 
 # An MCPToolset for a single registered MCP server. Adds the special
 # gcp.mcp.server.destination.id custom_metadata key on each returned tool. This special key is
@@ -232,13 +239,15 @@ def _get_connection_uri(
     for p in protocols:
       if protocol_type and p.get("type") != protocol_type:
         continue
+      protocol_version = p.get("protocolVersion")
       for i in p.get("interfaces", []):
-        if protocol_binding and i.get("protocolBinding") != protocol_binding:
+        mapped_binding = _TRANSPORT_MAPPING.get(i.get("protocolBinding"))
+        if protocol_binding and mapped_binding != protocol_binding:
           continue
         if url := i.get("url"):
-          return url
+          return url, protocol_version, mapped_binding
 
-    return None
+    return None, None, None
 
   def _clean_name(self, name: str) -> str:
     """Cleans a string to be a valid Python identifier for agent names."""
@@ -284,11 +293,13 @@ def get_mcp_toolset(
     if not isinstance(mcp_server_id, str):
       mcp_server_id = None
 
-    endpoint_uri = self._get_connection_uri(
+    endpoint_uri, _, _ = self._get_connection_uri(
         server_details, protocol_binding=A2ATransport.jsonrpc
-    ) or self._get_connection_uri(
-        server_details, protocol_binding=A2ATransport.http_json
     )
+    if not endpoint_uri:
+      endpoint_uri, _, _ = self._get_connection_uri(
+          server_details, protocol_binding=A2ATransport.http_json
+      )
     if not endpoint_uri:
       raise ValueError(
           f"MCP Server endpoint URI not found for: {mcp_server_name}"
@@ -339,7 +350,7 @@ def get_model_name(self, endpoint_name: str) -> str:
       projects/.../locations/.../publishers/google/models/...).
     """
     endpoint_details = self.get_endpoint(endpoint_name)
-    uri = self._get_connection_uri(endpoint_details)
+    uri, _, _ = self._get_connection_uri(endpoint_details)
     if not uri:
       raise ValueError(
           f"Connection URI not found for endpoint: {endpoint_name}"
@@ -378,7 +389,12 @@ def get_agent_info(self, name: str) -> Dict[str, Any]:
     """Retrieves detailed metadata of a specific A2A Agent."""
     return self._make_request(name)
 
-  def get_remote_a2a_agent(self, agent_name: str) -> RemoteA2aAgent:
+  def get_remote_a2a_agent(
+      self,
+      agent_name: str,
+      *,
+      httpx_client: httpx.AsyncClient | None = None,
+  ) -> RemoteA2aAgent:
     """Creates a RemoteA2aAgent instance for a registered A2A Agent."""
     agent_info = self.get_agent_info(agent_name)
 
@@ -389,17 +405,19 @@ def get_remote_a2a_agent(self, agent_name: str) -> RemoteA2aAgent:
       agent_card = AgentCard(**card_content)
       # Clean the name to be a valid identifier
       name = self._clean_name(agent_card.name)
+
       return RemoteA2aAgent(
           name=name,
           agent_card=agent_card,
           description=agent_card.description,
+          httpx_client=httpx_client,
       )
 
     name = self._clean_name(agent_info.get("displayName", agent_name))
     description = agent_info.get("description", "")
     version = agent_info.get("version", "")
 
-    url = self._get_connection_uri(
+    url, protocol_version, protocol_binding = self._get_connection_uri(
         agent_info, protocol_type=_ProtocolType.A2A_AGENT
     )
     if not url:
@@ -421,6 +439,8 @@ def get_remote_a2a_agent(self, agent_name: str) -> RemoteA2aAgent:
         name=name,
         description=description,
         version=version,
+        preferredTransport=protocol_binding or A2ATransport.http_json,
+        protocolVersion=protocol_version or "0.3.0",
         url=url,
         skills=skills,
         capabilities=AgentCapabilities(streaming=False, polling=False),
@@ -432,4 +452,5 @@ def get_remote_a2a_agent(self, agent_name: str) -> RemoteA2aAgent:
         name=name,
         agent_card=agent_card,
         description=description,
+        httpx_client=httpx_client,
     )

tests/unittests/integrations/agent_registry/test_agent_registry.py

@@ -60,7 +60,7 @@ async def test_get_mcp_toolset_adds_destination_id(
         ),
         "interfaces": [{
             "url": "https://mcp.com",
-            "protocolBinding": A2ATransport.jsonrpc,
+            "protocolBinding": "JSONRPC",
         }],
     }
     mock_httpx.return_value.__enter__.return_value.get.return_value = (
@@ -126,7 +126,7 @@ async def test_get_mcp_toolset_handles_missing_destination_id(
         # "mcpServerId" is intentionally omitted
         "interfaces": [{
             "url": "https://mcp.com",
-            "protocolBinding": A2ATransport.jsonrpc,
+            "protocolBinding": "JSONRPC",
         }],
     }
     mock_httpx.return_value.__enter__.return_value.get.return_value = (
@@ -176,25 +176,29 @@ def test_get_connection_uri_mcp_interfaces_top_level(self, registry):
             {"url": "https://mcp-v1main.com", "protocolBinding": "JSONRPC"}
         ]
     }
-    uri = registry._get_connection_uri(
+    uri, version, binding = registry._get_connection_uri(
         resource_details, protocol_binding=A2ATransport.jsonrpc
     )
     assert uri == "https://mcp-v1main.com"
+    assert version is None
+    assert binding == "JSONRPC"
 
   def test_get_connection_uri_agent_nested_protocols(self, registry):
     resource_details = {
         "protocols": [{
             "type": _ProtocolType.A2A_AGENT,
             "interfaces": [{
                 "url": "https://my-agent.com",
-                "protocolBinding": A2ATransport.jsonrpc,
+                "protocolBinding": "JSONRPC",
             }],
         }]
     }
-    uri = registry._get_connection_uri(
+    uri, version, binding = registry._get_connection_uri(
         resource_details, protocol_type=_ProtocolType.A2A_AGENT
     )
     assert uri == "https://my-agent.com"
+    assert version is None
+    assert binding == A2ATransport.jsonrpc
 
   def test_get_connection_uri_filtering(self, registry):
     resource_details = {
@@ -207,42 +211,52 @@ def test_get_connection_uri_filtering(self, registry):
                 "type": _ProtocolType.A2A_AGENT,
                 "interfaces": [{
                     "url": "https://my-agent.com",
-                    "protocolBinding": A2ATransport.http_json,
+                    "protocolBinding": "HTTP_JSON",
                 }],
             },
         ]
     }
     # Filter by type
-    uri = registry._get_connection_uri(
+    uri, version, binding = registry._get_connection_uri(
         resource_details, protocol_type=_ProtocolType.A2A_AGENT
     )
     assert uri == "https://my-agent.com"
+    assert version is None
+    assert binding == A2ATransport.http_json
 
     # Filter by binding
-    uri = registry._get_connection_uri(
+    uri, version, binding = registry._get_connection_uri(
         resource_details, protocol_binding=A2ATransport.http_json
     )
     assert uri == "https://my-agent.com"
+    assert version is None
+    assert binding == A2ATransport.http_json
 
     # No match
-    uri = registry._get_connection_uri(
+    uri, version, binding = registry._get_connection_uri(
         resource_details,
         protocol_type=_ProtocolType.A2A_AGENT,
         protocol_binding=A2ATransport.jsonrpc,
     )
     assert uri is None
+    assert version is None
+    assert binding is None
 
   def test_get_connection_uri_returns_none_if_no_interfaces(self, registry):
     resource_details = {}
-    uri = registry._get_connection_uri(resource_details)
+    uri, version, binding = registry._get_connection_uri(resource_details)
     assert uri is None
+    assert version is None
+    assert binding is None
 
   def test_get_connection_uri_returns_none_if_no_url_in_interfaces(
       self, registry
   ):
     resource_details = {"interfaces": [{"protocolBinding": "HTTP"}]}
-    uri = registry._get_connection_uri(resource_details)
+    uri, version, binding = registry._get_connection_uri(resource_details)
     assert uri is None
+    assert version is None
+    assert binding is None
 
   @patch("httpx.Client")
   def test_list_agents(self, mock_httpx, registry):
@@ -313,7 +327,7 @@ def test_get_mcp_toolset(self, mock_httpx, registry):
         "displayName": "TestPrefix",
         "interfaces": [{
             "url": "https://mcp.com",
-            "protocolBinding": A2ATransport.jsonrpc,
+            "protocolBinding": "JSONRPC",
         }],
     }
     mock_response.raise_for_status = MagicMock()
@@ -335,7 +349,7 @@ def test_get_mcp_toolset_with_auth(self, mock_httpx, registry):
         "displayName": "TestPrefix",
         "interfaces": [{
             "url": "https://mcp.com",
-            "protocolBinding": A2ATransport.jsonrpc,
+            "protocolBinding": "JSONRPC",
         }],
     }
     mock_response.raise_for_status = MagicMock()
@@ -370,9 +384,10 @@ def test_get_remote_a2a_agent(self, mock_httpx, registry):
         "version": "1.0",
         "protocols": [{
             "type": _ProtocolType.A2A_AGENT,
+            "protocolVersion": "0.4.0",
             "interfaces": [{
                 "url": "https://my-agent.com",
-                "protocolBinding": A2ATransport.jsonrpc,
+                "protocolBinding": "HTTP_JSON",
             }],
         }],
         "skills": [{"id": "s1", "name": "Skill 1", "description": "Desc 1"}],
@@ -393,6 +408,35 @@ def test_get_remote_a2a_agent(self, mock_httpx, registry):
     assert agent._agent_card.version == "1.0"
     assert len(agent._agent_card.skills) == 1
     assert agent._agent_card.skills[0].name == "Skill 1"
+    assert agent._agent_card.preferred_transport == A2ATransport.http_json
+    assert agent._agent_card.protocol_version == "0.4.0"
+
+  @patch("httpx.Client")
+  def test_get_remote_a2a_agent_defaults(self, mock_httpx, registry):
+    mock_response = MagicMock()
+    mock_response.json.return_value = {
+        "displayName": "TestAgent",
+        "description": "Test Desc",
+        "version": "1.0",
+        "protocols": [{
+            "type": _ProtocolType.A2A_AGENT,
+            "interfaces": [{
+                "url": "https://my-agent.com",
+            }],
+        }],
+    }
+    mock_response.raise_for_status = MagicMock()
+    mock_httpx.return_value.__enter__.return_value.get.return_value = (
+        mock_response
+    )
+
+    registry._credentials.token = "token"
+    registry._credentials.refresh = MagicMock()
+
+    agent = registry.get_remote_a2a_agent("test-agent")
+    assert isinstance(agent, RemoteA2aAgent)
+    assert agent._agent_card.preferred_transport == A2ATransport.http_json
+    assert agent._agent_card.protocol_version == "0.3.0"
 
   @patch("httpx.Client")
   def test_get_remote_a2a_agent_with_card(self, mock_httpx, registry):
@@ -436,6 +480,57 @@ def test_get_remote_a2a_agent_with_card(self, mock_httpx, registry):
     assert len(agent._agent_card.skills) == 1
     assert agent._agent_card.skills[0].name == "S1"
 
+  @patch("httpx.Client")
+  def test_get_remote_a2a_agent_with_httpx_client(self, mock_httpx, registry):
+    mock_response = MagicMock()
+    mock_response.json.return_value = {
+        "displayName": "TestAgent",
+        "description": "Test Desc",
+        "version": "1.0",
+        "protocols": [{
+            "type": _ProtocolType.A2A_AGENT,
+            "interfaces": [{
+                "url": "https://my-agent.com",
+            }],
+        }],
+    }
+    mock_response.raise_for_status = MagicMock()
+    mock_httpx.return_value.__enter__.return_value.get.return_value = (
+        mock_response
+    )
+
+    custom_client = httpx.AsyncClient()
+    agent = registry.get_remote_a2a_agent(
+        "test-agent", httpx_client=custom_client
+    )
+    assert agent._httpx_client is custom_client
+
+  @patch("httpx.Client")
+  def test_get_remote_a2a_agent_configures_transports(
+      self, mock_httpx, registry
+  ):
+    mock_response = MagicMock()
+    mock_response.json.return_value = {
+        "displayName": "TestAgent",
+        "protocols": [{
+            "type": _ProtocolType.A2A_AGENT,
+            "interfaces": [{
+                "url": "https://my-agent.com",
+                "protocolBinding": A2ATransport.jsonrpc,
+            }],
+        }],
+    }
+    mock_response.raise_for_status = MagicMock()
+    mock_httpx.return_value.__enter__.return_value.get.return_value = (
+        mock_response
+    )
+
+    registry._credentials.token = "token"
+    registry._credentials.refresh = MagicMock()
+
+    agent = registry.get_remote_a2a_agent("test-agent")
+    assert agent._agent_card.preferred_transport == A2ATransport.jsonrpc
+
   def test_get_auth_headers(self, registry):
     registry._credentials.token = "fake-token"
     registry._credentials.refresh = MagicMock()
@@ -472,7 +567,7 @@ def test_make_request_raises_request_error(self, mock_httpx, registry):
     registry._credentials.refresh = MagicMock()
 
     with pytest.raises(
-        RuntimeError, match="API request failed \(network error\)"
+        RuntimeError, match=r"API request failed \(network error\)"
     ):
       registry._make_request("test-path")