feat(adk): Implement tool wrapper & auth getters

packages/toolbox-adk/src/toolbox_adk/__init__.py

@@ -14,16 +14,16 @@
 
 from .client import ToolboxClient
 from .credentials import CredentialConfig, CredentialStrategy, CredentialType
-# from .tool import ToolboxContext, ToolboxTool
-# from .toolset import ToolboxToolset
+from .tool import ToolboxContext, ToolboxTool
+from .toolset import ToolboxToolset
 from .version import __version__
 
 __all__ = [
     "CredentialStrategy",
     "CredentialConfig",
     "CredentialType",
     "ToolboxClient",
-    # "ToolboxTool",
-    # "ToolboxContext",
-    # "ToolboxToolset",
+    "ToolboxTool",
+    "ToolboxContext",
+    "ToolboxToolset",
 ]

packages/toolbox-adk/src/toolbox_adk/tool.py

@@ -0,0 +1,179 @@
+# 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 typing import Any, Awaitable, Callable, Dict, Optional, cast
+
+import toolbox_core
+from fastapi.openapi.models import (
+    OAuth2,
+    OAuthFlowAuthorizationCode,
+    OAuthFlows,
+    SecurityScheme,
+)
+from google.adk.agents.readonly_context import ReadonlyContext
+from google.adk.auth.auth_credential import (
+    AuthCredential,
+    AuthCredentialTypes,
+    OAuth2Auth,
+)
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.tools.base_tool import BaseTool
+from toolbox_core.tool import ToolboxTool as CoreToolboxTool
+from typing_extensions import override
+
+from .client import USER_TOKEN_CONTEXT_VAR
+from .credentials import CredentialConfig, CredentialType
+
+
+class ToolboxContext:
+    """Context object passed to pre/post hooks."""
+
+    def __init__(self, arguments: Dict[str, Any], tool_context: ReadonlyContext):
+        self.arguments = arguments
+        self.tool_context = tool_context
+        self.result: Optional[Any] = None
+        self.error: Optional[Exception] = None
+
+
+class ToolboxTool(BaseTool):
+    """
+    A tool that delegates to a remote Toolbox tool, integrated with ADK.
+    """
+
+    def __init__(
+        self,
+        core_tool: CoreToolboxTool,
+        pre_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
+        post_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
+        auth_config: Optional[CredentialConfig] = None,
+    ):
+        """
+        Args:
+            core_tool: The underlying toolbox_core.py tool instance.
+            pre_hook: Async function called before execution. Can modify ctx.arguments.
+            post_hook: Async function called after execution (finally block). Can inspect ctx.result/error.
+            auth_config: Credential configuration to handle interactive flows.
+        """
+        # We act as a proxy.
+        # We need to extract metadata from the core tool to satisfy BaseTool's contract.
+
+        name = getattr(core_tool, "__name__", "unknown_tool")
+        description = (
+            getattr(core_tool, "__doc__", "No description provided.")
+            or "No description provided."
+        )
+
+        super().__init__(
+            name=name,
+            description=description,
+            # Pass empty custom_metadata as it is not currently used
+            custom_metadata={},
+        )
+        self._core_tool = core_tool
+        self._pre_hook = pre_hook
+        self._post_hook = post_hook
+        self._auth_config = auth_config
+
+    @override
+    async def run_async(
+        self,
+        args: Dict[str, Any],
+        tool_context: ReadonlyContext,
+    ) -> Any:
+        # Create context
+        ctx = ToolboxContext(arguments=args, tool_context=tool_context)
+
+        # 1. Pre-hook
+        if self._pre_hook:
+            await self._pre_hook(ctx)
+
+        # 2. ADK Auth Integration (3LO)
+        # Check if USER_IDENTITY is configured
+        reset_token = None
+
+        if self._auth_config and self._auth_config.type == CredentialType.USER_IDENTITY:
+            if not self._auth_config.client_id or not self._auth_config.client_secret:
+                raise ValueError("USER_IDENTITY requires client_id and client_secret")
+
+            # Construct ADK AuthConfig
+            scopes = self._auth_config.scopes or [
+                "https://www.googleapis.com/auth/cloud-platform"
+            ]
+            scope_dict = {s: "" for s in scopes}
+
+            auth_config_adk = AuthConfig(
+                auth_scheme=OAuth2(
+                    flows=OAuthFlows(
+                        authorizationCode=OAuthFlowAuthorizationCode(
+                            authorizationUrl="https://accounts.google.com/o/oauth2/auth",
+                            tokenUrl="https://oauth2.googleapis.com/token",
+                            scopes=scope_dict,
+                        )
+                    )
+                ),
+                raw_auth_credential=AuthCredential(
+                    auth_type=AuthCredentialTypes.OAUTH2,
+                    oauth2=OAuth2Auth(
+                        client_id=self._auth_config.client_id,
+                        client_secret=self._auth_config.client_secret,
+                    ),
+                ),
+            )
+
+            # Check if we already have credentials from a previous exchange
+            try:
+                # get_auth_response returns AuthCredential if found
+                ctx_any = cast(Any, tool_context)
+                creds = ctx_any.get_auth_response(auth_config_adk)
+                if creds and creds.oauth2 and creds.oauth2.access_token:
+                    reset_token = USER_TOKEN_CONTEXT_VAR.set(creds.oauth2.access_token)
+                else:
+                    # Request credentials and pause execution
+                    ctx_any.request_credential(auth_config_adk)
+                    return None
+            except Exception as e:
+                ctx.error = e
+                if "credential" in str(e).lower() or isinstance(e, ValueError):
+                    raise e
+                # Fallback to request logic
+                ctx_any = cast(Any, tool_context)
+                ctx_any.request_credential(auth_config_adk)
+                return None
+
+        try:
+            # Execute the core tool
+            result = await self._core_tool(**ctx.arguments)
+
+            ctx.result = result
+            return result
+
+        except Exception as e:
+            ctx.error = e
+            raise e
+        finally:
+            if reset_token:
+                USER_TOKEN_CONTEXT_VAR.reset(reset_token)
+            if self._post_hook:
+                await self._post_hook(ctx)
+
+    def bind_params(self, bounded_params: Dict[str, Any]) -> "ToolboxTool":
+        """Allows runtime binding of parameters, delegating to core tool."""
+        new_core_tool = self._core_tool.bind_params(bounded_params)
+        # Return a new wrapper
+        return ToolboxTool(
+            core_tool=new_core_tool,
+            pre_hook=self._pre_hook,
+            post_hook=self._post_hook,
+            auth_config=self._auth_config,
+        )

packages/toolbox-adk/src/toolbox_adk/toolset.py

@@ -0,0 +1,123 @@
+# 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 typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Union
+
+from google.adk.agents.readonly_context import ReadonlyContext
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.base_toolset import BaseToolset
+from typing_extensions import override
+
+from .client import ToolboxClient
+from .credentials import CredentialConfig
+from .tool import ToolboxContext, ToolboxTool
+
+
+class ToolboxToolset(BaseToolset):
+    """
+    A Toolset that provides tools from a remote Toolbox server.
+    """
+
+    def __init__(
+        self,
+        server_url: str,
+        toolset_name: Optional[str] = None,
+        tool_names: Optional[List[str]] = None,
+        credentials: Optional[CredentialConfig] = None,
+        additional_headers: Optional[
+            Dict[str, Union[str, Callable[[], str], Callable[[], Awaitable[str]]]]
+        ] = None,
+        bound_params: Optional[Mapping[str, Union[Callable[[], Any], Any]]] = None,
+        auth_token_getters: Optional[
+            Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]]]]
+        ] = None,
+        pre_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
+        post_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            server_url: The URL of the Toolbox server.
+            toolset_name: The name of the remote toolset to load.
+            tool_names: Specific tool names to load (alternative to toolset_name).
+            credentials: Authentication configuration.
+            additional_headers: Extra headers (static or dynamic).
+            bound_params: Parameters to bind globally to loaded tools.
+            auth_token_getters: Mapping of auth service names to token getters.
+            pre_hook: Hook to run before every tool execution.
+            post_hook: Hook to run after every tool execution.
+        """
+        super().__init__()
+        self._client = ToolboxClient(
+            server_url=server_url,
+            credentials=credentials,
+            additional_headers=additional_headers,
+            **kwargs,
+        )
+        self._toolset_name = toolset_name
+        self._tool_names = tool_names
+        self._bound_params = bound_params
+        self._auth_token_getters = auth_token_getters
+        self._pre_hook = pre_hook
+        self._post_hook = post_hook
+
+    @override
+    async def get_tools(
+        self, readonly_context: Optional[ReadonlyContext] = None
+    ) -> List[BaseTool]:
+        """Loads tools from the toolbox server and wraps them."""
+        # Note: We don't close the client after get_tools because tools might need it.
+
+        tools = []
+        # 1. Load specific toolset if requested
+        if self._toolset_name:
+            core_tools = await self._client.load_toolset(
+                self._toolset_name,
+                bound_params=self._bound_params or {},
+                auth_token_getters=self._auth_token_getters or {},
+            )
+            tools.extend(core_tools)
+
+        # 2. Load specific tools if requested
+        if self._tool_names:
+            for name in self._tool_names:
+                core_tool = await self._client.load_tool(
+                    name,
+                    bound_params=self._bound_params or {},
+                    auth_token_getters=self._auth_token_getters or {},
+                )
+                tools.append(core_tool)
+
+        # 3. If NO tools/toolsets were specified, default to loading everything (default toolset)
+        if not self._toolset_name and not self._tool_names:
+            core_tools = await self._client.load_toolset(
+                None,
+                bound_params=self._bound_params or {},
+                auth_token_getters=self._auth_token_getters or {},
+            )
+            tools.extend(core_tools)
+
+        # Wrap all core tools in ToolboxTool
+        return [
+            ToolboxTool(
+                core_tool=t,
+                pre_hook=self._pre_hook,
+                post_hook=self._post_hook,
+                auth_config=self._client.credential_config,
+            )
+            for t in tools
+        ]
+
+    async def close(self):
+        await self._client.close()