feat: Add Auth Provider support to agent registry

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

Author: wukath

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

@@ -36,6 +36,7 @@
 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.integrations.agent_identity.gcp_auth_provider_scheme import GcpAuthProviderScheme
 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
@@ -292,8 +293,24 @@ def get_mcp_toolset(
       mcp_server_name: str,
       auth_scheme: AuthScheme | None = None,
       auth_credential: AuthCredential | None = None,
+      *,
+      continue_uri: str | None = None,
   ) -> McpToolset:
-    """Constructs an McpToolset instance from a registered MCP Server."""
+    """Constructs an McpToolset from a registered MCP Server.
+
+    If `auth_scheme` is omitted, it is automatically resolved from the server's
+    IAM bindings via `GcpAuthProviderScheme`.
+
+    Args:
+      mcp_server_name: Resource name of the MCP Server.
+      auth_scheme: Optional auth scheme. Resolved via bindings if omitted.
+      auth_credential: Optional auth credential.
+      continue_uri: Optional continue URI to override what is in the auth
+        provider.
+
+    Returns:
+      An McpToolset for the MCP server.
+    """
     server_details = self.get_mcp_server(mcp_server_name)
     name = self._clean_name(server_details.get("displayName", mcp_server_name))
     mcp_server_id = server_details.get("mcpServerId")
@@ -313,6 +330,23 @@ def get_mcp_toolset(
       )
 
     headers = self._get_auth_headers() if _is_google_api(endpoint_uri) else None
+    if mcp_server_id and not auth_scheme:
+      try:
+        bindings_data = self._make_request("bindings")
+        for b in bindings_data.get("bindings", []):
+          target_id = b.get("target", {}).get("identifier", "")
+          if target_id.endswith(mcp_server_id):
+            auth_provider = b.get("authProviderBinding", {}).get("authProvider")
+            if auth_provider:
+              auth_scheme = GcpAuthProviderScheme(
+                  name=auth_provider, continue_uri=continue_uri
+              )
+              break
+      except Exception as e:
+        logger.warning(
+            f"Failed to fetch bindings for MCP Server {mcp_server_name}: {e}"
+        )
+
     connection_params = StreamableHTTPConnectionParams(
         url=endpoint_uri,
         headers=headers,

tests/unittests/integrations/agent_registry/test_agent_registry.py

@@ -637,3 +637,48 @@ def test_get_model_name_raises_value_error_if_no_uri(
     mock_get_endpoint.return_value = {}
     with pytest.raises(ValueError, match="Connection URI not found"):
       registry.get_model_name("test-endpoint")
+
+  @patch.object(AgentRegistry, "_make_request")
+  def test_get_mcp_toolset_with_binding(self, mock_make_request, registry):
+    def side_effect(*args, **kwargs):
+      if args[0] == "test-mcp":
+        return {
+            "displayName": "TestPrefix",
+            "mcpServerId": "server-456",
+            "interfaces": [{
+                "url": "https://mcp.com",
+                "protocolBinding": "JSONRPC",
+            }],
+        }
+      if args[0] == "bindings":
+        return {
+            "bindings": [{
+                "target": {
+                    "identifier": (
+                        "urn:mcp:projects-123:projects:123:locations:l:mcpServers:server-456"
+                    )
+                },
+                "authProviderBinding": {
+                    "authProvider": (
+                        "projects/123/locations/l/authProviders/ap-789"
+                    )
+                },
+            }]
+        }
+      return {}
+
+    mock_make_request.side_effect = side_effect
+
+    registry._credentials.token = "token"
+    registry._credentials.refresh = MagicMock()
+
+    toolset = registry.get_mcp_toolset(
+        "test-mcp", continue_uri="https://override.com/continue"
+    )
+    assert isinstance(toolset, McpToolset)
+    assert toolset._auth_scheme is not None
+    assert (
+        toolset._auth_scheme.name
+        == "projects/123/locations/l/authProviders/ap-789"
+    )
+    assert toolset._auth_scheme.continue_uri == "https://override.com/continue"