fix: usage limits returned in customer + other bugs

server/shell/g4.sh

@@ -10,7 +10,8 @@ $MOCHA_CMD 'tests/advanced/multiFeature/*.ts' \
            'tests/attach/updateQuantity/*.ts' \
            'tests/advanced/referrals/*.ts'
           
-# $MOCHA_CMD 'tests/attach/multiProduct/*.ts' 
+$MOCHA_CMD 'tests/attach/multiProduct/*.ts' \
+           'tests/advanced/usageLimit/*.ts' 
 
-# $MOCHA_CMD 'tests/advanced/usage/*.ts'
+$MOCHA_CMD 'tests/advanced/usage/*.ts'
            

server/src/external/stripe/webhookHandlers/handleCheckoutCompleted/handleRemainingSets.ts

@@ -1,20 +1,11 @@
+import Stripe from "stripe";
 import { createStripeSub } from "../../stripeSubUtils/createStripeSub.js";
-
 import { DrizzleCli } from "@/db/initDrizzle.js";
 import { AttachParams } from "@/internal/customers/cusProducts/AttachParams.js";
-import { getBillingType } from "@/internal/products/prices/priceUtils.js";
 import { findPriceFromStripeId } from "@/internal/products/prices/priceUtils/findPriceUtils.js";
 import { notNullish } from "@/utils/genUtils.js";
 import { ItemSet } from "@/utils/models/ItemSet.js";
-import {
-  APIVersion,
-  BillingType,
-  FullProduct,
-  Organization,
-  Price,
-} from "@autumn/shared";
-import Stripe from "stripe";
-import { getPlaceholderItem } from "../../stripePriceUtils.js";
+import { APIVersion, BillingType, Organization } from "@autumn/shared";
 import { getArrearItems } from "../../stripeSubUtils/getStripeSubItems/getArrearItems.js";
 
 const filterUsagePrices = ({

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

@@ -109,6 +109,7 @@ export const handleQuantityUpgrade = async ({
         stripeCusId: stripeSub.customer as string,
         stripeSubId: stripeSub.id,
         paymentMethod: paymentMethod || null,
+        logger,
       });
     }
   }

server/src/internal/customers/attach/attachUtils/handleAttachErrors.ts

@@ -106,6 +106,18 @@ const handlePrepaidErrors = async ({
           statusCode: 400,
         });
       }
+
+      let usageLimit = priceEnt.usage_limit;
+      let totalQuantity =
+        options?.quantity! * (price.config as UsagePriceConfig).billing_units!;
+
+      if (usageLimit && totalQuantity + priceEnt.allowance! > usageLimit) {
+        throw new RecaseError({
+          message: `Quantity + included usage exceeds usage limit of ${usageLimit} for feature ${priceEnt.feature_id}`,
+          code: ErrCode.InvalidOptions,
+          statusCode: 400,
+        });
+      }
     }
   }
 };

server/src/internal/customers/cusProducts/cusEnts/cusEntUtils.ts

@@ -315,13 +315,13 @@ export const getResetBalance = ({
 
   let billingType = getBillingType(config);
   if (billingType != BillingType.UsageInAdvance) {
-    return entitlement.allowance;
+    return entitlement.allowance || 0;
   }
 
   let quantity = options?.quantity;
   let billingUnits = (relatedPrice.config as UsagePriceConfig).billing_units;
   if (nullish(quantity) || nullish(billingUnits)) {
-    return entitlement.allowance;
+    return entitlement.allowance || 0;
   }
 
   try {
@@ -330,7 +330,7 @@ export const getResetBalance = ({
     console.log(
       "WARNING: Failed to return quantity * billing units, returning allowance...",
     );
-    return entitlement.allowance;
+    return entitlement.allowance || 0;
   }
 };
 

server/src/internal/customers/cusProducts/cusEnts/cusEntUtils/getExistingUsage.ts

@@ -6,6 +6,7 @@ import {
   Entity,
   Feature,
   FeatureType,
+  FullCusEntWithFullCusProduct,
   FullCusProduct,
   FullCustomerEntitlement,
   Price,
@@ -184,8 +185,8 @@ export const addExistingUsagesToCusEnts = ({
 
   let fullCusEnts = cusEnts.map((ce) => {
     let entitlement = entitlements.find((e) => e.id === ce.entitlement_id!);
-    return { ...ce, entitlement };
-  }) as FullCustomerEntitlement[];
+    return { ...ce, entitlement, customer_product: curCusProduct };
+  }) as FullCusEntWithFullCusProduct[];
 
   // Sort cusEnts
   sortCusEntsForDeduction(fullCusEnts);

server/src/internal/customers/cusUtils/cusUtils.ts

@@ -120,6 +120,8 @@ export const getCusEntsInFeatures = async ({
   reverseOrder?: boolean;
 }) => {
   let cusProducts = customer.customer_products;
+
+  // This is important, attaching customer_product to cus ent is used elsewhere, don't delete.
   let cusEnts = cusProducts.flatMap((cusProduct) => {
     return cusProduct.customer_entitlements.map((cusEnt) => ({
       ...cusEnt,

server/src/internal/customers/handlers/handleUpdateEntitlement.ts

@@ -73,6 +73,13 @@ export const handleUpdateEntitlement = async (req: any, res: any) => {
       withCusProduct: true,
     });
 
+    const cusProduct = await CusProductService.get({
+      db,
+      id: cusEnt.customer_product_id,
+      orgId: req.orgId,
+      env: req.env,
+    });
+
     if (balance < 0 && !cusEnt.usage_allowed) {
       throw new RecaseError({
         message: "Entitlement does not allow usage",
@@ -99,7 +106,10 @@ export const handleUpdateEntitlement = async (req: any, res: any) => {
     let originalBalance = structuredClone(masterBalance);
 
     let { newBalance, newEntities, newAdjustment } = performDeductionOnCusEnt({
-      cusEnt,
+      cusEnt: {
+        ...cusEnt,
+        customer_product: cusProduct!,
+      },
       toDeduct: deducted,
       addAdjustment: true,
       allowNegativeBalance: cusEnt.usage_allowed || false,

server/src/internal/products/product-items/compareItemUtils.ts

@@ -98,6 +98,10 @@ export const featurePriceItemsAreSame = ({
       condition: item1.included_usage == item2.included_usage,
       message: `Included usage different: ${item1.included_usage} != ${item2.included_usage}`,
     },
+    usage_limit: {
+      condition: item1.usage_limit == item2.usage_limit,
+      message: `Usage limit different: ${item1.usage_limit} !== ${item2.usage_limit}`,
+    },
     reset_usage_when_enabled: {
       condition:
         item1.reset_usage_when_enabled == item2.reset_usage_when_enabled,

server/src/internal/products/productUtils/compareProductUtils.ts

@@ -72,26 +72,18 @@ export const productsAreSame = ({
 
   // Check if any feature's usage limits have changed
   let usageLimitsChanged = false;
-  items1.some(item1 => {
-    const matchingItem2 = items2?.find(item2 => item2.feature_id === item1.feature_id);
+  items1.some((item1) => {
+    const matchingItem2 = items2?.find(
+      (item2) => item2.feature_id === item1.feature_id,
+    );
     if (!matchingItem2) return false;
-    
-    const feature = features.find(f => f.id === item1.feature_id);
+
+    const feature = features.find((f) => f.id === item1.feature_id);
     if (!feature) return false;
 
-    // Check if the usage limit has changed
-    if (item1.usage_limit !== matchingItem2.usage_limit) {
-      usageLimitsChanged = true;
-      return true;
-    }
     return false;
   });
 
-  if (usageLimitsChanged) {
-    itemsSame = false;
-    pricesChanged = true;
-  }
-
   items2 =
     curProductV2?.items ||
     mapToProductItems({

server/src/trigger/adjustAllowance.ts

@@ -96,15 +96,6 @@ export const adjustAllowance = async ({
 
   // TODO: TRACK
 
-  if (newBalance < -(cusEnt.entitlement.usage_limit || 0)) {
-    throw new RecaseError({
-      message: `Balance exceeds usage limit of ${cusEnt.entitlement.usage_limit}`,
-      code: ErrCode.InvalidInputs,
-      statusCode: StatusCodes.BAD_REQUEST,
-    });
-    // return;
-  }
-
   if (
     !cusProduct ||
     !cusPrice ||
@@ -114,6 +105,15 @@ export const adjustAllowance = async ({
     return { newReplaceables: [], invoice: null, deletedReplaceables: null };
   }
 
+  let ent = cusEnt.entitlement;
+  if (ent.usage_limit && newBalance < ent.allowance! - (ent.usage_limit || 0)) {
+    throw new RecaseError({
+      message: `Balance exceeds usage limit of ${cusEnt.entitlement.usage_limit}`,
+      code: ErrCode.InvalidInputs,
+      statusCode: StatusCodes.BAD_REQUEST,
+    });
+  }
+
   logger.info(`--------------------------------`);
   logger.info(`Updating arrear prorated usage: ${affectedFeature.name}`);
   logger.info(`Customer: ${customer.name}, Org: ${org.slug}`);

server/src/trigger/arrearProratedUsage/createUpgradeProrationInvoice.ts

@@ -148,6 +148,7 @@ export const createUpgradeProrationInvoice = async ({
       paymentMethod,
       stripeCusId: sub.customer as string,
       stripeSubId: sub.id,
+      logger,
     });
 
     logger.info(`Paid for invoice ${finalInvoice?.id}`);

server/src/trigger/arrearProratedUsage/handleProratedDowngrade.ts

@@ -110,6 +110,7 @@ export const createDowngradeProrationInvoice = async ({
       paymentMethod: null,
       stripeCusId: sub.customer as string,
       stripeSubId: sub.id,
+      logger,
     });
 
     invoice = finalInvoice;

server/src/trigger/deductUtils.ts

@@ -4,6 +4,7 @@ import {
   Event,
   FeatureType,
   FullCustomerEntitlement,
+  Entitlement,
 } from "@autumn/shared";
 
 import { AggregateType } from "@autumn/shared";
@@ -82,18 +83,39 @@ export const performDeduction = ({
   cusEntBalance,
   toDeduct,
   allowNegativeBalance = false,
+  ent,
+  resetBalance,
+  blockUsageLimit = true,
 }: {
   cusEntBalance: Decimal;
   toDeduct: number;
   allowNegativeBalance?: boolean;
+  ent: Entitlement;
+  resetBalance: number;
+  blockUsageLimit?: boolean;
 }) => {
   // Either deduct from balance or entity balance
-
   if (allowNegativeBalance) {
+    let usageLimit = ent.usage_limit;
+    let minBalance = usageLimit
+      ? new Decimal(resetBalance).minus(usageLimit).toNumber()
+      : undefined;
     let newBalance = cusEntBalance.minus(toDeduct).toNumber();
-    let deducted = toDeduct;
-    let toDeduct_ = 0;
-    return { newBalance, deducted, toDeduct: toDeduct_ };
+
+    if (
+      blockUsageLimit &&
+      minBalance &&
+      new Decimal(newBalance).lt(minBalance)
+    ) {
+      newBalance = minBalance;
+      let deducted = new Decimal(cusEntBalance).minus(minBalance).toNumber();
+      let toDeduct_ = new Decimal(toDeduct).minus(deducted).toNumber();
+      return { newBalance, deducted, toDeduct: toDeduct_ };
+    } else {
+      let deducted = toDeduct;
+      let toDeduct_ = 0;
+      return { newBalance, deducted, toDeduct: toDeduct_ };
+    }
   }
 
   if (cusEntBalance.lte(0) && toDeduct > 0) {

server/src/trigger/updateBalanceTask.ts

@@ -1,6 +1,7 @@
 import {
   AllowanceType,
   AppEnv,
+  FullCusProduct,
   CusProductStatus,
   Entity,
   Event,
@@ -8,6 +9,8 @@ import {
   FullCustomerEntitlement,
   FullCustomerPrice,
   Organization,
+  FullCusEntWithFullCusProduct,
+  BillingType,
 } from "@autumn/shared";
 import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService.js";
 import { Customer, FeatureType } from "@autumn/shared";
@@ -27,12 +30,18 @@ import {
 } from "@/internal/features/creditSystemUtils.js";
 import {
   getCusEntMasterBalance,
+  getRelatedCusPrice,
+  getResetBalance,
   getTotalNegativeBalance,
 } from "@/internal/customers/cusProducts/cusEnts/cusEntUtils.js";
 import { entityFeatureIdExists } from "@/internal/api/entities/entityUtils.js";
 import { CusService } from "@/internal/customers/CusService.js";
 import { DrizzleCli } from "@/db/initDrizzle.js";
 import { findCusEnt } from "@/internal/customers/cusProducts/cusEnts/cusEntUtils/findCusEntUtils.js";
+import {
+  getBillingType,
+  getEntOptions,
+} from "@/internal/products/prices/priceUtils.js";
 
 // Decimal.set({ precision: 12 }); // 12 DP precision
 
@@ -181,19 +190,36 @@ export const performDeductionOnCusEnt = ({
   allowNegativeBalance = false,
   addAdjustment = false,
   setZeroAdjustment = false,
+  blockUsageLimit = true,
 }: {
-  cusEnt: FullCustomerEntitlement;
+  cusEnt: FullCusEntWithFullCusProduct;
   toDeduct: number;
   entityId?: string | null;
   allowNegativeBalance?: boolean;
   addAdjustment?: boolean;
   setZeroAdjustment?: boolean;
+  blockUsageLimit?: boolean;
 }) => {
   let newEntities = structuredClone(cusEnt.entities);
   let newBalance = structuredClone(cusEnt.balance);
   let newAdjustment = structuredClone(cusEnt.adjustment);
   let deducted = 0;
 
+  let cusProduct = cusEnt.customer_product;
+  let options = notNullish(cusProduct)
+    ? getEntOptions(cusProduct.options, cusEnt.entitlement)
+    : undefined;
+  let cusPrice = notNullish(cusProduct)
+    ? getRelatedCusPrice(cusEnt, cusProduct.customer_prices)
+    : undefined;
+  let resetBalance = notNullish(cusProduct)
+    ? getResetBalance({
+        options,
+        relatedPrice: cusPrice?.price,
+        entitlement: cusEnt.entitlement,
+      })
+    : cusEnt.entitlement.allowance || 0;
+
   if (entityFeatureIdExists({ cusEnt })) {
     if (nullish(entityId)) {
       // 1. If no entity ID, deduct from all
@@ -216,6 +242,9 @@ export const performDeductionOnCusEnt = ({
           cusEntBalance: new Decimal(entityBalance),
           toDeduct: toDeductCursor,
           allowNegativeBalance,
+          ent: cusEnt.entitlement,
+          resetBalance,
+          blockUsageLimit,
         });
 
         newEntities[entityId].balance = newEntityBalance!;
@@ -246,6 +275,9 @@ export const performDeductionOnCusEnt = ({
         cusEntBalance: new Decimal(currentEntityBalance!),
         toDeduct,
         allowNegativeBalance,
+        ent: cusEnt.entitlement,
+        resetBalance,
+        blockUsageLimit,
       });
 
       newEntities![entityId!]!.balance = newEntityBalance!;
@@ -271,6 +303,9 @@ export const performDeductionOnCusEnt = ({
       cusEntBalance: new Decimal(cusEnt.balance!),
       toDeduct,
       allowNegativeBalance,
+      ent: cusEnt.entitlement,
+      resetBalance,
+      blockUsageLimit,
     });
 
     newBalance = newBalance_;
@@ -295,7 +330,7 @@ export const deductAllowanceFromCusEnt = async ({
 }: {
   toDeduct: number;
   deductParams: DeductParams;
-  cusEnt: FullCustomerEntitlement;
+  cusEnt: FullCusEntWithFullCusProduct;
   featureDeductions: any;
   willDeductCredits?: boolean;
   setZeroAdjustment?: boolean;
@@ -411,7 +446,7 @@ export const deductFromUsageBasedCusEnt = async ({
 }: {
   toDeduct: number;
   deductParams: DeductParams;
-  cusEnts: FullCustomerEntitlement[];
+  cusEnts: FullCusEntWithFullCusProduct[];
   setZeroAdjustment?: boolean;
 }) => {
   const { db, feature, env, org, cusPrices, customer, entity } = deductParams;
@@ -422,7 +457,7 @@ export const deductFromUsageBasedCusEnt = async ({
     feature,
     entity,
     onlyUsageAllowed: true,
-  });
+  }) as FullCusEntWithFullCusProduct;
 
   if (!usageBasedEnt) {
     console.log(
@@ -431,12 +466,20 @@ export const deductFromUsageBasedCusEnt = async ({
     return;
   }
 
+  let cusPrice = getRelatedCusPrice(usageBasedEnt, cusPrices);
+  let billingType = cusPrice?.price
+    ? getBillingType(cusPrice?.price.config!)
+    : undefined;
+  let blockUsageLimit =
+    billingType === BillingType.InArrearProrated ? false : true;
+
   let { newBalance, newEntities, deducted } = performDeductionOnCusEnt({
     cusEnt: usageBasedEnt,
     toDeduct,
     allowNegativeBalance: true,
     setZeroAdjustment,
     entityId: entity?.id,
+    blockUsageLimit,
   });
 
   let oldGrpBalance = getTotalNegativeBalance({

server/src/utils/scriptUtils/constructItem.ts

@@ -17,7 +17,7 @@ export const constructFeatureItem = ({
 }: {
   featureId: string;
   includedUsage?: number;
-  interval?: ProductItemInterval;
+  interval?: ProductItemInterval | null;
   entityFeatureId?: string;
   isBoolean?: boolean;
 }) => {
@@ -47,6 +47,7 @@ export const constructPrepaidItem = ({
     on_increase: OnIncrease.ProrateImmediately,
     on_decrease: OnDecrease.ProrateImmediately,
   },
+  usageLimit,
 }: {
   featureId: string;
   price?: number;
@@ -54,6 +55,7 @@ export const constructPrepaidItem = ({
   includedUsage?: number;
   isOneOff?: boolean;
   config?: ProductItemConfig;
+  usageLimit?: number;
 }) => {
   let item: ProductItem = {
     feature_id: featureId,
@@ -66,6 +68,7 @@ export const constructPrepaidItem = ({
     included_usage: includedUsage,
 
     config,
+    usage_limit: usageLimit,
   };
 
   return item;
@@ -81,6 +84,7 @@ export const constructArrearItem = ({
     on_decrease: OnDecrease.ProrateImmediately,
   },
   entityFeatureId,
+  usageLimit,
 }: {
   featureId: string;
   includedUsage?: number;
@@ -88,6 +92,7 @@ export const constructArrearItem = ({
   billingUnits?: number;
   config?: ProductItemConfig;
   entityFeatureId?: string;
+  usageLimit?: number;
 }) => {
   let item: ProductItem = {
     feature_id: featureId,
@@ -99,6 +104,7 @@ export const constructArrearItem = ({
     reset_usage_when_enabled: true,
     config,
     entity_feature_id: entityFeatureId,
+    usage_limit: usageLimit,
   };
 
   return item;

server/tests/advanced/usageLimit/usageLimit1.ts

@@ -28,7 +28,7 @@ export let pro = constructProduct({
 
 const testCase = "usageLimit1";
 
-describe(`${chalk.yellowBright(`${testCase}: Testing entities`)}`, () => {
+describe(`${chalk.yellowBright(`${testCase}: Testing usage limits for entities`)}`, () => {
   let customerId = testCase;
   let autumn: AutumnInt = new AutumnInt({ version: APIVersion.v1_4 });
   let testClockId: string;
@@ -134,9 +134,6 @@ describe(`${chalk.yellowBright(`${testCase}: Testing entities`)}`, () => {
     });
     const customer = await autumn.customers.get(customerId);
 
-    console.log(check);
-    console.log(customer);
-
     expect(check.balance).to.equal(-2);
     // @ts-ignore
     expect(check.usage_limit).to.equal(userItem.usage_limit);

server/tests/advanced/usageLimit/usageLimit2.ts

@@ -2,33 +2,48 @@ import chalk from "chalk";
 import Stripe from "stripe";
 import { AutumnInt } from "@/external/autumn/autumnCli.js";
 import { initCustomer } from "@/utils/scriptUtils/initCustomer.js";
-import { APIVersion, AppEnv, ErrCode, Organization } from "@autumn/shared";
+import { APIVersion, AppEnv, LimitedItem, Organization } from "@autumn/shared";
 
 import { DrizzleCli } from "@/db/initDrizzle.js";
 import { setupBefore } from "tests/before.js";
 import { createProducts } from "tests/utils/productUtils.js";
 import { constructProduct } from "@/utils/scriptUtils/createTestProducts.js";
-import { constructArrearProratedItem } from "@/utils/scriptUtils/constructItem.js";
+import {
+  constructArrearItem,
+  constructFeatureItem,
+} from "@/utils/scriptUtils/constructItem.js";
 import { TestFeature } from "tests/setup/v2Features.js";
 import { addPrefixToProducts, runAttachTest } from "tests/attach/utils.js";
-import { expectAutumnError } from "tests/utils/expectUtils/expectErrUtils.js";
 import { expect } from "chai";
+import { timeout } from "@/utils/genUtils.js";
 
-const userItem = constructArrearProratedItem({
-  featureId: TestFeature.Users,
-  pricePerUnit: 50,
-  includedUsage: 0,
-  usageLimit: 2,
-});
+const messageItem = constructArrearItem({
+  featureId: TestFeature.Messages,
+  includedUsage: 100,
+  billingUnits: 1,
+  price: 0.5,
+  usageLimit: 500,
+}) as LimitedItem;
 
 export let pro = constructProduct({
-  items: [userItem],
+  items: [messageItem],
   type: "pro",
 });
 
-const testCase = "entity1";
+const addOnMessages = constructFeatureItem({
+  featureId: TestFeature.Messages,
+  interval: null,
+  includedUsage: 250,
+}) as LimitedItem;
 
-describe(`${chalk.yellowBright(`${testCase}: Testing entities`)}`, () => {
+const messageAddOn = constructProduct({
+  type: "one_off",
+  items: [addOnMessages],
+});
+
+const testCase = "usageLimit2";
+
+describe(`${chalk.yellowBright(`${testCase}: Testing usage limits, usage prices`)}`, () => {
   let customerId = testCase;
   let autumn: AutumnInt = new AutumnInt({ version: APIVersion.v1_4 });
   let testClockId: string;
@@ -47,13 +62,13 @@ describe(`${chalk.yellowBright(`${testCase}: Testing entities`)}`, () => {
     stripeCli = this.stripeCli;
 
     addPrefixToProducts({
-      products: [pro],
+      products: [pro, messageAddOn],
       prefix: testCase,
     });
 
     await createProducts({
       autumn,
-      products: [pro],
+      products: [pro, messageAddOn],
       customerId,
       db,
       orgId: org.id,
@@ -72,29 +87,6 @@ describe(`${chalk.yellowBright(`${testCase}: Testing entities`)}`, () => {
     testClockId = testClockId1!;
   });
 
-  const entities = [
-    {
-      id: "1",
-      name: "Entity 1",
-      feature_id: TestFeature.Users,
-    },
-    {
-      id: "2",
-      name: "Entity 2",
-      feature_id: TestFeature.Users,
-    },
-    {
-      id: "3",
-      name: "Entity 3",
-      feature_id: TestFeature.Users,
-    },
-    {
-      id: "4",
-      name: "Entity 4",
-      feature_id: TestFeature.Users,
-    },
-  ];
-
   it("should attach pro product", async function () {
     await runAttachTest({
       autumn,
@@ -106,40 +98,91 @@ describe(`${chalk.yellowBright(`${testCase}: Testing entities`)}`, () => {
       env,
     });
   });
-  it("should create more entities than the limit and hit error", async function () {
-    await expectAutumnError({
-      errCode: ErrCode.FeatureLimitReached,
-      func: async () => {
-        await autumn.entities.create(customerId, entities);
-      },
-    });
-  });
 
-  it("should create entities one by one, then hit usage limit", async function () {
-    await autumn.entities.create(customerId, entities[0]);
-    await autumn.entities.create(customerId, entities[1]);
+  let initialUsage =
+    messageItem.included_usage + messageItem.usage_limit! + 1000;
 
-    await expectAutumnError({
-      errCode: ErrCode.FeatureLimitReached,
-      func: async () => {
-        await autumn.entities.create(customerId, entities[2]);
-      },
-    });
-  });
-
-  it("should have correct check and get customer value", async function () {
-    const check = await autumn.check({
+  it("should track more messages than limit and not surpass", async function () {
+    await autumn.track({
       customer_id: customerId,
-      feature_id: TestFeature.Users,
+      feature_id: TestFeature.Messages,
+      value: initialUsage,
     });
-    expect(check.balance).to.equal(-2);
-    // @ts-ignore
-    expect(check.usage_limit).to.equal(userItem.usage_limit);
 
-    const customer = await autumn.customers.get(customerId);
+    await timeout(2000);
+
+    let check = await autumn.check({
+      customer_id: customerId,
+      feature_id: TestFeature.Messages,
+    });
+    let customer = await autumn.customers.get(customerId);
+
+    let expectedBalance = messageItem.included_usage - messageItem.usage_limit!;
+
+    expect(check.balance).to.equal(expectedBalance);
+    expect(check.allowed).to.equal(false);
     // @ts-ignore
-    expect(customer.features[TestFeature.Users].usage_limit).to.equal(
-      userItem.usage_limit,
+    expect(check.usage_limit!).to.equal(messageItem.usage_limit!);
+    // @ts-ignore
+    expect(customer.features[TestFeature.Messages].usage_limit).to.equal(
+      messageItem.usage_limit!,
+    );
+  });
+
+  it("should purchase add ons and have correct check results", async function () {
+    await autumn.attach({
+      customer_id: customerId,
+      product_id: messageAddOn.id,
+    });
+
+    let check = await autumn.check({
+      customer_id: customerId,
+      feature_id: TestFeature.Messages,
+    });
+    let customer = await autumn.customers.get(customerId);
+    let expectedBalance =
+      messageItem.included_usage -
+      messageItem.usage_limit! +
+      addOnMessages.included_usage;
+
+    expect(check.balance).to.equal(expectedBalance);
+    expect(check.allowed).to.equal(true);
+
+    // @ts-ignore
+    expect(check.usage_limit!).to.equal(
+      messageItem.usage_limit! + addOnMessages.included_usage,
+    );
+    // @ts-ignore
+    expect(customer.features[TestFeature.Messages].usage_limit).to.equal(
+      messageItem.usage_limit! + addOnMessages.included_usage,
+    );
+  });
+
+  it("should use up all add ons and have correct check results", async function () {
+    await autumn.track({
+      customer_id: customerId,
+      feature_id: TestFeature.Messages,
+      value: addOnMessages.included_usage + 500,
+    });
+
+    await timeout(2000);
+
+    let check = await autumn.check({
+      customer_id: customerId,
+      feature_id: TestFeature.Messages,
+    });
+    let customer = await autumn.customers.get(customerId);
+
+    let expectedBalance = messageItem.included_usage - messageItem.usage_limit!;
+    expect(check.balance).to.equal(expectedBalance);
+    expect(check.allowed).to.equal(false);
+    // @ts-ignore
+    expect(check.usage_limit!).to.equal(
+      messageItem.usage_limit! + addOnMessages.included_usage,
+    );
+    // @ts-ignore
+    expect(customer.features[TestFeature.Messages].usage_limit).to.equal(
+      messageItem.usage_limit! + addOnMessages.included_usage,
     );
   });
 });

server/tests/advanced/usageLimit/usageLimit3.ts

@@ -0,0 +1,152 @@
+import chalk from "chalk";
+import Stripe from "stripe";
+import { AutumnInt } from "@/external/autumn/autumnCli.js";
+import { initCustomer } from "@/utils/scriptUtils/initCustomer.js";
+import {
+  APIVersion,
+  AppEnv,
+  ErrCode,
+  LimitedItem,
+  Organization,
+} from "@autumn/shared";
+
+import { DrizzleCli } from "@/db/initDrizzle.js";
+import { setupBefore } from "tests/before.js";
+import { createProducts } from "tests/utils/productUtils.js";
+import { constructProduct } from "@/utils/scriptUtils/createTestProducts.js";
+import {
+  constructFeatureItem,
+  constructPrepaidItem,
+} from "@/utils/scriptUtils/constructItem.js";
+import { TestFeature } from "tests/setup/v2Features.js";
+import { addPrefixToProducts, runAttachTest } from "tests/attach/utils.js";
+import { expect } from "chai";
+import { timeout } from "@/utils/genUtils.js";
+import { expectAutumnError } from "tests/utils/expectUtils/expectErrUtils.js";
+
+const messageItem = constructPrepaidItem({
+  featureId: TestFeature.Messages,
+  includedUsage: 50,
+  billingUnits: 100,
+  price: 8,
+  usageLimit: 500,
+}) as LimitedItem;
+
+export let pro = constructProduct({
+  items: [messageItem],
+  type: "pro",
+});
+
+// const addOnMessages = constructFeatureItem({
+//   featureId: TestFeature.Messages,
+//   interval: null,
+//   includedUsage: 250,
+// }) as LimitedItem;
+
+// const messageAddOn = constructProduct({
+//   type: "one_off",
+//   items: [addOnMessages],
+// });
+
+const testCase = "usageLimit3";
+
+describe(`${chalk.yellowBright(`${testCase}: Testing usage limits for prepaid`)}`, () => {
+  let customerId = testCase;
+  let autumn: AutumnInt = new AutumnInt({ version: APIVersion.v1_4 });
+  let testClockId: string;
+  let db: DrizzleCli, org: Organization, env: AppEnv;
+  let stripeCli: Stripe;
+
+  let curUnix = new Date().getTime();
+
+  before(async function () {
+    await setupBefore(this);
+    const { autumnJs } = this;
+    db = this.db;
+    org = this.org;
+    env = this.env;
+
+    stripeCli = this.stripeCli;
+
+    addPrefixToProducts({
+      products: [pro],
+      prefix: testCase,
+    });
+
+    await createProducts({
+      autumn,
+      products: [pro],
+      customerId,
+      db,
+      orgId: org.id,
+      env,
+    });
+
+    const { testClockId: testClockId1 } = await initCustomer({
+      autumn: autumnJs,
+      customerId,
+      db,
+      org,
+      env,
+      attachPm: "success",
+    });
+
+    testClockId = testClockId1!;
+  });
+
+  it("should attach pro product with quantity exceeding usage limit and get an error", async function () {
+    expectAutumnError({
+      errCode: ErrCode.InvalidOptions,
+      func: async () => {
+        return await runAttachTest({
+          autumn,
+          customerId,
+          product: pro,
+          stripeCli,
+          db,
+          org,
+          env,
+          options: [
+            {
+              feature_id: TestFeature.Messages,
+              quantity: 600,
+            },
+          ],
+        });
+      },
+    });
+  });
+  it("should attach pro product and update quantity with quantity exceeding usage limit and get an error", async function () {
+    await autumn.attach({
+      customer_id: customerId,
+      product_id: pro.id,
+      options: [
+        {
+          feature_id: TestFeature.Messages,
+          quantity: 100,
+        },
+      ],
+    });
+
+    expectAutumnError({
+      errCode: ErrCode.InvalidOptions,
+      func: async () => {
+        return await runAttachTest({
+          autumn,
+          customerId,
+          product: pro,
+          stripeCli,
+          db,
+          org,
+          env,
+          options: [
+            {
+              feature_id: TestFeature.Messages,
+              quantity: 600,
+            },
+          ],
+        });
+      },
+    });
+  });
+});

server/tests/advanced/usageLimit/usageLimit4.ts

@@ -0,0 +1,117 @@
+import chalk from "chalk";
+import Stripe from "stripe";
+import { AutumnInt } from "@/external/autumn/autumnCli.js";
+import { initCustomer } from "@/utils/scriptUtils/initCustomer.js";
+import {
+  APIVersion,
+  AppEnv,
+  ErrCode,
+  LimitedItem,
+  Organization,
+} from "@autumn/shared";
+
+import { DrizzleCli } from "@/db/initDrizzle.js";
+import { setupBefore } from "tests/before.js";
+import { createProducts } from "tests/utils/productUtils.js";
+import { constructProduct } from "@/utils/scriptUtils/createTestProducts.js";
+import {
+  constructArrearProratedItem,
+  constructFeatureItem,
+  constructPrepaidItem,
+} from "@/utils/scriptUtils/constructItem.js";
+import { TestFeature } from "tests/setup/v2Features.js";
+import { addPrefixToProducts, runAttachTest } from "tests/attach/utils.js";
+import { expect } from "chai";
+import { timeout } from "@/utils/genUtils.js";
+import { expectAutumnError } from "tests/utils/expectUtils/expectErrUtils.js";
+
+const messageItem = constructArrearProratedItem({
+  featureId: TestFeature.Users,
+  includedUsage: 1,
+  pricePerUnit: 10,
+  usageLimit: 3,
+}) as LimitedItem;
+
+export let pro = constructProduct({
+  items: [messageItem],
+  type: "pro",
+});
+
+const testCase = "usageLimit4";
+
+describe(`${chalk.yellowBright(`${testCase}: Testing usage limits for cont use item`)}`, () => {
+  let customerId = testCase;
+  let autumn: AutumnInt = new AutumnInt({ version: APIVersion.v1_4 });
+  let testClockId: string;
+  let db: DrizzleCli, org: Organization, env: AppEnv;
+  let stripeCli: Stripe;
+
+  let curUnix = new Date().getTime();
+
+  before(async function () {
+    await setupBefore(this);
+    const { autumnJs } = this;
+    db = this.db;
+    org = this.org;
+    env = this.env;
+
+    stripeCli = this.stripeCli;
+
+    addPrefixToProducts({
+      products: [pro],
+      prefix: testCase,
+    });
+
+    await createProducts({
+      autumn,
+      products: [pro],
+      customerId,
+      db,
+      orgId: org.id,
+      env,
+    });
+
+    const { testClockId: testClockId1 } = await initCustomer({
+      autumn: autumnJs,
+      customerId,
+      db,
+      org,
+      env,
+      attachPm: "success",
+    });
+
+    testClockId = testClockId1!;
+  });
+
+  it("should attach pro product with quantity exceeding usage limit and get an error", async function () {
+    await runAttachTest({
+      autumn,
+      customerId,
+      product: pro,
+      stripeCli,
+      db,
+      org,
+      env,
+    });
+  });
+  it("should attach pro product and update quantity with quantity exceeding usage limit and get an error", async function () {
+    await expectAutumnError({
+      errCode: ErrCode.InvalidInputs,
+      func: async () => {
+        await autumn.track({
+          customer_id: customerId,
+          feature_id: TestFeature.Users,
+          value: messageItem.usage_limit! + 1,
+        });
+      },
+    });
+
+    const check = await autumn.check({
+      customer_id: customerId,
+      feature_id: TestFeature.Users,
+    });
+
+    expect(check.balance).to.equal(0);
+    expect(check.allowed).to.equal(true);
+  });
+});

server/tests/attach/updateEnts/updateEnts1.ts

@@ -132,6 +132,7 @@ describe(`${chalk.yellowBright(`${testCase}: Testing update ents (changing inclu
       ],
     });
   });
+  return;
 
   it("should have correct invoice  next cycle", async function () {
     const invoiceTotal = await getExpectedInvoiceTotal({

shared/enums/ErrCode.ts

@@ -19,6 +19,7 @@ export const ErrCode = {
   InvalidInputs: "invalid_inputs",
   InvalidRequest: "invalid_request",
   InvalidExpand: "invalid_expand",
+  InvalidOptions: "invalid_options",
 
   // Org
   OrgNotFound: "org_not_found",

shared/models/productV2Models/productItemModels/featurePriceItem.ts

@@ -17,6 +17,7 @@ export const FeaturePriceItemSchema = ProductItemSchema.pick({
   billing_units: true,
 
   reset_usage_when_enabled: true,
+  usage_limit: true,
 }).extend({
   feature_id: z.string().nonempty(),
   included_usage: z.number().nonnegative().nullish(),

vite/src/views/products/product/product-item/product-item-config/advanced-config/AdvancedItemConfig.tsx

@@ -37,7 +37,7 @@ export const AdvancedItemConfig = () => {
 
       <div
         className={`overflow-hidden transition-all duration-150 ease-out ${
-          isOpen ? "max-h-60 opacity-100 mt-2" : "max-h-0 opacity-0"
+          isOpen ? "max-h-72 opacity-100 mt-2" : "max-h-0 opacity-0"
         }`}
       >
         <div className="flex flex-col gap-4 p-4 bg-stone-100">
@@ -54,8 +54,8 @@ export const AdvancedItemConfig = () => {
             disabled={usageType === FeatureUsageType.Continuous}
           />
 
-        <div className="relative flex flex-row items-center justify-between gap-3 min-h-[35px]">
-          <ToggleButton
+          <div className="relative flex flex-row items-center gap-3 min-h-[35px]">
+            <ToggleButton
               value={item.usage_limit != null}
               setValue={() => {
                 let usage_limit;
@@ -73,22 +73,21 @@ export const AdvancedItemConfig = () => {
               className="text-t3 h-fit"
             />
 
-              {item.usage_limit != null && (
-                <Input
-                  type="number"
-                  value={item.usage_limit || ""}
-                  className="ml-5"
-                  onChange={(e) => {
-                    setItem({
-                      ...item,
-                      usage_limit: parseInt(e.target.value),
-                    });
-                  }}
-                  placeholder="Enter usage limit"
-                />
-              )}
-            </div>
-        
+            {item.usage_limit != null && (
+              <Input
+                type="number"
+                value={item.usage_limit || ""}
+                className="ml-5 w-25"
+                onChange={(e) => {
+                  setItem({
+                    ...item,
+                    usage_limit: parseInt(e.target.value),
+                  });
+                }}
+                placeholder="eg. 100"
+              />
+            )}
+          </div>
 
           {showProrationConfig && (
             <>
@@ -98,8 +97,6 @@ export const AdvancedItemConfig = () => {
           )}
           {/* <div className="flex flex-col gap-2"></div>
           <div className="flex gap-2"></div> */}
-
-          
         </div>
       </div>
     </div>