fix(ledger): honor feature-restricted balance sources

openmeter/ledger/collector/collection_fbo.go

@@ -204,6 +204,10 @@ func (c *accrualCollector) listCustomerFBOSources(
 		planCopy := plan
 		expiresAt := plan.ExpiresAt
 		route := plan.FBOAddress.Route().Route()
+		if len(route.Features) > 0 && !lo.Contains(route.Features, featureKey) {
+			continue
+		}
+
 		breakageSources = append(breakageSources, fboCollectionSource{
 			address:           plan.FBOAddress,
 			available:         available,

openmeter/ledger/collector/collection_fbo_test.go

@@ -118,6 +118,46 @@ func TestCollectCustomerFBOFiltersByFeatureEligibility(t *testing.T) {
 	require.True(t, alpacadecimal.NewFromInt(10).Equal(unattributedSources[0].Amount))
 }
 
+func TestCollectCustomerFBOFiltersBreakageByFeatureEligibility(t *testing.T) {
+	env := ledgertestutils.NewIntegrationEnv(t, "collector")
+	breakageService := newTestBreakageService(t, env)
+	collector := newTestAccrualCollectorWithBreakage(env, breakageService)
+
+	expiresAt := env.Now().Add(10 * time.Hour)
+	bookExpiringCreditWithFeatures(t, env, breakageService, 1, 10, nil, expiresAt)
+	bookExpiringCreditWithFeatures(t, env, breakageService, 1, 30, []string{"api-calls"}, expiresAt)
+	bookExpiringCreditWithFeatures(t, env, breakageService, 1, 40, []string{"storage"}, expiresAt)
+
+	sources, err := collectCustomerFBOForFeatureForTest(
+		t,
+		env,
+		collector,
+		"api-calls",
+		alpacadecimal.NewFromInt(200),
+		env.Now(),
+	)
+	require.NoError(t, err)
+	require.Len(t, sources, 2)
+
+	require.Equal(t, []string{"api-calls"}, sources[0].Address.Route().Route().Features)
+	require.True(t, alpacadecimal.NewFromInt(30).Equal(sources[0].Amount), "restricted source amount: %s", sources[0].Amount)
+	require.Empty(t, sources[1].Address.Route().Route().Features)
+	require.True(t, alpacadecimal.NewFromInt(10).Equal(sources[1].Amount), "unrestricted source amount: %s", sources[1].Amount)
+
+	unattributedSources, err := collectCustomerFBOForFeatureForTest(
+		t,
+		env,
+		collector,
+		"",
+		alpacadecimal.NewFromInt(200),
+		env.Now(),
+	)
+	require.NoError(t, err)
+	require.Len(t, unattributedSources, 1)
+	require.Empty(t, unattributedSources[0].Address.Route().Route().Features)
+	require.True(t, alpacadecimal.NewFromInt(10).Equal(unattributedSources[0].Amount), "unrestricted source amount: %s", unattributedSources[0].Amount)
+}
+
 func TestCollectCustomerFBOUsesPriorityBeforeFeatureRestriction(t *testing.T) {
 	env := ledgertestutils.NewIntegrationEnv(t, "collector")
 	collector := newTestAccrualCollector(env)

openmeter/ledger/customerbalance/calculation.go

@@ -1,6 +1,8 @@
 package customerbalance
 
 import (
+	"slices"
+
 	"github.com/alpacahq/alpacadecimal"
 
 	"github.com/openmeterio/openmeter/openmeter/billing/charges"
@@ -85,6 +87,19 @@ func (i Impact) UnboundedAmount() alpacadecimal.Decimal {
 	return i.OutstandingAmount()
 }
 
+func (i Impact) FeatureKey() string {
+	switch i.Type() {
+	case meta.ChargeTypeFlatFee:
+		charge, _ := i.AsFlatFeeCharge()
+		return charge.Intent.GetEffectiveFeatureKey()
+	case meta.ChargeTypeUsageBased:
+		charge, _ := i.AsUsageBasedCharge()
+		return charge.Intent.GetEffectiveFeatureKey()
+	default:
+		return ""
+	}
+}
+
 type chargeLiveBalanceCalculator struct{}
 
 func (chargeLiveBalanceCalculator) CalculateLiveBalance(bookedBalance alpacadecimal.Decimal, impacts []Impact) alpacadecimal.Decimal {
@@ -96,6 +111,59 @@ func (chargeLiveBalanceCalculator) CalculateLiveBalance(bookedBalance alpacadeci
 	return liveBalance.Sub(unboundedAmount)
 }
 
+func (chargeLiveBalanceCalculator) CalculateLiveBalanceFromSources(settledBalance alpacadecimal.Decimal, sources []liveBalanceSource, impacts []Impact) alpacadecimal.Decimal {
+	liveBalance := settledBalance
+
+	for _, impact := range impacts {
+		if boundedAmount := impact.BoundedAmount(); boundedAmount.IsPositive() {
+			liveBalance = liveBalance.Sub(consumeLiveBalanceSources(sources, impact.FeatureKey(), boundedAmount))
+		}
+
+		// credit_only can create feature-attributed advance/negative balance, so
+		// it still changes live balance even when no positive eligible source
+		// exists for the impact.
+		liveBalance = liveBalance.Sub(impact.UnboundedAmount())
+	}
+
+	return liveBalance
+}
+
+func consumeLiveBalanceSources(sources []liveBalanceSource, featureKey string, target alpacadecimal.Decimal) alpacadecimal.Decimal {
+	remaining := target
+	consumed := alpacadecimal.Zero
+
+	for idx := range sources {
+		if !liveBalanceSourceMatchesFeature(sources[idx], featureKey) {
+			continue
+		}
+
+		amount := sources[idx].amount
+		if amount.GreaterThan(remaining) {
+			amount = remaining
+		}
+
+		sources[idx].amount = sources[idx].amount.Sub(amount)
+		remaining = remaining.Sub(amount)
+		consumed = consumed.Add(amount)
+		if remaining.IsZero() {
+			break
+		}
+	}
+
+	return consumed
+}
+
+// liveBalanceSourceMatchesFeature is allocability matching, not public balance
+// filter matching. Unrestricted credit sources can cover any charge, but
+// feature-restricted sources can only cover charges for that feature.
+func liveBalanceSourceMatchesFeature(source liveBalanceSource, featureKey string) bool {
+	if len(source.route.Features) == 0 {
+		return true
+	}
+
+	return featureKey != "" && slices.Contains(source.route.Features, featureKey)
+}
+
 func sumImpactAmounts(impacts []Impact) (bounded alpacadecimal.Decimal, unbounded alpacadecimal.Decimal) {
 	bounded = alpacadecimal.Zero
 	unbounded = alpacadecimal.Zero

openmeter/ledger/customerbalance/service.go

@@ -1,13 +1,15 @@
 package customerbalance
 
 import (
+	"cmp"
 	"context"
 	"errors"
 	"fmt"
 	"slices"
 	"time"
 
 	"github.com/alpacahq/alpacadecimal"
+	"github.com/samber/lo"
 	"github.com/samber/mo"
 
 	"github.com/openmeterio/openmeter/openmeter/billing/charges"
@@ -18,6 +20,7 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/ledger"
 	ledgerbreakage "github.com/openmeterio/openmeter/openmeter/ledger/breakage"
 	"github.com/openmeterio/openmeter/pkg/clock"
+	"github.com/openmeterio/openmeter/pkg/cmpx"
 	"github.com/openmeterio/openmeter/pkg/currencyx"
 	"github.com/openmeterio/openmeter/pkg/models"
 	"github.com/openmeterio/openmeter/pkg/pagination"
@@ -275,18 +278,110 @@ func (s *service) GetBalance(ctx context.Context, input GetBalanceServiceInput)
 		return nil, fmt.Errorf("get charge live balance impacts: %w", err)
 	}
 
+	live, err := s.calculateLiveBalance(ctx, input, settled, impacts)
+	if err != nil {
+		return nil, fmt.Errorf("calculate live balance: %w", err)
+	}
+
 	pending, err := s.getPendingGrantAmount(ctx, input.CustomerID, input.Currency, normalizeFeatureFilter(input.FeatureFilter), input.pendingGrantAsOf())
 	if err != nil {
 		return nil, fmt.Errorf("get pending grant amount: %w", err)
 	}
 
 	return balance{
 		settled: settled,
-		live:    s.balanceCalculator.CalculateLiveBalance(settled, impacts),
+		live:    live,
 		pending: pending,
 	}, nil
 }
 
+type liveBalanceSource struct {
+	route  ledger.Route
+	amount alpacadecimal.Decimal
+	cursor string
+}
+
+var _ cmpx.Comparable[liveBalanceSource] = liveBalanceSource{}
+
+func (s *service) calculateLiveBalance(ctx context.Context, input GetBalanceServiceInput, settled alpacadecimal.Decimal, impacts []Impact) (alpacadecimal.Decimal, error) {
+	// Live charge impacts must be applied against the same credit sources the
+	// collector could actually consume. An aggregate settled balance would let a
+	// charge for one feature reduce credit restricted to another feature, even
+	// though the eventual ledger collection would leave that credit untouched.
+	sources, err := s.getLiveBalanceSources(ctx, input)
+	if err != nil {
+		return alpacadecimal.Zero, err
+	}
+
+	return s.balanceCalculator.CalculateLiveBalanceFromSources(settled, sources, impacts), nil
+}
+
+func (s *service) getLiveBalanceSources(ctx context.Context, input GetBalanceServiceInput) ([]liveBalanceSource, error) {
+	customerAccounts, err := s.AccountResolver.GetCustomerAccounts(ctx, input.CustomerID)
+	if err != nil {
+		return nil, fmt.Errorf("get customer accounts: %w", err)
+	}
+
+	subAccounts, err := s.SubAccountService.ListSubAccounts(ctx, ledger.ListSubAccountsInput{
+		Namespace: customerAccounts.FBOAccount.ID().Namespace,
+		AccountID: customerAccounts.FBOAccount.ID().ID,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("list sub accounts: %w", err)
+	}
+
+	routeFilter := input.bookedRoute()
+	query := input.balanceQuery()
+	sources := make([]liveBalanceSource, 0, len(subAccounts))
+	for _, subAccount := range subAccounts {
+		route := subAccount.Route()
+		if !route.Matches(routeFilter) {
+			continue
+		}
+
+		sourceBalance, err := s.BalanceQuerier.GetSubAccountBalance(ctx, subAccount, query)
+		if err != nil {
+			return nil, fmt.Errorf("get sub account balance: %w", err)
+		}
+
+		if !sourceBalance.IsPositive() {
+			continue
+		}
+
+		sources = append(sources, liveBalanceSource{
+			route:  route,
+			amount: sourceBalance,
+			cursor: subAccount.Address().SubAccountID(),
+		})
+	}
+
+	slices.SortStableFunc(sources, cmpx.Compare[liveBalanceSource])
+
+	return sources, nil
+}
+
+// Compare keeps the source walk aligned with collection order. This matters for
+// live balance because source amounts are consumed in-memory as impacts are
+// applied, so a shared unrestricted source exhausted by one impact must not be
+// counted again for a later impact.
+func (s liveBalanceSource) Compare(other liveBalanceSource) int {
+	if c := cmp.Compare(lo.FromPtrOr(s.route.CreditPriority, ledger.DefaultCustomerFBOPriority), lo.FromPtrOr(other.route.CreditPriority, ledger.DefaultCustomerFBOPriority)); c != 0 {
+		return c
+	}
+
+	leftRestricted := len(s.route.Features) > 0
+	rightRestricted := len(other.route.Features) > 0
+	if leftRestricted != rightRestricted {
+		if leftRestricted {
+			return -1
+		}
+
+		return 1
+	}
+
+	return cmp.Compare(s.cursor, other.cursor)
+}
+
 func (s *service) GetSettledBalance(ctx context.Context, input GetBalanceServiceInput) (alpacadecimal.Decimal, error) {
 	if err := input.Validate(); err != nil {
 		return alpacadecimal.Zero, err
@@ -628,6 +723,10 @@ func (s *service) getUsageBasedChargePendingBalanceImpact(ctx context.Context, c
 	return newImpactOrNil(charges.NewCharge(currentTotals.Charge), currentTotals.DueTotals.Total)
 }
 
+// featureFilterMatchesChargeFeatureKey is query-scope matching: a feature
+// balance view includes unrestricted charge impacts so it can show the customer's
+// shared-credit exposure for that feature. Actual credit allocability is checked
+// separately when live impacts are applied to concrete credit sources.
 func featureFilterMatchesChargeFeatureKey(featureFilter mo.Option[creditpurchase.FeatureFilters], featureKey string) bool {
 	if featureFilter.IsAbsent() {
 		return true

openmeter/ledger/customerbalance/service_test.go

@@ -533,6 +533,43 @@ func TestGetBalanceFeatureFilterPendingChargeImpacts(t *testing.T) {
 	}
 }
 
+func TestGetBalanceAllFeatureFilterDoesNotApplyBoundedUsageToIneligibleRestrictedCredits(t *testing.T) {
+	env := newTestEnv(t)
+
+	env.bookFBOBalanceWithFeatures(t, alpacadecimal.NewFromInt(100), []string{"storage"})
+	env.fundOpenReceivableWithFeatures(t, alpacadecimal.NewFromInt(100), []string{"storage"})
+
+	env.addUsage(30, clock.Now().Add(-30*time.Minute))
+	env.createUsageBasedCharge(t, alpacadecimal.NewFromInt(1), productcatalog.CreditThenInvoiceSettlementMode, env.sp())
+
+	allBalance, err := env.Service.GetBalance(t.Context(), GetBalanceServiceInput{
+		CustomerID:    env.CustomerID,
+		Currency:      env.Currency,
+		FeatureFilter: AllFeatureFilter(),
+	})
+	require.NoError(t, err)
+	require.Equal(t, float64(100), allBalance.Settled().InexactFloat64())
+	require.Equal(t, float64(100), allBalance.Live().InexactFloat64())
+
+	storageBalance, err := env.Service.GetBalance(t.Context(), GetBalanceServiceInput{
+		CustomerID:    env.CustomerID,
+		Currency:      env.Currency,
+		FeatureFilter: NewFeatureFilter([]string{"storage"}),
+	})
+	require.NoError(t, err)
+	require.Equal(t, float64(100), storageBalance.Settled().InexactFloat64())
+	require.Equal(t, float64(100), storageBalance.Live().InexactFloat64())
+
+	apiRequestsBalance, err := env.Service.GetBalance(t.Context(), GetBalanceServiceInput{
+		CustomerID:    env.CustomerID,
+		Currency:      env.Currency,
+		FeatureFilter: NewFeatureFilter([]string{testFeatureKey}),
+	})
+	require.NoError(t, err)
+	require.Equal(t, float64(0), apiRequestsBalance.Settled().InexactFloat64())
+	require.Equal(t, float64(0), apiRequestsBalance.Live().InexactFloat64())
+}
+
 func TestGetBalancePendingGrants(t *testing.T) {
 	env := newTestEnv(t)