feat: Add 3LO auth support to ToolboxTool

Author: anubhav756

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

@@ -20,6 +20,11 @@
 from google.adk.tools.base_tool import BaseTool
 from google.adk.agents.readonly_context import ReadonlyContext
 
+from google_auth_oauthlib.flow import InstalledAppFlow
+from google.auth.transport.requests import Request
+from .credentials import CredentialConfig, CredentialType
+from .client import USER_TOKEN_CONTEXT_VAR
+
 
 class ToolboxContext:
     """Context object passed to pre/post hooks."""
@@ -40,35 +45,21 @@ def __init__(
         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.
-        # core_tool name, description etc. are available.
-        
-        # Note: BaseTool in adk-python typically requires tool_name, tool_description etc.
-        # We map them from the core tool.
-        # core_tool.tool_def is usually available or we use the wrapper's attributes if exposed.
-        # In toolbox-core 0.1.0+, tool instance is callable but metadata is also stored.
-        # Assuming core_tool has typical attributes or we introspect.
-        # For now, we rely on core_tool's string representation or assume it has name/doc.
 
         name = getattr(core_tool, "__name__", "unknown_tool")
         description = getattr(core_tool, "__doc__", "No description provided.") or "No description provided."
 
-        # toolbox_core tools might not expose explicit args schema in a way BaseTool expects 
-        # (dict of name -> type). ADK tools usually define `tool_args` as a dict.
-        # We might need to inspect the signature if not explicitly available.
-        # For this implementation, we allow dynamic arguments similar to FunctionTool.
-        # If BaseTool enforces providing `tool_args`, we might need to introspect.
-        # However, BaseTool logic in adk-python often derives it. 
-        # We'll call super with minimal known info or defaults.
-        
         super().__init__(
             name=name,
             description=description,
@@ -78,6 +69,8 @@ def __init__(
         self._core_tool = core_tool
         self._pre_hook = pre_hook
         self._post_hook = post_hook
+        self._auth_config = auth_config
+        self._user_creds: Optional[Any] = None
 
     @override
     async def run_async(
@@ -93,75 +86,52 @@ async def run_async(
             await self._pre_hook(ctx)
 
         # 2. ADK Auth Integration (3LO)
-        # We check if we need to inject a user token managed by ADK.
-        # This typically happens if the toolset was configured with USER_IDENTITY.
-        # We don't know the configuration here explicitly unless passed, 
-        # BUT the convention is: if we can get a token from tool_context, use it.
-        # Or better: The client (ToolboxClient) didn't set auth headers for USER_IDENTITY,
-        # so we must do it here. 
-        
-        # How do we know if we logic requires it? 
-        # We can try to fetch; if it exists, we assume we should use it?
-        # Or we rely on the fact that `core_tool` instance itself might have bound auth getters?
-        # Actually, toolbox-core tools support `add_auth_token_getters`.
-        # We need to bridge ADK's `get_auth_response` to toolbox-core's getter.
-        
-        # ADK 3LO flow:
-        # response = tool_context.get_auth_response()
-        # if not response:
-        #    tool_context.request_credential(...)
-        #    return "Please authenticate..."
-        # else:
-        #    token = response.token
-        
-        # If we have a way to detect 3LO requirement, we trigger it.
-        # For now, we'll try to retrieve an auth response if one is potentially available.
-        # But `toolbox-core` tool will fail if it needs auth and doesn't get it.
-        # The proper pattern is: wrap the core tool with a dynamically injected getter 
-        # that calls into ADK logic.
-        
-        # However, `toolbox-core`'s `auth_token_getters` expects a callable.
-        # Be careful: `tool_context` is available here in `run_async`, but `auth_token_getters`
-        # are usually bound locally.
-        
-        # We'll define a token getter that closes over `tool_context`.
-        # We'll define a token getter that closes over `tool_context` if needed in future.
-        # For now, unused.
+        # Check if USER_IDENTITY is configured
+        reset_token = None
+        if self._auth_config and self._auth_config.type == CredentialType.USER_IDENTITY:
+            # Handle interactive flow if credentials are missing or expired
+            if not self._user_creds or not self._user_creds.valid:
+                if self._user_creds and self._user_creds.expired and self._user_creds.refresh_token:
+                    try:
+                        self._user_creds.refresh(Request())
+                    except Exception:
+                        self._user_creds = None
+                
+                if not self._user_creds:
+                    # Trigger flow
+                    if not (self._auth_config.client_id and self._auth_config.client_secret):
+                        raise ValueError("USER_IDENTITY requires client_id and client_secret")
+                    
+                    config = {
+                        "installed": {
+                            "client_id": self._auth_config.client_id,
+                            "client_secret": self._auth_config.client_secret,
+                            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+                            "token_uri": "https://oauth2.googleapis.com/token",
+                        }
+                    }
+                    scopes = self._auth_config.scopes or ["https://www.googleapis.com/auth/cloud-platform"]
+                    
+                    flow = InstalledAppFlow.from_client_config(config, scopes=scopes)
+                    self._user_creds = flow.run_local_server(port=0)
+            
+            # Inject token into ContextVar
+            if self._user_creds:
+                reset_token = USER_TOKEN_CONTEXT_VAR.set(self._user_creds.token)
 
-        # Let's try to get a token if available, but fundamentally, ADK expects the tool to return 
-        # a request for credentials if missing.
-        auth_response = tool_context.get_auth_response()
-        
-        # NOTE: We can't easily tell if this specific tool *needs* user auth without metadata.
-        # Strategy: We assume that if the tool fails with an auth error, we might fallback?
-        # OR we rely on configuration. 
-        # Ideally, `ToolboxToolset` would have configured this tool with a special "trigger" if needed.
-        
-        # For now, we will proceed to execute. 
-        # If `USER_IDENTITY` was configured, we might have injected a special marker or getter.
-        # Revisit this logic if explicit "Force Auth" flag is needed on the tool.
-        
         try:
             # Execute the core tool
-            # mapping: toolbox-core expects (arg1=..., arg2=...)
-            # We unpack ctx.arguments
-            
-            # For 3LO support: We really need to know if we should be requesting credentials.
-            # For this MVP, we will rely on static auth or pre-existing headers.
-            # If we want to support 3LO properly, we need to handle the `tool_context.request_credential` flow.
-            # We'll check if `auth_token` argument is present or if we have a bound getter.
-            
-            # We execute:
             result = await self._core_tool(**ctx.arguments)
 
             ctx.result = result
             return result
 
         except Exception as e:
-            # TODO: Inspect e for Auth errors to trigger 3LO if configured?
             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)
 
@@ -172,5 +142,6 @@ def bind_params(self, bounded_params: Dict[str, Any]) -> 'ToolboxTool':
         return ToolboxTool(
             core_tool=new_core_tool, 
             pre_hook=self._pre_hook, 
-            post_hook=self._post_hook
+            post_hook=self._post_hook,
+            auth_config=self._auth_config
         )

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

@@ -98,7 +98,8 @@ async def get_tools(
             ToolboxTool(
                 core_tool=t, 
                 pre_hook=self._pre_hook, 
-                post_hook=self._post_hook
+                post_hook=self._post_hook,
+                auth_config=self._client.credential_config
             ) 
             for t in tools
         ]