feat: Add Authenticated Tool (Experimental)

PiperOrigin-RevId: 772992074

src/google/adk/tools/authenticated_function_tool.py

@@ -0,0 +1,107 @@
+# Copyright 2025 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.
+
+from __future__ import annotations
+
+import inspect
+import logging
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import Optional
+from typing import Union
+
+from typing_extensions import override
+
+from ..auth.auth_credential import AuthCredential
+from ..auth.auth_tool import AuthConfig
+from ..auth.credential_manager import CredentialManager
+from ..utils.feature_decorator import experimental
+from .function_tool import FunctionTool
+from .tool_context import ToolContext
+
+logger = logging.getLogger("google_adk." + __name__)
+
+
+@experimental
+class AuthenticatedFunctionTool(FunctionTool):
+  """A FunctionTool that handles authentication before the actual tool logic
+  gets called. Functions can accept a special `credential` argument which is the
+  credential ready for use.(Experimental)
+  """
+
+  def __init__(
+      self,
+      *,
+      func: Callable[..., Any],
+      auth_config: AuthConfig = None,
+      response_for_auth_required: Optional[Union[dict[str, Any], str]] = None,
+  ):
+    """Initializes the AuthenticatedFunctionTool.
+
+    Args:
+        func: The function to be called.
+        auth_config: The authentication configuration.
+        response_for_auth_required: The response to return when the tool is
+          requesting auth credential from the client. There could be two case,
+          the tool doesn't configure any credentials
+          (auth_config.raw_auth_credential is missing) or the credentials
+          configured is not enough to authenticate the tool (e.g. an OAuth
+          client id and client secrect is configured.) and needs client input
+          (e.g. client need to involve the end user in an oauth flow and get
+          back the oauth response.)
+    """
+    super().__init__(func=func)
+    self._ignore_params.append("credential")
+
+    if auth_config and auth_config.auth_scheme:
+      self._credentials_manager = CredentialManager(auth_config=auth_config)
+    else:
+      logger.warning(
+          "auth_config or auth_config.auth_scheme is missing. Will skip"
+          " authentication.Using FunctionTool instead if authentication is not"
+          " required."
+      )
+      self._credentials_manager = None
+    self._response_for_auth_required = response_for_auth_required
+
+  @override
+  async def run_async(
+      self, *, args: dict[str, Any], tool_context: ToolContext
+  ) -> Any:
+    credential = None
+    if self._credentials_manager:
+      credential = await self._credentials_manager.get_auth_credential(
+          tool_context
+      )
+      if not credential:
+        await self._credentials_manager.request_credential(tool_context)
+        return self._response_for_auth_required or "Pending User Authorization."
+
+    return await self._run_async_impl(
+        args=args, tool_context=tool_context, credential=credential
+    )
+
+  async def _run_async_impl(
+      self,
+      *,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+      credential: AuthCredential,
+  ) -> Any:
+    args_to_call = args.copy()
+    signature = inspect.signature(self.func)
+    if "credential" in signature.parameters:
+      args_to_call["credential"] = credential
+    return await super().run_async(args=args_to_call, tool_context=tool_context)

src/google/adk/tools/base_authenticated_tool.py

@@ -0,0 +1,107 @@
+# Copyright 2025 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.
+
+from __future__ import annotations
+
+from abc import abstractmethod
+import logging
+from typing import Any
+from typing import Optional
+from typing import Union
+
+from typing_extensions import override
+
+from ..auth.auth_credential import AuthCredential
+from ..auth.auth_tool import AuthConfig
+from ..auth.credential_manager import CredentialManager
+from ..utils.feature_decorator import experimental
+from .base_tool import BaseTool
+from .tool_context import ToolContext
+
+logger = logging.getLogger("google_adk." + __name__)
+
+
+@experimental
+class BaseAuthenticatedTool(BaseTool):
+  """A base tool class that handles authentication before the actual tool logic
+  gets called. Functions can accept a special `credential` argument which is the
+  credential ready for use.(Experimental)
+  """
+
+  def __init__(
+      self,
+      *,
+      name,
+      description,
+      auth_config: AuthConfig = None,
+      response_for_auth_required: Optional[Union[dict[str, Any], str]] = None,
+  ):
+    """
+    Args:
+      name: The name of the tool.
+      description: The description of the tool.
+      auth_config: The auth configuration of the tool.
+      response_for_auth_required: The response to return when the tool is
+          requesting auth credential from the client. There could be two case,
+          the tool doesn't configure any credentials
+          (auth_config.raw_auth_credential is missing) or the credentials
+          configured is not enough to authenticate the tool (e.g. an OAuth
+          client id and client secrect is configured.) and needs client input
+          (e.g. client need to involve the end user in an oauth flow and get
+          back the oauth response.)
+    """
+    super().__init__(
+        name=name,
+        description=description,
+    )
+
+    if auth_config and auth_config.auth_scheme:
+      self._credentials_manager = CredentialManager(auth_config=auth_config)
+    else:
+      logger.warning(
+          "auth_config or auth_config.auth_scheme is missing. Will skip"
+          " authentication.Using FunctionTool instead if authentication is not"
+          " required."
+      )
+      self._credentials_manager = None
+    self._response_for_auth_required = response_for_auth_required
+
+  @override
+  async def run_async(
+      self, *, args: dict[str, Any], tool_context: ToolContext
+  ) -> Any:
+    credential = None
+    if self._credentials_manager:
+      credential = await self._credentials_manager.get_auth_credential(
+          tool_context
+      )
+      if not credential:
+        await self._credentials_manager.request_credential(tool_context)
+        return self._response_for_auth_required or "Pending User Authorization."
+
+    return await self._run_async_impl(
+        args=args,
+        tool_context=tool_context,
+        credential=credential,
+    )
+
+  @abstractmethod
+  async def _run_async_impl(
+      self,
+      *,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+      credential: AuthCredential,
+  ) -> Any:
+    pass

tests/unittests/tools/test_base_authenticated_tool.py

@@ -0,0 +1,343 @@
+# Copyright 2025 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.
+
+from unittest.mock import AsyncMock
+from unittest.mock import Mock
+
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_schemes import AuthScheme
+from google.adk.auth.auth_schemes import AuthSchemeType
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.tools.base_authenticated_tool import BaseAuthenticatedTool
+from google.adk.tools.tool_context import ToolContext
+import pytest
+
+
+class _TestAuthenticatedTool(BaseAuthenticatedTool):
+  """Test implementation of BaseAuthenticatedTool for testing purposes."""
+
+  def __init__(
+      self,
+      name="test_auth_tool",
+      description="Test authenticated tool",
+      auth_config=None,
+      unauthenticated_response=None,
+  ):
+    super().__init__(
+        name=name,
+        description=description,
+        auth_config=auth_config,
+        response_for_auth_required=unauthenticated_response,
+    )
+    self.run_impl_called = False
+    self.run_impl_result = "test_result"
+
+  async def _run_async_impl(self, *, args, tool_context, credential):
+    """Test implementation of the abstract method."""
+    self.run_impl_called = True
+    self.last_args = args
+    self.last_tool_context = tool_context
+    self.last_credential = credential
+    return self.run_impl_result
+
+
+def _create_mock_auth_config():
+  """Creates a mock AuthConfig with proper structure."""
+  auth_scheme = Mock(spec=AuthScheme)
+  auth_scheme.type_ = AuthSchemeType.oauth2
+
+  auth_config = Mock(spec=AuthConfig)
+  auth_config.auth_scheme = auth_scheme
+
+  return auth_config
+
+
+def _create_mock_auth_credential():
+  """Creates a mock AuthCredential."""
+  credential = Mock(spec=AuthCredential)
+  credential.auth_type = AuthCredentialTypes.OAUTH2
+  return credential
+
+
+class TestBaseAuthenticatedTool:
+  """Test suite for BaseAuthenticatedTool."""
+
+  def test_init_with_auth_config(self):
+    """Test initialization with auth_config."""
+    auth_config = _create_mock_auth_config()
+    unauthenticated_response = {"error": "Not authenticated"}
+
+    tool = _TestAuthenticatedTool(
+        name="test_tool",
+        description="Test description",
+        auth_config=auth_config,
+        unauthenticated_response=unauthenticated_response,
+    )
+
+    assert tool.name == "test_tool"
+    assert tool.description == "Test description"
+    assert tool._credentials_manager is not None
+    assert tool._response_for_auth_required == unauthenticated_response
+
+  def test_init_with_no_auth_config(self):
+    """Test initialization without auth_config."""
+    tool = _TestAuthenticatedTool()
+
+    assert tool.name == "test_auth_tool"
+    assert tool.description == "Test authenticated tool"
+    assert tool._credentials_manager is None
+    assert tool._response_for_auth_required is None
+
+  def test_init_with_empty_auth_scheme(self):
+    """Test initialization with auth_config but no auth_scheme."""
+    auth_config = Mock(spec=AuthConfig)
+    auth_config.auth_scheme = None
+
+    tool = _TestAuthenticatedTool(auth_config=auth_config)
+
+    assert tool._credentials_manager is None
+
+  def test_init_with_default_unauthenticated_response(self):
+    """Test initialization with default unauthenticated response."""
+    auth_config = _create_mock_auth_config()
+
+    tool = _TestAuthenticatedTool(auth_config=auth_config)
+
+    assert tool._response_for_auth_required is None
+
+  @pytest.mark.asyncio
+  async def test_run_async_no_credentials_manager(self):
+    """Test run_async when no credentials manager is configured."""
+    tool = _TestAuthenticatedTool()
+    tool_context = Mock(spec=ToolContext)
+    args = {"param1": "value1"}
+
+    result = await tool.run_async(args=args, tool_context=tool_context)
+
+    assert result == "test_result"
+    assert tool.run_impl_called
+    assert tool.last_args == args
+    assert tool.last_tool_context == tool_context
+    assert tool.last_credential is None
+
+  @pytest.mark.asyncio
+  async def test_run_async_with_valid_credential(self):
+    """Test run_async when valid credential is available."""
+    auth_config = _create_mock_auth_config()
+    credential = _create_mock_auth_credential()
+
+    # Mock the credentials manager
+    mock_credentials_manager = AsyncMock()
+    mock_credentials_manager.get_auth_credential = AsyncMock(
+        return_value=credential
+    )
+
+    tool = _TestAuthenticatedTool(auth_config=auth_config)
+    tool._credentials_manager = mock_credentials_manager
+
+    tool_context = Mock(spec=ToolContext)
+    args = {"param1": "value1"}
+
+    result = await tool.run_async(args=args, tool_context=tool_context)
+
+    assert result == "test_result"
+    assert tool.run_impl_called
+    assert tool.last_args == args
+    assert tool.last_tool_context == tool_context
+    assert tool.last_credential == credential
+    mock_credentials_manager.get_auth_credential.assert_called_once_with(
+        tool_context
+    )
+
+  @pytest.mark.asyncio
+  async def test_run_async_no_credential_available(self):
+    """Test run_async when no credential is available."""
+    auth_config = _create_mock_auth_config()
+
+    # Mock the credentials manager to return None
+    mock_credentials_manager = AsyncMock()
+    mock_credentials_manager.get_auth_credential = AsyncMock(return_value=None)
+    mock_credentials_manager.request_credential = AsyncMock()
+
+    tool = _TestAuthenticatedTool(auth_config=auth_config)
+    tool._credentials_manager = mock_credentials_manager
+
+    tool_context = Mock(spec=ToolContext)
+    args = {"param1": "value1"}
+
+    result = await tool.run_async(args=args, tool_context=tool_context)
+
+    assert result == "Pending User Authorization."
+    assert not tool.run_impl_called
+    mock_credentials_manager.get_auth_credential.assert_called_once_with(
+        tool_context
+    )
+    mock_credentials_manager.request_credential.assert_called_once_with(
+        tool_context
+    )
+
+  @pytest.mark.asyncio
+  async def test_run_async_no_credential_with_custom_response(self):
+    """Test run_async when no credential is available with custom response."""
+    auth_config = _create_mock_auth_config()
+    custom_response = {
+        "status": "authentication_required",
+        "message": "Please login",
+    }
+
+    # Mock the credentials manager to return None
+    mock_credentials_manager = AsyncMock()
+    mock_credentials_manager.get_auth_credential = AsyncMock(return_value=None)
+    mock_credentials_manager.request_credential = AsyncMock()
+
+    tool = _TestAuthenticatedTool(
+        auth_config=auth_config, unauthenticated_response=custom_response
+    )
+    tool._credentials_manager = mock_credentials_manager
+
+    tool_context = Mock(spec=ToolContext)
+    args = {"param1": "value1"}
+
+    result = await tool.run_async(args=args, tool_context=tool_context)
+
+    assert result == custom_response
+    assert not tool.run_impl_called
+    mock_credentials_manager.get_auth_credential.assert_called_once_with(
+        tool_context
+    )
+    mock_credentials_manager.request_credential.assert_called_once_with(
+        tool_context
+    )
+
+  @pytest.mark.asyncio
+  async def test_run_async_no_credential_with_string_response(self):
+    """Test run_async when no credential is available with string response."""
+    auth_config = _create_mock_auth_config()
+    custom_response = "Custom authentication required message"
+
+    # Mock the credentials manager to return None
+    mock_credentials_manager = AsyncMock()
+    mock_credentials_manager.get_auth_credential = AsyncMock(return_value=None)
+    mock_credentials_manager.request_credential = AsyncMock()
+
+    tool = _TestAuthenticatedTool(
+        auth_config=auth_config, unauthenticated_response=custom_response
+    )
+    tool._credentials_manager = mock_credentials_manager
+
+    tool_context = Mock(spec=ToolContext)
+    args = {"param1": "value1"}
+
+    result = await tool.run_async(args=args, tool_context=tool_context)
+
+    assert result == custom_response
+    assert not tool.run_impl_called
+
+  @pytest.mark.asyncio
+  async def test_run_async_propagates_impl_exception(self):
+    """Test that run_async propagates exceptions from _run_async_impl."""
+    auth_config = _create_mock_auth_config()
+    credential = _create_mock_auth_credential()
+
+    # Mock the credentials manager
+    mock_credentials_manager = AsyncMock()
+    mock_credentials_manager.get_auth_credential = AsyncMock(
+        return_value=credential
+    )
+
+    tool = _TestAuthenticatedTool(auth_config=auth_config)
+    tool._credentials_manager = mock_credentials_manager
+
+    # Make the implementation raise an exception
+    async def failing_impl(*, args, tool_context, credential):
+      raise ValueError("Implementation failed")
+
+    tool._run_async_impl = failing_impl
+
+    tool_context = Mock(spec=ToolContext)
+    args = {"param1": "value1"}
+
+    with pytest.raises(ValueError, match="Implementation failed"):
+      await tool.run_async(args=args, tool_context=tool_context)
+
+  @pytest.mark.asyncio
+  async def test_run_async_with_different_args_types(self):
+    """Test run_async with different argument types."""
+    tool = _TestAuthenticatedTool()
+    tool_context = Mock(spec=ToolContext)
+
+    # Test with empty args
+    result = await tool.run_async(args={}, tool_context=tool_context)
+    assert result == "test_result"
+    assert tool.last_args == {}
+
+    # Test with complex args
+    complex_args = {
+        "string_param": "test",
+        "number_param": 42,
+        "list_param": [1, 2, 3],
+        "dict_param": {"nested": "value"},
+    }
+    result = await tool.run_async(args=complex_args, tool_context=tool_context)
+    assert result == "test_result"
+    assert tool.last_args == complex_args
+
+  @pytest.mark.asyncio
+  async def test_run_async_credentials_manager_exception(self):
+    """Test run_async when credentials manager raises an exception."""
+    auth_config = _create_mock_auth_config()
+
+    # Mock the credentials manager to raise an exception
+    mock_credentials_manager = AsyncMock()
+    mock_credentials_manager.get_auth_credential = AsyncMock(
+        side_effect=RuntimeError("Credential service error")
+    )
+
+    tool = _TestAuthenticatedTool(auth_config=auth_config)
+    tool._credentials_manager = mock_credentials_manager
+
+    tool_context = Mock(spec=ToolContext)
+    args = {"param1": "value1"}
+
+    with pytest.raises(RuntimeError, match="Credential service error"):
+      await tool.run_async(args=args, tool_context=tool_context)
+
+  def test_abstract_nature(self):
+    """Test that BaseAuthenticatedTool cannot be instantiated directly."""
+    with pytest.raises(TypeError):
+      # This should fail because _run_async_impl is abstract
+      BaseAuthenticatedTool(name="test", description="test")
+
+  @pytest.mark.asyncio
+  async def test_run_async_return_values(self):
+    """Test run_async with different return value types."""
+    tool = _TestAuthenticatedTool()
+    tool_context = Mock(spec=ToolContext)
+    args = {}
+
+    # Test with None return
+    tool.run_impl_result = None
+    result = await tool.run_async(args=args, tool_context=tool_context)
+    assert result is None
+
+    # Test with dict return
+    tool.run_impl_result = {"key": "value"}
+    result = await tool.run_async(args=args, tool_context=tool_context)
+    assert result == {"key": "value"}
+
+    # Test with list return
+    tool.run_impl_result = [1, 2, 3]
+    result = await tool.run_async(args=args, tool_context=tool_context)
+    assert result == [1, 2, 3]