fix: Support API-Key for MCP Tool authentication

PiperOrigin-RevId: 776641474

Author: seanzhougoogle

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

@@ -15,12 +15,11 @@
 from __future__ import annotations
 
 import base64
-import json
 import logging
 from typing import Optional
 
+from fastapi.openapi.models import APIKeyIn
 from google.genai.types import FunctionDeclaration
-from google.oauth2.credentials import Credentials
 from typing_extensions import override
 
 from .._gemini_schema_util import _to_gemini_schema
@@ -58,6 +57,9 @@ class MCPTool(BaseAuthenticatedTool):
 
  Internally, the tool initializes from a MCP Tool, and uses the MCP Session to
  call the tool.
+
+ Note: For API key authentication, only header-based API keys are supported.
+ Query and cookie-based API keys will result in authentication errors.
  """
 
  def __init__(
@@ -134,7 +136,19 @@ async def _run_async_impl(
  async def _get_headers(
  self, tool_context: ToolContext, credential: AuthCredential
  ) -> Optional[dict[str, str]]:
- headers = None
+ """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}"}
@@ -167,10 +181,33 @@ async def _get_headers(
  )
  }
  elif credential.api_key:
- # For API keys, we'll add them as headers since MCP typically uses header-based auth
- # The specific header name would depend on the API, using a common default
- # TODO Allow user to specify the header name for API keys.
- headers = {"X-API-Key": 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(

tests/unittests/tools/mcp_tool/test_mcp_tool.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 
 import sys
-from typing import Any
-from typing import Dict
 from unittest.mock import AsyncMock
 from unittest.mock import Mock
 from unittest.mock import patch
@@ -268,8 +266,102 @@ async def test_get_headers_http_basic(self):
  assert headers == {"Authorization": f"Basic {expected_encoded}"}
 
  @pytest.mark.asyncio
- async def test_get_headers_api_key(self):
- """Test header generation for API Key credentials."""
+ async def test_get_headers_api_key_with_valid_header_scheme(self):
+ """Test header generation for API Key credentials with header-based auth scheme."""
+ from fastapi.openapi.models import APIKey
+ from fastapi.openapi.models import APIKeyIn
+ from google.adk.auth.auth_schemes import AuthSchemeType
+
+ # Create auth scheme for header-based API key
+ auth_scheme = APIKey(**{
+ "type": AuthSchemeType.apiKey,
+ "in": APIKeyIn.header,
+ "name": "X-Custom-API-Key",
+ })
+ auth_credential = AuthCredential(
+ auth_type=AuthCredentialTypes.API_KEY, api_key="my_api_key"
+ )
+
+ tool = MCPTool(
+ mcp_tool=self.mock_mcp_tool,
+ mcp_session_manager=self.mock_session_manager,
+ auth_scheme=auth_scheme,
+ auth_credential=auth_credential,
+ )
+
+ tool_context = Mock(spec=ToolContext)
+ headers = await tool._get_headers(tool_context, auth_credential)
+
+ assert headers == {"X-Custom-API-Key": "my_api_key"}
+
+ @pytest.mark.asyncio
+ async def test_get_headers_api_key_with_query_scheme_raises_error(self):
+ """Test that API Key with query-based auth scheme raises ValueError."""
+ from fastapi.openapi.models import APIKey
+ from fastapi.openapi.models import APIKeyIn
+ from google.adk.auth.auth_schemes import AuthSchemeType
+
+ # Create auth scheme for query-based API key (not supported)
+ auth_scheme = APIKey(**{
+ "type": AuthSchemeType.apiKey,
+ "in": APIKeyIn.query,
+ "name": "api_key",
+ })
+ auth_credential = AuthCredential(
+ auth_type=AuthCredentialTypes.API_KEY, api_key="my_api_key"
+ )
+
+ tool = MCPTool(
+ mcp_tool=self.mock_mcp_tool,
+ mcp_session_manager=self.mock_session_manager,
+ auth_scheme=auth_scheme,
+ auth_credential=auth_credential,
+ )
+
+ tool_context = Mock(spec=ToolContext)
+
+ with pytest.raises(
+ ValueError,
+ match="MCPTool only supports header-based API key authentication",
+ ):
+ await tool._get_headers(tool_context, auth_credential)
+
+ @pytest.mark.asyncio
+ async def test_get_headers_api_key_with_cookie_scheme_raises_error(self):
+ """Test that API Key with cookie-based auth scheme raises ValueError."""
+ from fastapi.openapi.models import APIKey
+ from fastapi.openapi.models import APIKeyIn
+ from google.adk.auth.auth_schemes import AuthSchemeType
+
+ # Create auth scheme for cookie-based API key (not supported)
+ auth_scheme = APIKey(**{
+ "type": AuthSchemeType.apiKey,
+ "in": APIKeyIn.cookie,
+ "name": "session_id",
+ })
+ auth_credential = AuthCredential(
+ auth_type=AuthCredentialTypes.API_KEY, api_key="my_api_key"
+ )
+
+ tool = MCPTool(
+ mcp_tool=self.mock_mcp_tool,
+ mcp_session_manager=self.mock_session_manager,
+ auth_scheme=auth_scheme,
+ auth_credential=auth_credential,
+ )
+
+ tool_context = Mock(spec=ToolContext)
+
+ with pytest.raises(
+ ValueError,
+ match="MCPTool only supports header-based API key authentication",
+ ):
+ await tool._get_headers(tool_context, auth_credential)
+
+ @pytest.mark.asyncio
+ async def test_get_headers_api_key_without_auth_config_raises_error(self):
+ """Test that API Key without auth config raises ValueError."""
+ # Create tool without auth scheme/config
  tool = MCPTool(
  mcp_tool=self.mock_mcp_tool,
  mcp_session_manager=self.mock_session_manager,
@@ -278,11 +370,37 @@ async def test_get_headers_api_key(self):
  credential = AuthCredential(
  auth_type=AuthCredentialTypes.API_KEY, api_key="my_api_key"
  )
+ tool_context = Mock(spec=ToolContext)
+
+ with pytest.raises(
+ ValueError,
+ match="Cannot find corresponding auth scheme for API key credential",
+ ):
+ await tool._get_headers(tool_context, credential)
+
+ @pytest.mark.asyncio
+ async def test_get_headers_api_key_without_credentials_manager_raises_error(
+ self,
+ ):
+ """Test that API Key without credentials manager raises ValueError."""
+ tool = MCPTool(
+ mcp_tool=self.mock_mcp_tool,
+ mcp_session_manager=self.mock_session_manager,
+ )
 
+ # Manually set credentials manager to None to simulate error condition
+ tool._credentials_manager = None
+
+ credential = AuthCredential(
+ auth_type=AuthCredentialTypes.API_KEY, api_key="my_api_key"
+ )
  tool_context = Mock(spec=ToolContext)
- headers = await tool._get_headers(tool_context, credential)
 
- assert headers == {"X-API-Key": "my_api_key"}
+ with pytest.raises(
+ ValueError,
+ match="Cannot find corresponding auth scheme for API key credential",
+ ):
+ await tool._get_headers(tool_context, credential)
 
  @pytest.mark.asyncio
  async def test_get_headers_no_credential(self):
@@ -318,6 +436,48 @@ async def test_get_headers_service_account(self):
  # Should return None as service account credentials are not supported for direct header generation
  assert headers is None
 
+ @pytest.mark.asyncio
+ async def test_run_async_impl_with_api_key_header_auth(self):
+ """Test running tool with API key header authentication end-to-end."""
+ from fastapi.openapi.models import APIKey
+ from fastapi.openapi.models import APIKeyIn
+ from google.adk.auth.auth_schemes import AuthSchemeType
+
+ # Create auth scheme for header-based API key
+ auth_scheme = APIKey(**{
+ "type": AuthSchemeType.apiKey,
+ "in": APIKeyIn.header,
+ "name": "X-Service-API-Key",
+ })
+ auth_credential = AuthCredential(
+ auth_type=AuthCredentialTypes.API_KEY, api_key="test_service_key"
+ )
+
+ tool = MCPTool(
+ mcp_tool=self.mock_mcp_tool,
+ mcp_session_manager=self.mock_session_manager,
+ auth_scheme=auth_scheme,
+ auth_credential=auth_credential,
+ )
+
+ # Mock the session response
+ expected_response = {"result": "authenticated_success"}
+ self.mock_session.call_tool = AsyncMock(return_value=expected_response)
+
+ tool_context = Mock(spec=ToolContext)
+ args = {"param1": "test_value"}
+
+ result = await tool._run_async_impl(
+ args=args, tool_context=tool_context, credential=auth_credential
+ )
+
+ assert result == expected_response
+ # Check that headers were passed correctly with custom API key header
+ self.mock_session_manager.create_session.assert_called_once()
+ call_args = self.mock_session_manager.create_session.call_args
+ headers = call_args[1]["headers"]
+ assert headers == {"X-Service-API-Key": "test_service_key"}
+
  @pytest.mark.asyncio
  async def test_run_async_impl_retry_decorator(self):
  """Test that the retry decorator is applied correctly."""
@@ -350,6 +510,45 @@ async def test_get_headers_http_custom_scheme(self):
 
  assert headers == {"Authorization": "custom custom_token"}
 
+ @pytest.mark.asyncio
+ async def test_get_headers_api_key_error_logging(self):
+ """Test that API key errors are logged correctly."""
+ from fastapi.openapi.models import APIKey
+ from fastapi.openapi.models import APIKeyIn
+ from google.adk.auth.auth_schemes import AuthSchemeType
+
+ # Create auth scheme for query-based API key (not supported)
+ auth_scheme = APIKey(**{
+ "type": AuthSchemeType.apiKey,
+ "in": APIKeyIn.query,
+ "name": "api_key",
+ })
+ auth_credential = AuthCredential(
+ auth_type=AuthCredentialTypes.API_KEY, api_key="my_api_key"
+ )
+
+ tool = MCPTool(
+ mcp_tool=self.mock_mcp_tool,
+ mcp_session_manager=self.mock_session_manager,
+ auth_scheme=auth_scheme,
+ auth_credential=auth_credential,
+ )
+
+ tool_context = Mock(spec=ToolContext)
+
+ # Test with logging
+ with patch("google.adk.tools.mcp_tool.mcp_tool.logger") as mock_logger:
+ with pytest.raises(ValueError):
+ await tool._get_headers(tool_context, auth_credential)
+
+ # Verify error was logged
+ mock_logger.error.assert_called_once()
+ logged_message = mock_logger.error.call_args[0][0]
+ assert (
+ "MCPTool only supports header-based API key authentication"
+ in logged_message
+ )
+
  def test_init_validation(self):
  """Test that initialization validates required parameters."""
  # This test ensures that the MCPTool properly handles its dependencies