feat: Add support for auth related features

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

@@ -14,7 +14,7 @@
 import asyncio
 import os
 import uuid
-from typing import Any, Mapping, Optional, Union
+from typing import Any, Mapping, Optional
 
 from aiohttp import ClientSession
 
@@ -53,6 +53,10 @@ def base_url(self) -> str:
         return self.__mcp_base_url
 
     def __convert_tool_schema(self, tool_data: dict) -> ToolSchema:
+        meta = tool_data.get("_meta", {})
+        param_auth = meta.get("toolbox/authParams", {})
+        invoke_auth = meta.get("toolbox/authInvoke", [])
+
         parameters = []
         input_schema = tool_data.get("inputSchema", {})
         properties = input_schema.get("properties", {})
@@ -66,17 +70,25 @@ def __convert_tool_schema(self, tool_data: dict) -> ToolSchema:
                 )
             else:
                 additional_props = True
+
+            auth_sources = param_auth.get(name)
+
             parameters.append(
                 ParameterSchema(
                     name=name,
                     type=schema["type"],
                     description=schema.get("description", ""),
                     required=name in required,
+                    authSources=auth_sources,
                     additionalProperties=additional_props,
                 )
             )
 
-        return ToolSchema(description=tool_data["description"], parameters=parameters)
+        return ToolSchema(
+            description=tool_data["description"],
+            parameters=parameters,
+            authRequired=invoke_auth,
+        )
 
     async def __list_tools(
         self,

packages/toolbox-core/tests/test_e2e_mcp.py

@@ -133,6 +133,99 @@ async def test_bind_params_callable(
         assert "row4" not in response
 
 
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("toolbox_server")
+class TestAuth:
+    async def test_run_tool_unauth_with_auth(
+        self, toolbox: ToolboxClient, auth_token2: str
+    ):
+        """Tests running a tool that doesn't require auth, with auth provided."""
+
+        with pytest.raises(
+            ValueError,
+            match=rf"Validation failed for tool 'get-row-by-id': unused auth tokens: my-test-auth",
+        ):
+            await toolbox.load_tool(
+                "get-row-by-id",
+                auth_token_getters={"my-test-auth": lambda: auth_token2},
+            )
+
+    async def test_run_tool_no_auth(self, toolbox: ToolboxClient):
+        """Tests running a tool requiring auth without providing auth."""
+        tool = await toolbox.load_tool("get-row-by-id-auth")
+        with pytest.raises(
+            PermissionError,
+            match="One or more of the following authn services are required to invoke this tool: my-test-auth",
+        ):
+            await tool(id="2")
+
+    async def test_run_tool_wrong_auth(self, toolbox: ToolboxClient, auth_token2: str):
+        """Tests running a tool with incorrect auth. The tool
+        requires a different authentication than the one provided."""
+        tool = await toolbox.load_tool("get-row-by-id-auth")
+        auth_tool = tool.add_auth_token_getters({"my-test-auth": lambda: auth_token2})
+        with pytest.raises(
+            Exception,
+            match="Unauthorized",
+        ):
+            await auth_tool(id="2")
+
+    async def test_run_tool_auth(self, toolbox: ToolboxClient, auth_token1: str):
+        """Tests running a tool with correct auth."""
+        tool = await toolbox.load_tool("get-row-by-id-auth")
+        auth_tool = tool.add_auth_token_getters({"my-test-auth": lambda: auth_token1})
+        response = await auth_tool(id="2")
+        assert "row2" in response
+
+    @pytest.mark.asyncio
+    async def test_run_tool_async_auth(self, toolbox: ToolboxClient, auth_token1: str):
+        """Tests running a tool with correct auth using an async token getter."""
+        tool = await toolbox.load_tool("get-row-by-id-auth")
+
+        async def get_token_asynchronously():
+            return auth_token1
+
+        auth_tool = tool.add_auth_token_getters(
+            {"my-test-auth": get_token_asynchronously}
+        )
+        response = await auth_tool(id="2")
+        assert "row2" in response
+
+    async def test_run_tool_param_auth_no_auth(self, toolbox: ToolboxClient):
+        """Tests running a tool with a param requiring auth, without auth."""
+        tool = await toolbox.load_tool("get-row-by-email-auth")
+        with pytest.raises(
+            PermissionError,
+            match="One or more of the following authn services are required to invoke this tool: my-test-auth",
+        ):
+            await tool()
+
+    async def test_run_tool_param_auth(self, toolbox: ToolboxClient, auth_token1: str):
+        """Tests running a tool with a param requiring auth, with correct auth."""
+        tool = await toolbox.load_tool(
+            "get-row-by-email-auth",
+            auth_token_getters={"my-test-auth": lambda: auth_token1},
+        )
+        response = await tool()
+        assert "row4" in response
+        assert "row5" in response
+        assert "row6" in response
+
+    async def test_run_tool_param_auth_no_field(
+        self, toolbox: ToolboxClient, auth_token1: str
+    ):
+        """Tests running a tool with a param requiring auth, with insufficient auth."""
+        tool = await toolbox.load_tool(
+            "get-row-by-content-auth",
+            auth_token_getters={"my-test-auth": lambda: auth_token1},
+        )
+        with pytest.raises(
+            Exception,
+            match="no field named row_data in claims",
+        ):
+            await tool()
+
+
 @pytest.mark.asyncio
 @pytest.mark.usefixtures("toolbox_server")
 class TestOptionalParams: