fix: Add x-goog-user-project header to http calls in API Registry

This gcloud authentication header is required as a project override when using google ADC credentials. It is required for all OneMCP servers. Originally I manually put it in the big query sample, but since it is used for all OneMCP, it makes sense to just add it to the core API Registry logic. I also refactored the auth header logic a bit to deduplicate.

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

Author: wukath

contributing/samples/api_registry_agent/agent.py

@@ -21,11 +21,7 @@
 PROJECT_ID = "your-google-cloud-project-id"
 MCP_SERVER_NAME = "your-mcp-server-name"
 
-# Header required for BigQuery MCP server
-header_provider = lambda context: {
-    "x-goog-user-project": PROJECT_ID,
-}
-api_registry = ApiRegistry(PROJECT_ID, header_provider=header_provider)
+api_registry = ApiRegistry(PROJECT_ID)
 registry_tools = api_registry.get_toolset(
     mcp_server_name=MCP_SERVER_NAME,
 )

src/google/adk/tools/api_registry.py

@@ -16,6 +16,7 @@
 
 import sys
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import List
 from typing import Optional
@@ -59,12 +60,8 @@ def __init__(
 
     url = f"{API_REGISTRY_URL}/v1beta/projects/{self.api_registry_project_id}/locations/{self.location}/mcpServers"
     try:
-      request = google.auth.transport.requests.Request()
-      self._credentials.refresh(request)
-      headers = {
-          "Authorization": f"Bearer {self._credentials.token}",
-          "Content-Type": "application/json",
-      }
+      headers = self._get_auth_headers()
+      headers["Content-Type"] = "application/json"
       with httpx.Client() as client:
         response = client.get(url, headers=headers)
         response.raise_for_status()
@@ -107,11 +104,8 @@ def get_toolset(
       raise ValueError(f"MCP server {mcp_server_name} has no URLs.")
 
     mcp_server_url = server["urls"][0]
-    request = google.auth.transport.requests.Request()
-    self._credentials.refresh(request)
-    headers = {
-        "Authorization": f"Bearer {self._credentials.token}",
-    }
+    headers = self._get_auth_headers()
+
     return McpToolset(
         connection_params=StreamableHTTPConnectionParams(
             url="https://" + mcp_server_url,
@@ -121,3 +115,15 @@ def get_toolset(
         tool_name_prefix=tool_name_prefix,
         header_provider=self._header_provider,
     )
+
+  def _get_auth_headers(self) -> Dict[str, str]:
+    """Refreshes credentials and returns authorization headers."""
+    request = google.auth.transport.requests.Request()
+    self._credentials.refresh(request)
+    headers = {
+        "Authorization": f"Bearer {self._credentials.token}",
+    }
+    # Add quota project header if available in ADC
+    if self._credentials.quota_project_id:
+      headers["x-goog-user-project"] = self._credentials.quota_project_id
+    return headers

tests/unittests/tools/test_api_registry.py

@@ -14,6 +14,7 @@
 
 import sys
 import unittest
+from unittest.mock import create_autospec
 from unittest.mock import MagicMock
 from unittest.mock import patch
 
@@ -47,6 +48,7 @@ def setUp(self):
     self.mock_credentials = MagicMock()
     self.mock_credentials.token = "mock_token"
     self.mock_credentials.refresh = MagicMock()
+    self.mock_credentials.quota_project_id = None
     mock_auth_patcher = patch(
         "google.auth.default",
         return_value=(self.mock_credentials, None),
@@ -80,6 +82,32 @@ def test_init_success(self, MockHttpClient):
         },
     )
 
+  @patch("httpx.Client", autospec=True)
+  def test_init_with_quota_project_id_success(self, MockHttpClient):
+    self.mock_credentials.quota_project_id = "quota-project"
+    mock_response = create_autospec(httpx.Response, instance=True)
+    mock_response.json.return_value = MOCK_MCP_SERVERS_LIST
+    mock_client_instance = MockHttpClient.return_value
+    mock_client_instance.__enter__.return_value = mock_client_instance
+    mock_client_instance.get.return_value = mock_response
+
+    api_registry = ApiRegistry(
+        api_registry_project_id=self.project_id, location=self.location
+    )
+
+    self.assertEqual(len(api_registry._mcp_servers), 3)
+    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)
+    mock_client_instance.get.assert_called_once_with(
+        f"https://cloudapiregistry.googleapis.com/v1beta/projects/{self.project_id}/locations/{self.location}/mcpServers",
+        headers={
+            "Authorization": "Bearer mock_token",
+            "Content-Type": "application/json",
+            "x-goog-user-project": "quota-project",
+        },
+    )
+
   @patch("httpx.Client", autospec=True)
   def test_init_http_error(self, MockHttpClient):
     mock_client_instance = MockHttpClient.return_value
@@ -138,6 +166,38 @@ async def test_get_toolset_success(self, MockHttpClient, MockMcpToolset):
     )
     self.assertEqual(toolset, MockMcpToolset.return_value)
 
+  @patch("google.adk.tools.api_registry.McpToolset", autospec=True)
+  @patch("httpx.Client", autospec=True)
+  async def test_get_toolset_with_quota_project_id_success(
+      self, MockHttpClient, MockMcpToolset
+  ):
+    self.mock_credentials.quota_project_id = "quota-project"
+    mock_response = create_autospec(httpx.Response, instance=True)
+    mock_response.json.return_value = MOCK_MCP_SERVERS_LIST
+    mock_client_instance = MockHttpClient.return_value
+    mock_client_instance.__enter__.return_value = mock_client_instance
+    mock_client_instance.get.return_value = mock_response
+
+    api_registry = ApiRegistry(
+        api_registry_project_id=self.project_id, location=self.location
+    )
+
+    toolset = api_registry.get_toolset("test-mcp-server-1")
+
+    MockMcpToolset.assert_called_once_with(
+        connection_params=StreamableHTTPConnectionParams(
+            url="https://mcp.server1.com",
+            headers={
+                "Authorization": "Bearer mock_token",
+                "x-goog-user-project": "quota-project",
+            },
+        ),
+        tool_filter=None,
+        tool_name_prefix=None,
+        header_provider=None,
+    )
+    self.assertEqual(toolset, MockMcpToolset.return_value)
+
   @patch("google.adk.tools.api_registry.McpToolset", autospec=True)
   @patch("httpx.Client", autospec=True)
   async def test_get_toolset_with_filter_and_prefix(