feat: Support AgentRegistry association

google-genai-bot · copybara-github · commit 675476088b9f · 2026-03-27T12:55:43.000-07:00
PiperOrigin-RevId: 890595289
diff --git a/src/google/adk/integrations/agent_registry/agent_registry.py b/src/google/adk/integrations/agent_registry/agent_registry.py
@@ -37,17 +37,70 @@
 from a2a.types import TransportProtocol as A2ATransport
 from google.adk.agents.readonly_context import ReadonlyContext
 from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
+from google.adk.telemetry.tracing import GCP_MCP_SERVER_DESTINATION_ID
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
+from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
 from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
 from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
 import google.auth
 import google.auth.transport.requests
 import httpx
+from mcp import StdioServerParameters
+from typing_extensions import override
 
 logger = logging.getLogger("google_adk." + __name__)
 
 AGENT_REGISTRY_BASE_URL = "https://agentregistry.googleapis.com/v1alpha"
 
 
+# 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
+# added to execute_tool spans in google.adk.telemetry.tracing
+class AgentRegistrySingleMcpToolset(McpToolset):
+
+  def __init__(
+      self,
+      *,
+      destination_resource_id: str | None,
+      connection_params: (
+          StdioServerParameters
+          | StdioConnectionParams
+          | SseConnectionParams
+          | StreamableHTTPConnectionParams
+      ),
+      tool_name_prefix: str | None = None,
+      header_provider: (
+          Callable[[ReadonlyContext], Dict[str, str]] | None
+      ) = None,
+  ):
+    super().__init__(
+        connection_params=connection_params,
+        tool_name_prefix=tool_name_prefix,
+        header_provider=header_provider,
+    )
+    self.destination_resource_id = destination_resource_id
+
+  @override
+  async def get_tools(
+      self, readonly_context: ReadonlyContext | None = None
+  ) -> List[BaseTool]:
+    tools = await super().get_tools(readonly_context)
+
+    # Noop if there is no destination_resource_id
+    if self.destination_resource_id is None:
+      return tools
+
+    for tool in tools:
+      if not tool.custom_metadata:
+        tool.custom_metadata = {}
+
+      tool.custom_metadata[GCP_MCP_SERVER_DESTINATION_ID] = (
+          self.destination_resource_id
+      )
+    return tools
+
+
 class _ProtocolType(str, Enum):
   """Supported agent protocol types."""
 
@@ -196,6 +249,9 @@ def get_mcp_toolset(self, mcp_server_name: str) -> McpToolset:
     """Constructs an McpToolset instance from a registered MCP Server."""
     server_details = self.get_mcp_server(mcp_server_name)
     name = self._clean_name(server_details.get("displayName", mcp_server_name))
+    mcp_server_id = server_details.get("mcpServerId")
+    if not isinstance(mcp_server_id, str):
+      mcp_server_id = None
 
     endpoint_uri = self._get_connection_uri(
         server_details, protocol_binding=A2ATransport.jsonrpc
@@ -210,7 +266,8 @@ def get_mcp_toolset(self, mcp_server_name: str) -> McpToolset:
     connection_params = StreamableHTTPConnectionParams(
         url=endpoint_uri, headers=self._get_auth_headers()
     )
-    return McpToolset(
+    return AgentRegistrySingleMcpToolset(
+        destination_resource_id=mcp_server_id,
         connection_params=connection_params,
         tool_name_prefix=name,
         header_provider=self._header_provider,
diff --git a/src/google/adk/telemetry/tracing.py b/src/google/adk/telemetry/tracing.py
@@ -83,6 +83,11 @@
 
 USER_CONTENT_ELIDED = '<elided>'
 
+# Used to associate a span with a destination resource for AppHub. Tools with
+# this key in their BaseTool.custom_metadata will have the mapping added as a
+# span attribute
+GCP_MCP_SERVER_DESTINATION_ID = 'gcp.mcp.server.destination.id'
+
 # Needed to avoid circular imports
 if TYPE_CHECKING:
   from ..agents.base_agent import BaseAgent
@@ -190,6 +195,14 @@ def trace_tool_call(
     else:
       span.set_attribute(ERROR_TYPE, type(error).__name__)
 
+  # Special case for client side association with a remote tool call
+  if (
+      tool.custom_metadata
+      and GCP_MCP_SERVER_DESTINATION_ID in tool.custom_metadata
+  ):
+    destination_id = tool.custom_metadata[GCP_MCP_SERVER_DESTINATION_ID]
+    span.set_attribute(GCP_MCP_SERVER_DESTINATION_ID, destination_id)
+
   # Setting empty llm request and response (as UI expect these) while not
   # applicable for tool_response.
   span.set_attribute('gcp.vertex.agent.llm_request', '{}')
diff --git a/tests/unittests/integrations/agent_registry/test_agent_registry.py b/tests/unittests/integrations/agent_registry/test_agent_registry.py
@@ -12,15 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
+from unittest.mock import AsyncMock
 from unittest.mock import MagicMock
 from unittest.mock import patch
 
 from a2a.types import TransportProtocol as A2ATransport
 from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
 from google.adk.integrations.agent_registry import _ProtocolType
 from google.adk.integrations.agent_registry import AgentRegistry
+from google.adk.telemetry.tracing import GCP_MCP_SERVER_DESTINATION_ID
 from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
 import httpx
+from mcp import ClientSession
+from mcp.types import ListToolsResult
+from mcp.types import Tool
 import pytest
 
 
@@ -31,6 +37,130 @@ def registry(self):
     with patch("google.auth.default", return_value=(MagicMock(), "project-id")):
       return AgentRegistry(project_id="test-project", location="global")
 
+  @pytest.mark.asyncio
+  @patch("httpx.Client")
+  @patch(
+      "google.adk.tools.mcp_tool.mcp_session_manager.MCPSessionManager.create_session",
+      new_callable=AsyncMock,
+  )
+  async def test_get_mcp_toolset_adds_destination_id(
+      self, mock_create_session, mock_httpx, registry
+  ):
+    """Test that tools from get_mcp_toolset have the destination ID."""
+    # Arrange
+    mcp_server_name = "test-mcp-server"
+    mock_api_response = MagicMock()
+    mock_api_response.json.return_value = {
+        "displayName": "TestPrefix",
+        "mcpServerId": (
+            "urn:mcp:googleapis.com:projects:1234:locations:global:bigquery"
+        ),
+        "interfaces": [{
+            "url": "https://mcp.com",
+            "protocolBinding": A2ATransport.jsonrpc,
+        }],
+    }
+    mock_httpx.return_value.__enter__.return_value.get.return_value = (
+        mock_api_response
+    )
+
+    registry._credentials.token = "token"
+    registry._credentials.refresh = MagicMock()
+
+    mock_session = AsyncMock(spec=ClientSession)
+    mock_create_session.return_value = mock_session
+
+    # Mock the tools returned by list_tools
+    mock_session.list_tools.return_value = ListToolsResult(
+        tools=[
+            Tool(
+                name="tool1",
+                description="d1",
+                inputs={},
+                outputs={},
+                inputSchema={},
+            ),
+            Tool(
+                name="tool2",
+                description="d2",
+                inputs={},
+                outputs={},
+                inputSchema={},
+            ),
+        ]
+    )
+
+    # Act
+    toolset = registry.get_mcp_toolset(mcp_server_name)
+    tools = await toolset.get_tools()
+
+    # Assert
+    assert isinstance(toolset, McpToolset)
+    mock_session.list_tools.assert_called_once_with()
+    assert len(tools) == 2
+    for tool in tools:
+      assert tool.custom_metadata is not None
+      assert (
+          tool.custom_metadata.get(GCP_MCP_SERVER_DESTINATION_ID)
+          == "urn:mcp:googleapis.com:projects:1234:locations:global:bigquery"
+      )
+
+  @pytest.mark.asyncio
+  @patch("httpx.Client")
+  @patch(
+      "google.adk.tools.mcp_tool.mcp_session_manager.MCPSessionManager.create_session",
+      new_callable=AsyncMock,
+  )
+  async def test_get_mcp_toolset_handles_missing_destination_id(
+      self, mock_create_session, mock_httpx, registry
+  ):
+    """Test get_mcp_toolset when the destination ID is missing."""
+    # Arrange
+    mcp_server_name = "test-mcp-server"
+    mock_api_response = MagicMock()
+    mock_api_response.json.return_value = {
+        "displayName": "TestPrefix",
+        # "mcpServerId" is intentionally omitted
+        "interfaces": [{
+            "url": "https://mcp.com",
+            "protocolBinding": A2ATransport.jsonrpc,
+        }],
+    }
+    mock_httpx.return_value.__enter__.return_value.get.return_value = (
+        mock_api_response
+    )
+
+    registry._credentials.token = "token"
+    registry._credentials.refresh = MagicMock()
+
+    mock_session = AsyncMock(spec=ClientSession)
+    mock_create_session.return_value = mock_session
+
+    # Mock the tools returned by list_tools
+    mock_session.list_tools.return_value = ListToolsResult(
+        tools=[
+            Tool(
+                name="tool1",
+                description="d1",
+                inputs={},
+                outputs={},
+                inputSchema={},
+            ),
+        ]
+    )
+
+    # Act
+    toolset = registry.get_mcp_toolset(mcp_server_name)
+    tools = await toolset.get_tools()
+
+    # Assert
+    assert isinstance(toolset, McpToolset)
+    mock_session.list_tools.assert_called_once_with()
+    assert len(tools) == 1
+    for tool in tools:
+      # The custom_metadata shouldn't have been added
+      assert tool.custom_metadata is None
+
   def test_init_raises_value_error_if_params_missing(self):
     with pytest.raises(
         ValueError, match="project_id and location must be provided"
diff --git a/tests/unittests/telemetry/test_spans.py b/tests/unittests/telemetry/test_spans.py
