Add stat endpoint

aggregator/rpc_server.go

@@ -56,7 +56,7 @@ func (r *RpcServer) GetWallet(ctx context.Context, payload *avsproto.GetWalletRe
 		"factory", payload.FactoryAddress,
 	)
 
-	return r.engine.CreateSmartWallet(user, payload)
+	return r.engine.GetWallet(user, payload)
 }
 
 // Get nonce of an existing smart wallet of a given owner

core/taskengine/engine.go

@@ -226,7 +226,7 @@ func (n *Engine) GetSmartWallets(owner common.Address, payload *avsproto.ListWal
 	return wallets, nil
 }
 
-func (n *Engine) CreateSmartWallet(user *model.User, payload *avsproto.GetWalletReq) (*avsproto.GetWalletResp, error) {
+func (n *Engine) GetWallet(user *model.User, payload *avsproto.GetWalletReq) (*avsproto.GetWalletResp, error) {
 	// Verify data
 	// when user passing a custom factory address, we want to validate it
 	if payload.FactoryAddress != "" && !common.IsHexAddress(payload.FactoryAddress) {
@@ -266,10 +266,19 @@ func (n *Engine) CreateSmartWallet(user *model.User, payload *avsproto.GetWallet
 		return nil, status.Errorf(codes.Code(avsproto.Error_StorageWriteError), StorageWriteError)
 	}
 
+	statSvc := NewStatService(n.db)
+	stat, _ := statSvc.GetTaskCount(wallet)
+
 	return &avsproto.GetWalletResp{
 		Address:        sender.Hex(),
 		Salt:           salt.String(),
 		FactoryAddress: factoryAddress.Hex(),
+
+		TotalTaskCount:     stat.Total,
+		ActiveTaskCount:    stat.Active,
+		CompletedTaskCount: stat.Completed,
+		FailedTaskCount:    stat.Failed,
+		CanceledTaskCount:  stat.Canceled,
 	}, nil
 }
 

core/taskengine/engine_test.go

@@ -18,10 +18,10 @@ func TestListTasks(t *testing.T) {
 
 	config := testutil.GetAggregatorConfig()
 	n := New(db, config, nil, testutil.GetLogger())
-	n.CreateSmartWallet(testutil.TestUser1(), &avsproto.GetWalletReq{
+	n.GetWallet(testutil.TestUser1(), &avsproto.GetWalletReq{
 		Salt: "12345",
 	})
-	n.CreateSmartWallet(testutil.TestUser1(), &avsproto.GetWalletReq{
+	n.GetWallet(testutil.TestUser1(), &avsproto.GetWalletReq{
 		Salt: "6789",
 	})
 
@@ -80,10 +80,10 @@ func TestListTasksPagination(t *testing.T) {
 
 	config := testutil.GetAggregatorConfig()
 	n := New(db, config, nil, testutil.GetLogger())
-	n.CreateSmartWallet(testutil.TestUser1(), &avsproto.GetWalletReq{
+	n.GetWallet(testutil.TestUser1(), &avsproto.GetWalletReq{
 		Salt: "12345",
 	})
-	n.CreateSmartWallet(testutil.TestUser1(), &avsproto.GetWalletReq{
+	n.GetWallet(testutil.TestUser1(), &avsproto.GetWalletReq{
 		Salt: "6789",
 	})
 
@@ -241,10 +241,10 @@ func TestListWallets(t *testing.T) {
 	n := New(db, config, nil, testutil.GetLogger())
 	u := testutil.TestUser1()
 
-	n.CreateSmartWallet(u, &avsproto.GetWalletReq{
+	n.GetWallet(u, &avsproto.GetWalletReq{
 		Salt: "12345",
 	})
-	n.CreateSmartWallet(u, &avsproto.GetWalletReq{
+	n.GetWallet(u, &avsproto.GetWalletReq{
 		Salt: "9876",
 		// https://sepolia.etherscan.io/address/0x9406Cc6185a346906296840746125a0E44976454#readProxyContract
 		FactoryAddress: "0x9406Cc6185a346906296840746125a0E44976454",
@@ -773,3 +773,52 @@ func TestListSecrets(t *testing.T) {
 	}
 
 }
+
+func TestGetWalletReturnTaskStat(t *testing.T) {
+	db := testutil.TestMustDB()
+	defer storage.Destroy(db.(*storage.BadgerStorage))
+
+	config := testutil.GetAggregatorConfig()
+	n := New(db, config, nil, testutil.GetLogger())
+
+	user1 := testutil.TestUser1()
+	// Now create a test task
+	tr1 := testutil.RestTask()
+	tr1.Name = "t1"
+	// salt 0 wallet
+	tr1.SmartWalletAddress = "0x7c3a76086588230c7B3f4839A4c1F5BBafcd57C6"
+
+	result, _ := n.GetWallet(user1, &avsproto.GetWalletReq{
+		Salt: "0",
+	})
+
+	if result.TotalTaskCount > 0 {
+		t.Errorf("expect no task count yet but got :%d", result.TotalTaskCount)
+	}
+
+	taskResult, _ := n.CreateTask(testutil.TestUser1(), tr1)
+	result, _ = n.GetWallet(user1, &avsproto.GetWalletReq{
+		Salt: "0",
+	})
+
+	if result.TotalTaskCount != 1 || result.ActiveTaskCount != 1 || result.CompletedTaskCount != 0 {
+		t.Errorf("expect total=1 active=1 completed=0 but got %v", result)
+	}
+
+	// Make the task run to simulate completed count
+	n.TriggerTask(testutil.TestUser1(), &avsproto.UserTriggerTaskReq{
+		TaskId: taskResult.Id,
+		TriggerMetadata: &avsproto.TriggerMetadata{
+			BlockNumber: 101,
+		},
+		IsBlocking: true,
+	})
+
+	result, _ = n.GetWallet(user1, &avsproto.GetWalletReq{
+		Salt: "0",
+	})
+
+	if result.TotalTaskCount != 1 || result.ActiveTaskCount != 0 || result.CompletedTaskCount != 1 {
+		t.Errorf("expect total=1 active=0 completed=1 but got %v", result)
+	}
+}

core/taskengine/stats.go

@@ -0,0 +1,46 @@
+package taskengine
+
+import (
+	"strconv"
+
+	"github.com/AvaProtocol/ap-avs/model"
+	avsproto "github.com/AvaProtocol/ap-avs/protobuf"
+	"github.com/AvaProtocol/ap-avs/storage"
+)
+
+type StatService struct {
+	db storage.Storage
+}
+
+func NewStatService(db storage.Storage) *StatService {
+	return &StatService{
+		db: db,
+	}
+}
+
+func (svc *StatService) GetTaskCount(smartWalletAddress *model.SmartWallet) (*model.SmartWalletTaskStat, error) {
+	stat := &model.SmartWalletTaskStat{}
+
+	prefix := SmartWalletTaskStoragePrefix(*smartWalletAddress.Owner, *smartWalletAddress.Address)
+	items, err := svc.db.GetByPrefix(prefix)
+	if err != nil {
+		return stat, err
+	}
+
+	for _, item := range items {
+		taskStatus, _ := strconv.ParseInt(string(item.Value), 10, 32)
+		stat.Total += 1
+		switch avsproto.TaskStatus(taskStatus) {
+		case avsproto.TaskStatus_Active:
+			stat.Active += 1
+		case avsproto.TaskStatus_Completed:
+			stat.Completed += 1
+		case avsproto.TaskStatus_Failed:
+			stat.Failed += 1
+		case avsproto.TaskStatus_Canceled:
+			stat.Canceled += 1
+		}
+	}
+
+	return stat, nil
+}

core/taskengine/stats_test.go

@@ -0,0 +1,129 @@
+package taskengine
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+
+	"github.com/AvaProtocol/ap-avs/core/testutil"
+	"github.com/AvaProtocol/ap-avs/model"
+	avsproto "github.com/AvaProtocol/ap-avs/protobuf"
+	"github.com/AvaProtocol/ap-avs/storage"
+)
+
+func TestTaskStatCount(t *testing.T) {
+	db := testutil.TestMustDB()
+	defer storage.Destroy(db.(*storage.BadgerStorage))
+
+	config := testutil.GetAggregatorConfig()
+	n := New(db, config, nil, testutil.GetLogger())
+
+	user1 := testutil.TestUser1()
+
+	// Populate task
+	tr1 := testutil.RestTask()
+	tr1.Name = "t1"
+	tr1.MaxExecution = 1
+	// salt 0
+	tr1.SmartWalletAddress = "0x7c3a76086588230c7B3f4839A4c1F5BBafcd57C6"
+	n.CreateTask(testutil.TestUser1(), tr1)
+
+	statSvc := NewStatService(db)
+	result, _ := statSvc.GetTaskCount(user1.ToSmartWallet())
+
+	if !reflect.DeepEqual(
+		result, &model.SmartWalletTaskStat{
+			Total:  1,
+			Active: 1,
+		}) {
+		t.Errorf("expect task total is 1, but got %v", result)
+	}
+}
+
+func TestTaskStatCountCompleted(t *testing.T) {
+	db := testutil.TestMustDB()
+	defer storage.Destroy(db.(*storage.BadgerStorage))
+
+	user1 := testutil.TestUser1()
+
+	task1 := &model.Task{
+		&avsproto.Task{
+			Owner:              user1.Address.Hex(),
+			SmartWalletAddress: user1.SmartAccountAddress.Hex(),
+			Id:                 "t1",
+		},
+	}
+
+	db.Set(TaskUserKey(task1), []byte(fmt.Sprintf("%d", avsproto.TaskStatus_Completed)))
+
+	statSvc := NewStatService(db)
+	result, _ := statSvc.GetTaskCount(user1.ToSmartWallet())
+
+	if !reflect.DeepEqual(
+		result, &model.SmartWalletTaskStat{
+			Total:     1,
+			Active:    0,
+			Completed: 1,
+		}) {
+		t.Errorf("expect task total is 1, completed is 1 but got %v", result)
+	}
+}
+
+func TestTaskStatCountAllStatus(t *testing.T) {
+	db := testutil.TestMustDB()
+	defer storage.Destroy(db.(*storage.BadgerStorage))
+
+	user1 := testutil.TestUser1()
+
+	task1 := &model.Task{
+		&avsproto.Task{
+			Owner:              user1.Address.Hex(),
+			SmartWalletAddress: user1.SmartAccountAddress.Hex(),
+			Id:                 "t1",
+		},
+	}
+
+	task2 := &model.Task{
+		&avsproto.Task{
+			Owner:              user1.Address.Hex(),
+			SmartWalletAddress: user1.SmartAccountAddress.Hex(),
+			Id:                 "t2",
+		},
+	}
+
+	task3 := &model.Task{
+		&avsproto.Task{
+			Owner:              user1.Address.Hex(),
+			SmartWalletAddress: user1.SmartAccountAddress.Hex(),
+			Id:                 "t3",
+		},
+	}
+
+	task4 := &model.Task{
+		&avsproto.Task{
+			Owner:              user1.Address.Hex(),
+			SmartWalletAddress: user1.SmartAccountAddress.Hex(),
+			Id:                 "t4",
+		},
+	}
+
+	db.Set(TaskUserKey(task1), []byte(fmt.Sprintf("%d", avsproto.TaskStatus_Completed)))
+	db.Set(TaskUserKey(task2), []byte(fmt.Sprintf("%d", avsproto.TaskStatus_Failed)))
+	db.Set(TaskUserKey(task3), []byte(fmt.Sprintf("%d", avsproto.TaskStatus_Canceled)))
+	db.Set(TaskUserKey(task4), []byte(fmt.Sprintf("%d", avsproto.TaskStatus_Active)))
+
+	statSvc := NewStatService(db)
+	result, _ := statSvc.GetTaskCount(user1.ToSmartWallet())
+
+	if !reflect.DeepEqual(
+		result, &model.SmartWalletTaskStat{
+			Total:     4,
+			Active:    1,
+			Completed: 1,
+			Failed:    1,
+			Canceled:  1,
+		}) {
+		t.Errorf("expect task total=4, active=1, completed=1, failed=1, canceled=1, but got %v", result)
+	}
+
+}

model/user.go

@@ -24,6 +24,14 @@ func (u *User) LoadDefaultSmartWallet(rpcClient *ethclient.Client) error {
 	return nil
 }
 
+// Return the smartwallet struct re-present the default wallet for this user
+func (u *User) ToSmartWallet() *SmartWallet {
+	return &SmartWallet{
+		Owner:   &u.Address,
+		Address: u.SmartAccountAddress,
+	}
+}
+
 type SmartWallet struct {
 	Owner   *common.Address `json:"owner"`
 	Address *common.Address `json:"address"`
@@ -40,3 +48,11 @@ func (w *SmartWallet) FromStorageData(body []byte) error {
 
 	return err
 }
+
+type SmartWalletTaskStat struct {
+	Total     uint64
+	Active    uint64
+	Completed uint64
+	Failed    uint64
+	Canceled  uint64
+}