Add method to run any node with input variables

Co-Authored-By: Chris Li <chris.li.2046@gmail.com>

Author: devin-ai-integration[bot]

core/taskengine/engine.go

@@ -1060,3 +1060,76 @@ func (n *Engine) GetWorkflowCount(user *model.User, payload *avsproto.GetWorkflo
 		Total: total,
 	}, nil
 }
+
+func (n *Engine) RunNodeWithInputs(nodeType string, nodeConfig map[string]interface{}, inputVariables map[string]interface{}) (map[string]interface{}, error) {
+	vm, err := NewVMWithData(nil, nil, n.smartWalletConfig, nil)
+	if err != nil {
+		return nil, err
+	}
+	
+	vm.WithLogger(n.logger).WithDb(n.db)
+	
+	node, err := CreateNodeFromType(nodeType, nodeConfig, "")
+	if err != nil {
+		return nil, err
+	}
+	
+	// Run the node with input variables
+	executionStep, err := vm.RunNodeWithInputs(node, inputVariables)
+	if err != nil {
+		return nil, err
+	}
+	
+	if !executionStep.Success {
+		return nil, fmt.Errorf("execution failed: %s", executionStep.Error)
+	}
+	
+	result := make(map[string]interface{})
+	
+	switch nodeType {
+	case "blockTrigger":
+		if codeOutput := executionStep.GetCustomCode(); codeOutput != nil && codeOutput.Data != nil {
+			var data map[string]interface{}
+			if err := json.Unmarshal(codeOutput.Data.Value, &data); err == nil {
+				result = data
+			}
+		}
+	case "restApi":
+		if restOutput := executionStep.GetRestApi(); restOutput != nil && restOutput.Data != nil {
+			var data map[string]interface{}
+			if err := json.Unmarshal(restOutput.Data.Value, &data); err == nil {
+				result["data"] = data
+			} else {
+				result["data"] = restOutput.Data
+			}
+		}
+	case "contractRead":
+		if readOutput := executionStep.GetContractRead(); readOutput != nil {
+			result["data"] = readOutput.Data
+		}
+	case "customCode":
+		if codeOutput := executionStep.GetCustomCode(); codeOutput != nil && codeOutput.Data != nil {
+			var data map[string]interface{}
+			if err := json.Unmarshal(codeOutput.Data.Value, &data); err == nil {
+				result = data
+			} else {
+				result["data"] = codeOutput.Data
+			}
+		}
+	case "branch":
+		if branchOutput := executionStep.GetBranch(); branchOutput != nil {
+			result["conditionId"] = branchOutput.ConditionId
+		}
+	case "filter":
+		if filterOutput := executionStep.GetFilter(); filterOutput != nil && filterOutput.Data != nil {
+			var data interface{}
+			if err := json.Unmarshal(filterOutput.Data.Value, &data); err == nil {
+				result["data"] = data
+			} else {
+				result["data"] = filterOutput.Data
+			}
+		}
+	}
+	
+	return result, nil
+}

core/taskengine/run_node_test.go

@@ -0,0 +1,145 @@
+package taskengine
+
+import (
+	"testing"
+	
+	avsproto "github.com/AvaProtocol/EigenLayer-AVS/protobuf"
+	"github.com/AvaProtocol/EigenLayer-AVS/core/config"
+	"github.com/stretchr/testify/assert"
+	"github.com/oklog/ulid/v2"
+)
+
+func TestRunNodeWithInputs_BlockTrigger(t *testing.T) {
+	vm, err := NewVMWithData(nil, nil, &config.SmartWalletConfig{}, nil)
+	assert.NoError(t, err)
+	
+	node, err := CreateNodeFromType("blockTrigger", map[string]interface{}{
+		"blockNumber": float64(12345),
+	}, "")
+	assert.NoError(t, err)
+	
+	result, err := vm.RunNodeWithInputs(node, map[string]interface{}{})
+	
+	assert.NoError(t, err)
+	assert.NotNil(t, result)
+	assert.True(t, result.Success)
+	
+	codeOutput := result.GetCustomCode()
+	assert.NotNil(t, codeOutput)
+	assert.NotNil(t, codeOutput.Data)
+}
+
+func TestRunNodeWithInputs_CustomCode(t *testing.T) {
+	vm, err := NewVMWithData(nil, nil, &config.SmartWalletConfig{}, nil)
+	assert.NoError(t, err)
+	
+	nodeId := "test_" + ulid.Make().String()
+	node := &avsproto.TaskNode{
+		Id:   nodeId,
+		Name: "Test Custom Code",
+		TaskType: &avsproto.TaskNode_CustomCode{
+			CustomCode: &avsproto.CustomCodeNode{
+				Lang: avsproto.CustomCodeLang_JavaScript,
+				Source: `
+					if (typeof myVar === 'undefined') {
+						throw new Error("myVar is required but not provided");
+					}
+					return { result: myVar * 2 };
+				`,
+			},
+		},
+	}
+	
+	_, err = vm.RunNodeWithInputs(node, map[string]interface{}{})
+	assert.Error(t, err)
+	
+	result, err := vm.RunNodeWithInputs(node, map[string]interface{}{
+		"myVar": 5,
+	})
+	assert.NoError(t, err)
+	assert.NotNil(t, result)
+	assert.True(t, result.Success)
+	
+	codeOutput := result.GetCustomCode()
+	assert.NotNil(t, codeOutput)
+	assert.NotNil(t, codeOutput.Data)
+}
+
+func TestCreateNodeFromType(t *testing.T) {
+	nodeTypes := []string{"blockTrigger", "restApi", "contractRead", "customCode", "branch", "filter"}
+	
+	for _, nodeType := range nodeTypes {
+		config := map[string]interface{}{}
+		
+		switch nodeType {
+		case "blockTrigger":
+			config["blockNumber"] = float64(12345)
+		case "restApi":
+			config["url"] = "https://example.com"
+			config["method"] = "GET"
+		case "contractRead":
+			config["contractAddress"] = "0x1234567890123456789012345678901234567890"
+			config["callData"] = "0x12345678"
+		case "customCode":
+			config["source"] = "return { hello: 'world' };"
+		case "branch":
+			config["conditions"] = []interface{}{
+				map[string]interface{}{
+					"id":         "cond1",
+					"type":       "if",
+					"expression": "true",
+				},
+			}
+		case "filter":
+			config["expression"] = "item > 5"
+			config["input"] = "inputArray"
+		}
+		
+		node, err := CreateNodeFromType(nodeType, config, "")
+		assert.NoError(t, err)
+		assert.NotNil(t, node)
+		assert.NotEmpty(t, node.Id)
+		assert.Equal(t, "Single Node Execution", node.Name)
+		
+		switch nodeType {
+		case "blockTrigger":
+			assert.NotNil(t, node.GetCustomCode())
+		case "restApi":
+			assert.NotNil(t, node.GetRestApi())
+		case "contractRead":
+			assert.NotNil(t, node.GetContractRead())
+		case "customCode":
+			assert.NotNil(t, node.GetCustomCode())
+		case "branch":
+			assert.NotNil(t, node.GetBranch())
+		case "filter":
+			assert.NotNil(t, node.GetFilter())
+		}
+	}
+}
+
+func TestEngine_RunNodeWithInputs(t *testing.T) {
+	engine := New(nil, &config.Config{
+		SmartWallet: &config.SmartWalletConfig{},
+	}, nil, nil)
+	
+	result, err := engine.RunNodeWithInputs("blockTrigger", map[string]interface{}{
+		"blockNumber": float64(12345),
+	}, map[string]interface{}{})
+	
+	if err == nil {
+		assert.NotNil(t, result)
+		assert.Contains(t, result, "blockNumber")
+		assert.Equal(t, float64(12345), result["blockNumber"])
+	}
+	
+	result, err = engine.RunNodeWithInputs("customCode", map[string]interface{}{
+		"source": "return { message: 'Hello, World!' };",
+	}, map[string]interface{}{})
+	
+	if err == nil {
+		assert.NotNil(t, result)
+		assert.Contains(t, result, "message")
+		assert.Equal(t, "Hello, World!", result["message"])
+	}
+}