feat: Add 3LO auth support and fix client url param

Author: anubhav756

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

@@ -18,9 +18,13 @@
 from google.auth import compute_engine
 from google.oauth2 import id_token
 import google.auth
+from contextvars import ContextVar
 
 from .credentials import CredentialConfig, CredentialType
 
+# Global ContextVar for User Identity (3LO) tokens to be injected per-request
+USER_TOKEN_CONTEXT_VAR: ContextVar[Optional[str]] = ContextVar("toolbox_user_token", default=None)
+
 
 class ToolboxClient:
     """
@@ -60,7 +64,7 @@ def __init__(
             self._configure_auth(credentials)
 
         self._client = toolbox_core.ToolboxClient(
-            server_url=server_url,
+            url=server_url,
             client_headers=self._core_client_headers,
             **kwargs
         )
@@ -94,13 +98,18 @@ def _configure_auth(self, creds: CredentialConfig):
 
         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
+            # We use a ContextVar to inject the token per-request.
+            
+            def get_user_token() -> str:
+                token = USER_TOKEN_CONTEXT_VAR.get()
+                if not token:
+                    # If this is called but no token is set in context, it means 
+                    # the tool wrapper failed to set it or we are in a context where 
+                    # we expected it. We return empty string which might cause 401.
+                    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."""
@@ -127,8 +136,8 @@ def get_token() -> str:
                 # 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}"
+                curr_token = getattr(creds, "token", None)
+                return f"Bearer {curr_token}" if curr_token else ""
 
         return get_token
 
@@ -140,6 +149,10 @@ def get_token() -> str:
             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):
         return await self._client.load_toolset(toolset_name, **kwargs)