mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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 https://github.com/google/adk-python/issues/2168. Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com> PiperOrigin-RevId: 859201722
This commit is contained in:
Xiang (Sean) Zhou
committed by
Copybara-Service
Copybara-Service
parent
d62f9c896c
commit
e3d542a5ba
@@ -1,110 +0,0 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Utility functions for MCP tool authentication."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import logging
|
||||
from typing import Dict
|
||||
from typing import Optional
|
||||
|
||||
from fastapi.openapi import models as openapi_models
|
||||
from fastapi.openapi.models import APIKey
|
||||
from fastapi.openapi.models import HTTPBase
|
||||
|
||||
from ...auth.auth_credential import AuthCredential
|
||||
from ...auth.auth_schemes import AuthScheme
|
||||
|
||||
logger = logging.getLogger("google_adk." + __name__)
|
||||
|
||||
|
||||
def get_mcp_auth_headers(
|
||||
auth_scheme: Optional[AuthScheme], credential: Optional[AuthCredential]
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""Generates HTTP authentication headers for MCP calls.
|
||||
|
||||
Args:
|
||||
auth_scheme: The authentication scheme.
|
||||
credential: The resolved authentication credential.
|
||||
|
||||
Returns:
|
||||
A dictionary of headers, or None if no auth is applicable.
|
||||
|
||||
Raises:
|
||||
ValueError: If the auth scheme is unsupported or misconfigured.
|
||||
"""
|
||||
if not credential:
|
||||
return None
|
||||
|
||||
headers: Optional[Dict[str, str]] = None
|
||||
|
||||
if credential.oauth2:
|
||||
headers = {"Authorization": f"Bearer {credential.oauth2.access_token}"}
|
||||
elif credential.http:
|
||||
if not auth_scheme or not isinstance(auth_scheme, HTTPBase):
|
||||
logger.warning(
|
||||
"HTTP credential provided, but auth_scheme is missing or not"
|
||||
" HTTPBase."
|
||||
)
|
||||
return None
|
||||
|
||||
scheme = auth_scheme.scheme.lower()
|
||||
if scheme == "bearer" and credential.http.credentials.token:
|
||||
headers = {"Authorization": f"Bearer {credential.http.credentials.token}"}
|
||||
elif scheme == "basic":
|
||||
if (
|
||||
credential.http.credentials.username
|
||||
and credential.http.credentials.password
|
||||
):
|
||||
creds = f"{credential.http.credentials.username}:{credential.http.credentials.password}"
|
||||
encoded_creds = base64.b64encode(creds.encode()).decode()
|
||||
headers = {"Authorization": f"Basic {encoded_creds}"}
|
||||
else:
|
||||
logger.warning("Basic auth scheme missing username or password.")
|
||||
elif credential.http.credentials.token:
|
||||
# Handle other HTTP schemes like Digest, etc. if token is present
|
||||
headers = {
|
||||
"Authorization": (
|
||||
f"{auth_scheme.scheme} {credential.http.credentials.token}"
|
||||
)
|
||||
}
|
||||
else:
|
||||
logger.warning(f"Unsupported or incomplete HTTP auth scheme '{scheme}'.")
|
||||
elif credential.api_key:
|
||||
if not auth_scheme or not isinstance(auth_scheme, APIKey):
|
||||
logger.warning(
|
||||
"API key credential provided, but auth_scheme is missing or not"
|
||||
" APIKey."
|
||||
)
|
||||
return None
|
||||
|
||||
if auth_scheme.in_ != openapi_models.APIKeyIn.header:
|
||||
error_msg = (
|
||||
"MCP tools only support header-based API key authentication. "
|
||||
f"Configured location: {auth_scheme.in_}"
|
||||
)
|
||||
logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
headers = {auth_scheme.name: credential.api_key}
|
||||
elif credential.service_account:
|
||||
logger.warning(
|
||||
"Service account credentials should be exchanged for an access token "
|
||||
"before calling get_mcp_auth_headers."
|
||||
)
|
||||
else:
|
||||
logger.warning(f"Unsupported credential type: {type(credential)}")
|
||||
|
||||
return headers
|
||||
@@ -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 Optional
|
||||
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 .._gemini_schema_util import _to_gemini_schema
|
||||
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 @@ class McpTool(BaseAuthenticatedTool):
|
||||
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 @@ class McpTool(BaseAuthenticatedTool):
|
||||
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."""
|
||||
|
||||
@@ -33,14 +33,11 @@ from typing_extensions import override
|
||||
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 @@ class McpToolset(BaseToolset):
|
||||
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 = (
|
||||
|
||||
Reference in New Issue
Block a user