feat: Add auth scheme/credential support to MCP toolsets in Agent Registry

This is also so that we can add Agent Identity support when it's ready.

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

Author: wukath

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

@@ -25,10 +25,8 @@
 from typing import Dict
 from typing import List
 from typing import Mapping
-from typing import Optional
 from typing import Sequence
 from typing import TypedDict
-from typing import Union
 from urllib.parse import parse_qs
 from urllib.parse import urlparse
 
@@ -39,6 +37,8 @@
 from a2a.types import TransportProtocol as A2ATransport
 from google.adk.agents.readonly_context import ReadonlyContext
 from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_schemes import AuthScheme
 from google.adk.telemetry.tracing import GCP_MCP_SERVER_DESTINATION_ID
 from google.adk.tools.base_tool import BaseTool
 from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
@@ -75,11 +75,15 @@ def __init__(
       header_provider: (
           Callable[[ReadonlyContext], Dict[str, str]] | None
       ) = None,
+      auth_scheme: AuthScheme | None = None,
+      auth_credential: AuthCredential | None = None,
   ):
     super().__init__(
         connection_params=connection_params,
         tool_name_prefix=tool_name_prefix,
         header_provider=header_provider,
+        auth_scheme=auth_scheme,
+        auth_credential=auth_credential,
     )
     self.destination_resource_id = destination_resource_id
 
@@ -143,11 +147,11 @@ class AgentRegistry:
 
   def __init__(
       self,
-      project_id: Optional[str] = None,
-      location: Optional[str] = None,
-      header_provider: Optional[
-          Callable[[ReadonlyContext], Dict[str, str]]
-      ] = None,
+      project_id: str | None = None,
+      location: str | None = None,
+      header_provider: (
+          Callable[[ReadonlyContext], Dict[str, str]] | None
+      ) = None,
   ):
     """Initializes the AgentRegistry client.
 
@@ -190,7 +194,7 @@ def _get_auth_headers(self) -> Dict[str, str]:
       ) from e
 
   def _make_request(
-      self, path: str, params: Optional[Dict[str, Any]] = None
+      self, path: str, params: Dict[str, Any] | None = None
   ) -> Dict[str, Any]:
     """Helper function to make GET requests to the Agent Registry API."""
     if path.startswith("projects/"):
@@ -217,9 +221,9 @@ def _make_request(
   def _get_connection_uri(
       self,
       resource_details: Mapping[str, Any],
-      protocol_type: Optional[_ProtocolType] = None,
-      protocol_binding: Optional[A2ATransport] = None,
-  ) -> Optional[str]:
+      protocol_type: _ProtocolType | None = None,
+      protocol_binding: A2ATransport | None = None,
+  ) -> str | None:
     """Extracts the first matching URI based on type and binding filters."""
     protocols = list(resource_details.get("protocols", []))
     if "interfaces" in resource_details:
@@ -249,9 +253,9 @@ def _clean_name(self, name: str) -> str:
 
   def list_mcp_servers(
       self,
-      filter_str: Optional[str] = None,
-      page_size: Optional[int] = None,
-      page_token: Optional[str] = None,
+      filter_str: str | None = None,
+      page_size: int | None = None,
+      page_token: str | None = None,
   ) -> Dict[str, Any]:
     """Fetches a list of MCP Servers."""
     params = {}
@@ -267,7 +271,12 @@ def get_mcp_server(self, name: str) -> Dict[str, Any]:
     """Retrieves details of a specific MCP Server."""
     return self._make_request(name)
 
-  def get_mcp_toolset(self, mcp_server_name: str) -> McpToolset:
+  def get_mcp_toolset(
+      self,
+      mcp_server_name: str,
+      auth_scheme: AuthScheme | None = None,
+      auth_credential: AuthCredential | None = None,
+  ) -> McpToolset:
     """Constructs an McpToolset instance from a registered MCP Server."""
     server_details = self.get_mcp_server(mcp_server_name)
     name = self._clean_name(server_details.get("displayName", mcp_server_name))
@@ -293,15 +302,17 @@ def get_mcp_toolset(self, mcp_server_name: str) -> McpToolset:
         connection_params=connection_params,
         tool_name_prefix=name,
         header_provider=self._header_provider,
+        auth_scheme=auth_scheme,
+        auth_credential=auth_credential,
     )
 
   # --- Endpoint Methods ---
 
   def list_endpoints(
       self,
-      filter_str: Optional[str] = None,
-      page_size: Optional[int] = None,
-      page_token: Optional[str] = None,
+      filter_str: str | None = None,
+      page_size: int | None = None,
+      page_token: str | None = None,
   ) -> Dict[str, Any]:
     """Fetches a list of Endpoints."""
     params = {}
@@ -349,9 +360,9 @@ def get_model_name(self, endpoint_name: str) -> str:
 
   def list_agents(
       self,
-      filter_str: Optional[str] = None,
-      page_size: Optional[int] = None,
-      page_token: Optional[str] = None,
+      filter_str: str | None = None,
+      page_size: int | None = None,
+      page_token: str | None = None,
   ) -> Dict[str, Any]:
     """Fetches a list of registered A2A Agents."""
     params = {}

tests/unittests/integrations/agent_registry/test_agent_registry.py

@@ -18,7 +18,10 @@
 from unittest.mock import patch
 
 from a2a.types import TransportProtocol as A2ATransport
+from fastapi.openapi.models import OAuth2
 from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import OAuth2Auth
 from google.adk.integrations.agent_registry import _ProtocolType
 from google.adk.integrations.agent_registry import AgentRegistry
 from google.adk.telemetry.tracing import GCP_MCP_SERVER_DESTINATION_ID
@@ -325,6 +328,39 @@ def test_get_mcp_toolset(self, mock_httpx, registry):
     assert isinstance(toolset, McpToolset)
     assert toolset.tool_name_prefix == "TestPrefix"
 
+  @patch("httpx.Client")
+  def test_get_mcp_toolset_with_auth(self, mock_httpx, registry):
+    mock_response = MagicMock()
+    mock_response.json.return_value = {
+        "displayName": "TestPrefix",
+        "interfaces": [{
+            "url": "https://mcp.com",
+            "protocolBinding": A2ATransport.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()
+
+    auth_scheme = OAuth2(flows={})
+    auth_credential = AuthCredential(
+        auth_type="oauth2",
+        oauth2=OAuth2Auth(client_id="test_id", client_secret="test_secret"),
+    )
+
+    toolset = registry.get_mcp_toolset(
+        "test-mcp", auth_scheme=auth_scheme, auth_credential=auth_credential
+    )
+    assert isinstance(toolset, McpToolset)
+    auth_config = toolset.get_auth_config()
+    assert auth_config is not None
+    assert auth_config.auth_scheme == auth_scheme
+    assert auth_config.raw_auth_credential == auth_credential
+
   @patch("httpx.Client")
   def test_get_remote_a2a_agent(self, mock_httpx, registry):
     mock_response = MagicMock()