feat: unit config conversion support for entitlements - OM-400 (#4712)

Author: rolosp

openmeter/billing/generate.go

@@ -1,3 +1,9 @@
 package billing
 
-//go:generate go run github.com/awalterschulze/goderive
+// GODEBUG=gotypesalias=0 keeps type aliases transparent for goderive's go/types
+// analysis. On Go 1.22+ an alias (e.g. productcatalog.UnitConfig, which aliases
+// unitconfig.UnitConfig) is a first-class types.Alias node that this goderive
+// version does not unwrap; without the flag it fails to see the aliased type's
+// hand-written Equal method, auto-derives equality instead, and recurses into
+// alpacadecimal.Decimal until it OOMs. Remove once goderive unwraps aliases.
+//go:generate env GODEBUG=gotypesalias=0 go run github.com/awalterschulze/goderive

openmeter/billing/models/stddetailedline/generate.go

@@ -1,3 +1,6 @@
 package stddetailedline
 
-//go:generate go run github.com/awalterschulze/goderive
+// GODEBUG=gotypesalias=0 keeps type aliases transparent for goderive's go/types
+// analysis (Go 1.22+ types.Alias nodes this goderive version does not unwrap).
+// See openmeter/billing/generate.go for the full rationale.
+//go:generate env GODEBUG=gotypesalias=0 go run github.com/awalterschulze/goderive

openmeter/credit/adapter/balance_snapshot.go

@@ -69,6 +69,12 @@ func (b *balanceSnapshotRepo) Save(ctx context.Context, owner models.NamespacedI
 				SetGrantBalances(balance.Balances).
 				SetOverage(balance.Overage).
 				SetUsage(&balance.Usage)
+			// Record the conversion regime this snapshot was computed under (OM-400) so
+			// the resume path can refuse to reuse it under a different regime. Pointer
+			// GoType has no SetNillable*, so nil-guard the set (nil = raw).
+			if balance.UnitConfig != nil {
+				command = command.SetUnitConfig(balance.UnitConfig)
+			}
 			commands = append(commands, command)
 		}
 		_, err := rep.db.BalanceSnapshot.CreateBulk(commands...).Save(ctx)
@@ -85,5 +91,8 @@ func mapBalanceSnapshotEntity(entity *db.BalanceSnapshot) balance.Snapshot {
 	if entity.Usage != nil {
 		s.Usage = *entity.Usage
 	}
+	if entity.UnitConfig != nil {
+		s.UnitConfig = entity.UnitConfig
+	}
 	return s
 }

openmeter/credit/balance.go

@@ -109,10 +109,11 @@ func (m *connector) getBalanceSinceSnapshot(ctx context.Context, ownerID models.
 	// TODO: it might be the case that we don't save any snapshots as they require a history breakpoint. To solve this,
 	// we should introduce artificial history breakpoints in the engine, but that would result in more streaming.Query calls, so first lets improve the visibility of what's happening.
 	if err := m.snapshotEngineResult(ctx, snapshotParams{
-		grants:   grants,
-		owner:    ownerID,
-		notAfter: m.getSnapshotNotAfter(periodStart, clock.Now()),
-		meter:    owner.Meter,
+		grants:     grants,
+		owner:      ownerID,
+		notAfter:   m.getSnapshotNotAfter(periodStart, clock.Now()),
+		meter:      owner.Meter,
+		unitConfig: owner.UnitConfig,
 	}, result); err != nil {
 		return def, fmt.Errorf("failed to snapshot engine result: %w", err)
 	}
@@ -317,10 +318,11 @@ func (m *connector) ResetUsageForOwner(ctx context.Context, ownerID models.Names
 
 		// Let's save the snapshot
 		snap, err = m.saveSnapshot(ctx, snapshotParams{
-			grants:   grants,
-			owner:    ownerID,
-			notAfter: at,
-			meter:    owner.Meter,
+			grants:     grants,
+			owner:      ownerID,
+			notAfter:   at,
+			meter:      owner.Meter,
+			unitConfig: owner.UnitConfig,
 		}, snap)
 		if err != nil {
 			return nil, fmt.Errorf("failed to save snapshot: %w", err)

openmeter/credit/balance/balance.go

@@ -4,6 +4,7 @@ import (
 	"time"
 
 	"github.com/openmeterio/openmeter/openmeter/credit/grant"
+	"github.com/openmeterio/openmeter/openmeter/productcatalog/unitconfig"
 )
 
 func NewStartingMap(grants []grant.Grant, at time.Time) Map {
@@ -82,6 +83,13 @@ type Snapshot struct {
 	Balances Map
 	Overage  float64
 	At       time.Time
+
+	// UnitConfig records the conversion regime this snapshot's Usage/Balances were
+	// computed under (OM-400). The resume path only reuses a snapshot whose UnitConfig
+	// matches the owner's current one; a mismatch (e.g. a future backfill that sets an
+	// entitlement's unit_config after raw-unit snapshots already exist) forces a
+	// recompute rather than mixing raw and converted units. Nil = raw (no conversion).
+	UnitConfig *unitconfig.UnitConfig
 }
 
 func (g Snapshot) Balance() float64 {

openmeter/credit/balance/usage.go

@@ -100,7 +100,7 @@ func (u *usageQuerier) QueryUsage(ctx context.Context, ownerID models.Namespaced
 		vTo := alpacadecimal.NewFromFloat(valueTo)
 		vFrom := alpacadecimal.NewFromFloat(valueFrom)
 
-		return vTo.Sub(vFrom).InexactFloat64(), nil
+		return owner.ConvertUsage(vTo.Sub(vFrom).InexactFloat64()), nil
 
 	// For other aggregation types we can simply query the meter
 	case meter.MeterAggregationSum, meter.MeterAggregationCount, meter.MeterAggregationLatest:
@@ -118,7 +118,16 @@ func (u *usageQuerier) QueryUsage(ctx context.Context, ownerID models.Namespaced
 			return 0.0, fmt.Errorf("failed to query meter %s: %w", owner.Meter.Key, err)
 		}
 
-		return u.getValueFromRows(rows)
+		usage, err := u.getValueFromRows(rows)
+		if err != nil {
+			return 0.0, err
+		}
+
+		// Convert raw metered usage into the entitlement's billing unit (OM-400).
+		// Identity when the owner has no UnitConfig. Both the engine's usage source
+		// and the snapshot service's backfill querier route through here, so converted
+		// and raw usage can never be mixed for a given owner.
+		return owner.ConvertUsage(usage), nil
 	default:
 		return 0.0, fmt.Errorf("unsupported aggregation %s", owner.Meter.Aggregation)
 	}

openmeter/credit/grant/owner_connector.go

@@ -5,7 +5,10 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/alpacahq/alpacadecimal"
+
 	"github.com/openmeterio/openmeter/openmeter/meter"
+	"github.com/openmeterio/openmeter/openmeter/productcatalog/unitconfig"
 	"github.com/openmeterio/openmeter/openmeter/streaming"
 	"github.com/openmeterio/openmeter/pkg/models"
 	"github.com/openmeterio/openmeter/pkg/timeutil"
@@ -17,6 +20,27 @@ type Owner struct {
 	DefaultQueryParams streaming.QueryParams
 	ResetBehavior      ResetBehavior
 	StreamingCustomer  streaming.Customer
+
+	// UnitConfig, when set, is the rate card's UnitConfig snapshotted onto this
+	// owner's entitlement. ConvertUsage uses it to map raw metered usage into the
+	// unit the entitlement's limit and grants are authored in. Nil means identity
+	// (no unit_config, or the unitConfig feature is disabled), so balances are
+	// byte-identical to pre-UnitConfig behavior. It is the leaf unitconfig DTO, not
+	// productcatalog, so credit can reference it without an import cycle.
+	UnitConfig *unitconfig.UnitConfig
+}
+
+// ConvertUsage maps a raw metered usage value into the owner's billing unit using
+// its UnitConfig — conversion only, no rounding, so balance checks always use the
+// precise value. Returns the raw value unchanged when the owner has no UnitConfig.
+func (o Owner) ConvertUsage(rawUsage float64) float64 {
+	if o.UnitConfig == nil {
+		return rawUsage
+	}
+
+	converted, _ := o.UnitConfig.Apply(alpacadecimal.NewFromFloat(rawUsage))
+
+	return converted.InexactFloat64()
 }
 
 type EndCurrentUsagePeriodParams struct {

openmeter/credit/grant/owner_connector_test.go

@@ -0,0 +1,43 @@
+package grant_test
+
+import (
+	"testing"
+
+	"github.com/alpacahq/alpacadecimal"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/openmeterio/openmeter/openmeter/credit/grant"
+	"github.com/openmeterio/openmeter/openmeter/productcatalog/unitconfig"
+)
+
+// TestOwnerConvertUsage covers the OM-400 conversion at the credit boundary: identity
+// when the owner has no UnitConfig, and conversion without rounding (balance checks
+// always use the precise value) when it does.
+func TestOwnerConvertUsage(t *testing.T) {
+	t.Run("no unit config is identity", func(t *testing.T) {
+		owner := grant.Owner{}
+		assert.Equal(t, 99_300_000_000.0, owner.ConvertUsage(99_300_000_000))
+	})
+
+	t.Run("divide converts", func(t *testing.T) {
+		owner := grant.Owner{
+			UnitConfig: &unitconfig.UnitConfig{
+				Operation:        unitconfig.UnitConfigOperationDivide,
+				ConversionFactor: alpacadecimal.NewFromInt(1_000_000_000),
+			},
+		}
+		assert.InDelta(t, 99.3, owner.ConvertUsage(99_300_000_000), 1e-6)
+	})
+
+	t.Run("rounding is ignored — balance uses the precise converted value", func(t *testing.T) {
+		owner := grant.Owner{
+			UnitConfig: &unitconfig.UnitConfig{
+				Operation:        unitconfig.UnitConfigOperationDivide,
+				ConversionFactor: alpacadecimal.NewFromInt(1_000_000_000),
+				Rounding:         unitconfig.UnitConfigRoundingModeCeiling,
+			},
+		}
+		// ceiling would invoice 100; the balance check must see 99.3.
+		assert.InDelta(t, 99.3, owner.ConvertUsage(99_300_000_000), 1e-6)
+	})
+}

openmeter/credit/helper.go

@@ -15,6 +15,7 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/credit/engine"
 	"github.com/openmeterio/openmeter/openmeter/credit/grant"
 	"github.com/openmeterio/openmeter/openmeter/meter"
+	"github.com/openmeterio/openmeter/openmeter/productcatalog/unitconfig"
 	"github.com/openmeterio/openmeter/openmeter/streaming"
 	"github.com/openmeterio/openmeter/pkg/framework/transaction"
 	"github.com/openmeterio/openmeter/pkg/models"
@@ -33,29 +34,56 @@ func (m *connector) GetLastValidSnapshotAt(ctx context.Context, owner models.Nam
 			// if no snapshot is found we have to calculate from start of time on all grants and usage
 			m.Logger.Debug("no saved balance found for owner, calculating from start of time", "owner", owner.ID, "at", at)
 
-			startOfMeasurement, err := m.OwnerConnector.GetStartOfMeasurement(ctx, owner)
-			if err != nil {
-				return bal, err
-			}
+			return m.startOfMeasurementSnapshot(ctx, owner, at)
+		}
 
-			grants, err := m.GrantRepo.ListActiveGrantsBetween(ctx, owner, startOfMeasurement, at)
-			if err != nil {
-				return bal, err
-			}
+		return bal, fmt.Errorf("failed to get latest valid grant balance at %s for owner %s: %w", at, owner.ID, err)
+	}
 
-			bal = balance.Snapshot{
-				At:       startOfMeasurement,
-				Balances: balance.NewStartingMap(grants, startOfMeasurement),
-				Overage:  0.0, // There cannot be overage at the start of measurement
-			}
-		} else {
-			return bal, fmt.Errorf("failed to get latest valid grant balance at %s for owner %s: %w", at, owner.ID, err)
-		}
+	// OM-400: only resume from a snapshot computed under the owner's CURRENT conversion
+	// regime. If they differ — e.g. a future backfill set the entitlement's unit_config
+	// after raw-unit snapshots already existed — resuming would combine raw baseline
+	// state with converted incremental usage. Recompute from the start of measurement in
+	// the current regime instead. This is the enforced invariant behind Option A's
+	// "an entitlement converts for its whole life or never" (it also covers a flag flip).
+	regime, err := m.OwnerConnector.DescribeOwner(ctx, owner)
+	if err != nil {
+		return bal, fmt.Errorf("failed to describe owner %s: %w", owner.ID, err)
+	}
+
+	if !bal.UnitConfig.Equal(regime.UnitConfig) {
+		m.Logger.WarnContext(ctx, "balance snapshot unit_config regime does not match owner; recomputing from start of measurement",
+			"owner", owner.ID, "at", at)
+
+		return m.startOfMeasurementSnapshot(ctx, owner, at)
 	}
 
 	return bal, nil
 }
 
+// startOfMeasurementSnapshot builds a fresh zero snapshot at the owner's start of
+// measurement, used when no usable snapshot exists (none saved, or the latest one was
+// computed under a different unit_config regime — see GetLastValidSnapshotAt). Its
+// UnitConfig is left nil because it carries no usage yet; the recomputed result saved
+// afterwards is stamped with the current regime (see saveSnapshot / snapshotParams).
+func (m *connector) startOfMeasurementSnapshot(ctx context.Context, owner models.NamespacedID, at time.Time) (balance.Snapshot, error) {
+	startOfMeasurement, err := m.OwnerConnector.GetStartOfMeasurement(ctx, owner)
+	if err != nil {
+		return balance.Snapshot{}, err
+	}
+
+	grants, err := m.GrantRepo.ListActiveGrantsBetween(ctx, owner, startOfMeasurement, at)
+	if err != nil {
+		return balance.Snapshot{}, err
+	}
+
+	return balance.Snapshot{
+		At:       startOfMeasurement,
+		Balances: balance.NewStartingMap(grants, startOfMeasurement),
+		Overage:  0.0, // There cannot be overage at the start of measurement
+	}, nil
+}
+
 func (m *connector) runEngineInSpan(ctx context.Context, eng engine.Engine, runParams engine.RunParams) (engine.RunResult, error) {
 	ctx, span := m.Tracer.Start(ctx, "credit.runEngine", cTrace.WithEngineParams(runParams))
 	defer span.End()
@@ -170,6 +198,9 @@ type snapshotParams struct {
 	owner models.NamespacedID
 	// Snapshot is saved if the segment is not after this time & the start of the current usage period (at time of snapshot)
 	notAfter time.Time
+	// unitConfig is the owner's current conversion regime, stamped onto the saved
+	// snapshot so the resume path can detect a later regime change (OM-400). Nil = raw.
+	unitConfig *unitconfig.UnitConfig
 }
 
 // It is assumed that there are no snapshots persisted during the length of the history (as engine.Run starts with a snapshot that should be the last valid snapshot)
@@ -228,6 +259,10 @@ func (m *connector) saveSnapshot(ctx context.Context, params snapshotParams, sna
 		return snap, fmt.Errorf("failed to remove inactive grants from snapshot: %w", err)
 	}
 
+	// Stamp the regime this snapshot was computed under so the resume path can reject it
+	// if the owner's unit_config later changes (OM-400).
+	snap.UnitConfig = params.unitConfig
+
 	if err := m.BalanceSnapshotService.Save(ctx, params.owner, []balance.Snapshot{snap}); err != nil {
 		return snap, fmt.Errorf("failed to save snapshot: %w", err)
 	}