License pooled balances: lifecycle, batch transitions, and runtime

server/src/internal/billing/v2/actions/attachLicense/compute/computeAttachLicensePlan.ts

@@ -5,6 +5,7 @@ import {
 	findFeatureById,
 } from "@autumn/shared";
 import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
+import { computePooledBalanceTransitionPlan } from "@/internal/billing/v2/pooledBalances/compute/computePooledBalanceTransitionPlan.js";
 import { initFullCustomerProductFromCustomerLicense } from "@/internal/billing/v2/utils/initFullCustomerProduct/initFullCustomerProductFromCustomerLicense.js";
 import { constructEntity } from "@/internal/entities/entityUtils/entityUtils.js";
 import type { AttachLicenseContext, AttachLicensePlan } from "../types.js";
@@ -86,12 +87,20 @@ export const computeAttachLicensePlan = ({
 			}),
 		}));
 
+	const incomingCustomerProducts = insertedAssignments.map(
+		(assignment) => assignment.customerProduct,
+	);
+	const { pooledBalancePlan } = computePooledBalanceTransitionPlan({
+		ctx,
+		fullCustomer,
+		incomingCustomerProducts,
+		now: context.currentEpochMs,
+	});
+
 	const billingPlan: AutumnBillingPlan = {
 		customerId: fullCustomer.id ?? fullCustomer.internal_id,
 		insertEntities: newEntities,
-		insertCustomerProducts: insertedAssignments.map(
-			(assignment) => assignment.customerProduct,
-		),
+		insertCustomerProducts: incomingCustomerProducts,
 		updateCustomerProducts: computeReusedCustomerProductUpdates({
 			reusedAssignments,
 		}),
@@ -101,6 +110,7 @@ export const computeAttachLicensePlan = ({
 				remainingChange: -assignmentEntities.length,
 			},
 		],
+		pooledBalancePlan,
 	};
 
 	return {

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

@@ -1,9 +1,13 @@
 import {
+	EntInterval,
 	type EntitlementPrice,
 	type EntitlementWithFeature,
+	entToPooledBalanceIdentity,
 	entsAreSame,
+	entsHaveSamePooledIdentity,
 	type InitCustomerEntitlementContext,
 	type InitFullCustomerProductOptions,
+	PooledBalanceResetMode,
 } from "@autumn/shared";
 import { initCustomerEntitlementFields } from "@/internal/billing/v2/utils/initFullCustomerProduct/initCustomerEntitlement/initCustomerEntitlementFields";
 import type {
@@ -16,7 +20,10 @@ import type {
 	ComputedEntitlementPriceTransitions,
 	EntitlementPriceTransition,
 } from "../../transitions/computeEntitlementPriceTransitions";
-import { computeCustomerEntitlementPatch } from "./computeCustomerEntitlementPatch";
+import {
+	computeCustomerEntitlementInitialState,
+	computeCustomerEntitlementPatch,
+} from "./computeCustomerEntitlementPatch";
 
 const findCandidateEntitlementIds = ({
 	candidateOutgoingEntitlements,
@@ -57,16 +64,39 @@ const computeReplaceOperation = ({
 	);
 	if (fromEntitlementIds.length === 0) return undefined;
 
+	const customerEntitlementPatch = computeCustomerEntitlementPatch({
+		fromEntitlement,
+		toEntitlement,
+	});
+	const fromIsPooled = fromEntitlement.pooled === true;
+	const toIsPooled = toEntitlement.pooled === true;
+	const isPooledReplace = fromIsPooled && toIsPooled;
+	if (!isPooledReplace) {
+		return {
+			type: "replace",
+			fromEntitlementIds,
+			toEntitlementId: toEntitlement.id,
+			fromEntitlementPrice,
+			toEntitlementPrice,
+			customerEntitlementPatch,
+		};
+	}
+
+	const incrementAmount =
+		customerEntitlementPatch.balance?.type === "increment"
+			? customerEntitlementPatch.balance.amount
+			: 0;
+
 	return {
 		type: "replace",
 		fromEntitlementIds,
 		toEntitlementId: toEntitlement.id,
 		fromEntitlementPrice,
 		toEntitlementPrice,
-		customerEntitlementPatch: computeCustomerEntitlementPatch({
-			fromEntitlement,
-			toEntitlement,
-		}),
+		customerEntitlementPatch: {
+			unlimited: customerEntitlementPatch.unlimited,
+		},
+		pooledContributionPatch: { type: "increment", amount: incrementAmount },
 	};
 };
 
@@ -87,15 +117,61 @@ const computeAddOperation = ({
 	});
 	existingEntitlementIds.push(entitlementPrice.entitlement.id);
 
+	const entitlement = entitlementPrice.entitlement;
+	const customerEntitlement = initCustomerEntitlementFields({
+		initContext,
+		initOptions,
+		entitlement,
+	});
+	if (entitlement.pooled !== true) {
+		return {
+			type: "add",
+			entitlementPrice,
+			existingEntitlementIds: [...new Set(existingEntitlementIds)],
+			customerEntitlement,
+		};
+	}
+
+	const customerLicenseLinkId = initOptions.customerLicenseLinkId;
+	if (!customerLicenseLinkId) {
+		throw new Error(
+			"Pooled entitlement addition requires a customer license link",
+		);
+	}
+
+	const initialState = computeCustomerEntitlementInitialState({ entitlement });
+	const entIdentity = entToPooledBalanceIdentity({ entitlement });
+	const resetMode =
+		entIdentity.interval === EntInterval.Lifetime
+			? PooledBalanceResetMode.Lifetime
+			: PooledBalanceResetMode.Lazy;
+	const isLifetimeReset = resetMode === PooledBalanceResetMode.Lifetime;
 	return {
 		type: "add",
 		entitlementPrice,
 		existingEntitlementIds: [...new Set(existingEntitlementIds)],
-		customerEntitlement: initCustomerEntitlementFields({
-			initContext,
-			initOptions,
-			entitlement: entitlementPrice.entitlement,
-		}),
+		customerEntitlement: {
+			...customerEntitlement,
+			balance: 0,
+		},
+		pooledAdd: {
+			contributionAmount: initialState.granted,
+			nextResetAt: isLifetimeReset
+				? null
+				: (customerEntitlement.next_reset_at ?? null),
+			featureId: entitlement.feature.id,
+			rollover: entitlement.rollover ?? null,
+			identity: {
+				...entIdentity,
+				internalCustomerId: customerEntitlement.internal_customer_id,
+				resetCycleAnchor: isLifetimeReset
+					? null
+					: (customerEntitlement.reset_cycle_anchor ?? null),
+				resetMode,
+				stripeSubscriptionId: null,
+				customerLicenseLinkId,
+			},
+		},
 	};
 };
 
@@ -115,6 +191,51 @@ const computeRemoveOperation = ({
 	return { type: "remove", entitlementPrice, fromEntitlementIds };
 };
 
+const computeTransitionOperations = ({
+	candidateOutgoingEntitlements,
+	transition,
+	initContext,
+	initOptions,
+}: {
+	candidateOutgoingEntitlements: EntitlementWithFeature[];
+	transition: EntitlementPriceTransition;
+	initContext: InitCustomerEntitlementContext;
+	initOptions: InitFullCustomerProductOptions;
+}): EntitlementPriceOperation[] => {
+	const fromEntitlement = transition.fromEntitlementPrice.entitlement;
+	const toEntitlement = transition.toEntitlementPrice.entitlement;
+	const fromIsPooled = fromEntitlement.pooled === true;
+	const toIsPooled = toEntitlement.pooled === true;
+	const isPooledAmountChange =
+		fromIsPooled &&
+		toIsPooled &&
+		entsHaveSamePooledIdentity(fromEntitlement, toEntitlement);
+
+	if (isPooledAmountChange || (!fromIsPooled && !toIsPooled)) {
+		const operation = computeReplaceOperation({
+			candidateOutgoingEntitlements,
+			transition,
+		});
+		return operation ? [operation] : [];
+	}
+
+	const operations: EntitlementPriceOperation[] = [];
+	const remove = computeRemoveOperation({
+		candidateOutgoingEntitlements,
+		entitlementPrice: transition.fromEntitlementPrice,
+	});
+	if (remove) operations.push(remove);
+	operations.push(
+		computeAddOperation({
+			candidateOutgoingEntitlements,
+			entitlementPrice: transition.toEntitlementPrice,
+			initContext,
+			initOptions,
+		}),
+	);
+	return operations;
+};
+
 const hasPrice = (entitlementPrice: EntitlementPrice) =>
 	Boolean(entitlementPrice.price);
 
@@ -148,11 +269,14 @@ export const computeEntitlementPriceOperations = ({
 			continue;
 		}
 
-		const operation = computeReplaceOperation({
-			candidateOutgoingEntitlements,
-			transition,
-		});
-		if (operation) operations.push(operation);
+		operations.push(
+			...computeTransitionOperations({
+				candidateOutgoingEntitlements,
+				transition,
+				initContext: customerEntitlementInitContext,
+				initOptions: customerEntitlementInitOptions,
+			}),
+		);
 	}
 
 	for (const entitlementPrice of entitlementPriceTransitions.added) {

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

@@ -12,6 +12,7 @@ import { BATCH_TRANSITION_OPERATION_CONCURRENCY } from "../utils/batchTransition
 import { executeBatchedMutation } from "./executeBatchedMutation";
 import { addCustomerEntitlementsBatch } from "./sql/addCustomerEntitlementsBatch";
 import { deleteCustomerEntitlementsBatch } from "./sql/deleteCustomerEntitlementsBatch";
+import { insertPooledBalanceGraph } from "./sql/insertPooledBalanceGraph";
 import { replaceCustomerEntitlementsBatch } from "./sql/replaceCustomerEntitlementsBatch";
 
 const executeReplacement = async ({
@@ -55,8 +56,19 @@ const executeAddition = async ({
 	return executeBatchedMutation({
 		db: ctx.db,
 		operationName: "Customer entitlement addition",
-		executeBatch: ({ db, batchSize }) =>
-			addCustomerEntitlementsBatch({
+		executeBatch: async ({ db, batchSize }) => {
+			const pooledBalanceId = operation.pooledAdd
+				? await insertPooledBalanceGraph({
+						db,
+						pooledAdd: operation.pooledAdd,
+						customerId: operation.customerEntitlement.customer_id ?? null,
+						orgId: ctx.org.id,
+						env: ctx.env,
+						now: Date.now(),
+					})
+				: undefined;
+
+			return addCustomerEntitlementsBatch({
 				db,
 				customerLicenseLinkId: batchTransition.customerLicenseLinkId,
 				assignmentCutoffMs: batchTransition.assignmentCutoffMs,
@@ -65,7 +77,14 @@ const executeAddition = async ({
 				),
 				operation,
 				batchSize,
-			}),
+				pooledBalanceId,
+				contributionIds: pooledBalanceId
+					? Array.from({ length: batchSize }, () =>
+							generateId("pool_contribution"),
+						)
+					: undefined,
+			});
+		},
 	});
 };