feat: Develop the internal ToolboxClient wrapper

anubhav756 · anubhav756 · commit 269dc5c41775 · 2025-12-10T17:34:05.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,150 @@
+# 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 Optional, Dict, Callable, Any, Union
+import toolbox_core
+from google.auth import transport
+from google.auth import compute_engine
+from google.oauth2 import id_token
+import google.auth
+
+from .credentials import CredentialConfig, CredentialType
+
+
+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, str]] = None,
+        **kwargs
+    ):
+        """
+        Args:
+            server_url: The URL of the Toolbox server.
+            credentials: The CredentialConfig object (from CredentialStrategy).
+            additional_headers: A dictionary of static headers to include in every request.
+            **kwargs: Additional arguments passed to toolbox_core.ToolboxClient.
+        """
+        self._server_url = server_url
+        self._credentials = credentials
+        self._additional_headers = additional_headers or {}
+        
+        # Prepare auth_token_getters for toolbox-core
+        # toolbox_core expects: dict[str, Callable[[], str | Awaitable[str]]]
+        # However, for general headers (like Authorization), we can pass them in client_headers
+        # if they are static or simpler. Toolbox-core supports `client_headers` which can be dynamic.
+        
+        self._core_client_headers: Dict[str, Union[str, Callable[[], 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(
+            server_url=server_url,
+            client_headers=self._core_client_headers,
+            **kwargs
+        )
+
+    def _configure_auth(self, creds: CredentialConfig):
+        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
+            # We wrap it to match the signature expected by toolbox-core headers
+            # (which accepts callables)
+            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.
+            # The client itself doesn't set a global header because the token is per-user 
+            # and passed via the tool's execution context (in the future) or handled by the tool wrapper.
+            # The ToolboxTool wrapper will need to inject the token per-request or we rely on 
+            # toolbox-core's per-request auth support if widely available, but ADK flow involves 
+            # getting the token in `run_async`. 
+            # For now, we leave client-level headers empty for this strategy.
+            pass
+
+    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:
+            # Note: This is a synchronous call. Toolbox-core supports sync callables in headers.
+            # Ideally we would use async but google-auth is primarily sync for these helpers.
+            request = transport.requests.Request()
+            # Try to get ID token directly (e.g. on Cloud Run)
+            try:
+                token = id_token.fetch_id_token(request, audience)
+                return f"Bearer {token}"
+            except Exception:
+                # Fallback to default credentials (e.g. local gcloud)
+                creds, _ = google.auth.default()
+                if not creds.valid:
+                    creds.refresh(request)
+                # If specific ID token credentials, use them, otherwise this might be Access Token (scoped)
+                # For Toolbox we usually need ID Tokens.
+                # If the user is locally auth'd via `gcloud auth login`, fetch_id_token is preferred.
+                # If falling back to service account file:
+                if hasattr(creds, 'id_token') and creds.id_token:
+                     return f"Bearer {creds.id_token}"
+                
+                # If we are here, we might need to manually sign via IAM or similar if it's a generic SA.
+                # For simplicity in this v1, we assume fetch_id_token works or standard creds work.
+                # Re-attempt fetch_id_token on the credentials object if possible
+                curr_token = creds.token
+                return f"Bearer {curr_token}"
+
+        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
+
+    async def load_toolset(self, toolset_name: str, **kwargs):
+        return await self._client.load_toolset(toolset_name, **kwargs)
+
+    async def load_tool(self, tool_name: str, **kwargs):
+        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,165 @@
+# 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 pytest
+from unittest.mock import MagicMock, patch, ANY, AsyncMock
+from toolbox_adk.client import ToolboxClient
+from toolbox_adk.credentials import CredentialStrategy
+
+class TestToolboxClient:
+    
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    def test_init_no_auth(self, mock_core_client):
+        creds = CredentialStrategy.TOOLBOX_IDENTITY()
+        client = ToolboxClient("http://server", credentials=creds)
+        
+        mock_core_client.assert_called_with(
+            server_url="http://server",
+            client_headers={}
+        )
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    def test_init_manual_token(self, mock_core_client):
+        creds = CredentialStrategy.MANUAL_TOKEN("abc")
+        client = ToolboxClient("http://server", credentials=creds)
+        
+        mock_core_client.assert_called_with(
+            server_url="http://server",
+            client_headers={"Authorization": "Bearer abc"}
+        )
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    def test_init_additional_headers(self, mock_core_client):
+        creds = CredentialStrategy.TOOLBOX_IDENTITY()
+        headers = {"X-Custom": "Val"}
+        client = ToolboxClient(
+            "http://server", 
+            credentials=creds, 
+            additional_headers=headers
+        )
+        
+        mock_core_client.assert_called_with(
+            server_url="http://server",
+            client_headers={"X-Custom": "Val"}
+        )
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    @patch("toolbox_adk.client.id_token.fetch_id_token")
+    def test_adc_auth_flow_success(self, mock_fetch_token, mock_core_client):
+        mock_fetch_token.return_value = "id_token_123"
+        
+        creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS("http://aud")
+        client = ToolboxClient("http://server", credentials=creds)
+        
+        # Verify a callable was passed
+        args, kwargs = mock_core_client.call_args
+        assert "Authorization" in kwargs["client_headers"]
+        token_getter = kwargs["client_headers"]["Authorization"]
+        assert callable(token_getter)
+        
+        # Verify callable behavior
+        token = token_getter()
+        assert token == "Bearer id_token_123"
+        mock_fetch_token.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")
+    def test_adc_auth_flow_fallback(self, mock_default, mock_fetch_token, mock_core_client):
+        # unexpected error on fetch_id_token
+        mock_fetch_token.side_effect = Exception("No metadata")
+        
+        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("http://aud")
+        client = ToolboxClient("http://server", 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()
+    
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    def test_manual_creds(self, mock_core_client):
+        mock_g_creds = MagicMock()
+        mock_g_creds.valid = False
+        mock_g_creds.token = "oauth_token"
+        
+        creds = CredentialStrategy.MANUAL_CREDS(mock_g_creds)
+        client = ToolboxClient("http://server", credentials=creds)
+        
+        token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
+        token = token_getter()
+        
+        assert token == "Bearer oauth_token"
+        assert token == "Bearer oauth_token"
+        mock_g_creds.refresh.assert_called()
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    def test_init_validation_errors(self, mock_core_client):
+        # ADC missing audience
+        with pytest.raises(ValueError, match="target_audience is required"):
+            # Fix: only pass target_audience as keyword arg OR positional, not both mixed in a way that causes overlap if defined so
+            # Actually simpler: just pass raw None 
+            ToolboxClient("url", credentials=CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(None))
+            
+        # Manual token missing token
+        with pytest.raises(ValueError, match="token is required"):
+            ToolboxClient("url", credentials=CredentialStrategy.MANUAL_TOKEN(None))
+            
+        # Manual creds missing credentials
+        with pytest.raises(ValueError, match="credentials object is required"):
+            ToolboxClient("url", credentials=CredentialStrategy.MANUAL_CREDS(None))
+
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    @patch("toolbox_adk.client.id_token.fetch_id_token")
+    @patch("toolbox_adk.client.google.auth.default")
+    def test_adc_auth_flow_fallback_access_token(self, mock_default, mock_fetch_token, mock_core_client):
+        # fetch_id_token fails
+        mock_fetch_token.side_effect = Exception("No metadata")
+        
+        mock_creds = MagicMock()
+        mock_creds.valid = False
+        mock_creds.id_token = None # No ID token
+        mock_creds.token = "access_token_123"
+        mock_default.return_value = (mock_creds, "proj")
+        
+        creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS("http://aud")
+        client = ToolboxClient("http://server", credentials=creds)
+        
+        token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
+        token = token_getter()
+        
+        assert token == "Bearer access_token_123"
+        mock_creds.refresh.assert_called()
+
+    @pytest.mark.asyncio
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    async def test_delegation(self, mock_core_client):
+        mock_instance = mock_core_client.return_value
+        mock_instance.load_toolset = AsyncMock(return_value=["t1"])
+        mock_instance.close = AsyncMock()
+        
+        client = ToolboxClient("http://server")
+        tools = await client.load_toolset("my-set", extra="arg")
+        
+        mock_instance.load_toolset.assert_awaited_with("my-set", extra="arg")
+        assert tools == ["t1"]
+        
+        await client.close()
+        mock_instance.close.assert_awaited()
