fix(ledger): correct feature-filtered transaction balances

Author: GAlexIHU

openmeter/ledger/customerbalance/expired_loader.go

@@ -31,6 +31,7 @@ func (l *expiredCreditTransactionLoader) Load(ctx context.Context, input creditT
 
 	items := make([]CreditTransaction, 0, len(result.Items))
 	for _, impact := range result.Items {
+		balanceAsOf := impact.BookedAt
 		items = append(items, CreditTransaction{
 			ID:          impact.ID,
 			CreatedAt:   impact.CreatedAt,
@@ -40,6 +41,7 @@ func (l *expiredCreditTransactionLoader) Load(ctx context.Context, input creditT
 			Amount:      impact.Amount,
 			Name:        "Expired credits",
 			Annotations: impact.Annotations,
+			balanceAsOf: &balanceAsOf,
 		})
 	}
 

openmeter/ledger/customerbalance/expired_loader_test.go

@@ -310,10 +310,10 @@ func TestListCreditTransactionsFeatureFilter(t *testing.T) {
 		FeatureFilter: NewFeatureFilter([]string{testFeatureKey}),
 	})
 	require.NoError(t, err)
-	requireCreditTransactionEvents(t, issuedAt, beforeConsumption.Items, []expectedCreditTransaction{
-		{txType: CreditTransactionTypeFunded, bookedAfter: 2 * time.Minute, amount: 20},
-		{txType: CreditTransactionTypeFunded, bookedAfter: time.Minute, amount: 10},
-		{txType: CreditTransactionTypeFunded, bookedAfter: 0, amount: 100},
+	requireCreditTransactions(t, issuedAt, beforeConsumption.Items, []expectedCreditTransaction{
+		{txType: CreditTransactionTypeFunded, bookedAfter: 2 * time.Minute, amount: 20, balanceBefore: 110, balanceAfter: 130},
+		{txType: CreditTransactionTypeFunded, bookedAfter: time.Minute, amount: 10, balanceBefore: 100, balanceAfter: 110},
+		{txType: CreditTransactionTypeFunded, bookedAfter: 0, amount: 100, balanceBefore: 0, balanceAfter: 100},
 	})
 
 	consumedAt := issuedAt.Add(time.Hour)

openmeter/ledger/customerbalance/funded_loader.go

@@ -2,6 +2,7 @@ package customerbalance
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/creditpurchase"
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/meta"
@@ -33,20 +34,26 @@ func (l *fundedCreditTransactionLoader) Load(ctx context.Context, input creditTr
 
 	items := make([]CreditTransaction, 0, len(result.Items))
 	for _, activity := range result.Items {
+		balanceCursor, err := l.balanceCursorForFundedActivity(ctx, input.CustomerID.Namespace, activity)
+		if err != nil {
+			return creditTransactionLoaderResult{}, err
+		}
+
 		annotations := models.Annotations{
 			ledger.AnnotationChargeID: activity.ChargeID.ID,
 		}
 
 		items = append(items, CreditTransaction{
-			ID:          models.NamespacedID(activity.ChargeID),
-			CreatedAt:   activity.ChargeCreatedAt,
-			BookedAt:    activity.FundedAt,
-			Type:        CreditTransactionTypeFunded,
-			Currency:    activity.Currency,
-			Amount:      activity.Amount,
-			Name:        activity.Name,
-			Description: activity.Description,
-			Annotations: annotations,
+			ID:            models.NamespacedID(activity.ChargeID),
+			CreatedAt:     activity.ChargeCreatedAt,
+			BookedAt:      activity.FundedAt,
+			Type:          CreditTransactionTypeFunded,
+			Currency:      activity.Currency,
+			Amount:        activity.Amount,
+			Name:          activity.Name,
+			Description:   activity.Description,
+			Annotations:   annotations,
+			balanceCursor: balanceCursor,
 		})
 	}
 
@@ -56,6 +63,38 @@ func (l *fundedCreditTransactionLoader) Load(ctx context.Context, input creditTr
 	}, nil
 }
 
+func (l *fundedCreditTransactionLoader) balanceCursorForFundedActivity(
+	ctx context.Context,
+	namespace string,
+	activity creditpurchase.FundedCreditActivity,
+) (*ledger.TransactionCursor, error) {
+	if activity.TransactionGroupID == "" {
+		return nil, nil
+	}
+
+	group, err := l.service.Ledger.GetTransactionGroup(ctx, models.NamespacedID{
+		Namespace: namespace,
+		ID:        activity.TransactionGroupID,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("get funded credit transaction group %s: %w", activity.TransactionGroupID, err)
+	}
+
+	for _, tx := range group.Transactions() {
+		impact, currency, err := creditTransactionFBOImpact(tx)
+		if err != nil {
+			continue
+		}
+
+		if currency == activity.Currency && impact.Equal(activity.Amount) {
+			cursor := tx.Cursor()
+			return &cursor, nil
+		}
+	}
+
+	return nil, fmt.Errorf("funded credit transaction group %s has no matching customer FBO transaction", activity.TransactionGroupID)
+}
+
 func toFundedCreditActivityCursor(cursor *ledger.TransactionCursor) *creditpurchase.FundedCreditActivityCursor {
 	if cursor == nil {
 		return nil

openmeter/ledger/customerbalance/transactions.go

@@ -111,6 +111,9 @@ type CreditTransaction struct {
 	Name        string
 	Description *string
 	Annotations models.Annotations
+
+	balanceCursor *ledger.TransactionCursor
+	balanceAsOf   *time.Time
 }
 
 type CreditTransactionBalance struct {
@@ -182,9 +185,7 @@ func (s *service) ListCreditTransactions(ctx context.Context, input ListCreditTr
 			CustomerID:    input.CustomerID,
 			Currency:      items[0].Currency,
 			FeatureFilter: normalizeFeatureFilter(input.FeatureFilter),
-			BalanceQuery: ledger.BalanceQuery{
-				After: lo.ToPtr(creditTransactionCursor(items[0])),
-			},
+			BalanceQuery:  items[0].balanceQuery(),
 		})
 		if err != nil {
 			return ListCreditTransactionsResult{}, fmt.Errorf("get FBO balance after transaction %s: %w", items[0].ID.ID, err)
@@ -264,16 +265,18 @@ func creditTransactionFromLedgerTransaction(tx ledger.Transaction) (CreditTransa
 	if err != nil {
 		return CreditTransaction{}, err
 	}
+	cursor := tx.Cursor()
 
 	return CreditTransaction{
-		ID:          tx.ID(),
-		CreatedAt:   tx.Cursor().CreatedAt,
-		BookedAt:    tx.BookedAt(),
-		Type:        creditTransactionType(fboImpact),
-		Currency:    currency,
-		Amount:      fboImpact,
-		Name:        "",
-		Annotations: tx.Annotations(),
+		ID:            tx.ID(),
+		CreatedAt:     tx.Cursor().CreatedAt,
+		BookedAt:      tx.BookedAt(),
+		Type:          creditTransactionType(fboImpact),
+		Currency:      currency,
+		Amount:        fboImpact,
+		Name:          "",
+		Annotations:   tx.Annotations(),
+		balanceCursor: &cursor,
 	}, nil
 }
 
@@ -314,6 +317,19 @@ func applyCreditTransactionBalances(items []CreditTransaction, after alpacadecim
 	}
 }
 
+func (tx CreditTransaction) balanceQuery() ledger.BalanceQuery {
+	if tx.balanceCursor != nil {
+		return ledger.BalanceQuery{After: tx.balanceCursor}
+	}
+
+	if tx.balanceAsOf != nil {
+		return ledger.BalanceQuery{AsOf: tx.balanceAsOf}
+	}
+
+	asOf := tx.BookedAt
+	return ledger.BalanceQuery{AsOf: &asOf}
+}
+
 func creditTransactionType(fboImpact alpacadecimal.Decimal) CreditTransactionType {
 	if fboImpact.IsPositive() {
 		return CreditTransactionTypeFunded

openmeter/ledger/primitives.go

@@ -274,35 +274,14 @@ func (i ListTransactionsInput) Validate() error {
 		})
 	}
 
-	if i.Route.Features.IsPresent() && i.Route.MatchFeature != "" {
+	if err := validateListTransactionsRouteFilter(i.Route); err != nil {
 		return ErrListTransactionsInputInvalid.WithAttrs(models.Attributes{
 			"reason": "route_invalid",
 			"route":  i.Route,
-			"error":  errors.New("features and match feature filters cannot be combined"),
+			"error":  err,
 		})
 	}
 
-	if i.Route.Features.IsPresent() {
-		features, _ := i.Route.Features.Get()
-		if err := validateFeatures(features); err != nil {
-			return ErrListTransactionsInputInvalid.WithAttrs(models.Attributes{
-				"reason": "route_invalid",
-				"route":  i.Route,
-				"error":  fmt.Errorf("features: %w", err),
-			})
-		}
-	}
-
-	if i.Route.MatchFeature != "" {
-		if err := validateFeatures([]string{i.Route.MatchFeature}); err != nil {
-			return ErrListTransactionsInputInvalid.WithAttrs(models.Attributes{
-				"reason": "route_invalid",
-				"route":  i.Route,
-				"error":  fmt.Errorf("match feature: %w", err),
-			})
-		}
-	}
-
 	switch i.CreditMovement {
 	case ListTransactionsCreditMovementUnspecified, ListTransactionsCreditMovementPositive, ListTransactionsCreditMovementNegative:
 	default:
@@ -315,6 +294,49 @@ func (i ListTransactionsInput) Validate() error {
 	return nil
 }
 
+func validateListTransactionsRouteFilter(route RouteFilter) error {
+	var errs []error
+
+	if route.TaxCode.IsPresent() {
+		errs = append(errs, errors.New("tax code filter is not supported"))
+	}
+
+	if route.TaxBehavior.IsPresent() {
+		errs = append(errs, errors.New("tax behavior filter is not supported"))
+	}
+
+	if route.CostBasis.IsPresent() {
+		errs = append(errs, errors.New("cost basis filter is not supported"))
+	}
+
+	if route.CreditPriority != nil {
+		errs = append(errs, errors.New("credit priority filter is not supported"))
+	}
+
+	if route.TransactionAuthorizationStatus != nil {
+		errs = append(errs, errors.New("transaction authorization status filter is not supported"))
+	}
+
+	if route.Features.IsPresent() && route.MatchFeature != "" {
+		errs = append(errs, errors.New("features and match feature filters cannot be combined"))
+	}
+
+	if route.Features.IsPresent() {
+		features, _ := route.Features.Get()
+		if err := validateFeatures(features); err != nil {
+			errs = append(errs, fmt.Errorf("features: %w", err))
+		}
+	}
+
+	if route.MatchFeature != "" {
+		if err := validateFeatures([]string{route.MatchFeature}); err != nil {
+			errs = append(errs, fmt.Errorf("match feature: %w", err))
+		}
+	}
+
+	return errors.Join(errs...)
+}
+
 type ListTransactionsCreditMovement uint8
 
 const (

openmeter/ledger/validations_test.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/alpacahq/alpacadecimal"
+	"github.com/samber/mo"
 	"github.com/stretchr/testify/require"
 
 	"github.com/openmeterio/openmeter/openmeter/ledger"
@@ -91,6 +92,100 @@ func TestValidateTransactionInputEntryAmountPrecision(t *testing.T) {
 	}
 }
 
+func TestListTransactionsInputValidateRouteFilter(t *testing.T) {
+	costBasis := alpacadecimal.NewFromFloat(0.7)
+	taxCode := "vat"
+	taxBehavior := ledger.TaxBehaviorInclusive
+	creditPriority := 1
+	authStatus := ledger.TransactionAuthorizationStatusAuthorized
+
+	tests := []struct {
+		name    string
+		route   ledger.RouteFilter
+		wantErr bool
+	}{
+		{
+			name: "currency route filter is supported",
+			route: ledger.RouteFilter{
+				Currency: currencyx.Code("USD"),
+			},
+		},
+		{
+			name: "exact features route filter is supported",
+			route: ledger.RouteFilter{
+				Features: mo.Some([]string{"feature-a"}),
+			},
+		},
+		{
+			name: "match feature route filter is supported",
+			route: ledger.RouteFilter{
+				MatchFeature: "feature-a",
+			},
+		},
+		{
+			name: "cost basis route filter is rejected",
+			route: ledger.RouteFilter{
+				CostBasis: mo.Some(&costBasis),
+			},
+			wantErr: true,
+		},
+		{
+			name: "tax code route filter is rejected",
+			route: ledger.RouteFilter{
+				TaxCode: mo.Some(&taxCode),
+			},
+			wantErr: true,
+		},
+		{
+			name: "tax behavior route filter is rejected",
+			route: ledger.RouteFilter{
+				TaxBehavior: mo.Some(&taxBehavior),
+			},
+			wantErr: true,
+		},
+		{
+			name: "credit priority route filter is rejected",
+			route: ledger.RouteFilter{
+				CreditPriority: &creditPriority,
+			},
+			wantErr: true,
+		},
+		{
+			name: "transaction authorization route filter is rejected",
+			route: ledger.RouteFilter{
+				TransactionAuthorizationStatus: &authStatus,
+			},
+			wantErr: true,
+		},
+		{
+			name: "exact features and match feature cannot be combined",
+			route: ledger.RouteFilter{
+				Features:     mo.Some([]string{"feature-a"}),
+				MatchFeature: "feature-a",
+			},
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := ledger.ListTransactionsInput{
+				Namespace: "ns-test",
+				Limit:     1,
+				Route:     tt.route,
+			}.Validate()
+
+			if tt.wantErr {
+				require.Error(t, err)
+				require.ErrorIs(t, err, ledger.ErrListTransactionsInputInvalid)
+				return
+			}
+
+			require.NoError(t, err)
+		})
+	}
+}
+
 func mustPostingAddress(t *testing.T, currency currencyx.Code) ledger.PostingAddress {
 	t.Helper()