diff --git a/openmeter/ledger/collector/collection_fbo.go b/openmeter/ledger/collector/collection_fbo.go index 77a49762a4..f29c1124c2 100644 --- a/openmeter/ledger/collector/collection_fbo.go +++ b/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, diff --git a/openmeter/ledger/collector/collection_fbo_test.go b/openmeter/ledger/collector/collection_fbo_test.go index c4dc003419..3a101868bd 100644 --- a/openmeter/ledger/collector/collection_fbo_test.go +++ b/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) diff --git a/openmeter/ledger/customerbalance/calculation.go b/openmeter/ledger/customerbalance/calculation.go index a2b4350d48..9ad4a9bf71 100644 --- a/openmeter/ledger/customerbalance/calculation.go +++ b/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 diff --git a/openmeter/ledger/customerbalance/service.go b/openmeter/ledger/customerbalance/service.go index f408d85677..66399b3334 100644 --- a/openmeter/ledger/customerbalance/service.go +++ b/openmeter/ledger/customerbalance/service.go @@ -1,6 +1,7 @@ package customerbalance import ( + "cmp" "context" "errors" "fmt" @@ -8,6 +9,7 @@ import ( "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,6 +278,11 @@ 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) @@ -282,11 +290,98 @@ func (s *service) GetBalance(ctx context.Context, input GetBalanceServiceInput) 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 diff --git a/openmeter/ledger/customerbalance/service_test.go b/openmeter/ledger/customerbalance/service_test.go index 60a5972f92..6fc4c44304 100644 --- a/openmeter/ledger/customerbalance/service_test.go +++ b/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)