From db5efc61bb79b8c3b78b51c5e92495026e560472 Mon Sep 17 00:00:00 2001 From: John Yeo Date: Wed, 11 Mar 2026 17:03:53 +0000 Subject: [PATCH] feat: worked on check for spend limits --- .../api/check/checkTypes/CheckData.tsx | 3 +- .../api/check/checkUtils/getCheckData.ts | 13 +- .../check/checkUtils/getV2CheckResponse.ts | 11 +- .../helpers/fullCustomerToSpendLimit.ts | 34 + .../balances/utils/handleThresholdReached.ts | 4 + .../apiCusUtils/getApiCustomerBase.ts | 5 +- .../check-entity-product-spend-limit.test.ts | 609 ++++++++++++++++++ .../check-per-entity-spend-limit.test.ts | 397 ++++++++++++ .../update-entity-billing-controls.test.ts | 57 ++ .../billingControls/entityBillingControls.ts | 39 ++ .../{entities => }/billingControls/index.ts | 2 +- shared/api/billingControls/spendLimit.ts | 6 + shared/api/common/customerData.ts | 2 +- shared/api/common/entityData.ts | 2 +- shared/api/customers/apiSubjectV0.ts | 7 + shared/api/customers/baseApiCustomer.ts | 2 +- .../utils/check/getFeatureToUseForCheck.ts | 10 +- .../utils/convert/apiBalanceToAllowed.ts | 9 +- .../convert/apiBalanceV1ToAvailableOverage.ts | 24 +- shared/api/customers/index.ts | 1 + .../customers/utils/apiSubjectToSpendLimit.ts | 28 + shared/api/customers/utils/index.ts | 1 + shared/api/entities/apiEntityV2.ts | 2 +- .../billingControls/entityBillingControls.ts | 15 - .../billingControls/entitySpendLimit.ts | 6 - .../api/entities/crud/createEntityParams.ts | 4 +- .../api/entities/crud/updateEntityParams.ts | 4 +- shared/api/entities/index.ts | 2 +- shared/index.ts | 4 +- .../customerBillingControls.ts} | 20 +- .../billingControls/entityBillingControls.ts | 4 +- .../{entitySpendLimit.ts => spendLimit.ts} | 4 +- shared/models/cusModels/cusModels.ts | 6 +- shared/models/cusModels/cusTable.ts | 2 +- .../cusModels/entityModels/entityModels.ts | 4 +- .../cusModels/entityModels/entityTable.ts | 4 +- shared/models/cusModels/index.ts | 4 +- 37 files changed, 1278 insertions(+), 73 deletions(-) create mode 100644 server/src/internal/balances/autoTopUp/helpers/fullCustomerToSpendLimit.ts create mode 100644 server/tests/integration/balances/check/spend-limit/check-entity-product-spend-limit.test.ts create mode 100644 server/tests/integration/balances/check/spend-limit/check-per-entity-spend-limit.test.ts create mode 100644 shared/api/billingControls/entityBillingControls.ts rename shared/api/{entities => }/billingControls/index.ts (53%) create mode 100644 shared/api/billingControls/spendLimit.ts create mode 100644 shared/api/customers/apiSubjectV0.ts create mode 100644 shared/api/customers/utils/apiSubjectToSpendLimit.ts delete mode 100644 shared/api/entities/billingControls/entityBillingControls.ts delete mode 100644 shared/api/entities/billingControls/entitySpendLimit.ts rename shared/models/cusModels/{billingControlModels.ts => billingControls/customerBillingControls.ts} (79%) rename shared/models/cusModels/billingControls/{entitySpendLimit.ts => spendLimit.ts} (85%) diff --git a/server/src/internal/api/check/checkTypes/CheckData.tsx b/server/src/internal/api/check/checkTypes/CheckData.tsx index ebcaf06b5..9d2416b0f 100644 --- a/server/src/internal/api/check/checkTypes/CheckData.tsx +++ b/server/src/internal/api/check/checkTypes/CheckData.tsx @@ -1,9 +1,10 @@ -import type { ApiBalanceV1, Feature } from "@autumn/shared"; +import type { ApiBalanceV1, ApiCustomerV5, ApiEntityV2, Feature } from "@autumn/shared"; export interface CheckData { customerId: string; entityId?: string; apiBalance?: ApiBalanceV1; + apiSubject: ApiCustomerV5 | ApiEntityV2; originalFeature: Feature; featureToUse: Feature; } diff --git a/server/src/internal/api/check/checkUtils/getCheckData.ts b/server/src/internal/api/check/checkUtils/getCheckData.ts index 3db6620fa..6b56ab78c 100644 --- a/server/src/internal/api/check/checkUtils/getCheckData.ts +++ b/server/src/internal/api/check/checkUtils/getCheckData.ts @@ -56,7 +56,7 @@ export const getCheckData = async ({ throw new FeatureNotFoundError({ featureId: feature_id }); } - let apiEntity: ApiCustomerV5 | ApiEntityV2 | undefined; + let apiSubject: ApiCustomerV5 | ApiEntityV2 | undefined; const start = performance.now(); const fullCustomer = await getOrCreateCachedFullCustomer({ ctx, @@ -73,7 +73,7 @@ export const getCheckData = async ({ `[check] getOrCreateCachedFullCustomer took ${performance.now() - start}ms`, ); - apiEntity = apiCustomer; + apiSubject = apiCustomer; if (entity_id && fullCustomer.entity) { const { apiEntity: apiEntityResult } = await getApiEntityBase({ ctx, @@ -81,10 +81,10 @@ export const getCheckData = async ({ fullCus: fullCustomer, }); - apiEntity = apiEntityResult; + apiSubject = apiEntityResult; } - if (!apiEntity) { + if (!apiSubject) { throw new InternalError({ message: "failed to get entity object from cache", }); @@ -93,7 +93,7 @@ export const getCheckData = async ({ const featureToUseMin = getFeatureToUseForCheck({ creditSystems, feature, - apiEntity, + apiSubject, requiredBalance, }); @@ -112,12 +112,13 @@ export const getCheckData = async ({ ctx.logger.error(`[getCheckData] Failed to trigger auto top-up: ${error}`); }); - const apiBalance = apiEntity.balances?.[featureToUse.id]; + const apiBalance = apiSubject.balances?.[featureToUse.id]; return { customerId: customer_id, entityId: entity_id, apiBalance, + apiSubject, originalFeature: feature, featureToUse, }; diff --git a/server/src/internal/api/check/checkUtils/getV2CheckResponse.ts b/server/src/internal/api/check/checkUtils/getV2CheckResponse.ts index a775ace86..d5ceb3de6 100644 --- a/server/src/internal/api/check/checkUtils/getV2CheckResponse.ts +++ b/server/src/internal/api/check/checkUtils/getV2CheckResponse.ts @@ -13,8 +13,14 @@ export const getV2CheckResponse = async ({ checkData: CheckData; requiredBalance: number; }) => { - const { customerId, entityId, apiBalance, originalFeature, featureToUse } = - checkData; + const { + customerId, + entityId, + apiBalance, + apiSubject, + originalFeature, + featureToUse, + } = checkData; // If credit system used, need to convert required balance to credit system required balance if ( @@ -40,6 +46,7 @@ export const getV2CheckResponse = async ({ const allowed = apiBalanceToAllowed({ apiBalance, + apiSubject, feature: featureToUse, requiredBalance, }); diff --git a/server/src/internal/balances/autoTopUp/helpers/fullCustomerToSpendLimit.ts b/server/src/internal/balances/autoTopUp/helpers/fullCustomerToSpendLimit.ts new file mode 100644 index 000000000..4fc5c42cc --- /dev/null +++ b/server/src/internal/balances/autoTopUp/helpers/fullCustomerToSpendLimit.ts @@ -0,0 +1,34 @@ +import type { DbSpendLimit, FullCustomer } from "@autumn/shared"; + +/** Extract the enabled spend limit for a given feature from a FullCustomer. Returns undefined if none found. */ +export const fullCustomerToSpendLimit = ({ + fullCustomer, + featureId, + internalEntityId, +}: { + fullCustomer: FullCustomer; + featureId: string; + internalEntityId?: string; +}): DbSpendLimit | undefined => { + if (internalEntityId) { + fullCustomer.entity = fullCustomer.entities.find( + (entity) => entity.id === internalEntityId, + ); + } + + if (fullCustomer.entity) { + return fullCustomer.entity.spend_limits?.find( + (spendLimit) => + spendLimit.feature_id === featureId && + spendLimit.enabled && + spendLimit.overage_limit !== undefined, + ); + } + + return fullCustomer.spend_limits?.find( + (spendLimit) => + spendLimit.feature_id === featureId && + spendLimit.enabled && + spendLimit.overage_limit !== undefined, + ); +}; diff --git a/server/src/internal/balances/utils/handleThresholdReached.ts b/server/src/internal/balances/utils/handleThresholdReached.ts index 2a13f47a8..21ae90224 100644 --- a/server/src/internal/balances/utils/handleThresholdReached.ts +++ b/server/src/internal/balances/utils/handleThresholdReached.ts @@ -70,12 +70,14 @@ const handleAllowanceUsed = async ({ const oldAllowed = apiBalanceToAllowed({ apiBalance: prevCusFeature, + apiSubject: prevApiCustomer, feature, requiredBalance: 1, }); const newAllowed = apiBalanceToAllowed({ apiBalance: newCusFeature, + apiSubject: newApiCustomer, feature, requiredBalance: 1, }); @@ -139,12 +141,14 @@ export const handleThresholdReached = async ({ const oldAllowed = apiBalanceToAllowed({ apiBalance: prevApiCustomer.balances[feature.id], + apiSubject: prevApiCustomer, feature, requiredBalance: 1, }); const newAllowed = apiBalanceToAllowed({ apiBalance: newApiCustomer.balances[feature.id], + apiSubject: newApiCustomer, feature, requiredBalance: 1, }); diff --git a/server/src/internal/customers/cusUtils/apiCusUtils/getApiCustomerBase.ts b/server/src/internal/customers/cusUtils/apiCusUtils/getApiCustomerBase.ts index 81c015f46..f03213024 100644 --- a/server/src/internal/customers/cusUtils/apiCusUtils/getApiCustomerBase.ts +++ b/server/src/internal/customers/cusUtils/apiCusUtils/getApiCustomerBase.ts @@ -61,7 +61,10 @@ export const getApiCustomerBase = async ({ purchases: apiPurchases, balances: apiBalances, send_email_receipts: fullCus.send_email_receipts ?? false, - billing_controls: { auto_topups: fullCus.auto_topups ?? undefined }, + billing_controls: { + auto_topups: fullCus.auto_topups ?? undefined, + spend_limits: fullCus.spend_limits ?? undefined, + }, invoices: fullCus.invoices && ctx.expand.includes(CustomerExpand.Invoices) diff --git a/server/tests/integration/balances/check/spend-limit/check-entity-product-spend-limit.test.ts b/server/tests/integration/balances/check/spend-limit/check-entity-product-spend-limit.test.ts new file mode 100644 index 000000000..d599fe3a3 --- /dev/null +++ b/server/tests/integration/balances/check/spend-limit/check-entity-product-spend-limit.test.ts @@ -0,0 +1,609 @@ +import { expect, test } from "bun:test"; +import type { CheckResponseV3, EntityBillingControls } from "@autumn/shared"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { items } from "@tests/utils/fixtures/items.js"; +import { products } from "@tests/utils/fixtures/products.js"; +import { timeout } from "@tests/utils/genUtils.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; +import { getCreditCost } from "@/internal/features/creditSystemUtils"; + +type AutumnV2_1Client = Awaited>["autumnV2_1"]; + +const normalizeCheckResponse = (response: CheckResponseV3) => ({ + allowed: response.allowed, + customer_id: response.customer_id, + entity_id: response.entity_id ?? null, + required_balance: response.required_balance ?? null, + balance: response.balance + ? { + feature_id: response.balance.feature_id, + granted: response.balance.granted, + remaining: response.balance.remaining, + usage: response.balance.usage, + unlimited: response.balance.unlimited, + overage_allowed: response.balance.overage_allowed, + max_purchase: response.balance.max_purchase, + breakdown: + response.balance.breakdown?.map((item) => ({ + plan_id: item.plan_id, + included_grant: item.included_grant, + prepaid_grant: item.prepaid_grant, + remaining: item.remaining, + usage: item.usage, + unlimited: item.unlimited, + billing_method: item.price?.billing_method ?? null, + max_purchase: item.price?.max_purchase ?? null, + reset_interval: item.reset?.interval ?? null, + })) ?? [], + } + : null, +}); + +const setEntitySpendLimit = async ({ + autumn, + customerId, + entityId, + featureId, + overageLimit, + enabled = true, +}: { + autumn: AutumnV2_1Client; + customerId: string; + entityId: string; + featureId: string; + overageLimit: number; + enabled?: boolean; +}) => { + const billingControls: EntityBillingControls = { + spend_limits: [ + { + feature_id: featureId, + enabled, + overage_limit: overageLimit, + }, + ], + }; + + await autumn.entities.update(customerId, entityId, { + billing_controls: billingControls, + }); +}; + +const getActionUnitsForCreditAmount = ({ + creditAmount, + creditCostPerActionUnit, +}: { + creditAmount: number; + creditCostPerActionUnit: number; +}) => creditAmount / creditCostPerActionUnit; + +const expectBoundaryAndParity = async ({ + autumn, + customerId, + entityId, + featureId, + allowedRequiredBalance, + blockedRequiredBalance, + expectedFeatureId = featureId, + expectedAllowedResponseRequiredBalance = allowedRequiredBalance, + expectedBlockedResponseRequiredBalance = blockedRequiredBalance, +}: { + autumn: AutumnV2_1Client; + customerId: string; + entityId: string; + featureId: string; + allowedRequiredBalance: number; + blockedRequiredBalance: number; + expectedFeatureId?: string; + expectedAllowedResponseRequiredBalance?: number; + expectedBlockedResponseRequiredBalance?: number; +}) => { + const allowedCached = await autumn.check({ + customer_id: customerId, + entity_id: entityId, + feature_id: featureId, + required_balance: allowedRequiredBalance, + }); + + const blockedCached = await autumn.check({ + customer_id: customerId, + entity_id: entityId, + feature_id: featureId, + required_balance: blockedRequiredBalance, + }); + + expect(allowedCached.allowed).toBe(true); + expect(blockedCached.allowed).toBe(false); + expect(allowedCached.balance?.feature_id).toBe(expectedFeatureId); + expect(blockedCached.balance?.feature_id).toBe(expectedFeatureId); + expect(allowedCached.required_balance).toBe( + expectedAllowedResponseRequiredBalance, + ); + expect(blockedCached.required_balance).toBe( + expectedBlockedResponseRequiredBalance, + ); + + await timeout(4000); + + const allowedUncached = await autumn.check({ + customer_id: customerId, + entity_id: entityId, + feature_id: featureId, + required_balance: allowedRequiredBalance, + skip_cache: true, + }); + + const blockedUncached = await autumn.check({ + customer_id: customerId, + entity_id: entityId, + feature_id: featureId, + required_balance: blockedRequiredBalance, + skip_cache: true, + }); + + expect(normalizeCheckResponse(allowedUncached)).toEqual( + normalizeCheckResponse(allowedCached), + ); + expect(normalizeCheckResponse(blockedUncached)).toEqual( + normalizeCheckResponse(blockedCached), + ); +}; + +test.concurrent(`${chalk.yellowBright("check-entity-product-spend-limit1: lifetime + consumable entity product respects spend limit and cache parity")}`, async () => { + const entityProduct = products.base({ + id: "entity-product-lifetime-consumable", + items: [ + items.lifetimeMessages({ + includedUsage: 1000, + }), + items.consumableMessages({ + includedUsage: 100, + maxPurchase: 300, + price: 0.5, + }), + ], + }); + + const { autumnV2_1, customerId, entities } = await initScenario({ + customerId: "check-entity-product-spend-limit-1", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [entityProduct] }), + s.entities({ count: 1, featureId: TestFeature.Users }), + ], + actions: [s.attach({ productId: entityProduct.id, entityIndex: 0 })], + }); + + await setEntitySpendLimit({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Messages, + overageLimit: 25, + }); + + await autumnV2_1.track({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + value: 1120, + }); + + await expectBoundaryAndParity({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Messages, + allowedRequiredBalance: 5, + blockedRequiredBalance: 6, + }); +}); + +test.concurrent(`${chalk.yellowBright("check-entity-product-spend-limit2: prepaid + consumable entity product respects spend limit and cache parity")}`, async () => { + const entityProduct = products.base({ + id: "entity-product-prepaid-consumable", + items: [ + items.prepaidMessages({ + includedUsage: 100, + billingUnits: 100, + price: 8.5, + }), + items.consumableMessages({ + includedUsage: 200, + maxPurchase: 300, + price: 0.5, + }), + ], + }); + + const prepaidQuantity = 500; + const { autumnV2_1, customerId, entities } = await initScenario({ + customerId: "check-entity-product-spend-limit-2", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [entityProduct] }), + s.entities({ count: 1, featureId: TestFeature.Users }), + ], + actions: [ + s.attach({ + productId: entityProduct.id, + entityIndex: 0, + options: [ + { + feature_id: TestFeature.Messages, + quantity: prepaidQuantity, + }, + ], + }), + ], + }); + + await setEntitySpendLimit({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Messages, + overageLimit: 25, + }); + + await autumnV2_1.track({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + value: 820, + }); + + await expectBoundaryAndParity({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Messages, + allowedRequiredBalance: 5, + blockedRequiredBalance: 6, + }); +}); + +test.concurrent(`${chalk.yellowBright("check-entity-product-spend-limit3: two entities with different spend limits stay isolated and match skip_cache")}`, async () => { + const entityProduct = products.base({ + id: "entity-product-two-entities", + items: [ + items.prepaidMessages({ + includedUsage: 100, + billingUnits: 100, + price: 8.5, + }), + items.consumableMessages({ + includedUsage: 200, + maxPurchase: 300, + price: 0.5, + }), + ], + }); + + const prepaidQuantity = 500; + const { autumnV2_1, customerId, entities } = await initScenario({ + customerId: "check-entity-product-spend-limit-3", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [entityProduct] }), + s.entities({ count: 2, featureId: TestFeature.Users }), + ], + actions: [ + s.attach({ + productId: entityProduct.id, + entityIndex: 0, + options: [ + { + feature_id: TestFeature.Messages, + quantity: prepaidQuantity, + }, + ], + }), + s.attach({ + productId: entityProduct.id, + entityIndex: 1, + options: [ + { + feature_id: TestFeature.Messages, + quantity: prepaidQuantity, + }, + ], + }), + ], + }); + + await setEntitySpendLimit({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Messages, + overageLimit: 25, + }); + + await setEntitySpendLimit({ + autumn: autumnV2_1, + customerId, + entityId: entities[1].id, + featureId: TestFeature.Messages, + overageLimit: 40, + }); + + await autumnV2_1.track({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + value: 820, + }); + + await autumnV2_1.track({ + customer_id: customerId, + entity_id: entities[1].id, + feature_id: TestFeature.Messages, + value: 820, + }); + + const entity1AllowedCached = await autumnV2_1.check({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + required_balance: 5, + }); + const entity1BlockedCached = await autumnV2_1.check({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + required_balance: 6, + }); + const entity2AllowedCached = await autumnV2_1.check({ + customer_id: customerId, + entity_id: entities[1].id, + feature_id: TestFeature.Messages, + required_balance: 20, + }); + const entity2BlockedCached = await autumnV2_1.check({ + customer_id: customerId, + entity_id: entities[1].id, + feature_id: TestFeature.Messages, + required_balance: 21, + }); + + expect(entity1AllowedCached.allowed).toBe(true); + expect(entity1BlockedCached.allowed).toBe(false); + expect(entity2AllowedCached.allowed).toBe(true); + expect(entity2BlockedCached.allowed).toBe(false); + + await timeout(4000); + + const entity1AllowedUncached = await autumnV2_1.check({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + required_balance: 5, + skip_cache: true, + }); + const entity1BlockedUncached = await autumnV2_1.check({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + required_balance: 6, + skip_cache: true, + }); + const entity2AllowedUncached = await autumnV2_1.check({ + customer_id: customerId, + entity_id: entities[1].id, + feature_id: TestFeature.Messages, + required_balance: 20, + skip_cache: true, + }); + const entity2BlockedUncached = await autumnV2_1.check({ + customer_id: customerId, + entity_id: entities[1].id, + feature_id: TestFeature.Messages, + required_balance: 21, + skip_cache: true, + }); + + expect(normalizeCheckResponse(entity1AllowedUncached)).toEqual( + normalizeCheckResponse(entity1AllowedCached), + ); + expect(normalizeCheckResponse(entity1BlockedUncached)).toEqual( + normalizeCheckResponse(entity1BlockedCached), + ); + expect(normalizeCheckResponse(entity2AllowedUncached)).toEqual( + normalizeCheckResponse(entity2AllowedCached), + ); + expect(normalizeCheckResponse(entity2BlockedUncached)).toEqual( + normalizeCheckResponse(entity2BlockedCached), + ); +}); + +test.concurrent(`${chalk.yellowBright("check-entity-product-spend-limit4: allocated workflows entity product respects spend limit and cache parity")}`, async () => { + const entityProduct = products.base({ + id: "entity-product-workflows", + items: [items.allocatedWorkflows({ includedUsage: 1 })], + }); + + const { autumnV2_1, customerId, entities } = await initScenario({ + customerId: "check-entity-product-spend-limit-4", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [entityProduct] }), + s.entities({ count: 1, featureId: TestFeature.Users }), + ], + actions: [s.attach({ productId: entityProduct.id, entityIndex: 0 })], + }); + + await setEntitySpendLimit({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Workflows, + overageLimit: 2, + }); + + await autumnV2_1.track({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Workflows, + value: 2, + }); + + await expectBoundaryAndParity({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Workflows, + allowedRequiredBalance: 1, + blockedRequiredBalance: 2, + }); +}); + +test.concurrent(`${chalk.yellowBright("check-entity-product-spend-limit5: credit-system entity product uses converted credits and cache parity")}`, async () => { + const includedCredits = 100; + const spendLimitCredits = 25; + const existingOverageCredits = 20; + + const entityProduct = products.base({ + id: "entity-product-credits", + items: [ + items.consumable({ + featureId: TestFeature.Credits, + includedUsage: includedCredits, + maxPurchase: 300, + price: 0.5, + }), + ], + }); + + const { autumnV2_1, customerId, entities, ctx } = await initScenario({ + customerId: "check-entity-product-spend-limit-5", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [entityProduct] }), + s.entities({ count: 1, featureId: TestFeature.Users }), + ], + actions: [s.attach({ productId: entityProduct.id, entityIndex: 0 })], + }); + + const creditsFeature = ctx.features.find( + (f) => f.id === TestFeature.Credits, + )!; + const action1CreditCost = getCreditCost({ + featureId: TestFeature.Action1, + creditSystem: creditsFeature, + amount: 1, + }); + const creditsRemainingUntilLimit = spendLimitCredits - existingOverageCredits; + const usageToReachOverageBoundary = getActionUnitsForCreditAmount({ + creditAmount: includedCredits + existingOverageCredits, + creditCostPerActionUnit: action1CreditCost, + }); + const allowedActionUnits = getActionUnitsForCreditAmount({ + creditAmount: creditsRemainingUntilLimit, + creditCostPerActionUnit: action1CreditCost, + }); + const blockedActionUnits = getActionUnitsForCreditAmount({ + creditAmount: creditsRemainingUntilLimit + 1, + creditCostPerActionUnit: action1CreditCost, + }); + + await setEntitySpendLimit({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Credits, + overageLimit: spendLimitCredits, + }); + + await autumnV2_1.track({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Action1, + value: usageToReachOverageBoundary, + }); + + await expectBoundaryAndParity({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Action1, + allowedRequiredBalance: allowedActionUnits, + blockedRequiredBalance: blockedActionUnits, + expectedFeatureId: TestFeature.Credits, + expectedAllowedResponseRequiredBalance: creditsRemainingUntilLimit, + expectedBlockedResponseRequiredBalance: creditsRemainingUntilLimit + 1, + }); +}); + +test.concurrent(`${chalk.yellowBright("check-entity-product-spend-limit6: disabled spend limit falls back to max_purchase and matches skip_cache")}`, async () => { + const entityProduct = products.base({ + id: "entity-product-disabled-limit", + items: [ + items.consumableMessages({ + includedUsage: 100, + maxPurchase: 300, + price: 0.5, + }), + ], + }); + + const { autumnV2_1, customerId, entities } = await initScenario({ + customerId: "check-entity-product-spend-limit-6", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [entityProduct] }), + s.entities({ count: 1, featureId: TestFeature.Users }), + ], + actions: [s.attach({ productId: entityProduct.id, entityIndex: 0 })], + }); + + await setEntitySpendLimit({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Messages, + overageLimit: 25, + }); + + await autumnV2_1.track({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + value: 120, + }); + + await setEntitySpendLimit({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Messages, + overageLimit: 25, + enabled: false, + }); + + const cached = await autumnV2_1.check({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + required_balance: 50, + }); + + expect(cached.allowed).toBe(true); + expect(cached.balance?.feature_id).toBe(TestFeature.Messages); + expect(cached.required_balance).toBe(50); + + await timeout(4000); + + const uncached = await autumnV2_1.check({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + required_balance: 50, + skip_cache: true, + }); + + expect(normalizeCheckResponse(uncached)).toEqual( + normalizeCheckResponse(cached), + ); +}); diff --git a/server/tests/integration/balances/check/spend-limit/check-per-entity-spend-limit.test.ts b/server/tests/integration/balances/check/spend-limit/check-per-entity-spend-limit.test.ts new file mode 100644 index 000000000..7ae2adaa4 --- /dev/null +++ b/server/tests/integration/balances/check/spend-limit/check-per-entity-spend-limit.test.ts @@ -0,0 +1,397 @@ +import { expect, test } from "bun:test"; +import type { CheckResponseV3, EntityBillingControls } from "@autumn/shared"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { items } from "@tests/utils/fixtures/items.js"; +import { products } from "@tests/utils/fixtures/products.js"; +import { timeout } from "@tests/utils/genUtils.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; +import { getCreditCost } from "@/internal/features/creditSystemUtils"; +import { constructArrearProratedItem } from "@/utils/scriptUtils/constructItem.js"; + +type AutumnV2_1Client = Awaited>["autumnV2_1"]; + +const normalizeCheckResponse = (response: CheckResponseV3) => ({ + allowed: response.allowed, + customer_id: response.customer_id, + entity_id: response.entity_id ?? null, + required_balance: response.required_balance ?? null, + balance: response.balance + ? { + feature_id: response.balance.feature_id, + granted: response.balance.granted, + remaining: response.balance.remaining, + usage: response.balance.usage, + unlimited: response.balance.unlimited, + overage_allowed: response.balance.overage_allowed, + max_purchase: response.balance.max_purchase, + breakdown: + response.balance.breakdown?.map((item) => ({ + plan_id: item.plan_id, + included_grant: item.included_grant, + prepaid_grant: item.prepaid_grant, + remaining: item.remaining, + usage: item.usage, + unlimited: item.unlimited, + billing_method: item.price?.billing_method ?? null, + max_purchase: item.price?.max_purchase ?? null, + reset_interval: item.reset?.interval ?? null, + })) ?? [], + } + : null, +}); + +const setEntitySpendLimit = async ({ + autumn, + customerId, + entityId, + featureId, + overageLimit, + enabled = true, +}: { + autumn: AutumnV2_1Client; + customerId: string; + entityId: string; + featureId: string; + overageLimit: number; + enabled?: boolean; +}) => { + const billingControls: EntityBillingControls = { + spend_limits: [ + { + feature_id: featureId, + enabled, + overage_limit: overageLimit, + }, + ], + }; + + await autumn.entities.update(customerId, entityId, { + billing_controls: billingControls, + }); +}; + +const getActionUnitsForCreditAmount = ({ + creditAmount, + creditCostPerActionUnit, +}: { + creditAmount: number; + creditCostPerActionUnit: number; +}) => creditAmount / creditCostPerActionUnit; + +const expectBoundaryAndParity = async ({ + autumn, + customerId, + entityId, + featureId, + allowedRequiredBalance, + blockedRequiredBalance, + expectedFeatureId = featureId, + expectedAllowedResponseRequiredBalance = allowedRequiredBalance, + expectedBlockedResponseRequiredBalance = blockedRequiredBalance, +}: { + autumn: AutumnV2_1Client; + customerId: string; + entityId: string; + featureId: string; + allowedRequiredBalance: number; + blockedRequiredBalance: number; + expectedFeatureId?: string; + expectedAllowedResponseRequiredBalance?: number; + expectedBlockedResponseRequiredBalance?: number; +}) => { + const allowedCached = await autumn.check({ + customer_id: customerId, + entity_id: entityId, + feature_id: featureId, + required_balance: allowedRequiredBalance, + }); + + const blockedCached = await autumn.check({ + customer_id: customerId, + entity_id: entityId, + feature_id: featureId, + required_balance: blockedRequiredBalance, + }); + + expect(allowedCached.allowed).toBe(true); + expect(blockedCached.allowed).toBe(false); + expect(allowedCached.balance?.feature_id).toBe(expectedFeatureId); + expect(blockedCached.balance?.feature_id).toBe(expectedFeatureId); + expect(allowedCached.required_balance).toBe( + expectedAllowedResponseRequiredBalance, + ); + expect(blockedCached.required_balance).toBe( + expectedBlockedResponseRequiredBalance, + ); + + await timeout(4000); + + const allowedUncached = await autumn.check({ + customer_id: customerId, + entity_id: entityId, + feature_id: featureId, + required_balance: allowedRequiredBalance, + skip_cache: true, + }); + + const blockedUncached = await autumn.check({ + customer_id: customerId, + entity_id: entityId, + feature_id: featureId, + required_balance: blockedRequiredBalance, + skip_cache: true, + }); + + expect(normalizeCheckResponse(allowedUncached)).toEqual( + normalizeCheckResponse(allowedCached), + ); + expect(normalizeCheckResponse(blockedUncached)).toEqual( + normalizeCheckResponse(blockedCached), + ); +}; + +test.concurrent(`${chalk.yellowBright("check-per-entity-spend-limit1: lifetime + consumable per-entity messages respect spend limit and cache parity")}`, async () => { + const perEntityProduct = products.base({ + id: "per-entity-lifetime-consumable", + items: [ + items.lifetimeMessages({ + includedUsage: 1000, + entityFeatureId: TestFeature.Users, + }), + items.consumableMessages({ + includedUsage: 100, + maxPurchase: 300, + price: 0.5, + entityFeatureId: TestFeature.Users, + }), + ], + }); + + const { autumnV2_1, customerId, entities } = await initScenario({ + customerId: "check-per-entity-spend-limit-1", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [perEntityProduct] }), + s.entities({ count: 1, featureId: TestFeature.Users }), + ], + actions: [s.attach({ productId: perEntityProduct.id })], + }); + + await setEntitySpendLimit({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Messages, + overageLimit: 25, + }); + + await autumnV2_1.track({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + value: 1120, + }); + + await expectBoundaryAndParity({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Messages, + allowedRequiredBalance: 5, + blockedRequiredBalance: 6, + }); +}); + +test.concurrent(`${chalk.yellowBright("check-per-entity-spend-limit2: prepaid + consumable per-entity messages respect spend limit and cache parity")}`, async () => { + const perEntityProduct = products.base({ + id: "per-entity-prepaid-consumable", + items: [ + items.prepaidMessages({ + includedUsage: 100, + billingUnits: 100, + price: 8.5, + entityFeatureId: TestFeature.Users, + }), + items.consumableMessages({ + includedUsage: 200, + maxPurchase: 300, + price: 0.5, + entityFeatureId: TestFeature.Users, + }), + ], + }); + + const prepaidQuantity = 500; + const { autumnV2_1, customerId, entities } = await initScenario({ + customerId: "check-per-entity-spend-limit-2", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [perEntityProduct] }), + s.entities({ count: 1, featureId: TestFeature.Users }), + ], + actions: [ + s.attach({ + productId: perEntityProduct.id, + options: [ + { + feature_id: TestFeature.Messages, + quantity: prepaidQuantity, + }, + ], + }), + ], + }); + + await setEntitySpendLimit({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Messages, + overageLimit: 25, + }); + + await autumnV2_1.track({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + value: 820, + }); + + await expectBoundaryAndParity({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Messages, + allowedRequiredBalance: 5, + blockedRequiredBalance: 6, + }); +}); + +test.concurrent(`${chalk.yellowBright("check-per-entity-spend-limit3: allocated workflows per entity respect spend limit and cache parity")}`, async () => { + const workflowItem = { + ...constructArrearProratedItem({ + featureId: TestFeature.Workflows, + pricePerUnit: 10, + includedUsage: 1, + }), + entity_feature_id: TestFeature.Users, + }; + + const perEntityProduct = products.base({ + id: "per-entity-workflows", + items: [workflowItem], + }); + + const { autumnV2_1, customerId, entities } = await initScenario({ + customerId: "check-per-entity-spend-limit-3", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [perEntityProduct] }), + s.entities({ count: 1, featureId: TestFeature.Users }), + ], + actions: [s.attach({ productId: perEntityProduct.id })], + }); + + await setEntitySpendLimit({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Workflows, + overageLimit: 2, + }); + + await autumnV2_1.track({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Workflows, + value: 2, + }); + + await expectBoundaryAndParity({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Workflows, + allowedRequiredBalance: 1, + blockedRequiredBalance: 2, + }); +}); + +test.concurrent(`${chalk.yellowBright("check-per-entity-spend-limit4: credit-system checks use converted credits and respect spend limit with cache parity")}`, async () => { + const includedCredits = 100; + const spendLimitCredits = 25; + const existingOverageCredits = 20; + + const perEntityCredits = items.consumable({ + featureId: TestFeature.Credits, + includedUsage: 100, + maxPurchase: 300, + price: 0.5, + entityFeatureId: TestFeature.Users, + }); + + const perEntityProduct = products.base({ + id: "per-entity-credits", + items: [perEntityCredits], + }); + + const { autumnV2_1, customerId, entities, ctx } = await initScenario({ + customerId: "check-per-entity-spend-limit-4", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [perEntityProduct] }), + s.entities({ count: 1, featureId: TestFeature.Users }), + ], + actions: [s.attach({ productId: perEntityProduct.id })], + }); + + const creditsFeature = ctx.features.find( + (f) => f.id === TestFeature.Credits, + )!; + const action1CreditCost = getCreditCost({ + featureId: TestFeature.Action1, + creditSystem: creditsFeature, + amount: 1, + }); + const creditsRemainingUntilLimit = spendLimitCredits - existingOverageCredits; + const usageToReachOverageBoundary = getActionUnitsForCreditAmount({ + creditAmount: includedCredits + existingOverageCredits, + creditCostPerActionUnit: action1CreditCost, + }); + const allowedActionUnits = getActionUnitsForCreditAmount({ + creditAmount: creditsRemainingUntilLimit, + creditCostPerActionUnit: action1CreditCost, + }); + const blockedActionUnits = getActionUnitsForCreditAmount({ + creditAmount: creditsRemainingUntilLimit + 1, + creditCostPerActionUnit: action1CreditCost, + }); + + await setEntitySpendLimit({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Credits, + overageLimit: spendLimitCredits, + }); + + await autumnV2_1.track({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Action1, + value: usageToReachOverageBoundary, + }); + + await expectBoundaryAndParity({ + autumn: autumnV2_1, + customerId, + entityId: entities[0].id, + featureId: TestFeature.Action1, + allowedRequiredBalance: allowedActionUnits, + blockedRequiredBalance: blockedActionUnits, + expectedFeatureId: TestFeature.Credits, + expectedAllowedResponseRequiredBalance: creditsRemainingUntilLimit, + expectedBlockedResponseRequiredBalance: creditsRemainingUntilLimit + 1, + }); +}); diff --git a/server/tests/integration/crud/entities/update-entity-billing-controls.test.ts b/server/tests/integration/crud/entities/update-entity-billing-controls.test.ts index 9d3d4bc0c..adc6727de 100644 --- a/server/tests/integration/crud/entities/update-entity-billing-controls.test.ts +++ b/server/tests/integration/crud/entities/update-entity-billing-controls.test.ts @@ -114,3 +114,60 @@ test.concurrent(`${chalk.yellowBright("entity billing controls: require feature_ }), }); }); + +test.concurrent(`${chalk.yellowBright("entity billing controls: reject duplicate spend limit feature ids on create and update")}`, async () => { + const { customerId, autumnV2_1 } = await initScenario({ + customerId: "entity-billing-controls-3", + setup: [s.customer({})], + actions: [], + }); + + await expectAutumnError({ + func: async () => + await autumnV2_1.entities.create(customerId, { + id: "entity-3", + name: "Entity 3", + feature_id: TestFeature.Users, + billing_controls: { + spend_limits: [ + { + feature_id: TestFeature.Messages, + enabled: true, + overage_limit: 10, + }, + { + feature_id: TestFeature.Messages, + enabled: false, + overage_limit: 20, + }, + ], + }, + }), + }); + + await autumnV2_1.entities.create(customerId, { + id: "entity-4", + name: "Entity 4", + feature_id: TestFeature.Users, + }); + + await expectAutumnError({ + func: async () => + await autumnV2_1.entities.update(customerId, "entity-4", { + billing_controls: { + spend_limits: [ + { + feature_id: TestFeature.Credits, + enabled: true, + overage_limit: 15, + }, + { + feature_id: TestFeature.Credits, + enabled: true, + overage_limit: 30, + }, + ], + }, + }), + }); +}); diff --git a/shared/api/billingControls/entityBillingControls.ts b/shared/api/billingControls/entityBillingControls.ts new file mode 100644 index 000000000..f20db714e --- /dev/null +++ b/shared/api/billingControls/entityBillingControls.ts @@ -0,0 +1,39 @@ +import { z } from "zod/v4"; +import { ApiSpendLimitSchema } from "./spendLimit.js"; + +export const ApiEntityBillingControlsSchema = z.object({ + spend_limits: z.array(ApiSpendLimitSchema).optional().meta({ + description: "List of overage spend limits per feature.", + }), +}); + +export const ApiEntityBillingControlsInputSchema = + ApiEntityBillingControlsSchema.superRefine((billingControls, ctx) => { + const featureIds = new Set(); + + for (const [index, spendLimit] of ( + billingControls.spend_limits ?? [] + ).entries()) { + if (!spendLimit.feature_id) { + continue; + } + + if (featureIds.has(spendLimit.feature_id)) { + ctx.addIssue({ + code: "custom", + message: "Only one spend limit entry is allowed per feature_id", + path: ["spend_limits", index, "feature_id"], + }); + return; + } + + featureIds.add(spendLimit.feature_id); + } + }); + +export type ApiEntityBillingControls = z.infer< + typeof ApiEntityBillingControlsSchema +>; +export type ApiEntityBillingControlsInput = z.input< + typeof ApiEntityBillingControlsInputSchema +>; diff --git a/shared/api/entities/billingControls/index.ts b/shared/api/billingControls/index.ts similarity index 53% rename from shared/api/entities/billingControls/index.ts rename to shared/api/billingControls/index.ts index d4f81579e..3e1cd4fa2 100644 --- a/shared/api/entities/billingControls/index.ts +++ b/shared/api/billingControls/index.ts @@ -1,2 +1,2 @@ export * from "./entityBillingControls.js"; -export * from "./entitySpendLimit.js"; +export * from "./spendLimit.js"; diff --git a/shared/api/billingControls/spendLimit.ts b/shared/api/billingControls/spendLimit.ts new file mode 100644 index 000000000..15cf7121d --- /dev/null +++ b/shared/api/billingControls/spendLimit.ts @@ -0,0 +1,6 @@ +import type { z } from "zod/v4"; +import { DbSpendLimitSchema } from "../../models/cusModels/billingControls/spendLimit.js"; + +export const ApiSpendLimitSchema = DbSpendLimitSchema; + +export type ApiSpendLimit = z.infer; diff --git a/shared/api/common/customerData.ts b/shared/api/common/customerData.ts index 586328720..b2ff6d586 100644 --- a/shared/api/common/customerData.ts +++ b/shared/api/common/customerData.ts @@ -1,5 +1,5 @@ import { z } from "zod/v4"; -import { CustomerBillingControlsSchema } from "../../models/cusModels/billingControlModels"; +import { CustomerBillingControlsSchema } from "../../models/cusModels/billingControls/customerBillingControls"; import { ExternalProcessorsSchema } from "../../models/genModels/processorSchemas"; // for internal use only diff --git a/shared/api/common/entityData.ts b/shared/api/common/entityData.ts index fce43157c..6b416ba24 100644 --- a/shared/api/common/entityData.ts +++ b/shared/api/common/entityData.ts @@ -1,5 +1,5 @@ import { z } from "zod/v4"; -import { ApiEntityBillingControlsSchema } from "../entities/billingControls/entityBillingControls.js"; +import { ApiEntityBillingControlsSchema } from "../billingControls/entityBillingControls.js"; export const EntityDataSchema = z .object({ diff --git a/shared/api/customers/apiSubjectV0.ts b/shared/api/customers/apiSubjectV0.ts new file mode 100644 index 000000000..94369a30b --- /dev/null +++ b/shared/api/customers/apiSubjectV0.ts @@ -0,0 +1,7 @@ +import type { z } from "zod/v4"; +import { ApiEntityV2Schema } from "../entities/apiEntityV2"; +import { ApiCustomerV5Schema } from "./apiCustomerV5"; + +export const ApiSubjectV0Schema = ApiCustomerV5Schema.or(ApiEntityV2Schema); + +export type ApiSubjectV0 = z.infer; diff --git a/shared/api/customers/baseApiCustomer.ts b/shared/api/customers/baseApiCustomer.ts index 130eda24e..e7ef0efca 100644 --- a/shared/api/customers/baseApiCustomer.ts +++ b/shared/api/customers/baseApiCustomer.ts @@ -1,4 +1,4 @@ -import { CustomerBillingControlsSchema } from "@models/cusModels/billingControlModels"; +import { CustomerBillingControlsSchema } from "@models/cusModels/billingControls/customerBillingControls"; import { AppEnv } from "@models/genModels/genEnums"; import { z } from "zod/v4"; diff --git a/shared/api/customers/cusFeatures/utils/check/getFeatureToUseForCheck.ts b/shared/api/customers/cusFeatures/utils/check/getFeatureToUseForCheck.ts index 2b1ae0904..1fbba4000 100644 --- a/shared/api/customers/cusFeatures/utils/check/getFeatureToUseForCheck.ts +++ b/shared/api/customers/cusFeatures/utils/check/getFeatureToUseForCheck.ts @@ -8,12 +8,12 @@ import { export const getFeatureToUseForCheck = ({ creditSystems, feature, - apiEntity, + apiSubject, requiredBalance, }: { creditSystems: Feature[]; feature: Feature; - apiEntity: ApiCustomerV5 | ApiEntityV2; + apiSubject: ApiCustomerV5 | ApiEntityV2; requiredBalance: number; }) => { // 1. If there's a credit system & cusEnts for that credit system -> return credit system @@ -21,12 +21,13 @@ export const getFeatureToUseForCheck = ({ // 3. Otherwise, feature to use is credit system if exists, otherwise return feature if (creditSystems.length === 0) return feature; - const mainBalance = apiEntity?.balances?.[feature.id]; + const mainBalance = apiSubject.balances?.[feature.id]; if ( mainBalance && apiBalanceToAllowed({ apiBalance: mainBalance, + apiSubject, feature, requiredBalance, }) @@ -35,12 +36,13 @@ export const getFeatureToUseForCheck = ({ } for (const creditSystem of creditSystems) { - const apiBalance = apiEntity?.balances?.[creditSystem.id]; + const apiBalance = apiSubject.balances?.[creditSystem.id]; if (!apiBalance) continue; if ( apiBalanceToAllowed({ apiBalance, + apiSubject, feature: creditSystem, requiredBalance, }) diff --git a/shared/api/customers/cusFeatures/utils/convert/apiBalanceToAllowed.ts b/shared/api/customers/cusFeatures/utils/convert/apiBalanceToAllowed.ts index 2c7aeb425..449d83dc7 100644 --- a/shared/api/customers/cusFeatures/utils/convert/apiBalanceToAllowed.ts +++ b/shared/api/customers/cusFeatures/utils/convert/apiBalanceToAllowed.ts @@ -1,3 +1,4 @@ +import type { ApiSubjectV0 } from "@api/customers/apiSubjectV0"; import type { ApiBalanceV1 } from "@api/customers/cusFeatures/apiBalanceV1"; import { apiBalanceV1ToAvailableOverage } from "@api/customers/cusFeatures/utils/convert/apiBalanceV1ToAvailableOverage"; import type { Feature } from "@models/featureModels/featureModels"; @@ -6,10 +7,12 @@ import { Decimal } from "decimal.js"; export const apiBalanceToAllowed = ({ apiBalance, + apiSubject, feature, requiredBalance, }: { apiBalance: ApiBalanceV1; + apiSubject: ApiSubjectV0; feature: Feature; requiredBalance: number; }) => { @@ -35,7 +38,11 @@ export const apiBalanceToAllowed = ({ // 3. Overage allowed if (apiBalance.overage_allowed) { // 1. Available overage - const availableOverage = apiBalanceV1ToAvailableOverage({ apiBalance }); + const availableOverage = apiBalanceV1ToAvailableOverage({ + apiBalance, + apiSubject, + feature, + }); if (notNullish(availableOverage)) { return new Decimal(availableOverage) diff --git a/shared/api/customers/cusFeatures/utils/convert/apiBalanceV1ToAvailableOverage.ts b/shared/api/customers/cusFeatures/utils/convert/apiBalanceV1ToAvailableOverage.ts index abe2ac1b2..ced0cbd19 100644 --- a/shared/api/customers/cusFeatures/utils/convert/apiBalanceV1ToAvailableOverage.ts +++ b/shared/api/customers/cusFeatures/utils/convert/apiBalanceV1ToAvailableOverage.ts @@ -1,3 +1,7 @@ +import type { ApiSpendLimit } from "@api/billingControls"; +import type { ApiSubjectV0 } from "@api/customers/apiSubjectV0"; +import { apiSubjectToSpendLimit } from "@api/customers/utils/apiSubjectToSpendLimit"; +import type { Feature } from "@models/featureModels/featureModels"; import { sumValues } from "@utils/utils"; import { Decimal } from "decimal.js"; import type { ApiBalanceBreakdownV1, ApiBalanceV1 } from "../../apiBalanceV1"; @@ -37,11 +41,29 @@ export const apiBalanceV1ToMaxOverage = ({ export const apiBalanceV1ToAvailableOverage = ({ apiBalance, + apiSubject, + feature, }: { apiBalance: ApiBalanceV1; + apiSubject: ApiSubjectV0; + feature: Feature; }): number | undefined => { - const maxOverage = apiBalanceV1ToMaxOverage({ apiBalance }); const overage = apiBalanceV1ToOverage({ apiBalance }); + const spendLimit: ApiSpendLimit | undefined = apiSubject + ? apiSubjectToSpendLimit({ + subject: apiSubject, + feature, + }) + : undefined; + + if (spendLimit?.overage_limit !== undefined) { + return Math.max( + 0, + new Decimal(spendLimit.overage_limit).sub(overage).toNumber(), + ); + } + + const maxOverage = apiBalanceV1ToMaxOverage({ apiBalance }); if (maxOverage === undefined) { return undefined; diff --git a/shared/api/customers/index.ts b/shared/api/customers/index.ts index 5434ffcc2..e150e6234 100644 --- a/shared/api/customers/index.ts +++ b/shared/api/customers/index.ts @@ -1,6 +1,7 @@ // Main customer schemas export * from "./apiCustomer"; export * from "./apiCustomerV5"; +export * from "./apiSubjectV0"; export * from "./baseApiCustomer"; // Submodules export * from "./components/index"; diff --git a/shared/api/customers/utils/apiSubjectToSpendLimit.ts b/shared/api/customers/utils/apiSubjectToSpendLimit.ts new file mode 100644 index 000000000..09981045b --- /dev/null +++ b/shared/api/customers/utils/apiSubjectToSpendLimit.ts @@ -0,0 +1,28 @@ +import type { ApiSpendLimit } from "@api/billingControls"; +import type { Feature } from "@models/featureModels/featureModels"; +import type { ApiSubjectV0 } from "../apiSubjectV0"; + +export const apiSubjectToSpendLimit = ({ + subject, + feature, +}: { + subject: ApiSubjectV0; + feature: Feature; +}): ApiSpendLimit | undefined => { + if (!("billing_controls" in subject) || !subject.billing_controls) { + return undefined; + } + + if (!("spend_limits" in subject.billing_controls)) { + return undefined; + } + + const spendLimits = subject.billing_controls.spend_limits ?? []; + + return spendLimits.find( + (spendLimit) => + spendLimit.enabled && + spendLimit.feature_id === feature.id && + spendLimit.overage_limit !== undefined, + ); +}; diff --git a/shared/api/customers/utils/index.ts b/shared/api/customers/utils/index.ts index 8dac49239..fadac01e1 100644 --- a/shared/api/customers/utils/index.ts +++ b/shared/api/customers/utils/index.ts @@ -1 +1,2 @@ export * from "./apiCustomerToFeatures.js"; +export * from "./apiSubjectToSpendLimit.js"; diff --git a/shared/api/entities/apiEntityV2.ts b/shared/api/entities/apiEntityV2.ts index 663b49adc..2f1570fd6 100644 --- a/shared/api/entities/apiEntityV2.ts +++ b/shared/api/entities/apiEntityV2.ts @@ -1,12 +1,12 @@ import { ApiBalanceV1Schema } from "@api/customers/cusFeatures/apiBalanceV1.js"; import { z } from "zod/v4"; +import { ApiEntityBillingControlsSchema } from "../billingControls/entityBillingControls.js"; import { ApiPurchaseV0Schema, ApiSubscriptionV1Schema, } from "../customers/cusPlans/apiSubscriptionV1.js"; import { ApiInvoiceV1Schema } from "../others/apiInvoice/apiInvoiceV1.js"; import { ApiBaseEntitySchema } from "./apiBaseEntity.js"; -import { ApiEntityBillingControlsSchema } from "./billingControls/entityBillingControls.js"; // V2 base entity - uses V1 subscriptions (single array with status field) export const BaseApiEntityV2Schema = ApiBaseEntitySchema.extend({ diff --git a/shared/api/entities/billingControls/entityBillingControls.ts b/shared/api/entities/billingControls/entityBillingControls.ts deleted file mode 100644 index f059732f8..000000000 --- a/shared/api/entities/billingControls/entityBillingControls.ts +++ /dev/null @@ -1,15 +0,0 @@ -import { z } from "zod/v4"; -import { ApiEntitySpendLimitSchema } from "./entitySpendLimit.js"; - -export const ApiEntityBillingControlsSchema = z.object({ - spend_limits: z.array(ApiEntitySpendLimitSchema).optional().meta({ - description: "List of overage spend limits per feature.", - }), -}); - -export type ApiEntityBillingControls = z.infer< - typeof ApiEntityBillingControlsSchema ->; -export type ApiEntityBillingControlsInput = z.input< - typeof ApiEntityBillingControlsSchema ->; diff --git a/shared/api/entities/billingControls/entitySpendLimit.ts b/shared/api/entities/billingControls/entitySpendLimit.ts deleted file mode 100644 index 5025cf60b..000000000 --- a/shared/api/entities/billingControls/entitySpendLimit.ts +++ /dev/null @@ -1,6 +0,0 @@ -import type { z } from "zod/v4"; -import { EntitySpendLimitSchema } from "../../../models/cusModels/billingControls/entitySpendLimit.js"; - -export const ApiEntitySpendLimitSchema = EntitySpendLimitSchema; - -export type ApiEntitySpendLimit = z.infer; diff --git a/shared/api/entities/crud/createEntityParams.ts b/shared/api/entities/crud/createEntityParams.ts index 9dd26d8ec..c665322d0 100644 --- a/shared/api/entities/crud/createEntityParams.ts +++ b/shared/api/entities/crud/createEntityParams.ts @@ -1,6 +1,6 @@ import { z } from "zod/v4"; +import { ApiEntityBillingControlsInputSchema } from "../../billingControls/entityBillingControls.js"; import { CustomerDataSchema } from "../../common/customerData.js"; -import { ApiEntityBillingControlsSchema } from "../billingControls/entityBillingControls.js"; export const CreateEntityParamsV0Schema = z.object({ id: z @@ -18,7 +18,7 @@ export const CreateEntityParamsV0Schema = z.object({ feature_id: z.string().meta({ description: "The ID of the feature this entity is associated with", }), - billing_controls: ApiEntityBillingControlsSchema.optional().meta({ + billing_controls: ApiEntityBillingControlsInputSchema.optional().meta({ description: "Billing controls for the entity.", }), customer_data: CustomerDataSchema.optional().meta({ diff --git a/shared/api/entities/crud/updateEntityParams.ts b/shared/api/entities/crud/updateEntityParams.ts index 5138aa8a3..6589751ee 100644 --- a/shared/api/entities/crud/updateEntityParams.ts +++ b/shared/api/entities/crud/updateEntityParams.ts @@ -1,5 +1,5 @@ import { z } from "zod/v4"; -import { ApiEntityBillingControlsSchema } from "../billingControls/entityBillingControls.js"; +import { ApiEntityBillingControlsInputSchema } from "../../billingControls/entityBillingControls.js"; export const UpdateEntityParamsSchema = z.object({ customer_id: z.string().optional().meta({ @@ -8,7 +8,7 @@ export const UpdateEntityParamsSchema = z.object({ entity_id: z.string().meta({ description: "The ID of the entity.", }), - billing_controls: ApiEntityBillingControlsSchema.optional().meta({ + billing_controls: ApiEntityBillingControlsInputSchema.optional().meta({ description: "Billing controls to replace on the entity.", }), }); diff --git a/shared/api/entities/index.ts b/shared/api/entities/index.ts index d1cc85d43..e1ec38741 100644 --- a/shared/api/entities/index.ts +++ b/shared/api/entities/index.ts @@ -1,2 +1,2 @@ -export * from "./billingControls/index.js"; +export * from "../billingControls/index.js"; export * from "./crud/index.js"; diff --git a/shared/index.ts b/shared/index.ts index 9f9d537f2..e5d30782d 100644 --- a/shared/index.ts +++ b/shared/index.ts @@ -46,9 +46,9 @@ export * from "./models/chatResultModels/chatResultFeature"; // 4. Chat Result Models export * from "./models/chatResultModels/chatResultTable"; export * from "./models/checkModels/checkPreviewModels"; -// 8. Customer Models -export * from "./models/cusModels/billingControlModels"; export * from "./models/cusModels/billingControls/autoTopupLimitTable"; +// 8. Customer Models +export * from "./models/cusModels/billingControls/customerBillingControls"; export * from "./models/cusModels/billingControls/purchaseLimitInterval"; export * from "./models/cusModels/cusModels"; // Processor Models diff --git a/shared/models/cusModels/billingControlModels.ts b/shared/models/cusModels/billingControls/customerBillingControls.ts similarity index 79% rename from shared/models/cusModels/billingControlModels.ts rename to shared/models/cusModels/billingControls/customerBillingControls.ts index 44cfabe59..a9cbe489a 100644 --- a/shared/models/cusModels/billingControlModels.ts +++ b/shared/models/cusModels/billingControls/customerBillingControls.ts @@ -3,12 +3,9 @@ import { type EntityBillingControls, type EntityBillingControlsInput, EntityBillingControlsSchema, -} from "./billingControls/entityBillingControls.js"; -import { - type EntitySpendLimit, - EntitySpendLimitSchema, -} from "./billingControls/entitySpendLimit.js"; -import { PurchaseLimitIntervalEnum } from "./billingControls/purchaseLimitInterval.js"; +} from "./entityBillingControls.js"; +import { PurchaseLimitIntervalEnum } from "./purchaseLimitInterval.js"; +import { type DbSpendLimit, DbSpendLimitSchema } from "./spendLimit.js"; export const AutoTopupPurchaseLimitSchema = z.object({ interval: PurchaseLimitIntervalEnum.meta({ @@ -45,6 +42,9 @@ export const CustomerBillingControlsSchema = z.object({ auto_topups: z.array(AutoTopupSchema).optional().meta({ description: "List of auto top-up configurations per feature.", }), + spend_limits: z.array(DbSpendLimitSchema).optional().meta({ + description: "List of overage spend limits per feature.", + }), }); export type AutoTopupPurchaseLimit = z.infer< @@ -59,9 +59,5 @@ export type CustomerBillingControlsInput = z.input< typeof CustomerBillingControlsSchema >; -export { EntityBillingControlsSchema, EntitySpendLimitSchema }; -export type { - EntityBillingControls, - EntityBillingControlsInput, - EntitySpendLimit, -}; +export { EntityBillingControlsSchema, DbSpendLimitSchema }; +export type { EntityBillingControls, EntityBillingControlsInput, DbSpendLimit }; diff --git a/shared/models/cusModels/billingControls/entityBillingControls.ts b/shared/models/cusModels/billingControls/entityBillingControls.ts index 5f37735f4..66b33419a 100644 --- a/shared/models/cusModels/billingControls/entityBillingControls.ts +++ b/shared/models/cusModels/billingControls/entityBillingControls.ts @@ -1,8 +1,8 @@ import { z } from "zod/v4"; -import { EntitySpendLimitSchema } from "./entitySpendLimit.js"; +import { DbSpendLimitSchema } from "./spendLimit.js"; export const EntityBillingControlsSchema = z.object({ - spend_limits: z.array(EntitySpendLimitSchema).optional().meta({ + spend_limits: z.array(DbSpendLimitSchema).optional().meta({ description: "List of overage spend limits per feature.", }), }); diff --git a/shared/models/cusModels/billingControls/entitySpendLimit.ts b/shared/models/cusModels/billingControls/spendLimit.ts similarity index 85% rename from shared/models/cusModels/billingControls/entitySpendLimit.ts rename to shared/models/cusModels/billingControls/spendLimit.ts index 8d3fc2d69..7e94ee15d 100644 --- a/shared/models/cusModels/billingControls/entitySpendLimit.ts +++ b/shared/models/cusModels/billingControls/spendLimit.ts @@ -1,6 +1,6 @@ import { z } from "zod/v4"; -export const EntitySpendLimitSchema = z +export const DbSpendLimitSchema = z .object({ feature_id: z.string().optional().meta({ description: "Optional feature ID this spend limit applies to.", @@ -26,4 +26,4 @@ export const EntitySpendLimitSchema = z }, ); -export type EntitySpendLimit = z.infer; +export type DbSpendLimit = z.infer; diff --git a/shared/models/cusModels/cusModels.ts b/shared/models/cusModels/cusModels.ts index 89f2f23cf..f51c0b704 100644 --- a/shared/models/cusModels/cusModels.ts +++ b/shared/models/cusModels/cusModels.ts @@ -1,7 +1,10 @@ import { z } from "zod/v4"; import { AppEnv } from "../genModels/genEnums.js"; import { ExternalProcessorsSchema } from "../genModels/processorSchemas.js"; -import { AutoTopupSchema } from "./billingControlModels.js"; +import { + AutoTopupSchema, + DbSpendLimitSchema, +} from "./billingControls/customerBillingControls.js"; export const CustomerSchema = z.object({ id: z.string().nullish(), // given by user @@ -19,6 +22,7 @@ export const CustomerSchema = z.object({ metadata: z.record(z.any(), z.any()).nullish().default({}), send_email_receipts: z.boolean().default(false), auto_topups: z.array(AutoTopupSchema).nullish(), + spend_limits: z.array(DbSpendLimitSchema).nullish(), }); export type Customer = z.infer; diff --git a/shared/models/cusModels/cusTable.ts b/shared/models/cusModels/cusTable.ts index 0d0798f2d..d19f572f9 100644 --- a/shared/models/cusModels/cusTable.ts +++ b/shared/models/cusModels/cusTable.ts @@ -12,7 +12,7 @@ import { import { collatePgColumn } from "../../db/utils.js"; import type { ExternalProcessors } from "../genModels/processorSchemas.js"; import { organizations } from "../orgModels/orgTable.js"; -import type { AutoTopup } from "./billingControlModels.js"; +import type { AutoTopup } from "./billingControls/customerBillingControls.js"; export type CustomerProcessor = { type: "stripe"; diff --git a/shared/models/cusModels/entityModels/entityModels.ts b/shared/models/cusModels/entityModels/entityModels.ts index 4ca03ef7c..f63ffbace 100644 --- a/shared/models/cusModels/entityModels/entityModels.ts +++ b/shared/models/cusModels/entityModels/entityModels.ts @@ -1,6 +1,6 @@ import { z } from "zod/v4"; import type { Feature } from "../../featureModels/featureModels.js"; -import { EntitySpendLimitSchema } from "../billingControlModels.js"; +import { DbSpendLimitSchema } from "../billingControls/customerBillingControls.js"; export const EntitySchema = z.object({ id: z.string().nullable(), @@ -13,7 +13,7 @@ export const EntitySchema = z.object({ deleted: z.boolean(), feature_id: z.string(), internal_feature_id: z.string(), - spend_limits: z.array(EntitySpendLimitSchema).nullish(), + spend_limits: z.array(DbSpendLimitSchema).nullish(), }); // export const CreateEntitySchema = z.object({ diff --git a/shared/models/cusModels/entityModels/entityTable.ts b/shared/models/cusModels/entityModels/entityTable.ts index 4f3f5da0f..fc828a65d 100644 --- a/shared/models/cusModels/entityModels/entityTable.ts +++ b/shared/models/cusModels/entityModels/entityTable.ts @@ -10,7 +10,7 @@ import { } from "drizzle-orm/pg-core"; import { features } from "../../featureModels/featureTable.js"; import { organizations } from "../../orgModels/orgTable.js"; -import type { EntitySpendLimit } from "../billingControlModels.js"; +import type { DbSpendLimit } from "../billingControls/customerBillingControls.js"; import { customers } from "../cusTable.js"; export const entities = pgTable( @@ -25,7 +25,7 @@ export const entities = pgTable( name: text(), deleted: boolean().default(false).notNull(), internal_feature_id: text("internal_feature_id"), - spend_limits: jsonb().$type(), + spend_limits: jsonb().$type(), // Optional... feature_id: text("feature_id"), diff --git a/shared/models/cusModels/index.ts b/shared/models/cusModels/index.ts index e1770ed36..6b06a2daf 100644 --- a/shared/models/cusModels/index.ts +++ b/shared/models/cusModels/index.ts @@ -1,4 +1,4 @@ -export * from "./billingControlModels.js"; +export * from "./billingControls/customerBillingControls.js"; export * from "./billingControls/entityBillingControls.js"; -export * from "./billingControls/entitySpendLimit.js"; export * from "./billingControls/purchaseLimitInterval.js"; +export * from "./billingControls/spendLimit.js";