fix(licenses): reset transition usage unless carry-over is enabled

Run smaller license transitions synchronously and preserve explicit carry-over semantics across retained and pooled entitlements.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: johnyeocx

server/src/internal/billing/v2/actions/attach/compute/computeAttachPlan.ts

@@ -73,6 +73,7 @@ export const computeAttachPlan = ({
 				customerLicenseBillingContext:
 					attachBillingContext.customerLicenseBillingContext,
 				carryCustomerLicenseState: planTiming === "immediate",
+				carryOverUsages: params.carry_over_usages,
 			})
 		: [];
 	const customerLicenseTransitions =

server/src/internal/billing/v2/actions/batchTransition/batchTransition.ts

@@ -43,6 +43,7 @@ export const batchTransition = async ({
 
 	const hasEntitlementPriceTransitions =
 		entitlementPriceTransitions.transitions.length > 0 ||
+		entitlementPriceTransitions.retained.length > 0 ||
 		entitlementPriceTransitions.added.length > 0 ||
 		entitlementPriceTransitions.deleted.length > 0;
 

server/src/internal/billing/v2/actions/batchTransition/compute/operations/computeBatchTransitionOperations.ts

@@ -1,4 +1,5 @@
 import type {
+	CarryOverUsages,
 	EntitlementWithFeature,
 	InitCustomerEntitlementContext,
 	InitFullCustomerProductOptions,
@@ -20,12 +21,14 @@ export const computeBatchTransitionOperations = ({
 	productTransitions,
 	customerEntitlementInitContext,
 	customerEntitlementInitOptions,
+	carryOverUsages,
 }: {
 	candidateOutgoingEntitlements: EntitlementWithFeature[];
 	candidateOutgoingBasePrices: Price[];
 	productTransitions: ProductTransitions;
 	customerEntitlementInitContext: InitCustomerEntitlementContext;
 	customerEntitlementInitOptions: InitFullCustomerProductOptions;
+	carryOverUsages?: CarryOverUsages;
 }): Pick<
 	CustomerEntitlementBatchTransition,
 	"operations" | "unhandledTransitions"
@@ -36,6 +39,7 @@ export const computeBatchTransitionOperations = ({
 			entitlementPriceTransitions: productTransitions.entitlementPrices,
 			customerEntitlementInitContext,
 			customerEntitlementInitOptions,
+			carryOverUsages,
 		});
 	const basePriceOperation = computeBasePriceOperation({
 		basePriceTransition: productTransitions.basePrice,

server/src/internal/billing/v2/actions/batchTransition/compute/operations/entitlementPriceOperations/computeCustomerEntitlementPatch.ts

@@ -1,5 +1,7 @@
 import {
+	type CarryOverUsages,
 	type EntitlementWithFeature,
+	featureUtils,
 	getStartingBalance,
 	isBooleanEntitlement,
 	isUnlimitedEntitlement,
@@ -32,14 +34,35 @@ export const computeCustomerEntitlementInitialState = ({
 	};
 };
 
+/** Mirrors attach: allocated usage always carries, `carry_from_previous`
+ * carries its own entitlement, and the param carries listed consumables. */
+export const shouldCarryOverUsage = ({
+	toEntitlement,
+	carryOverUsages,
+}: {
+	toEntitlement: EntitlementWithFeature;
+	carryOverUsages: CarryOverUsages;
+}): boolean => {
+	if (featureUtils.isAllocated(toEntitlement.feature)) return true;
+	if (toEntitlement.carry_from_previous) return true;
+	if (!carryOverUsages?.enabled) return false;
+	if (!carryOverUsages.feature_ids) return true;
+	return carryOverUsages.feature_ids.includes(toEntitlement.feature.id);
+};
+
 const computeBalancePatch = ({
 	fromInitialState,
 	toInitialState,
+	carryUsage,
 }: {
 	fromInitialState: CustomerEntitlementInitialState;
 	toInitialState: CustomerEntitlementInitialState;
+	carryUsage: boolean;
 }): CustomerEntitlementBalancePatch | undefined => {
 	if (fromInitialState.tracksBalance && toInitialState.tracksBalance) {
+		if (!carryUsage) {
+			return { type: "set", amount: toInitialState.granted };
+		}
 		const amount = new Decimal(toInitialState.granted).sub(
 			fromInitialState.granted,
 		);
@@ -57,9 +80,11 @@ const computeBalancePatch = ({
 export const computeCustomerEntitlementPatch = ({
 	fromEntitlement,
 	toEntitlement,
+	carryOverUsages,
 }: {
 	fromEntitlement: EntitlementWithFeature;
 	toEntitlement: EntitlementWithFeature;
+	carryOverUsages?: CarryOverUsages;
 }): CustomerEntitlementPatch => {
 	if (
 		isBooleanEntitlement({ entitlement: fromEntitlement }) ||
@@ -75,7 +100,11 @@ export const computeCustomerEntitlementPatch = ({
 		entitlement: toEntitlement,
 	});
 	const patch: CustomerEntitlementPatch = {};
-	const balance = computeBalancePatch({ fromInitialState, toInitialState });
+	const balance = computeBalancePatch({
+		fromInitialState,
+		toInitialState,
+		carryUsage: shouldCarryOverUsage({ toEntitlement, carryOverUsages }),
+	});
 
 	if (balance) patch.balance = balance;
 	if (fromInitialState.unlimited !== toInitialState.unlimited) {

server/src/internal/billing/v2/actions/batchTransition/compute/operations/entitlementPriceOperations/computeEntitlementPriceOperations.ts

@@ -1,12 +1,14 @@
 import {
+	type CarryOverUsages,
 	EntInterval,
 	type EntitlementPrice,
 	type EntitlementWithFeature,
-	entToPooledBalanceIdentity,
 	entsAreSame,
 	entsHaveSamePooledIdentity,
+	entToPooledBalanceIdentity,
 	type InitCustomerEntitlementContext,
 	type InitFullCustomerProductOptions,
+	isBooleanEntitlement,
 	PooledBalanceResetMode,
 } from "@autumn/shared";
 import { initCustomerEntitlementFields } from "@/internal/billing/v2/utils/initFullCustomerProduct/initCustomerEntitlement/initCustomerEntitlementFields";
@@ -23,6 +25,7 @@ import type {
 import {
 	computeCustomerEntitlementInitialState,
 	computeCustomerEntitlementPatch,
+	shouldCarryOverUsage,
 } from "./computeCustomerEntitlementPatch";
 
 const findCandidateEntitlementIds = ({
@@ -47,9 +50,11 @@ const findCandidateEntitlementIds = ({
 const computeReplaceOperation = ({
 	candidateOutgoingEntitlements,
 	transition,
+	carryOverUsages,
 }: {
 	candidateOutgoingEntitlements: EntitlementWithFeature[];
 	transition: EntitlementPriceTransition;
+	carryOverUsages?: CarryOverUsages;
 }): ReplaceEntitlementPriceOperation | undefined => {
 	const { fromEntitlementPrice, toEntitlementPrice } = transition;
 	const fromEntitlement = fromEntitlementPrice.entitlement;
@@ -62,15 +67,19 @@ const computeReplaceOperation = ({
 		(entitlementId) =>
 			!definitionsAreSame || entitlementId !== toEntitlement.id,
 	);
-	if (fromEntitlementIds.length === 0) return undefined;
+	if (fromEntitlementIds.length === 0) {
+		if (!definitionsAreSame) return undefined;
+		fromEntitlementIds.push(toEntitlement.id);
+	}
 
+	const fromIsPooled = fromEntitlement.pooled === true;
+	const toIsPooled = toEntitlement.pooled === true;
+	const isPooledReplace = fromIsPooled && toIsPooled;
 	const customerEntitlementPatch = computeCustomerEntitlementPatch({
 		fromEntitlement,
 		toEntitlement,
+		carryOverUsages,
 	});
-	const fromIsPooled = fromEntitlement.pooled === true;
-	const toIsPooled = toEntitlement.pooled === true;
-	const isPooledReplace = fromIsPooled && toIsPooled;
 	if (!isPooledReplace) {
 		return {
 			type: "replace",
@@ -82,10 +91,15 @@ const computeReplaceOperation = ({
 		};
 	}
 
-	const incrementAmount =
+	const pooledContributionPatch =
 		customerEntitlementPatch.balance?.type === "increment"
-			? customerEntitlementPatch.balance.amount
-			: 0;
+			? customerEntitlementPatch.balance
+			: {
+					type: "set" as const,
+					amount: computeCustomerEntitlementInitialState({
+						entitlement: toEntitlement,
+					}).granted,
+				};
 
 	return {
 		type: "replace",
@@ -96,7 +110,7 @@ const computeReplaceOperation = ({
 		customerEntitlementPatch: {
 			unlimited: customerEntitlementPatch.unlimited,
 		},
-		pooledContributionPatch: { type: "increment", amount: incrementAmount },
+		pooledContributionPatch,
 	};
 };
 
@@ -196,11 +210,13 @@ const computeTransitionOperations = ({
 	transition,
 	initContext,
 	initOptions,
+	carryOverUsages,
 }: {
 	candidateOutgoingEntitlements: EntitlementWithFeature[];
 	transition: EntitlementPriceTransition;
 	initContext: InitCustomerEntitlementContext;
 	initOptions: InitFullCustomerProductOptions;
+	carryOverUsages?: CarryOverUsages;
 }): EntitlementPriceOperation[] => {
 	const fromEntitlement = transition.fromEntitlementPrice.entitlement;
 	const toEntitlement = transition.toEntitlementPrice.entitlement;
@@ -215,6 +231,7 @@ const computeTransitionOperations = ({
 		const operation = computeReplaceOperation({
 			candidateOutgoingEntitlements,
 			transition,
+			carryOverUsages,
 		});
 		return operation ? [operation] : [];
 	}
@@ -244,18 +261,21 @@ export const computeEntitlementPriceOperations = ({
 	entitlementPriceTransitions,
 	customerEntitlementInitContext,
 	customerEntitlementInitOptions,
+	carryOverUsages,
 }: {
 	candidateOutgoingEntitlements: EntitlementWithFeature[];
 	entitlementPriceTransitions: ComputedEntitlementPriceTransitions;
 	customerEntitlementInitContext: InitCustomerEntitlementContext;
 	customerEntitlementInitOptions: InitFullCustomerProductOptions;
+	carryOverUsages?: CarryOverUsages;
 }): {
 	operations: EntitlementPriceOperation[];
 	unhandled: ComputedEntitlementPriceTransitions;
 } => {
 	const operations: EntitlementPriceOperation[] = [];
 	const unhandled: ComputedEntitlementPriceTransitions = {
 		transitions: [],
+		retained: [],
 		added: [],
 		deleted: [],
 	};
@@ -275,10 +295,44 @@ export const computeEntitlementPriceOperations = ({
 				transition,
 				initContext: customerEntitlementInitContext,
 				initOptions: customerEntitlementInitOptions,
+				carryOverUsages,
 			}),
 		);
 	}
 
+	for (const transition of entitlementPriceTransitions.retained) {
+		if (
+			isBooleanEntitlement({
+				entitlement: transition.toEntitlementPrice.entitlement,
+			})
+		) {
+			continue;
+		}
+		if (
+			hasPrice(transition.fromEntitlementPrice) ||
+			hasPrice(transition.toEntitlementPrice)
+		) {
+			unhandled.retained.push(transition);
+			continue;
+		}
+
+		if (
+			shouldCarryOverUsage({
+				toEntitlement: transition.toEntitlementPrice.entitlement,
+				carryOverUsages,
+			})
+		) {
+			continue;
+		}
+
+		const operation = computeReplaceOperation({
+			candidateOutgoingEntitlements,
+			transition,
+			carryOverUsages,
+		});
+		if (operation) operations.push(operation);
+	}
+
 	for (const entitlementPrice of entitlementPriceTransitions.added) {
 		if (hasPrice(entitlementPrice)) {
 			unhandled.added.push(entitlementPrice);

server/src/internal/billing/v2/actions/batchTransition/compute/transitions/computeEntitlementPriceTransitions.ts

@@ -14,6 +14,8 @@ export type EntitlementPriceTransition = {
 
 export type ComputedEntitlementPriceTransitions = {
 	transitions: EntitlementPriceTransition[];
+	/** Same-definition survivors. They still need a replace when usage resets. */
+	retained: EntitlementPriceTransition[];
 	added: EntitlementPrice[];
 	deleted: EntitlementPrice[];
 };
@@ -53,6 +55,7 @@ export const computeEntitlementPriceTransitions = ({
 	}
 
 	const transitions: EntitlementPriceTransition[] = [];
+	const retained: EntitlementPriceTransition[] = [];
 	const deleted: EntitlementPrice[] = [];
 	fromEntitlementPrices.forEach((fromEntitlementPrice, fromIndex) => {
 		const toEntitlementPrice = matchedByFromIndex.get(fromIndex);
@@ -73,13 +76,16 @@ export const computeEntitlementPriceTransitions = ({
 			})
 		) {
 			transitions.push({ fromEntitlementPrice, toEntitlementPrice });
+			return;
 		}
+
+		retained.push({ fromEntitlementPrice, toEntitlementPrice });
 	});
 
 	const added = toEntitlementPrices.filter(
 		(toEntitlementPrice) =>
 			!claimedToEntitlementIds.has(toEntitlementPrice.entitlement.id),
 	);
 
-	return { transitions, added, deleted };
+	return { transitions, retained, added, deleted };
 };

server/src/internal/billing/v2/actions/batchTransition/execute/executeCustomerEntitlementOperations.ts

@@ -25,11 +25,6 @@ const executeReplacement = async ({
 	operation: ReplaceEntitlementPriceOperation;
 }) => {
 	if (operation.fromEntitlementIds.length === 0) return 0;
-	if (operation.fromEntitlementIds.includes(operation.toEntitlementId)) {
-		throw new Error(
-			"Batch replacement requires different outgoing and incoming entitlement IDs",
-		);
-	}
 
 	return executeBatchedMutation({
 		db: ctx.db,