From c41e1da0753d420ea4fc43d351e3e400639df35c Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Thu, 20 Aug 2026 11:16:56 +0530 Subject: [PATCH] feat(core): client tool and toolset loading with secure parameters --- .../toolbox-core/src/toolbox_core/client.py | 52 +++++- .../src/toolbox_core/sync_client.py | 20 ++- packages/toolbox-core/tests/test_client.py | 150 ++++++++++++++++++ .../toolbox-core/tests/test_sync_client.py | 41 ++++- 4 files changed, 251 insertions(+), 12 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 8f9d0cd75..3e2fad78c 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -292,7 +292,10 @@ def __parse_tool( client_headers: Mapping[ str, Union[Callable[[], str], Callable[[], Awaitable[str]], str] ], - ) -> tuple[ToolboxTool, set[str], set[str]]: + secure_params: Mapping[ + str, Union[Callable[[], Any], Callable[[], Awaitable[Any]], Any] + ] = {}, + ) -> tuple[ToolboxTool, set[str], set[str], set[str]]: """Internal helper to create a callable tool from its schema.""" # sort into reg, authn, and bound params params = [] @@ -308,6 +311,17 @@ def __parse_tool( else: # regular parameter params.append(p) + remaining_secure_params = [] + bound_sec_params: dict[ + str, Union[Callable[[], Any], Callable[[], Awaitable[Any]], Any] + ] = {} + for sp in schema.secure_parameters: + if sp.name in secure_params: + bound_sec_params[sp.name] = secure_params[sp.name] + else: + remaining_secure_params.append(sp) + used_secure_keys = set(bound_sec_params.keys()) + authn_params, authz_tokens, used_auth_keys = identify_auth_requirements( authn_params, schema.authRequired, @@ -325,11 +339,13 @@ def __parse_tool( auth_service_token_getters=MappingProxyType(auth_token_getters), bound_params=MappingProxyType(bound_params), client_headers=MappingProxyType(client_headers), + secure_params=tuple(remaining_secure_params), + bound_secure_params=MappingProxyType(bound_sec_params), ) used_bound_keys = set(bound_params.keys()) - return tool, used_auth_keys, used_bound_keys + return tool, used_auth_keys, used_bound_keys, used_secure_keys async def __aenter__(self): """ @@ -371,6 +387,9 @@ async def load_tool( bound_params: Mapping[ str, Union[Callable[[], Any], Callable[[], Awaitable[Any]], Any] ] = {}, + secure_params: Mapping[ + str, Union[Callable[[], Any], Callable[[], Awaitable[Any]], Any] + ] = {}, ) -> ToolboxTool: """ Asynchronously loads a tool from the server. @@ -385,6 +404,8 @@ async def load_tool( callables that return the corresponding authentication token. bound_params: A mapping of parameter names to bind to specific values or callables that are called to produce values as needed. + secure_params: A mapping of secure parameter names to bind to specific + values or callables. Returns: ToolboxTool: A callable object representing the loaded tool, ready @@ -393,7 +414,7 @@ async def load_tool( Raises: ValueError: If the loaded tool instance fails to utilize at least - one provided parameter or auth token (if any provided). + one provided parameter, auth token, or secure parameter (if any provided). """ # Resolve client headers resolved_headers = { @@ -409,16 +430,18 @@ async def load_tool( if name not in manifest.tools: # TODO: Better exception raise ValueError(f"Tool '{name}' not found!") - tool, used_auth_keys, used_bound_keys = self.__parse_tool( + tool, used_auth_keys, used_bound_keys, used_secure_keys = self.__parse_tool( name, manifest.tools[name], auth_token_getters, bound_params, self.__client_headers, + secure_params, ) provided_auth_keys = set(auth_token_getters.keys()) provided_bound_keys = set(bound_params.keys()) + provided_secure_keys = set(secure_params.keys()) validate_unused_requirements( provided_auth_keys, @@ -427,6 +450,8 @@ async def load_tool( used_bound_keys, name, is_toolset=False, + provided_secure_keys=provided_secure_keys, + used_secure_keys=used_secure_keys, ) return tool @@ -441,6 +466,9 @@ async def load_toolset( str, Union[Callable[[], Any], Callable[[], Awaitable[Any]], Any] ] = {}, strict: bool = False, + secure_params: Mapping[ + str, Union[Callable[[], Any], Callable[[], Awaitable[Any]], Any] + ] = {}, ) -> list[ToolboxTool]: """ Asynchronously fetches a toolset and loads all tools defined within it. @@ -452,10 +480,12 @@ async def load_toolset( bound_params: A mapping of parameter names to bind to specific values or callables that are called to produce values as needed. strict: If True, raises an error if *any* loaded tool instance fails - to utilize all of the given parameters or auth tokens. (if any + to utilize all of the given parameters, auth tokens, or secure parameters (if any provided). If False (default), raises an error only if a - user-provided parameter or auth token cannot be applied to *any* + user-provided parameter, auth token, or secure parameter cannot be applied to *any* loaded tool across the set. + secure_params: A mapping of secure parameter names to bind to specific values or + callables that are called to produce values as needed. Returns: list[ToolboxTool]: A list of callables, one for each tool defined @@ -479,17 +509,20 @@ async def load_toolset( tools: list[ToolboxTool] = [] overall_used_auth_keys: set[str] = set() overall_used_bound_params: set[str] = set() + overall_used_secure_params: set[str] = set() provided_auth_keys = set(auth_token_getters.keys()) provided_bound_keys = set(bound_params.keys()) + provided_secure_keys = set(secure_params.keys()) # parse each tool's name and schema into a list of ToolboxTools for tool_name, schema in manifest.tools.items(): - tool, used_auth_keys, used_bound_keys = self.__parse_tool( + tool, used_auth_keys, used_bound_keys, used_secure_keys = self.__parse_tool( tool_name, schema, auth_token_getters, bound_params, self.__client_headers, + secure_params, ) tools.append(tool) @@ -501,10 +534,13 @@ async def load_toolset( used_bound_keys, tool_name, is_toolset=False, + provided_secure_keys=provided_secure_keys, + used_secure_keys=used_secure_keys, ) else: overall_used_auth_keys.update(used_auth_keys) overall_used_bound_params.update(used_bound_keys) + overall_used_secure_params.update(used_secure_keys) validate_unused_requirements( provided_auth_keys, @@ -513,6 +549,8 @@ async def load_toolset( overall_used_bound_params, name or "default", is_toolset=True, + provided_secure_keys=provided_secure_keys, + used_secure_keys=overall_used_secure_params, ) return tools diff --git a/packages/toolbox-core/src/toolbox_core/sync_client.py b/packages/toolbox-core/src/toolbox_core/sync_client.py index e438fd503..cc96768df 100644 --- a/packages/toolbox-core/src/toolbox_core/sync_client.py +++ b/packages/toolbox-core/src/toolbox_core/sync_client.py @@ -103,6 +103,9 @@ def load_tool( bound_params: Mapping[ str, Union[Callable[[], Any], Callable[[], Awaitable[Any]], Any] ] = {}, + secure_params: Mapping[ + str, Union[Callable[[], Any], Callable[[], Awaitable[Any]], Any] + ] = {}, ) -> ToolboxSyncTool: """ Synchronously loads a tool from the server. @@ -117,13 +120,17 @@ def load_tool( callables that return the corresponding authentication token. bound_params: A mapping of parameter names to bind to specific values or callables that are called to produce values as needed. + secure_params: A mapping of secure parameter names to bind to specific values or + callables that are called to produce values as needed. Returns: ToolboxSyncTool: A callable object representing the loaded tool, ready for execution. The specific arguments and behavior of the callable depend on the tool itself. """ - coro = self.__async_client.load_tool(name, auth_token_getters, bound_params) + coro = self.__async_client.load_tool( + name, auth_token_getters, bound_params, secure_params + ) if not self.__loop or not self.__thread: raise ValueError("Background loop or thread cannot be None.") @@ -141,6 +148,9 @@ def load_toolset( str, Union[Callable[[], Any], Callable[[], Awaitable[Any]], Any] ] = {}, strict: bool = False, + secure_params: Mapping[ + str, Union[Callable[[], Any], Callable[[], Awaitable[Any]], Any] + ] = {}, ) -> list[ToolboxSyncTool]: """ Synchronously fetches a toolset and loads all tools defined within it. @@ -152,10 +162,12 @@ def load_toolset( bound_params: A mapping of parameter names to bind to specific values or callables that are called to produce values as needed. strict: If True, raises an error if *any* loaded tool instance fails - to utilize all of the given parameters or auth tokens. (if any + to utilize all of the given parameters, auth tokens, or secure parameters (if any provided). If False (default), raises an error only if a - user-provided parameter or auth token cannot be applied to *any* + user-provided parameter, auth token, or secure parameter cannot be applied to *any* loaded tool across the set. + secure_params: A mapping of secure parameter names to bind to specific values or + callables that are called to produce values as needed. Returns: list[ToolboxSyncTool]: A list of callables, one for each tool defined @@ -165,7 +177,7 @@ def load_toolset( ValueError: If validation fails based on the `strict` flag. """ coro = self.__async_client.load_toolset( - name, auth_token_getters, bound_params, strict + name, auth_token_getters, bound_params, strict, secure_params ) if not self.__loop or not self.__thread: diff --git a/packages/toolbox-core/tests/test_client.py b/packages/toolbox-core/tests/test_client.py index bd67fde9b..6730d0066 100644 --- a/packages/toolbox-core/tests/test_client.py +++ b/packages/toolbox-core/tests/test_client.py @@ -68,6 +68,7 @@ async def tool_invoke( arguments: dict, headers: Mapping[str, str], telemetry_attributes: Optional[TelemetryAttributes] = None, + secure_arguments: Optional[dict] = None, ) -> str: return await self.tool_invoke_mock(tool_name, arguments, headers) @@ -1030,3 +1031,152 @@ async def test_multistep_cascading_fallback(): assert mock_create.call_count == 2 assert mock_create.call_args_list[0][0][0] == Protocol.MCP_v20251125 assert mock_create.call_args_list[1][0][0] == Protocol.MCP_v20241105 + + +class TestClientSecureParams: + """Tests for secure parameters support in ToolboxClient.""" + + @pytest.mark.asyncio + async def test_load_tool_with_secure_params_success(self, mock_transport): + schema = ToolSchema( + description="A tool with secure parameters", + parameters=[ + ParameterSchema( + name="query", + type="string", + description="search query", + required=True, + ) + ], + secure_parameters=[ + ParameterSchema( + name="api_token", + type="string", + description="API Token", + required=True, + ) + ], + ) + mock_transport.tool_get_mock.return_value = ManifestSchema( + serverVersion="1.0.0", tools={"my_tool": schema} + ) + + client = ToolboxClient(TEST_BASE_URL) + client._ToolboxClient__transport = mock_transport + + tool = await client.load_tool( + "my_tool", + secure_params={"api_token": "token-123"}, + ) + + assert tool._bound_secure_params == {"api_token": "token-123"} + assert len(tool._secure_params) == 0 + + @pytest.mark.asyncio + async def test_load_tool_with_unused_secure_params_raises(self, mock_transport): + schema = ToolSchema( + description="A tool with secure parameters", + parameters=[ + ParameterSchema( + name="query", + type="string", + description="search query", + required=True, + ) + ], + secure_parameters=[ + ParameterSchema( + name="api_token", + type="string", + description="API Token", + required=True, + ) + ], + ) + mock_transport.tool_get_mock.return_value = ManifestSchema( + serverVersion="1.0.0", tools={"my_tool": schema} + ) + + client = ToolboxClient(TEST_BASE_URL) + client._ToolboxClient__transport = mock_transport + + with pytest.raises( + ValueError, + match=r"Validation failed for tool 'my_tool': unused secure parameters: extra_token\.", + ): + await client.load_tool( + "my_tool", + secure_params={"api_token": "token-123", "extra_token": "extra"}, + ) + + @pytest.mark.asyncio + async def test_load_toolset_with_secure_params_success(self, mock_transport): + schema1 = ToolSchema( + description="Tool 1", + parameters=[], + secure_parameters=[ + ParameterSchema( + name="api_token", + type="string", + description="API Token", + required=True, + ) + ], + ) + schema2 = ToolSchema( + description="Tool 2", + parameters=[], + secure_parameters=[ + ParameterSchema( + name="db_pass", + type="string", + description="DB Password", + required=False, + ) + ], + ) + mock_transport.tools_list_mock.return_value = ManifestSchema( + serverVersion="1.0.0", tools={"tool1": schema1, "tool2": schema2} + ) + + client = ToolboxClient(TEST_BASE_URL) + client._ToolboxClient__transport = mock_transport + + tools = await client.load_toolset( + "my_set", + secure_params={"api_token": "token-123", "db_pass": "pass-123"}, + ) + + assert len(tools) == 2 + assert tools[0]._bound_secure_params == {"api_token": "token-123"} + assert tools[1]._bound_secure_params == {"db_pass": "pass-123"} + + @pytest.mark.asyncio + async def test_load_toolset_with_unused_secure_params_raises(self, mock_transport): + schema1 = ToolSchema( + description="Tool 1", + parameters=[], + secure_parameters=[ + ParameterSchema( + name="api_token", + type="string", + description="API Token", + required=True, + ) + ], + ) + mock_transport.tools_list_mock.return_value = ManifestSchema( + serverVersion="1.0.0", tools={"tool1": schema1} + ) + + client = ToolboxClient(TEST_BASE_URL) + client._ToolboxClient__transport = mock_transport + + with pytest.raises( + ValueError, + match=r"Validation failed for toolset 'my_set': unused secure parameters could not be applied to any tool: unused_sec\.", + ): + await client.load_toolset( + "my_set", + secure_params={"api_token": "token-123", "unused_sec": "unused"}, + ) diff --git a/packages/toolbox-core/tests/test_sync_client.py b/packages/toolbox-core/tests/test_sync_client.py index 6a6649b9f..603f1285e 100644 --- a/packages/toolbox-core/tests/test_sync_client.py +++ b/packages/toolbox-core/tests/test_sync_client.py @@ -15,7 +15,7 @@ import inspect from typing import Mapping, Optional -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, create_autospec, patch import pytest @@ -29,6 +29,7 @@ ) from toolbox_core.sync_client import ToolboxSyncClient from toolbox_core.sync_tool import ToolboxSyncTool +from toolbox_core.tool import ToolboxTool TEST_BASE_URL = "http://toolbox.example.com" @@ -65,6 +66,7 @@ async def tool_invoke( arguments: dict, headers: Mapping[str, str], telemetry_attributes: Optional[TelemetryAttributes] = None, + secure_arguments: Optional[dict] = None, ) -> str: return await self.tool_invoke_mock(tool_name, arguments, headers) @@ -578,3 +580,40 @@ def test_sync_client_init_with_client_info(sync_client_environment): _, kwargs = mock_async_client_cls.call_args assert kwargs["client_name"] == client_name assert kwargs["client_version"] == client_version + + +def test_sync_client_load_tool_with_secure_params(sync_client_environment): + """Tests load_tool forwards secure_params to async client.""" + client = ToolboxSyncClient(TEST_BASE_URL) + with patch.object( + client._ToolboxSyncClient__async_client, "load_tool", new_callable=AsyncMock + ) as mock_load: + mock_tool = create_autospec(ToolboxTool, instance=True) + mock_tool.__name__ = "my_tool" + mock_load.return_value = mock_tool + + sec_params = {"api_key": "secret123"} + tool = client.load_tool("my_tool", secure_params=sec_params) + + mock_load.assert_called_once_with("my_tool", {}, {}, sec_params) + assert isinstance(tool, ToolboxSyncTool) + + +def test_sync_client_load_toolset_with_secure_params(sync_client_environment): + """Tests load_toolset forwards secure_params to async client.""" + client = ToolboxSyncClient(TEST_BASE_URL) + with patch.object( + client._ToolboxSyncClient__async_client, + "load_toolset", + new_callable=AsyncMock, + ) as mock_load: + mock_tool = create_autospec(ToolboxTool, instance=True) + mock_tool.__name__ = "my_tool" + mock_load.return_value = [mock_tool] + + sec_params = {"api_key": "secret123"} + tools = client.load_toolset("my_set", secure_params=sec_params) + + mock_load.assert_called_once_with("my_set", {}, {}, False, sec_params) + assert len(tools) == 1 + assert isinstance(tools[0], ToolboxSyncTool)