fix: child unique reference handling (#4385)

Author: turip

.agents/skills/test/SKILL.md

@@ -38,6 +38,12 @@ For running a specific test directly:
 POSTGRES_HOST=127.0.0.1 go test -tags=dynamic ./openmeter/<domain>/...
 ```
 
+When a Postgres-backed test fails, the suite output often includes a `testdbconf:` URL for the per-test database. Use that URL with `psql` to inspect the failed test state before rerunning, because the database is still useful for RCA. Quote the connection string if it contains query parameters, for example:
+
+```bash
+psql 'postgres://pgtdbuser:pgtdbpass@127.0.0.1:5432/testdb_tpl_...?sslmode=disable'
+```
+
 ## Key Test Utilities
 
 From `openmeter/testutils/`:

openmeter/billing/charges/charge.go

@@ -137,6 +137,31 @@ func (c Charge) GetChargeID() (meta.ChargeID, error) {
 	return meta.ChargeID{}, fmt.Errorf("invalid charge type: %s", c.t)
 }
 
+func (c Charge) GetUniqueReferenceID() (*string, error) {
+	switch c.t {
+	case meta.ChargeTypeFlatFee:
+		if c.flatFee == nil {
+			return nil, fmt.Errorf("flat fee charge is nil")
+		}
+
+		return c.flatFee.Intent.UniqueReferenceID, nil
+	case meta.ChargeTypeCreditPurchase:
+		if c.creditPurchase == nil {
+			return nil, fmt.Errorf("credit purchase charge is nil")
+		}
+
+		return c.creditPurchase.Intent.UniqueReferenceID, nil
+	case meta.ChargeTypeUsageBased:
+		if c.usageBased == nil {
+			return nil, fmt.Errorf("usage based charge is nil")
+		}
+
+		return c.usageBased.Intent.UniqueReferenceID, nil
+	}
+
+	return nil, fmt.Errorf("invalid charge type: %s", c.t)
+}
+
 func (c Charge) GetCustomerID() (customer.CustomerID, error) {
 	switch c.t {
 	case meta.ChargeTypeFlatFee:

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

@@ -191,10 +191,9 @@ func buildFlatFeeGatheringLine(input buildFlatFeeGatheringLineInput) (billing.Ga
 
 			TaxConfig: lineIntent.TaxConfig.ToTaxConfig(),
 
-			Engine:                 billing.LineEngineTypeChargeFlatFee,
-			ChargeID:               lo.ToPtr(flatFee.ID),
-			ChildUniqueReferenceID: lineIntent.UniqueReferenceID,
-			Subscription:           subscription,
+			Engine:       billing.LineEngineTypeChargeFlatFee,
+			ChargeID:     lo.ToPtr(flatFee.ID),
+			Subscription: subscription,
 		},
 	}
 

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

@@ -127,10 +127,9 @@ func gatheringLineFromUsageBasedChargeForPeriod(charge usagebased.Charge, servic
 
 			TaxConfig: intent.TaxConfig.ToTaxConfig(),
 
-			ChargeID:               lo.ToPtr(charge.ID),
-			Engine:                 billing.LineEngineTypeChargeUsageBased,
-			ChildUniqueReferenceID: intent.UniqueReferenceID,
-			Subscription:           subscription,
+			ChargeID:     lo.ToPtr(charge.ID),
+			Engine:       billing.LineEngineTypeChargeUsageBased,
+			Subscription: subscription,
 		},
 	}
 

openmeter/billing/httpdriver/invoice.go

@@ -317,7 +317,7 @@ func (h *handler) ProgressInvoice(action ProgressAction) ProgressInvoiceHandler
 			case InvoiceProgressActionAdvance:
 				invoice, err = h.service.AdvanceInvoice(ctx, request.Invoice)
 			case InvoiceProgressActionSnapshotQuantities:
-				invoice, err = h.service.SnapshotQuantities(ctx, request.Invoice)
+				invoice, err = h.service.ForceCollectInvoice(ctx, request.Invoice)
 			default:
 				return ProgressInvoiceResponse{}, fmt.Errorf("invalid action: %s", action)
 			}

openmeter/billing/service.go

@@ -88,7 +88,8 @@ type InvoiceService interface {
 	// - the invoice is in a state that cannot be advanced (e.g. waiting for draft period to expire)
 	// - the invoice is advanced to the final state
 	AdvanceInvoice(ctx context.Context, input AdvanceInvoiceInput) (StandardInvoice, error)
-	SnapshotQuantities(ctx context.Context, input SnapshotQuantitiesInput) (StandardInvoice, error)
+	// ForceCollectInvoice bypasses the invoice collection period and moves the invoice into collection immediately.
+	ForceCollectInvoice(ctx context.Context, input ForceCollectInvoiceInput) (StandardInvoice, error)
 	ApproveInvoice(ctx context.Context, input ApproveInvoiceInput) (StandardInvoice, error)
 	PaymentAuthorized(ctx context.Context, input PaymentAuthorizedInput) (StandardInvoice, error)
 	RetryInvoice(ctx context.Context, input RetryInvoiceInput) (StandardInvoice, error)

openmeter/billing/service/invoice.go

@@ -514,8 +514,9 @@ func (s *Service) PaymentAuthorized(ctx context.Context, input billing.PaymentAu
 	return s.executeTriggerOnInvoice(ctx, input, billing.TriggerAuthorized)
 }
 
-func (s *Service) SnapshotQuantities(ctx context.Context, input billing.SnapshotQuantitiesInput) (billing.StandardInvoice, error) {
-	return s.executeTriggerOnInvoice(ctx, input, billing.TriggerSnapshotQuantities)
+// ForceCollectInvoice bypasses the invoice collection period and moves the invoice into collection immediately.
+func (s *Service) ForceCollectInvoice(ctx context.Context, input billing.ForceCollectInvoiceInput) (billing.StandardInvoice, error) {
+	return s.executeTriggerOnInvoice(ctx, input, billing.TriggerForceCollect)
 }
 
 func (s *Service) RetryInvoice(ctx context.Context, input billing.RetryInvoiceInput) (billing.StandardInvoice, error) {

openmeter/billing/service/stdinvoicestate.go

@@ -101,7 +101,7 @@ func allocateStateMachine() *InvoiceStateMachine {
 		).
 		Permit(billing.TriggerDelete, billing.StandardInvoiceStatusDeleteInProgress).
 		Permit(billing.TriggerUpdated, billing.StandardInvoiceStatusDraftUpdating).
-		Permit(billing.TriggerSnapshotQuantities, billing.StandardInvoiceStatusDraftCollecting)
+		Permit(billing.TriggerForceCollect, billing.StandardInvoiceStatusDraftCollecting)
 
 	stateMachine.Configure(billing.StandardInvoiceStatusDraftCollecting).
 		Permit(billing.TriggerNext, billing.StandardInvoiceStatusDraftValidating).
@@ -409,7 +409,7 @@ func (m *InvoiceStateMachine) StatusDetails(ctx context.Context) (billing.Standa
 		outErr = errors.Join(outErr, err)
 	}
 
-	if availableActions.SnapshotQuantities, err = m.calculateAvailableActionDetails(ctx, billing.TriggerSnapshotQuantities); err != nil {
+	if availableActions.SnapshotQuantities, err = m.calculateAvailableActionDetails(ctx, billing.TriggerForceCollect); err != nil {
 		outErr = errors.Join(outErr, err)
 	}
 

openmeter/billing/stdinvoice.go

@@ -752,11 +752,11 @@ type AssociatedLineCountsAdapterResponse struct {
 }
 
 type (
-	AdvanceInvoiceInput     = InvoiceID
-	ApproveInvoiceInput     = InvoiceID
-	PaymentAuthorizedInput  = InvoiceID
-	RetryInvoiceInput       = InvoiceID
-	SnapshotQuantitiesInput = InvoiceID
+	AdvanceInvoiceInput      = InvoiceID
+	ApproveInvoiceInput      = InvoiceID
+	ForceCollectInvoiceInput = InvoiceID
+	PaymentAuthorizedInput   = InvoiceID
+	RetryInvoiceInput        = InvoiceID
 )
 
 type UpdateStandardInvoiceAdapterInput = StandardInvoice

openmeter/billing/stdinvoicestate.go

@@ -21,8 +21,8 @@ var (
 	// TriggerUpdated is used to trigger a change in the invoice (we are using this to calculate the immutable states
 	// and trigger re-validation)
 	TriggerUpdated InvoiceTrigger = "trigger_updated"
-	// TriggerSnapshotQuantities is used to snapshot the quantities for usage based line items
-	TriggerSnapshotQuantities InvoiceTrigger = "trigger_snapshot_quantities"
+	// TriggerForceCollect skips the collection wait and forces an invoice into collection.
+	TriggerForceCollect InvoiceTrigger = "trigger_force_collect"
 	// triggerDelete is used to delete the invoice
 	TriggerDelete InvoiceTrigger = "trigger_delete"