⚡️ Update TX and api code

internal/types/fees.go

@@ -8,32 +8,23 @@ import (
 
 // DB FEE Types
 
-type FeeRunState string
+type FeeBatchState string
 
 const (
-	FeeRunStateDraft   FeeRunState = "draft"
-	FeeRunStateSent    FeeRunState = "sent"
-	FeeRunStateSuccess FeeRunState = "completed"
-	FeeRunStateFailed  FeeRunState = "failed"
+	FeeBatchStateDraft   FeeBatchState = "draft"
+	FeeBatchStateSent    FeeBatchState = "sent"
+	FeeBatchStateSuccess FeeBatchState = "completed"
+	FeeBatchStateFailed  FeeBatchState = "failed"
 )
 
 // individual fee record in the db
-type Fee struct {
-	ID        uuid.UUID `db:"id"`
-	FeeRunID  uuid.UUID `db:"fee_run_id"`
-	Amount    int       `db:"amount"`
-	CreatedAt time.Time `db:"created_at"`
-}
-
-// fee table or fee_run_with_totals
-type FeeRun struct {
-	ID          uuid.UUID   `db:"id"`
-	Status      FeeRunState `db:"status"`
-	CreatedAt   time.Time   `db:"created_at"`
-	UpdatedAt   time.Time   `db:"updated_at"`
-	TxHash      *string     `db:"tx_hash"`
-	PolicyID    uuid.UUID   `db:"policy_id"`
-	TotalAmount int         `db:"total_amount"`
-	FeeCount    int         `db:"fee_count"`
-	Fees        []Fee       `db:"fees"`
+type FeeBatch struct {
+	ID        uuid.UUID     `db:"id"`
+	BatchID   uuid.UUID     `db:"batch_id"`
+	PublicKey string        `db:"public_key"`
+	Status    FeeBatchState `db:"status"`
+	Amount    uint64        `db:"amount"`
+	CreatedAt time.Time     `db:"created_at"`
+	UpdatedAt time.Time     `db:"updated_at"`
+	TxHash    *string       `db:"tx_hash"`
 }

internal/verifierapi/fees.go

@@ -4,9 +4,9 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
-	"time"
 
 	"github.com/google/uuid"
+	"github.com/vultisig/plugin/internal/types"
 )
 
 type FeeDto struct {
@@ -27,16 +27,47 @@ type FeeHistoryDto struct {
 	FeesPendingCollection int       `json:"fees_pending_collection" validate:"required"` // Total fees pending collection in the smallest unit, e.g., "1000000" for 0.01 VULTI
 }
 
-func (v *VerifierApi) GetPublicKeysFees(ecdsaPublicKey string) (*FeeHistoryDto, error) {
-	response, err := v.getAuth(fmt.Sprintf("/fees/publickey/%s", ecdsaPublicKey))
+type FeeBalanceDto struct {
+	Balance   int64  `json:"balance" validate:"required"`
+	PublicKey string `json:"public_key" validate:"required"`
+}
+
+type FeeBatchCreateResponseDto struct {
+	PublicKey string    `json:"public_key" validate:"required"`
+	Amount    uint64    `json:"amount" validate:"required"`
+	BatchID   uuid.UUID `json:"batch_id" validate:"required"`
+}
+
+type FeeBatchUpdateRequestResponseDto struct {
+	PublicKey string              `json:"public_key" validate:"required"`
+	BatchID   uuid.UUID           `json:"batch_id" validate:"required"`
+	TxHash    string              `json:"tx_hash" validate:"required"`
+	Status    types.FeeBatchState `json:"status" validate:"required"`
+}
+
+func (v *VerifierApi) CreateFeeBatch(publicKey string) (*FeeBatchCreateResponseDto, error) {
+	response, err := v.postAuth("/fees/batch", map[string]interface{}{
+		"public_key": publicKey,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to create fee batch: %w", err)
+	}
+	defer response.Body.Close()
+
+	var feeBatchResponse APIResponse[FeeBatchCreateResponseDto]
+	if err := json.NewDecoder(response.Body).Decode(&feeBatchResponse); err != nil {
+		return nil, fmt.Errorf("failed to decode fee batch response: %w", err)
+	}
+
+	return &feeBatchResponse.Data, nil
+}
+
+func (v *VerifierApi) GetFeeHistory(ecdsaPublicKey string) (*FeeHistoryDto, error) {
+	response, err := v.getAuth(fmt.Sprintf("/fees/history/%s", ecdsaPublicKey))
 	if err != nil {
 		return nil, fmt.Errorf("failed to get public key fees: %w", err)
 	}
-	defer func() {
-		if err := response.Body.Close(); err != nil {
-			v.logger.WithError(err).Error("Failed to close response body")
-		}
-	}()
+	defer response.Body.Close()
 	if response.StatusCode == http.StatusNotFound {
 		return nil, fmt.Errorf("public key not found")
 	}
@@ -57,28 +88,64 @@ func (v *VerifierApi) GetPublicKeysFees(ecdsaPublicKey string) (*FeeHistoryDto,
 	return &feeHistory.Data, nil
 }
 
-func (v *VerifierApi) MarkFeeAsCollected(txHash string, collectedAt time.Time, feeIds ...uuid.UUID) error {
+func (v *VerifierApi) GetFeeBalance(ecdsaPublicKey string) (*FeeBalanceDto, error) {
+	response, err := v.getAuth(fmt.Sprintf("/fees/balance/%s", ecdsaPublicKey))
+	if err != nil {
+		return nil, fmt.Errorf("failed to get public key fees: %w", err)
+	}
+	defer response.Body.Close()
+
+	if response.StatusCode == http.StatusNotFound {
+		return nil, fmt.Errorf("public key not found")
+	}
+
+	if response.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("failed to get public key fees, status code: %d", response.StatusCode)
+	}
 
-	var body = struct {
-		IDs         []uuid.UUID `json:"ids"`
-		TxHash      string      `json:"tx_hash"`
-		CollectedAt time.Time   `json:"collected_at"`
-	}{
-		IDs:         feeIds,
-		TxHash:      txHash,
-		CollectedAt: collectedAt,
+	var feeBalance APIResponse[FeeBalanceDto]
+	if err := json.NewDecoder(response.Body).Decode(&feeBalance); err != nil {
+		return nil, fmt.Errorf("failed to decode public key fees response: %w", err)
 	}
 
-	url := "/fees/collected"
-	response, err := v.postAuth(url, body)
+	return &feeBalance.Data, nil
+}
+
+func (v *VerifierApi) CreateFeeCredit(id uuid.UUID, amount int64, publicKey string) error {
+	response, err := v.postAuth("/fees/credit", map[string]interface{}{
+		"id":         id,
+		"amount":     amount,
+		"public_key": publicKey,
+	})
 	if err != nil {
-		return fmt.Errorf("failed to mark fee as collected: %w", err)
+		return fmt.Errorf("failed to create fee credit: %w", err)
 	}
 	defer response.Body.Close()
 
-	if response.StatusCode != http.StatusOK {
-		return fmt.Errorf("failed to mark fee as collected, status code: %d", response.StatusCode)
+	return nil
+}
+
+func (v *VerifierApi) RevertFeeCredit(txHash string, batchId uuid.UUID) error {
+	response, err := v.postAuth(fmt.Sprintf("/fees/revert/%s", batchId), struct{}{})
+	if err != nil {
+		return fmt.Errorf("failed to revert fee credit: %w", err)
 	}
+	defer response.Body.Close()
 
 	return nil
 }
+
+func (v *VerifierApi) UpdateFeeBatchTxHash(publickey string, batchId uuid.UUID, hash string) (*FeeBatchCreateResponseDto, error) {
+	response, err := v.putAuth("/fees/batch", FeeBatchUpdateRequestResponseDto{
+		PublicKey: publickey,
+		BatchID:   batchId,
+		TxHash:    hash,
+		Status:    types.FeeBatchStateSent,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to get fee batch: %w", err)
+	}
+	defer response.Body.Close()
+
+	return nil, nil
+}

internal/verifierapi/verifierapi.go

@@ -53,12 +53,20 @@ func (v *VerifierApi) getAuth(endpoint string) (*http.Response, error) {
 }
 
 func (v *VerifierApi) postAuth(endpoint string, body any) (*http.Response, error) {
+	return v.bodyRequest(endpoint, body, http.MethodPost)
+}
+
+func (v *VerifierApi) putAuth(endpoint string, body any) (*http.Response, error) {
+	return v.bodyRequest(endpoint, body, http.MethodPut)
+}
+
+func (v *VerifierApi) bodyRequest(endpoint string, body any, httpMethod string) (*http.Response, error) {
 	jsonBody, err := json.Marshal(body)
 	if err != nil {
 		return nil, err
 	}
 
-	request, err := http.NewRequest(http.MethodPost, v.url+endpoint, bytes.NewBuffer(jsonBody))
+	request, err := http.NewRequest(httpMethod, v.url+endpoint, bytes.NewBuffer(jsonBody))
 	if err != nil {
 		return nil, err
 	}

plugin/fees/load.go

@@ -7,7 +7,6 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/hibiken/asynq"
-	"github.com/sirupsen/logrus"
 	"github.com/vultisig/plugin/internal/types"
 	vtypes "github.com/vultisig/verifier/types"
 	"golang.org/x/sync/errgroup"
@@ -46,7 +45,11 @@ func (fp *FeePlugin) LoadFees(ctx context.Context, task *asynq.Task) error {
 				return fmt.Errorf("failed to acquire semaphore: %w", err)
 			}
 			defer sem.Release(1)
-			return fp.executeFeeLoading(ctx, feePolicy)
+			err := fp.executeFeeLoading(ctx, feePolicy)
+			if err != nil {
+				fp.logger.WithError(err).WithField("public_key", feePolicy.PublicKey).Error("Failed to execute fee loading")
+			}
+			return err
 		})
 	}
 
@@ -60,80 +63,28 @@ func (fp *FeePlugin) LoadFees(ctx context.Context, task *asynq.Task) error {
 func (fp *FeePlugin) executeFeeLoading(ctx context.Context, feePolicy vtypes.PluginPolicy) error {
 
 	// Get list of fees from the verifier connected to the fee policy
-	feesResponse, err := fp.verifierApi.GetPublicKeysFees(feePolicy.PublicKey)
+	batch, err := fp.verifierApi.CreateFeeBatch(feePolicy.PublicKey)
 	if err != nil {
 		return fmt.Errorf("failed to get plugin policy fees: %w", err)
 	}
 
-	// Early return if no fees to collect
-	if feesResponse.FeesPendingCollection <= 0 {
-		fp.logger.WithField("publicKey", feePolicy.PublicKey).Info("No fees pending collection")
-		return nil
+	if err != nil {
+		return fmt.Errorf("failed to create fee batch: %w", err)
 	}
 
-	// If fees are greater than 0, we need to collect them
-	fp.logger.WithFields(logrus.Fields{
-		"publicKey": feePolicy.PublicKey,
-	}).Info("Fees pending collection: ", feesResponse.FeesPendingCollection)
-
-	checkAmount := 0
-	for _, fee := range feesResponse.Fees {
-		if !fee.Collected {
-			checkAmount += fee.Amount
-		}
-	}
-	if checkAmount != feesResponse.FeesPendingCollection {
-		return fmt.Errorf("fees pending collection amount does not match the sum of the fees")
+	if batch.Amount == 0 || batch.BatchID == uuid.Nil {
+		fp.logger.WithField("public_key", feePolicy.PublicKey).Info("No fees to load")
+		return nil
 	}
 
-	for _, fee := range feesResponse.Fees {
-		if !fee.Collected {
-
-			// Check if the fee has already been loaded and added to a fee run, if so, skip it
-			existingFee, err := fp.db.GetFees(ctx, fee.ID)
-			if err != nil {
-				return fmt.Errorf("failed to get fee: %w", err)
-			}
-			if len(existingFee) > 0 {
-				fp.logger.WithFields(logrus.Fields{
-					"publicKey": feePolicy.PublicKey,
-					"feeId":     fee.ID,
-					"runId":     existingFee[0].FeeRunID,
-				}).Info("Fee already added to a fee run")
-				continue
-			}
-
-			// If the fee hasn't been loaded, look for a draft run and add it to it
-			run, err := fp.db.GetPendingFeeRun(ctx, feePolicy.ID)
-			if err != nil {
-				return fmt.Errorf("failed to get pending fee run: %w", err)
-			}
+	_, err = fp.db.CreateFeeBatch(ctx, nil, types.FeeBatch{
+		ID:        uuid.New(),
+		BatchID:   batch.BatchID,
+		PublicKey: feePolicy.PublicKey,
+		Status:    types.FeeBatchStateDraft,
+		TxHash:    nil,
+		Amount:    uint64(batch.Amount),
+	})
 
-			// If no draft run is found, create a new one and add the fee to it
-			if run == nil {
-				run, err = fp.db.CreateFeeRun(ctx, feePolicy.ID, types.FeeRunStateDraft, fee)
-				if err != nil {
-					return fmt.Errorf("failed to create fee run: %w", err)
-				}
-				fp.logger.WithFields(logrus.Fields{
-					"publicKey": feePolicy.PublicKey,
-					"feeIds":    []uuid.UUID{fee.ID},
-					"runId":     run.ID,
-				}).Info("Fee run created")
-
-				// If a draft run is found, add the fee to it
-			} else {
-				if err := fp.db.CreateFee(ctx, run.ID, fee); err != nil {
-					return fmt.Errorf("failed to create fee: %w", err)
-				}
-				fp.logger.WithFields(logrus.Fields{
-					"publicKey": feePolicy.PublicKey,
-					"feeIds":    []uuid.UUID{fee.ID},
-					"runId":     run.ID,
-				}).Info("Fee added to fee run")
-			}
-		}
-	}
-
-	return nil
+	return err
 }