From 798f65df86b1bbe33d864e30c5b1f9e155e89810 Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Sat, 31 Jan 2026 16:25:13 -0800 Subject: [PATCH] feat(tools): Implement toolset auth for McpToolset, OpenAPIToolset, and others Update existing toolsets to utilize the new toolset authentication framework. Key changes: - McpToolset: Add _auth_config instance variable, _get_auth_headers() method to build auth headers from exchanged credentials, and get_auth_config() override. Auth headers are now included when creating MCP sessions. - OpenAPIToolset: Add _auth_config and get_auth_config() to expose auth configuration to the framework. - ApplicationIntegrationToolset: Add _auth_config and get_auth_config(). - APIHubToolset: Add _auth_config and get_auth_config(). When ADK resolves toolset auth before calling get_tools(), it populates exchanged_auth_credential on the auth_config. Toolsets can then use this credential when making authenticated requests. Also update test fixtures in test_apihub_toolset.py to use real auth objects instead of mocks that fail pydantic validation. Co-authored-by: Xiang (Sean) Zhou PiperOrigin-RevId: 863764941 --- .../adk/tools/apihub_tool/apihub_toolset.py | 26 +- .../application_integration_toolset.py | 26 +- src/google/adk/tools/mcp_tool/mcp_toolset.py | 112 +++++++- .../openapi_spec_parser/openapi_toolset.py | 26 +- .../tools/apihub_tool/test_apihub_toolset.py | 23 +- .../tools/mcp_tool/test_mcp_toolset_auth.py | 269 ++++++++++++++++++ 6 files changed, 443 insertions(+), 39 deletions(-) create mode 100644 tests/unittests/tools/mcp_tool/test_mcp_toolset_auth.py diff --git a/src/google/adk/tools/apihub_tool/apihub_toolset.py b/src/google/adk/tools/apihub_tool/apihub_toolset.py index 57514b11..5560c7ce 100644 --- a/src/google/adk/tools/apihub_tool/apihub_toolset.py +++ b/src/google/adk/tools/apihub_tool/apihub_toolset.py @@ -145,6 +145,16 @@ class APIHubToolset(BaseToolset): self._openapi_toolset = None self._auth_scheme = auth_scheme self._auth_credential = auth_credential + # Store auth config as instance variable so ADK can populate + # exchanged_auth_credential in-place before calling get_tools() + self._auth_config: Optional[AuthConfig] = ( + AuthConfig( + auth_scheme=auth_scheme, + raw_auth_credential=auth_credential, + ) + if auth_scheme + else None + ) if not self._lazy_load_spec: self._prepare_toolset() @@ -191,11 +201,11 @@ class APIHubToolset(BaseToolset): await self._openapi_toolset.close() @override - def get_auth_config(self) -> AuthConfig | None: - """Returns the auth config for this toolset.""" - if self._auth_scheme is None: - return None - return AuthConfig( - auth_scheme=self._auth_scheme, - raw_auth_credential=self._auth_credential, - ) + def get_auth_config(self) -> Optional[AuthConfig]: + """Returns the auth config for this toolset. + + ADK will populate exchanged_auth_credential on this config before calling + get_tools(). The toolset can then access the ready-to-use credential via + self._auth_config.exchanged_auth_credential. + """ + return self._auth_config diff --git a/src/google/adk/tools/application_integration_tool/application_integration_toolset.py b/src/google/adk/tools/application_integration_tool/application_integration_toolset.py index 6dc72d38..5e068fba 100644 --- a/src/google/adk/tools/application_integration_tool/application_integration_toolset.py +++ b/src/google/adk/tools/application_integration_tool/application_integration_toolset.py @@ -143,6 +143,16 @@ class ApplicationIntegrationToolset(BaseToolset): self._service_account_json = service_account_json self._auth_scheme = auth_scheme self._auth_credential = auth_credential + # Store auth config as instance variable so ADK can populate + # exchanged_auth_credential in-place before calling get_tools() + self._auth_config: Optional[AuthConfig] = ( + AuthConfig( + auth_scheme=auth_scheme, + raw_auth_credential=auth_credential, + ) + if auth_scheme + else None + ) integration_client = IntegrationClient( project, @@ -281,11 +291,11 @@ class ApplicationIntegrationToolset(BaseToolset): await self._openapi_toolset.close() @override - def get_auth_config(self) -> AuthConfig | None: - """Returns the auth config for this toolset.""" - if self._auth_scheme is None: - return None - return AuthConfig( - auth_scheme=self._auth_scheme, - raw_auth_credential=self._auth_credential, - ) + def get_auth_config(self) -> Optional[AuthConfig]: + """Returns the auth config for this toolset. + + ADK will populate exchanged_auth_credential on this config before calling + get_tools(). The toolset can then access the ready-to-use credential via + self._auth_config.exchanged_auth_credential. + """ + return self._auth_config diff --git a/src/google/adk/tools/mcp_tool/mcp_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index 30228813..598452c7 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -149,6 +149,82 @@ class McpToolset(BaseToolset): self._auth_scheme = auth_scheme self._auth_credential = auth_credential self._require_confirmation = require_confirmation + # Store auth config as instance variable so ADK can populate + # exchanged_auth_credential in-place before calling get_tools() + self._auth_config: Optional[AuthConfig] = ( + AuthConfig( + auth_scheme=auth_scheme, + raw_auth_credential=auth_credential, + ) + if auth_scheme + else None + ) + + def _get_auth_headers(self) -> Optional[Dict[str, str]]: + """Build authentication headers from exchanged credential. + + Returns: + Dictionary of auth headers, or None if no auth configured. + """ + if not self._auth_config or not self._auth_config.exchanged_auth_credential: + return None + + credential = self._auth_config.exchanged_auth_credential + headers: Optional[Dict[str, str]] = None + + 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 + 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 + and credential.http.credentials.username + and credential.http.credentials.password + ): + credentials_str = ( + f"{credential.http.credentials.username}" + f":{credential.http.credentials.password}" + ) + encoded_credentials = base64.b64encode( + credentials_str.encode() + ).decode() + headers = {"Authorization": f"Basic {encoded_credentials}"} + elif credential.http.credentials and credential.http.credentials.token: + # Handle other HTTP schemes with token + headers = { + "Authorization": ( + f"{credential.http.scheme} {credential.http.credentials.token}" + ) + } + elif credential.api_key: + # For API key, use the auth scheme to determine header name + if self._auth_config.auth_scheme: + from fastapi.openapi.models import APIKeyIn + + if hasattr(self._auth_config.auth_scheme, "in_"): + if self._auth_config.auth_scheme.in_ == APIKeyIn.header: + headers = {self._auth_config.auth_scheme.name: credential.api_key} + else: + logger.warning( + "McpToolset only supports header-based API key authentication." + " Configured location: %s", + self._auth_config.auth_scheme.in_, + ) + else: + # Default to using scheme name as header + headers = {self._auth_config.auth_scheme.name: credential.api_key} + + return headers async def _execute_with_session( self, @@ -157,12 +233,22 @@ class McpToolset(BaseToolset): readonly_context: Optional[ReadonlyContext] = None, ) -> T: """Creates a session and executes a coroutine with it.""" - headers = ( - self._header_provider(readonly_context) - if self._header_provider and readonly_context - else None + headers: Dict[str, str] = {} + + # Add headers from header_provider if available + if self._header_provider and readonly_context: + provider_headers = self._header_provider(readonly_context) + if provider_headers: + headers.update(provider_headers) + + # Add auth headers from exchanged credential if available + auth_headers = self._get_auth_headers() + if auth_headers: + headers.update(auth_headers) + + session = await self._mcp_session_manager.create_session( + headers=headers if headers else None ) - session = await self._mcp_session_manager.create_session(headers=headers) timeout_in_seconds = ( self._connection_params.timeout if hasattr(self._connection_params, "timeout") @@ -274,14 +360,14 @@ class McpToolset(BaseToolset): print(f"Warning: Error during McpToolset cleanup: {e}", file=self._errlog) @override - def get_auth_config(self) -> AuthConfig | None: - """Returns the auth config for this toolset.""" - if self._auth_scheme is None: - return None - return AuthConfig( - auth_scheme=self._auth_scheme, - raw_auth_credential=self._auth_credential, - ) + def get_auth_config(self) -> Optional[AuthConfig]: + """Returns the auth config for this toolset. + + ADK will populate exchanged_auth_credential on this config before calling + get_tools(). The toolset can then access the ready-to-use credential via + self._auth_config.exchanged_auth_credential. + """ + return self._auth_config @override @classmethod diff --git a/src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py b/src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py index 5f490c2b..15317ff3 100644 --- a/src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py +++ b/src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py @@ -131,6 +131,16 @@ class OpenAPIToolset(BaseToolset): self._header_provider = header_provider self._auth_scheme = auth_scheme self._auth_credential = auth_credential + # Store auth config as instance variable so ADK can populate + # exchanged_auth_credential in-place before calling get_tools() + self._auth_config: Optional[AuthConfig] = ( + AuthConfig( + auth_scheme=auth_scheme, + raw_auth_credential=auth_credential, + ) + if auth_scheme + else None + ) if not spec_dict: spec_dict = self._load_spec(spec_str, spec_str_type) self._ssl_verify = ssl_verify @@ -216,11 +226,11 @@ class OpenAPIToolset(BaseToolset): pass @override - def get_auth_config(self) -> AuthConfig | None: - """Returns the auth config for this toolset.""" - if self._auth_scheme is None: - return None - return AuthConfig( - auth_scheme=self._auth_scheme, - raw_auth_credential=self._auth_credential, - ) + def get_auth_config(self) -> Optional[AuthConfig]: + """Returns the auth config for this toolset. + + ADK will populate exchanged_auth_credential on this config before calling + get_tools(). The toolset can then access the ready-to-use credential via + self._auth_config.exchanged_auth_credential. + """ + return self._auth_config diff --git a/tests/unittests/tools/apihub_tool/test_apihub_toolset.py b/tests/unittests/tools/apihub_tool/test_apihub_toolset.py index 064055fb..a8399801 100644 --- a/tests/unittests/tools/apihub_tool/test_apihub_toolset.py +++ b/tests/unittests/tools/apihub_tool/test_apihub_toolset.py @@ -14,7 +14,12 @@ from unittest.mock import MagicMock +from fastapi.openapi.models import OAuth2 +from fastapi.openapi.models import OAuthFlowAuthorizationCode +from fastapi.openapi.models import OAuthFlows from google.adk.auth.auth_credential import AuthCredential +from google.adk.auth.auth_credential import AuthCredentialTypes +from google.adk.auth.auth_credential import OAuth2Auth from google.adk.auth.auth_schemes import AuthScheme from google.adk.tools.apihub_tool.apihub_toolset import APIHubToolset from google.adk.tools.apihub_tool.clients.apihub_client import BaseAPIHubClient @@ -67,13 +72,27 @@ def lazy_apihub_toolset(): # Fixture for auth scheme @pytest.fixture def mock_auth_scheme(): - return MagicMock(spec=AuthScheme) + return OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl='https://example.com/auth', + tokenUrl='https://example.com/token', + scopes={'read': 'Read access'}, + ) + ) + ) # Fixture for auth credential @pytest.fixture def mock_auth_credential(): - return MagicMock(spec=AuthCredential) + return AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id='test_client_id', + client_secret='test_client_secret', + ), + ) # Test cases diff --git a/tests/unittests/tools/mcp_tool/test_mcp_toolset_auth.py b/tests/unittests/tools/mcp_tool/test_mcp_toolset_auth.py new file mode 100644 index 00000000..895dbd45 --- /dev/null +++ b/tests/unittests/tools/mcp_tool/test_mcp_toolset_auth.py @@ -0,0 +1,269 @@ +# 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. + +"""Tests for MCPToolset authentication functionality.""" + +import base64 +from unittest.mock import Mock + +from fastapi.openapi.models import APIKey as APIKeyScheme +from fastapi.openapi.models import APIKeyIn +from fastapi.openapi.models import OAuth2 +from fastapi.openapi.models import OAuthFlowAuthorizationCode +from fastapi.openapi.models import OAuthFlows +from google.adk.auth.auth_credential import AuthCredential +from google.adk.auth.auth_credential import AuthCredentialTypes +from google.adk.auth.auth_credential import HttpAuth +from google.adk.auth.auth_credential import HttpCredentials +from google.adk.auth.auth_credential import OAuth2Auth +from google.adk.auth.auth_tool import AuthConfig +from google.adk.tools.mcp_tool.mcp_toolset import McpToolset +from mcp import StdioServerParameters +import pytest + + +class TestMcpToolsetGetAuthConfig: + """Tests for McpToolset.get_auth_config method.""" + + def test_get_auth_config_returns_none_without_auth_scheme(self): + """Test that get_auth_config returns None when no auth configured.""" + toolset = McpToolset( + connection_params=StdioServerParameters(command="echo", args=["test"]) + ) + + assert toolset.get_auth_config() is None + + def test_get_auth_config_returns_config_with_auth_scheme(self): + """Test that get_auth_config returns AuthConfig when auth configured.""" + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/auth", + tokenUrl="https://example.com/token", + scopes={"read": "Read access"}, + ) + ) + ) + auth_credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="test_client_id", + client_secret="test_client_secret", + ), + ) + + toolset = McpToolset( + connection_params=StdioServerParameters(command="echo", args=["test"]), + auth_scheme=auth_scheme, + auth_credential=auth_credential, + ) + + 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 + + def test_get_auth_config_returns_same_instance(self): + """Test that get_auth_config returns the same instance each time.""" + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/auth", + tokenUrl="https://example.com/token", + scopes={}, + ) + ) + ) + + toolset = McpToolset( + connection_params=StdioServerParameters(command="echo", args=["test"]), + auth_scheme=auth_scheme, + ) + + # Should return the same instance + config1 = toolset.get_auth_config() + config2 = toolset.get_auth_config() + assert config1 is config2 + + +class TestMcpToolsetGetAuthHeaders: + """Tests for McpToolset._get_auth_headers method.""" + + @pytest.fixture + def toolset_with_oauth2(self): + """Create a toolset with OAuth2 auth configured.""" + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/auth", + tokenUrl="https://example.com/token", + scopes={"read": "Read access"}, + ) + ) + ) + auth_credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="test_client_id", + client_secret="test_client_secret", + ), + ) + + return McpToolset( + connection_params=StdioServerParameters(command="echo", args=["test"]), + auth_scheme=auth_scheme, + auth_credential=auth_credential, + ) + + def test_get_auth_headers_returns_none_without_auth_config(self): + """Test that _get_auth_headers returns None without auth config.""" + toolset = McpToolset( + connection_params=StdioServerParameters(command="echo", args=["test"]) + ) + + assert toolset._get_auth_headers() is None + + def test_get_auth_headers_returns_none_without_exchanged_credential( + self, toolset_with_oauth2 + ): + """Test that _get_auth_headers returns None without exchanged credential.""" + # No exchanged credential set yet + assert toolset_with_oauth2._get_auth_headers() is None + + def test_get_auth_headers_oauth2_bearer_token(self, toolset_with_oauth2): + """Test that _get_auth_headers returns Bearer token for OAuth2.""" + # Set exchanged credential with access token + toolset_with_oauth2._auth_config.exchanged_auth_credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth(access_token="test-access-token"), + ) + + headers = toolset_with_oauth2._get_auth_headers() + + assert headers is not None + assert headers["Authorization"] == "Bearer test-access-token" + + def test_get_auth_headers_http_bearer_token(self): + """Test that _get_auth_headers returns Bearer token for HTTP bearer.""" + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/auth", + tokenUrl="https://example.com/token", + scopes={}, + ) + ) + ) + + toolset = McpToolset( + connection_params=StdioServerParameters(command="echo", args=["test"]), + auth_scheme=auth_scheme, + ) + + # Set exchanged credential with HTTP bearer token + toolset._auth_config.exchanged_auth_credential = AuthCredential( + auth_type=AuthCredentialTypes.HTTP, + http=HttpAuth( + scheme="bearer", + credentials=HttpCredentials(token="test-bearer-token"), + ), + ) + + headers = toolset._get_auth_headers() + + assert headers is not None + assert headers["Authorization"] == "Bearer test-bearer-token" + + def test_get_auth_headers_http_basic_auth(self): + """Test that _get_auth_headers returns Basic auth for HTTP basic.""" + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/auth", + tokenUrl="https://example.com/token", + scopes={}, + ) + ) + ) + + toolset = McpToolset( + connection_params=StdioServerParameters(command="echo", args=["test"]), + auth_scheme=auth_scheme, + ) + + # Set exchanged credential with HTTP basic auth + toolset._auth_config.exchanged_auth_credential = AuthCredential( + auth_type=AuthCredentialTypes.HTTP, + http=HttpAuth( + scheme="basic", + credentials=HttpCredentials( + username="testuser", + password="testpass", + ), + ), + ) + + headers = toolset._get_auth_headers() + + assert headers is not None + expected_credentials = base64.b64encode(b"testuser:testpass").decode() + assert headers["Authorization"] == f"Basic {expected_credentials}" + + def test_get_auth_headers_api_key_header(self): + """Test that _get_auth_headers returns API key in header.""" + # Note: fastapi's APIKey model uses 'in' not 'in_', but accepts both + auth_scheme = APIKeyScheme(**{ + "in": APIKeyIn.header, + "name": "X-API-Key", + }) + + toolset = McpToolset( + connection_params=StdioServerParameters(command="echo", args=["test"]), + auth_scheme=auth_scheme, + ) + + # Set exchanged credential with API key + toolset._auth_config.exchanged_auth_credential = AuthCredential( + auth_type=AuthCredentialTypes.API_KEY, + api_key="test-api-key-12345", + ) + + headers = toolset._get_auth_headers() + + assert headers is not None + assert headers["X-API-Key"] == "test-api-key-12345" + + def test_get_auth_headers_api_key_non_header_logs_warning(self, caplog): + """Test that non-header API key logs a warning.""" + # Note: fastapi's APIKey model uses 'in' not 'in_' + auth_scheme = APIKeyScheme(**{ + "in": APIKeyIn.query, # Query param, not header + "name": "api_key", + }) + + toolset = McpToolset( + connection_params=StdioServerParameters(command="echo", args=["test"]), + auth_scheme=auth_scheme, + ) + + # Set exchanged credential with API key + toolset._auth_config.exchanged_auth_credential = AuthCredential( + auth_type=AuthCredentialTypes.API_KEY, + api_key="test-api-key", + ) + + headers = toolset._get_auth_headers() + + # Should return None for non-header API key + assert headers is None