feat: Support authentication for MCP tool listing

Currently only tool calling supports MCP auth. This refactors the auth logic into a auth_utils file and uses it for tool listing as well. Fixes #2168.

Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com>
PiperOrigin-RevId: 859201722

Author: seanzhougoogle

src/google/adk/tools/mcp_tool/mcp_auth_utils.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import base64
 import inspect
 import logging
 from typing import Any
@@ -23,6 +24,7 @@
 from typing import Union
 import warnings
 
+from fastapi.openapi.models import APIKeyIn
 from google.genai.types import FunctionDeclaration
 from mcp.types import Tool as McpBaseTool
 from typing_extensions import override
@@ -37,7 +39,6 @@
 from ..base_authenticated_tool import BaseAuthenticatedTool
 #  import
 from ..tool_context import ToolContext
-from .mcp_auth_utils import get_mcp_auth_headers
 from .mcp_session_manager import MCPSessionManager
 from .mcp_session_manager import retry_on_errors
 
@@ -194,12 +195,7 @@ async def _run_async_impl(
         Any: The response from the tool.
     """
     # Extract headers from credential for session pooling
-    auth_scheme = (
-        self._auth_config.auth_scheme
-        if hasattr(self, "_auth_config") and self._auth_config
-        else None
-    )
-    auth_headers = get_mcp_auth_headers(auth_scheme, credential)
+    auth_headers = await self._get_headers(tool_context, credential)
     dynamic_headers = None
     if self._header_provider:
       dynamic_headers = self._header_provider(
@@ -221,6 +217,90 @@ async def _run_async_impl(
     response = await session.call_tool(self._mcp_tool.name, arguments=args)
     return response.model_dump(exclude_none=True, mode="json")
 
+  async def _get_headers(
+      self, tool_context: ToolContext, credential: AuthCredential
+  ) -> Optional[dict[str, str]]:
+    """Extracts authentication headers from credentials.
+
+    Args:
+        tool_context: The tool context of the current invocation.
+        credential: The authentication credential to process.
+
+    Returns:
+        Dictionary of headers to add to the request, or None if no auth.
+
+    Raises:
+        ValueError: If API key authentication is configured for non-header location.
+    """
+    headers: Optional[dict[str, str]] = None
+    if credential:
+      if credential.oauth2:
+        headers = {"Authorization": f"Bearer {credential.oauth2.access_token}"}
+      elif credential.http:
+        # Handle HTTP authentication schemes
+        if (
+            credential.http.scheme.lower() == "bearer"
+            and credential.http.credentials.token
+        ):
+          headers = {
+              "Authorization": f"Bearer {credential.http.credentials.token}"
+          }
+        elif credential.http.scheme.lower() == "basic":
+          # Handle basic auth
+          if (
+              credential.http.credentials.username
+              and credential.http.credentials.password
+          ):
+
+            credentials = f"{credential.http.credentials.username}:{credential.http.credentials.password}"
+            encoded_credentials = base64.b64encode(
+                credentials.encode()
+            ).decode()
+            headers = {"Authorization": f"Basic {encoded_credentials}"}
+        elif credential.http.credentials.token:
+          # Handle other HTTP schemes with token
+          headers = {
+              "Authorization": (
+                  f"{credential.http.scheme} {credential.http.credentials.token}"
+              )
+          }
+      elif credential.api_key:
+        if (
+            not self._credentials_manager
+            or not self._credentials_manager._auth_config
+        ):
+          error_msg = (
+              "Cannot find corresponding auth scheme for API key credential"
+              f" {credential}"
+          )
+          logger.error(error_msg)
+          raise ValueError(error_msg)
+        elif (
+            self._credentials_manager._auth_config.auth_scheme.in_
+            != APIKeyIn.header
+        ):
+          error_msg = (
+              "McpTool only supports header-based API key authentication."
+              " Configured location:"
+              f" {self._credentials_manager._auth_config.auth_scheme.in_}"
+          )
+          logger.error(error_msg)
+          raise ValueError(error_msg)
+        else:
+          headers = {
+              self._credentials_manager._auth_config.auth_scheme.name: (
+                  credential.api_key
+              )
+          }
+      elif credential.service_account:
+        # Service accounts should be exchanged for access tokens before reaching this point
+        logger.warning(
+            "Service account credentials should be exchanged before MCP"
+            " session creation"
+        )
+
+    return headers
+
 
 class MCPTool(McpTool):
   """Deprecated name, use `McpTool` instead."""

src/google/adk/tools/mcp_tool/mcp_tool.py

@@ -33,14 +33,11 @@
 from ...agents.readonly_context import ReadonlyContext
 from ...auth.auth_credential import AuthCredential
 from ...auth.auth_schemes import AuthScheme
-from ...auth.auth_tool import AuthConfig
-from ...auth.credential_manager import CredentialManager
 from ..base_tool import BaseTool
 from ..base_toolset import BaseToolset
 from ..base_toolset import ToolPredicate
 from ..tool_configs import BaseToolConfig
 from ..tool_configs import ToolArgsConfig
-from .mcp_auth_utils import get_mcp_auth_headers
 from .mcp_session_manager import MCPSessionManager
 from .mcp_session_manager import retry_on_errors
 from .mcp_session_manager import SseConnectionParams
@@ -157,50 +154,13 @@ async def get_tools(
     Returns:
         List[BaseTool]: A list of tools available under the specified context.
     """
-    provided_headers = (
+    headers = (
         self._header_provider(readonly_context)
         if self._header_provider and readonly_context
-        else {}
+        else None
     )
-
-    auth_headers = {}
-    if self._auth_scheme:
-      try:
-        # Instantiate CredentialsManager to resolve credentials
-        auth_config = AuthConfig(
-            auth_scheme=self._auth_scheme,
-            raw_auth_credential=self._auth_credential,
-        )
-        credentials_manager = CredentialManager(auth_config)
-
-        # Resolve the credential
-        resolved_credential = await credentials_manager.get_auth_credential(
-            readonly_context
-        )
-
-        if resolved_credential:
-          auth_headers = get_mcp_auth_headers(
-              self._auth_scheme, resolved_credential
-          )
-        else:
-          logger.warning(
-              "Failed to resolve credential for tool listing, proceeding"
-              " without auth headers."
-          )
-      except Exception as e:
-        logger.warning(
-            "Error generating auth headers for tool listing: %s, proceeding"
-            " without auth headers.",
-            e,
-            exc_info=True,
-        )
-
-    merged_headers = {**(provided_headers or {}), **(auth_headers or {})}
-
     # Get session from session manager
-    session = await self._mcp_session_manager.create_session(
-        headers=merged_headers
-    )
+    session = await self._mcp_session_manager.create_session(headers=headers)
 
     # Fetch available tools from the MCP server
     timeout_in_seconds = (