diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/contextUtilsV2.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/contextUtilsV2.lua index 8e30a8a22..25b75b3e0 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/contextUtilsV2.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/contextUtilsV2.lua @@ -24,6 +24,7 @@ local function init_context(params) env = params.env, customer_id = params.customer_id, customer_entitlement_deductions = params.customer_entitlement_deductions, + anchor_entitlements = params.anchor_entitlements, balance_keys_by_feature_id = params.balance_keys_by_feature_id, }) @@ -96,6 +97,23 @@ local function init_context(params) end end + -- Register usage-window anchor cus_ents that are not in the deduction set so + -- their counter can be read/mutated and persisted (HSET) on apply. + for customer_entitlement_id, balance_entry in pairs(read_result.balances_by_id) do + if balance_entry.anchor_only + and is_nil(context.customer_entitlements[customer_entitlement_id]) + then + context.customer_entitlements[customer_entitlement_id] = { + base_path = customer_entitlement_id, + balance_key = balance_entry.balance_key, + subject_balance = balance_entry.subject_balance, + customer_entitlement_id = customer_entitlement_id, + feature_id = balance_entry.feature_id, + is_anchor_only = true, + } + end + end + return context end diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromSubjectBalances.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromSubjectBalances.lua index ede381166..9d3abc3c7 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromSubjectBalances.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromSubjectBalances.lua @@ -111,6 +111,30 @@ local idempotency_ttl_ms = params.idempotency_ttl_ms local lock = params.lock local unwind_value = params.unwind_value local lock_receipt_key = lock_receipt_key_from_keys +local usage_window_limits = params.usage_window_limits +local usage_window_now = params.usage_window_now +local is_consumption = params.is_consumption + +-- Distinct usage-window anchor cus_ents to force-load into context (they own +-- the counters and may not be in the deduction set). +local anchor_entitlements = {} +if not is_nil(usage_window_limits) then + local seen_anchor_ids = {} + for _, usage_window_limit in ipairs(usage_window_limits) do + local anchor_id = usage_window_limit.anchor_customer_entitlement_id + local anchor_feature_id = usage_window_limit.anchor_feature_id + if not is_nil(anchor_id) + and not is_nil(anchor_feature_id) + and not seen_anchor_ids[anchor_id] + then + seen_anchor_ids[anchor_id] = true + table.insert(anchor_entitlements, { + customer_entitlement_id = anchor_id, + feature_id = anchor_feature_id, + }) + end + end +end if not is_nil(idempotency_key) then if redis.call('EXISTS', idempotency_key) == 1 then @@ -144,6 +168,7 @@ local context = init_context({ env = env, customer_id = customer_id, customer_entitlement_deductions = customer_entitlement_deductions, + anchor_entitlements = anchor_entitlements, balance_keys_by_feature_id = params.balance_keys_by_feature_id, debug = params.debug, }) @@ -235,11 +260,6 @@ for _, cus_ent_id in ipairs(unwind_modified_cus_ent_ids) do end end -local modified_customer_entitlement_ids = collect_modified_customer_entitlement_ids({ - context = context, - extra_customer_entitlement_ids = unwind_modified_cus_ent_ids, -}) - logger.log(" remaining_amount: %s", tostring(remaining_amount or "nil")) logger.log(" is_refund: %s", tostring(remaining_amount < 0 or false)) local mutation_logs = context.mutation_logs @@ -259,6 +279,54 @@ if remaining_amount > 0 and overage_behaviour == 'reject' then }) end +-- Hard windowed usage-limit enforcement, on ACTUAL consumed amounts, before any +-- writes. Only for positive consumption (refunds / target_balance / granted +-- balance edits never trip or move counters). v1 also excludes lock-based and +-- unwind flows: counter reversal on partial unwind is not implemented yet, so +-- enforcing there could drift the counter. +local enforce_usage_windows = is_consumption + and is_nil(unwind_value) + and (is_nil(lock) or not lock.enabled) + and not is_nil(usage_window_limits) + and #usage_window_limits > 0 + +if enforce_usage_windows then + local exceeded_feature_id = check_usage_window_limits({ + context = context, + usage_window_limits = usage_window_limits, + updates = updates, + amount_to_deduct = amount_to_deduct, + remaining_amount = remaining_amount, + }) + + if not is_nil(exceeded_feature_id) then + return cjson.encode({ + error = 'USAGE_LIMIT_EXCEEDED', + feature_id = exceeded_feature_id, + remaining = remaining_amount, + updates = {}, + rollover_updates = {}, + modified_customer_entitlement_ids = new_empty_array(), + mutation_logs = mutation_logs, + logs = context.logs, + }) + end + + increment_usage_window_counters({ + context = context, + usage_window_limits = usage_window_limits, + updates = updates, + amount_to_deduct = amount_to_deduct, + remaining_amount = remaining_amount, + now = usage_window_now, + }) +end + +local modified_customer_entitlement_ids = collect_modified_customer_entitlement_ids({ + context = context, + extra_customer_entitlement_ids = unwind_modified_cus_ent_ids, +}) + if not is_nil(lock) and not is_nil(lock.enabled) and lock.enabled diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/readSubjectBalances.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/readSubjectBalances.lua index d4e485815..b118c39b5 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/readSubjectBalances.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/readSubjectBalances.lua @@ -18,29 +18,61 @@ local function read_subject_balances(params) local balances_by_id = {} local missing_customer_entitlement_ids = {} local entries_by_balance_key = {} + local seen_ids_by_balance_key = {} + -- Anchor-only cus_ents (usage-window owners not in the deduction set) must + -- not abort the deduction when absent; a missing anchor fails closed later. + local anchor_only_ids = {} local balance_keys_by_feature_id = safe_table(params.balance_keys_by_feature_id) + local function queue_balance_read(customer_entitlement_id, feature_id) + if not (customer_entitlement_id and feature_id) then + return false + end + + local balance_key = balance_keys_by_feature_id[feature_id] + if not balance_key then + return false + end + + if entries_by_balance_key[balance_key] == nil then + entries_by_balance_key[balance_key] = { + feature_id = feature_id, + customer_entitlement_ids = {}, + } + seen_ids_by_balance_key[balance_key] = {} + end + + if seen_ids_by_balance_key[balance_key][customer_entitlement_id] then + return true + end + seen_ids_by_balance_key[balance_key][customer_entitlement_id] = true + + table.insert( + entries_by_balance_key[balance_key].customer_entitlement_ids, + customer_entitlement_id + ) + return true + end + + local deduction_ids = {} for _, ent_obj in ipairs(params.customer_entitlement_deductions or {}) do local customer_entitlement_id = ent_obj.customer_entitlement_id - local feature_id = ent_obj.feature_id - - if customer_entitlement_id and feature_id then - local balance_key = balance_keys_by_feature_id[feature_id] - if not balance_key then + if customer_entitlement_id then + local queued = queue_balance_read(customer_entitlement_id, ent_obj.feature_id) + if not queued then table.insert(missing_customer_entitlement_ids, customer_entitlement_id) - else - if entries_by_balance_key[balance_key] == nil then - entries_by_balance_key[balance_key] = { - feature_id = feature_id, - customer_entitlement_ids = {}, - } - end - - table.insert( - entries_by_balance_key[balance_key].customer_entitlement_ids, - customer_entitlement_id - ) end + deduction_ids[customer_entitlement_id] = true + end + end + + -- A deduction target must never be marked anchor_only, or a cache miss on it is + -- swallowed instead of surfacing as missing + triggering the Postgres fallback. + for _, anchor in ipairs(params.anchor_entitlements or {}) do + local customer_entitlement_id = anchor.customer_entitlement_id + if customer_entitlement_id and not deduction_ids[customer_entitlement_id] then + anchor_only_ids[customer_entitlement_id] = true + queue_balance_read(customer_entitlement_id, anchor.feature_id) end end @@ -57,16 +89,19 @@ local function read_subject_balances(params) local subject_balance = decode_subject_balance(raw_value) if subject_balance == nil then - table.insert( - missing_customer_entitlement_ids, - customer_entitlement_id - ) + if not anchor_only_ids[customer_entitlement_id] then + table.insert( + missing_customer_entitlement_ids, + customer_entitlement_id + ) + end else balances_by_id[customer_entitlement_id] = { balance_key = balance_key, customer_entitlement_id = customer_entitlement_id, feature_id = entry.feature_id, subject_balance = subject_balance, + anchor_only = anchor_only_ids[customer_entitlement_id] or nil, } end end diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/runDeductionOnContextV2.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/runDeductionOnContextV2.lua index e092ee3a3..5e9901b33 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/runDeductionOnContextV2.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/runDeductionOnContextV2.lua @@ -274,6 +274,11 @@ local function run_deduction_on_context(params) update.adjustment = ent_data.adjustment or 0 update.additional_balance = 0 + + if ent_data.subject_balance + and type(ent_data.subject_balance.usage_windows) == 'table' then + update.usage_windows = ent_data.subject_balance.usage_windows + end end end diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/usageWindowUtilsV2.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/usageWindowUtilsV2.lua new file mode 100644 index 000000000..9d7b003fc --- /dev/null +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/usageWindowUtilsV2.lua @@ -0,0 +1,163 @@ +-- ============================================================================ +-- USAGE WINDOW UTILITIES (V2) +-- Hard windowed usage-limit enforcement, evaluated at the orchestration layer +-- against ACTUAL consumed amounts (post-deduction, pre-write). +-- +-- Counters live inline on the anchor cus_ent's subject_balance.usage_windows, +-- keyed by the deterministic window key built in TS. The window key includes +-- window_start_at, so the current window's counter is found-or-created at +-- limit.key and a rolled window is simply a different (absent) key. +-- ============================================================================ + +-- Tolerance for float drift (credit-ratio conversions leave sub-nano noise). +local USAGE_WINDOW_EPSILON = 1e-9 + +local function get_anchor_usage_windows(context, anchor_customer_entitlement_id) + if is_nil(anchor_customer_entitlement_id) then + return nil + end + + local ent_data = context.customer_entitlements[anchor_customer_entitlement_id] + if not ent_data or not ent_data.subject_balance then + return nil + end + + if type(ent_data.subject_balance.usage_windows) ~= 'table' then + ent_data.subject_balance.usage_windows = {} + end + + return ent_data.subject_balance.usage_windows +end + +-- Actually-consumed amount for a limit, in its native unit. metered_feature +-- counts feature units (the tracked total); balance counts credits drained from +-- the anchor pool (its `deducted`, which is in credits). +local function usage_window_consumed(params) + local limit = params.limit + local updates = params.updates + + if limit.dimension_type == 'balance' then + local anchor_update = updates[limit.anchor_customer_entitlement_id] + return anchor_update and safe_number(anchor_update.deducted) or 0 + end + + return safe_number(params.amount_to_deduct) - safe_number(params.remaining_amount) +end + +-- Returns the feature_id of the first limit that would be exceeded (so the +-- caller can hard-reject), or nil if every limit has room. Null/missing anchor +-- fails closed: a cap that cannot resolve an owner must not silently allow. +local function check_usage_window_limits(params) + local context = params.context + local limits = params.usage_window_limits or {} + + for _, limit in ipairs(limits) do + local windows = get_anchor_usage_windows( + context, + limit.anchor_customer_entitlement_id + ) + if is_nil(windows) then + return limit.feature_id + end + + local consumed = usage_window_consumed({ + limit = limit, + updates = params.updates, + amount_to_deduct = params.amount_to_deduct, + remaining_amount = params.remaining_amount, + }) + + if consumed > USAGE_WINDOW_EPSILON then + local existing = windows[limit.key] + local current_usage = existing and safe_number(existing.usage_amount) or 0 + if current_usage + consumed + > safe_number(limit.limit) + USAGE_WINDOW_EPSILON then + return limit.feature_id + end + end + end + + return nil +end + +-- Applies the consumed amount to each anchor counter (find-or-create at the +-- current window key), prunes closed sibling windows, and marks the anchor dirty +-- so apply_pending_writes persists it (even when the anchor's balance did not +-- change, or when only a prune happened). +local function increment_usage_window_counters(params) + local context = params.context + local limits = params.usage_window_limits or {} + local now = params.now + + for _, limit in ipairs(limits) do + local windows = get_anchor_usage_windows( + context, + limit.anchor_customer_entitlement_id + ) + if not is_nil(windows) then + -- Prune closed windows every pass (not only when consuming) so the map + -- does not grow for sporadically-active features (lifetime never closes). + local pruned = false + for window_key, window in pairs(windows) do + if window_key ~= limit.key + and type(window) == 'table' + and safe_number(window.window_end_at) < now + then + windows[window_key] = nil + pruned = true + end + end + + local consumed = usage_window_consumed({ + limit = limit, + updates = params.updates, + amount_to_deduct = params.amount_to_deduct, + remaining_amount = params.remaining_amount, + }) + + if consumed > USAGE_WINDOW_EPSILON then + -- balance_amount is audit-only; for metered caps it captures just the + -- anchor pool's credits, not every pool the track touched. + local consumed_credits = 0 + local anchor_update = params.updates[limit.anchor_customer_entitlement_id] + if anchor_update then + consumed_credits = safe_number(anchor_update.deducted) + end + + local existing = windows[limit.key] + if is_nil(existing) then + existing = { + key = limit.key, + dimension_type = limit.dimension_type, + dimension_feature_id = limit.dimension_feature_id or cjson.null, + scope_type = limit.scope_type, + entity_id = limit.entity_id or cjson.null, + internal_entity_id = limit.internal_entity_id or cjson.null, + interval = limit.interval, + window_start_at = limit.window_start_at, + window_end_at = limit.window_end_at, + usage_amount = 0, + balance_amount = 0, + } + windows[limit.key] = existing + end + + existing.usage_amount = safe_number(existing.usage_amount) + consumed + existing.balance_amount = + safe_number(existing.balance_amount) + consumed_credits + existing.limit_snapshot = safe_number(limit.limit) + existing.updated_at = now + + mark_customer_entitlement_for_update( + context, + limit.anchor_customer_entitlement_id + ) + elseif pruned then + mark_customer_entitlement_for_update( + context, + limit.anchor_customer_entitlement_id + ) + end + end + end +end diff --git a/server/src/_luaScriptsV2/luaScriptsV2.ts b/server/src/_luaScriptsV2/luaScriptsV2.ts index e08cb37bd..93d91c7b3 100644 --- a/server/src/_luaScriptsV2/luaScriptsV2.ts +++ b/server/src/_luaScriptsV2/luaScriptsV2.ts @@ -44,6 +44,7 @@ import READ_SUBJECT_BALANCES from "./fullSubjectDeduction/readSubjectBalances.lu import RUN_DEDUCTION_ON_CONTEXT_V2 from "./fullSubjectDeduction/runDeductionOnContextV2.lua"; import SPEND_LIMIT_UTILS_V2 from "./fullSubjectDeduction/spendLimitUtilsV2.lua"; import UPDATE_AGGREGATED_BALANCES from "./fullSubjectDeduction/updateAggregatedBalances.lua"; +import USAGE_WINDOW_UTILS_V2 from "./fullSubjectDeduction/usageWindowUtilsV2.lua"; // ============================================================================ // UPDATE SUBJECT BALANCES HELPERS (V2 cache — per-feature hash updates) @@ -205,6 +206,7 @@ ${GET_TOTAL_BALANCE} ${DEDUCT_FROM_ROLLOVERS_V2} ${DEDUCT_FROM_MAIN_BALANCE_V2} ${SPEND_LIMIT_UTILS_V2} +${USAGE_WINDOW_UTILS_V2} ${RUN_DEDUCTION_ON_CONTEXT_V2} ${MUTATION_ITEM_UTILS} ${LOCK_RECEIPT_UTILS_V2} diff --git a/server/src/internal/balances/check/runCheckWithTrackV2.ts b/server/src/internal/balances/check/runCheckWithTrackV2.ts index 331194177..13f20c567 100644 --- a/server/src/internal/balances/check/runCheckWithTrackV2.ts +++ b/server/src/internal/balances/check/runCheckWithTrackV2.ts @@ -10,6 +10,7 @@ import { type ParsedCheckParams, RecaseError, type TrackParams, + UsageLimitExceededError, } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { getTrackFeatureDeductions } from "@/internal/balances/track/utils/getFeatureDeductions.js"; @@ -94,7 +95,10 @@ export const runCheckWithTrackV2 = async ({ checkData.evaluationApiBalance = trackedBalance ?? undefined; trackBalances = response.balances; } catch (error) { - if (error instanceof InsufficientBalanceError) { + if ( + error instanceof InsufficientBalanceError || + error instanceof UsageLimitExceededError + ) { allowed = false; } else { throw error; diff --git a/server/src/internal/balances/track/v3/handleRedisTrackErrorV3.ts b/server/src/internal/balances/track/v3/handleRedisTrackErrorV3.ts index 5df489ac0..62b79af28 100644 --- a/server/src/internal/balances/track/v3/handleRedisTrackErrorV3.ts +++ b/server/src/internal/balances/track/v3/handleRedisTrackErrorV3.ts @@ -5,9 +5,10 @@ import { RecaseError, type TrackParams, type TrackResponseV3, + UsageLimitExceededError, } from "@autumn/shared"; -import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { RedisUnavailableError } from "@/external/redis/utils/errors.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; import { RedisDeductionError, @@ -39,6 +40,12 @@ export const handleRedisTrackErrorV3 = async ({ }); } + if (error.code === RedisDeductionErrorCode.UsageLimitExceeded) { + throw new UsageLimitExceededError({ + featureId: error.featureId ?? body.feature_id, + }); + } + if (error.code === RedisDeductionErrorCode.LockAlreadyExists) { throw new RecaseError({ message: "A lock with this ID already exists", diff --git a/server/src/internal/balances/utils/deductionV2/applyDeductionUpdateToFullSubject.ts b/server/src/internal/balances/utils/deductionV2/applyDeductionUpdateToFullSubject.ts index 4a4fe11a7..26de4db00 100644 --- a/server/src/internal/balances/utils/deductionV2/applyDeductionUpdateToFullSubject.ts +++ b/server/src/internal/balances/utils/deductionV2/applyDeductionUpdateToFullSubject.ts @@ -45,6 +45,7 @@ const applyUpdate = ({ additional_balance: update.additional_balance, adjustment: update.adjustment, entities: update.entities, + usage_windows: update.usage_windows ?? customerEntitlement.usage_windows, replaceables: getUpdatedReplaceables({ replaceables: customerEntitlement.replaceables, update, diff --git a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts index ac26b243b..c0c67a96c 100644 --- a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts +++ b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts @@ -272,6 +272,7 @@ export const executeRedisDeductionV2 = async ({ throw new RedisDeductionError({ message: `Redis deduction failed: ${resultJson.error}`, code: resultJson.error as RedisDeductionErrorCode, + featureId: resultJson.feature_id, }); } diff --git a/server/src/internal/balances/utils/types/deductionUpdate.ts b/server/src/internal/balances/utils/types/deductionUpdate.ts index d06d00dee..47e3d107c 100644 --- a/server/src/internal/balances/utils/types/deductionUpdate.ts +++ b/server/src/internal/balances/utils/types/deductionUpdate.ts @@ -2,6 +2,7 @@ import type { EntityBalance, InsertReplaceable, Replaceable, + UsageWindows, } from "@autumn/shared"; export interface DeductionUpdate { @@ -14,6 +15,7 @@ export interface DeductionUpdate { additional_deducted?: number; newReplaceables?: InsertReplaceable[]; deletedReplaceables?: Replaceable[]; + usage_windows?: UsageWindows | null; } export type DeductionUpdates = Record; diff --git a/server/src/internal/balances/utils/types/redisDeductionError.ts b/server/src/internal/balances/utils/types/redisDeductionError.ts index 2487c23cc..83449fbbb 100644 --- a/server/src/internal/balances/utils/types/redisDeductionError.ts +++ b/server/src/internal/balances/utils/types/redisDeductionError.ts @@ -9,6 +9,7 @@ export enum RedisDeductionErrorCode { SkipCache = "SKIP_CACHE", LockAlreadyExists = "LOCK_ALREADY_EXISTS", DuplicateIdempotencyKey = "DUPLICATE_IDEMPOTENCY_KEY", + UsageLimitExceeded = "USAGE_LIMIT_EXCEEDED", } /** Errors that should trigger a fallback to Postgres */ @@ -23,17 +24,21 @@ export const FALLBACK_ERROR_CODES = [ /** Error thrown by Redis deduction operations */ export class RedisDeductionError extends Error { code: RedisDeductionErrorCode; + featureId?: string; constructor({ message, code, + featureId, }: { message: string; code: RedisDeductionErrorCode; + featureId?: string; }) { super(message); this.name = "RedisDeductionError"; this.code = code; + this.featureId = featureId; } isRedisUnavailable(): boolean { diff --git a/server/tests/integration/balances/track/usage-limit/track-customer-usage-limit.test.ts b/server/tests/integration/balances/track/usage-limit/track-customer-usage-limit.test.ts new file mode 100644 index 000000000..601ab0c4f --- /dev/null +++ b/server/tests/integration/balances/track/usage-limit/track-customer-usage-limit.test.ts @@ -0,0 +1,319 @@ +import { expect, test } from "bun:test"; +import { type CustomerBillingControls, EntInterval } 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"; + +type AutumnV2_1Client = Awaited>["autumnV2_1"]; + +// Arms a windowed usage cap via spend_limits[].usage_limit (overage off); +// `interval` sets the explicit window override. +const setCustomerUsageLimit = async ({ + autumn, + customerId, + featureId, + limit, + interval = EntInterval.Month, +}: { + autumn: AutumnV2_1Client; + customerId: string; + featureId: string; + limit: number; + interval?: EntInterval; +}) => { + const billingControls: CustomerBillingControls = { + spend_limits: [ + { + feature_id: featureId, + enabled: false, + usage_limit: limit, + usage_limit_interval: interval, + }, + ], + }; + + await timeout(2000); + await autumn.customers.update(customerId, { + billing_controls: billingControls, + }); + await timeout(3000); +}; + +// Credit system: 100 credits, 1 action1 = 0.2 credits (see v2Features.ts). +// A cap of 5 action1 units consumes only 1 credit, so the cap must block the +// 6th unit while ~99 credits remain, proving it's a second, independent +// dimension, not a balance check. +test.concurrent( + `${chalk.yellowBright("track-customer-usage-limit1: per-feature cap blocks deduction while credits remain")}`, + async () => { + const customerProduct = products.base({ + id: "track-customer-usage-limit", + items: [items.monthlyCredits({ includedUsage: 100 })], + }); + + const customerId = `track-customer-usage-limit-1-${Date.now()}`; + const { autumnV2_1 } = await initScenario({ + customerId, + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [customerProduct] }), + ], + actions: [s.billing.attach({ productId: customerProduct.id })], + }); + + await setCustomerUsageLimit({ + autumn: autumnV2_1, + customerId, + featureId: TestFeature.Action1, + limit: 5, + }); + + // Consume exactly up to the cap: 5 action1 units = 1 credit deducted. Assert + // the synchronous track response; a re-read races the async write-through. + const consumed = await autumnV2_1.track({ + customer_id: customerId, + feature_id: TestFeature.Action1, + value: 5, + }); + expect(consumed.balances?.[TestFeature.Credits]).toMatchObject({ + feature_id: TestFeature.Credits, + granted: 100, + remaining: 99, + usage: 1, + }); + + // The 6th unit exceeds the cap. It must be hard-blocked BEFORE any + // deduction, even though ~99 credits remain. + let blocked = false; + let blockedCode: string | undefined; + try { + await autumnV2_1.track({ + customer_id: customerId, + feature_id: TestFeature.Action1, + value: 1, + }); + } catch (error) { + blocked = true; + blockedCode = (error as { code?: string }).code; + } + + expect(blocked).toBe(true); + // 400 not 429: clients flatten a 429 to a generic rate_limit_exceeded. + expect(blockedCode).toBe("usage_limit_exceeded"); + }, +); + +test.concurrent( + `${chalk.yellowBright("track-customer-usage-limit2: credit-pool sub-interval cap (1 credit/day) blocks while monthly credits remain")}`, + async () => { + const customerProduct = products.base({ + id: "track-customer-credit-day-cap", + items: [items.monthlyCredits({ includedUsage: 100 })], + }); + + const customerId = `track-customer-credit-day-cap-1-${Date.now()}`; + const { autumnV2_1 } = await initScenario({ + customerId, + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [customerProduct] }), + ], + actions: [s.billing.attach({ productId: customerProduct.id })], + }); + + // 1 action1 = 0.2 credits, so 5 action1 = exactly 1 credit (the daily cap). + await setCustomerUsageLimit({ + autumn: autumnV2_1, + customerId, + featureId: TestFeature.Credits, + limit: 1, + interval: EntInterval.Day, + }); + + const consumed = await autumnV2_1.track({ + customer_id: customerId, + feature_id: TestFeature.Action1, + value: 5, + }); + expect(consumed.balances?.[TestFeature.Credits]).toMatchObject({ + feature_id: TestFeature.Credits, + granted: 100, + remaining: 99, + usage: 1, + }); + + let blocked = false; + let blockedCode: string | undefined; + try { + await autumnV2_1.track({ + customer_id: customerId, + feature_id: TestFeature.Action1, + value: 1, + }); + } catch (error) { + blocked = true; + blockedCode = (error as { code?: string }).code; + } + + expect(blocked).toBe(true); + expect(blockedCode).toBe("usage_limit_exceeded"); + }, +); + +// set_usage must be rejected when the feature has an enforced usage window; +// otherwise it bypasses the hard cap (it carries no window provenance). +test.concurrent( + `${chalk.yellowBright("track-customer-usage-limit3: set_usage is rejected when the feature has a usage window")}`, + async () => { + const customerProduct = products.base({ + id: "track-customer-setusage-guard", + items: [items.monthlyMessages({ includedUsage: 100 })], + }); + + const customerId = `track-customer-setusage-guard-1-${Date.now()}`; + const { autumnV2_1 } = await initScenario({ + customerId, + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [customerProduct] }), + ], + actions: [s.billing.attach({ productId: customerProduct.id })], + }); + + await setCustomerUsageLimit({ + autumn: autumnV2_1, + customerId, + featureId: TestFeature.Messages, + limit: 5, + }); + + let blockedCode: string | undefined; + try { + await autumnV2_1.balances.update({ + customer_id: customerId, + feature_id: TestFeature.Messages, + current_balance: 50, + }); + } catch (error) { + blockedCode = (error as { code?: string }).code; + } + + expect(blockedCode).toBe("set_usage_not_allowed_with_usage_limit"); + }, +); + +// A single spend_limit entry carrying BOTH an overage_limit and a windowed usage +// cap must still enforce the window (the two caps are independent). +test.concurrent( + `${chalk.yellowBright("track-customer-usage-limit4: a spend_limit with both overage_limit and a usage window still enforces the window")}`, + async () => { + const customerProduct = products.base({ + id: "track-customer-compound-cap", + items: [items.monthlyCredits({ includedUsage: 100 })], + }); + + const customerId = `track-customer-compound-cap-1-${Date.now()}`; + const { autumnV2_1 } = await initScenario({ + customerId, + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [customerProduct] }), + ], + actions: [s.billing.attach({ productId: customerProduct.id })], + }); + + const billingControls: CustomerBillingControls = { + spend_limits: [ + { + feature_id: TestFeature.Action1, + enabled: true, + overage_limit: 20, + usage_limit: 5, + usage_limit_interval: EntInterval.Month, + }, + ], + }; + await timeout(2000); + await autumnV2_1.customers.update(customerId, { + billing_controls: billingControls, + }); + await timeout(3000); + + await autumnV2_1.track({ + customer_id: customerId, + feature_id: TestFeature.Action1, + value: 5, + }); + + let blockedCode: string | undefined; + try { + await autumnV2_1.track({ + customer_id: customerId, + feature_id: TestFeature.Action1, + value: 1, + }); + } catch (error) { + blockedCode = (error as { code?: string }).code; + } + + expect(blockedCode).toBe("usage_limit_exceeded"); + }, +); + +// Two concurrent tracks on the SAME customer's SAME window must serialize (Redis +// runs each deduction Lua atomically): combined value exceeds the cap, so exactly +// one succeeds and one is rejected, and the counter reflects only the winner. +test.concurrent( + `${chalk.yellowBright("track-customer-usage-limit6: concurrent tracks on one window serialize, one rejected")}`, + async () => { + const customerProduct = products.base({ + id: "track-customer-concurrent-cap", + items: [items.monthlyCredits({ includedUsage: 100 })], + }); + + const customerId = `track-customer-concurrent-cap-1-${Date.now()}`; + const { autumnV2_1 } = await initScenario({ + customerId, + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [customerProduct] }), + ], + actions: [s.billing.attach({ productId: customerProduct.id })], + }); + + // Cap action1 at 5/month; two concurrent tracks of 5 each => combined 10 > 5. + await setCustomerUsageLimit({ + autumn: autumnV2_1, + customerId, + featureId: TestFeature.Action1, + limit: 5, + }); + + const results = await Promise.allSettled([ + autumnV2_1.track({ + customer_id: customerId, + feature_id: TestFeature.Action1, + value: 5, + }), + autumnV2_1.track({ + customer_id: customerId, + feature_id: TestFeature.Action1, + value: 5, + }), + ]); + + const fulfilled = results.filter((result) => result.status === "fulfilled"); + const rejected = results.filter( + (result): result is PromiseRejectedResult => result.status === "rejected", + ); + + expect(fulfilled).toHaveLength(1); + expect(rejected).toHaveLength(1); + expect((rejected[0].reason as { code?: string }).code).toBe( + "usage_limit_exceeded", + ); + }, +); diff --git a/shared/api/errors/classes/balancesErrClasses.ts b/shared/api/errors/classes/balancesErrClasses.ts index e768d9af1..86c46ad7f 100644 --- a/shared/api/errors/classes/balancesErrClasses.ts +++ b/shared/api/errors/classes/balancesErrClasses.ts @@ -18,3 +18,22 @@ export class InsufficientBalanceError extends RecaseError { this.name = "InsufficientBalanceError"; } } + +export class UsageLimitExceededError extends RecaseError { + constructor(opts?: { + message?: string; + featureId?: string; + limit?: number; + }) { + super({ + message: + opts?.message || + `Usage limit exceeded${opts?.featureId ? ` for feature ${opts.featureId}` : ""}${opts?.limit !== undefined ? ` (limit ${opts.limit})` : ""}`, + code: BalancesErrorCode.UsageLimitExceeded, + // 400 (mirrors InsufficientBalanceError): clients flatten any 429 to a + // generic rate-limit error, which would hide the usage_limit_exceeded code. + statusCode: 400, + }); + this.name = "UsageLimitExceededError"; + } +} diff --git a/shared/api/errors/codes/balancesErrCodes.ts b/shared/api/errors/codes/balancesErrCodes.ts index 2a05243d0..1bfde8070 100644 --- a/shared/api/errors/codes/balancesErrCodes.ts +++ b/shared/api/errors/codes/balancesErrCodes.ts @@ -1,5 +1,6 @@ export const BalancesErrorCode = { InsufficientBalance: "insufficient_balance", + UsageLimitExceeded: "usage_limit_exceeded", } as const; export type BalancesErrorCode =