feat: Add search agents and search MCP severs to agent registry

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

Author: wukath

contributing/samples/integrations/agent_registry_agent/README.md

@@ -44,6 +44,8 @@ The sample uses `AgentRegistry` to:
 
 - List registered agents using `list_agents()`.
 - List registered MCP servers using `list_mcp_servers()`.
+- Search registered agents using `search_agents(search_string)`.
+- Search registered MCP servers using `search_mcp_servers(search_string)`.
 
 It also shows (in comments) how to:
 

contributing/samples/integrations/agent_registry_agent/agent.py

@@ -45,6 +45,22 @@
 for endpoint in endpoints.get("endpoints", []):
   print(f"- Endpoint: {endpoint.get('displayName')} ({endpoint.get('name')})")
 
+# Search agents and MCP servers matching a query
+print(f"\nSearching agents matching 'Workspace' in {project_id}/{location}...")
+matching_agents = registry.search_agents(search_string="Workspace")
+for agent in matching_agents.get("agents", []):
+  print(f"- Found Agent: {agent.get('displayName')} ({agent.get('name')})")
+
+print(
+    "\nSearching MCP servers matching 'agentregistry' in"
+    f" {project_id}/{location}..."
+)
+matching_servers = registry.search_mcp_servers(search_string="agentregistry")
+for server in matching_servers.get("mcpServers", []):
+  print(
+      f"- Found 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)
 

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

@@ -24,6 +24,7 @@
 from typing import Callable
 from typing import Dict
 from typing import List
+from typing import Literal
 from typing import Mapping
 from typing import TypedDict
 from urllib.parse import urlparse
@@ -239,9 +240,13 @@ def _get_auth_headers(self) -> Dict[str, str]:
       ) from e
 
   def _make_request(
-      self, path: str, params: Dict[str, Any] | None = None
+      self,
+      path: str,
+      method: str = "GET",
+      params: Dict[str, Any] | None = None,
+      json_data: Dict[str, Any] | None = None,
   ) -> Dict[str, Any]:
-    """Helper function to make GET requests to the Agent Registry API."""
+    """Helper function to make requests to the Agent Registry API."""
     if path.startswith("projects/"):
       url = f"{self._base_url}/{path}"
     else:
@@ -254,7 +259,10 @@ def _make_request(
     )
     try:
       # Using AuthorizedSession for internal API calls to handle mTLS/Auth.
-      response = self._session.get(url, headers=headers, params=params)
+      if method == "POST":
+        response = self._session.post(url, headers=headers, json=json_data)
+      else:
+        response = self._session.get(url, headers=headers, params=params)
       response.raise_for_status()
       return response.json()
     except requests.exceptions.HTTPError as e:
@@ -267,12 +275,41 @@ def _make_request(
     except Exception as e:
       raise RuntimeError(f"API request failed: {e}") from e
 
+  def _search(
+      self,
+      resource_type: str,
+      *,
+      search_string: str | None = None,
+      search_type: Literal["KEYWORD", "SEMANTIC"] | None = None,
+      filter_str: str | None = None,
+      order_by: str | None = None,
+      page_size: int | None = None,
+      page_token: str | None = None,
+  ) -> Dict[str, Any]:
+    """Helper function to execute search requests."""
+    json_data: dict[str, Any] = {}
+    if search_string is not None:
+      json_data["searchString"] = search_string
+    if search_type is not None:
+      json_data["searchType"] = search_type
+    if filter_str is not None:
+      json_data["filter"] = filter_str
+    if order_by is not None:
+      json_data["orderBy"] = order_by
+    if page_size is not None:
+      json_data["pageSize"] = page_size
+    if page_token is not None:
+      json_data["pageToken"] = page_token
+    return self._make_request(
+        f"{resource_type}:search", method="POST", json_data=json_data
+    )
+
   def _get_connection_uri(
       self,
       resource_details: Mapping[str, Any],
       protocol_type: _ProtocolType | None = None,
       protocol_binding: _compat.TransportProtocol | None = None,
-  ) -> tuple[Any, Any, Any]:
+  ) -> tuple[str | None, str | None, _compat.TransportProtocol | None]:
     """Extracts the first matching URI based on type and binding filters."""
     protocols = list(resource_details.get("protocols", []))
     if "interfaces" in resource_details:
@@ -318,6 +355,27 @@ def list_mcp_servers(
       params["pageToken"] = page_token
     return self._make_request("mcpServers", params=params)
 
+  def search_mcp_servers(
+      self,
+      *,
+      search_string: str | None = None,
+      search_type: Literal["KEYWORD", "SEMANTIC"] | None = None,
+      filter_str: str | None = None,
+      order_by: str | None = None,
+      page_size: int | None = None,
+      page_token: str | None = None,
+  ) -> Dict[str, Any]:
+    """Searches registered MCP Servers."""
+    return self._search(
+        "mcpServers",
+        search_string=search_string,
+        search_type=search_type,
+        filter_str=filter_str,
+        order_by=order_by,
+        page_size=page_size,
+        page_token=page_token,
+    )
+
   def get_mcp_server(self, name: str) -> Dict[str, Any]:
     """Retrieves details of a specific MCP Server."""
     return self._make_request(name)
@@ -473,6 +531,27 @@ def list_agents(
       params["pageToken"] = page_token
     return self._make_request("agents", params=params)
 
+  def search_agents(
+      self,
+      *,
+      search_string: str | None = None,
+      search_type: Literal["KEYWORD", "SEMANTIC"] | None = None,
+      filter_str: str | None = None,
+      order_by: str | None = None,
+      page_size: int | None = None,
+      page_token: str | None = None,
+  ) -> Dict[str, Any]:
+    """Searches registered A2A Agents."""
+    return self._search(
+        "agents",
+        search_string=search_string,
+        search_type=search_type,
+        filter_str=filter_str,
+        order_by=order_by,
+        page_size=page_size,
+        page_token=page_token,
+    )
+
   def get_agent_info(self, name: str) -> Dict[str, Any]:
     """Retrieves detailed metadata of a specific A2A Agent."""
     return self._make_request(name)

tests/unittests/integrations/agent_registry/test_agent_registry.py

@@ -302,6 +302,74 @@ def test_list_agents(self, registry):
     agents = registry.list_agents()
     assert agents == {"agents": []}
 
+  def test_search_agents(self, registry):
+    """Tests search_agents API call."""
+    # pylint: disable=protected-access
+    mock_response = MagicMock()
+    mock_response.json.return_value = {"agents": [{"name": "agent-1"}]}
+    mock_response.raise_for_status = MagicMock()
+    registry._session.post.return_value = mock_response
+
+    registry._credentials.token = "token"
+    registry._credentials.refresh = MagicMock()
+
+    agents = registry.search_agents(
+        search_string="test-agent",
+        search_type="KEYWORD",
+        filter_str="display_name:test",
+        order_by="name",
+        page_size=10,
+        page_token="next-token",
+    )
+    assert agents == {"agents": [{"name": "agent-1"}]}
+    registry._session.post.assert_called_once_with(
+        f"{registry._base_url}/projects/test-project/locations/global/agents:search",
+        headers={"x-goog-user-project": "test-project"},
+        json={
+            "searchString": "test-agent",
+            "searchType": "KEYWORD",
+            "filter": "display_name:test",
+            "orderBy": "name",
+            "pageSize": 10,
+            "pageToken": "next-token",
+        },
+    )
+    # pylint: enable=protected-access
+
+  def test_search_mcp_servers(self, registry):
+    """Tests search_mcp_servers API call."""
+    # pylint: disable=protected-access
+    mock_response = MagicMock()
+    mock_response.json.return_value = {"mcpServers": [{"name": "mcp-1"}]}
+    mock_response.raise_for_status = MagicMock()
+    registry._session.post.return_value = mock_response
+
+    registry._credentials.token = "token"
+    registry._credentials.refresh = MagicMock()
+
+    mcp_servers = registry.search_mcp_servers(
+        search_string="test-mcp",
+        search_type="KEYWORD",
+        filter_str="display_name:test",
+        order_by="name",
+        page_size=10,
+        page_token="next-token",
+    )
+    assert mcp_servers == {"mcpServers": [{"name": "mcp-1"}]}
+    registry._session.post.assert_called_once_with(
+        f"{registry._base_url}/projects/test-project/locations/global/mcpServers:search",
+        headers={"x-goog-user-project": "test-project"},
+        json={
+            "searchString": "test-mcp",
+            "searchType": "KEYWORD",
+            "filter": "display_name:test",
+            "orderBy": "name",
+            "pageSize": 10,
+            "pageToken": "next-token",
+        },
+    )
+    # pylint: enable=protected-access
+
   def test_get_mcp_server(self, registry):
     mock_response = MagicMock()
     mock_response.json.return_value = {"name": "test-mcp"}
@@ -606,8 +674,8 @@ def test_get_auth_headers(self, registry):
 
   def test_make_request_raises_http_status_error(self, registry):
     mock_response = MagicMock()
-    mock_response.status_code = 404
-    mock_response.text = "Not Found"
+    mock_response.status_code = 500
+    mock_response.text = "Internal Server Error"
     error = requests.exceptions.HTTPError(
         "Error", request=MagicMock(), response=mock_response
     )
@@ -617,7 +685,7 @@ def test_make_request_raises_http_status_error(self, registry):
     registry._credentials.refresh = MagicMock()
 
     with pytest.raises(
-        RuntimeError, match="API request failed with status 404"
+        RuntimeError, match="API request failed with status 500"
     ):
       registry._make_request("test-path")