From beabfbb4a7dc5b7a872b5138420548b60e7c2b94 Mon Sep 17 00:00:00 2001 From: John Yeo Date: Mon, 10 Mar 2025 23:54:55 +0800 Subject: [PATCH] feat: added usage endpoint --- frontend/src/app/demo/[slug]/page.tsx | 12 + frontend/src/views/demo/CursorDemo.tsx | 33 +- frontend/src/views/demo/CustomerBalances.tsx | 3 +- frontend/src/views/demo/autumnBackend.tsx | 3 + server/src/external/stripe/stripeCusUtils.ts | 3 + .../src/external/stripe/stripePriceUtils.ts | 16 +- .../webhookHandlers/handleInvoiceCreated.ts | 18 +- server/src/internal/api/apiRouter.ts | 2 + .../handlers/handleUpdateBalances.ts | 1 - .../api/customers/products/attachRouter.ts | 3 + server/src/internal/api/events/eventRouter.ts | 7 + server/src/internal/api/events/usageRouter.ts | 188 ++++++++++ .../customers/entitlements/cusEntUtils.ts | 16 + server/src/internal/prices/priceUtils.ts | 8 +- server/src/internal/products/productUtils.ts | 4 +- server/src/internal/public/publicRouter.ts | 1 + server/src/queue/JobName.ts | 4 + server/src/queue/queue.ts | 8 +- server/src/trigger/updateBalanceTask.ts | 45 ++- server/src/trigger/updateUsageTask.ts | 354 ++++++++++++++++++ server/test.sh | 1 + shared/models/eventModels/eventModels.ts | 2 + 22 files changed, 691 insertions(+), 41 deletions(-) create mode 100644 server/src/internal/api/events/usageRouter.ts create mode 100644 server/src/queue/JobName.ts create mode 100644 server/src/trigger/updateUsageTask.ts diff --git a/frontend/src/app/demo/[slug]/page.tsx b/frontend/src/app/demo/[slug]/page.tsx index f67eefa1e..f1062eeda 100644 --- a/frontend/src/app/demo/[slug]/page.tsx +++ b/frontend/src/app/demo/[slug]/page.tsx @@ -1,4 +1,5 @@ import { createClient } from "@/utils/supabase/server"; +import CursorDemo from "@/views/demo/CursorDemo"; import DynamicDemoView from "@/views/demo/DynamicDemo"; import MintDemoView from "@/views/demo/MintDemo"; @@ -18,6 +19,17 @@ export default async function DemoPage({ params }) { return
Not Found...
; } + if (params.slug == "cursor") { + return ( + + ); + } + if (params.slug == "mint") { return ( - {/*
- {buttons.map((button, index) => ( +
+ + + {/* {buttons.map((button, index) => ( - ))} -
*/} + ))} */} +
{/* Dark search input component */}
diff --git a/frontend/src/views/demo/CustomerBalances.tsx b/frontend/src/views/demo/CustomerBalances.tsx index b48fe3bed..441fc4fae 100644 --- a/frontend/src/views/demo/CustomerBalances.tsx +++ b/frontend/src/views/demo/CustomerBalances.tsx @@ -29,7 +29,8 @@ function CustomerBalances({ customer }: { customer: any }) { {entitlement.unlimited === true ? "Unlimited" : entitlement.balance !== undefined - ? entitlement.balance + entitlement.used === 0 + ? // ? entitlement.balance + entitlement.used === 0 + entitlement.balance < 0 ? `Used: ${-entitlement.balance}` : `${entitlement.balance} / ${ entitlement.balance + entitlement.used diff --git a/frontend/src/views/demo/autumnBackend.tsx b/frontend/src/views/demo/autumnBackend.tsx index 18d7c442a..d4df98221 100644 --- a/frontend/src/views/demo/autumnBackend.tsx +++ b/frontend/src/views/demo/autumnBackend.tsx @@ -13,14 +13,17 @@ export const checkAccess = async ({ axiosInstance, customerId, featureId, + userId, }: { axiosInstance: AxiosInstance; customerId: string; featureId: string; + userId?: string; }) => { const { data } = await axiosInstance.post("/entitled", { customer_id: customerId, feature_id: featureId, + group: userId, }); return data; }; diff --git a/server/src/external/stripe/stripeCusUtils.ts b/server/src/external/stripe/stripeCusUtils.ts index b0cc6d6db..8c98d4782 100644 --- a/server/src/external/stripe/stripeCusUtils.ts +++ b/server/src/external/stripe/stripeCusUtils.ts @@ -12,11 +12,13 @@ export const createStripeCusIfNotExists = async ({ org, env, customer, + logger, }: { sb: SupabaseClient; org: Organization; env: AppEnv; customer: Customer; + logger: any; }) => { let createNew = false; const stripeCli = createStripeCli({ org, env }); @@ -31,6 +33,7 @@ export const createStripeCusIfNotExists = async ({ } if (createNew) { + logger.info(`Creating new stripe customer for ${customer.id}`); const stripeCustomer = await createStripeCustomer({ org, env, diff --git a/server/src/external/stripe/stripePriceUtils.ts b/server/src/external/stripe/stripePriceUtils.ts index ec8de909c..20586005f 100644 --- a/server/src/external/stripe/stripePriceUtils.ts +++ b/server/src/external/stripe/stripePriceUtils.ts @@ -440,6 +440,7 @@ export const createStripePriceIFNotExist = async ({ entitlements, product, org, + logger, }: { sb: SupabaseClient; stripeCli: Stripe; @@ -447,6 +448,7 @@ export const createStripePriceIFNotExist = async ({ entitlements: EntitlementWithFeature[]; product: Product; org: Organization; + logger: any; }) => { const billingType = getBillingType(price.config!); @@ -474,15 +476,15 @@ export const createStripePriceIFNotExist = async ({ } } } catch (error: any) { - console.log("Stripe price not found / inactive"); - console.log("Error:", error.message); + logger.info("Stripe price not found / inactive"); + logger.info("Error:", error.message); config.stripe_price_id = undefined; config.stripe_meter_id = undefined; } if (billingType == BillingType.FixedCycle) { if (!config.stripe_price_id) { - console.log("Creating stripe fixed cycle price"); + logger.info("Creating stripe fixed cycle price"); await createStripeFixedCyclePrice({ sb, stripeCli, @@ -504,7 +506,7 @@ export const createStripePriceIFNotExist = async ({ }); if (!productId) { - console.log( + logger.info( "Creating stripe product for in advance price, one off & tiered" ); await createStripeInAdvancePrice({ @@ -520,7 +522,7 @@ export const createStripePriceIFNotExist = async ({ // For the rest if (!isOneOffAndTiered && !config.stripe_price_id) { - console.log("Creating stripe price for in advance price"); + logger.info("Creating stripe price for in advance price"); await createStripeInAdvancePrice({ sb, stripeCli, @@ -532,7 +534,7 @@ export const createStripePriceIFNotExist = async ({ } } else if (billingType == BillingType.UsageInArrear) { if (!config.stripe_price_id) { - console.log("Creating stripe price for in arrear price"); + logger.info("Creating stripe price for in arrear price"); await createStripeInArrearPrice({ sb, stripeCli, @@ -544,7 +546,7 @@ export const createStripePriceIFNotExist = async ({ } } else if (billingType == BillingType.InArrearProrated) { if (!config.stripe_price_id) { - console.log("Creating stripe price for in arrear prorated price"); + logger.info("Creating stripe price for in arrear prorated price"); await createStripeInAdvancePrice({ sb, stripeCli, diff --git a/server/src/external/stripe/webhookHandlers/handleInvoiceCreated.ts b/server/src/external/stripe/webhookHandlers/handleInvoiceCreated.ts index 12dc87af2..d5121349f 100644 --- a/server/src/external/stripe/webhookHandlers/handleInvoiceCreated.ts +++ b/server/src/external/stripe/webhookHandlers/handleInvoiceCreated.ts @@ -31,7 +31,10 @@ import { createStripeInvoiceItem } from "@/internal/customers/invoices/invoiceIt import { getRelatedCusEnt } from "@/internal/customers/prices/cusPriceUtils.js"; import { getNextEntitlementReset } from "@/utils/timeUtils.js"; import { formatUnixToDateTime, generateId } from "@/utils/genUtils.js"; -import { getMinCusEntBalance } from "@/internal/customers/entitlements/cusEntUtils.js"; +import { + getMinCusEntBalance, + getTotalNegativeBalance, +} from "@/internal/customers/entitlements/cusEntUtils.js"; import { getResetBalancesUpdate } from "@/internal/customers/entitlements/groupByUtils.js"; import { createLogtailWithContext } from "@/external/logtail/logtailUtils.js"; @@ -178,10 +181,13 @@ const handleUsageInArrear = async ({ return; } - const usage = new Decimal(allowance).minus(minBalance).toNumber(); + const totalNegativeBalance = getTotalNegativeBalance(relatedCusEnt); + const finalBalance = Math.max(totalNegativeBalance, minBalance); + const totalQuantity = new Decimal(allowance).minus(finalBalance).toNumber(); const billingUnits = (price.config as UsagePriceConfig).billing_units || 1; - const roundedUsage = - Math.ceil(new Decimal(usage).div(billingUnits).toNumber()) * billingUnits; + const roundedQuantity = + Math.ceil(new Decimal(totalQuantity).div(billingUnits).toNumber()) * + billingUnits; const usageTimestamp = Math.round( subDays(new Date(invoice.created * 1000), 1).getTime() / 1000 @@ -191,7 +197,7 @@ const handleUsageInArrear = async ({ event_name: price.id!, payload: { stripe_customer_id: customer.processor.id, - value: roundedUsage.toString(), + value: roundedQuantity.toString(), }, timestamp: usageTimestamp, }); @@ -201,7 +207,7 @@ const handleUsageInArrear = async ({ `✅ Submitted meter event for customer ${customer.id}, feature: ${feature.id}` ); logger.info( - `Allowance: ${allowance}, Min Balance: ${minBalance}, Usage: ${usage}, Rounded: ${roundedUsage}` + `Allowance: ${allowance}, Min Balance: ${minBalance}, Quantity: ${totalQuantity}, Rounded: ${roundedQuantity}` ); let invoiceCreatedStr = formatUnixToDateTime(invoice.created * 1000); diff --git a/server/src/internal/api/apiRouter.ts b/server/src/internal/api/apiRouter.ts index 88b65cbb6..632914b5f 100644 --- a/server/src/internal/api/apiRouter.ts +++ b/server/src/internal/api/apiRouter.ts @@ -10,6 +10,7 @@ import { featureApiRouter } from "./features/featureApiRouter.js"; import { entitledRouter } from "./entitled/entitledRouter.js"; import { attachRouter } from "./customers/products/attachRouter.js"; import { pricingMiddleware } from "@/middleware/pricingMiddleware.js"; +import { usageRouter } from "./events/usageRouter.js"; const apiRouter = Router(); @@ -46,5 +47,6 @@ apiRouter.use("/entitlements", entitlementApiRouter); apiRouter.use("/events", eventsRouter); apiRouter.use("/prices", priceRouter); apiRouter.use("/entitled", entitledRouter); +apiRouter.use("/usage", usageRouter); export { apiRouter }; diff --git a/server/src/internal/api/customers/handlers/handleUpdateBalances.ts b/server/src/internal/api/customers/handlers/handleUpdateBalances.ts index d5e7c52ed..b07846370 100644 --- a/server/src/internal/api/customers/handlers/handleUpdateBalances.ts +++ b/server/src/internal/api/customers/handlers/handleUpdateBalances.ts @@ -262,7 +262,6 @@ export const handleUpdateBalances = async (req: any, res: any) => { await deductFromUsageBasedCusEnt({ toDeduct, cusEnts, - features, deductParams: { sb, feature, diff --git a/server/src/internal/api/customers/products/attachRouter.ts b/server/src/internal/api/customers/products/attachRouter.ts index 91f0f4154..993aecfb5 100644 --- a/server/src/internal/api/customers/products/attachRouter.ts +++ b/server/src/internal/api/customers/products/attachRouter.ts @@ -195,6 +195,7 @@ export const checkStripeConnections = async ({ org, env, customer, + logger: req.logger, }), ]; for (const product of products) { @@ -204,6 +205,7 @@ export const checkStripeConnections = async ({ org, env, product, + logger: req.logger, }) ); } @@ -220,6 +222,7 @@ export const checkStripeConnections = async ({ entitlements, product: product!, org, + logger: req.logger, }) ); } diff --git a/server/src/internal/api/events/eventRouter.ts b/server/src/internal/api/events/eventRouter.ts index 61dd2f181..e0c14c927 100644 --- a/server/src/internal/api/events/eventRouter.ts +++ b/server/src/internal/api/events/eventRouter.ts @@ -19,6 +19,7 @@ import { OrgService } from "@/internal/orgs/OrgService.js"; import { QueueManager } from "@/queue/QueueManager.js"; import { subDays } from "date-fns"; import { FeatureService } from "@/internal/features/FeatureService.js"; +import { handleUsageEvent } from "./usageRouter.js"; export const eventsRouter = Router(); @@ -132,6 +133,12 @@ export const handleEventSent = async ({ customer_data: any; event_data: any; }) => { + if (event_data.feature_id) { + return handleUsageEvent({ + req, + }); + } + const { sb, pg, orgId, env } = req; const org = await OrgService.getFullOrg({ diff --git a/server/src/internal/api/events/usageRouter.ts b/server/src/internal/api/events/usageRouter.ts new file mode 100644 index 000000000..17e8729f7 --- /dev/null +++ b/server/src/internal/api/events/usageRouter.ts @@ -0,0 +1,188 @@ +import { Router } from "express"; +import { Customer, ErrCode, Event } from "@autumn/shared"; +import RecaseError, { handleRequestError } from "@/utils/errorUtils.js"; +import { generateId, nullish } from "@/utils/genUtils.js"; + +import { EventService } from "./EventService.js"; +import { CusService } from "@/internal/customers/CusService.js"; +import { createNewCustomer } from "../customers/cusUtils.js"; +import { OrgService } from "@/internal/orgs/OrgService.js"; +import { FeatureService } from "@/internal/features/FeatureService.js"; +import { StatusCodes } from "http-status-codes"; +import { QueueManager } from "@/queue/QueueManager.js"; +import { z } from "zod"; +import { JobName } from "@/queue/JobName.js"; +export const eventsRouter = Router(); +export const usageRouter = Router(); + +const getCusFeatureAndOrg = async ({ + req, + customerId, + featureId, + customerData, +}: { + req: any; + customerId: string; + featureId: string; + customerData: any; +}) => { + // 1. Get customer + let [customer, featureRes, org] = await Promise.all([ + CusService.getById({ + sb: req.sb, + id: customerId, + orgId: req.orgId, + env: req.env, + }), + FeatureService.getWithCreditSystems({ + sb: req.sb, + featureId, + orgId: req.orgId, + env: req.env, + }), + OrgService.getFullOrg({ + sb: req.sb, + orgId: req.minOrg.id, + }), + ]); + + let { feature, creditSystems } = featureRes; + if (!feature) { + throw new RecaseError({ + message: `Feature ${featureId} not found`, + code: ErrCode.FeatureNotFound, + statusCode: StatusCodes.NOT_FOUND, + }); + } + + if (!customer) { + customer = await createNewCustomer({ + sb: req.sb, + orgId: req.orgId, + env: req.env, + customer: { + id: customerId, + name: customerData?.name, + email: customerData?.email, + fingerprint: customerData?.fingerprint, + }, + }); + } + + return { customer, org, feature, creditSystems }; +}; + +const createAndInsertEvent = async ({ + req, + customer, + featureId, + value, + set_usage, + properties, +}: { + req: any; + customer: Customer; + featureId: string; + value?: number; + set_usage?: boolean; + properties: any; +}) => { + const newEvent: Event = { + id: generateId("evt"), + org_id: req.orgId, + env: req.env, + internal_customer_id: customer.internal_id, + timestamp: Date.now(), + customer_id: customer.id, + event_name: featureId, + properties, + value, + set_usage: set_usage || false, + }; + + await EventService.insertEvent(req.sb, newEvent); +}; + +export const handleUsageEvent = async ({ + req, + setUsage = false, +}: { + req: any; + setUsage?: boolean; +}) => { + let { customer_id, customer_data, properties, feature_id, value } = req.body; + properties = properties || {}; + + const { customer, org, feature, creditSystems } = await getCusFeatureAndOrg({ + req, + customerId: customer_id, + featureId: feature_id, + customerData: customer_data, + }); + + await createAndInsertEvent({ + req, + customer, + featureId: feature_id, + value, + set_usage: setUsage, + properties, + }); + + const features = [feature, ...creditSystems]; + + const queue = await QueueManager.getQueue({ useBackup: false }); + + if (nullish(value) || isNaN(parseFloat(value))) { + value = 1; + } else { + value = parseFloat(value); + } + + const payload = { + customerId: customer.internal_id, + customer, + features, + org, + env: req.env, + properties, + value, + set_usage: setUsage, + }; + + 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; +}; + +usageRouter.post("", async (req: any, res: any) => { + try { + await handleUsageEvent({ req, setUsage: true }); + res.status(StatusCodes.OK).json({ success: true }); + } catch (error) { + return handleRequestError({ + req, + res, + error, + action: "handleUsageEvent", + }); + } +}); diff --git a/server/src/internal/customers/entitlements/cusEntUtils.ts b/server/src/internal/customers/entitlements/cusEntUtils.ts index 4d6772914..e12030a4c 100644 --- a/server/src/internal/customers/entitlements/cusEntUtils.ts +++ b/server/src/internal/customers/entitlements/cusEntUtils.ts @@ -29,6 +29,7 @@ import { } from "@/utils/genUtils.js"; import { getGroupbalanceFromParams } from "./groupByUtils.js"; +import { Decimal } from "decimal.js"; export const getBalanceForFeature = async ({ sb, @@ -584,6 +585,21 @@ export const getMinCusEntBalance = ({ return Math.min(...balances); }; +export const getTotalNegativeBalance = (cusEnt: FullCustomerEntitlement) => { + let starting = cusEnt.balance! < 0 ? cusEnt.balance! : 0; + if (!cusEnt.balances) { + return starting; + } + + starting = starting - cusEnt.balance!; + for (const group in cusEnt.balances) { + if (cusEnt.balances[group].balance < 0) { + starting = starting - cusEnt.balances[group].balance; + } + } + return starting; +}; + // GET EXISTING USAGE export const getExistingUsageFromCusProducts = ({ entitlement, diff --git a/server/src/internal/prices/priceUtils.ts b/server/src/internal/prices/priceUtils.ts index 65be89bca..2dff554fe 100644 --- a/server/src/internal/prices/priceUtils.ts +++ b/server/src/internal/prices/priceUtils.ts @@ -310,7 +310,13 @@ export const getPriceForOverage = (price: Price, overage: number) => { } let amount = 0; - let remainingUsage = overage; + let billingUnits = usageConfig.billing_units || 1; + let remainingUsage = new Decimal( + Math.ceil(new Decimal(overage).div(billingUnits).toNumber()) + ) + .mul(billingUnits) + .toNumber(); + for (let i = 0; i < usageConfig.usage_tiers.length; i++) { let tier = usageConfig.usage_tiers[i]; diff --git a/server/src/internal/products/productUtils.ts b/server/src/internal/products/productUtils.ts index 2a0e77a81..bf6f20fad 100644 --- a/server/src/internal/products/productUtils.ts +++ b/server/src/internal/products/productUtils.ts @@ -163,11 +163,13 @@ export const checkStripeProductExists = async ({ org, env, product, + logger, }: { sb: SupabaseClient; org: Organization; env: AppEnv; product: FullProduct; + logger: any; }) => { let createNew = false; let stripeCli = createStripeCli({ @@ -191,7 +193,7 @@ export const checkStripeProductExists = async ({ } if (createNew) { - console.log("Creating new product in Stripe"); + logger.info(`Creating new product in Stripe for ${product.name}`); const stripeProduct = await stripeCli.products.create({ name: product.name, }); diff --git a/server/src/internal/public/publicRouter.ts b/server/src/internal/public/publicRouter.ts index 603bb9333..3d9276565 100644 --- a/server/src/internal/public/publicRouter.ts +++ b/server/src/internal/public/publicRouter.ts @@ -106,6 +106,7 @@ publicRouter.get("/customers/:customer_id", async (req: any, res: any) => { sb: req.sb, orgId: req.org.id, env: req.env, + params: req.query, }); res.status(200).json({ diff --git a/server/src/queue/JobName.ts b/server/src/queue/JobName.ts new file mode 100644 index 000000000..5ac42ab88 --- /dev/null +++ b/server/src/queue/JobName.ts @@ -0,0 +1,4 @@ +export enum JobName { + UpdateBalance = "update-balance", + UpdateUsage = "update-usage", +} diff --git a/server/src/queue/queue.ts b/server/src/queue/queue.ts index 8dd3ceb88..7b1a13fa4 100644 --- a/server/src/queue/queue.ts +++ b/server/src/queue/queue.ts @@ -2,6 +2,8 @@ import { Job, Queue, Worker } from "bullmq"; import { runUpdateBalanceTask } from "@/trigger/updateBalanceTask.js"; import { QueueManager } from "./QueueManager.js"; import { createLogtail } from "@/external/logtail/logtailUtils.js"; +import { runUpdateUsageTask } from "@/trigger/updateUsageTask.js"; +import { JobName } from "./JobName.js"; const NUM_WORKERS = 5; @@ -79,7 +81,11 @@ const initWorker = ({ } try { - await runUpdateBalanceTask({ payload: job.data, logger: logtail }); + if (job.name === JobName.UpdateBalance) { + await runUpdateBalanceTask({ payload: job.data, logger: logtail }); + } else if (job.name === JobName.UpdateUsage) { + await runUpdateUsageTask({ payload: job.data, logger: logtail }); + } } catch (error) { console.error("Error updating balance:", error); } finally { diff --git a/server/src/trigger/updateBalanceTask.ts b/server/src/trigger/updateBalanceTask.ts index 52177ccd3..0a42bec54 100644 --- a/server/src/trigger/updateBalanceTask.ts +++ b/server/src/trigger/updateBalanceTask.ts @@ -29,11 +29,11 @@ import { initGroupBalancesForEvent, } from "@/internal/customers/entitlements/groupByUtils.js"; import { notNullish, nullish, nullOrUndefined } from "@/utils/genUtils.js"; -import { getMinCusEntBalance } from "@/internal/customers/entitlements/cusEntUtils.js"; import { creditSystemContainsFeature, featureToCreditSystem, } from "@/internal/features/creditSystemUtils.js"; +import { getTotalNegativeBalance } from "@/internal/customers/entitlements/cusEntUtils.js"; // Decimal.set({ precision: 12 }); // 12 DP precision @@ -75,7 +75,9 @@ const getFeatureDeductions = ({ // Check if unlimited exists let unlimitedExists = cusEnts.some( - (cusEnt) => cusEnt.entitlement.allowance_type === AllowanceType.Unlimited + (cusEnt) => + cusEnt.entitlement.allowance_type === AllowanceType.Unlimited && + cusEnt.entitlement.internal_feature_id == feature.internal_id ); if (unlimitedExists || !deduction) { @@ -109,13 +111,13 @@ const getFeatureDeductions = ({ return featureDeductions; }; -const logBalanceUpdate = ({ +export const logBalanceUpdate = ({ timeTaken, customer, features, cusEnts, featureDeductions, - event, + properties, org, }: { timeTaken: string; @@ -123,7 +125,7 @@ const logBalanceUpdate = ({ features: Feature[]; cusEnts: FullCustomerEntitlement[]; featureDeductions: any; - event: Event; + properties: any; org: Organization; }) => { console.log(` - getCusEntsInFeatures: ${timeTaken}ms`); @@ -132,13 +134,13 @@ const logBalanceUpdate = ({ org.slug } | Features: ${features.map((f) => f.id).join(", ")}` ); - console.log(" - Properties:", event.properties); + console.log(" - Properties:", properties); console.log( " - CusEnts:", cusEnts.map((cusEnt: any) => { let balanceStr = cusEnt.balance; let { groupVal, balance } = getGroupBalanceFromProperties({ - properties: event.properties, + properties, cusEnt, features, }); @@ -229,6 +231,12 @@ export const deductAllowanceFromCusEnt = async ({ toDeduct = 0; } + const totalNegativeBalance = getTotalNegativeBalance(cusEnt); + const originalGrpBalance = Math.max(totalNegativeBalance, balance!); + const newGrpBalance = new Decimal(originalGrpBalance) + .minus(deducted) + .toNumber(); + await CustomerEntitlementService.update({ sb, id: cusEnt.id, @@ -247,8 +255,8 @@ export const deductAllowanceFromCusEnt = async ({ customer, affectedFeature: feature, cusEnt: cusEnt as any, - originalBalance: getMinCusEntBalance({ cusEnt }), - newBalance: getMinCusEntBalance({ cusEnt, newBalance, groupVal }), + originalBalance: originalGrpBalance, + newBalance: newGrpBalance, deduction: deducted, }); @@ -292,12 +300,10 @@ export const deductFromUsageBasedCusEnt = async ({ toDeduct, deductParams, cusEnts, - features, }: { toDeduct: number; deductParams: DeductParams; cusEnts: FullCustomerEntitlement[]; - features: Feature[]; }) => { const { sb, feature, env, org, cusPrices, customer, properties } = deductParams; @@ -343,6 +349,12 @@ export const deductFromUsageBasedCusEnt = async ({ }), }); + const totalNegativeBalance = getTotalNegativeBalance(usageBasedEnt); + const originalGrpBalance = Math.max(totalNegativeBalance, balance!); + const newGrpBalance = new Decimal(originalGrpBalance) + .minus(toDeduct) + .toNumber(); + await adjustAllowance({ sb, env, @@ -351,12 +363,8 @@ export const deductFromUsageBasedCusEnt = async ({ cusEnt: usageBasedEnt as any, cusPrices: cusPrices as any, customer, - originalBalance: getMinCusEntBalance({ cusEnt: usageBasedEnt }), - newBalance: getMinCusEntBalance({ - cusEnt: usageBasedEnt, - newBalance, - groupVal, - }), + originalBalance: originalGrpBalance, + newBalance: newGrpBalance, deduction: toDeduct, }); }; @@ -401,7 +409,7 @@ export const updateCustomerBalance = async ({ features, cusEnts, featureDeductions, - event, + properties: event.properties, org, }); @@ -453,7 +461,6 @@ export const updateCustomerBalance = async ({ await deductFromUsageBasedCusEnt({ toDeduct, cusEnts, - features, deductParams: { sb, feature, diff --git a/server/src/trigger/updateUsageTask.ts b/server/src/trigger/updateUsageTask.ts new file mode 100644 index 000000000..0347a3241 --- /dev/null +++ b/server/src/trigger/updateUsageTask.ts @@ -0,0 +1,354 @@ +import { createSupabaseClient } from "@/external/supabaseUtils.js"; +import { + getBelowThresholdPrice, + handleBelowThresholdInvoicing, +} from "./invoiceThresholdUtils.js"; +import { + AllowanceType, + AppEnv, + CusProductStatus, + Customer, + Feature, + FeatureType, + FullCustomerEntitlement, + Organization, +} from "@autumn/shared"; +import { SupabaseClient } from "@supabase/supabase-js"; +import { getCusEntsInFeatures } from "@/internal/api/customers/cusUtils.js"; + +import { featureToCreditSystem } from "@/internal/features/creditSystemUtils.js"; +import { getFeatureBalance } from "@/internal/customers/entitlements/cusEntUtils.js"; +import { Decimal } from "decimal.js"; +import { + getGroupBalanceFromProperties, + initGroupBalancesForEvent, +} from "@/internal/customers/entitlements/groupByUtils.js"; +import { + deductAllowanceFromCusEnt, + deductFromUsageBasedCusEnt, +} from "./updateBalanceTask.js"; +import { JobName } from "@/queue/JobName.js"; + +// 2. Get deductions for each feature +const getFeatureDeductions = ({ + cusEnts, + value, + features, + shouldSet, +}: { + cusEnts: FullCustomerEntitlement[]; + value: number; + features: Feature[]; + shouldSet: boolean; +}) => { + let meteredFeature = features.find((f) => f.type === FeatureType.Metered)!; + const featureDeductions = []; + for (const feature of features) { + let unlimitedExists = cusEnts.some( + (cusEnt) => + cusEnt.entitlement.allowance_type === AllowanceType.Unlimited && + cusEnt.entitlement.internal_feature_id == feature.internal_id + ); + + if (unlimitedExists) { + continue; + } + + if (feature.type === FeatureType.CreditSystem) { + value = featureToCreditSystem({ + featureId: meteredFeature.id, + creditSystem: feature, + amount: value, + }); + } + + // If it's set + let deduction = value; + + if (shouldSet) { + let totalAllowance = cusEnts.reduce((acc, curr) => { + return acc + (curr.entitlement.allowance || 0); + }, 0); + + let targetBalance = new Decimal(totalAllowance).sub(value).toNumber(); + + let totalBalance = getFeatureBalance({ + cusEnts, + internalFeatureId: feature.internal_id!, + })!; + + deduction = new Decimal(totalBalance).sub(targetBalance).toNumber(); + } + + if (deduction == 0) { + console.log(` - Skipping feature ${feature.id} -- deduction is 0`); + continue; + } + + featureDeductions.push({ + feature, + deduction, + }); + } + + featureDeductions.sort((a, b) => { + if ( + a.feature.type === FeatureType.CreditSystem && + b.feature.type !== FeatureType.CreditSystem + ) { + return 1; + } + + if ( + a.feature.type !== FeatureType.CreditSystem && + b.feature.type === FeatureType.CreditSystem + ) { + return -1; + } + + return a.feature.id.localeCompare(b.feature.id); + }); + + return featureDeductions; +}; + +const logUsageUpdate = ({ + timeTaken, + customer, + features, + cusEnts, + featureDeductions, + properties, + org, + setUsage, +}: { + timeTaken: string; + customer: Customer; + features: Feature[]; + cusEnts: FullCustomerEntitlement[]; + featureDeductions: any; + properties: any; + org: Organization; + setUsage: boolean; +}) => { + console.log(` - getCusEntsInFeatures: ${timeTaken}ms`); + console.log( + ` - Customer: ${customer.id} (${customer.env}) | Org: ${ + org.slug + } | Features: ${features.map((f) => f.id).join(", ")} | Set Usage: ${ + setUsage ? "true" : "false" + }` + ); + + console.log( + " - 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"; + } else if (groupVal) { + balanceStr = `${balance} [${groupVal}]`; + } + } catch (error) { + balanceStr = "failed_to_get_balance"; + } + + return `${cusEnt.feature_id} - ${balanceStr} (${ + cusEnt.customer_product ? cusEnt.customer_product.product_id : "" + })`; + }), + "| Deductions:", + featureDeductions.map((f: any) => `${f.feature.id}: ${f.deduction}`) + ); +}; + +// Main function to update customer balance +export const updateUsage = async ({ + sb, + customer, + features, + org, + env, + value, + properties, + setUsage, +}: { + sb: SupabaseClient; + customer: Customer; + features: Feature[]; + org: Organization; + env: AppEnv; + value: number; + properties: any; + setUsage: boolean; +}) => { + const startTime = performance.now(); + const { cusEnts, cusPrices } = await getCusEntsInFeatures({ + sb, + internalCustomerId: customer.internal_id, + internalFeatureIds: features.map((f) => f.internal_id!), + inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue], + withPrices: true, + }); + const endTime = performance.now(); + + // 1. Get deductions for each feature + const featureDeductions = getFeatureDeductions({ + cusEnts, + value, + shouldSet: setUsage, + features, + }); + + logUsageUpdate({ + timeTaken: (endTime - startTime).toFixed(2), + customer, + features, + cusEnts, + featureDeductions, + properties, + org, + setUsage, + }); + + // 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; + + for (const cusEnt of cusEnts) { + if (cusEnt.entitlement.internal_feature_id != feature.internal_id) { + continue; + } + + toDeduct = await deductAllowanceFromCusEnt({ + toDeduct, + cusEnt, + features, + deductParams: { + sb, + feature, + env, + org, + cusPrices: cusPrices as any[], + customer, + properties, + }, + featureDeductions, + willDeductCredits: true, + }); + } + + if (toDeduct == 0) { + continue; + } + + await deductFromUsageBasedCusEnt({ + toDeduct, + cusEnts, + deductParams: { + sb, + feature, + env, + org, + cusPrices: cusPrices as any[], + customer, + properties, + }, + }); + } + + return cusEnts; +}; + +// MAIN FUNCTION +export const runUpdateUsageTask = async ({ + payload, + logger, +}: { + payload: any; + logger: any; +}) => { + try { + const sb = createSupabaseClient(); + + // 1. Update customer balance + const { customer, features, value, set_usage, properties, org, env } = + payload; + + console.log("--------------------------------"); + console.log( + `HANDLING USAGE TASK FOR CUSTOMER (${customer.id}), ORG: ${org.slug}` + ); + + const cusEnts: any = await updateUsage({ + sb, + customer, + features, + value, + properties, + org, + env, + setUsage: set_usage, + }); + + 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 handleBelowThresholdInvoicing({ + sb, + internalCustomerId: payload.internalCustomerId, + belowThresholdPrice, + }); + } else { + console.log(" ✅ No below threshold price found"); + } + } catch (error) { + if (logger) { + logger.use((log: any) => { + return { + ...log, + task: JobName.UpdateUsage, + data: payload, + }; + }); + + logger.error(`ERROR UPDATING USAGE`); + logger.error(error); + } else { + console.log(error); + } + } +}; diff --git a/server/test.sh b/server/test.sh index 81a5cad32..1e6ded3af 100755 --- a/server/test.sh +++ b/server/test.sh @@ -13,6 +13,7 @@ elif [ "$1" == "advanced-parallel" ]; then 'tests/advanced/**/*.ts' --ignore 'tests/advanced/11_arrear_prorated.ts' \ --ignore 'tests/00_setup.ts' \ --ignore 'tests/alex/**/*.ts' + # --ignore 'tests/advanced/11_arrear_prorated2.ts' \ elif [ "$1" == "alex-parallel" ]; then npx mocha 'tests/alex/00_setup.ts' && npx mocha --parallel --timeout 10000000 \ diff --git a/shared/models/eventModels/eventModels.ts b/shared/models/eventModels/eventModels.ts index b816af645..73bee3072 100644 --- a/shared/models/eventModels/eventModels.ts +++ b/shared/models/eventModels/eventModels.ts @@ -14,6 +14,7 @@ export const EventSchema = z.object({ timestamp: z.number(), internal_customer_id: z.string(), value: z.number().nullish(), + set_usage: z.boolean().nullish(), }); export const CreateEventSchema = z.object({ @@ -23,6 +24,7 @@ export const CreateEventSchema = z.object({ timestamp: z.number().nullish(), idempotency_key: z.string().nullish(), value: z.number().nullish(), + set_usage: z.boolean().nullish(), }); export type Event = z.infer;