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 <seanzhougoogle@google.com>
PiperOrigin-RevId: 863764941

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

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

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

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

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

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