feature: entities can now have products of their own

2025-05-07 20:40:12 +01:00
parent 08be4737a4
commit a2ec951efb
64 changed files with 1411 additions and 1022 deletions
--- a/package-lock.json
+++ b/package-lock.json
@@ -5644,6 +5644,15 @@
        "node": ">=4"
      }
    },
+    "node_modules/async-listen": {
+      "version": "3.1.0",
+      "resolved": "https://registry.npmjs.org/async-listen/-/async-listen-3.1.0.tgz",
+      "integrity": "sha512-TkOhqze98lP+6e7SPbrBpyhTpfvqqX8VYKGn4uckrgPan4WQIHnTaUD2zZzZS18eVVDj4rHPcIZa1PGgvo1DfA==",
+      "license": "MIT",
+      "engines": {
+        "node": ">= 14"
+      }
+    },
    "node_modules/asynckit": {
      "version": "0.4.0",
      "resolved": "https://registry.npmjs.org/asynckit/-/asynckit-0.4.0.tgz",
@@ -14848,7 +14857,7 @@
        "@types/express": "^5.0.0",
        "@unkey/api": "^0.29.0",
        "ai": "^4.3.10",
-        "autumn-js": "^0.0.13",
+        "autumn-js": "^0.0.27",
        "body-parser": "^1.20.3",
        "bullmq": "^5.31.1",
        "chai": "^5.1.2",
@@ -15442,18 +15451,16 @@
        "node": ">=8"
      }
    },
-    "server/node_modules/async-listen": {
-      "version": "3.0.1",
-      "license": "MIT",
-      "engines": {
-        "node": ">= 14"
-      }
-    },
    "server/node_modules/autumn-js": {
-      "version": "0.0.13",
-      "resolved": "https://registry.npmjs.org/autumn-js/-/autumn-js-0.0.13.tgz",
-      "integrity": "sha512-R+rfwELDdIUL09ESwryxOiAdM81JCWxc/G0XM/mDMGBPyA+yEE76Qzei7fpJgo+IxBVUaYphxefTO9srb304Gg==",
+      "version": "0.0.27",
+      "resolved": "https://registry.npmjs.org/autumn-js/-/autumn-js-0.0.27.tgz",
+      "integrity": "sha512-K5DCyZd67oIj7zk5QBbNOvG+Kw9B37ktuvWD+VczkPjs/RH0mleFp66WQhGS4PBeWh0o5pwq3nCufmqh9Lj57w==",
      "license": "MIT",
+      "dependencies": {
+        "async-listen": "^3.1.0",
+        "commander": "^13.1.0",
+        "ora": "^8.2.0"
+      },
      "peerDependencies": {
        "@clerk/nextjs": "^6.16.0",
        "@supabase/ssr": "^0.6.1",
@@ -15473,6 +15480,15 @@
        }
      }
    },
+    "server/node_modules/autumn-js/node_modules/commander": {
+      "version": "13.1.0",
+      "resolved": "https://registry.npmjs.org/commander/-/commander-13.1.0.tgz",
+      "integrity": "sha512-/rFeCpNJQbhSZjGVwO9RFV3xPqbnERS8MmIQzCtD/zl6gpJuV/bMLuN92oG3F7d8oDEHHRrujSXNUr8fpjntKw==",
+      "license": "MIT",
+      "engines": {
+        "node": ">=18"
+      }
+    },
    "server/node_modules/body-parser": {
      "version": "1.20.3",
      "license": "MIT",
--- a/server/package.json
+++ b/server/package.json
@@ -45,7 +45,7 @@
    "@types/express": "^5.0.0",
    "@unkey/api": "^0.29.0",
    "ai": "^4.3.10",
-    "autumn-js": "^0.0.13",
+    "autumn-js": "^0.0.27",
    "body-parser": "^1.20.3",
    "bullmq": "^5.31.1",
    "chai": "^5.1.2",
--- a/server/src/errors/errCodes.ts
+++ b/server/src/errors/errCodes.ts
@@ -70,4 +70,7 @@ export const ErrCode = {

  // Payment errors
  CardDeclinedError: "card_declined_error",
+
+  // Entity
+  EntityNotFound: "entity_not_found",
 };
--- a/server/src/external/stripe/createStripePrice/createStripeFixedPrice.ts
+++ b/server/src/external/stripe/createStripePrice/createStripeFixedPrice.ts
@@ -20,10 +20,6 @@ export const createStripeFixedPrice = async ({
  org: Organization;
  curStripePrice: Stripe.Price | null;
 }) => {
-  if (curStripePrice) {
-    return;
-  }
-
  const config = price.config as FixedPriceConfig;

  let amount = new Decimal(config.amount).mul(100).toNumber();
--- a/server/src/external/stripe/createStripePrice/createStripePrice.ts
+++ b/server/src/external/stripe/createStripePrice/createStripePrice.ts
@@ -19,7 +19,10 @@ import { createStripeFixedPrice } from "./createStripeFixedPrice.js";
 import { createStripePrepaid } from "./createStripePrepaid.js";
 import { createStripeOneOffTieredProduct } from "./createStripeOneOffTiered.js";
 import { createStripeInArrearPrice } from "./createStripeInArrear.js";
-import { createStripeArrearProrated } from "./createStripeArrearProrated.js";
+import {
+  createStripeArrearProrated,
+  createStripeMeteredPrice,
+} from "./createStripeArrearProrated.js";

 export const checkCurStripePrice = async ({
  price,
@@ -100,8 +103,9 @@ export const createStripePriceIFNotExist = async ({
    stripeCli,
  });

-  price.config!.stripe_price_id = stripePrice?.id;
-  (price.config! as UsagePriceConfig).stripe_product_id = stripeProd?.id;
+  let config = price.config! as UsagePriceConfig;
+  config.stripe_price_id = stripePrice?.id;
+  config.stripe_product_id = stripeProd?.id;

  let relatedEnt = getPriceEntitlement(price, entitlements);
  let isOneOffAndTiered = priceIsOneOffAndTiered(price, relatedEnt);
@@ -111,15 +115,17 @@ export const createStripePriceIFNotExist = async ({
    billingType == BillingType.FixedCycle ||
    billingType == BillingType.OneOff
  ) {
-    logger.info("Creating stripe fixed price");
-    await createStripeFixedPrice({
-      sb,
-      stripeCli,
-      price,
-      product,
-      org,
-      curStripePrice: stripePrice,
-    });
+    if (!stripePrice) {
+      logger.info("Creating stripe fixed price");
+      await createStripeFixedPrice({
+        sb,
+        stripeCli,
+        price,
+        product,
+        org,
+        curStripePrice: stripePrice,
+      });
+    }
  }

  // 2. If prepaid
@@ -161,6 +167,23 @@ export const createStripePriceIFNotExist = async ({
        org,
        curStripeProd: stripeProd,
      });
+    } else if (!config.stripe_placeholder_price_id) {
+      logger.info(`Creating stripe placeholder price`);
+      let placeholderPrice = await createStripeMeteredPrice({
+        sb,
+        stripeCli,
+        price,
+        entitlements,
+        product,
+        org,
+      });
+      config.stripe_placeholder_price_id = placeholderPrice.id;
+      await sb
+        .from("prices")
+        .update({
+          config,
+        })
+        .eq("id", price.id);
    }
  }

--- a/server/src/external/stripe/priceToStripeItem/priceToArrearProrated.ts
+++ b/server/src/external/stripe/priceToStripeItem/priceToArrearProrated.ts
@@ -12,6 +12,7 @@ export const priceToInArrearProrated = ({
 }) => {
  const config = price.config as UsagePriceConfig;
  let quantity = existingUsage || 0;
+
  if (quantity == 0 && isCheckout) {
    return {
      price: config.stripe_placeholder_price_id,
--- a/server/src/external/stripe/stripeErrorUtils.ts
+++ b/server/src/external/stripe/stripeErrorUtils.ts
@@ -0,0 +1,4 @@
+export const stripeErrToCusMsg = (error: any) => {
+  let code = error.code;
+  let msg = error.message;
+};
--- a/server/src/external/stripe/stripeMeterUtils.ts
+++ b/server/src/external/stripe/stripeMeterUtils.ts
@@ -1 +1,49 @@
+import { Customer, Feature } from "@autumn/shared";
+import { getBillingType } from "@/internal/prices/priceUtils.js";
+import { BillingType, Price, UsagePriceConfig } from "@autumn/shared";
 import { SupabaseClient } from "@supabase/supabase-js";
+import Stripe from "stripe";
+
+export const submitUsageToStripe = async ({
+  price,
+  stripeCli,
+  usage,
+  customer,
+  usageTimestamp,
+  feature,
+  logger,
+}: {
+  stripeCli: Stripe;
+  price: Price;
+  usage: number;
+  customer: Customer;
+  usageTimestamp?: number;
+  feature: Feature;
+  logger: any;
+}) => {
+  let config = price.config as UsagePriceConfig;
+  let billingType = getBillingType(config);
+
+  if (billingType != BillingType.UsageInArrear) {
+    logger.warn(
+      `Price ${price.id} is not usage in arrear type, can't send usage`
+    );
+  }
+  const stripeMeter = await stripeCli.billing.meters.retrieve(
+    config.stripe_meter_id!
+  );
+
+  await stripeCli.billing.meterEvents.create({
+    // event_name: price.id!,
+    event_name: stripeMeter.event_name,
+    payload: {
+      stripe_customer_id: customer.processor.id,
+      value: usage.toString(),
+    },
+    timestamp: usageTimestamp || Math.floor(Date.now() / 1000),
+  });
+
+  logger.info(
+    `🌟🌟🌟 Submitted meter event for customer ${customer.id}, feature: ${feature.name}, usage: ${usage}`
+  );
+};
--- a/server/src/external/stripe/stripePriceUtils.ts
+++ b/server/src/external/stripe/stripePriceUtils.ts
@@ -4,6 +4,10 @@ import {
  Price,
  AllowanceType,
  Feature,
+  Customer,
+  Organization,
+  FullCusProduct,
+  UsagePriceConfig,
 } from "@autumn/shared";

 import Stripe from "stripe";
@@ -13,6 +17,7 @@ import {
  getEntOptions,
  getPriceAmount,
  getPriceEntitlement,
+  getPriceForOverage,
  getPriceOptions,
  getProductForPrice,
 } from "@/internal/prices/priceUtils.js";
@@ -21,6 +26,7 @@ import { SupabaseClient } from "@supabase/supabase-js";
 import { AttachParams } from "@/internal/customers/products/AttachParams.js";
 import { getExistingUsageFromCusProducts } from "@/internal/customers/entitlements/cusEntUtils.js";
 import { priceToStripeItem } from "./priceToStripeItem/priceToStripeItem.js";
+import { getFeatureName } from "@/internal/features/displayUtils.js";

 export const createSubMeta = ({ features }: { features: Feature[] }) => {
  const usageFeatures = features.map((f) => ({
@@ -118,6 +124,7 @@ export const getStripeSubItems = async ({
        cusProducts: attachParams.cusProducts,
        entities: attachParams.entities,
        carryExistingUsages,
+        internalEntityId: attachParams.internalEntityId,
      });

      if (
@@ -179,48 +186,64 @@ export const getStripeSubItems = async ({
  return itemSets;
 };

-// Can delete
-export const pricesToInvoiceItems = async ({
-  sb,
-  stripeCli,
-  attachParams,
+export const getInvoiceItemForUsage = ({
  stripeInvoiceId,
+  price,
+  feature,
+  totalUsage,
+  overage,
+  currency,
+  customer,
+  cusProduct,
+  logger,
+  periodStart,
+  periodEnd,
 }: {
-  sb: SupabaseClient;
-  stripeCli: Stripe;
-  attachParams: AttachParams;
  stripeInvoiceId: string;
+  price: Price;
+  feature: Feature;
+  totalUsage: number;
+  overage: number;
+  currency: string;
+  customer: Customer;
+  cusProduct: FullCusProduct;
+  logger: any;
+  periodStart: number;
+  periodEnd: number;
 }) => {
-  const { prices, optionsList, entitlements, products, customer } =
-    attachParams;
-  for (const price of prices) {
-    // Calculate amount
-    const options = getPriceOptions(price, optionsList);
-    const entitlement = getPriceEntitlement(price, entitlements);
-    const amount = getPriceAmount({
-      price,
-      options,
-      relatedEnt: entitlement,
-    });
+  let priceAmount = getPriceForOverage(price, overage);
+  let featureName = getFeatureName({
+    feature,
+    plural: totalUsage == 1 ? false : true,
+    capitalize: true,
+  });

-    let allowanceStr = "";
-    if (entitlement) {
-      allowanceStr =
-        entitlement.allowance_type == AllowanceType.Unlimited
-          ? "Unlimited"
-          : entitlement.allowance_type == AllowanceType.None
-          ? "None"
-          : `${entitlement.allowance}`;
-      allowanceStr = `x ${allowanceStr} (${entitlement.feature.name})`;
-    }
+  let config = price.config! as UsagePriceConfig;
+  let invoiceItem: Stripe.InvoiceItemCreateParams = {
+    invoice: stripeInvoiceId,
+    customer: customer.processor.id,
+    currency,

-    let product = getProductForPrice(price, products)!;
+    description: `${cusProduct.product.name} - ${featureName} x ${Math.round(
+      totalUsage
+    )}`,

-    await stripeCli.invoiceItems.create({
-      customer: customer.processor.id,
-      amount: amount * 100,
-      invoice: stripeInvoiceId,
-      description: `${product.name}${allowanceStr}`,
-    });
-  }
+    price_data: {
+      product: config.stripe_product_id!,
+      unit_amount: Math.max(Math.round(priceAmount * 100), 0),
+      currency,
+    },
+    period: {
+      start: periodStart,
+      end: periodEnd,
+    },
+  };
+
+  logger.info(
+    `🌟🌟 Created invoice item for ${
+      feature.name
+    } usage. Amount: ${priceAmount.toFixed(2)}, Total Usage: ${totalUsage}`
+  );
+
+  return invoiceItem;
 };
--- a/server/src/external/stripe/stripeWebhooks.ts
+++ b/server/src/external/stripe/stripeWebhooks.ts
@@ -76,6 +76,7 @@ stripeWebhookRouter.post(
            org,
            subscription: event.data.object,
            env,
+            logger,
          });
          break;

--- a/server/src/external/stripe/utils.ts
+++ b/server/src/external/stripe/utils.ts
@@ -72,6 +72,18 @@ export const calculateMetered1Price = ({
  return totalPrice;
 };

+export const subToAutumnInterval = (sub: Stripe.Subscription) => {
+  let recuringItem = sub.items.data.find((i) => i.price.recurring != null);
+  if (!recuringItem) {
+    return BillingInterval.OneOff;
+  }
+
+  return stripeToAutumnInterval({
+    interval: recuringItem.price.recurring!.interval,
+    intervalCount: recuringItem.price.recurring!.interval_count,
+  });
+};
+
 export const stripeToAutumnInterval = ({
  interval,
  intervalCount,
--- a/server/src/external/stripe/webhookHandlers/handleInvoiceCreated.ts
+++ b/server/src/external/stripe/webhookHandlers/handleInvoiceCreated.ts
@@ -30,6 +30,10 @@ import { getResetBalancesUpdate } from "@/internal/customers/entitlements/groupB
 import { createLogtailWithContext } from "@/external/logtail/logtailUtils.js";
 import { EntityService } from "@/internal/api/entities/EntityService.js";
 import { Client } from "pg";
+import { FeatureService } from "@/internal/features/FeatureService.js";
+import { getFeatureName } from "@/internal/features/displayUtils.js";
+import { submitUsageToStripe } from "../stripeMeterUtils.js";
+import { getInvoiceItemForUsage } from "../stripePriceUtils.js";

 // Format invoice nicely

@@ -239,16 +243,19 @@ const handleUsageInArrear = async ({
    return;
  }

-  const balance = getTotalNegativeBalance({
+  const totalNegativeBalance = getTotalNegativeBalance({
    cusEnt: relatedCusEnt as any,
    balance: relatedCusEnt.balance!,
    entities: relatedCusEnt.entities!,
    billingUnits: (price.config as UsagePriceConfig).billing_units || 1,
  });

-  const totalQuantity = new Decimal(allowance).minus(balance).toNumber();
+  const totalQuantity = new Decimal(allowance)
+    .minus(totalNegativeBalance)
+    .toNumber();

  const billingUnits = (price.config as UsagePriceConfig).billing_units || 1;
+
  const roundedQuantity =
    Math.ceil(new Decimal(totalQuantity).div(billingUnits).toNumber()) *
    billingUnits;
@@ -257,43 +264,48 @@ const handleUsageInArrear = async ({
    subDays(new Date(invoice.created * 1000), 1).getTime() / 1000
  );

-  // 1. Get stripe meter
-  const stripeMeter = await stripeCli.billing.meters.retrieve(
-    config.stripe_meter_id!
-  );
-
-  await stripeCli.billing.meterEvents.create({
-    // event_name: price.id!,
-    event_name: stripeMeter.event_name,
-    payload: {
-      stripe_customer_id: customer.processor.id,
-      value: roundedQuantity.toString(),
-    },
-    timestamp: usageTimestamp,
-  });
-
  let feature = relatedCusEnt.entitlement.feature;
-  logger.info(
-    `✅ Submitted meter event for customer ${customer.id}, feature: ${feature.id}, stripe event: ${stripeMeter.event_name}`
-  );
-  logger.info(
-    `Allowance: ${allowance}, Balance: ${balance}, Quantity: ${totalQuantity}, Rounded: ${roundedQuantity}`
-  );
+  if (activeProduct.internal_entity_id) {
+    let currency = invoice.currency;
+    // console.log("Total negative balance: ", totalNegativeBalance);
+    // console.log("Total quantity: ", roundedQuantity);
+    let invoiceItem = getInvoiceItemForUsage({
+      stripeInvoiceId: invoice.id,
+      price,
+      overage: -totalNegativeBalance,
+      customer,
+      currency,
+      cusProduct: activeProduct,
+      feature,
+      totalUsage: totalQuantity,
+      logger,
+      periodStart: invoice.period_start,
+      periodEnd: invoice.period_end,
+    });

-  let invoiceCreatedStr = formatUnixToDateTime(invoice.created * 1000);
-  let usageTimestampStr = formatUnixToDateTime(usageTimestamp * 1000);
-  logger.info(
-    `Invoice created: ${invoiceCreatedStr}, Usage timestamp: ${usageTimestampStr}`
-  );
+    await stripeCli.invoiceItems.create(invoiceItem);
+  } else {
+    await submitUsageToStripe({
+      price,
+      stripeCli,
+      usage: roundedQuantity,
+      customer,
+      usageTimestamp,
+      feature: relatedCusEnt.entitlement.feature,
+      logger,
+    });
+  }
+
+  // let invoiceCreatedStr = formatUnixToDateTime(invoice.created * 1000);
+  // let usageTimestampStr = formatUnixToDateTime(usageTimestamp * 1000);
+  // logger.info(
+  //   `Invoice created: ${invoiceCreatedStr}, Usage timestamp: ${usageTimestampStr}`
+  // );

-  // reset balance
-  // TODO: If lifetime, reset to 0...
  if (relatedCusEnt.entitlement.interval == EntInterval.Lifetime) {
-    logger.info(
-      `Feature ${feature.id} has lifetime interval, skipping reset...`
-    );
    return;
  }
+
  let ent = relatedCusEnt.entitlement;
  let resetBalancesUpdate = getResetBalancesUpdate({
    cusEnt: relatedCusEnt,
@@ -480,6 +492,52 @@ export const handleInvoiceCreated = async ({
      return;
    }

+    let internalEntityId = activeProducts.find(
+      (p) => p.internal_entity_id
+    )?.internal_entity_id;
+
+    let features = await FeatureService.getFeatures({
+      sb,
+      orgId: org.id,
+      env,
+    });
+
+    if (internalEntityId) {
+      // Add memo to invoice
+      try {
+        let stripeCli = createStripeCli({ org, env });
+        let entity = await EntityService.getByInternalId({
+          sb,
+          internalId: internalEntityId,
+          orgId: org.id,
+          env,
+        });
+
+        if (entity) {
+          logger.info(`Entity: ${entity.name}`);
+        }
+
+        // Add memo to invoice
+        let feature = features.find(
+          (f) => f.internal_id == entity?.internal_feature_id
+        );
+
+        await stripeCli.invoices.update(invoice.id, {
+          description: `${getFeatureName({
+            feature,
+            plural: false,
+            capitalize: true,
+          })}: ${entity?.name} (ID: ${entity?.id})`,
+        });
+      } catch (error: any) {
+        if (
+          error.message != "Finalized invoices can't be updated in this way"
+        ) {
+          logger.error(`Failed to add entity ID to invoice description`, error);
+        }
+      }
+    }
+
    const stripeSubs = await getStripeSubs({
      stripeCli: createStripeCli({ org, env }),
      subIds: activeProducts.map((p) => p.subscription_ids).flat(),
@@ -514,3 +572,20 @@ export const handleInvoiceCreated = async ({
 //   });
 //   continue;
 // }
+
+// // 1. Get stripe meter
+// const stripeMeter = await stripeCli.billing.meters.retrieve(
+//   config.stripe_meter_id!
+// );
+
+// await stripeCli.billing.meterEvents.create({
+//   // event_name: price.id!,
+//   event_name: stripeMeter.event_name,
+//   payload: {
+//     stripe_customer_id: customer.processor.id,
+//     value: roundedQuantity.toString(),
+//   },
+//   timestamp: usageTimestamp,
+// });
+
+// let feature = relatedCusEnt.entitlement.feature;
--- a/server/src/external/stripe/webhookHandlers/handleSubCreated.ts
+++ b/server/src/external/stripe/webhookHandlers/handleSubCreated.ts
@@ -1,5 +1,12 @@
 import { CusProductService } from "@/internal/customers/products/CusProductService.js";
-import { CusProductStatus, Organization, ProcessorType } from "@autumn/shared";
+import {
+  BillingType,
+  CusProductStatus,
+  FullCusProduct,
+  Organization,
+  Price,
+  ProcessorType,
+} from "@autumn/shared";
 import { AppEnv } from "@autumn/shared";
 import { SupabaseClient } from "@supabase/supabase-js";
 import Stripe from "stripe";
@@ -8,17 +15,20 @@ import { InvoiceService } from "@/internal/customers/invoices/InvoiceService.js"
 import { getStripeExpandedInvoice } from "../stripeInvoiceUtils.js";
 import { SubService } from "@/internal/subscriptions/SubService.js";
 import { generateId } from "@/utils/genUtils.js";
+import { getBillingType } from "@/internal/prices/priceUtils.js";

 export const handleSubCreated = async ({
  sb,
  subscription,
  org,
  env,
+  logger,
 }: {
  sb: SupabaseClient;
  subscription: Stripe.Subscription;
  org: Organization;
  env: AppEnv;
+  logger: any;
 }) => {
  if (subscription.schedule) {
    const cusProds = await CusProductService.getByStripeScheduledId({
@@ -115,4 +125,71 @@ export const handleSubCreated = async ({

    await Promise.all(batchUpdate);
  }
+
+  // Get cus prods for sub
+  let cusProds = await CusProductService.getByStripeSubId({
+    sb,
+    stripeSubId: subscription.id,
+    orgId: org.id,
+    env,
+    withCusEnts: true,
+    withCusPrices: true,
+  });
+
+  let stripeCli = createStripeCli({ org, env });
+  let handleInArrearWithEntity = async (cusProd: FullCusProduct) => {
+    if (!cusProd.internal_entity_id) {
+      return;
+    }
+
+    let arrearPrices = cusProd.customer_prices
+      .map((cp) => cp.price)
+      .filter(
+        (p: Price) =>
+          getBillingType(p.config as any) == BillingType.UsageInArrear
+      );
+
+    if (arrearPrices.length == 0) {
+      return;
+    }
+
+    let itemsToDelete = [];
+    for (const arrearPrice of arrearPrices) {
+      let subItem = subscription.items.data.find(
+        (i) => i.price.id == arrearPrice.config?.stripe_price_id
+      );
+
+      if (!subItem) {
+        continue;
+      }
+
+      itemsToDelete.push({
+        id: subItem.id,
+        deleted: true,
+      });
+    }
+
+    if (itemsToDelete.length > 0) {
+      try {
+        await stripeCli.subscriptions.update(subscription.id, {
+          items: itemsToDelete,
+        });
+        console.log(
+          `sub.created, cus product with entity: deleted ${itemsToDelete.length} items`
+        );
+      } catch (error) {
+        logger.error(
+          `sub.created, cus product with entity: failed to delete items`,
+          error
+        );
+      }
+    }
+  };
+
+  let batchUpdate = [];
+  for (const cusProd of cusProds) {
+    batchUpdate.push(handleInArrearWithEntity(cusProd));
+  }
+
+  await Promise.all(batchUpdate);
 };
--- a/server/src/internal/api/apiRouter.ts
+++ b/server/src/internal/api/apiRouter.ts
@@ -106,7 +106,7 @@ apiRouter.use("/redemptions", redemptionRouter);

 // Cus Product
 apiRouter.use(attachRouter);
-apiRouter.use("/expire", expireRouter);
+apiRouter.use("/cancel", expireRouter);
 apiRouter.use("/entitled", entitledRouter);
 apiRouter.use("/check", entitledRouter);

--- a/server/src/internal/api/customers/cusUtils.ts
+++ b/server/src/internal/api/customers/cusUtils.ts
@@ -5,6 +5,7 @@ import {
  Customer,
  CustomerData,
  CustomerSchema,
+  FullCustomer,
  FullCustomerEntitlement,
  Organization,
  ProductSchema,
@@ -25,6 +26,7 @@ import {
 import { processInvoice } from "@/internal/customers/invoices/InvoiceService.js";
 import { InvoiceService } from "@/internal/customers/invoices/InvoiceService.js";
 import { handleCreateCustomer } from "./handlers/handleCreateCustomer.js";
+import { sortCusEntsForDeduction } from "@/internal/customers/entitlements/cusEntUtils.js";

 export const updateCustomerDetails = async ({
  sb,
@@ -191,58 +193,67 @@ export const processFullCusProducts = ({
 // IMPORTANT FUNCTION
 export const getCusEntsInFeatures = async ({
  sb,
-  internalCustomerId,
+  customer,
  internalFeatureIds,
-  inStatuses = [CusProductStatus.Active],
-  withPrices = false,
-  withProduct = false,
  logger,
  reverseOrder = false,
 }: {
  sb: SupabaseClient;
-  internalCustomerId: string;
+  customer: FullCustomer;
  internalFeatureIds?: string[];
-  inStatuses?: CusProductStatus[];
-  withPrices?: boolean;
-  withProduct?: boolean;
  logger: any;
  reverseOrder?: boolean;
 }) => {
-  const fullCusProducts = await CusService.getFullCusProducts({
-    sb,
-    internalCustomerId,
-    inStatuses: inStatuses,
-    withPrices: withPrices,
-    withProduct: withProduct,
-    logger,
+  let cusProducts = customer.customer_products;
+  let cusEnts = cusProducts.flatMap((cusProduct) => {
+    return cusProduct.customer_entitlements.map((cusEnt) => ({
+      ...cusEnt,
+      customer_product: cusProduct,
+    }));
  });

-  const cusEntsWithCusProduct = fullCusProductToCusEnts(
-    fullCusProducts!,
-    inStatuses,
-    reverseOrder
-  );
+  let cusPrices = cusProducts.flatMap((cusProduct) => {
+    return cusProduct.customer_prices || [];
+  });

-  if (!cusEntsWithCusProduct) {
-    return { cusEnts: [] };
-  }
-
-  let cusEnts: FullCustomerEntitlement[] = [];
  if (internalFeatureIds) {
-    cusEnts = cusEntsWithCusProduct.filter((cusEnt) =>
+    cusEnts = cusEnts.filter((cusEnt) =>
      internalFeatureIds.includes(cusEnt.internal_feature_id)
    );
-  } else {
-    cusEnts = cusEntsWithCusProduct;
  }

-  // sortCusEntsForDeduction(cusEnts, reverseOrder);
-
-  if (!withPrices) {
-    return { cusEnts, cusPrices: undefined };
+  if (customer.entity) {
+    let entity = customer.entity;
+    cusEnts = cusEnts.filter(
+      (cusEnt) =>
+        cusEnt.customer_product.internal_entity_id === entity.internal_id ||
+        cusEnt.entities
+    );
  }

-  const cusPrices = fullCusProductToCusPrices(fullCusProducts, inStatuses);
+  sortCusEntsForDeduction(cusEnts, reverseOrder);
+  // // const cusEntsWithCusProduct = fullCusProductToCusEnts(
+  // //   customer.customer_products,
+  // //   [CusProductStatus.Active, CusProductStatus.PastDue],
+  // //   reverseOrder
+  // // );
+
+  // if (!cusEntsWithCusProduct) {
+  //   return { cusEnts: [] };
+  // }
+
+  // let cusEnts: FullCustomerEntitlement[] = [];
+  // if (internalFeatureIds) {
+  //   cusEnts = cusEntsWithCusProduct.filter((cusEnt) =>
+  //     internalFeatureIds.includes(cusEnt.internal_feature_id)
+  //   );
+  // } else {
+  //   cusEnts = cusEntsWithCusProduct;
+  // }
+
+  // // sortCusEntsForDeduction(cusEnts, reverseOrder);
+
+  // const cusPrices = fullCusProductToCusPrices(fullCusProducts, inStatuses);

  return { cusEnts, cusPrices };
 };
--- a/server/src/internal/api/customers/handlers/handleCreateCustomer.ts
+++ b/server/src/internal/api/customers/handlers/handleCreateCustomer.ts
@@ -75,6 +75,7 @@ export const createNewCustomer = async ({
  nextResetAt,
  processor,
  logger,
+  createDefaultProducts = true,
 }: {
  sb: SupabaseClient;
  org: Organization;
@@ -83,6 +84,7 @@ export const createNewCustomer = async ({
  nextResetAt?: number;
  processor?: any;
  logger: any;
+  createDefaultProducts?: boolean;
 }) => {
  logger.info(`Creating new customer: ${customer.id}`);
  logger.info(`Org ID: ${org.id}`);
@@ -137,6 +139,10 @@ export const createNewCustomer = async ({
    customer: customerData,
  });

+  if (!createDefaultProducts) {
+    return newCustomer;
+  }
+
  if (nonFreeProds.length > 0) {
    await initStripeCusAndProducts({
      sb,
@@ -202,6 +208,7 @@ const handleIdIsNull = async ({
  newCus,
  logger,
  processor,
+  createDefaultProducts,
 }: {
  sb: SupabaseClient;
  org: Organization;
@@ -209,6 +216,7 @@ const handleIdIsNull = async ({
  newCus: CreateCustomer;
  logger: any;
  processor?: any;
+  createDefaultProducts?: boolean;
 }) => {
  // 1. ID is null
  if (!newCus.email) {
@@ -252,6 +260,7 @@ const handleIdIsNull = async ({
    customer: newCus,
    logger,
    processor,
+    createDefaultProducts,
  });

  return createdCustomer;
@@ -265,6 +274,7 @@ export const handleCreateCustomerWithId = async ({
  logger,
  newCus,
  processor,
+  createDefaultProducts = true,
 }: {
  sb: SupabaseClient;
  org: Organization;
@@ -272,6 +282,7 @@ export const handleCreateCustomerWithId = async ({
  logger: any;
  newCus: CreateCustomer;
  processor?: any;
+  createDefaultProducts?: boolean;
 }) => {
  // 1. Get by ID
  let existingCustomer = await CusService.getById({
@@ -327,6 +338,7 @@ export const handleCreateCustomerWithId = async ({
    customer: newCus,
    logger,
    processor,
+    createDefaultProducts,
  });
 };

@@ -339,6 +351,7 @@ export const handleCreateCustomer = async ({
  params = {},
  processor,
  getDetails = true,
+  createDefaultProducts = true,
 }: {
  cusData: CreateCustomer;
  sb: SupabaseClient;
@@ -348,6 +361,7 @@ export const handleCreateCustomer = async ({
  params?: any;
  processor?: any;
  getDetails?: boolean;
+  createDefaultProducts?: boolean;
 }) => {
  const newCus = CreateCustomerSchema.parse(cusData);

@@ -361,6 +375,7 @@ export const handleCreateCustomer = async ({
      newCus,
      logger,
      processor,
+      createDefaultProducts,
    });
  } else {
    createdCustomer = await handleCreateCustomerWithId({
@@ -370,6 +385,7 @@ export const handleCreateCustomer = async ({
      logger,
      newCus,
      processor,
+      createDefaultProducts,
    });
  }

--- a/server/src/internal/api/customers/handlers/handleCusProductExpired.ts
+++ b/server/src/internal/api/customers/handlers/handleCusProductExpired.ts
@@ -8,6 +8,7 @@ import {
  fullCusProductToProduct,
 } from "@/internal/customers/products/cusProductUtils.js";
 import { OrgService } from "@/internal/orgs/OrgService.js";
+import { isOneOff } from "@/internal/products/productUtils.js";
 import RecaseError, { handleRequestError } from "@/utils/errorUtils.js";
 import {
  ErrCode,
@@ -15,14 +16,15 @@ import {
  FullCusProduct,
  Organization,
  AppEnv,
+  FullCustomer,
 } from "@autumn/shared";
 import { SupabaseClient } from "@supabase/supabase-js";
 import { StatusCodes } from "http-status-codes";

-export const expireCusProduct = async ({
+export const removeScheduledProduct = async ({
  sb,
-  cusProduct, // cus product to expire
-  cusProducts, // other cus products
+  cusProduct,
+  cusProducts,
  org,
  env,
  logger,
@@ -33,10 +35,57 @@ export const expireCusProduct = async ({
  org: Organization;
  env: AppEnv;
  logger: any;
+}) => {
+  const stripeCli = createStripeCli({ org: org, env: env });
+
+  // Get full product from cus product
+  let fullProduct = fullCusProductToProduct(cusProduct);
+
+  // 1. Cancel future product schedule
+  await cancelFutureProductSchedule({
+    sb,
+    org,
+    cusProducts,
+    product: fullProduct,
+    stripeCli,
+    logger,
+    env,
+  });
+
+  // 2. Delete scheduled product
+  await CusProductService.delete({
+    sb,
+    cusProductId: cusProduct.id,
+  });
+  return;
+};
+
+export const expireCusProduct = async ({
+  sb,
+  cusProduct, // cus product to expire
+  cusProducts, // other cus products
+  org,
+  env,
+  logger,
+  customer,
+  expireImmediately = true,
+}: {
+  sb: SupabaseClient;
+  cusProduct: FullCusProduct;
+  cusProducts: FullCusProduct[];
+  org: Organization;
+  env: AppEnv;
+  logger: any;
+  customer: FullCustomer;
+  expireImmediately: boolean;
 }) => {
  logger.info("--------------------------------");
-  logger.info(`🔔 Handling CusProduct Expired`);
-  logger.info(`Customer: ${cusProduct.customer.id} (${env}), Org: ${org.id}`);
+  logger.info(
+    `🔔 Expiring cutomer product (${
+      expireImmediately ? "immediately" : "end of cycle"
+    })`
+  );
+  logger.info(`Customer: ${customer.id} (${env}), Org: ${org.id}`);
  logger.info(
    `Product: ${cusProduct.product.name}, Status: ${cusProduct.status}`
  );
@@ -56,50 +105,22 @@ export const expireCusProduct = async ({
  // If current product is scheduled

  if (cusProduct.status == CusProductStatus.Scheduled) {
-    const stripeCli = createStripeCli({ org: org, env: env });
-
-    // Get full product from cus product
-    let fullProduct = fullCusProductToProduct(cusProduct);
-
-    // 1. Cancel future product schedule
-    await cancelFutureProductSchedule({
+    await removeScheduledProduct({
      sb,
-      org,
+      cusProduct,
      cusProducts,
-      product: fullProduct,
-      stripeCli,
-      logger,
+      org,
      env,
+      logger,
    });
+  }

-    // 2. Delete scheduled product
-    await CusProductService.delete({
-      sb,
-      cusProductId: cusProduct.id,
-    });
-  } else {
-    if (cusProduct.product.is_add_on) {
-      await cancelCusProductSubscriptions({
-        sb,
-        cusProduct,
-        org,
-        env,
-      });
-
-      await CusProductService.update({
-        sb,
-        cusProductId: cusProduct.id,
-        updates: {
-          status: CusProductStatus.Expired,
-          ended_at: Date.now(),
-        },
-      });
-
-      return;
-    }
+  // 1. If main product, can't expire if there's scheduled product
+  let isMain = !cusProduct.product.is_add_on;
+  if (isMain) {
    const futureProduct = await CusProductService.getFutureProduct({
      sb,
-      internalCustomerId: cusProduct.customer.internal_id,
+      internalCustomerId: customer.internal_id,
      productGroup: cusProduct.product.group,
    });

@@ -110,25 +131,67 @@ export const expireCusProduct = async ({
        statusCode: StatusCodes.BAD_REQUEST,
      });
    }
-    // For regular products
-    // 1. Cancel stripe subscriptions
-    const cancelled = await cancelCusProductSubscriptions({
+  }
+
+  // 2. If expire at cycle end, just cancel subscriptions
+  if (!expireImmediately) {
+    await cancelCusProductSubscriptions({
+      sb,
+      cusProduct,
+      org,
+      env,
+      expireImmediately,
+    });
+
+    if (isOneOff(cusProduct.customer_prices.map((p) => p.price))) {
+      await CusProductService.update({
+        sb,
+        cusProductId: cusProduct.id,
+        updates: { status: CusProductStatus.Expired },
+      });
+    }
+
+    return;
+  }
+
+  if (cusProduct.product.is_add_on) {
+    await cancelCusProductSubscriptions({
      sb,
      cusProduct,
      org,
      env,
    });

-    if (!cancelled) {
-      await expireAndActivate({
-        sb,
-        env,
-        cusProduct,
-        org,
-      });
-    } // else will be handled by webhook
+    await CusProductService.update({
+      sb,
+      cusProductId: cusProduct.id,
+      updates: {
+        status: CusProductStatus.Expired,
+        ended_at: Date.now(),
+      },
+    });
+
+    return;
  }

+  // For regular products
+  // 1. Cancel stripe subscriptions
+  const cancelled = await cancelCusProductSubscriptions({
+    sb,
+    cusProduct,
+    org,
+    env,
+  });
+
+  if (!cancelled) {
+    await expireAndActivate({
+      sb,
+      env,
+      cusProduct,
+      org,
+    });
+  } // else will be handled by webhook
+
  return;
 };

@@ -162,6 +225,8 @@ export const handleCusProductExpired = async (req: any, res: any) => {
      org,
      env: req.env,
      logger: req.logtail,
+      customer: cusProduct.customer,
+      expireImmediately: true,
    });

    res.status(200).json({ message: "Product expired" });
--- a/server/src/internal/api/customers/handlers/handleUpdateBalances.ts
+++ b/server/src/internal/api/customers/handlers/handleUpdateBalances.ts
@@ -25,11 +25,12 @@ import { notNullish } from "@/utils/genUtils.js";
 const getCusFeaturesAndOrg = async (req: any, customerId: string) => {
  // 1. Get customer
  const [customer, features, org] = await Promise.all([
-    CusService.getByIdOrInternalId({
+    CusService.getWithProducts({
      sb: req.sb,
      idOrInternalId: customerId,
      orgId: req.orgId,
      env: req.env,
+      entityId: req.params.entity_id,
    }),
    FeatureService.getFromReq(req),
    OrgService.getFromReq(req),
@@ -79,10 +80,8 @@ export const handleUpdateBalances = async (req: any, res: any) => {

    const { cusEnts, cusPrices } = await getCusEntsInFeatures({
      sb: req.sb,
-      internalCustomerId: customer.internal_id,
+      customer,
      internalFeatureIds: featuresToUpdate.map((f) => f.internal_id!),
-      inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue],
-      withPrices: true,
      logger: req.logtail,
    });

--- a/server/src/internal/api/customers/products/attachRouter.ts
+++ b/server/src/internal/api/customers/products/attachRouter.ts
@@ -256,7 +256,6 @@ export const checkStripeConnections = async ({
    );
  }
  await Promise.all(batchPriceUpdates);
-  // throw new Error("test");
 };

 export const customerHasPm = async ({
@@ -279,6 +278,7 @@ attachRouter.post("/attach", async (req: any, res) => {
    const {
      customer_id,
      product_id,
+      entity_id,
      customer_data,

      is_custom,
@@ -333,6 +333,7 @@ attachRouter.post("/attach", async (req: any, res) => {
      sb,
      customerId: customer_id,
      productId: product_id,
+      entityId: entity_id,
      customerData: customer_data,
      org,
      orgId: org.id,
--- a/server/src/internal/api/customers/products/expireRouter.ts
+++ b/server/src/internal/api/customers/products/expireRouter.ts
@@ -15,28 +15,36 @@ expireRouter.post("", async (req, res) =>
    action: "expire",
    handler: async (req, res) => {
      let { sb, orgId, env, logtail: logger } = req;
-      let { customer_id, product_id } = req.body;
+      let { customer_id, product_id, entity_id, cancel_immediately } = req.body;
+
+      let expireImmediately = cancel_immediately || false;

      let [customer, org] = await Promise.all([
-        CusService.getById({ sb, orgId, id: customer_id, env, logger }),
+        CusService.getWithProducts({
+          sb,
+          orgId,
+          idOrInternalId: customer_id,
+          env,
+          withEntities: true,
+          entityId: entity_id,
+          inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue],
+        }),
        OrgService.getFromReq(req),
      ]);

-      let cusProducts = await CusService.getFullCusProducts({
-        sb,
-        internalCustomerId: customer.internal_id,
-        withProduct: true,
-        withPrices: true,
-        logger,
-        inStatuses: [CusProductStatus.Active, CusProductStatus.Scheduled],
-      });
-
-      for (const cusProduct of cusProducts) {
-        cusProduct.customer = customer;
+      if (entity_id && !customer.entity) {
+        throw new RecaseError({
+          code: ErrCode.EntityNotFound,
+          message: `Entity ${entity_id} not found for customer ${customer_id}`,
+        });
      }

+      let cusProducts = customer.customer_products;
+
      let cusProductsToExpire = cusProducts.filter(
-        (cusProduct: FullCusProduct) => cusProduct.product.id == product_id
+        (cusProduct: FullCusProduct) =>
+          cusProduct.product.id == product_id &&
+          (entity_id ? cusProduct.entity_id == entity_id : true)
      );

      if (cusProductsToExpire.length == 0) {
@@ -54,6 +62,8 @@ expireRouter.post("", async (req, res) =>
          org,
          env,
          logger,
+          customer,
+          expireImmediately,
        });
      }

--- a/server/src/internal/api/entities/EntityService.ts
+++ b/server/src/internal/api/entities/EntityService.ts
@@ -41,6 +41,32 @@ export class EntityService {
    return data;
  }

+  static async getByInternalId({
+    sb,
+    internalId,
+    orgId,
+    env,
+  }: {
+    sb: SupabaseClient;
+    internalId: string;
+    orgId: string;
+    env: string;
+  }) {
+    const { data, error } = await sb
+      .from("entities")
+      .select("*")
+      .eq("internal_id", internalId)
+      .eq("org_id", orgId)
+      .eq("env", env)
+      .single();
+
+    if (error) {
+      throw error;
+    }
+
+    return data;
+  }
+
  static async get({
    sb,
    orgId,
--- a/server/src/internal/api/entities/entityUtils.ts
+++ b/server/src/internal/api/entities/entityUtils.ts
@@ -1,16 +1,23 @@
+import { submitUsageToStripe } from "@/external/stripe/stripeMeterUtils.js";
+import { createStripeCli } from "@/external/stripe/utils.js";
 import { CustomerEntitlementService } from "@/internal/customers/entitlements/CusEntitlementService.js";
 import {
  getBillingInterval,
  getBillingType,
+  roundUsage,
 } from "@/internal/prices/priceUtils.js";
 import { notNullish } from "@/utils/genUtils.js";
 import {
+  AppEnv,
  BillingType,
+  Customer,
  Entitlement,
  Entity,
  Feature,
  FullCustomerEntitlement,
  FullCustomerPrice,
+  Organization,
+  UsagePriceConfig,
 } from "@autumn/shared";
 import { SupabaseClient } from "@supabase/supabase-js";

@@ -78,12 +85,18 @@ export const removeEntityFromCusEnt = async ({
  entity,
  logger,
  cusPrice,
+  customer,
+  org,
+  env,
 }: {
  sb: SupabaseClient;
  cusEnt: FullCustomerEntitlement;
  entity: Entity;
  logger: any;
  cusPrice?: FullCustomerPrice;
+  customer: Customer;
+  org: Organization;
+  env: AppEnv;
 }) => {
  // isLinked
  let isLinked = isLinkedToEntity({
@@ -104,9 +117,30 @@ export const removeEntityFromCusEnt = async ({
  let newEntities = structuredClone(cusEnt.entities!);

  // TODO: Send usage to stripe if cus price exists
+  let stripeCli = createStripeCli({
+    org,
+    env,
+  });
  if (cusPrice) {
-    let billingType = getBillingType(cusPrice.price.config!);
-    console.log("Usage to send:", -newEntities[entity.id]?.balance);
+    let config = cusPrice.price.config as UsagePriceConfig;
+    let billingType = getBillingType(config);
+    if (billingType == BillingType.UsageInArrear) {
+      let usage = -newEntities[entity.id]?.balance;
+
+      usage = roundUsage({
+        usage,
+        billingUnits: config.billing_units!,
+      });
+
+      await submitUsageToStripe({
+        price: cusPrice.price,
+        usage,
+        customer,
+        feature: entitlement.feature,
+        logger,
+        stripeCli,
+      });
+    }
  }

  delete newEntities[entity.id];
--- a/server/src/internal/api/entities/handleCreateEntity.ts
+++ b/server/src/internal/api/entities/handleCreateEntity.ts
@@ -157,6 +157,7 @@ export const handleCreateEntity = async (req: any, res: any) => {
        orgId,
        env,
        inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue],
+        withEntities: true,
      }),
      FeatureService.getFromReq(req),
      OrgService.getFromReq(req),
@@ -195,15 +196,7 @@ export const handleCreateEntity = async (req: any, res: any) => {
    }

    let cusProducts = await customer.customer_products;
-
-    // Fetch existing
-    let existingEntities = await EntityService.get({
-      sb,
-      orgId,
-      env,
-      internalFeatureId: feature.internal_id,
-      internalCustomerId: customer.internal_id,
-    });
+    let existingEntities = customer.entities;

    logger.info("Existing entities:");
    logger.info(
@@ -275,10 +268,6 @@ export const handleCreateEntity = async (req: any, res: any) => {
      const newBalance =
        cusEnt.balance - (newCount + replacedCount) + (unused || 0);

-      // console.log("originalBalance", originalBalance);
-      // console.log("newBalance", newBalance);
-      // throw new Error("test");
-
      await adjustAllowance({
        sb,
        env,
@@ -292,6 +281,7 @@ export const handleCreateEntity = async (req: any, res: any) => {
        deduction: newCount + replacedCount,
        product,
        replacedCount,
+        fromEntities: true,
      });

      await req.pg.query(
--- a/server/src/internal/api/entities/handleDeleteEntity.ts
+++ b/server/src/internal/api/entities/handleDeleteEntity.ts
@@ -145,6 +145,9 @@ export const handleDeleteEntity = async (req: any, res: any) => {
          entity,
          logger,
          cusPrice: relatedCusPrice,
+          customer,
+          org,
+          env,
        });
      }

--- a/server/src/internal/api/entitled/entitledRouter.ts
+++ b/server/src/internal/api/entitled/entitledRouter.ts
@@ -101,8 +101,6 @@ const getMeteredEntitledResult = ({
    // 1. Skip if feature not among cusEnt

    if (!cusEntsContainFeature({ cusEnts, feature })) {
-      console.log("Feature not found", feature.id);
-
      continue;
    }

@@ -246,7 +244,8 @@ const getCusEntsAndFeatures = async ({
  sb: SupabaseClient;
  logger: any;
 }) => {
-  let { customer_id, feature_id, customer_data } = req.body;
+  let { customer_id, feature_id, customer_data, entity_id } = req.body;
+
  let { sb, orgId, env } = req;

  // 1. Get org and features
@@ -267,9 +266,18 @@ const getCusEntsAndFeatures = async ({
      customerData: customer_data,
      inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue],
      logger,
+      entityId: entity_id,
    }),
  ]);

+  if (entity_id && !customer.entity) {
+    throw new RecaseError({
+      message: `Entity ${entity_id} not found for customer ${customer_id}`,
+      code: ErrCode.EntityNotFound,
+      statusCode: StatusCodes.BAD_REQUEST,
+    });
+  }
+
  const { feature, creditSystems, allFeatures } = featureRes;

  const duration = Date.now() - startTime;
@@ -300,7 +308,25 @@ const getCusEntsAndFeatures = async ({
    });
  });

-  return { cusEnts, feature, creditSystems, org, cusProducts, allFeatures };
+  if (customer.entity) {
+    cusEnts = cusEnts.filter((cusEnt) => {
+      return (
+        notNullish(cusEnt.entities) ||
+        cusEnt.customer_product.internal_entity_id ===
+          customer.entity.internal_id
+      );
+    });
+  }
+
+  return {
+    cusEnts,
+    feature,
+    creditSystems,
+    org,
+    cusProducts,
+    allFeatures,
+    entity: customer.entity,
+  };
 };

 entitledRouter.post("", async (req: any, res: any) => {
--- a/server/src/internal/api/events/eventRouter.ts
+++ b/server/src/internal/api/events/eventRouter.ts
@@ -3,11 +3,13 @@ import {
  APIVersion,
  AppEnv,
  CreateEventSchema,
+  CusProductStatus,
  Customer,
  ErrCode,
  Event,
  Feature,
  FeatureType,
+  FullCustomer,
  Organization,
 } from "@autumn/shared";
 import RecaseError, { handleRequestError } from "@/utils/errorUtils.js";
@@ -24,6 +26,8 @@ import { StatusCodes } from "http-status-codes";
 import { getOrCreateCustomer } from "@/internal/customers/cusUtils/getOrCreateCustomer.js";
 import { FeatureService } from "@/internal/features/FeatureService.js";
 import { creditSystemContainsFeature } from "@/internal/features/creditSystemUtils.js";
+import { addTaskToQueue } from "@/queue/queueUtils.js";
+import { JobName } from "@/queue/JobName.js";

 export const eventsRouter = Router();

@@ -35,6 +39,7 @@ const getEventAndCustomer = async ({
  customer_data,
  event_data,
  logger,
+  entityId,
 }: {
  sb: SupabaseClient;
  org: Organization;
@@ -42,6 +47,7 @@ const getEventAndCustomer = async ({
  customer_id: string;
  customer_data: any;
  event_data: any;
+  entityId: string;
  logger: any;
 }) => {
  if (!customer_id) {
@@ -52,7 +58,7 @@ const getEventAndCustomer = async ({
    });
  }

-  let customer: Customer;
+  let customer: FullCustomer;

  // 2. Check if customer ID is valid
  customer = await getOrCreateCustomer({
@@ -62,6 +68,8 @@ const getEventAndCustomer = async ({
    customerId: customer_id,
    customerData: customer_data,
    logger,
+    entityId,
+    inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue],
  });

  // 3. Insert event
@@ -92,16 +100,10 @@ const getEventAndCustomer = async ({

 const getAffectedFeatures = async ({
  req,
-  pg,
  event,
-  orgId,
-  env,
 }: {
  req: any;
-  pg: Client;
  event: Event;
-  orgId: string;
-  env: AppEnv;
 }) => {
  let features = await FeatureService.getFromReq(req);

@@ -127,27 +129,6 @@ const getAffectedFeatures = async ({
  });

  return [...featuresWithEvent, ...creditSystems];
-
-  // const { rows }: { rows: Feature[] } = await pg.query(`
-  //   with features_with_event as (
-  //     select * from features
-  //     where org_id = '${orgId}'
-  //     and env = '${env}'
-  //     and config -> 'filters' @> '[{"value": ["${event.event_name}"]}]'::jsonb
-  //   )
-
-  //   select * from features WHERE
-  //   org_id = '${orgId}'
-  //   and env = '${env}'
-  //   and EXISTS (
-  //     SELECT 1 FROM jsonb_array_elements(config->'schema') as schema_element WHERE
-  //     schema_element->>'metered_feature_id' IN (SELECT id FROM features_with_event)
-  //   )
-  //   UNION all
-  //   select * from features_with_event
-  // `);
-
-  // return rows;
 };

 export const handleEventSent = async ({
@@ -179,14 +160,12 @@ export const handleEventSent = async ({
    customer_data,
    event_data,
    logger: req.logtail,
+    entityId: event_data.entity_id,
  });

  const affectedFeatures = await getAffectedFeatures({
    req,
-    pg: pg,
    event,
-    orgId,
-    env,
  });

  if (affectedFeatures.length == 0) {
@@ -199,36 +178,19 @@ export const handleEventSent = async ({

  if (affectedFeatures.length > 0) {
    const payload = {
-      customerId: customer.internal_id,
-      customer,
+      internalCustomerId: customer.internal_id,
+      customerId: customer.id,
+      entityId: event_data.entity_id,
      features: affectedFeatures,
      event,
      org,
      env,
    };

-    const queue = await QueueManager.getQueue({ useBackup: false });
-
-    try {
-      // Add timeout to queue operation
-      await queue.add("update-balance", payload);
-      // console.log("Added update-balance to queue");
-    } catch (error: any) {
-      try {
-        console.log("Adding update-balance to backup queue");
-        const backupQueue = await QueueManager.getQueue({ useBackup: true });
-        await backupQueue.add("update-balance", payload);
-      } catch (error: any) {
-        throw new RecaseError({
-          message: "Failed to add update-balance to queue (backup)",
-          code: "EVENT_QUEUE_ERROR",
-          statusCode: 500,
-          data: {
-            message: error.message,
-          },
-        });
-      }
-    }
+    await addTaskToQueue({
+      jobName: JobName.UpdateBalance,
+      payload,
+    });

    return { event, affectedFeatures, org };
  }
@@ -267,6 +229,4 @@ eventsRouter.post("", async (req: any, res: any) => {
    handleRequestError({ req, res, error, action: "POST event failed" });
    return;
  }
-
-  return;
 });
--- a/server/src/internal/api/events/usageRouter.ts
+++ b/server/src/internal/api/events/usageRouter.ts
@@ -1,5 +1,11 @@
 import { Router } from "express";
-import { Customer, ErrCode, Event, FeatureType } from "@autumn/shared";
+import {
+  CusProductStatus,
+  Customer,
+  ErrCode,
+  Event,
+  FeatureType,
+} from "@autumn/shared";
 import RecaseError, { handleRequestError } from "@/utils/errorUtils.js";
 import { generateId, nullish } from "@/utils/genUtils.js";

@@ -12,6 +18,7 @@ import { QueueManager } from "@/queue/QueueManager.js";
 import { JobName } from "@/queue/JobName.js";
 import { getOrCreateCustomer } from "@/internal/customers/cusUtils/getOrCreateCustomer.js";
 import { creditSystemContainsFeature } from "@/internal/features/creditSystemUtils.js";
+import { addTaskToQueue } from "@/queue/queueUtils.js";
 export const eventsRouter = Router();
 export const usageRouter = Router();

@@ -19,11 +26,13 @@ const getCusFeatureAndOrg = async ({
  req,
  customerId,
  featureId,
+  entityId,
  customerData,
 }: {
  req: any;
  customerId: string;
  featureId: string;
+  entityId: string;
  customerData: any;
 }) => {
  // 1. Get customer
@@ -36,6 +45,8 @@ const getCusFeatureAndOrg = async ({
      customerId,
      customerData,
      logger: req.logtail,
+      entityId,
+      inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue],
    }),
    FeatureService.getFromReq(req),
  ]);
@@ -126,6 +137,7 @@ export const handleUsageEvent = async ({
    customerId: customer_id,
    featureId: feature_id,
    customerData: customer_data,
+    entityId: entity_id,
  });

  let newEvent = await createAndInsertEvent({
@@ -148,8 +160,8 @@ export const handleUsageEvent = async ({
  }

  const payload = {
-    customerId: customer.internal_id,
-    customer,
+    customerId: customer.id,
+    internalCustomerId: customer.internal_id,
    features,
    org,
    env: req.env,
@@ -159,25 +171,30 @@ export const handleUsageEvent = async ({
    entityId: entity_id,
  };

-  try {
-    // Add timeout to queue operation
-    await queue.add(JobName.UpdateUsage, payload);
-  } catch (error: any) {
-    try {
-      console.log("Adding update-balance to backup queue");
-      const backupQueue = await QueueManager.getQueue({ useBackup: true });
-      await backupQueue.add(JobName.UpdateUsage, payload);
-    } catch (error: any) {
-      throw new RecaseError({
-        message: "Failed to add update-usage to queue (backup)",
-        code: "EVENT_QUEUE_ERROR",
-        statusCode: 500,
-        data: {
-          message: error.message,
-        },
-      });
-    }
-  }
+  await addTaskToQueue({
+    jobName: JobName.UpdateUsage,
+    payload,
+  });
+
+  // try {
+  //   // Add timeout to queue operation
+  //   await queue.add(JobName.UpdateUsage, payload);
+  // } catch (error: any) {
+  //   try {
+  //     console.log("Adding update-balance to backup queue");
+  //     const backupQueue = await QueueManager.getQueue({ useBackup: true });
+  //     await backupQueue.add(JobName.UpdateUsage, payload);
+  //   } catch (error: any) {
+  //     throw new RecaseError({
+  //       message: "Failed to add update-usage to queue (backup)",
+  //       code: "EVENT_QUEUE_ERROR",
+  //       statusCode: 500,
+  //       data: {
+  //         message: error.message,
+  //       },
+  //     });
+  //   }
+  // }

  return { event: newEvent, affectedFeatures: features, org };
 };
--- a/server/src/internal/customers/CusService.ts
+++ b/server/src/internal/customers/CusService.ts
@@ -38,18 +38,24 @@ export class CusService {
      CusProductStatus.PastDue,
      CusProductStatus.Scheduled,
    ],
+    withEntities = false,
+    entityId,
  }: {
    sb: SupabaseClient;
    idOrInternalId: string;
    orgId: string;
    env: AppEnv;
    inStatuses?: CusProductStatus[];
+    withEntities?: boolean;
+    entityId?: string;
  }) {
    const { data, error } = await sb.rpc("get_cus_with_products", {
      p_cus_id: idOrInternalId,
      p_org_id: orgId,
      p_env: env,
      p_statuses: inStatuses,
+      p_with_entities: withEntities,
+      p_entity_id: entityId,
    });

    if (error) {
@@ -60,7 +66,7 @@ export class CusService {
      return null;
    }

-    let { customer, products } = data;
+    let { customer, products, entities, entity } = data;

    if (!products) {
      products = [];
@@ -78,6 +84,8 @@ export class CusService {
    return {
      ...customer,
      customer_products: products,
+      entities: entities,
+      entity: entity,
    };
  }

--- a/server/src/internal/customers/add-product/createFullCusProduct.ts
+++ b/server/src/internal/customers/add-product/createFullCusProduct.ts
@@ -76,6 +76,8 @@ export const initCusProduct = ({
  subscriptionIds,
  subscriptionScheduleIds,
  isCustom,
+  entityId,
+  internalEntityId,
 }: {
  customer: Customer;
  product: FullProduct;
@@ -94,6 +96,8 @@ export const initCusProduct = ({
  subscriptionIds?: string[];
  subscriptionScheduleIds?: string[];
  isCustom?: boolean;
+  entityId?: string;
+  internalEntityId?: string;
 }) => {
  let isFuture = startsAt && startsAt > Date.now();

@@ -133,6 +137,8 @@ export const initCusProduct = ({
    scheduled_ids: subscriptionScheduleIds,
    is_custom: isCustom || false,
    quantity: 1,
+    internal_entity_id: internalEntityId,
+    entity_id: entityId,
  };
 };

@@ -191,18 +197,21 @@ export const expireOrDeleteCusProduct = async ({
  startsAt,
  product,
  cusProducts,
+  internalEntityId,
 }: {
  sb: SupabaseClient;
  startsAt?: number;
  product: FullProduct;
  cusProducts?: FullCusProduct[];
+  internalEntityId?: string;
 }) => {
  // 1. If startsAt
  if (startsAt && startsAt > Date.now()) {
    let curScheduledProduct = cusProducts?.find(
      (cp) =>
        cp.product.group === product.group &&
-        cp.status === CusProductStatus.Scheduled
+        cp.status === CusProductStatus.Scheduled &&
+        (internalEntityId ? cp.internal_entity_id === internalEntityId : true)
    );

    if (curScheduledProduct) {
@@ -215,6 +224,7 @@ export const expireOrDeleteCusProduct = async ({
    let { curMainProduct } = await getExistingCusProducts({
      product,
      cusProducts: cusProducts as FullCusProduct[],
+      internalEntityId,
    });

    if (curMainProduct) {
@@ -234,11 +244,13 @@ export const getExistingCusProduct = async ({
  cusProducts,
  product,
  internalCustomerId,
+  internalEntityId,
 }: {
  sb?: SupabaseClient;
  cusProducts?: FullCusProduct[];
  product: FullProduct;
  internalCustomerId: string;
+  internalEntityId?: string;
 }) => {
  if (!cusProducts) {
    cusProducts = await CusService.getFullCusProducts({
@@ -250,6 +262,7 @@ export const getExistingCusProduct = async ({
  const { curMainProduct } = await getExistingCusProducts({
    product,
    cusProducts: cusProducts as FullCusProduct[],
+    internalEntityId,
  });

  return curMainProduct;
@@ -319,6 +332,7 @@ export const createFullCusProduct = async ({
      cusProducts: attachParams.cusProducts,
      product,
      internalCustomerId: customer.internal_id,
+      internalEntityId: attachParams.internalEntityId,
    });
  } catch (error) {}

@@ -350,12 +364,6 @@ export const createFullCusProduct = async ({
  for (const entitlement of entitlements) {
    const options = getEntOptions(optionsList, entitlement);
    const relatedPrice = getEntRelatedPrice(entitlement, prices);
-    // const existingCusEnt = curCusProduct?.customer_entitlements.find(
-    //   (ce) => ce.internal_feature_id === entitlement.internal_feature_id
-    // );
-
-    // Update existing entitlement if one off
-
    const cusEnt: any = initCusEntitlement({
      entitlement,
      customer,
@@ -424,6 +432,8 @@ export const createFullCusProduct = async ({
    subscriptionIds,
    subscriptionScheduleIds,
    isCustom: attachParams.isCustom || false,
+    entityId: attachParams.entityId,
+    internalEntityId: attachParams.internalEntityId,
  });

  // Expire previous product if not one off
@@ -433,6 +443,7 @@ export const createFullCusProduct = async ({
      startsAt,
      product,
      cusProducts: attachParams.cusProducts,
+      internalEntityId: attachParams.internalEntityId,
    });
  }

--- a/server/src/internal/customers/add-product/handleAddProduct.ts
+++ b/server/src/internal/customers/add-product/handleAddProduct.ts
@@ -12,7 +12,11 @@ import chalk from "chalk";

 import { SupabaseClient } from "@supabase/supabase-js";
 import { createFullCusProduct } from "../add-product/createFullCusProduct.js";
-import { createStripeCli } from "@/external/stripe/utils.js";
+import {
+  createStripeCli,
+  stripeToAutumnInterval,
+  subToAutumnInterval,
+} from "@/external/stripe/utils.js";
 import { AttachParams, AttachResultSchema } from "../products/AttachParams.js";
 import { getPriceAmount } from "../../prices/priceUtils.js";
 import {
@@ -40,6 +44,8 @@ import {
  getNextStartOfMonthUnix,
 } from "@/internal/prices/billingIntervalUtils.js";
 import { SuccessCode } from "@autumn/shared";
+import { getStripeSubs } from "@/external/stripe/stripeSubUtils.js";
+import { nullish } from "@/utils/genUtils.js";

 export const handleBillNowPrices = async ({
  sb,
@@ -49,6 +55,7 @@ export const handleBillNowPrices = async ({
  fromRequest = true,
  carryExistingUsages = false,
  shouldPreview = false,
+  disableMerge = false,
 }: {
  sb: any;
  attachParams: AttachParams;
@@ -57,9 +64,11 @@ export const handleBillNowPrices = async ({
  fromRequest?: boolean;
  carryExistingUsages?: boolean;
  shouldPreview?: boolean;
+  disableMerge?: boolean;
 }) => {
  const logger = req.logtail;
-  let { org, customer, products, freeTrial, invoiceOnly } = attachParams;
+  let { org, customer, products, freeTrial, invoiceOnly, cusProducts } =
+    attachParams;

  if (attachParams.disableFreeTrial) {
    freeTrial = null;
@@ -75,30 +84,46 @@ export const handleBillNowPrices = async ({
  let subscriptions: Stripe.Subscription[] = [];
  let invoiceIds: string[] = [];

+  // Merge billing cycles...
+  let mergeCusProduct =
+    !disableMerge && org.config.merge_billing_cycles
+      ? cusProducts?.find((cp) =>
+          products.some((p) => p.group == cp.product.group)
+        )
+      : undefined;
+
+  let mergeSubs = await getStripeSubs({
+    stripeCli,
+    subIds: mergeCusProduct?.subscription_ids,
+  });
+
  for (const itemSet of itemSets) {
    if (itemSet.interval === BillingInterval.OneOff) {
      continue;
    }

-    const { items } = itemSet;
+    let mergeWithSub = mergeSubs.find(
+      (sub) => subToAutumnInterval(sub) == itemSet.interval
+    );

    let subscription;
    try {
-      // Should create 2 subscriptions
-
      let billingCycleAnchorUnix;
      if (org.config.anchor_start_of_month) {
        billingCycleAnchorUnix = getNextStartOfMonthUnix(itemSet.interval);
      }

      if (attachParams.billingAnchor) {
-        // Add interval to now
        billingCycleAnchorUnix = getAlignedIntervalUnix(
          attachParams.billingAnchor,
          itemSet.interval
        );
      }

+      if (mergeWithSub) {
+        billingCycleAnchorUnix = mergeWithSub.current_period_end * 1000;
+      }
+
      subscription = await createStripeSub({
        sb,
        stripeCli,
@@ -410,6 +435,7 @@ export const handleAddProduct = async ({
  fromRequest = true,
  carryExistingUsages = false,
  keepResetIntervals = false,
+  disableMerge = false,
 }: {
  req: {
    sb: SupabaseClient;
@@ -420,6 +446,7 @@ export const handleAddProduct = async ({
  fromRequest?: boolean;
  carryExistingUsages?: boolean;
  keepResetIntervals?: boolean;
+  disableMerge?: boolean;
 }) => {
  const logger = req.logtail;
  const { customer, products, prices } = attachParams;
@@ -462,6 +489,7 @@ export const handleAddProduct = async ({
      res,
      fromRequest,
      carryExistingUsages,
+      disableMerge,
    });

    return;
--- a/server/src/internal/customers/add-product/handleExistingProduct.ts
+++ b/server/src/internal/customers/add-product/handleExistingProduct.ts
@@ -29,9 +29,11 @@ import { ACTIVE_STATUSES } from "../products/CusProductService.js";
 export const getExistingCusProducts = async ({
  product,
  cusProducts,
+  internalEntityId,
 }: {
  product: Product;
  cusProducts: FullCusProduct[];
+  internalEntityId?: string;
 }) => {
  if (!cusProducts || cusProducts.length === 0) {
    return {
@@ -41,25 +43,34 @@ export const getExistingCusProducts = async ({
    };
  }

-  let curMainProduct = cusProducts.find(
-    (cp: any) =>
-      cp.product.group === product.group &&
-      !cp.product.is_add_on &&
-      // ACTIVE_STATUSES.includes(cp.status) &&
-      (cp.status == CusProductStatus.Active ||
-        cp.status == CusProductStatus.PastDue) &&
-      !isOneOff(cp.customer_prices.map((cp: any) => cp.price))
-  );
+  let curMainProduct = cusProducts.find((cp: any) => {
+    let sameGroup = cp.product.group === product.group;
+    let isMain = !cp.product.is_add_on;
+    let isActive =
+      cp.status == CusProductStatus.Active ||
+      cp.status == CusProductStatus.PastDue;
+
+    let oneOff = isOneOff(cp.customer_prices.map((cp: any) => cp.price));
+
+    let sameEntity = internalEntityId
+      ? cp.internal_entity_id === internalEntityId
+      : true;
+
+    return sameGroup && isMain && isActive && !oneOff && sameEntity;
+  });

  const curSameProduct = cusProducts!.find(
-    (cp: any) => cp.product.internal_id === product.internal_id
+    (cp: any) =>
+      cp.product.internal_id === product.internal_id &&
+      (internalEntityId ? cp.internal_entity_id === internalEntityId : true)
  );

  const curScheduledProduct = cusProducts!.find(
    (cp: any) =>
      cp.status === CusProductStatus.Scheduled &&
      cp.product.group === product.group &&
-      !cp.product.is_add_on
+      !cp.product.is_add_on &&
+      (internalEntityId ? cp.internal_entity_id === internalEntityId : true)
  );

  return { curMainProduct, curSameProduct, curScheduledProduct };
@@ -164,6 +175,7 @@ export const handleExistingProduct = async ({
    await getExistingCusProducts({
      product,
      cusProducts: cusProducts || [],
+      internalEntityId: attachParams.internalEntityId,
    });

  if (isOneOff(product.prices)) {
--- a/server/src/internal/customers/add-product/handleInvoiceOnly.ts
+++ b/server/src/internal/customers/add-product/handleInvoiceOnly.ts
@@ -1,224 +0,0 @@
-import { createStripeCli } from "@/external/stripe/utils.js";
-import { pricesOnlyOneOff } from "@/internal/prices/priceUtils.js";
-import { createFullCusProduct } from "./createFullCusProduct.js";
-import { InvoiceService } from "../invoices/InvoiceService.js";
-import {
-  AppEnv,
-  CollectionMethod,
-  CusProduct,
-  CusProductStatus,
-  Customer,
-  Organization,
-} from "@autumn/shared";
-import { SupabaseClient } from "@supabase/supabase-js";
-import { CusProductService } from "../products/CusProductService.js";
-import Stripe from "stripe";
-import {
-  getStripeSubItems,
-  pricesToInvoiceItems,
-} from "@/external/stripe/stripePriceUtils.js";
-import { AttachParams } from "../products/AttachParams.js";
-import { attachToInsertParams } from "@/internal/products/productUtils.js";
-import {
-  getInvoiceExpansion,
-  getStripeExpandedInvoice,
-} from "@/external/stripe/stripeInvoiceUtils.js";
-
-// export const removeCurrentProduct = async ({
-//   sb,
-//   customer,
-//   curCusProduct,
-//   org,
-//   env,
-// }: {
-//   sb: SupabaseClient;
-//   customer: Customer;
-//   curCusProduct: CusProduct;
-//   org: Organization;
-//   env: AppEnv;
-// }) => {
-//   console.log("   - Removing current product");
-//   // 1. Expire current product
-//   CusProductService.update({
-//     sb,
-//     cusProductId: curCusProduct.id,
-//     updates: {
-//       status: CusProductStatus.Expired,
-//     },
-//   });
-
-//   // Cancel stripe subscription
-//   const stripeCli = createStripeCli({ org, env: customer.env });
-
-//   for (const subId of curCusProduct.subscription_ids || []) {
-//     await stripeCli.subscriptions.cancel(subId, {
-//       prorate: true,
-//     });
-//   }
-// };
-
-// export const invoiceOnlyOneOff = async ({
-//   res,
-//   sb,
-//   attachParams,
-// }: {
-//   res: any;
-//   sb: SupabaseClient;
-//   attachParams: any;
-// }) => {
-//   // 1. Create stripe subscription (with invoice)
-
-//   const { org, env, customer, prices, product } = attachParams;
-
-//   const stripeCli = createStripeCli({ org, env });
-
-//   // Create invoice
-//   console.log("   - Creating stripe invoice");
-//   const invoice = await stripeCli.invoices.create({
-//     customer: customer.processor.id,
-//     collection_method: "send_invoice",
-//     days_until_due: 30,
-//   });
-
-//   // Add to invoice items
-//   await pricesToInvoiceItems({
-//     attachParams,
-//     stripeInvoiceId: invoice.id,
-//     sb,
-//     stripeCli,
-//   });
-
-//   // 2. Create full cus product
-//   console.log("   - Adding full cus product");
-//   await createFullCusProduct({
-//     sb,
-//     attachParams,
-
-//     lastInvoiceId: invoice.id,
-//   });
-
-//   // 3. Finalize invoice
-//   const finalizedInvoice = await stripeCli.invoices.finalizeInvoice(
-//     invoice.id,
-//     getInvoiceExpansion()
-//   );
-
-//   // 4. Create invoice from stripe
-//   await InvoiceService.createInvoiceFromStripe({
-//     sb,
-//     stripeInvoice: finalizedInvoice,
-//     internalCustomerId: customer.internal_id,
-//     org,
-//     productIds: [product.id],
-//     internalProductIds: [product.internal_id],
-//   });
-
-//   console.log("   ✅ Done");
-
-//   res.status(200).json({
-//     invoice_url: finalizedInvoice.hosted_invoice_url,
-//   });
-// };
-
-// export const handleInvoiceOnly = async ({
-//   req,
-//   res,
-//   attachParams,
-//   curCusProduct,
-// }: {
-//   req: any;
-//   res: any;
-//   attachParams: AttachParams;
-//   curCusProduct: any;
-// }) => {
-//   console.log("SCENARIO: INVOICE ONLY");
-//   const { org, customer, prices, products } = attachParams;
-
-//   // If current product, expire and cancel stripe subscription
-//   // if (curCusProduct && !product.is_add_on) {
-//   //   // Handle removal of current product
-//   // }
-
-//   // Handle one off prices
-//   if (pricesOnlyOneOff(prices)) {
-//     console.log("Handling one-off priced product (invoice only)");
-//     await invoiceOnlyOneOff({
-//       sb: req.sb,
-//       attachParams,
-//       res,
-//     });
-//     return;
-//   }
-
-//   // 1. Create stripe subscription (with invoice)
-//   console.log("   - Creating stripe subscription");
-//   const stripeCli = createStripeCli({ org, env: customer.env });
-//   const itemSets = await getStripeSubItems({
-//     attachParams,
-//   });
-
-//   let stripeSubs: Stripe.Subscription[] = [];
-//   for (const itemSet of itemSets) {
-//     const { items, subMeta } = itemSet;
-//     // Create subscription
-
-//     const stripeSub = await stripeCli.subscriptions.create({
-//       customer: customer.processor.id,
-//       collection_method: "send_invoice",
-//       days_until_due: 30,
-//       items,
-//       metadata: subMeta,
-//     });
-//     stripeSubs.push(stripeSub);
-
-//     await stripeCli.subscriptions.update(stripeSub.id, {
-//       collection_method: "charge_automatically",
-//     });
-//   }
-
-//   // 1. Add full cus product
-//   console.log("   - Adding full cus product");
-//   for (const product of products) {
-//     await createFullCusProduct({
-//       sb: req.sb,
-//       attachParams: attachToInsertParams(attachParams, product),
-//       subscriptionId: stripeSubs[0].id,
-//       subscriptionIds: stripeSubs.map((s) => s.id),
-//       lastInvoiceId: stripeSubs[0].latest_invoice as string,
-//     });
-//   }
-
-//   let firstInvoice;
-//   for (const stripeSub of stripeSubs) {
-//     // Get stripe invoice
-//     console.log("   - Inserting stripe invoice into db");
-//     // 1. Finalize invoice
-//     await stripeCli.invoices.finalizeInvoice(
-//       stripeSub.latest_invoice as string
-//     );
-
-//     const stripeInvoice = await getStripeExpandedInvoice({
-//       stripeCli,
-//       stripeInvoiceId: stripeSub.latest_invoice as string,
-//     });
-
-//     if (!firstInvoice) {
-//       firstInvoice = stripeInvoice;
-//     }
-
-//     await InvoiceService.createInvoiceFromStripe({
-//       sb: req.sb,
-//       stripeInvoice,
-//       internalCustomerId: customer.internal_id,
-//       org,
-//       productIds: products.map((p) => p.id),
-//       internalProductIds: products.map((p) => p.internal_id),
-//     });
-//   }
-
-//   console.log("   ✅ Done");
-
-//   res.status(200).json({
-//     invoice_url: firstInvoice?.hosted_invoice_url,
-//   });
-// };
--- a/server/src/internal/customers/add-product/handleSameProduct.ts
+++ b/server/src/internal/customers/add-product/handleSameProduct.ts
@@ -4,6 +4,7 @@ import {
  ErrCode,
  Feature,
  FullCusProduct,
+  FullCustomer,
  FullCustomerEntitlement,
  FullCustomerPrice,
  Organization,
@@ -31,6 +32,7 @@ import {
 import { fullCusProductToProduct } from "../products/cusProductUtils.js";
 import { isFreeProduct } from "@/internal/products/productUtils.js";
 import { SuccessCode } from "@autumn/shared";
+import { notNullish } from "@/utils/genUtils.js";

 const getOptionsToUpdate = (oldOptionsList: any[], newOptionsList: any[]) => {
  let differentOptionsExist = false;
@@ -297,7 +299,9 @@ export const handleSameMainProduct = async ({
    }
  }

-  if (optionsToUpdate.length === 0 && !curScheduledProduct) {
+  let isCanceled = notNullish(curMainProduct.canceled_at);
+
+  if (optionsToUpdate.length === 0 && !curScheduledProduct && !isCanceled) {
    // Update options
    throw new RecaseError({
      message: `Customer already has product ${product.name}, can't attach again`,
@@ -308,13 +312,13 @@ export const handleSameMainProduct = async ({

  let messages: string[] = [];

-  if (curScheduledProduct) {
-    // 1. Delete future product
-    const stripeCli = createStripeCli({
-      org,
-      env: customer.env,
-    });
+  // 1. Delete future product
+  const stripeCli = createStripeCli({
+    org,
+    env: customer.env,
+  });

+  if (curScheduledProduct) {
    await cancelFutureProductSchedule({
      sb,
      org,
@@ -334,6 +338,25 @@ export const handleSameMainProduct = async ({
    messages.push(
      `Removed scheduled product ${curScheduledProduct.product.name}`
    );
+  } else if (isCanceled) {
+    for (const subId of curMainProduct.subscription_ids || []) {
+      await stripeCli.subscriptions.update(subId, {
+        cancel_at: null,
+      });
+    }
+
+    let entities = attachParams.entities;
+    let entity = curMainProduct.internal_entity_id
+      ? entities.find(
+          (e) => e.internal_id === curMainProduct.internal_entity_id
+        )
+      : undefined;
+
+    messages.push(
+      `Successfully renewed product ${product.name}${
+        entity ? ` for entity ${entity.name || entity.id}` : ""
+      }`
+    );
  }

  // 2. Update quantities
@@ -359,8 +382,8 @@ export const handleSameMainProduct = async ({
      customer_id: customer.id,
      product_ids: products.map((p) => p.id),

-      code: SuccessCode.PrepaidQuantityUpdated,
-      message: `Successfully updated prepaid quantities for ${products
+      code: SuccessCode.RenewedProduct,
+      message: `Successfully renewed product ${products
        .map((p) => p.name)
        .join(", ")}`,
    })
--- a/server/src/internal/customers/change-product/changeProductUtils.ts
+++ b/server/src/internal/customers/change-product/changeProductUtils.ts
@@ -23,6 +23,7 @@ export const cancelScheduledProductIfExists = async ({
  let { curScheduledProduct } = await getExistingCusProducts({
    product: curFullProduct,
    cusProducts: attachParams.cusProducts!,
+    internalEntityId: attachParams.internalEntityId,
  });

  if (curScheduledProduct) {
--- a/server/src/internal/customers/change-product/handleDowngrade.ts
+++ b/server/src/internal/customers/change-product/handleDowngrade.ts
@@ -137,6 +137,11 @@ export const handleDowngrade = async ({

  // 1. Cancel all current subscriptions
  const intervalToOtherSubs: Record<string, any> = {};
+  console.log(
+    "Subs: ",
+    curSubscriptions.map((s) => s.id)
+  );
+
  for (const sub of curSubscriptions) {
    let latestEndDate = new Date(latestPeriodEnd * 1000);
    let curEndDate = new Date(sub.current_period_end * 1000);
--- a/server/src/internal/customers/change-product/handleUpgrade.ts
+++ b/server/src/internal/customers/change-product/handleUpgrade.ts
@@ -400,6 +400,7 @@ export const handleUpgrade = async ({
      fromRequest: fromReq,
      carryExistingUsages,
      keepResetIntervals: newVersion, // keep reset intervals if upgrading version (migrations)
+      disableMerge: true,
    });

    if (notNullish(curCusProduct.subscription_ids)) {
--- a/server/src/internal/customers/cusUtils/getOrCreateCustomer.ts
+++ b/server/src/internal/customers/cusUtils/getOrCreateCustomer.ts
@@ -6,9 +6,12 @@ import {
  AppEnv,
  CusProductStatus,
  CustomerData,
+  ErrCode,
  FullCustomer,
  Organization,
 } from "@autumn/shared";
+import { StatusCodes } from "http-status-codes";
+import RecaseError from "@/utils/errorUtils.js";

 export const getOrCreateCustomer = async ({
  sb,
@@ -23,6 +26,8 @@ export const getOrCreateCustomer = async ({
    CusProductStatus.Scheduled,
  ],
  skipGet = false,
+  withEntities = false,
+  entityId,
 }: {
  sb: SupabaseClient;
  org: Organization;
@@ -32,6 +37,8 @@ export const getOrCreateCustomer = async ({
  logger: any;
  inStatuses?: CusProductStatus[];
  skipGet?: boolean;
+  withEntities?: boolean;
+  entityId?: string;
 }): Promise<FullCustomer> => {
  let customer;

@@ -42,6 +49,8 @@ export const getOrCreateCustomer = async ({
      orgId: org.id,
      env,
      inStatuses,
+      withEntities,
+      entityId,
    });
  }

@@ -68,6 +77,8 @@ export const getOrCreateCustomer = async ({
        orgId: org.id,
        env,
        inStatuses,
+        withEntities,
+        entityId,
      });
    } catch (error: any) {
      if (error?.data?.code == "23505") {
@@ -77,6 +88,8 @@ export const getOrCreateCustomer = async ({
          orgId: org.id,
          env,
          inStatuses,
+          withEntities,
+          entityId,
        });
      } else {
        throw error;
@@ -91,5 +104,12 @@ export const getOrCreateCustomer = async ({
    logger,
  });

+  if (entityId && !customer.entity) {
+    throw new RecaseError({
+      message: `Entity ${entityId} not found for customer ${customerId}`,
+      code: ErrCode.EntityNotFound,
+      statusCode: StatusCodes.BAD_REQUEST,
+    });
+  }
  return customer as FullCustomer;
 };
--- a/server/src/internal/customers/entitlements/cusEntUtils.ts
+++ b/server/src/internal/customers/entitlements/cusEntUtils.ts
@@ -176,7 +176,9 @@ export const getCusEntBalance = ({
 };

 export const sortCusEntsForDeduction = (
-  cusEnts: FullCustomerEntitlement[],
+  cusEnts: (FullCustomerEntitlement & {
+    customer_product?: FullCusProduct;
+  })[],
  reverseOrder: boolean = false
 ) => {
  let intervalOrder: Record<EntInterval, number> = {
@@ -263,8 +265,7 @@ export const sortCusEntsForDeduction = (
    }

    // 3. Sort by interval
-
-    if (aEnt.interval && bEnt.interval) {
+    if (aEnt.interval && bEnt.interval && aEnt.interval != bEnt.interval) {
      if (reverseOrder) {
        return intervalOrder[bEnt.interval] - intervalOrder[aEnt.interval];
      } else {
@@ -272,6 +273,18 @@ export const sortCusEntsForDeduction = (
      }
    }

+    // Check if a is main product
+    let aIsAddOn = a.customer_product?.product?.is_add_on;
+    let bIsAddOn = b.customer_product?.product?.is_add_on;
+
+    if (aIsAddOn && !bIsAddOn) {
+      return 1;
+    }
+
+    if (!aIsAddOn && bIsAddOn) {
+      return -1;
+    }
+
    // 4. Sort by created_at
    return a.created_at - b.created_at;
  });
@@ -494,7 +507,7 @@ export const getTotalNegativeBalance = ({
          .mul(billingUnits)
          .toNumber();
      }
-      totalNegative += entities[group].balance;
+      totalNegative += balance;
    }
  }

@@ -507,11 +520,13 @@ export const getExistingUsageFromCusProducts = ({
  cusProducts,
  entities,
  carryExistingUsages = false,
+  internalEntityId,
 }: {
  entitlement: EntitlementWithFeature;
  cusProducts?: FullCusProduct[];
  entities: Entity[];
  carryExistingUsages?: boolean;
+  internalEntityId?: string;
 }) => {
  if (!entitlement || entitlement.feature.type === FeatureType.Boolean) {
    return 0;
@@ -535,7 +550,11 @@ export const getExistingUsageFromCusProducts = ({
  // NOTE: Assuming that feature entitlements are unique to each main product...
  let existingCusEnt = cusProducts
    ?.filter(
-      (cp) => cp.status === CusProductStatus.Active && !cp.product.is_add_on
+      (cp) =>
+        (cp.status === CusProductStatus.Active ||
+          cp.status === CusProductStatus.PastDue) &&
+        !cp.product.is_add_on &&
+        (internalEntityId ? cp.internal_entity_id === internalEntityId : true)
    )
    .flatMap((cp) => cp.customer_entitlements)
    .find((ce) => ce.internal_feature_id === entitlement.internal_feature_id);
--- a/server/src/internal/customers/products/AttachParams.ts
+++ b/server/src/internal/customers/products/AttachParams.ts
@@ -42,6 +42,9 @@ export type AttachParams = {
  isCustom?: boolean;
  disableFreeTrial?: boolean;
  features: Feature[];
+
+  entityId?: string;
+  internalEntityId?: string;
 };

 export type InsertCusProductParams = {
@@ -69,6 +72,9 @@ export type InsertCusProductParams = {
  isCustom?: boolean;
  disableFreeTrial?: boolean;
  features: Feature[];
+
+  entityId?: string;
+  internalEntityId?: string;
 };

 export const AttachResultSchema = z.object({
--- a/server/src/internal/customers/products/CusProductService.ts
+++ b/server/src/internal/customers/products/CusProductService.ts
@@ -254,6 +254,7 @@ export class CusProductService {
    env,
    inStatuses,
    withCusEnts = false,
+    withCusPrices = false,
  }: {
    sb: SupabaseClient;
    stripeSubId: string;
@@ -261,6 +262,7 @@ export class CusProductService {
    env: AppEnv;
    inStatuses?: string[];
    withCusEnts?: boolean;
+    withCusPrices?: boolean;
  }) {
    const query = sb
      .from("customer_products")
@@ -269,6 +271,10 @@ export class CusProductService {
          withCusEnts
            ? ", customer_entitlements:customer_entitlements(*, entitlement:entitlements!inner(*, feature:features!inner(*)))"
            : ""
+        }${
+          withCusPrices
+            ? ", customer_prices:customer_prices(*, price:prices!inner(*))"
+            : ""
        }` as "*"
      )
      .or(
--- a/server/src/internal/customers/products/attachUtils.ts
+++ b/server/src/internal/customers/products/attachUtils.ts
@@ -6,6 +6,7 @@ import {
  CustomerData,
  Entitlement,
  EntitlementWithFeature,
+  Entity,
  Feature,
  FeatureOptions,
  FreeTrial,
@@ -20,9 +21,7 @@ import { SupabaseClient } from "@supabase/supabase-js";
 import { ErrCode } from "@/errors/errCodes.js";
 import RecaseError from "@/utils/errorUtils.js";
 import { ProductService } from "@/internal/products/ProductService.js";
-import { OrgService } from "@/internal/orgs/OrgService.js";
 import { notNullish, nullish } from "@/utils/genUtils.js";
-import { FeatureService } from "@/internal/features/FeatureService.js";

 import {
  getFreeTrialAfterFingerprint,
@@ -30,7 +29,6 @@ import {
 } from "@/internal/products/free-trials/freeTrialUtils.js";

 import { StatusCodes } from "http-status-codes";
-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";
@@ -39,41 +37,6 @@ import { handleNewProductItems } from "@/internal/products/product-items/product
 import { getBillingType } from "@/internal/prices/priceUtils.js";
 import { Decimal } from "decimal.js";

-const getOrCreateCustomerAndProducts = async ({
-  sb,
-  customerId,
-  customerData,
-  org,
-  env,
-  logger,
-}: {
-  sb: SupabaseClient;
-  customerId: string;
-  customerData?: CustomerData;
-  org: Organization;
-  env: AppEnv;
-  logger: any;
-}) => {
-  const customer = await getOrCreateCustomer({
-    sb,
-    org,
-    env,
-    customerId,
-    customerData,
-    logger,
-    inStatuses: [
-      CusProductStatus.Active,
-      CusProductStatus.Scheduled,
-      CusProductStatus.PastDue,
-    ],
-  });
-
-  // Handle existing cus product...
-  const cusProducts = customer.customer_products;
-
-  return { customer, cusProducts };
-};
-
 const getProducts = async ({
  sb,
  productId,
@@ -187,7 +150,7 @@ const getProducts = async ({
  return [];
 };

-const getCustomerProductsFeaturesAndOrg = async ({
+const getCustomerAndProducts = async ({
  sb,
  org,
  customerId,
@@ -198,6 +161,7 @@ const getCustomerProductsFeaturesAndOrg = async ({
  env,
  logger,
  version,
+  entityId,
 }: {
  sb: SupabaseClient;
  org: Organization;
@@ -208,20 +172,30 @@ const getCustomerProductsFeaturesAndOrg = async ({
  env: AppEnv;
  logger: any;
  version?: number;
+  entityId?: string;
 }) => {
-  const [cusRes, products] = await Promise.all([
-    getOrCreateCustomerAndProducts({
+  const [customer, products] = await Promise.all([
+    getOrCreateCustomer({
      sb,
-      customerId,
-      customerData,
      org,
      env,
+      customerId,
+      customerData,
      logger,
+      inStatuses: [
+        CusProductStatus.Active,
+        CusProductStatus.Scheduled,
+        CusProductStatus.PastDue,
+      ],
+      entityId,
+      withEntities: true,
    }),
    getProducts({ sb, productId, productIds, orgId: org.id, env, version }),
  ]);

-  return { ...cusRes, products };
+  let cusProducts = customer.customer_products;
+
+  return { customer, cusProducts, products };
 };

 const getEntsWithFeature = (ents: Entitlement[], features: Feature[]) => {
@@ -300,6 +274,7 @@ export const getFullCusProductData = async ({
  customerId,
  customerData,
  productId,
+  entityId,
  productIds,
  orgId,
  itemsInput,
@@ -315,6 +290,7 @@ export const getFullCusProductData = async ({
  sb: SupabaseClient;
  customerId: string;
  customerData?: Customer;
+  entityId?: string;
  productId?: string;
  productIds?: string[];
  orgId: string;
@@ -327,24 +303,17 @@ export const getFullCusProductData = async ({
  version?: number;
 }) => {
  // 1. Get customer, product, org & features
-  const { customer, products, cusProducts } =
-    await getCustomerProductsFeaturesAndOrg({
-      org,
-      sb,
-      customerId,
-      customerData,
-      productId,
-      productIds,
-      env,
-      logger,
-      version,
-    });
-
-  const entities = await EntityService.get({
+  const { customer, products, cusProducts } = await getCustomerAndProducts({
+    org,
    sb,
-    internalCustomerId: customer.internal_id,
-    orgId,
+    customerId,
+    customerData,
+    productId,
+    productIds,
    env,
+    logger,
+    version,
+    entityId,
  });

  if (!isCustom) {
@@ -359,8 +328,6 @@ export const getFullCusProductData = async ({
      });
    }

-
-
    return {
      customer,
      products,
@@ -377,8 +344,11 @@ export const getFullCusProductData = async ({
        .flat() as EntitlementWithFeature[],
      freeTrial,
      cusProducts,
-      entities,
-
+      entities: customer.entities,
+      entityId: entityId,
+      internalEntityId: entityId
+        ? customer.entities.find((e) => e.id === entityId)?.internal_id
+        : undefined,
    };
  }

@@ -454,6 +424,10 @@ export const getFullCusProductData = async ({
    entitlements: entitlements as EntitlementWithFeature[],
    freeTrial: uniqueFreeTrial,
    cusProducts,
-    entities,
+    entities: customer.entities,
+    entityId: entityId,
+    internalEntityId: entityId
+      ? customer.entities.find((e) => e.internal_id === entityId)?.internal_id
+      : undefined,
  };
 };
--- a/server/src/internal/customers/products/cusProductUtils.ts
+++ b/server/src/internal/customers/products/cusProductUtils.ts
@@ -23,6 +23,7 @@ import { createFullCusProduct } from "../add-product/createFullCusProduct.js";
 import Stripe from "stripe";
 import {
  deleteScheduledIds,
+  getStripeSubs,
  subIsPrematurelyCanceled,
 } from "@/external/stripe/stripeSubUtils.js";
 import { sortCusEntsForDeduction } from "../entitlements/cusEntUtils.js";
@@ -89,12 +90,14 @@ export const cancelCusProductSubscriptions = async ({
  org,
  env,
  excludeIds,
+  expireImmediately = true,
 }: {
  sb: SupabaseClient;
  cusProduct: FullCusProduct;
  org: Organization;
  env: AppEnv;
  excludeIds?: string[];
+  expireImmediately?: boolean;
 }) => {
  // 1. Cancel all subscriptions
  const stripeCli = createStripeCli({
@@ -102,13 +105,31 @@ export const cancelCusProductSubscriptions = async ({
    env: env,
  });

+  let latestSubEnd: number | undefined;
+  if (cusProduct.subscription_ids && cusProduct.subscription_ids.length > 0) {
+    let stripeSubs = await getStripeSubs({
+      stripeCli,
+      subIds: cusProduct.subscription_ids,
+    });
+
+    latestSubEnd = stripeSubs[0].current_period_end;
+  }
+
  const cancelStripeSub = async (subId: string) => {
    if (excludeIds && excludeIds.includes(subId)) {
      return;
    }

    try {
-      await stripeCli.subscriptions.cancel(subId);
+      if (expireImmediately) {
+        await stripeCli.subscriptions.cancel(subId);
+      } else {
+        await stripeCli.subscriptions.update(subId, {
+          cancel_at: latestSubEnd || undefined,
+          cancel_at_period_end: latestSubEnd ? undefined : true,
+        });
+      }
+
      console.log(
        `Cancelled stripe subscription ${subId}, org: ${org.slug}, product: ${cusProduct.product.name}, customer: ${cusProduct.customer.id}`
      );
--- a/server/src/internal/features/displayUtils.ts
+++ b/server/src/internal/features/displayUtils.ts
@@ -5,20 +5,27 @@ import { Organization } from "@autumn/shared";
 export const getFeatureName = ({
  feature,
  plural,
+  capitalize = false,
 }: {
  feature: Feature;
  plural: boolean;
+  capitalize: boolean;
 }) => {
  let featureName = feature.name;
-  if (!feature.display) {
-    return featureName;
+
+  if (feature.display) {
+    if (plural) {
+      featureName = feature.display.plural || featureName;
+    } else {
+      featureName = feature.display.singular || featureName;
+    }
  }

-  if (plural) {
-    return feature.display.plural || featureName;
+  if (capitalize) {
+    featureName = featureName.charAt(0).toUpperCase() + featureName.slice(1);
  }

-  return feature.display.singular || featureName;
+  return featureName;
 };

 export const getFeatureNameWithCapital = ({
--- a/server/src/internal/prices/priceUtils.ts
+++ b/server/src/internal/prices/priceUtils.ts
@@ -390,3 +390,21 @@ export const priceToAmountOrTiers = (price: Price) => {
    }
  }
 };
+
+export const roundUsage = ({
+  usage,
+  billingUnits,
+}: {
+  usage: number;
+  billingUnits: number;
+}) => {
+  if (!billingUnits || billingUnits == 1) {
+    return usage;
+  }
+
+  return new Decimal(usage)
+    .div(billingUnits)
+    .ceil()
+    .mul(billingUnits)
+    .toNumber();
+};
--- a/server/src/internal/products/product-items/productItemUtils.ts
+++ b/server/src/internal/products/product-items/productItemUtils.ts
@@ -113,17 +113,20 @@ export const constructFeatureItem = ({
  included_usage,
  interval,
  entitlement_id,
+  entity_feature_id,
 }: {
  feature_id: string;
  included_usage?: number | string;
  interval?: EntInterval;
  entitlement_id?: string;
+  entity_feature_id?: string;
 }) => {
  let item: ProductItem = {
    feature_id,
    included_usage: included_usage as number,
    interval: entToItemInterval(interval),
    entitlement_id,
+    entity_feature_id,
  };

  return item;
@@ -146,6 +149,7 @@ export const constructPriceItem = ({

 export const constructFeaturePriceItem = ({
  feature_id,
+
  feature_type,
  included_usage,
  price,
@@ -153,15 +157,17 @@ export const constructFeaturePriceItem = ({
  usage_model,
  billing_units = 1,
  reset_usage_when_enabled = false,
+  entity_feature_id,
 }: {
  feature_id: string;
-  feature_type: ProductItemFeatureType;
+  feature_type?: ProductItemFeatureType;
  included_usage?: number;
  price: number;
  interval: BillingInterval;
  usage_model?: UsageModel;
  billing_units?: number;
  reset_usage_when_enabled?: boolean;
+  entity_feature_id?: string;
 }) => {
  let item: ProductItem & {
    included_usage: number;
@@ -174,6 +180,7 @@ export const constructFeaturePriceItem = ({
    usage_model,
    billing_units,
    reset_usage_when_enabled,
+    entity_feature_id,
  };

  return item;
--- a/server/src/queue/queue.ts
+++ b/server/src/queue/queue.ts
@@ -119,9 +119,15 @@ const initWorker = ({
        return;
      }

-      const { customerId } = job.data; // customerId is internal customer id
+      const { internalCustomerId } = job.data; // customerId is internal customer id

-      while (!(await acquireLock({ customerId, timeout: 10000, useBackup }))) {
+      while (
+        !(await acquireLock({
+          customerId: internalCustomerId,
+          timeout: 10000,
+          useBackup,
+        }))
+      ) {
        await queue.add(job.name, job.data, {
          delay: 50,
        });
@@ -141,7 +147,7 @@ const initWorker = ({
      } catch (error) {
        console.error("Error processing job:", error);
      } finally {
-        await releaseLock({ customerId, useBackup });
+        await releaseLock({ customerId: internalCustomerId, useBackup });
      }
    },
    {
--- a/server/src/trigger/adjustAllowance.ts
+++ b/server/src/trigger/adjustAllowance.ts
@@ -153,7 +153,9 @@ export const adjustAllowance = async ({
    `   - New quantity = ${paidUsage} (paid) + ${cusEnt.entitlement.allowance} (allowance) = ${quantity} `
  );

-  let prorationBehaviour = "create_prorations";
+  let prorationBehaviour = org.config.bill_upgrade_immediately
+    ? "always_invoice"
+    : "create_prorations";

  // If prorate unused is false, then remove end of cycle
  if (!org.config.prorate_unused) {
@@ -256,12 +258,22 @@ export const adjustAllowance = async ({
    await stripeCli.subscriptionItems.update(subItem.id, {
      quantity: quantity,
      proration_behavior: prorationBehaviour as any,
+      payment_behavior: fromEntities ? "error_if_incomplete" : undefined,
    });
    logger.info(`   ✅ Adjusted sub item ${subItem.id} to ${quantity}`);
  } catch (error: any) {
-    logger.error(`❗️ Error updating subscription item`);
-    logger.error(error);
-    return;
+    if (fromEntities) {
+      throw new RecaseError({
+        message: `Failed to update subscription subscription: ${error.message}`,
+        code: ErrCode.StripeUpdateSubscriptionFailed,
+        statusCode: error.statusCode,
+      });
+    } else {
+      logger.error(
+        `❗️ adjustAllowance: Error updating subscription item (from event)`
+      );
+      logger.error(error);
+    }
  }

  return;
--- a/server/src/trigger/updateBalanceTask.ts
+++ b/server/src/trigger/updateBalanceTask.ts
@@ -5,6 +5,7 @@ import {
  CusProductStatus,
  Event,
  Feature,
+  FullCustomer,
  FullCustomerEntitlement,
  FullCustomerPrice,
  Organization,
@@ -31,6 +32,7 @@ import {
  getTotalNegativeBalance,
 } from "@/internal/customers/entitlements/cusEntUtils.js";
 import { entityFeatureIdExists } from "@/internal/api/entities/entityUtils.js";
+import { CusService } from "@/internal/customers/CusService.js";

 // Decimal.set({ precision: 12 }); // 12 DP precision

@@ -127,7 +129,6 @@ export const logBalanceUpdate = ({
  entityId?: string | null;
  org: Organization;
 }) => {
-  console.log(`   - getCusEntsInFeatures: ${timeTaken}ms`);
  console.log(
    `   - Customer: ${customer.id} (${customer.env}) | Org: ${
      org.slug
@@ -307,21 +308,12 @@ export const deductAllowanceFromCusEnt = async ({
  entityId?: string | null;
  setZeroAdjustment?: boolean;
 }) => {
-  const { sb, feature, env, org, cusPrices, customer, properties } =
-    deductParams;
+  const { sb, feature, env, org, cusPrices, customer } = deductParams;

  if (toDeduct == 0) {
    return 0;
  }

-  // Either deduct from balance or entity balance
-
-  // let newBalance = structuredClone(cusEnt.balance);
-  // let newEntities = structuredClone(cusEnt.entities);
-  // let deducted = 0;
-
-  // console.log("entityId", entityId);
-
  let {
    newBalance,
    newEntities,
@@ -346,10 +338,6 @@ export const deductAllowanceFromCusEnt = async ({
    balance: newBalance!,
    entities: newEntities!,
  });
-  // console.log("Saving:", {
-  //   balance: newBalance,
-  //   entities: newEntities,
-  // });

  let updates: any = {
    balance: newBalance,
@@ -358,6 +346,7 @@ export const deductAllowanceFromCusEnt = async ({
  if (setZeroAdjustment) {
    updates.adjustment = 0;
  }
+
  await CustomerEntitlementService.update({
    sb,
    id: cusEnt.id,
@@ -518,7 +507,8 @@ export const deductFromUsageBasedCusEnt = async ({
 // Main function to update customer balance
 export const updateCustomerBalance = async ({
  sb,
-  customer,
+  customerId,
+  entityId,
  event,
  features,
  org,
@@ -526,7 +516,8 @@ export const updateCustomerBalance = async ({
  logger,
 }: {
  sb: SupabaseClient;
-  customer: Customer;
+  customerId: string;
+  entityId: string;
  event: Event;
  features: Feature[];
  org: Organization;
@@ -535,13 +526,19 @@ export const updateCustomerBalance = async ({
 }) => {
  const startTime = performance.now();
  console.log("REVERSE DEDUCTION ORDER", org.config.reverse_deduction_order);
+  const customer = await CusService.getWithProducts({
+    sb,
+    idOrInternalId: customerId,
+    orgId: org.id,
+    env,
+    inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue],
+    entityId,
+  });

  const { cusEnts, cusPrices } = await getCusEntsInFeatures({
    sb,
-    internalCustomerId: customer.internal_id,
+    customer,
    internalFeatureIds: features.map((f) => f.internal_id!),
-    inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue],
-    withPrices: true,
    logger,
    reverseOrder: org.config.reverse_deduction_order,
  });
@@ -566,14 +563,6 @@ export const updateCustomerBalance = async ({
    entityId: event.entity_id,
  });

-  // // 2. Handle group_by initialization
-  // await initGroupBalancesForEvent({
-  //   sb,
-  //   features,
-  //   cusEnts,
-  //   properties: event.properties,
-  // });
-
  // 3. Return if no customer entitlements or features found
  if (cusEnts.length === 0 || features.length === 0) {
    console.log("   - No customer entitlements or features found");
@@ -643,49 +632,28 @@ export const runUpdateBalanceTask = async ({
 }) => {
  try {
    // 1. Update customer balance
-    const { customer, features, event, org, env } = payload;
+    const { customerId, features, event, org, env, entityId } = payload;

    console.log("--------------------------------");
    console.log(
-      `UPDATING BALANCE FOR CUSTOMER (${customer.id}), ORG: ${org.slug}`
+      `UPDATING BALANCE FOR CUSTOMER (${customerId}), ORG: ${org.slug}`
    );

    const cusEnts: any = await updateCustomerBalance({
      sb,
-      customer,
+      customerId,
      features,
      event,
      org,
      env,
      logger,
+      entityId,
    });

    if (!cusEnts || cusEnts.length === 0) {
      return;
    }
    console.log("   ✅ Customer balance updated");
-
-    // // 2. Check if there's below threshold price
-    // const belowThresholdPrice = await getBelowThresholdPrice({
-    //   sb,
-    //   internalCustomerId: customer.internal_id,
-    //   cusEnts,
-    // });
-
-    // if (belowThresholdPrice) {
-    //   console.log("2. Below threshold price found");
-
-    //   // await new Promise((resolve) => setTimeout(resolve, 1000));
-
-    //   await handleBelowThresholdInvoicing({
-    //     sb,
-    //     internalCustomerId: customer.internal_id,
-    //     belowThresholdPrice,
-    //     logger,
-    //   });
-    // } else {
-    //   console.log("   ✅ No below threshold price found");
-    // }
  } catch (error) {
    if (logger) {
      logger.use((log: any) => {
--- a/server/src/trigger/updateUsageTask.ts
+++ b/server/src/trigger/updateUsageTask.ts
@@ -10,6 +10,7 @@ import {
  Customer,
  Feature,
  FeatureType,
+  FullCustomer,
  FullCustomerEntitlement,
  Organization,
 } from "@autumn/shared";
@@ -28,6 +29,7 @@ import {
  deductFromUsageBasedCusEnt,
 } from "./updateBalanceTask.js";
 import { JobName } from "@/queue/JobName.js";
+import { CusService } from "@/internal/customers/CusService.js";

 // 2. Get deductions for each feature
 const getFeatureDeductions = ({
@@ -114,25 +116,22 @@ const getFeatureDeductions = ({
 };

 const logUsageUpdate = ({
-  timeTaken,
  customer,
  features,
  cusEnts,
  featureDeductions,
-  properties,
  org,
  setUsage,
+  entityId,
 }: {
-  timeTaken: string;
  customer: Customer;
  features: Feature[];
  cusEnts: FullCustomerEntitlement[];
  featureDeductions: any;
-  properties: any;
  org: Organization;
  setUsage: boolean;
+  entityId?: string;
 }) => {
-  console.log(`   - getCusEntsInFeatures: ${timeTaken}ms`);
  console.log(
    `   - Customer: ${customer.id} (${customer.env}) | Org: ${
      org.slug
@@ -145,12 +144,6 @@ const logUsageUpdate = ({
    "   - CusEnts:",
    cusEnts.map((cusEnt: any) => {
      let balanceStr = cusEnt.balance;
-      // let { groupVal, balance } = getGroupBalanceFromProperties({
-      //   properties,
-      //   cusEnt,
-      //   features,
-      // });
-
      try {
        if (cusEnt.entitlement.allowance_type === AllowanceType.Unlimited) {
          balanceStr = "Unlimited";
@@ -159,6 +152,10 @@ const logUsageUpdate = ({
        balanceStr = "failed_to_get_balance";
      }

+      if (entityId && cusEnt.entities) {
+        balanceStr = `${cusEnt.entities?.[entityId!]?.balance} [${entityId}]`;
+      }
+
      return `${cusEnt.feature_id} - ${balanceStr} (${
        cusEnt.customer_product ? cusEnt.customer_product.product_id : ""
      })`;
@@ -171,7 +168,7 @@ const logUsageUpdate = ({
 // Main function to update customer balance
 export const updateUsage = async ({
  sb,
-  customer,
+  customerId,
  features,
  org,
  env,
@@ -182,7 +179,7 @@ export const updateUsage = async ({
  entityId,
 }: {
  sb: SupabaseClient;
-  customer: Customer;
+  customerId: string;
  features: Feature[];
  org: Organization;
  env: AppEnv;
@@ -190,19 +187,24 @@ export const updateUsage = async ({
  properties: any;
  setUsage: boolean;
  logger: any;
-  entityId: string;
+  entityId?: string;
 }) => {
-  const startTime = performance.now();
+  const customer = await CusService.getWithProducts({
+    sb,
+    idOrInternalId: customerId,
+    orgId: org.id,
+    env,
+    inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue],
+    entityId,
+  });
+
  const { cusEnts, cusPrices } = await getCusEntsInFeatures({
    sb,
-    internalCustomerId: customer.internal_id,
+    customer,
    internalFeatureIds: features.map((f) => f.internal_id!),
-    inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue],
-    withPrices: true,
    logger,
    reverseOrder: org.config?.reverse_deduction_order,
  });
-  const endTime = performance.now();

  // 1. Get deductions for each feature
  const featureDeductions = getFeatureDeductions({
@@ -213,32 +215,21 @@ export const updateUsage = async ({
  });

  logUsageUpdate({
-    timeTaken: (endTime - startTime).toFixed(2),
    customer,
    features,
    cusEnts,
    featureDeductions,
-    properties,
    org,
    setUsage,
+    entityId,
  });

-  // // 2. Handle group_by initialization
-  // await initGroupBalancesForEvent({
-  //   sb,
-  //   features,
-  //   cusEnts,
-  //   properties,
-  // });
-
  // 3. Return if no customer entitlements or features found
  if (cusEnts.length === 0 || features.length === 0) {
    console.log("   - No customer entitlements or features found");
    return;
  }

-  // 4. Perform deductions and update customer balance
-
  for (const obj of featureDeductions) {
    let { feature, deduction: toDeduct } = obj;

@@ -304,7 +295,8 @@ export const runUpdateUsageTask = async ({
  try {
    // 1. Update customer balance
    const {
-      customer,
+      internalCustomerId,
+      customerId,
      features,
      value,
      set_usage,
@@ -316,12 +308,12 @@ export const runUpdateUsageTask = async ({

    console.log("--------------------------------");
    console.log(
-      `HANDLING USAGE TASK FOR CUSTOMER (${customer.id}), ORG: ${org.slug}`
+      `HANDLING USAGE TASK FOR CUSTOMER (${customerId}), ORG: ${org.slug}`
    );

    const cusEnts: any = await updateUsage({
      sb,
-      customer,
+      customerId,
      features,
      value,
      properties,
@@ -336,26 +328,6 @@ export const runUpdateUsageTask = async ({
      return;
    }
    console.log("   ✅ Customer balance updated");
-
-    // // 2. Check if there's below threshold price
-    // const belowThresholdPrice = await getBelowThresholdPrice({
-    //   sb,
-    //   internalCustomerId: customer.internal_id,
-    //   cusEnts,
-    // });
-
-    // if (belowThresholdPrice) {
-    //   console.log("2. Below threshold price found");
-
-    //   await handleBelowThresholdInvoicing({
-    //     sb,
-    //     internalCustomerId: customer.internal_id,
-    //     belowThresholdPrice,
-    //     logger,
-    //   });
-    // } else {
-    //   console.log("   ✅ No below threshold price found");
-    // }
  } catch (error) {
    if (logger) {
      logger.use((log: any) => {
--- a/server/test.sh
+++ b/server/test.sh
@@ -15,8 +15,8 @@ if [ "$1" == "basic-parallel" ]; then
 elif [ "$1" == "advanced-parallel" ]; then
    MOCHA_PARALLEL=true  \
    $MOCHA_SETUP \
-    && $MOCHA_CMD 'tests/advanced/usage/*.ts' \
    && $MOCHA_CMD 'tests/advanced/arrear_prorated/*.ts' 'tests/advanced/coupons/*.ts'\
+    # && $MOCHA_CMD 'tests/advanced/usage/*.ts' \
    # && $MOCHA_CMD 'tests/advanced/coupons/*.ts'\


--- a/server/tests/advanced/arrear_prorated/arrear_prorated1.ts
+++ b/server/tests/advanced/arrear_prorated/arrear_prorated1.ts
@@ -1,254 +0,0 @@
-// import { assert, expect } from "chai";
-// import { AutumnCli } from "tests/cli/AutumnCli.js";
-// import { advanceProducts, features } from "tests/global.js";
-// import { compareMainProduct } from "tests/utils/compare.js";
-// import {
-//   advanceClockForInvoice,
-//   advanceTestClock,
-// } from "../../utils/stripeUtils.js";
-// import { timeout } from "../../utils/genUtils.js";
-// import { createStripeCli } from "@/external/stripe/utils.js";
-// import { addDays } from "date-fns";
-// import chalk from "chalk";
-// import Stripe from "stripe";
-
-// import { initCustomerWithTestClock } from "tests/utils/testInitUtils.js";
-
-// const advanceArrearProratedCycle = async ({
-//   stripeSub,
-//   stripeCli,
-//   testClockId,
-//   customerId,
-//   billingUnits,
-//   startingFrom,
-//   startingBalance,
-// }: {
-//   stripeSub: Stripe.Subscription;
-//   stripeCli: Stripe;
-//   testClockId: string;
-//   customerId: string;
-//   billingUnits: number;
-//   startingFrom?: number;
-//   startingBalance?: number;
-// }) => {
-//   // 1. Get total period
-//   let totalPeriod =
-//     (stripeSub.current_period_end - stripeSub.current_period_start) * 1000;
-
-//   // 2. Get allowance
-//   let allowance =
-//     advanceProducts.proratedArrearSeats.entitlements.seats.allowance!;
-
-//   // 3. Get starting balance
-//   let balance = startingBalance || allowance;
-
-//   // 4. Get price per seat
-//   let pricePerSeat =
-//     advanceProducts.proratedArrearSeats.prices[1].config.usage_tiers[0].amount;
-
-//   let skipDays = 2;
-//   // 5. Get accrued price
-//   let accruedPrice = 0;
-//   if (startingBalance) {
-//     accruedPrice =
-//       (-startingBalance *
-//         pricePerSeat *
-//         (startingFrom! - stripeSub.current_period_start * 1000)) /
-//       totalPeriod;
-//     accruedPrice = Math.max(accruedPrice, 0);
-//     console.log("   🔍 Starting balance: ", startingBalance);
-//     console.log("   🔍 Starting price: ", accruedPrice);
-//   }
-
-//   let curTime = startingFrom || stripeSub.current_period_start * 1000;
-//   let numberOfEvents = 2;
-
-//   console.group();
-//   console.group();
-//   for (let i = 0; i < numberOfEvents; i++) {
-//     let sign = balance > 0 ? 1 : Math.random() > 0.5 ? 1 : -1;
-
-//     let currentUsage = allowance - balance;
-
-//     let nextBoundary =
-//       Math.ceil((currentUsage + 1) / billingUnits) * billingUnits;
-//     let prevBoundary = nextBoundary - billingUnits;
-
-//     let valueNeeded = 0;
-//     if (sign > 0) {
-//       // Add random amount to push above next boundary
-//       const valueToGetToNegative = balance + 1;
-//       valueNeeded =
-//         Math.floor(Math.random() * 10) + (nextBoundary - currentUsage + 1);
-
-//       valueNeeded = Math.max(valueNeeded, valueToGetToNegative);
-//     } else {
-//       valueNeeded = -(
-//         Math.floor(Math.random() * 10) +
-//         (currentUsage - prevBoundary + 1)
-//       );
-//     }
-
-//     let firstHalf = Math.floor(valueNeeded / 2);
-//     let secondHalf = valueNeeded - firstHalf;
-
-//     // Test event sending in quick succession
-//     await AutumnCli.sendEvent({
-//       customerId,
-//       eventName: features.seats.id,
-//       properties: {
-//         value: firstHalf,
-//       },
-//     });
-
-//     await AutumnCli.sendEvent({
-//       customerId,
-//       eventName: features.seats.id,
-//       properties: {
-//         value: secondHalf,
-//       },
-//     });
-
-//     await timeout(2000);
-
-//     let prevBalance = balance;
-//     balance -= valueNeeded;
-
-//     // Calculate prorated price only when crossing boundary
-//     let newPrice = Math.max(0, -balance * pricePerSeat);
-//     let prevCurTime = curTime;
-//     curTime = addDays(curTime, 2).getTime();
-
-//     if (i === numberOfEvents - 1) {
-//       curTime = stripeSub.current_period_end * 1000;
-//     }
-
-//     let proratedPrice = (newPrice * (curTime - prevCurTime)) / totalPeriod;
-//     accruedPrice += Number(proratedPrice.toFixed(2));
-
-//     console.log(`Event ${i + 1}:`);
-//     console.log(`  - Value added: ${valueNeeded}`);
-//     console.log(`  - Balance: ${prevBalance} -> ${balance}`);
-//     console.log(`  - Prorated price: ${proratedPrice.toFixed(2)}`);
-//     console.log(`  - Accrued price: ${accruedPrice.toFixed(2)}`);
-
-//     await advanceTestClock({
-//       stripeCli,
-//       testClockId,
-//       numberOfDays: 2,
-//       startingFrom: new Date(prevCurTime),
-//     });
-//   }
-
-//   console.groupEnd();
-//   console.groupEnd();
-
-//   // Advance test clock to end of period
-//   // let advanceTo = addDays(addMonths(new Date(), 1), 2);
-
-//   let advanceToStart = startingFrom ? new Date(startingFrom) : new Date();
-//   let advanceTo = await advanceClockForInvoice({
-//     stripeCli,
-//     testClockId,
-//     waitForMeterUpdate: false,
-//     // numberOfDays: 2,
-//     startingFrom: advanceToStart,
-//   });
-
-//   // Check invoice amount
-//   const res = await AutumnCli.getCustomer(customerId);
-//   let invoice = res.invoices[0];
-
-//   let basePrice = advanceProducts.proratedArrearSeats.prices[0].config.amount;
-//   let nextMonthUsagePrice = Math.max(-balance * pricePerSeat, 0);
-//   console.log("   🔍 Next month usage price: ", nextMonthUsagePrice);
-
-//   let expectedTotal = Number(
-//     (accruedPrice + basePrice + nextMonthUsagePrice).toFixed(2)
-//   );
-//   expect(invoice.total).to.be.greaterThan(expectedTotal - 0.01);
-//   expect(invoice.total).to.be.lessThan(expectedTotal + 0.01);
-
-//   return {
-//     balance,
-//     advancedTo: advanceTo,
-//   };
-// };
-
-// describe(`${chalk.yellowBright(
-//   "Testing in_arrear_prorated -- update via events"
-// )}`, () => {
-//   const customerId = "arrear-prorated-events";
-
-//   let testClockId = "";
-//   let stripeCli: Stripe;
-//   let subId = "";
-//   let stripeSub: Stripe.Subscription;
-//   let billingUnits =
-//     advanceProducts.proratedArrearSeats.prices[1].config.billing_units || 1;
-
-//   before(async function () {
-//     const { testClockId: createdTestClockId } = await initCustomerWithTestClock(
-//       {
-//         customerId,
-//         org: this.org,
-//         env: this.env,
-//         sb: this.sb,
-//       }
-//     );
-
-//     stripeCli = createStripeCli({
-//       org: this.org,
-//       env: this.env,
-//     });
-
-//     testClockId = createdTestClockId;
-//   });
-
-//   it("should attach in arrear prorated seats", async () => {
-//     await AutumnCli.attach({
-//       customerId,
-//       productId: advanceProducts.proratedArrearSeats.id,
-//     });
-//   });
-
-//   it("should have correct product", async () => {
-//     const res = await AutumnCli.getCustomer(customerId);
-//     compareMainProduct({
-//       sent: advanceProducts.proratedArrearSeats,
-//       cusRes: res,
-//     });
-
-//     // 2. Get subscription period start and period end
-//     subId = res.products[0].subscription_ids[0];
-//     stripeSub = await stripeCli.subscriptions.retrieve(subId);
-//   });
-
-//   it("should run two cycles and have correct invoice / balance", async () => {
-//     // Do it again
-//     let { advancedTo, balance } = await advanceArrearProratedCycle({
-//       stripeSub,
-//       stripeCli,
-//       testClockId,
-//       customerId,
-//       billingUnits,
-//     });
-
-//     // console.log(`   Advanced to ${format(new Date(advancedTo), "yyyy-MM-dd")}`);
-
-//     // let newStripeSub = await stripeCli.subscriptions.retrieve(subId);
-//     // await advanceArrearProratedCycle({
-//     //   stripeSub: newStripeSub,
-//     //   stripeCli,
-//     //   testClockId,
-//     //   customerId,
-//     //   billingUnits,
-//     //   startingFrom: advancedTo,
-//     //   startingBalance: balance,
-//     // });
-//   });
-
-//   // TODO: Test reset at for in arrear prorated with Ent Interval = Lifetime
-//   // TODO: Test in arrear prorated for entitlements with billing units > 1
-//   return;
-// });
--- a/server/tests/advanced/arrear_prorated/arrear_prorated2.ts
+++ b/server/tests/advanced/arrear_prorated/arrear_prorated2.ts
@@ -11,6 +11,9 @@ import Stripe from "stripe";

 import { initCustomerWithTestClock } from "tests/utils/testInitUtils.js";
 import { checkSubscriptionContainsProducts } from "tests/utils/scheduleCheckUtils.js";
+import { CacheManager } from "@/external/caching/CacheManager.js";
+import { CacheType } from "@/external/caching/cacheActions.js";
+import { hashApiKey } from "@/internal/dev/api-keys/apiKeyUtils.js";

 const advanceAPThroughBalances = async ({
  stripeSub,
@@ -169,7 +172,7 @@ const advanceAPThroughBalances = async ({
 };

 describe(`${chalk.yellowBright(
-  "Testing update in arrear prorated through /balances"
+  "arrear_prorated2: testing update in arrear prorated through /balances"
 )}`, () => {
  const customerId = "arrear-prorated-balances";

@@ -196,6 +199,23 @@ describe(`${chalk.yellowBright(
    });

    testClockId = createdTestClockId;
+
+    // Update org config
+    await this.sb
+      .from("organizations")
+      .update({
+        config: {
+          ...this.org.config,
+          bill_upgrade_immediately: false,
+        },
+      })
+      .eq("id", this.org.id);
+
+    await CacheManager.invalidate({
+      action: CacheType.SecretKey,
+      value: hashApiKey(process.env.UNIT_TEST_AUTUMN_SECRET_KEY!),
+    });
+    await CacheManager.disconnect();
  });

  it("should attach in arrear prorated seats", async () => {
@@ -258,6 +278,21 @@ describe(`${chalk.yellowBright(
    });
  });

+  after(async function () {
+    await this.sb
+      .from("organizations")
+      .update({
+        config: {
+          ...this.org.config,
+          bill_upgrade_immediately: true,
+        },
+      })
+      .eq("id", this.org.id);
+    void CacheManager.invalidate({
+      action: CacheType.SecretKey,
+      value: hashApiKey(process.env.UNIT_TEST_AUTUMN_SECRET_KEY!),
+    });
+  });
  // TODO: Test reset at for in arrear prorated with Ent Interval = Lifetime
  // TODO: Test in arrear prorated for entitlements with billing units > 1
 });
--- a/server/tests/basic/04_entitled.ts
+++ b/server/tests/basic/04_entitled.ts
@@ -20,7 +20,8 @@ const checkEntitledOnProduct = async ({
 }) => {
  // 1. Send events
  const allowance = totalAllowance || product.entitlements.metered1.allowance;
-  const randomNum = Math.floor(Math.random() * (allowance - 1));
+  // const randomNum = Math.floor(Math.random() * (allowance - 1));
+  const randomNum = 3;

  const batchUpdates = [];
  for (let i = 0; i < randomNum; i++) {
--- a/server/tests/basic/entities/entities2.ts
+++ b/server/tests/basic/entities/entities2.ts
@@ -0,0 +1,225 @@
+import { compareMainProduct } from "../../utils/compare.js";
+
+import { entityProducts, features } from "../../global.js";
+
+import { assert, expect } from "chai";
+import chalk from "chalk";
+import { Autumn } from "@/external/autumn/autumnCli.js";
+import { setupBefore } from "tests/before.js";
+import {
+  BillingInterval,
+  CusProductStatus,
+  EntInterval,
+  ErrCode,
+  ProductItemFeatureType,
+  UsageModel,
+} from "@autumn/shared";
+import { getFeaturePrice, getUsagePriceTiers } from "tests/utils/genUtils.js";
+
+import { Stripe } from "stripe";
+import { CusService } from "@/internal/customers/CusService.js";
+import { SupabaseClient } from "@supabase/supabase-js";
+import { checkBalance } from "tests/utils/autumnUtils.js";
+import { initCustomerWithTestClock } from "tests/utils/testInitUtils.js";
+import { advanceTestClock } from "tests/utils/stripeUtils.js";
+import { addHours, addMonths } from "date-fns";
+import { CacheManager } from "@/external/caching/CacheManager.js";
+import { CacheType } from "@/external/caching/cacheActions.js";
+import { hashApiKey } from "@/internal/dev/api-keys/apiKeyUtils.js";
+import {
+  constructFeatureItem,
+  constructFeaturePriceItem,
+} from "@/internal/products/product-items/productItemUtils.js";
+import { createProduct } from "tests/utils/productUtils.js";
+
+// Check balance and stripe quantity
+const checkEntAndStripeQuantity = async ({
+  sb,
+  autumn,
+  stripeCli,
+  featureId,
+  customerId,
+  expectedBalance,
+  expectedUsage,
+  expectedStripeQuantity,
+}: {
+  sb: SupabaseClient;
+  autumn: Autumn;
+  stripeCli: Stripe;
+  featureId: string;
+  customerId: string;
+  expectedBalance: number;
+  expectedUsage?: number;
+  expectedStripeQuantity: number;
+}) => {
+  let { customer, entitlements, products } = await autumn.customers.get(
+    customerId
+  );
+
+  let cusProducts = await CusService.getFullCusProducts({
+    sb,
+    internalCustomerId: customer.internal_id,
+    withPrices: true,
+    withProduct: true,
+    inStatuses: [CusProductStatus.Active],
+  });
+
+  let entitlement = entitlements.find((e: any) => e.feature_id == featureId);
+
+  expect(entitlement.balance).to.equal(expectedBalance);
+  if (expectedUsage) {
+    expect(entitlement.used).to.equal(
+      expectedUsage,
+      `Get customer ${customerId} returned incorrect "used" for feature ${featureId}`
+    );
+  }
+
+  if (products.length == 0) {
+    assert.fail(`Get customer ${customerId} returned no products`);
+  }
+
+  // 2. Get stripe quantity
+  let mainProduct = products[0];
+
+  if (mainProduct.subscription_ids.length == 0) {
+    assert.fail(`Get customer ${customerId} returned no subscriptions`);
+  }
+
+  let price = getFeaturePrice({
+    product: mainProduct,
+    featureId: featureId,
+    cusProducts,
+  });
+
+  if (!price) {
+    assert.fail(
+      `Get customer ${customerId} returned no price for feature ${featureId}`
+    );
+  }
+
+  let stripeSub = await stripeCli.subscriptions.retrieve(
+    mainProduct.subscription_ids[0]
+  );
+  let subItem = stripeSub.items.data.find(
+    (item: any) => item.price.id == price.config!.stripe_price_id
+  );
+
+  if (!subItem) {
+    assert.fail(
+      `Get customer ${customerId} returned no sub item for feature ${featureId}`
+    );
+  }
+
+  expect(subItem.quantity).to.equal(
+    expectedStripeQuantity,
+    `Get customer ${customerId} returned incorrect stripe quantity for feature ${featureId}`
+  );
+};
+
+// UNCOMMENT FROM HERE
+let entity2Pro = {
+  id: "entity2Pro",
+  name: "Entity 2 Pro",
+  items: {
+    seats: constructFeaturePriceItem({
+      feature_id: features.seats.id,
+      included_usage: 0,
+      price: 150,
+      interval: BillingInterval.Month,
+      usage_model: UsageModel.PayPerUse,
+    }),
+    metered1: constructFeaturePriceItem({
+      feature_id: features.metered1.id,
+      included_usage: 50_000,
+      billing_units: 50_000,
+      price: 10,
+      interval: BillingInterval.Month,
+      entity_feature_id: features.seats.id,
+    }),
+    metered2: constructFeatureItem({
+      feature_id: features.metered2.id,
+      included_usage: 4000,
+      interval: EntInterval.Month,
+      entity_feature_id: features.seats.id,
+    }),
+  },
+};
+
+describe(`${chalk.yellowBright(
+  "entities2: Testing entities with prorate_unused: true"
+)}`, () => {
+  let customerId = "entity2";
+  let autumn: Autumn;
+  let stripeCli: Stripe;
+  let testClockId: string;
+
+  before(async function () {
+    await setupBefore(this);
+    autumn = this.autumn;
+    stripeCli = this.stripeCli;
+
+    const { testClockId: testClockId1 } = await initCustomerWithTestClock({
+      customerId,
+      sb: this.sb,
+      org: this.org,
+      env: this.env,
+    });
+
+    await createProduct({
+      autumn: this.autumn,
+      product: entity2Pro,
+    });
+
+    testClockId = testClockId1;
+
+    // await this.sb
+    //   .from("organizations")
+    //   .update({
+    //     config: {
+    //       ...this.org.config,
+    //       prorate_unused: true,
+    //     },
+    //   })
+    //   .eq("id", this.org.id);
+
+    // await CacheManager.invalidate({
+    //   action: CacheType.SecretKey,
+    //   value: hashApiKey(process.env.UNIT_TEST_AUTUMN_SECRET_KEY!),
+    // });
+    // await CacheManager.disconnect();
+  });
+
+  it("should create entity, then attach pro product", async function () {
+    await autumn.entities.create(customerId, {
+      id: "1",
+      name: "seat_1",
+      featureId: features.seats.id,
+    });
+
+    await autumn.attach({
+      customerId,
+      productId: entity2Pro.id,
+    });
+
+    let { customer, invoices } = await autumn.customers.get(customerId);
+
+    expect(invoices.length).to.equal(1);
+    expect(invoices[0].total).to.equal(entity2Pro.items.seats.price);
+  });
+
+  after(async function () {
+    // await this.sb
+    //   .from("organizations")
+    //   .update({
+    //     config: {
+    //       ...this.org.config,
+    //       prorate_unused: true,
+    //     },
+    //   })
+    //   .eq("id", this.org.id);
+    // void CacheManager.invalidate({
+    //   action: CacheType.SecretKey,
+    //   value: hashApiKey(process.env.UNIT_TEST_AUTUMN_SECRET_KEY!),
+    // });
+  });
+});
--- a/server/tests/before.ts
+++ b/server/tests/before.ts
@@ -5,6 +5,7 @@ import { createSupabaseClient } from "@/external/supabaseUtils.js";
 import { AppEnv } from "@autumn/shared";
 import { OrgService } from "@/internal/orgs/OrgService.js";
 import { Autumn } from "@/external/autumn/autumnCli.js";
+import { Autumn as AutumnJS } from "autumn-js";
 import { createStripeCli } from "@/external/stripe/utils.js";

 const ORG_SLUG = "unit-test-org";
@@ -12,14 +13,21 @@ const DEFAULT_ENV = AppEnv.Sandbox;

 export const setupBefore = async (instance: any) => {
  const sb = createSupabaseClient();
-  const org = await OrgService.getBySlug({sb, slug: ORG_SLUG});
+  const org = await OrgService.getBySlug({ sb, slug: ORG_SLUG });
  const env = DEFAULT_ENV;
  const autumnSecretKey = process.env.UNIT_TEST_AUTUMN_SECRET_KEY!;
  const autumn = new Autumn(autumnSecretKey);
-  const stripeCli = createStripeCli({org, env});
+
+  const autumnJs = new AutumnJS({
+    secretKey: autumnSecretKey,
+    url: "http://localhost:8080/v1",
+  });
+
+  const stripeCli = createStripeCli({ org, env });
  instance.sb = sb;
  instance.org = org;
  instance.env = env;
  instance.autumn = autumn;
  instance.stripeCli = stripeCli;
-}
+  instance.autumnJs = autumnJs;
+};
--- a/server/tests/utils/productUtils.ts
+++ b/server/tests/utils/productUtils.ts
@@ -16,6 +16,7 @@ export const createProduct = async ({
  if (typeof clone.items === "object") {
    clone.items = Object.values(clone.items);
  }
+
  await autumn.products.create(clone);

  // await autumn.products.update(product.id, clone);
--- a/shared/errors/SuccessCode.ts
+++ b/shared/errors/SuccessCode.ts
@@ -20,4 +20,6 @@ export enum SuccessCode {
  NewProductAttached = "new_product_attached",
  OneOffProductAttached = "one_off_product_attached",
  ProductFound = "product_found",
+
+  RenewedProduct = "renewed_product",
 }
--- a/shared/models/cusModels/cusProductModels.ts
+++ b/shared/models/cusModels/cusProductModels.ts
@@ -45,10 +45,12 @@ export enum CusProductStatus {

 export const CusProductSchema = z.object({
  id: z.string(),
-  internal_customer_id: z.string(),
  internal_product_id: z.string(),
-  customer_id: z.string().nullish(),
  product_id: z.string(),
+  internal_customer_id: z.string(),
+  customer_id: z.string().nullish(),
+  internal_entity_id: z.string().nullish(),
+  entity_id: z.string().nullish(),
  created_at: z.number(),

  // Useful for event-driven subscriptions (and usage-based to check limits)
--- a/shared/models/cusModels/fullCusModel.ts
+++ b/shared/models/cusModels/fullCusModel.ts
@@ -1,6 +1,9 @@
 import { Customer } from "./cusModels.js";
 import { FullCusProduct } from "./cusProductModels.js";
+import { Entity } from "./entityModels/entityModels.js";

 export type FullCustomer = Customer & {
  customer_products: FullCusProduct[];
+  entities: Entity[];
+  entity: Entity;
 };
--- a/shared/models/orgModels/orgConfigModels.ts
+++ b/shared/models/orgModels/orgConfigModels.ts
@@ -13,6 +13,7 @@ export const OrgConfigSchema = z.object({

  include_past_due: z.boolean().default(false),
  sync_status: z.boolean().default(true),
+  merge_billing_cycles: z.boolean().default(true),
 });

 export type OrgConfig = z.infer<typeof OrgConfigSchema>;
--- a/vite/src/views/customers/customer/CustomerProductList.tsx
+++ b/vite/src/views/customers/customer/CustomerProductList.tsx
@@ -123,6 +123,10 @@ export const CustomerProductList = ({
                    key: "Stripe Subscription ID (1)",
                    value: cusProduct.subscription_ids?.[0] || "N/A",
                  },
+                  {
+                    key: "Entity ID",
+                    value: cusProduct.entity_id || "N/A",
+                  },
                ]}
              >
                <div className="flex items-center gap-2">
--- a/vite/src/views/customers/customer/CustomerView.tsx
+++ b/vite/src/views/customers/customer/CustomerView.tsx
@@ -21,6 +21,13 @@ import ErrorScreen from "@/views/general/ErrorScreen";
 import { InvoicesTable } from "./InvoicesTable";
 import { CustomerDetails } from "./CustomerDetails";
 import { AdminHover } from "@/components/general/AdminHover";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";

 export default function CustomerView({ env }: { env: AppEnv }) {
  const { customer_id } = useParams();
@@ -139,7 +146,18 @@ export default function CustomerView({ env }: { env: AppEnv }) {
                )}
              </h2>
            </div>
-            {/* <CustomerToolbar customer={customer} /> */}
+            <Select>
+              <SelectTrigger>
+                <SelectValue placeholder="Select entity" />
+              </SelectTrigger>
+              <SelectContent>
+                {entities.map((entity: any) => (
+                  <SelectItem key={entity.id} value={entity.id}>
+                    {entity.name}
+                  </SelectItem>
+                ))}
+              </SelectContent>
+            </Select>
          </div>
          <div className="flex w-full !pb-[50px]">
            {/* main content */}
--- a/vite/src/views/customers/customer/entitlements/CustomerEntitlementsList.tsx
+++ b/vite/src/views/customers/customer/entitlements/CustomerEntitlementsList.tsx
@@ -61,6 +61,16 @@ export const CustomerEntitlementsList = () => {
    return product?.name;
  };

+  const getEntity = (cusEnt: FullCustomerEntitlement) => {
+    const cusProduct = customer.products.find(
+      (p: any) => p.id === cusEnt.customer_product_id
+    );
+    const entity = entities.find(
+      (e: any) => e.internal_id === cusProduct?.internal_entity_id
+    );
+    return entity;
+  };
+
  const sortedEntitlements = filteredEntitlements;

  const handleSelectCusEntitlement = (cusEnt: FullCustomerEntitlement) => {
@@ -178,6 +188,7 @@ export const CustomerEntitlementsList = () => {
        (cusEnt: FullCustomerEntitlement & { unused: number }) => {
          const entitlement = cusEnt.entitlement;
          const allowanceType = entitlement.allowance_type;
+          // const entity = getEntity(cusEnt);
          return (
            <Row
              key={cusEnt.id}
--- a/vite/src/views/products/rewards/RewardConfig.tsx
+++ b/vite/src/views/products/rewards/RewardConfig.tsx
@@ -5,13 +5,7 @@ import { useEffect, useState } from "react";
 import { Select, SelectContent, SelectItem } from "@/components/ui/select";
 import { SelectTrigger, SelectValue } from "@/components/ui/select";
 import { keyToTitle, slugify } from "@/utils/formatUtils/formatTextUtils";
-import {
-  Reward,
-  CouponDurationType,
-  RewardType,
-  Product,
-  FullProduct,
-} from "@autumn/shared";
+import { Reward, RewardType, Product, FullProduct } from "@autumn/shared";
 import { useProductsContext } from "../ProductsContext";
 import { DiscountConfig } from "./DiscountConfig";
 import { notNullish } from "@/utils/genUtils";