fix: updated create and remove entity

package-lock.json

@@ -9886,6 +9886,18 @@
       "integrity": "sha512-l3nz3euub2QMg5ouu5U09Ew9Wf6/wQ8I++ch1loQ0ljmzhmfZYrH9fflS22i/PQEvsPvxCwxgz5q7UB8K1JO4Q==",
       "license": "MIT"
     },
+    "node_modules/csv-parser": {
+      "version": "3.2.0",
+      "resolved": "https://registry.npmjs.org/csv-parser/-/csv-parser-3.2.0.tgz",
+      "integrity": "sha512-fgKbp+AJbn1h2dcAHKIdKNSSjfp43BZZykXsCjzALjKy80VXQNHPFJ6T9Afwdzoj24aMkq8GwDS7KGcDPpejrA==",
+      "license": "MIT",
+      "bin": {
+        "csv-parser": "bin/csv-parser"
+      },
+      "engines": {
+        "node": ">= 10"
+      }
+    },
     "node_modules/currency-symbol-map": {
       "version": "5.1.0",
       "resolved": "https://registry.npmjs.org/currency-symbol-map/-/currency-symbol-map-5.1.0.tgz",
@@ -21282,6 +21294,7 @@
         "cors": "^2.8.5",
         "cron": "^3.5.0",
         "csv-parse": "^5.6.0",
+        "csv-parser": "^3.2.0",
         "currency-symbol-map": "^5.1.0",
         "date-fns": "^4.1.0",
         "decimal.js": "^10.5.0",

server/package.json

@@ -51,6 +51,7 @@
     "cors": "^2.8.5",
     "cron": "^3.5.0",
     "csv-parse": "^5.6.0",
+    "csv-parser": "^3.2.0",
     "currency-symbol-map": "^5.1.0",
     "date-fns": "^4.1.0",
     "decimal.js": "^10.5.0",

server/src/external/redis/redisUtils.ts

@@ -49,6 +49,58 @@ export const handleAttachRaceCondition = async ({
   }
 };
 
+export const handleCustomerRaceCondition = async ({
+  action,
+  customerId,
+  orgId,
+  env,
+  res,
+  logger,
+}: {
+  action: any;
+  customerId: string;
+  orgId: string;
+  env: string;
+  res: any;
+  logger: any;
+}) => {
+  const redisConn = await QueueManager.getConnection({ useBackup: false });
+  try {
+    const lockKey = `${action}_${customerId}_${orgId}_${env}`;
+    const existingLock = await redisConn.get(lockKey);
+    if (existingLock) {
+      throw new RecaseError({
+        message: `Action ${action} already running for customer ${customerId}, try again in a few seconds`,
+        code: ErrCode.InvalidRequest,
+        statusCode: 400,
+      });
+    }
+    // Create lock with 5 second timeout
+    await redisConn.set(lockKey, "1", "PX", 5000, "NX");
+
+    let originalJson = res.json;
+    res.json = async function (body: any) {
+      try {
+        await clearLock({ lockKey, logger });
+      } catch (error) {
+        logger.warn("❗️❗️ Error clearing lock");
+        logger.warn(error);
+      }
+      originalJson.call(this, body);
+    };
+
+    return lockKey;
+  } catch (error) {
+    if (error instanceof RecaseError) {
+      throw error;
+    }
+
+    logger.warn("❗️❗️ Error acquiring lock");
+    logger.warn(error);
+    return null;
+  }
+};
+
 export const clearLock = async ({
   lockKey,
   logger,

server/src/external/stripe/stripePriceUtils.ts

@@ -11,6 +11,7 @@ import {
   AllowanceType,
   EntitlementWithFeature,
   Feature,
+  FullCustomerEntitlement,
 } from "@autumn/shared";
 
 import RecaseError from "@/utils/errorUtils.js";
@@ -39,7 +40,7 @@ import {
   createStripeOneOffTieredProduct,
 } from "./createStripePrice.js";
 
-import { getExistingUsageFromCusProducts } from "@/internal/customers/entitlements/cusEntUtils.js";
+import { getCusEntMasterBalance, getExistingUsageFromCusProducts } from "@/internal/customers/entitlements/cusEntUtils.js";
 import {
   priceToInArrearProrated,
   priceToUsageInAdvance,
@@ -266,8 +267,11 @@ export const getStripeSubItems = async ({
       const existingUsage = getExistingUsageFromCusProducts({
         entitlement: priceEnt,
         cusProducts: attachParams.cusProducts,
+        entities: attachParams.entities,
       });
 
+
+
       if (
         billingType == BillingType.UsageInArrear ||
         billingType == BillingType.InArrearProrated ||

server/src/index.ts

@@ -19,6 +19,7 @@ import { AppEnv } from "@autumn/shared";
 import { createSupabaseClient } from "./external/supabaseUtils.js";
 import { createLogtail } from "./external/logtail/logtailUtils.js";
 import { format } from "date-fns";
+import { handleRequestError } from "./utils/errorUtils.js";
 
 const init = async () => {
   const app = express();
@@ -91,3 +92,14 @@ const init = async () => {
 };
 
 init();
+
+process.on("unhandledRejection", (reason, promise) => {
+  try {
+    const logtail = createLogtail();
+    logtail.error("❗️❗️❗️ UNHANDLED REJECTION");
+    logtail.error(reason);
+    logtail.flush();
+  } catch (error) {
+    console.log("Unhandled rejection", error);
+  }
+});

server/src/internal/api/customers/cusRouter.ts

@@ -29,6 +29,7 @@ import { handleAddCouponToCus } from "./handlers/handleAddCouponToCus.js";
 
 import { handlePostCustomerRequest } from "./handlers/handleCreateCustomer.js";
 import { notNullish } from "@/utils/genUtils.js";
+import { entityRouter } from "../entities/entityRouter.js";
 
 export const cusRouter = Router();
 
@@ -388,3 +389,5 @@ cusRouter.get("/:customer_id/billing_portal", async (req: any, res: any) => {
 // Invoice
 
 cusRouter.post("/:customer_id/coupons/:coupon_id", handleAddCouponToCus);
+
+cusRouter.use("/:customer_id/entities", entityRouter);

server/src/internal/api/customers/handlers/handleCreateCustomer.ts

@@ -174,6 +174,7 @@ export const createNewCustomer = async ({
         optionsList: [],
         cusProducts: [],
         invoiceOnly: true,
+        entities: [],
       },
       fromRequest: false,
     });

server/src/internal/api/entities/EntityService.ts

@@ -5,11 +5,13 @@ export class EntityService {
   static async getById({
     sb,
     entityId,
+    internalCustomerId,
     orgId,
     env,
   }: {
     sb: SupabaseClient;
     entityId: string;
+    internalCustomerId?: string;
     orgId: string;
     env: string;
   }) {
@@ -17,6 +19,7 @@ export class EntityService {
       .from("entities")
       .select("*")
       .eq("id", entityId)
+      .eq("internal_customer_id", internalCustomerId)
       .eq("org_id", orgId)
       .eq("env", env)
       .single();
@@ -38,26 +41,31 @@ export class EntityService {
     return data;
   }
 
-  static async getInIds({
+  static async get({
     sb,
-    ids,
     orgId,
     internalFeatureId,
+    internalCustomerId,
     env,
   }: {
     sb: SupabaseClient;
-    ids: string[];
     orgId: string;
     env: string;
     internalFeatureId?: string;
+    internalCustomerId?: string;
   }) {
-    const { data, error } = await sb
+    let query = sb
       .from("entities")
       .select("*")
-      .in("id", ids)
       .eq("org_id", orgId)
       .eq("env", env)
-      .eq("internal_feature_id", internalFeatureId);
+      .eq("internal_customer_id", internalCustomerId);
+
+    if (internalFeatureId) {
+      query = query.eq("internal_feature_id", internalFeatureId);
+    }
+
+    const { data, error } = await query;
 
     if (error) {
       throw error;

server/src/internal/api/entities/entityRouter.ts

@@ -2,7 +2,7 @@ import { Router } from "express";
 import { handleCreateEntity } from "./handleCreateEntity.js";
 import { handleDeleteEntity } from "./handleDeleteEntity.js";
 
-export const entityRouter = Router();
+export const entityRouter = Router({ mergeParams: true });
 
 // 1. Create entity
 entityRouter.post("", handleCreateEntity);

server/src/internal/api/entities/handleCreateEntity.ts

@@ -6,9 +6,16 @@ import { OrgService } from "@/internal/orgs/OrgService.js";
 import RecaseError, { handleRequestError } from "@/utils/errorUtils.js";
 import { getLinkedCusEnt } from "./entityUtils.js";
 import { EntityService } from "./EntityService.js";
-import { AppEnv, CusProductStatus, Entity } from "@autumn/shared";
+import { AppEnv, CusProductStatus, Customer, Entity, FullCusProduct, FullCustomerEntitlement, FullCustomerPrice, Product } from "@autumn/shared";
 import { generateId } from "@/utils/genUtils.js";
 import { adjustAllowance } from "@/trigger/adjustAllowance.js";
+import { getActiveCusProductStatuses } from "@/utils/constants.js";
+import { isTrialing } from "@/internal/customers/products/cusProductUtils.js";
+import {Decimal} from "decimal.js";
+import { getPriceForOverage } from "@/internal/prices/priceUtils.js";
+import { Logger } from "@slack/web-api";
+import Stripe from "stripe";
+import { getCusEntMasterBalance } from "@/internal/customers/entitlements/cusEntUtils.js";
 
 export const constructEntity = ({
   inputEntity,
@@ -55,25 +62,30 @@ export const getEntityToAction = ({
   // 1. GET ENTITY TO ACTION
   let entityToAction: any = {};
   let createCount = 0;
+  let replacedEntities: string[] = [];
   for (const inputEntity of inputEntities) {
     let curEntity = existingEntities.find((e: any) => e.id === inputEntity.id);
 
-    if (curEntity && !curEntity.deleted) {
+    if (curEntity && curEntity.deleted) {
       // Replace
       entityToAction[inputEntity.id] = {
         action: "replace",
         replace: curEntity,
         entity: inputEntity,
       };
+      replacedEntities.push(curEntity.id);
+      continue;
     }
 
     let replaced = false;
     for (const entity of existingEntities) {
       if (
         entity.deleted &&
-        !Object.keys(entityToAction).some((id) => id === entity.id)
+        !replacedEntities.includes(entity.id)
       ) {
         replaced = true;
+        replacedEntities.push(entity.id);
+
         entityToAction[inputEntity.id] = {
           action: "replace",
           replace: entity,
@@ -87,7 +99,7 @@ export const getEntityToAction = ({
       // Create
       entityToAction[inputEntity.id] = {
         action: "create",
-        inputEntity,
+        entity: inputEntity,
       };
       createCount++;
     }
@@ -121,12 +133,106 @@ export const getEntityToAction = ({
   return entityToAction;
 };
 
+// export const payForEntitiesImmediately = async ({
+//   sb,
+//   env,
+//   org,
+//   cusProduct,
+//   cusEnt,
+//   cusPrice,
+//   logger,
+//   oldUsage,
+//   createdNumber,
+//   stripeCli,
+//   product,
+//   customer,
+// }:{
+//   cusProduct: FullCusProduct;
+//   cusEnt: FullCustomerEntitlement;
+//   cusPrice: FullCustomerPrice;
+//   logger: Logger;
+//   oldUsage: number;
+//   createdNumber: number;
+//   stripeCli: Stripe;
+//   product: Product;
+//   customer: Customer;
+// }) => {
+//   if (!isTrialing(cusProduct as FullCusProduct)) {
+//     // let entitlement = cusEnt.entitlement;
+//     // let newUsage = entitlement.allowance! - newBalance;
+//     // let oldUsage = entitlement.allowance! - originalBalance;
+//     // newUsage = newUsage - (replacedCount || 0);
+
+//     // let newAmount = getPriceForOverage(cusPrice.price, newUsage);
+//     // let oldAmount = getPriceForOverage(cusPrice.price, oldUsage);
+
+//     const stripeAmount = new Decimal(newAmount)
+//       .sub(oldAmount)
+//       .mul(100)
+//       .round()
+//       .toNumber();
+
+//     logger.info(`   - Stripe amount: ${stripeAmount}`);
+
+//     if (stripeAmount > 0) {
+//       const invoice = await stripeCli.invoices.create({
+//         customer: customer.processor.id,
+//         auto_advance: false,
+//         subscription: sub.id,
+//       });
+
+//       await stripeCli.invoiceItems.create({
+//         customer: customer.processor.id,
+//         invoice: invoice.id,
+//         quantity: 1,
+//         description: `${product!.name} - ${
+//           affectedFeature.name
+//         } x ${Math.round(newUsage - oldUsage)}`,
+
+//         price_data: {
+//           product: config.stripe_product_id!,
+//           unit_amount: stripeAmount,
+//           currency: org.default_currency,
+//         },
+//       });
+      
+
+//       const { paid, error } = await payForInvoice({
+//         fullOrg: org,
+//         env,
+//         customer,
+//         invoice,
+//         logger,
+//       });
+
+      
+//       // console.log("Invoice paid result:", paid, error);
+//       const latestInvoice = await stripeCli.invoices.retrieve(invoice.id, {
+//         ...getInvoiceExpansion()
+//       });
+
+//       await InvoiceService.createInvoiceFromStripe({
+//         sb,
+//         stripeInvoice: latestInvoice,
+//         internalCustomerId: customer.internal_id,
+//         org,
+//         productIds: [product!.id],
+//         internalProductIds: [product!.internal_id],
+//       });
+
+//       if (!paid) {
+//         logger.warn("❗️ Failed to pay for invoice!");
+//       }
+//     }
+//   }
+// };
+
 export const handleCreateEntity = async (req: any, res: any) => {
   try {
     // Create entity!
 
     const { sb, env, orgId, logtail: logger } = req;
-    const { customer_id, feature_id, entity: inputEntities } = req.body;
+    const { customer_id } = req.params;
 
     let [customer, features, org] = await Promise.all([
       CusService.getByIdOrInternalId({
@@ -139,6 +245,22 @@ export const handleCreateEntity = async (req: any, res: any) => {
       OrgService.getFromReq(req),
     ]);
 
+    let inputEntities: any[] = [];
+    if (Array.isArray(req.body)) {
+      inputEntities = req.body;
+    } else {
+      inputEntities = [req.body];
+    }
+
+    let featureIds = [...new Set(inputEntities.map((e: any) => e.feature_id))];
+    if (featureIds.length > 1) {
+      throw new RecaseError({
+        message: "Multiple features not supported",
+        code: "MULTIPLE_FEATURES_NOT_SUPPORTED",
+      });
+    }
+
+    let feature_id = featureIds[0];
     let feature = features.find((f: any) => f.id === feature_id);
 
     let cusProducts = await CusService.getFullCusProducts({
@@ -146,21 +268,25 @@ export const handleCreateEntity = async (req: any, res: any) => {
       internalCustomerId: customer.internal_id,
       withProduct: true,
       withPrices: true,
-      inStatuses: [CusProductStatus.Active],
+      inStatuses: getActiveCusProductStatuses(),
       logger,
     });
 
     // Fetch existing
-    let existingEntities = await EntityService.getInIds({
+    let existingEntities = await EntityService.get({
       sb,
-      ids: inputEntities.map((e: any) => e.id),
       orgId,
       env,
       internalFeatureId: feature.internal_id,
+      internalCustomerId: customer.internal_id,
     });
 
+
+    console.log("existingEntities", existingEntities.map((e: any) => `${e.id} - ${e.name}, deleted: ${e.deleted}`));
+    
+
     for (const entity of existingEntities) {
-      if (entity && !entity.deleted) {
+      if (inputEntities.some((e: any) => e.id === entity.id) && !entity.deleted) {
         throw new RecaseError({
           message: `Entity ${entity.id} already exists`,
           code: "ENTITY_ALREADY_EXISTS",
@@ -171,6 +297,8 @@ export const handleCreateEntity = async (req: any, res: any) => {
       }
     }
     
+
+
     const entityToAction = getEntityToAction({
       inputEntities,
       existingEntities,
@@ -179,39 +307,12 @@ export const handleCreateEntity = async (req: any, res: any) => {
       cusProducts,
     });
 
-    // 3. CREATE ENTITIES
-    for (const id in entityToAction) {
-      let { action, inputEntity, replace } = entityToAction[id];
 
-      // Create and add to customer entitlement?
-      if (action === "create") {
-        await EntityService.insert({
-          sb,
-          data: constructEntity({
-            inputEntity,
-            feature,
-            internalCustomerId: customer.internal_id,
-            orgId,
-            env,
-          }),
-        });
-      } else if (action === "replace") {
-        await EntityService.update({
-          sb,
-          internalId: replace.internal_id,
-          update: {
-            deleted: false,
-          },
-        });
-      }
-    }
-    logger.info(`  Created / replaced entities!`);
     
-    // 4. CREATE LINKED CUSTOMER ENTITLEMENTS
+    // 3. CREATE LINKED CUSTOMER ENTITLEMENTS
     for (const cusProduct of cusProducts) {
       let cusEnts = cusProduct.customer_entitlements;
       let product = cusProduct.product;
-
       let cusEnt = cusEnts.find(
         (e: any) => e.entitlement.feature.id === feature_id
       );
@@ -225,14 +326,52 @@ export const handleCreateEntity = async (req: any, res: any) => {
         (e: any) => e.entitlement.entity_feature_id === feature.id
       );
 
+      // 1. Pay for new seats
+      let replacedCount = Object.keys(entityToAction).filter(
+        (id) => entityToAction[id].action === "replace"
+      ).length;
+      let newCount = Object.keys(entityToAction).filter(
+        (id) => entityToAction[id].action === "create"
+      ).length;
+
+      let { unused } = getCusEntMasterBalance({
+        cusEnt,
+        entities: existingEntities,
+      });
+      
+      // const originalBalance = cusEnt.balance - (replacedCount || 0) + (unused || 0);
+      // const newBalance = cusEnt.balance - (newCount + replacedCount) + (unused || 0);
+      const originalBalance = cusEnt.balance + (unused || 0);
+      const newBalance = cusEnt.balance - (newCount + replacedCount) + (unused || 0);
+
+      // console.log("originalBalance", originalBalance);
+      // console.log("newBalance", newBalance);
+      // console.log("Replaced count", replacedCount);
+      // throw new Error("test");
+
+      await adjustAllowance({
+        sb,
+        env,
+        org,
+        cusPrices: cusProducts.flatMap((p: any) => p.customer_prices),
+        customer,
+        affectedFeature: feature,
+        cusEnt: { ...cusEnt, customer_product: cusProduct },
+        originalBalance,
+        newBalance,
+        deduction: newCount + replacedCount,
+        product,
+        replacedCount,
+      });
+
+      await req.pg.query(
+        `UPDATE customer_entitlements SET balance = balance - $1 WHERE id = $2`,
+        [newCount, cusEnt.id]
+      );
+
       // For each linked feature, create customer entitlement for entity...
       for (const linkedCusEnt of linkedCusEnts) {
-        // let linkedCusEnt = getLinkedCusEnt({
-        //   linkedFeature,
-        //   cusEnts,
-        // });
 
-        // console.log("linkedCusEnt", linkedCusEnt?.entitlement.feature.id);
         let allowance = linkedCusEnt?.entitlement.allowance;
         let newEntities = linkedCusEnt?.entities || {};
 
@@ -241,13 +380,19 @@ export const handleCreateEntity = async (req: any, res: any) => {
 
           if (entityAction.action === "create") {
             newEntities[entity.id] = {
+              id: entity.id,
               balance: allowance,
               adjustment: 0,
             };
           } else if (entityAction.action === "replace") {
             let tmp = newEntities[entityAction.replace.id];
             delete newEntities[entityAction.replace.id];
-            newEntities[entity.id] = tmp;
+            newEntities[entity.id] = {
+              id: entity.id,
+              ...tmp,
+            };
+
+            
           }
         }
         
@@ -257,36 +402,38 @@ export const handleCreateEntity = async (req: any, res: any) => {
           updates: { entities: newEntities },
         });
       }
+    }
 
-      // 2. Update main customer entitlement (decrement balance)
+    // 4. CREATE ENTITIES
+    for (const id in entityToAction) {
+      let { action, entity, replace } = entityToAction[id];
 
-      let replacedCount = Object.keys(entityToAction).filter(
+      // Create and add to customer entitlement?
-        (id) => entityToAction[id].action === "replace"
+      if (action === "create") {
-      ).length;
+        await EntityService.insert({
-      let newCount = Object.keys(entityToAction).filter(
-        (id) => entityToAction[id].action === "create"
-      ).length;
-
-      await req.pg.query(
-        `UPDATE customer_entitlements SET balance = balance - $1 WHERE id = $2`,
-        [newCount, cusEnt.id]
-      );
-
-      adjustAllowance({
           sb,
+          data: constructEntity({
+            inputEntity: entity,
+            feature,
+            internalCustomerId: customer.internal_id,
+            orgId,
             env,
-        org,
+          }),
-        cusPrices: cusProducts.flatMap((p: any) => p.customer_prices),
+        });
-        customer,
+      } else if (action === "replace") {
-        affectedFeature: feature,
+        await EntityService.update({
-        cusEnt: { ...cusEnt, customer_product: cusProduct },
+          sb,
-        originalBalance: cusEnt.balance,
+          internalId: replace.internal_id,
-        newBalance: cusEnt.balance - (newCount + replacedCount),
+          update: {
-        deduction: newCount + replacedCount,
+            id: entity.id,
-        product,
+            name: entity.name,
-        replacedCount,
+            deleted: false,
+          },
         });
       }
+    }
+    logger.info(`  Created / replaced entities!`);
+
 
     res.status(200).json({
       success: true,

server/src/internal/api/entities/handleDeleteEntity.ts

@@ -5,28 +5,66 @@ import { ErrCode } from "@autumn/shared";
 import { CusService } from "@/internal/customers/CusService.js";
 import { adjustAllowance } from "@/trigger/adjustAllowance.js";
 import { OrgService } from "@/internal/orgs/OrgService.js";
+import { handleCustomerRaceCondition } from "@/external/redis/redisUtils.js";
+import { getCusEntMasterBalance } from "@/internal/customers/entitlements/cusEntUtils.js";
 
 export const handleDeleteEntity = async (req: any, res: any) => {
   try {
     const { orgId, env, logtail: logger, sb } = req;
-    const entityId = req.params.entity_id;
+    const { customer_id, entity_id } = req.params;
 
-    const entity = await EntityService.getById({
+    
+    await handleCustomerRaceCondition({
+      action: "entity",
+      customerId: customer_id,
+      orgId,
+      env,
+      res,
+      logger,
+    });
+
+    // console.log("Handling race condition for:", customer_id);
+    // console.log("Customer ID:", customer_id);
+    // console.log("Entity ID:", entity_id);
+
+
+    const customer = await CusService.getById({
       sb: req.sb,
-      entityId,
+      id: customer_id,
+      orgId: req.orgId,
+      env: req.env,
+      logger,
+    });
+
+    if (!customer) {
+      throw new RecaseError({
+        message: `Customer ${customer_id} not found`,
+        code: ErrCode.CustomerNotFound,
+        statusCode: StatusCodes.NOT_FOUND,
+      });
+    }
+
+    const existingEntities = await EntityService.get({
+      sb: req.sb,
+      internalCustomerId: customer.internal_id,
       orgId: req.orgId,
       env: req.env,
     });
     
+
+    const entity = existingEntities.find(
+      (e: any) => e.id === entity_id
+    );
+
     if (!entity) {
       throw new RecaseError({
-        message: `Entity ${entityId} not found`,
+        message: `Entity ${entity_id} not found`,
         code: ErrCode.EntityNotFound,
         statusCode: StatusCodes.NOT_FOUND,
       });
     } else if (entity.deleted) {
       throw new RecaseError({
-        message: `Entity ${entityId} already deleted`,
+        message: `Entity ${entity_id} already deleted`,
         code: ErrCode.EntityAlreadyDeleted,
         statusCode: StatusCodes.BAD_REQUEST,
       });
@@ -40,18 +78,14 @@ export const handleDeleteEntity = async (req: any, res: any) => {
       logger,
     });
 
-    const [customer, org] = await Promise.all([
+    const org = await OrgService.getFromReq(req);
-      CusService.getByInternalId({
-        sb: req.sb,
-        internalId: entity.internal_customer_id,
-      }),
-      OrgService.getFromReq(req),
-    ]);
 
     for (const cusProduct of cusProducts) {
       let cusEnts = cusProduct.customer_entitlements;
       let product = cusProduct.product;
 
+      
+
       let cusEnt = cusEnts.find(
         (e: any) =>
           e.entitlement.feature.internal_id === entity.internal_feature_id
@@ -61,8 +95,14 @@ export const handleDeleteEntity = async (req: any, res: any) => {
         continue;
       }
 
-      let newBalance = cusEnt.balance + 1;
+      let {unused} = getCusEntMasterBalance({
-      adjustAllowance({
+        cusEnt,
+        entities: existingEntities,
+      });
+      
+      let newBalance = (cusEnt.balance + 1) + (unused || 0);
+      
+      await adjustAllowance({
         sb,
         env,
         org,
@@ -70,7 +110,7 @@ export const handleDeleteEntity = async (req: any, res: any) => {
         customer,
         affectedFeature: cusEnt.entitlement.feature,
         cusEnt: { ...cusEnt, customer_product: cusProduct },
-        originalBalance: cusEnt.balance,
+        originalBalance: cusEnt.balance + (unused || 0),
         newBalance,
         deduction: 1,
         product,
@@ -85,38 +125,9 @@ export const handleDeleteEntity = async (req: any, res: any) => {
       },
     });
 
-    // If not X, delete entity AND entitlements...
+    logger.info(` ✅ Finished deleting entity ${entity_id}`);
 
-    // await EntityService.update({
-    //   sb: req.sb,
-    //   internalId: entity.internal_id,
-    //   update: {
-    //     deleted: true,
-    //   },
-    // });
 
-    // // console.log("Deleting entity:", entity);
-    // const customer = await Promise.all([
-    //   CusService.getByInternalId({
-    //     sb: req.sb,
-    //     internalId: entity.internal_customer_id,
-    //   }),
-    //   CusService.getFullCusProducts({
-    //     sb: req.sb,
-    //     internalCustomerId: entity.internal_customer_id,
-    //     withProduct: true,
-    //     withPrices: true,
-    //     logger,
-    //   }),
-    // ]);
-
-    // if (!customer) {
-    //   throw new RecaseError({
-    //     message: `Customer ${entity.internal_customer_id} not found`,
-    //     code: ErrCode.CustomerNotFound,
-    //     statusCode: StatusCodes.NOT_FOUND,
-    //   });
-    // }
 
     res.status(200).json({
       success: true,

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

@@ -6,6 +6,7 @@ import {
   CusProductStatus,
   EntitlementWithFeature,
   FeatureOptions,
+  FeatureType,
   FullCustomerEntitlement,
   Organization,
   Price,
@@ -151,6 +152,44 @@ export const updateOneTimeCusProduct = async ({
     } 
   }
 
+
+  // Handle adding quantity to base entitlements if cus product purchased multiple times.
+  for (const entitlement of attachParams.entitlements) {
+    const relatedPrice = getEntRelatedPrice(entitlement, attachParams.prices);
+    const feature = entitlement.feature;
+
+    if (relatedPrice || feature.type == FeatureType.Boolean || entitlement.allowance_type === AllowanceType.Unlimited) {
+      continue;
+    }
+
+    const newOptionIndex = newOptionsList.findIndex(
+      (o) => o.internal_feature_id === entitlement.internal_feature_id
+    );
+
+    if (newOptionIndex === -1) {
+      // Get existing option
+      const existingOption = existingCusProduct.options.find(
+        (o) => o.internal_feature_id === entitlement.internal_feature_id
+      );
+
+      if (existingOption) {
+        newOptionsList.push({
+          feature_id: entitlement.feature.id,
+          quantity: (existingOption?.quantity || 0) + 1,
+          internal_feature_id: entitlement.internal_feature_id,
+        });
+      } else {
+        newOptionsList.push({
+          feature_id: entitlement.feature.id,
+          quantity: 2,
+          internal_feature_id: entitlement.internal_feature_id,
+        });
+      }
+    }
+  }
+
+
+
   await CusProductService.update({
     sb,
     cusProductId: existingCusProduct.id,

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

@@ -214,8 +214,9 @@ const handleOneOffPrices = async ({
   const stripeCli = createStripeCli({ org, env: customer.env });
 
   logger.info("   1. Creating invoice");
-  const stripeInvoice = await stripeCli.invoices.create({
+  let stripeInvoice = await stripeCli.invoices.create({
     customer: customer.processor.id,
+    auto_advance: false,
   });
 
   // 2. Create invoice items
@@ -245,7 +246,7 @@ const handleOneOffPrices = async ({
   }
 
   if (!attachParams.invoiceOnly) {
-    const finalizedInvoice = await stripeCli.invoices.finalizeInvoice(
+    stripeInvoice = await stripeCli.invoices.finalizeInvoice(
       stripeInvoice.id,
       getInvoiceExpansion()
     );
@@ -259,19 +260,20 @@ const handleOneOffPrices = async ({
       logger,
     });
 
-    if (!paid && fromRequest) {
+    if (!paid) {
       await stripeCli.invoices.voidInvoice(stripeInvoice.id);
+      if (fromRequest && org.config.checkout_on_failed_payment) {
         await handleCreateCheckout({
           sb,
           req,
           res,
           attachParams,
         });
-      return;
+      } else {
-    } else if (!paid) {
         throw error;
       }
     }
+  }
 
   // Insert full customer product
   logger.info("   3. Creating full customer product");

server/src/internal/customers/change-product/scheduleUtils.ts

@@ -157,6 +157,7 @@ export const cancelFutureProductSchedule = async ({
       entitlements: fullCurProduct.entitlements,
       freeTrial: null,
       optionsList: [],
+      entities: [],
     },
   });
 

server/src/internal/customers/entitlements/cusEntUtils.ts

@@ -351,13 +351,15 @@ export const getCusBalancesByEntitlement = async ({
 
     data[key].balance += balance || 0;
     data[key].adjustment += adjustment || 0;
-    data[key].total +=
+    let total = (getResetBalance({
-      (getResetBalance({
       entitlement: ent,
       options: getEntOptions(cusProduct.options, ent),
       relatedPrice: getRelatedCusPrice(cusEnt, cusPrices)?.price,
     }) || 0) * count;
 
+    data[key].total += total;
+      
+
     data[key].unused += unused || 0;
   }
 
@@ -684,9 +686,11 @@ export const getTotalNegativeBalance = ({
 export const getExistingUsageFromCusProducts = ({
   entitlement,
   cusProducts,
+  entities,
 }: {
   entitlement: EntitlementWithFeature;
   cusProducts?: FullCusProduct[];
+  entities: Entity[];
 }) => {
   if (!entitlement || entitlement.feature.type === FeatureType.Boolean) {
     return 0;
@@ -715,5 +719,16 @@ export const getExistingUsageFromCusProducts = ({
 
   // // Calculate existing usage
   let existingAllowance = existingCusEnt.entitlement.allowance!;
-  return existingAllowance - existingCusEnt.balance!;
+
+
+  let { balance, adjustment, count, unused } = getCusEntMasterBalance({
+    cusEnt: existingCusEnt as any,
+    entities: entities,
+  });
+  existingUsage = existingAllowance - balance!;
+  if (unused && unused > 0) {
+    existingUsage -= unused;
+  }
+
+  return existingUsage;
 };

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

@@ -1,6 +1,7 @@
 import {
   Customer,
   EntitlementWithFeature,
+  Entity,
   FeatureOptions,
   FreeTrial,
   FullCusProduct,
@@ -33,6 +34,8 @@ export type AttachParams = {
   invoiceOnly?: boolean | undefined;
   billingAnchor?: number | undefined;
   metadata?: Record<string, string> | undefined;
+
+  entities: Entity[];
 };
 
 export type InsertCusProductParams = {

server/src/internal/customers/products/attachUtils.ts

@@ -33,6 +33,7 @@ import { createNewCustomer } from "@/internal/api/customers/handlers/handleCreat
 import { CusService } from "../CusService.js";
 import { getExistingCusProducts } from "../add-product/handleExistingProduct.js";
 import { getPricesForCusProduct } from "../change-product/scheduleUtils.js";
+import { EntityService } from "@/internal/api/entities/EntityService.js";
 
 const getOrCreateCustomerAndProducts = async ({
   sb,
@@ -309,6 +310,13 @@ export const getFullCusProductData = async ({
       logger,
     });
   
+  const entities = await EntityService.get({
+    sb,
+    internalCustomerId: customer.internal_id,
+    orgId,
+    env,
+  });
+
   let newOptionsList: FeatureOptions[] = [];
 
   for (const options of optionsListInput) {
@@ -356,6 +364,7 @@ export const getFullCusProductData = async ({
         .flat() as EntitlementWithFeature[],
       freeTrial,
       cusProducts,
+      entities,
     };
   }
 
@@ -441,5 +450,6 @@ export const getFullCusProductData = async ({
     entitlements: entitlementsWithFeature as EntitlementWithFeature[],
     freeTrial: uniqueFreeTrial,
     cusProducts,
+    entities,
   };
 };

server/src/internal/customers/products/cusProductUtils.ts

@@ -317,6 +317,7 @@ export const processFullCusProduct = ({
 }) => {
   // Process prices
 
+
   const prices = cusProduct.customer_prices.map((cp) => {
     let price = cp.price;
 
@@ -376,6 +377,8 @@ export const processFullCusProduct = ({
       }
     }
   });
+
+
   const trialing =
     cusProduct.trial_ends_at && cusProduct.trial_ends_at > Date.now();
 
@@ -392,6 +395,9 @@ export const processFullCusProduct = ({
       current_period_end: baseSub?.current_period_end
         ? baseSub.current_period_end * 1000
         : null,
+      current_period_start: baseSub?.current_period_start
+        ? baseSub.current_period_start * 1000
+        : null,
     };
   }
 

server/src/internal/prices/priceUtils.ts

@@ -279,10 +279,11 @@ export const getPriceAmount = ({
   options?: FeatureOptions;
   relatedEnt?: EntitlementWithFeature;
 }) => {
-  if (price.billing_type == BillingType.OneOff) {
+  let billingType = getBillingType(price.config!);
+  if (billingType == BillingType.OneOff) {
     let config = price.config as FixedPriceConfig;
     return Number(config.amount.toFixed(2));
-  } else if (price.billing_type == BillingType.UsageInAdvance) {
+  } else if (billingType == BillingType.UsageInAdvance) {
     let quantity = options?.quantity!;
     let config = price.config as UsagePriceConfig;
 

server/src/internal/public/publicRouter.ts

@@ -81,6 +81,7 @@ const publicRouterMiddleware = async (req: any, res: any, next: any) => {
 publicRouter.use(publicRouterMiddleware);
 
 publicRouter.get("/customers/:customer_id", async (req: any, res: any) => {
+  
   try {
     const customerId = req.params.customer_id;
     console.log("Getting customer (public)", customerId);
@@ -126,16 +127,33 @@ publicRouter.get("/customers/:customer_id", async (req: any, res: any) => {
 publicRouter.get(
   "/customers/:customerId/products",
   async (req: any, res: any) => {
+    try {
       const customerId = req.params.customerId;
 
-    const cusProducts = await CusProductService.getFullByCustomerId({
+    const customer = await CusService.getById({
       sb: req.sb,
-      customerId,
+      id: customerId,
       orgId: req.org.id,
       env: req.env,
-      inStatuses: [CusProductStatus.Active, CusProductStatus.Scheduled],
+      logger: req.logtail,
     });
 
+    if (!customer) {
+      return res.status(404).json({
+        message: `Customer ${customerId} not found`,
+      });
+    }
+
+    const cusProducts = await CusService.getFullCusProducts({
+      sb: req.sb,
+      internalCustomerId: customer.internal_id,
+      inStatuses: [CusProductStatus.Active, CusProductStatus.Scheduled],
+      withProduct: true,
+      withPrices: true,
+    });
+
+    
+
     if (!cusProducts || cusProducts.length === 0) {
       return res.status(200).json({
         main: [],
@@ -147,7 +165,12 @@ publicRouter.get(
     let addOns = [];
 
     for (const cusProduct of cusProducts) {
-      let processed = processFullCusProduct(cusProduct);
+      
+      let processed = processFullCusProduct({
+        cusProduct,
+        org: req.org,
+        subs: [],
+      });
 
       if (processed.status == CusProductStatus.Trialing) {
         processed.status = CusProductStatus.Active;
@@ -161,18 +184,21 @@ publicRouter.get(
       }
     }
 
-    // console.log("main", main);
 
     res.status(200).json({
       main,
         add_ons: addOns,
       });
+    } catch (error) {
+      handleRequestError({ req, error, res, action: "get customer products" });
+    }
   }
 );
 
 publicRouter.get(
   "/products/:product_id/options",
   async (req: any, res: any) => {
+    try {
       const product = await ProductService.getFullProductStrict({
         sb: req.sb,
         productId: req.params.product_id,
@@ -191,6 +217,9 @@ publicRouter.get(
       const options = getOptionsFromPrices(prices, features);
   
       res.status(200).json(options);
+    } catch (error) {
+      handleRequestError({ req, error, res, action: "get product options" });
+    }
   }
 );
 

server/src/trigger/adjustAllowance.ts

@@ -1,4 +1,5 @@
 import {
+  ErrCode,
   FullCusProduct,
   FullCustomerEntitlement,
   Product,
@@ -31,9 +32,11 @@ import { generateId } from "@/utils/genUtils.js";
 import { createStripeInvoiceItem } from "@/internal/customers/invoices/invoiceItemUtils.js";
 import { createLogtailWithContext } from "@/external/logtail/logtailUtils.js";
 import { LoggerAction } from "@autumn/shared";
-import { payForInvoice } from "@/external/stripe/stripeInvoiceUtils.js";
+import { getInvoiceExpansion, payForInvoice } from "@/external/stripe/stripeInvoiceUtils.js";
 import { isTrialing } from "@/internal/customers/products/cusProductUtils.js";
 import { ProductService } from "@/internal/products/ProductService.js";
+import { InvoiceService } from "@/internal/customers/invoices/InvoiceService.js";
+import RecaseError from "@/utils/errorUtils.js";
 
 type CusEntWithCusProduct = FullCustomerEntitlement & {
   customer_product: CusProduct;
@@ -321,8 +324,9 @@ export const adjustAllowance = async ({
     return;
   }
 
-  let quantity = newUsage + cusEnt.entitlement.allowance!;
   
+  
+  let quantity = newUsage + cusEnt.entitlement.allowance!;
   let prorationBehaviour = "create_prorations";
 
   // If prorate unused is false, then remove end of cycle
@@ -330,12 +334,10 @@ export const adjustAllowance = async ({
     prorationBehaviour = "none";
 
     const downgrade = quantity < (subItem.quantity || 0);
-
     if (!downgrade && !isTrialing(cusProduct as FullCusProduct)) {
       let entitlement = cusEnt.entitlement;
       let newUsage = entitlement.allowance! - newBalance;
-      let oldUsage = entitlement.allowance! - originalBalance;
+      let oldUsage = entitlement.allowance! - originalBalance + (replacedCount || 0);
-      newUsage = newUsage - (replacedCount || 0);
 
       let newAmount = getPriceForOverage(cusPrice.price, newUsage);
       let oldAmount = getPriceForOverage(cusPrice.price, oldUsage);
@@ -379,6 +381,7 @@ export const adjustAllowance = async ({
           },
         });
         
+
         const { paid, error } = await payForInvoice({
           fullOrg: org,
           env,
@@ -387,6 +390,27 @@ export const adjustAllowance = async ({
           logger,
         });
 
+        if (!paid) {
+          await stripeCli.invoices.voidInvoice(invoice.id);
+          throw new RecaseError({
+            message: "Failed to pay for invoice",
+            code: ErrCode.PayInvoiceFailed,
+          })
+        }
+        
+        const latestInvoice = await stripeCli.invoices.retrieve(invoice.id, {
+          ...getInvoiceExpansion()
+        });
+
+        await InvoiceService.createInvoiceFromStripe({
+          sb,
+          stripeInvoice: latestInvoice,
+          internalCustomerId: customer.internal_id,
+          org,
+          productIds: [product!.id],
+          internalProductIds: [product!.internal_id],
+        });
+
         if (!paid) {
           logger.warn("❗️ Failed to pay for invoice!");
         }
@@ -394,6 +418,8 @@ export const adjustAllowance = async ({
     } 
   }
   
+
+
   if (quantity < 0) {
     quantity = 0;
     logger.warn("❗️ Warning: quantity is negative, setting to 0");

server/src/trigger/updateBalanceTask.ts

@@ -186,12 +186,14 @@ export const performDeductionOnCusEnt = ({
   entityId,
   allowNegativeBalance = false,
   addAdjustment = false,
+  setZeroAdjustment = false,
 }: {
   cusEnt: FullCustomerEntitlement;
   toDeduct: number;
   entityId?: string | null;
   allowNegativeBalance?: boolean;
   addAdjustment?: boolean;
+  setZeroAdjustment?: boolean;
 }) => {
   let newEntities = structuredClone(cusEnt.entities);
   let newBalance = structuredClone(cusEnt.balance);
@@ -229,6 +231,10 @@ export const performDeductionOnCusEnt = ({
           newEntities![entityId!]!.adjustment = adjustment - newDeducted!;
         }
 
+        if (setZeroAdjustment) {
+          newEntities![entityId!]!.adjustment = 0;
+        }
+
         toDeductCursor = newToDeduct!;
         deducted += newDeducted!;
       }
@@ -255,6 +261,10 @@ export const performDeductionOnCusEnt = ({
         newEntities![entityId!]!.adjustment = adjustment - newDeducted!;
       }
 
+      if (setZeroAdjustment) {
+        newEntities![entityId!]!.adjustment = 0;
+      }
+
       toDeduct = newToDeduct!;
       deducted += newDeducted!;
     }
@@ -294,6 +304,7 @@ export const deductAllowanceFromCusEnt = async ({
   featureDeductions,
   willDeductCredits = false,
   entityId,
+  setZeroAdjustment = false,
 }: {
   toDeduct: number;
   deductParams: DeductParams;
@@ -302,6 +313,7 @@ export const deductAllowanceFromCusEnt = async ({
   featureDeductions: any;
   willDeductCredits?: boolean;
   entityId?: string | null;
+  setZeroAdjustment?: boolean;
 }) => {
   const { sb, feature, env, org, cusPrices, customer, properties } =
     deductParams;
@@ -328,6 +340,7 @@ export const deductAllowanceFromCusEnt = async ({
     toDeduct,
     entityId,
     allowNegativeBalance: false,
+    setZeroAdjustment,
   });
 
   let originalGrpBalance = getTotalNegativeBalance({
@@ -346,13 +359,17 @@ export const deductAllowanceFromCusEnt = async ({
   //   entities: newEntities,
   // });
 
+  let updates: any = {
+    balance: newBalance,
+    entities: newEntities,
+  }
+  if (setZeroAdjustment) {
+    updates.adjustment = 0;
+  }
   await CustomerEntitlementService.update({
     sb,
     id: cusEnt.id,
-    updates: {
+    updates,
-      balance: newBalance,
-      entities: newEntities,
-    },
   });
 
   await adjustAllowance({
@@ -405,11 +422,13 @@ export const deductFromUsageBasedCusEnt = async ({
   deductParams,
   cusEnts,
   entityId,
+  setZeroAdjustment = false,
 }: {
   toDeduct: number;
   deductParams: DeductParams;
   cusEnts: FullCustomerEntitlement[];
   entityId?: string | null;
+  setZeroAdjustment?: boolean;
 }) => {
   const { sb, feature, env, org, cusPrices, customer, properties } =
     deductParams;
@@ -452,6 +471,7 @@ export const deductFromUsageBasedCusEnt = async ({
     toDeduct,
     entityId,
     allowNegativeBalance: true,
+    setZeroAdjustment,
   });
 
   // console.log("NEW BALANCE", newBalance);
@@ -469,13 +489,18 @@ export const deductFromUsageBasedCusEnt = async ({
     entities: newEntities!,
   });
 
+  let updates: any = {
+    balance: newBalance,
+    entities: newEntities,
+  }
+  if (setZeroAdjustment) {
+    updates.adjustment = 0;
+  }
+
   await CustomerEntitlementService.update({
     sb,
     id: usageBasedEnt.id,
-    updates: {
+    updates
-      balance: newBalance,
-      entities: newEntities,
-    },
   });
 
   // const totalNegativeBalance = getTotalNegativeBalance(usageBasedEnt);

server/src/trigger/updateUsageTask.ts

@@ -220,13 +220,13 @@ export const updateUsage = async ({
     setUsage,
   });
 
-  // 2. Handle group_by initialization
+  // // 2. Handle group_by initialization
-  await initGroupBalancesForEvent({
+  // await initGroupBalancesForEvent({
-    sb,
+  //   sb,
-    features,
+  //   features,
-    cusEnts,
+  //   cusEnts,
-    properties,
+  //   properties,
-  });
+  // });
 
   // 3. Return if no customer entitlements or features found
   if (cusEnts.length === 0 || features.length === 0) {
@@ -258,6 +258,7 @@ export const updateUsage = async ({
         },
         featureDeductions,
         willDeductCredits: true,
+        setZeroAdjustment: true,
       });
     }
 
@@ -277,6 +278,7 @@ export const updateUsage = async ({
         customer,
         properties,
       },
+      setZeroAdjustment: true,
     });
   }
 

server/src/utils/constants.ts

@@ -1 +1,8 @@
+import { CusProductStatus } from "@autumn/shared";
+
 export const BREAK_API_VERSION = 0.2;
+
+export const getActiveCusProductStatuses = () => [
+  CusProductStatus.Active,
+  CusProductStatus.PastDue,
+];

server/test.sh

@@ -9,10 +9,11 @@ if [ "$1" == "basic-parallel" ]; then
     tests/attach/**/*.ts \
 
 elif [ "$1" == "advanced-parallel" ]; then
-    MOCHA_PARALLEL=true  $MOCHA_SETUP && $MOCHA_CMD \
+    MOCHA_PARALLEL=true  \
-    'tests/advanced/usage/*.ts' \
     && $MOCHA_CMD 'tests/advanced/coupons/*.ts' \
-    && $MOCHA_CMD 'tests/advanced/arrear_prorated/*.ts'
+    # $MOCHA_SETUP \
+    # && $MOCHA_CMD 'tests/advanced/arrear_prorated/*.ts' \
+    # && $MOCHA_CMD 'tests/advanced/usage/*.ts' \
 
 
 elif [ "$1" == "alex-parallel" ]; then

shared/errors/errCode.ts

@@ -100,7 +100,7 @@ export const ErrCode = {
   GetCusPriceFailed: "get_cus_price_failed",
 
   // Pay for invoice
-  PayInvoiceFailed: "pay_invoice_failed",
+  PayInvoiceFailed: "invoice_payment_failed",
 
   // COUPONS
   PromoCodeAlreadyExistsInStripe: "promo_code_already_exists_in_stripe",

shared/models/orgModels/orgConfigModels.ts

@@ -8,6 +8,7 @@ export const OrgConfigSchema = z.object({
   prorate_unused: z.boolean().default(true),
 
   api_version: z.number().default(0.2),
+  checkout_on_failed_payment: z.boolean().default(true),
 });
 
 export type OrgConfig = z.infer<typeof OrgConfigSchema>;