feat: add authenticated parameters support

Author: kurtisvg

packages/toolbox-core/src/toolbox_core/client.py

@@ -11,13 +11,13 @@
 # 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
+import types
+from typing import Any, Callable, Optional
 
 from aiohttp import ClientSession
 
 from .protocol import ManifestSchema, ToolSchema
-from .tool import ToolboxTool
+from .tool import ToolboxTool, filter_required_authn_params
 
 
 class ToolboxClient:
@@ -53,14 +53,34 @@ def __init__(
             session = ClientSession()
         self.__session = session
 
-    def __parse_tool(self, name: str, schema: ToolSchema) -> ToolboxTool:
+    def __parse_tool(
+        self,
+        name: str,
+        schema: ToolSchema,
+        auth_token_getters: dict[str, Callable[[], str]],
+    ) -> ToolboxTool:
         """Internal helper to create a callable tool from its schema."""
+        # sort into authenticated and reg params
+        params = []
+        authn_params: dict[str, list[str]] = {}
+        auth_sources: set[str] = set()
+        for p in schema.parameters:
+            if not p.authSources:
+                params.append(p)
+            else:
+                authn_params[p.name] = p.authSources
+                auth_sources.update(p.authSources)
+
+        authn_params = filter_required_authn_params(authn_params, auth_sources)
+
         tool = ToolboxTool(
             session=self.__session,
             base_url=self.__base_url,
             name=name,
             desc=schema.description,
-            params=[p.to_param() for p in schema.parameters],
+            params=[p.to_param() for p in params],
+            required_authn_params=types.MappingProxyType(authn_params),
+            auth_service_token_getters=auth_token_getters,
         )
         return tool
 
@@ -99,6 +119,7 @@ async def close(self):
     async def load_tool(
         self,
         name: str,
+        auth_service_tokens: dict[str, Callable[[], str]] = {},
     ) -> ToolboxTool:
         """
         Asynchronously loads a tool from the server.
@@ -127,13 +148,14 @@ async def load_tool(
         if name not in manifest.tools:
             # TODO: Better exception
             raise Exception(f"Tool '{name}' not found!")
-        tool = self.__parse_tool(name, manifest.tools[name])
+        tool = self.__parse_tool(name, manifest.tools[name], auth_service_tokens)
 
         return tool
 
     async def load_toolset(
         self,
         name: str,
+        auth_token_getters: dict[str, Callable[[], str]] = {},
     ) -> list[ToolboxTool]:
         """
         Asynchronously fetches a toolset and loads all tools defined within it.
@@ -152,5 +174,8 @@ async def load_toolset(
         manifest: ManifestSchema = ManifestSchema(**json)
 
         # parse each tools name and schema into a list of ToolboxTools
-        tools = [self.__parse_tool(n, s) for n, s in manifest.tools.items()]
+        tools = [
+            self.__parse_tool(n, s, auth_token_getters)
+            for n, s in manifest.tools.items()
+        ]
         return tools

packages/toolbox-core/src/toolbox_core/tool.py

@@ -13,10 +13,13 @@
 # limitations under the License.
 
 
+import types
+from collections import defaultdict
 from inspect import Parameter, Signature
-from typing import Any
+from typing import Any, Callable, DefaultDict, Iterable, Mapping, Optional, Sequence
 
 from aiohttp import ClientSession
+from pytest import Session
 
 
 class ToolboxTool:
@@ -32,20 +35,19 @@ class ToolboxTool:
     and `inspect` work as expected.
     """
 
-    __url: str
-    __session: ClientSession
-    __signature__: Signature
-
     def __init__(
         self,
         session: ClientSession,
         base_url: str,
         name: str,
         desc: str,
-        params: list[Parameter],
+        params: Sequence[Parameter],
+        required_authn_params: Mapping[str, list[str]],
+        auth_service_token_getters: Mapping[str, Callable[[], str]],
     ):
         """
-        Initializes a callable that will trigger the tool invocation through the Toolbox server.
+        Initializes a callable that will trigger the tool invocation through the
+        Toolbox server.
 
         Args:
             session: The `aiohttp.ClientSession` used for making API requests.
@@ -54,19 +56,69 @@ def __init__(
             desc: The description of the remote tool (used as its docstring).
             params: A list of `inspect.Parameter` objects defining the tool's
                 arguments and their types/defaults.
+            required_authn_params: A dict of required authenticated parameters that
+                need a auth_service_token_getter set for them yet.
+            auth_service_tokens: A dict of authService -> token (or callables that
+                produce a token)
         """
 
         # used to invoke the toolbox API
-        self.__session = session
+        self.__session: ClientSession = session
+        self.__base_url: str = base_url
         self.__url = f"{base_url}/api/tool/{name}/invoke"
 
-        # the following properties are set to help anyone that might inspect it determine
+        self.__desc = desc
+        self.__params = params
+
+        # the following properties are set to help anyone that might inspect it determine usage
         self.__name__ = name
         self.__doc__ = desc
         self.__signature__ = Signature(parameters=params, return_annotation=str)
         self.__annotations__ = {p.name: p.annotation for p in params}
         # TODO: self.__qualname__ ??
 
+        # map of parameter name to auth service required by it
+        self.__required_authn_params = required_authn_params
+        # map of authService -> token_getter
+        self.__auth_service_token_getters = auth_service_token_getters
+
+    def __copy(
+        self,
+        session: Optional[ClientSession] = None,
+        base_url: Optional[str] = None,
+        name: Optional[str] = None,
+        desc: Optional[str] = None,
+        params: Optional[list[Parameter]] = None,
+        required_authn_params: Optional[Mapping[str, list[str]]] = None,
+        auth_service_token_getters: Optional[Mapping[str, Callable[[], str]]] = None,
+    ):
+        """
+        Creates a copy of the ToolboxTool, overriding specific fields.
+
+        Args:
+            session: The `aiohttp.ClientSession` used for making API requests.
+            base_url: The base URL of the Toolbox server API.
+            name: The name of the remote tool.
+            desc: The description of the remote tool (used as its docstring).
+            params: A list of `inspect.Parameter` objects defining the tool's
+                arguments and their types/defaults.
+            required_authn_params: A dict of required authenticated parameters that need
+                a auth_service_token_getter set for them yet.
+            auth_service_token_getters: A dict of authService -> token (or callables
+                that produce a token)
+
+        """
+        return ToolboxTool(
+            session=session or self.__session,
+            base_url=base_url or self.__base_url,
+            name=name or self.__name__,
+            desc=desc or self.__desc,
+            params=params or self.__params,
+            required_authn_params=required_authn_params or self.__required_authn_params,
+            auth_service_token_getters=auth_service_token_getters
+            or self.__auth_service_token_getters,
+        )
+
     async def __call__(self, *args: Any, **kwargs: Any) -> str:
         """
         Asynchronously calls the remote tool with the provided arguments.
@@ -81,16 +133,96 @@ async def __call__(self, *args: Any, **kwargs: Any) -> str:
         Returns:
             The string result returned by the remote tool execution.
         """
+
+        # check if any auth services need to be specified yet
+        if len(self.__required_authn_params) > 0:
+            req_auth_services = set(l for l in self.__required_authn_params.keys())
+            raise Exception(
+                f"One of more of the following authn services are required to invoke this tool: {','.join(req_auth_services)}"
+            )
+
+        # validate inputs to this call using the signature
         all_args = self.__signature__.bind(*args, **kwargs)
         all_args.apply_defaults()  # Include default values if not provided
         payload = all_args.arguments
 
+        # create headers for auth services
+        headers = {}
+        for auth_service, token_getter in self.__auth_service_token_getters.items():
+            headers[f"{auth_service}_token"] = token_getter()
+
         async with self.__session.post(
             self.__url,
             json=payload,
+            headers=headers,
         ) as resp:
-            ret = await resp.json()
-            if "error" in ret:
-                # TODO: better error
-                raise Exception(ret["error"])
-        return ret.get("result", ret)
+            body = await resp.json()
+            if resp.status < 200 or resp.status >= 300:
+                err = body.get("error", f"unexpected status from server: {resp.status}")
+                raise Exception(err)
+        return body.get("result", body)
+
+    def add_auth_token_getters(
+        self,
+        auth_token_getters: Mapping[str, Callable[[], str]],
+    ) -> "ToolboxTool":
+        """
+        Registers a auth token getter function that is used for AuthServices when tools
+        are invoked.
+
+        Args:
+            auth_token_getters: A mapping of authentication service names to
+                callables that return the corresponding authentication token.
+
+        Returns:
+            A new ToolboxTool instance with the specified authentication token
+            getters registered.
+        """
+
+        # throw an error if the authentication source is already registered
+        dupes = auth_token_getters.keys() & self.__auth_service_token_getters.keys()
+        if dupes:
+            raise ValueError(
+                f"Authentication source(s) `{', '.join(dupes)}` already registered in tool `{self.__name__}`."
+            )
+
+        # create a read-only updated value for new_getters
+        new_getters = types.MappingProxyType(
+            dict(self.__auth_service_token_getters, **auth_token_getters)
+        )
+        # create a read-only updated for params that are still required
+        new_req_authn_params = types.MappingProxyType(
+            filter_required_authn_params(
+                self.__required_authn_params, auth_token_getters.keys()
+            )
+        )
+
+        return self.__copy(
+            auth_service_token_getters=new_getters,
+            required_authn_params=new_req_authn_params,
+        )
+
+
+def filter_required_authn_params(
+    req_authn_params: Mapping[str, list[str]], auth_services: Iterable[str]
+) -> dict[str, list[str]]:
+    """
+    Utility function for reducing 'req_authn_params' to a subset of parameters that aren't supplied by a least one service in auth_services.
+
+    Args:
+        req_authn_params: A mapping of parameter names to sets of required
+            authentication services.
+        auth_services: An iterable of authentication service names for which
+            token getters are available.
+
+    Returns:
+        A new dictionary representing the subset of required authentication
+        parameters that are not covered by the provided `auth_services`.
+    """
+    req_params = {}
+    for param, services in req_authn_params.items():
+        # if we don't have a token_getter for any of the services required by the param, the param is still required
+        required = not any(s in services for s in auth_services)
+        if required:
+            req_params[param] = services
+    return req_params