fix: properly use feature filters in charge impact calculations

GAlexIHU · GAlexIHU · commit 90e752f4105e · 2026-06-17T16:18:53.000+02:00
diff --git a/openmeter/ledger/customerbalance/service.go b/openmeter/ledger/customerbalance/service.go
@@ -232,7 +232,7 @@ func (s *service) GetBalance(ctx context.Context, input GetBalanceServiceInput)
 
 	// Pending balance remains a current projection from open charges.
 	// Historical cursor/as-of filtering only affects the booked/settled side for now.
-	impacts, err := s.getChargePendingBalanceImpacts(ctx, input.CustomerID, input.Currency)
+	impacts, err := s.getChargePendingBalanceImpacts(ctx, input.CustomerID, input.Currency, normalizeFeatureFilter(input.FeatureFilter))
 	if err != nil {
 		return nil, fmt.Errorf("get charge pending balance impacts: %w", err)
 	}
@@ -275,7 +275,7 @@ func (s *service) GetFBOCurrencies(ctx context.Context, customerID customer.Cust
 	return codes, nil
 }
 
-func (s *service) getChargePendingBalanceImpacts(ctx context.Context, customerID customer.CustomerID, currency currencyx.Code) ([]Impact, error) {
+func (s *service) getChargePendingBalanceImpacts(ctx context.Context, customerID customer.CustomerID, currency currencyx.Code, featureFilter mo.Option[creditpurchase.FeatureFilters]) ([]Impact, error) {
 	items, err := pagination.CollectAll(
 		ctx,
 		pagination.NewPaginator(func(ctx context.Context, page pagination.Page) (pagination.Result[charges.Charge], error) {
@@ -299,7 +299,7 @@ func (s *service) getChargePendingBalanceImpacts(ctx context.Context, customerID
 
 	impacts := make([]Impact, 0, len(items))
 	for _, charge := range items {
-		impact, err := s.getChargePendingBalanceImpact(ctx, charge, currency)
+		impact, err := s.getChargePendingBalanceImpact(ctx, charge, currency, featureFilter)
 		if err != nil {
 			return nil, err
 		}
@@ -314,22 +314,22 @@ func (s *service) getChargePendingBalanceImpacts(ctx context.Context, customerID
 	return impacts, nil
 }
 
-func (s *service) getChargePendingBalanceImpact(ctx context.Context, charge charges.Charge, currency currencyx.Code) (*Impact, error) {
+func (s *service) getChargePendingBalanceImpact(ctx context.Context, charge charges.Charge, currency currencyx.Code, featureFilter mo.Option[creditpurchase.FeatureFilters]) (*Impact, error) {
 	if !chargeHasStarted(charge) {
 		return nil, nil
 	}
 
 	switch charge.Type() {
 	case meta.ChargeTypeFlatFee:
-		return getFlatFeeChargePendingBalanceImpact(charge, currency)
+		return getFlatFeeChargePendingBalanceImpact(charge, currency, featureFilter)
 	case meta.ChargeTypeUsageBased:
-		return s.getUsageBasedChargePendingBalanceImpact(ctx, charge, currency)
+		return s.getUsageBasedChargePendingBalanceImpact(ctx, charge, currency, featureFilter)
 	default:
 		return nil, nil
 	}
 }
 
-func getFlatFeeChargePendingBalanceImpact(charge charges.Charge, currency currencyx.Code) (*Impact, error) {
+func getFlatFeeChargePendingBalanceImpact(charge charges.Charge, currency currencyx.Code, featureFilter mo.Option[creditpurchase.FeatureFilters]) (*Impact, error) {
 	flatFeeCharge, err := charge.AsFlatFeeCharge()
 	if err != nil {
 		return nil, fmt.Errorf("map flat fee charge: %w", err)
@@ -339,10 +339,14 @@ func getFlatFeeChargePendingBalanceImpact(charge charges.Charge, currency curren
 		return nil, nil
 	}
 
+	if !featureFilterMatchesChargeFeatureKey(featureFilter, flatFeeCharge.Intent.FeatureKey) {
+		return nil, nil
+	}
+
 	return newImpactOrNil(charge, flatFeeCharge.State.AmountAfterProration)
 }
 
-func (s *service) getUsageBasedChargePendingBalanceImpact(ctx context.Context, charge charges.Charge, currency currencyx.Code) (*Impact, error) {
+func (s *service) getUsageBasedChargePendingBalanceImpact(ctx context.Context, charge charges.Charge, currency currencyx.Code, featureFilter mo.Option[creditpurchase.FeatureFilters]) (*Impact, error) {
 	usageBasedCharge, err := charge.AsUsageBasedCharge()
 	if err != nil {
 		return nil, fmt.Errorf("map usage based charge: %w", err)
@@ -352,6 +356,10 @@ func (s *service) getUsageBasedChargePendingBalanceImpact(ctx context.Context, c
 		return nil, nil
 	}
 
+	if !featureFilterMatchesChargeFeatureKey(featureFilter, usageBasedCharge.Intent.FeatureKey) {
+		return nil, nil
+	}
+
 	currentTotals, err := s.UsageBasedService.GetCurrentTotals(ctx, usagebased.GetCurrentTotalsInput{
 		ChargeID: usageBasedCharge.GetChargeID(),
 	})
@@ -362,6 +370,23 @@ func (s *service) getUsageBasedChargePendingBalanceImpact(ctx context.Context, c
 	return newImpactOrNil(charges.NewCharge(currentTotals.Charge), currentTotals.DueTotals.Total)
 }
 
+func featureFilterMatchesChargeFeatureKey(featureFilter mo.Option[creditpurchase.FeatureFilters], featureKey string) bool {
+	if featureFilter.IsAbsent() {
+		return true
+	}
+
+	features := featureFilter.OrEmpty()
+	if features == nil {
+		return featureKey == ""
+	}
+
+	if featureKey == "" {
+		return true
+	}
+
+	return len(features) == 1 && features[0] == featureKey
+}
+
 func chargeHasStarted(charge charges.Charge) bool {
 	now := clock.Now()
 
diff --git a/openmeter/ledger/customerbalance/service_test.go b/openmeter/ledger/customerbalance/service_test.go
@@ -460,3 +460,65 @@ func TestGetBalanceFeatureFilter(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, float64(10), exactFeatureABalance.Settled().InexactFloat64())
 }
+
+func TestGetBalanceFeatureFilterPendingChargeImpacts(t *testing.T) {
+	env := newTestEnv(t)
+
+	env.bookFBOBalanceWithFeatures(t, alpacadecimal.NewFromInt(100), nil)
+	env.fundOpenReceivableWithFeatures(t, alpacadecimal.NewFromInt(100), nil)
+	env.bookFBOBalanceWithFeatures(t, alpacadecimal.NewFromInt(10), []string{testFeatureKey})
+	env.fundOpenReceivableWithFeatures(t, alpacadecimal.NewFromInt(10), []string{testFeatureKey})
+	env.bookFBOBalanceWithFeatures(t, alpacadecimal.NewFromInt(20), []string{"storage"})
+	env.fundOpenReceivableWithFeatures(t, alpacadecimal.NewFromInt(20), []string{"storage"})
+
+	env.addUsage(30, clock.Now().Add(-30*time.Minute))
+	env.createFlatFeeCharge(t, alpacadecimal.NewFromInt(5), productcatalog.CreditOnlySettlementMode, env.sp())
+	env.createFlatFeeCharge(t, alpacadecimal.NewFromInt(7), productcatalog.CreditOnlySettlementMode, env.sp(), testFeatureKey)
+	env.createUsageBasedCharge(t, alpacadecimal.NewFromInt(1), productcatalog.CreditOnlySettlementMode, env.sp())
+
+	tests := []struct {
+		name        string
+		filter      mo.Option[creditpurchase.FeatureFilters]
+		wantSettled float64
+		wantPending float64
+	}{
+		{
+			name:        "omitted filter includes every charge impact",
+			filter:      AllFeatureFilter(),
+			wantSettled: 130,
+			wantPending: 88,
+		},
+		{
+			name:        "unrestricted filter includes only unrestricted charge impacts",
+			filter:      NewUnrestrictedFeatureFilter(),
+			wantSettled: 100,
+			wantPending: 95,
+		},
+		{
+			name:        "matching feature filter includes unrestricted and matching charge impacts",
+			filter:      NewFeatureFilter([]string{testFeatureKey}),
+			wantSettled: 110,
+			wantPending: 68,
+		},
+		{
+			name:        "non-matching feature filter excludes restricted charge impacts for other features",
+			filter:      NewFeatureFilter([]string{"storage"}),
+			wantSettled: 120,
+			wantPending: 115,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			balance, err := env.Service.GetBalance(t.Context(), GetBalanceServiceInput{
+				CustomerID:    env.CustomerID,
+				Currency:      env.Currency,
+				FeatureFilter: tt.filter,
+			})
+			require.NoError(t, err)
+
+			require.Equal(t, tt.wantSettled, balance.Settled().InexactFloat64())
+			require.Equal(t, tt.wantPending, balance.Pending().InexactFloat64())
+		})
+	}
+}
diff --git a/openmeter/ledger/customerbalance/testenv_test.go b/openmeter/ledger/customerbalance/testenv_test.go
@@ -447,15 +447,23 @@ func (e *testEnv) createUsageBasedChargeInCurrency(t *testing.T, unitPrice alpac
 	return createdCharges[0].Charge
 }
 
-func (e *testEnv) createFlatFeeCharge(t *testing.T, amount alpacadecimal.Decimal, settlementMode productcatalog.SettlementMode, servicePeriod timeutil.ClosedPeriod) flatfee.Charge {
-	return e.createFlatFeeChargeInCurrency(t, amount, settlementMode, servicePeriod, e.Currency)
+func (e *testEnv) createFlatFeeCharge(t *testing.T, amount alpacadecimal.Decimal, settlementMode productcatalog.SettlementMode, servicePeriod timeutil.ClosedPeriod, featureKeys ...string) flatfee.Charge {
+	return e.createFlatFeeChargeInCurrency(t, amount, settlementMode, servicePeriod, e.Currency, featureKeys...)
 }
 
-func (e *testEnv) createFlatFeeChargeInCurrency(t *testing.T, amount alpacadecimal.Decimal, settlementMode productcatalog.SettlementMode, servicePeriod timeutil.ClosedPeriod, currency currencyx.Code) flatfee.Charge {
+func (e *testEnv) createFlatFeeChargeInCurrency(t *testing.T, amount alpacadecimal.Decimal, settlementMode productcatalog.SettlementMode, servicePeriod timeutil.ClosedPeriod, currency currencyx.Code, featureKeys ...string) flatfee.Charge {
 	t.Helper()
 
+	require.LessOrEqual(t, len(featureKeys), 1)
+
+	featureKey := ""
+	if len(featureKeys) == 1 {
+		featureKey = featureKeys[0]
+	}
+
 	createdCharges, err := e.flatFeeService.Create(t.Context(), flatfee.CreateInput{
-		Namespace: e.Namespace,
+		Namespace:     e.Namespace,
+		FeatureMeters: e.featureMeters,
 		Intents: []flatfee.Intent{
 			{
 				Intent: chargemeta.Intent{
@@ -470,6 +478,7 @@ func (e *testEnv) createFlatFeeChargeInCurrency(t *testing.T, amount alpacadecim
 				InvoiceAt:             e.Now().Add(-time.Minute),
 				SettlementMode:        settlementMode,
 				PaymentTerm:           productcatalog.InAdvancePaymentTerm,
+				FeatureKey:            featureKey,
 				AmountBeforeProration: amount,
 			},
 		},
