added below threshold invoicing using trigger

2025-01-23 10:15:33 +00:00
parent c478a950a7
commit 77deaedf1b
11 changed files with 1936 additions and 434 deletions
--- a/package-lock.json
+++ b/package-lock.json
--- a/server/.gitignore
+++ b/server/.gitignore
@@ -0,0 +1 @@
+.trigger
--- a/server/package.json
+++ b/server/package.json
@@ -15,13 +15,14 @@
    "prod:start": "node dist/index.js",
    "cron": "tsx src/cron.ts",
    "clean": "tsx src/clean.ts",
-    "invoiceCron": "tsx src/invoiceCron.ts"
+    "invoiceCron": "tsx src/invoiceCron.ts",
+    "trigger": "npx trigger.dev@latest dev"
  },
  "author": "",
  "license": "ISC",
  "dependencies": {
-    "@autumn/shared": "*",
    "@anthropic-ai/sdk": "^0.32.1",
+    "@autumn/shared": "*",
    "@aws-sdk/client-secrets-manager": "^3.726.0",
    "@clerk/express": "^1.3.22",
    "@clickhouse/client": "^1.10.0",
@@ -29,6 +30,7 @@
    "@slack/events-api": "^3.0.1",
    "@slack/web-api": "^7.8.0",
    "@supabase/supabase-js": "^2.46.2",
+    "@trigger.dev/sdk": "^3.3.11",
    "@types/cors": "^2.8.17",
    "@types/express": "^5.0.0",
    "@unkey/api": "^0.29.0",
@@ -53,6 +55,7 @@
    "tsx": "^4.19.2"
  },
  "devDependencies": {
+    "@trigger.dev/build": "^3.3.11",
    "@types/pg": "^8.11.10"
  }
 }
--- a/server/src/external/stripe/webhookHandlers/handleInvoicePaid.ts
+++ b/server/src/external/stripe/webhookHandlers/handleInvoicePaid.ts
@@ -44,13 +44,11 @@ export const handleInvoicePaid = async ({
    return;
  }

-  console.log("Threshold invoice:", invoice.id);
-  const cusProduct = await CusProductService.getPastDueByInvoiceId({
-    sb,
-    invoiceId: invoice.id,
-  });
+  // console.log("Threshold invoice:", invoice.id);
+  // const cusProduct = await CusProductService.getPastDueByInvoiceId({
+  //   sb,
+  //   invoiceId: invoice.id,
+  // });

-  // If invoice.paid
-  // 1
-  console.log("Customer product:", cusProduct);
+  // console.log("Customer product:", cusProduct);
 };
--- a/server/src/internal/api/events/eventRouter.ts
+++ b/server/src/internal/api/events/eventRouter.ts
@@ -1,5 +1,6 @@
 import { Router } from "express";
 import {
+  AppEnv,
  Customer,
  ErrCode,
  Event,
@@ -17,11 +18,9 @@ import { EventService } from "./EventService.js";
 import { StatusCodes } from "http-status-codes";
 import { CustomerEntitlementService } from "../../customers/entitlements/CusEntitlementService.js";
 import { CusService } from "@/internal/customers/CusService.js";
-
-import {
-  getBelowThresholdPrice,
-  handleBelowThresholdInvoicing,
-} from "./invoiceThresholdUtils.js";
+import { getBelowThresholdPrice } from "../../../trigger/invoiceThresholdUtils.js";
+import { updateBalanceTask } from "@/trigger/updateBalanceTask.js";
+import { Client } from "pg";

 export const eventsRouter = Router();

@@ -132,6 +131,36 @@ const getFeaturesAndCustomerEnts = async ({
  return { customerEntitlements: cusEnts, features: rows };
 };

+const getAffectedFeatures = async ({
+  pg,
+  event,
+  orgId,
+  env,
+}: {
+  pg: Client;
+  event: Event;
+  orgId: string;
+  env: AppEnv;
+}) => {
+  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 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;
+};
+
 eventsRouter.post("", async (req: any, res: any) => {
  const body = req.body;
  const orgId = req.orgId;
@@ -140,81 +169,32 @@ eventsRouter.post("", async (req: any, res: any) => {
  try {
    const { customer, event } = await getEventAndCustomer(req);

-    const { customerEntitlements, features } = await getFeaturesAndCustomerEnts(
-      { req, customer, event }
-    );
+    const affectedFeatures = await getAffectedFeatures({
+      pg: req.pg,
+      event,
+      orgId,
+      env,
+    });

-    if (features.length === 0 || customerEntitlements.length === 0) {
-      res.status(200).json({ success: true, event_id: event.id });
-      return;
-    }
-
-    const featureIdToDeduction: any = {};
-    const meteredFeatures = features.filter(
-      (feature) => feature.type === FeatureType.Metered
-    );
-
-    for (const cusEnt of customerEntitlements) {
-      const internalFeatureId = cusEnt.internal_feature_id;
-      if (featureIdToDeduction[internalFeatureId]) {
-        continue;
-      }
-
-      const feature = features.find(
-        (feature) => feature.internal_id === internalFeatureId
-      );
-
-      if (feature?.type === FeatureType.Metered) {
-        featureIdToDeduction[internalFeatureId] = {
-          cusEntId: cusEnt.id,
-          deduction: 1,
-        };
-      }
-
-      if (feature?.type === FeatureType.CreditSystem) {
-        const deduction = getCreditSystemDeduction(meteredFeatures, feature);
-        if (deduction) {
-          featureIdToDeduction[internalFeatureId] = {
-            cusEntId: cusEnt.id,
-            deduction: deduction,
-          };
+    if (affectedFeatures.length > 0) {
+      console.log("Queued update balance task...");
+      await updateBalanceTask.trigger(
+        {
+          customer,
+          features: affectedFeatures,
+        },
+        {
+          queue: {
+            name: "customer",
+            concurrencyLimit: 1,
+          },
+          concurrencyKey: customer.internal_id,
        }
-      }
+      );
    }

-    const updateQuery = `UPDATE customer_entitlements SET balance = balance - CASE
-  ${Object.entries(featureIdToDeduction)
-    .map(
-      ([featureId, deduction]: [string, any]) =>
-        `WHEN id = '${deduction.cusEntId}' THEN ${deduction.deduction}`
-    )
-    .join("\n")}
-  END
-  WHERE id IN (${Object.values(featureIdToDeduction)
-    .map((deduction: any) => `'${deduction.cusEntId}'`)
-    .join(",")});`;
-
-    console.log("Successfully updated customer entitlements");
-
-    // UPDATE CUSTOMER ENTITLEMENTS
-    // const belowThresholdPrice = await getBelowThresholdPrice({
-    //   sb: req.sb,
-    //   internalCustomerId: customer.internal_id,
-    //   cusEnts: customerEntitlements,
-    // });
-
-    // if (belowThresholdPrice) {
-    //   console.log("Below threshold price found, queuing check...");
-    //   await handleBelowThresholdInvoicing({
-    //     sb: req.sb,
-    //     internalCustomerId: customer.internal_id,
-    //     belowThresholdPrice,
-    //   });
-    // }
-
-    await req.pg.query(updateQuery);
-
    res.status(200).json({ success: true, event_id: event.id });
+    return;
  } catch (error) {
    handleRequestError({ res, error, action: "POST event failed" });
    return;
--- a/server/src/internal/customers/products/CusProductService.ts
+++ b/server/src/internal/customers/products/CusProductService.ts
@@ -171,7 +171,7 @@ export class CusProductService {
    const { data, error } = await sb
      .from("customer_products")
      .select("*")
-      .eq("processor->>invoice_id", invoiceId)
+      .eq("processor->>last_invoice_id", invoiceId)
      .eq("status", CusProductStatus.PastDue)
      .single();

--- a/server/src/internal/api/events/invoiceThresholdUtils.ts
+++ b/server/src/internal/api/events/invoiceThresholdUtils.ts
@@ -1,43 +1,31 @@
 import Stripe from "stripe";
-import { getCusPaymentMethod } from "../../../external/stripe/stripeCusUtils.js";
-import { createStripeCli } from "../../../external/stripe/utils.js";
-import { createSupabaseClient } from "../../../external/supabaseUtils.js";
-import { CusService } from "../../customers/CusService.js";
-import { getFeatureBalance } from "../../customers/entitlements/cusEntUtils.js";
-import { OrgService } from "../../orgs/OrgService.js";
-import { ProductService } from "../../products/ProductService.js";
-import { createPgClient } from "../../../middleware/envMiddleware.js";
+import { getCusPaymentMethod } from "@/external/stripe/stripeCusUtils.js";
+import { createStripeCli } from "@/external/stripe/utils.js";
+import { OrgService } from "@/internal/orgs/OrgService.js";
 import {
  AppEnv,
-  BillingType,
-  CusProduct,
  Customer,
-  Entitlement,
-  EntitlementWithFeatureSchema,
  Organization,
-  Price,
  UsagePriceConfig,
  FullCustomerPrice,
  FullCusProduct,
  CusProductStatus,
-  EntitlementWithFeature,
  CustomerEntitlement,
-  FullCustomerEntitlement,
 } from "@autumn/shared";

 import dotenv from "dotenv";
 import { SupabaseClient } from "@supabase/supabase-js";
-import { Client } from "pg";
-import { InvoiceService } from "../../customers/invoices/InvoiceService.js";
-import { Invoice } from "@autumn/shared";
-import { generateId } from "../../../utils/genUtils.js";
-import { CusProductService } from "../../customers/products/CusProductService.js";
-import { getEntOptions, getPriceEntitlement } from "../../prices/priceUtils.js";
-import { CustomerEntitlementService } from "../../customers/entitlements/CusEntitlementService.js";
-import chalk from "chalk";
+import { InvoiceService } from "@/internal/customers/invoices/InvoiceService.js";
+import { CusProductService } from "@/internal/customers/products/CusProductService.js";
+import {
+  getEntOptions,
+  getPriceEntitlement,
+} from "@/internal/prices/priceUtils.js";
+import { CustomerEntitlementService } from "@/internal/customers/entitlements/CusEntitlementService.js";

 dotenv.config();

+// FUNCTION 3: INVOICE CUSTOMER FOR BELOW THRESHOLD PRICE
 const createBelowThresholdInvoice = async ({
  stripeCli,
  customer,
@@ -229,7 +217,9 @@ const invoiceCustomer = async ({
    "Current balance:",
    cusEnt.balance,
    "| Update amount:",
-    cusEnt.entitlement.allowance
+    cusEnt.entitlement.allowance,
+    "| New balance:",
+    newBalance
  );

  await CustomerEntitlementService.update({
@@ -287,7 +277,11 @@ const checkBalanceBelowThreshold = async ({
    internalFeatureId: priceEnt.feature.internal_id!,
  });

-  return options?.threshold && featureBalance < options?.threshold;
+  return {
+    threshold: options?.threshold,
+    balance: featureBalance,
+    below: options?.threshold && featureBalance < options?.threshold,
+  };
 };

 // FUNCTION 2: QUEUE CHECK BELOW THRESHOLD PRICE
@@ -313,19 +307,20 @@ export const handleBelowThresholdInvoicing = async ({
  }

  // 4. Check if feature balance is below threshold
-  const belowThreshold = await checkBalanceBelowThreshold({
+  const { threshold, balance, below } = await checkBalanceBelowThreshold({
    sb,
    fullCusProduct,
    belowThresholdPrice,
  });

-  console.log("Below threshold:", belowThreshold);
-  if (!belowThreshold) {
+  console.log(
+    `Feature balance: ${balance}, threshold: ${threshold}, below: ${below}`
+  );
+
+  if (!below) {
    return;
  }

-  console.log("Feature balance < threshold, creating invoice...");
-
  // 1. Invoice customer
  await invoiceCustomer({
    sb,
@@ -334,6 +329,7 @@ export const handleBelowThresholdInvoicing = async ({
  });
 };

+// NON QUEUE BASED
 // FUNCTION 1: CHECK IF THERE'S A BELOW THRESHOLD PRICE
 export const getBelowThresholdPrice = async ({
  sb,
--- a/server/src/trigger/updateBalanceTask.ts
+++ b/server/src/trigger/updateBalanceTask.ts
@@ -0,0 +1,46 @@
+import { task } from "@trigger.dev/sdk/v3";
+import { createSupabaseClient } from "@/external/supabaseUtils.js";
+import { updateCustomerBalance } from "./updateBalanceUtils.js";
+import { handleBelowThresholdInvoicing } from "./invoiceThresholdUtils.js";
+import { getBelowThresholdPrice } from "./invoiceThresholdUtils.js";
+
+export const updateBalanceTask = task({
+  id: "update-customer-balance",
+
+  maxDuration: 300, // Stop executing after 300 secs (5 mins) of compute
+
+  run: async (payload: any, { ctx }) => {
+    try {
+      const sb = createSupabaseClient();
+
+      // 1. Update customer balance
+      const { customer, features } = payload;
+
+      const cusEnts: any = await updateCustomerBalance({
+        sb,
+        customer,
+        features,
+      });
+
+      // 2. Check if there's below threshold price
+      const belowThresholdPrice = await getBelowThresholdPrice({
+        sb,
+        internalCustomerId: customer.internal_id,
+        cusEnts,
+      });
+
+      if (belowThresholdPrice) {
+        console.log("--------------------------------");
+        console.log("Below threshold price found");
+        await handleBelowThresholdInvoicing({
+          sb,
+          internalCustomerId: payload.internalCustomerId,
+          belowThresholdPrice,
+        });
+      }
+    } catch (error) {
+      console.log(`Error updating customer balance: ${error}`);
+      console.log(error);
+    }
+  },
+});
--- a/server/src/trigger/updateBalanceUtils.ts
+++ b/server/src/trigger/updateBalanceUtils.ts
@@ -0,0 +1,122 @@
+import { CustomerEntitlementService } from "@/internal/customers/entitlements/CusEntitlementService.js";
+import { Customer, FeatureType, Event, AppEnv } from "@autumn/shared";
+
+import { CustomerEntitlement, Feature } from "@autumn/shared";
+import { SupabaseClient } from "@supabase/supabase-js";
+
+const getCustomerEntitlements = async ({
+  sb,
+  internalCustomerId,
+  features,
+}: {
+  sb: SupabaseClient;
+  internalCustomerId: string;
+  features: Feature[];
+}) => {
+  const internalFeatureIds = features.map((feature) => feature.internal_id);
+  const cusEnts = await CustomerEntitlementService.getActiveInFeatureIds({
+    sb,
+    internalCustomerId,
+    internalFeatureIds: internalFeatureIds as string[],
+  });
+
+  cusEnts.sort((a, b) => {
+    if (a.balance <= 0) return 1;
+    if (b.balance <= 0) return -1;
+
+    return a.created_at - b.created_at;
+  });
+
+  return cusEnts;
+};
+
+const getCreditSystemDeduction = (
+  meteredFeatures: Feature[],
+  creditSystem: Feature
+) => {
+  let creditsUpdate = 0;
+  let meteredFeatureIds = meteredFeatures.map((feature) => feature.id);
+
+  for (const schema of creditSystem.config.schema) {
+    if (meteredFeatureIds.includes(schema.metered_feature_id)) {
+      creditsUpdate += (1 / schema.feature_amount) * schema.credit_amount;
+    }
+  }
+
+  return creditsUpdate;
+};
+
+export const updateCustomerBalance = async ({
+  sb,
+  customer,
+  features,
+}: {
+  sb: SupabaseClient;
+  customer: Customer;
+  features: Feature[];
+}) => {
+  const cusEnts = await getCustomerEntitlements({
+    sb,
+    internalCustomerId: customer.internal_id,
+    features,
+  });
+
+  if (cusEnts.length === 0 || features.length === 0) {
+    return;
+  }
+
+  // Update customer balance
+  const featureIdToDeduction: any = {};
+  const meteredFeatures = features.filter(
+    (feature) => feature.type === FeatureType.Metered
+  );
+
+  for (const cusEnt of cusEnts) {
+    const internalFeatureId = cusEnt.internal_feature_id;
+    if (featureIdToDeduction[internalFeatureId]) {
+      continue;
+    }
+
+    const feature = features.find(
+      (feature) => feature.internal_id === internalFeatureId
+    );
+
+    if (feature?.type === FeatureType.Metered) {
+      featureIdToDeduction[internalFeatureId] = {
+        cusEntId: cusEnt.id,
+        deduction: 1,
+        feature: feature,
+      };
+    }
+
+    if (feature?.type === FeatureType.CreditSystem) {
+      const deduction = getCreditSystemDeduction(meteredFeatures, feature);
+      if (deduction) {
+        featureIdToDeduction[internalFeatureId] = {
+          cusEntId: cusEnt.id,
+          deduction: deduction,
+          feature: feature,
+        };
+      }
+    }
+
+    let deduction = featureIdToDeduction[internalFeatureId]?.deduction;
+    let curBalance = cusEnt.balance!;
+
+    if (!curBalance) {
+      continue;
+    }
+
+    await sb
+      .from("customer_entitlements")
+      .update({ balance: curBalance - deduction })
+      .eq("id", cusEnt.id);
+  }
+
+  let featuresUpdated = Object.values(featureIdToDeduction).map(
+    (obj: any) => `(${obj.feature.id}: ${obj.deduction})`
+  );
+
+  console.log(`Deducted ${featuresUpdated}`);
+  return cusEnts;
+};
--- a/server/trigger.config.ts
+++ b/server/trigger.config.ts
@@ -0,0 +1,22 @@
+import { defineConfig } from "@trigger.dev/sdk/v3";
+
+export default defineConfig({
+  project: "proj_yqrybepbgrhzmnbaccat",
+  runtime: "node",
+  logLevel: "log",
+  // The max compute seconds a task is allowed to run. If the task run exceeds this duration, it will be stopped.
+  // You can override this on an individual task.
+  // See https://trigger.dev/docs/runs/max-duration
+  maxDuration: 3600,
+  retries: {
+    enabledInDev: true,
+    default: {
+      maxAttempts: 3,
+      minTimeoutInMs: 1000,
+      maxTimeoutInMs: 10000,
+      factor: 2,
+      randomize: true,
+    },
+  },
+  dirs: ["./src/trigger"],
+});
--- a/server/tsconfig.json
+++ b/server/tsconfig.json
@@ -21,7 +21,7 @@
      "@shared/*": ["../shared/*"]
    }
  },
-  "include": ["src"],
+  "include": ["src", "trigger.config.ts"],
  "references": [
    {
      "path": "../shared"