feat(adk): Implement client wrapper

anubhav756 · anubhav756 · commit 9f43f3ae158f · 2025-12-11T05:02:54.000+05:30
diff --git a/packages/toolbox-adk/src/toolbox_adk/client.py b/packages/toolbox-adk/src/toolbox_adk/client.py
@@ -0,0 +1,154 @@
+# 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 contextvars import ContextVar
+from typing import Any, Awaitable, Callable, Dict, Optional, Union
+
+import google.auth
+import toolbox_core
+from google.auth import compute_engine, transport
+from google.oauth2 import id_token
+
+from .credentials import CredentialConfig, CredentialType
+
+USER_TOKEN_CONTEXT_VAR: ContextVar[Optional[str]] = ContextVar(
+    "toolbox_user_token", default=None
+)
+
+
+class ToolboxClient:
+    """
+    Wraps toolbox_core.ToolboxClient to provide ADK-native authentication strategy support.
+    """
+
+    def __init__(
+        self,
+        server_url: str,
+        credentials: Optional[CredentialConfig] = None,
+        additional_headers: Optional[
+            Dict[str, Union[str, Callable[[], str], Callable[[], Awaitable[str]]]]
+        ] = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            server_url: The URL of the Toolbox server.
+            credentials: The CredentialConfig object (from CredentialStrategy).
+            additional_headers: Dictionary of headers (static or dynamic callables).
+            **kwargs: Additional arguments passed to toolbox_core.ToolboxClient.
+        """
+        self._server_url = server_url
+        self._credentials = credentials
+        self._additional_headers = additional_headers or {}
+
+        self._core_client_headers: Dict[
+            str, Union[str, Callable[[], str], Callable[[], Awaitable[str]]]
+        ] = {}
+
+        # Add static additional headers
+        for k, v in self._additional_headers.items():
+            self._core_client_headers[k] = v
+
+        if credentials:
+            self._configure_auth(credentials)
+
+        self._client = toolbox_core.ToolboxClient(
+            url=server_url, client_headers=self._core_client_headers, **kwargs
+        )
+
+    def _configure_auth(self, creds: CredentialConfig) -> None:
+        if creds.type == CredentialType.TOOLBOX_IDENTITY:
+            # No auth headers needed
+            pass
+
+        elif creds.type == CredentialType.APPLICATION_DEFAULT_CREDENTIALS:
+            if not creds.target_audience:
+                raise ValueError(
+                    "target_audience is required for APPLICATION_DEFAULT_CREDENTIALS"
+                )
+
+            # Create an async capable token getter
+            self._core_client_headers["Authorization"] = self._create_adc_token_getter(
+                creds.target_audience
+            )
+
+        elif creds.type == CredentialType.MANUAL_TOKEN:
+            if not creds.token:
+                raise ValueError("token is required for MANUAL_TOKEN")
+            scheme = creds.scheme or "Bearer"
+            self._core_client_headers["Authorization"] = f"{scheme} {creds.token}"
+
+        elif creds.type == CredentialType.MANUAL_CREDS:
+            if not creds.credentials:
+                raise ValueError("credentials object is required for MANUAL_CREDS")
+
+            # Adapter for google-auth credentials object to callable
+            self._core_client_headers["Authorization"] = (
+                self._create_creds_token_getter(creds.credentials)
+            )
+
+        elif creds.type == CredentialType.USER_IDENTITY:
+            # For USER_IDENTITY (3LO), the *Tool* handles the interactive flow at runtime.
+
+            def get_user_token() -> str:
+                token = USER_TOKEN_CONTEXT_VAR.get()
+                if not token:
+                    return ""
+                return f"Bearer {token}"
+
+            self._core_client_headers["Authorization"] = get_user_token
+
+    def _create_adc_token_getter(self, audience: str) -> Callable[[], str]:
+        """Returns a callable that fetches a fresh ID token using ADC."""
+
+        def get_token() -> str:
+            request = transport.requests.Request()
+            try:
+                token = id_token.fetch_id_token(request, audience)
+                return f"Bearer {token}"
+            except Exception:
+                # Fallback to default credentials
+                creds, _ = google.auth.default()
+                if not creds.valid:
+                    creds.refresh(request)
+
+                if hasattr(creds, "id_token") and creds.id_token:
+                    return f"Bearer {creds.id_token}"
+
+                curr_token = getattr(creds, "token", None)
+                return f"Bearer {curr_token}" if curr_token else ""
+
+        return get_token
+
+    def _create_creds_token_getter(self, credentials: Any) -> Callable[[], str]:
+        def get_token() -> str:
+            request = transport.requests.Request()
+            if not credentials.valid:
+                credentials.refresh(request)
+            return f"Bearer {credentials.token}"
+
+        return get_token
+
+    @property
+    def credential_config(self) -> Optional[CredentialConfig]:
+        return self._credentials
+
+    async def load_toolset(self, toolset_name: str, **kwargs: Any) -> Any:
+        return await self._client.load_toolset(toolset_name, **kwargs)
+
+    async def load_tool(self, tool_name: str, **kwargs: Any) -> Any:
+        return await self._client.load_tool(tool_name, **kwargs)
+
+    async def close(self):
+        await self._client.close()
diff --git a/packages/toolbox-adk/tests/unit/test_client.py b/packages/toolbox-adk/tests/unit/test_client.py
@@ -0,0 +1,191 @@
+# 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.
+
+import unittest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from toolbox_adk import CredentialStrategy, ToolboxClient
+from toolbox_adk.client import CredentialType
+
+
+@pytest.mark.asyncio
+class TestToolboxClientAuth:
+    """Unit tests for Client Auth logic."""
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    async def test_init_toolbox_identity(self, mock_core_client):
+        """Test init with TOOLBOX_IDENTITY (no auth headers)."""
+        creds = CredentialStrategy.TOOLBOX_IDENTITY()
+        client = ToolboxClient(server_url="http://test", credentials=creds)
+
+        # Verify core client created with empty headers for auth
+        _, kwargs = mock_core_client.call_args
+        assert "client_headers" in kwargs
+        headers = kwargs["client_headers"]
+        assert "Authorization" not in headers
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    @patch("toolbox_adk.client.id_token.fetch_id_token")
+    @patch("toolbox_adk.client.google.auth.default")
+    @patch("toolbox_adk.client.transport.requests.Request")
+    async def test_init_adc_success_fetch_id_token(
+        self, mock_req, mock_default, mock_fetch_id, mock_core_client
+    ):
+        """Test ADC strategy where fetch_id_token succeeds."""
+        mock_fetch_id.return_value = "id-token-123"
+
+        creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(
+            target_audience="aud"
+        )
+        client = ToolboxClient(server_url="http://test", credentials=creds)
+
+        _, kwargs = mock_core_client.call_args
+        headers = kwargs["client_headers"]
+        assert "Authorization" in headers
+        token_getter = headers["Authorization"]
+        assert callable(token_getter)
+
+        # Call the getter
+        token_val = token_getter()
+        assert token_val == "Bearer id-token-123"
+        mock_fetch_id.assert_called()
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    @patch("toolbox_adk.client.id_token.fetch_id_token")
+    @patch("toolbox_adk.client.google.auth.default")
+    @patch("toolbox_adk.client.transport.requests.Request")
+    async def test_init_adc_fallback_creds(
+        self, mock_req, mock_default, mock_fetch_id, mock_core_client
+    ):
+        """Test ADC strategy fallback to default() when fetch_id_token fails."""
+        mock_fetch_id.side_effect = Exception("No metadata server")
+
+        # Mock default creds
+        mock_creds = MagicMock()
+        mock_creds.valid = False
+        mock_creds.id_token = "fallback-id-token"
+        mock_default.return_value = (mock_creds, "proj")
+
+        creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(
+            target_audience="aud"
+        )
+        client = ToolboxClient(server_url="http://test", credentials=creds)
+
+        token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
+        token = token_getter()
+        assert token == "Bearer fallback-id-token"
+        mock_creds.refresh.assert_called()  # Because we set valid=False
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    @patch("toolbox_adk.client.id_token.fetch_id_token")
+    @patch("toolbox_adk.client.google.auth.default")
+    @patch("toolbox_adk.client.transport.requests.Request")
+    async def test_init_adc_fallback_creds_token(
+        self, mock_req, mock_default, mock_fetch_id, mock_core_client
+    ):
+        """Test ADC fallback when creds have .token but no .id_token."""
+        mock_fetch_id.side_effect = Exception("No metadata server")
+
+        mock_creds = MagicMock()
+        mock_creds.valid = True
+        del mock_creds.id_token  # Simulate no id_token attr or None
+        mock_creds.token = "access-token-123"  # e.g. user creds
+        mock_default.return_value = (mock_creds, "proj")
+
+        creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(
+            target_audience="aud"
+        )
+        client = ToolboxClient(server_url="http://test", credentials=creds)
+        token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
+        assert token_getter() == "Bearer access-token-123"
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    async def test_init_manual_token(self, mock_core_client):
+        creds = CredentialStrategy.MANUAL_TOKEN(token="abc")
+        client = ToolboxClient("http://test", credentials=creds)
+        headers = mock_core_client.call_args[1]["client_headers"]
+        assert headers["Authorization"] == "Bearer abc"
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    async def test_init_manual_creds(self, mock_core_client):
+        mock_google_creds = MagicMock()
+        mock_google_creds.valid = True
+        mock_google_creds.token = "creds-token"
+
+        creds = CredentialStrategy.MANUAL_CREDS(credentials=mock_google_creds)
+        client = ToolboxClient("http://test", credentials=creds)
+
+        token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
+        assert token_getter() == "Bearer creds-token"
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    async def test_init_user_identity(self, mock_core_client):
+        creds = CredentialStrategy.USER_IDENTITY(client_id="c", client_secret="s")
+        client = ToolboxClient("http://test", credentials=creds)
+
+        token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
+        # Should be empty initially
+        assert token_getter() == ""
+
+        # Set context
+        from toolbox_adk.client import USER_TOKEN_CONTEXT_VAR
+
+        token = USER_TOKEN_CONTEXT_VAR.set("user-tok")
+        try:
+            assert token_getter() == "Bearer user-tok"
+        finally:
+            USER_TOKEN_CONTEXT_VAR.reset(token)
+
+    async def test_validation_errors(self):
+        with pytest.raises(ValueError):
+            # ADC requires audience
+            creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(
+                target_audience=""
+            )
+            ToolboxClient("http://test", credentials=creds)
+
+        with pytest.raises(ValueError):
+            creds = CredentialStrategy.MANUAL_TOKEN(token="")
+            ToolboxClient("http://test", credentials=creds)
+
+        with pytest.raises(ValueError):
+            creds = CredentialStrategy.MANUAL_CREDS(credentials=None)
+            ToolboxClient("http://test", credentials=creds)
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    async def test_load_methods(self, mock_core_client_class):
+        # Setup mock instance
+        mock_instance = AsyncMock()
+        mock_core_client_class.return_value = mock_instance
+
+        client = ToolboxClient(
+            "http://test", credentials=CredentialStrategy.TOOLBOX_IDENTITY()
+        )
+
+        # Test load_toolset
+        await client.load_toolset("ts", foo="bar")
+        mock_instance.load_toolset.assert_called_with("ts", foo="bar")
+
+        # Test load_tool
+        await client.load_tool("t", baz="qux")
+        mock_instance.load_tool.assert_called_with("t", baz="qux")
+
+        # Test close
+        await client.close()
+        mock_instance.close.assert_called_once()
+
+        # Test property
+        assert client.credential_config is not None
