diff --git a/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_tool_calling_node.py b/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_tool_calling_node.py index b5d12e6e5c..f8fd9270d4 100644 --- a/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_tool_calling_node.py +++ b/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_tool_calling_node.py @@ -246,6 +246,16 @@ class Workflow(BaseWorkflow): "environment_variable": "my-api-key-header-value", }, }, + { + "id": "0ce4327e-72ab-4fb3-83ab-9cce5d398388", + "key": "include_tools", + "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}}, + }, + { + "id": "9ee4fd57-b73a-4b97-991e-14cad22645fa", + "key": "exclude_tools", + "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}}, + }, ], "definition": {"name": "MCPServer", "module": ["vellum", "workflows", "types", "definition"]}, } @@ -254,7 +264,60 @@ class Workflow(BaseWorkflow): } +def test_serialize_node__tool_calling_node__mcp_server_with_tool_filtering(): + """Tests that MCPServer with include_tools and exclude_tools serializes correctly.""" + + # GIVEN a tool calling node with an mcp server that has tool filtering configured + class MyToolCallingNode(ToolCallingNode): + functions = [ + MCPServer( + name="my-mcp-server", + url="https://my-mcp-server.com", + authorization_type=AuthorizationType.API_KEY, + api_key_header_key="my-api-key-header-key", + api_key_header_value=EnvironmentVariableReference(name="my-api-key-header-value"), + include_tools=["tool-a", "tool-b"], + exclude_tools=["tool-c"], + ) + ] + + # AND a workflow with the tool calling node + class Workflow(BaseWorkflow): + graph = MyToolCallingNode + + # WHEN the workflow is serialized + workflow_display = get_workflow_display(workflow_class=Workflow) + serialized_workflow: dict = workflow_display.serialize() + + # THEN the node should properly serialize the mcp server with tool filtering + my_tool_calling_node = next( + node + for node in serialized_workflow["workflow_raw_data"]["nodes"] + if node["id"] == str(MyToolCallingNode.__id__) + ) + + functions_attribute = next( + attribute for attribute in my_tool_calling_node["attributes"] if attribute["name"] == "functions" + ) + + # AND the include_tools and exclude_tools should be serialized as JSON arrays + entries = functions_attribute["value"]["items"][0]["entries"] + include_tools_entry = next(entry for entry in entries if entry["key"] == "include_tools") + exclude_tools_entry = next(entry for entry in entries if entry["key"] == "exclude_tools") + + assert include_tools_entry["value"] == { + "type": "CONSTANT_VALUE", + "value": {"type": "JSON", "value": ["tool-a", "tool-b"]}, + } + assert exclude_tools_entry["value"] == { + "type": "CONSTANT_VALUE", + "value": {"type": "JSON", "value": ["tool-c"]}, + } + + def test_serialize_node__tool_calling_node__mcp_server_no_authorization(): + """Tests that MCPServer without authorization serializes correctly.""" + # GIVEN a tool calling node with an mcp server class MyToolCallingNode(ToolCallingNode): functions = [ @@ -300,6 +363,8 @@ class Workflow(BaseWorkflow): "bearer_token_value": None, "api_key_header_key": None, "api_key_header_value": None, + "include_tools": None, + "exclude_tools": None, } ], }, diff --git a/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_mcp_serialization.py b/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_mcp_serialization.py index 15bc998c48..69da255879 100644 --- a/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_mcp_serialization.py +++ b/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_mcp_serialization.py @@ -96,6 +96,16 @@ def test_serialize_workflow(): "key": "api_key_header_value", "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}}, }, + { + "id": "90981dfe-9dc8-4f40-a5d6-b8b42edfc075", + "key": "include_tools", + "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}}, + }, + { + "id": "d48a1658-0d94-49fc-9499-5f400079eebd", + "key": "exclude_tools", + "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}}, + }, ], "definition": {"name": "MCPServer", "module": ["vellum", "workflows", "types", "definition"]}, } diff --git a/src/vellum/workflows/integrations/mcp_service.py b/src/vellum/workflows/integrations/mcp_service.py index 68df3f4c1d..e3d767d117 100644 --- a/src/vellum/workflows/integrations/mcp_service.py +++ b/src/vellum/workflows/integrations/mcp_service.py @@ -264,10 +264,28 @@ def execute_tool(self, tool_def: MCPToolDefinition, arguments: Dict[str, Any]) - ) def hydrate_tool_definitions(self, server_def: MCPServer) -> List[MCPToolDefinition]: - """Hydrate an MCPToolDefinition with detailed information from the MCP server.""" + """Hydrate an MCPToolDefinition with detailed information from the MCP server. + + Tool filtering is applied based on include_tools and exclude_tools: + - include_tools acts as a whitelist: only tools in this list are included + - exclude_tools acts as a blacklist: tools in this list are filtered out + - exclude_tools takes precedence: if a tool appears in both lists, it is excluded + """ try: tools = self.list_tools(server_def) + filtered_tools = [] + for tool in tools: + tool_name = tool["name"] + + if server_def.include_tools is not None and tool_name not in server_def.include_tools: + continue + + if server_def.exclude_tools is not None and tool_name in server_def.exclude_tools: + continue + + filtered_tools.append(tool) + return [ MCPToolDefinition( name=tool["name"], @@ -275,7 +293,7 @@ def hydrate_tool_definitions(self, server_def: MCPServer) -> List[MCPToolDefinit description=tool["description"], parameters=tool["inputSchema"], ) - for tool in tools + for tool in filtered_tools ] except Exception as e: logger.warning(f"Failed to hydrate MCP server '{server_def.name}': {e}") diff --git a/src/vellum/workflows/integrations/tests/test_mcp_service.py b/src/vellum/workflows/integrations/tests/test_mcp_service.py index 39fd6b6d53..cfdffa335b 100644 --- a/src/vellum/workflows/integrations/tests/test_mcp_service.py +++ b/src/vellum/workflows/integrations/tests/test_mcp_service.py @@ -271,3 +271,125 @@ def test_mcp_service_call_tool_includes_stacktrace_and_raw_data_on_error(): assert exc_info.value.raw_data["operation"] == "call_tool" assert exc_info.value.raw_data["error_type"] == "RuntimeError" assert exc_info.value.raw_data["error_message"] == "Tool execution failed" + + +def test_mcp_service_hydrate_tool_definitions__include_tools_filters_to_whitelist(): + """ + Tests that include_tools acts as a whitelist, only including specified tools. + """ + # GIVEN an MCP server with include_tools specified + sample_mcp_server = MCPServer( + name="test-server", + url="https://test.mcp.server.com/mcp", + include_tools=["tool-a", "tool-c"], + ) + + # AND a mock MCP service that returns multiple tools + with mock.patch("vellum.workflows.integrations.mcp_service.asyncio.run") as mock_run: + mock_run.return_value = [ + {"name": "tool-a", "description": "Tool A", "inputSchema": {}}, + {"name": "tool-b", "description": "Tool B", "inputSchema": {}}, + {"name": "tool-c", "description": "Tool C", "inputSchema": {}}, + ] + + # WHEN we hydrate tool definitions + service = MCPService() + tool_definitions = service.hydrate_tool_definitions(sample_mcp_server) + + # THEN only the tools in include_tools should be returned + assert len(tool_definitions) == 2 + tool_names = [t.name for t in tool_definitions] + assert "tool-a" in tool_names + assert "tool-c" in tool_names + assert "tool-b" not in tool_names + + +def test_mcp_service_hydrate_tool_definitions__exclude_tools_filters_out_blacklist(): + """ + Tests that exclude_tools acts as a blacklist, filtering out specified tools. + """ + # GIVEN an MCP server with exclude_tools specified + sample_mcp_server = MCPServer( + name="test-server", + url="https://test.mcp.server.com/mcp", + exclude_tools=["tool-b"], + ) + + # AND a mock MCP service that returns multiple tools + with mock.patch("vellum.workflows.integrations.mcp_service.asyncio.run") as mock_run: + mock_run.return_value = [ + {"name": "tool-a", "description": "Tool A", "inputSchema": {}}, + {"name": "tool-b", "description": "Tool B", "inputSchema": {}}, + {"name": "tool-c", "description": "Tool C", "inputSchema": {}}, + ] + + # WHEN we hydrate tool definitions + service = MCPService() + tool_definitions = service.hydrate_tool_definitions(sample_mcp_server) + + # THEN all tools except the excluded one should be returned + assert len(tool_definitions) == 2 + tool_names = [t.name for t in tool_definitions] + assert "tool-a" in tool_names + assert "tool-c" in tool_names + assert "tool-b" not in tool_names + + +def test_mcp_service_hydrate_tool_definitions__exclude_takes_precedence_over_include(): + """ + Tests that exclude_tools takes precedence when a tool appears in both lists. + """ + # GIVEN an MCP server with both include_tools and exclude_tools + # AND tool-a appears in both lists + sample_mcp_server = MCPServer( + name="test-server", + url="https://test.mcp.server.com/mcp", + include_tools=["tool-a", "tool-b"], + exclude_tools=["tool-a"], + ) + + # AND a mock MCP service that returns multiple tools + with mock.patch("vellum.workflows.integrations.mcp_service.asyncio.run") as mock_run: + mock_run.return_value = [ + {"name": "tool-a", "description": "Tool A", "inputSchema": {}}, + {"name": "tool-b", "description": "Tool B", "inputSchema": {}}, + {"name": "tool-c", "description": "Tool C", "inputSchema": {}}, + ] + + # WHEN we hydrate tool definitions + service = MCPService() + tool_definitions = service.hydrate_tool_definitions(sample_mcp_server) + + # THEN tool-a should be excluded even though it's in include_tools + assert len(tool_definitions) == 1 + assert tool_definitions[0].name == "tool-b" + + +def test_mcp_service_hydrate_tool_definitions__no_filtering_when_both_none(): + """ + Tests that all tools are returned when neither include_tools nor exclude_tools is specified. + """ + # GIVEN an MCP server with no filtering configured + sample_mcp_server = MCPServer( + name="test-server", + url="https://test.mcp.server.com/mcp", + ) + + # AND a mock MCP service that returns multiple tools + with mock.patch("vellum.workflows.integrations.mcp_service.asyncio.run") as mock_run: + mock_run.return_value = [ + {"name": "tool-a", "description": "Tool A", "inputSchema": {}}, + {"name": "tool-b", "description": "Tool B", "inputSchema": {}}, + {"name": "tool-c", "description": "Tool C", "inputSchema": {}}, + ] + + # WHEN we hydrate tool definitions + service = MCPService() + tool_definitions = service.hydrate_tool_definitions(sample_mcp_server) + + # THEN all tools should be returned + assert len(tool_definitions) == 3 + tool_names = [t.name for t in tool_definitions] + assert "tool-a" in tool_names + assert "tool-b" in tool_names + assert "tool-c" in tool_names diff --git a/src/vellum/workflows/types/definition.py b/src/vellum/workflows/types/definition.py index cf63f9ffe9..a23c1c3ae6 100644 --- a/src/vellum/workflows/types/definition.py +++ b/src/vellum/workflows/types/definition.py @@ -205,6 +205,8 @@ class MCPServer(UniversalBaseModel): bearer_token_value: Optional[Union[str, EnvironmentVariableReference]] = None api_key_header_key: Optional[str] = None api_key_header_value: Optional[Union[str, EnvironmentVariableReference]] = None + include_tools: Optional[List[str]] = None + exclude_tools: Optional[List[str]] = None model_config = {"arbitrary_types_allowed": True}