feat: Support page token in API Registry

API registry uses a nextPageToken for pagination, so we should loop through all the pages by sending the received nextPageToken to subsequent requests until no token is returned. The first page contains 1P MCP servers, and later pages contain 3P customer servers.

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

Author: wukath

src/google/adk/tools/api_registry.py

@@ -14,7 +14,6 @@
 
 from __future__ import annotations
 
-import sys
 from typing import Any
 from typing import Callable
 
@@ -56,17 +55,29 @@ def __init__(
     self._header_provider = header_provider
 
     url = f"{API_REGISTRY_URL}/v1beta/projects/{self.api_registry_project_id}/locations/{self.location}/mcpServers"
+
     try:
       headers = self._get_auth_headers()
       headers["Content-Type"] = "application/json"
+      page_token = None
       with httpx.Client() as client:
-        response = client.get(url, headers=headers)
-        response.raise_for_status()
-        mcp_servers_list = response.json().get("mcpServers", [])
-        for server in mcp_servers_list:
-          server_name = server.get("name", "")
-          if server_name:
-            self._mcp_servers[server_name] = server
+        while True:
+          params = {}
+          if page_token:
+            params["pageToken"] = page_token
+
+          response = client.get(url, headers=headers, params=params)
+          response.raise_for_status()
+          data = response.json()
+          mcp_servers_list = data.get("mcpServers", [])
+          for server in mcp_servers_list:
+            server_name = server.get("name", "")
+            if server_name:
+              self._mcp_servers[server_name] = server
+
+          page_token = data.get("nextPageToken")
+          if not page_token:
+            break
     except (httpx.HTTPError, ValueError) as e:
       # Handle error in fetching or parsing tool definitions
       raise RuntimeError(

tests/unittests/tools/test_api_registry.py

@@ -91,6 +91,7 @@ def test_init_success(self, MockHttpClient):
             "Authorization": "Bearer mock_token",
             "Content-Type": "application/json",
         },
+        params={},
     )
 
   @patch("httpx.Client", autospec=True)
@@ -119,6 +120,71 @@ def test_init_with_quota_project_id_success(self, MockHttpClient):
             "Content-Type": "application/json",
             "x-goog-user-project": "quota-project",
         },
+        params={},
+    )
+
+  @patch("httpx.Client", autospec=True)
+  def test_init_with_pagination_success(self, MockHttpClient):
+    mock_response1 = create_autospec(httpx.Response, instance=True)
+    mock_response1.json.return_value = {
+        "mcpServers": [
+            {
+                "name": "test-mcp-server-1",
+                "urls": ["mcp.server1.com"],
+            },
+            {
+                "name": "test-mcp-server-2",
+                "urls": ["mcp.server2.com"],
+            },
+        ],
+        "nextPageToken": "next_page_token",
+    }
+    mock_response2 = create_autospec(httpx.Response, instance=True)
+    mock_response2.json.return_value = {
+        "mcpServers": [
+            {
+                "name": "test-mcp-server-no-url",
+            },
+            {
+                "name": "test-mcp-server-http",
+                "urls": ["http://mcp.server_http.com"],
+            },
+            {
+                "name": "test-mcp-server-https",
+                "urls": ["https://mcp.server_https.com"],
+            },
+        ]
+    }
+    mock_client_instance = MockHttpClient.return_value
+    mock_client_instance.__enter__.return_value = mock_client_instance
+    mock_client_instance.get.side_effect = [mock_response1, mock_response2]
+
+    api_registry = ApiRegistry(
+        api_registry_project_id=self.project_id, location=self.location
+    )
+
+    self.assertEqual(len(api_registry._mcp_servers), 5)
+    self.assertIn("test-mcp-server-1", api_registry._mcp_servers)
+    self.assertIn("test-mcp-server-2", api_registry._mcp_servers)
+    self.assertIn("test-mcp-server-no-url", api_registry._mcp_servers)
+    self.assertIn("test-mcp-server-http", api_registry._mcp_servers)
+    self.assertIn("test-mcp-server-https", api_registry._mcp_servers)
+    self.assertEqual(mock_client_instance.get.call_count, 2)
+    mock_client_instance.get.assert_any_call(
+        f"https://cloudapiregistry.googleapis.com/v1beta/projects/{self.project_id}/locations/{self.location}/mcpServers",
+        headers={
+            "Authorization": "Bearer mock_token",
+            "Content-Type": "application/json",
+        },
+        params={},
+    )
+    mock_client_instance.get.assert_called_with(
+        f"https://cloudapiregistry.googleapis.com/v1beta/projects/{self.project_id}/locations/{self.location}/mcpServers",
+        headers={
+            "Authorization": "Bearer mock_token",
+            "Content-Type": "application/json",
+        },
+        params={"pageToken": "next_page_token"},
     )
 
   @patch("httpx.Client", autospec=True)