diff --git a/contributing/samples/agent_registry_agent/README.md b/contributing/samples/agent_registry_agent/README.md new file mode 100644 index 00000000..b9370b64 --- /dev/null +++ b/contributing/samples/agent_registry_agent/README.md @@ -0,0 +1,49 @@ +# Agent Registry Sample + +This sample demonstrates how to use the `AgentRegistry` client to discover agents and MCP servers registered in Google Cloud. + +## Setup + +1. Ensure you have Google Cloud credentials configured (e.g., `gcloud auth application-default login`). +2. Set the following environment variables: + +```bash +export GOOGLE_CLOUD_PROJECT=your-project-id +export GOOGLE_CLOUD_LOCATION=global # or your specific region +``` + +3. Obtain the full resource names for the agents and MCP servers you want to use. You can do this by running the sample script once to list them: + + ```bash + python3 agent.py + ``` + + Alternatively, use `gcloud` to list them: + + ```bash + # For agents + gcloud alpha agent-registry agents list --project=$GOOGLE_CLOUD_PROJECT --location=$GOOGLE_CLOUD_LOCATION + + # For MCP servers + gcloud alpha agent-registry mcp-servers list --project=$GOOGLE_CLOUD_PROJECT --location=$GOOGLE_CLOUD_LOCATION + ``` + +4. Replace `AGENT_NAME` and `MCP_SERVER_NAME` in `agent.py` with the last part of the resource names (e.g., if the name is `projects/.../agents/my-agent`, use `my-agent`). + +## Running the Sample + +Run the sample script to list available agents and MCP servers: + +```bash +python3 agent.py +``` + +## How it Works + +The sample uses `AgentRegistry` to: +- List registered agents using `list_agents()`. +- List registered MCP servers using `list_mcp_servers()`. + +It also shows (in comments) how to: +- Get a `RemoteA2aAgent` instance using `get_remote_a2a_agent(name)`. +- Get an `McpToolset` instance using `get_mcp_toolset(name)`. diff --git a/contributing/samples/agent_registry_agent/__init__.py b/contributing/samples/agent_registry_agent/__init__.py new file mode 100644 index 00000000..4015e47d --- /dev/null +++ b/contributing/samples/agent_registry_agent/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from . import agent diff --git a/contributing/samples/agent_registry_agent/agent.py b/contributing/samples/agent_registry_agent/agent.py new file mode 100644 index 00000000..38036dea --- /dev/null +++ b/contributing/samples/agent_registry_agent/agent.py @@ -0,0 +1,63 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sample agent demonstrating Agent Registry discovery.""" + +import os + +from google.adk.agents.llm_agent import LlmAgent +from google.adk.integrations.agent_registry import AgentRegistry + +# Project and location can be set via environment variables: +# GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION +project_id = os.environ.get("GOOGLE_CLOUD_PROJECT") +location = os.environ.get("GOOGLE_CLOUD_LOCATION", "global") + +# Initialize Agent Registry client +registry = AgentRegistry(project_id=project_id, location=location) + +print(f"Listing agents in {project_id}/{location}...") +agents = registry.list_agents() +for agent in agents.get("agents", []): + print(f"- Agent: {agent.get('displayName')} ({agent.get('name')})") + +print(f"\nListing MCP servers in {project_id}/{location}...") +mcp_servers = registry.list_mcp_servers() +for server in mcp_servers.get("mcpServers", []): + print(f"- MCP Server: {server.get('displayName')} ({server.get('name')})") + +# Example of using a specific agent or MCP server from the registry: +# (Note: These names should be full resource names as returned by list methods) + +# 1. Using a Remote A2A Agent as a sub-agent +# TODO: Replace AGENT_NAME with your agent name +remote_agent = registry.get_remote_a2a_agent( + f"projects/{project_id}/locations/{location}/agents/AGENT_NAME" +) + +# 2. Using an MCP Server in a toolset +# TODO: Replace MCP_SERVER_NAME with your MCP server name +mcp_toolset = registry.get_mcp_toolset( + f"projects/{project_id}/locations/{location}/mcpServers/MCP_SERVER_NAME" +) + +root_agent = LlmAgent( + model="gemini-2.5-flash", + name="discovery_agent", + instruction=( + "You have access to tools and sub-agents discovered via Registry." + ), + tools=[mcp_toolset], + sub_agents=[remote_agent], +) diff --git a/tests/unittests/integrations/agent_registry/__init__.py b/tests/unittests/integrations/agent_registry/__init__.py new file mode 100644 index 00000000..58d482ea --- /dev/null +++ b/tests/unittests/integrations/agent_registry/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unittests/integrations/agent_registry/test_agent_registry.py b/tests/unittests/integrations/agent_registry/test_agent_registry.py new file mode 100644 index 00000000..f54cdb67 --- /dev/null +++ b/tests/unittests/integrations/agent_registry/test_agent_registry.py @@ -0,0 +1,236 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock +from unittest.mock import patch + +from google.adk.agents.remote_a2a_agent import RemoteA2aAgent +from google.adk.integrations.agent_registry import AgentRegistry +from google.adk.tools.mcp_tool.mcp_toolset import McpToolset +import httpx +import pytest + + +class TestAgentRegistry: + + @pytest.fixture + def registry(self): + with patch("google.auth.default", return_value=(MagicMock(), "project-id")): + return AgentRegistry(project_id="test-project", location="global") + + def test_init_raises_value_error_if_params_missing(self): + with pytest.raises( + ValueError, match="project_id and location must be provided" + ): + AgentRegistry(project_id=None, location=None) + + def test_get_connection_uri_mcp_interfaces_top_level(self, registry): + resource_details = { + "interfaces": [ + {"url": "https://mcp-v1main.com", "protocolBinding": "JSONRPC"} + ] + } + uri = registry._get_connection_uri( + resource_details, protocol_binding="JSONRPC" + ) + assert uri == "https://mcp-v1main.com" + + def test_get_connection_uri_agent_nested_protocols(self, registry): + resource_details = { + "protocols": [{ + "type": "A2A_AGENT", + "interfaces": [{ + "url": "https://my-agent.com", + "protocolBinding": "JSONRPC", + }], + }] + } + uri = registry._get_connection_uri( + resource_details, protocol_type="A2A_AGENT" + ) + assert uri == "https://my-agent.com" + + def test_get_connection_uri_filtering(self, registry): + resource_details = { + "protocols": [ + { + "type": "CUSTOM", + "interfaces": [{"url": "https://custom.com"}], + }, + { + "type": "A2A_AGENT", + "interfaces": [{ + "url": "https://my-agent.com", + "protocolBinding": "HTTP_JSON", + }], + }, + ] + } + # Filter by type + uri = registry._get_connection_uri( + resource_details, protocol_type="A2A_AGENT" + ) + assert uri == "https://my-agent.com" + + # Filter by binding + uri = registry._get_connection_uri( + resource_details, protocol_binding="HTTP_JSON" + ) + assert uri == "https://my-agent.com" + + # No match + uri = registry._get_connection_uri( + resource_details, protocol_type="A2A_AGENT", protocol_binding="JSONRPC" + ) + assert uri is None + + def test_get_connection_uri_returns_none_if_no_interfaces(self, registry): + resource_details = {} + uri = registry._get_connection_uri(resource_details) + assert uri 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) + assert uri is None + + @patch("httpx.Client") + def test_list_agents(self, mock_httpx, registry): + mock_response = MagicMock() + mock_response.json.return_value = {"agents": []} + mock_response.raise_for_status = MagicMock() + mock_httpx.return_value.__enter__.return_value.get.return_value = ( + mock_response + ) + + # Mock auth refresh + registry._credentials.token = "token" + registry._credentials.refresh = MagicMock() + + agents = registry.list_agents() + assert agents == {"agents": []} + + @patch("httpx.Client") + def test_get_mcp_server(self, mock_httpx, registry): + mock_response = MagicMock() + mock_response.json.return_value = {"name": "test-mcp"} + 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() + + server = registry.get_mcp_server("test-mcp") + assert server == {"name": "test-mcp"} + + @patch("httpx.Client") + def test_get_mcp_toolset(self, mock_httpx, registry): + mock_response = MagicMock() + mock_response.json.return_value = { + "displayName": "TestPrefix", + "interfaces": [ + {"url": "https://mcp.com", "protocolBinding": "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() + + toolset = registry.get_mcp_toolset("test-mcp") + assert isinstance(toolset, McpToolset) + assert toolset.tool_name_prefix == "TestPrefix" + + @patch("httpx.Client") + def test_get_remote_a2a_agent(self, mock_httpx, registry): + mock_response = MagicMock() + mock_response.json.return_value = { + "displayName": "TestAgent", + "description": "Test Desc", + "agentSpec": { + "a2aAgentCardUrl": "https://my-agent.com/agent-card.json" + }, + } + 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.name == "TestAgent" + assert agent.description == "Test Desc" + assert agent._agent_card_source == "https://my-agent.com/agent-card.json" + + def test_get_auth_headers(self, registry): + registry._credentials.token = "fake-token" + registry._credentials.refresh = MagicMock() + registry._credentials.quota_project_id = "quota-project" + + headers = registry._get_auth_headers() + assert headers["Authorization"] == "Bearer fake-token" + assert headers["x-goog-user-project"] == "quota-project" + + @patch("httpx.Client") + def test_make_request_raises_http_status_error(self, mock_httpx, registry): + mock_response = MagicMock() + mock_response.status_code = 404 + mock_response.text = "Not Found" + error = httpx.HTTPStatusError( + "Error", request=MagicMock(), response=mock_response + ) + mock_httpx.return_value.__enter__.return_value.get.side_effect = error + + registry._credentials.token = "token" + registry._credentials.refresh = MagicMock() + + with pytest.raises( + RuntimeError, match="API request failed with status 404" + ): + registry._make_request("test-path") + + @patch("httpx.Client") + def test_make_request_raises_request_error(self, mock_httpx, registry): + error = httpx.RequestError("Connection failed", request=MagicMock()) + mock_httpx.return_value.__enter__.return_value.get.side_effect = error + + registry._credentials.token = "token" + registry._credentials.refresh = MagicMock() + + with pytest.raises( + RuntimeError, match="API request failed \(network error\)" + ): + registry._make_request("test-path") + + @patch("httpx.Client") + def test_make_request_raises_generic_exception(self, mock_httpx, registry): + mock_httpx.return_value.__enter__.return_value.get.side_effect = Exception( + "Generic error" + ) + + registry._credentials.token = "token" + registry._credentials.refresh = MagicMock() + + with pytest.raises(RuntimeError, match="API request failed: Generic error"): + registry._make_request("test-path")