fix: one time product race condition

server/src/internal/customers/add-product/createOneTimeCusProduct.ts

@@ -13,6 +13,8 @@ import {
 } from "@autumn/shared";
 import type { DrizzleCli } from "@/db/initDrizzle.js";
 import { addProductsUpdatedWebhookTask } from "@/internal/analytics/handlers/handleProductsUpdated.js";
+import { executeRedisDeduction } from "@/internal/balances/utils/deduction/executeRedisDeduction.js";
+import type { FeatureDeduction } from "@/internal/balances/utils/types/featureDeduction.js";
 import { triggerVerifyCacheConsistency } from "@/internal/billing/v2/workflows/verifyCacheConsistency/triggerVerifyCacheConsistency.js";
 import { getEntRelatedPrice } from "@/internal/products/entitlements/entitlementUtils.js";
 import { getEntOptions } from "@/internal/products/prices/priceUtils.js";
@@ -32,6 +34,7 @@ const updateOneOffExistingEntitlement = async ({
 	options,
 	relatedPrice,
 	logger,
+	attachParams,
 }: {
 	db: DrizzleCli;
 	cusEnt: FullCustomerEntitlement;
@@ -41,6 +44,7 @@ const updateOneOffExistingEntitlement = async ({
 	options?: FeatureOptions;
 	relatedPrice?: Price;
 	logger: any;
+	attachParams: InsertCusProductParams;
 }) => {
 	if (entitlement.allowance_type === AllowanceType.Unlimited) {
 		return;
@@ -76,6 +80,29 @@ const updateOneOffExistingEntitlement = async ({
 		},
 	});
 
+	const context = attachParams.req;
+	if (context) {
+		const featureDeductions: FeatureDeduction[] = [
+			{
+				feature: entitlement.feature,
+				deduction: -resetBalance,
+			},
+		];
+
+		try {
+			await executeRedisDeduction({
+				ctx: context,
+				fullCustomer: attachParams.customer,
+				deductions: featureDeductions,
+				deductionOptions: {
+					overageBehaviour: "allow",
+				},
+			});
+		} catch (error) {
+			logger.warn(`Failed to execute Redis deduction: ${error}`);
+		}
+	}
+
 	return;
 };
 
@@ -128,6 +155,7 @@ export const updateOneTimeCusProduct = async ({
 				options: options || undefined,
 				relatedPrice,
 				logger,
+				attachParams,
 			});
 		} else {
 			const newCusEnt = initCusEntitlement({

server/src/internal/customers/attach/attachFunctions/updateQuantityFlow/handleQuantityDowngrade.ts

@@ -17,6 +17,7 @@ import {
 import { Decimal } from "decimal.js";
 import type { Stripe } from "stripe";
 import { subToPeriodStartEnd } from "@/external/stripe/stripeSubUtils/convertSubUtils.js";
+import { executeRedisDeduction } from "@/internal/balances/utils/deduction/executeRedisDeduction";
 import type { AttachParams } from "@/internal/customers/cusProducts/AttachParams.js";
 import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService.js";
 import { InvoiceService } from "@/internal/invoices/InvoiceService.js";
@@ -193,5 +194,23 @@ export const handleQuantityDowngrade = async ({
 			id: cusEnt.id,
 			amount: decrementBy,
 		});
+
+		try {
+			await executeRedisDeduction({
+				ctx,
+				fullCustomer: attachParams.customer,
+				deductions: [
+					{
+						feature: cusEnt.entitlement.feature,
+						deduction: decrementBy,
+					},
+				],
+				deductionOptions: {
+					overageBehaviour: "allow",
+				},
+			});
+		} catch (error) {
+			logger.warn(`Failed to execute Redis deduction: ${error}`);
+		}
 	}
 };

server/src/internal/customers/attach/attachFunctions/updateQuantityFlow/handleQuantityUpgrade.ts

@@ -17,6 +17,7 @@ import {
 import { Decimal } from "decimal.js";
 import type { Stripe } from "stripe";
 import { subToPeriodStartEnd } from "@/external/stripe/stripeSubUtils/convertSubUtils.js";
+import { executeRedisDeduction } from "@/internal/balances/utils/deduction/executeRedisDeduction";
 import type { AttachParams } from "@/internal/customers/cusProducts/AttachParams.js";
 import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService.js";
 import { InvoiceService } from "@/internal/invoices/InvoiceService.js";
@@ -201,6 +202,24 @@ export const handleQuantityUpgrade = async ({
 			id: cusEnt.id,
 			amount: incrementBy,
 		});
+
+		try {
+			await executeRedisDeduction({
+				ctx,
+				fullCustomer: attachParams.customer,
+				deductions: [
+					{
+						feature: cusEnt.entitlement.feature,
+						deduction: -incrementBy,
+					},
+				],
+				deductionOptions: {
+					overageBehaviour: "allow",
+				},
+			});
+		} catch (error) {
+			logger.warn(`Failed to execute Redis deduction: ${error}`);
+		}
 	}
 	return { invoice };
 };

server/src/internal/customers/cusProducts/AttachParams.ts

@@ -4,7 +4,6 @@ import type {
 	AttachConfig,
 	AttachReplaceable,
 	AttachScenario,
-	Customer,
 	EntitlementWithFeature,
 	Entity,
 	Feature,
@@ -96,7 +95,7 @@ export type InsertCusProductParams = {
 	req?: AutumnContext;
 	now?: number;
 
-	customer: Customer;
+	customer: FullCustomer;
 	org: Organization;
 	product: FullProduct;
 	prices: Price[];

server/tests/attach/migrations/migration3.test.ts

@@ -7,6 +7,7 @@ import {
 import { defaultApiVersion } from "@tests/constants.js";
 import { TestFeature } from "@tests/setup/v2Features.js";
 import { attachAndExpectCorrect } from "@tests/utils/expectUtils/expectAttach.js";
+import { timeout } from "@tests/utils/genUtils.js";
 import { advanceTestClock } from "@tests/utils/stripeUtils.js";
 import ctx from "@tests/utils/testInitUtils/createTestContext.js";
 import chalk from "chalk";
@@ -100,6 +101,7 @@ describe(`${chalk.yellowBright(`${testCase}: Testing migration for pro with tria
 
 	test("should attach track usage and get correct balance", async () => {
 		const wordsUsage = 120000;
+		await timeout(2000);
 		await autumn.track({
 			customer_id: customerId,
 			value: wordsUsage,

server/tests/attach/others/others2.test.ts

@@ -38,6 +38,7 @@ describe(`${chalk.yellowBright(`${testCase}: Testing one-off`)}`, () => {
 			ctx,
 			products: [oneOff],
 			prefix: testCase,
+			customerId,
 		});
 
 		const { testClockId: testClockId1 } = await initCustomerV3({

server/tests/balances/track/race-condition/track-race-condition2.test.ts

@@ -8,11 +8,7 @@ import { currentRegion } from "@/external/redis/initRedis.js";
 import { executeRedisDeduction } from "@/internal/balances/utils/deduction/executeRedisDeduction.js";
 import { syncItemV3 } from "@/internal/balances/utils/sync/syncItemV3.js";
 import { getOrSetCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.js";
-import { constructFeatureItem } from "@/utils/scriptUtils/constructItem.js";
-import {
-	constructProduct,
-	constructRawProduct,
-} from "@/utils/scriptUtils/createTestProducts.js";
+import { constructRawProduct } from "@/utils/scriptUtils/createTestProducts.js";
 import { initCustomerV3 } from "@/utils/scriptUtils/testUtils/initCustomerV3.js";
 import { initProductsV0 } from "@/utils/scriptUtils/testUtils/initProductsV0.js";
 import { deleteCachedFullCustomer } from "../../../../src/internal/customers/cusUtils/fullCustomerCacheUtils/deleteCachedFullCustomer.js";
@@ -23,15 +19,15 @@ import {
 import { constructPrepaidItem } from "../../../../src/utils/scriptUtils/constructItem.js";
 import { timeout } from "../../../utils/genUtils";
 
-const pro = constructProduct({
-	type: "pro",
-	items: [
-		constructFeatureItem({
-			featureId: TestFeature.Messages,
-			includedUsage: 100,
-		}),
-	],
-});
+// const pro = constructProduct({
+// 	type: "pro",
+// 	items: [
+// 		constructFeatureItem({
+// 			featureId: TestFeature.Messages,
+// 			includedUsage: 100,
+// 		}),
+// 	],
+// });
 
 const oneOffCredits = constructRawProduct({
 	id: "one_off_messages",
@@ -72,13 +68,13 @@ describe(`${chalk.yellowBright("track-race-condition2: sync should not wipe out
 
 		await initProductsV0({
 			ctx,
-			products: [pro, oneOffCredits],
+			products: [oneOffCredits],
 			prefix: testCase,
 		});
 
 		await autumnV2.attach({
 			customer_id: customerId,
-			product_ids: [pro.id, oneOffCredits.id],
+			product_ids: [oneOffCredits.id],
 			options: [
 				{
 					feature_id: TestFeature.Messages,
@@ -122,7 +118,7 @@ describe(`${chalk.yellowBright("track-race-condition2: sync should not wipe out
 			deductions: [
 				{
 					feature: messagesFeature,
-					deduction: 5,
+					deduction: 95, // use up a BIT of one-off credits
 				},
 			],
 			fullCustomer,
@@ -209,7 +205,7 @@ describe(`${chalk.yellowBright("track-race-condition2: sync should not wipe out
 
 		// Expected: 100 (pro) + 100 (initial one-off) - 5 (tracked) + 100 (attached one-off) = 295
 		expect(cachedCustomer.balances[TestFeature.Messages].current_balance).toBe(
-			295,
+			105,
 		);
 
 		const customerAfterSync = await autumnV2.customers.get<ApiCustomer>(
@@ -220,6 +216,6 @@ describe(`${chalk.yellowBright("track-race-condition2: sync should not wipe out
 		);
 		expect(
 			customerAfterSync.balances[TestFeature.Messages].current_balance,
-		).toBe(295);
+		).toBe(105);
 	});
 });

server/tests/integration/billing/legacy/attach/attach-update-quantity.test.ts

@@ -175,7 +175,7 @@ test.concurrent(`${chalk.yellowBright("attach: quantity decrease → increase
 	// Decrease to 200 (on_decrease: none - sets upcoming_quantity, no immediate change)
 	await autumnV1.attach({
 		customer_id: customerId,
-		product_id: `${pro.id}_${customerId}`,
+		product_id: pro.id,
 		options: [{ feature_id: TestFeature.Messages, quantity: 200 }],
 	});
 
@@ -190,14 +190,14 @@ test.concurrent(`${chalk.yellowBright("attach: quantity decrease → increase
 		customer: customerAfterDecrease,
 		productId: pro.id,
 		featureId: TestFeature.Messages,
-		quantity: 3, // Still 300 / 100
-		upcomingQuantity: 2, // 200 / 100
+		quantity: 300, // Still 300 / 100
+		upcomingQuantity: 200, // 200 / 100
 	});
 
 	// Increase to 400 (prorate_immediately - creates invoice, immediate change)
 	await autumnV1.attach({
 		customer_id: customerId,
-		product_id: `${pro.id}_${customerId}`,
+		product_id: pro.id,
 		options: [{ feature_id: TestFeature.Messages, quantity: 400 }],
 	});
 
@@ -219,7 +219,7 @@ test.concurrent(`${chalk.yellowBright("attach: quantity decrease → increase
 	// Decrease back to 200 (sets upcoming_quantity again)
 	await autumnV1.attach({
 		customer_id: customerId,
-		product_id: `${pro.id}_${customerId}`,
+		product_id: pro.id,
 		options: [{ feature_id: TestFeature.Messages, quantity: 200 }],
 	});
 
@@ -231,8 +231,8 @@ test.concurrent(`${chalk.yellowBright("attach: quantity decrease → increase
 		customer: customerFinal,
 		productId: pro.id,
 		featureId: TestFeature.Messages,
-		quantity: 4, // Still 400 / 100
-		upcomingQuantity: 2, // 200 / 100
+		quantity: 400, // Still 400 / 100
+		upcomingQuantity: 200, // 200 / 100
 	});
 });
 
@@ -358,8 +358,8 @@ test.concurrent(`${chalk.yellowBright("attach: prepaid add-on with entities - up
 		customer: entity2,
 		productId: prepaidAddOn.id,
 		featureId: TestFeature.Messages,
-		quantity: entity2OriginalQuantity / 100, // billingUnits = 100
-		upcomingQuantity: entity2DowngradedQuantity / 100,
+		quantity: entity2OriginalQuantity, // billingUnits = 100
+		upcomingQuantity: entity2DowngradedQuantity,
 	});
 });
 
@@ -438,7 +438,7 @@ test.concurrent(`${chalk.yellowBright("attach: quantity upgrade with prorate-nex
 		customer: customerFinal,
 		productId: pro.id,
 		featureId: TestFeature.Messages,
-		quantity: 4, // 400 / 100 billingUnits
+		quantity: 400,
 		upcomingQuantity: "undefined", // No upcoming_quantity since it's an upgrade
 	});