fix(billing): support credit-only period patches (#4597)

Author: turip

openmeter/billing/charges/flatfee/service/creditsonly.go

@@ -4,11 +4,14 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/alpacahq/alpacadecimal"
 	"github.com/samber/lo"
+	"github.com/samber/mo"
 
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/flatfee"
 	flatfeerealizations "github.com/openmeterio/openmeter/openmeter/billing/charges/flatfee/service/realizations"
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/meta"
+	"github.com/openmeterio/openmeter/openmeter/billing/models/totals"
 	"github.com/openmeterio/openmeter/openmeter/productcatalog"
 	"github.com/openmeterio/openmeter/pkg/clock"
 	"github.com/openmeterio/openmeter/pkg/statelessx"
@@ -44,19 +47,25 @@ func (s *CreditsOnlyStateMachine) configureStates() {
 	s.Configure(flatfee.StatusCreated).
 		Permit(meta.TriggerNext, flatfee.StatusActive, statelessx.BoolFn(s.IsInsideServicePeriod)).
 		InternalTransition(meta.TriggerDelete, statelessx.WithParameters(s.DeleteCharge)).
+		InternalTransition(meta.TriggerExtend, statelessx.WithParameters(s.ExtendCharge)).
+		InternalTransition(meta.TriggerShrink, statelessx.WithParameters(s.ShrinkCharge)).
 		OnActive(
 			s.AdvanceAfterServicePeriodFrom,
 		)
 
 	s.Configure(flatfee.StatusActive).
 		Permit(meta.TriggerNext, flatfee.StatusFinal, statelessx.BoolFn(s.IsAfterInvoiceAt)).
 		InternalTransition(meta.TriggerDelete, statelessx.WithParameters(s.DeleteCharge)).
+		InternalTransition(meta.TriggerExtend, statelessx.WithParameters(s.ExtendCharge)).
+		InternalTransition(meta.TriggerShrink, statelessx.WithParameters(s.ShrinkCharge)).
 		OnActive(
 			s.AdvanceAfterInvoiceAt,
 		)
 
 	s.Configure(flatfee.StatusFinal).
 		InternalTransition(meta.TriggerDelete, statelessx.WithParameters(s.DeleteCharge)).
+		InternalTransition(meta.TriggerExtend, statelessx.WithParameters(s.ExtendCharge)).
+		InternalTransition(meta.TriggerShrink, statelessx.WithParameters(s.ShrinkCharge)).
 		OnActive(
 			statelessx.AllOf(
 				s.AllocateCredits,
@@ -102,6 +111,10 @@ func (s *CreditsOnlyStateMachine) AllocateCredits(ctx context.Context) error {
 		}
 	}
 
+	if s.Charge.Realizations.CurrentRun != nil && len(s.Charge.Realizations.CurrentRun.CreditRealizations) > 0 {
+		return s.reconcileCurrentRunCredits(ctx, amount)
+	}
+
 	result, err := s.Realizations.AllocateCreditsOnly(ctx, flatfeerealizations.AllocateCreditsOnlyInput{
 		Charge:             s.Charge,
 		Amount:             amount,
@@ -115,6 +128,104 @@ func (s *CreditsOnlyStateMachine) AllocateCredits(ctx context.Context) error {
 	return nil
 }
 
+func (s *CreditsOnlyStateMachine) ExtendCharge(ctx context.Context, patch meta.PatchExtend) error {
+	return s.applyPeriodPatch(ctx, patch)
+}
+
+func (s *CreditsOnlyStateMachine) ShrinkCharge(ctx context.Context, patch meta.PatchShrink) error {
+	return s.applyPeriodPatch(ctx, patch)
+}
+
+func (s *CreditsOnlyStateMachine) applyPeriodPatch(ctx context.Context, patch periodPatch) error {
+	targetIntent, err := s.Charge.Intent.GetIntentForTarget(patch.GetTarget())
+	if err != nil {
+		return fmt.Errorf("getting %s intent: %w", patch.GetTarget(), err)
+	}
+
+	if err := patch.ValidateWith(targetIntent.IntentMutableFields.IntentMutableFields); err != nil {
+		return fmt.Errorf("validate %s patch: %w", patch.Op(), err)
+	}
+
+	intent := s.Charge.Intent
+	if err := intent.Mutate(patch.GetTarget(), func(fields *flatfee.IntentMutableFields) {
+		fields.ServicePeriod.To = patch.GetNewServicePeriodTo()
+		fields.FullServicePeriod.To = patch.GetNewFullServicePeriodTo()
+		fields.BillingPeriod.To = patch.GetNewBillingPeriodTo()
+		fields.InvoiceAt = patch.GetNewInvoiceAt()
+	}); err != nil {
+		return fmt.Errorf("mutating %s intent: %w", patch.GetTarget(), err)
+	}
+
+	s.Charge.Intent = intent
+
+	if patch.GetTarget() == meta.ChangeTargetBase && s.Charge.Intent.HasOverrideLayer() {
+		// Subscription sync targets the base intent. When an override is active,
+		// customer-facing credit allocations remain owned by the override.
+		return nil
+	}
+
+	amountAfterProration, err := intent.CalculateAmountAfterProration()
+	if err != nil {
+		return fmt.Errorf("calculating amount after proration: %w", err)
+	}
+	s.Charge.State.AmountAfterProration = amountAfterProration
+
+	if s.Charge.Realizations.CurrentRun == nil {
+		return nil
+	}
+
+	return s.reconcileCurrentRunCredits(ctx, amountAfterProration)
+}
+
+func (s *CreditsOnlyStateMachine) reconcileCurrentRunCredits(ctx context.Context, amount alpacadecimal.Decimal) error {
+	currentRun := s.Charge.Realizations.CurrentRun
+	if currentRun == nil {
+		return nil
+	}
+
+	currencyCalculator, err := s.Charge.Intent.GetCurrency().Calculator()
+	if err != nil {
+		return fmt.Errorf("get currency calculator: %w", err)
+	}
+
+	amount = currencyCalculator.RoundToPrecision(amount)
+	servicePeriod := s.Charge.Intent.GetEffectiveServicePeriod()
+	run := *currentRun
+	run.ServicePeriod = servicePeriod
+
+	reconcileResult, err := s.Realizations.ReconcileCredits(ctx, flatfeerealizations.ReconcileCreditRealizationsInput{
+		Charge:             s.Charge,
+		Run:                run,
+		AllocateAt:         flatfee.UsageBookedAt(s.Charge.Intent.GetEffectivePaymentTerm(), servicePeriod),
+		TargetAmount:       amount,
+		CurrencyCalculator: currencyCalculator,
+	})
+	if err != nil {
+		return fmt.Errorf("reconcile credits for run %s: %w", run.ID.ID, err)
+	}
+
+	run.CreditRealizations = append(run.CreditRealizations, reconcileResult.Realizations...)
+
+	runBase, err := s.Adapter.UpdateRealizationRun(ctx, flatfee.UpdateRealizationRunInput{
+		ID:                   run.ID,
+		ServicePeriod:        mo.Some(servicePeriod),
+		AmountAfterProration: mo.Some(amount),
+		Totals: mo.Some(totals.Totals{
+			Amount:       amount,
+			CreditsTotal: amount,
+			Total:        alpacadecimal.Zero,
+		}),
+		NoFiatTransactionRequired: mo.Some(true),
+	})
+	if err != nil {
+		return fmt.Errorf("update credit-only run: %w", err)
+	}
+
+	run.RealizationRunBase = runBase
+	s.Charge.Realizations.CurrentRun = &run
+	return nil
+}
+
 func (s *CreditsOnlyStateMachine) DeleteCharge(ctx context.Context, patch meta.PatchDelete) error {
 	if err := s.Charge.Intent.Mutate(patch.GetTarget(), func(fields *flatfee.IntentMutableFields) {
 		fields.IntentDeletedAt = lo.ToPtr(clock.Now())

openmeter/billing/charges/meta/patch.go

@@ -1,6 +1,7 @@
 package meta
 
 import (
+	"context"
 	"fmt"
 	"slices"
 
@@ -53,3 +54,11 @@ type TriggerPatchResult[T any] struct {
 	Charge         *T
 	InvoicePatches []invoiceupdater.Patch
 }
+
+// PatchAction adapts a generic Patch action to a concrete patch action when
+// statelessx.AllOfWithParameters requires strict typing for composed actions.
+func PatchAction[T Patch](fn func(context.Context, Patch) error) func(context.Context, T) error {
+	return func(ctx context.Context, patch T) error {
+		return fn(ctx, patch)
+	}
+}

openmeter/billing/charges/usagebased/service/creditheninvoice.go

@@ -36,11 +36,6 @@ type periodPatch interface {
 	ValidateWith(meta.IntentMutableFields) error
 }
 
-type applyPeriodPatchResult struct {
-	ShouldReconcile  bool
-	OldServicePeriod timeutil.ClosedPeriod
-}
-
 var (
 	_ periodPatch = meta.PatchExtend{}
 	_ periodPatch = meta.PatchShrink{}
@@ -398,14 +393,19 @@ func (s *CreditThenInvoiceStateMachine) ShrinkCharge(_ context.Context, patch me
 	return nil
 }
 
-func (s *CreditThenInvoiceStateMachine) applyPeriodPatch(patch periodPatch) (applyPeriodPatchResult, error) {
+type creditThenInvoiceApplyPeriodPatchResult struct {
+	ShouldReconcile  bool
+	OldServicePeriod timeutil.ClosedPeriod
+}
+
+func (s *CreditThenInvoiceStateMachine) applyPeriodPatch(patch periodPatch) (creditThenInvoiceApplyPeriodPatchResult, error) {
 	targetIntent, err := s.Charge.Intent.GetIntentForTarget(patch.GetTarget())
 	if err != nil {
-		return applyPeriodPatchResult{}, fmt.Errorf("getting %s intent: %w", patch.GetTarget(), err)
+		return creditThenInvoiceApplyPeriodPatchResult{}, fmt.Errorf("getting %s intent: %w", patch.GetTarget(), err)
 	}
 
 	if err := patch.ValidateWith(targetIntent.IntentMutableFields.IntentMutableFields); err != nil {
-		return applyPeriodPatchResult{}, fmt.Errorf("validate %s patch: %w", patch.Op(), err)
+		return creditThenInvoiceApplyPeriodPatchResult{}, fmt.Errorf("validate %s patch: %w", patch.Op(), err)
 	}
 
 	oldServicePeriod := meta.NormalizeClosedPeriod(targetIntent.IntentMutableFields.ServicePeriod)
@@ -416,16 +416,16 @@ func (s *CreditThenInvoiceStateMachine) applyPeriodPatch(patch periodPatch) (app
 		fields.BillingPeriod.To = patch.GetNewBillingPeriodTo()
 		fields.InvoiceAt = patch.GetNewInvoiceAt()
 	}); err != nil {
-		return applyPeriodPatchResult{}, fmt.Errorf("mutating %s intent: %w", patch.GetTarget(), err)
+		return creditThenInvoiceApplyPeriodPatchResult{}, fmt.Errorf("mutating %s intent: %w", patch.GetTarget(), err)
 	}
 
 	if patch.GetTarget() == meta.ChangeTargetBase && s.Charge.Intent.HasOverrideLayer() {
 		// Subscription sync targets the base intent. When an override is active,
 		// customer-facing invoice/run state remains owned by the override layer.
-		return applyPeriodPatchResult{}, nil
+		return creditThenInvoiceApplyPeriodPatchResult{}, nil
 	}
 
-	return applyPeriodPatchResult{
+	return creditThenInvoiceApplyPeriodPatchResult{
 		ShouldReconcile:  true,
 		OldServicePeriod: oldServicePeriod,
 	}, nil