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

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 @@ class ApiRegistry:
     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 @@ class TestApiRegistry(unittest.IsolatedAsyncioTestCase):
             "Authorization": "Bearer mock_token",
             "Content-Type": "application/json",
         },
+        params={},
     )
 
   @patch("httpx.Client", autospec=True)
@@ -119,6 +120,71 @@ class TestApiRegistry(unittest.IsolatedAsyncioTestCase):
             "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)