diff --git a/server/tests/integration/balances/track/basic/track-tokens-limits.test.ts b/server/tests/integration/balances/track/basic/track-tokens-limits.test.ts index 35afcd6dd..824f68709 100644 --- a/server/tests/integration/balances/track/basic/track-tokens-limits.test.ts +++ b/server/tests/integration/balances/track/basic/track-tokens-limits.test.ts @@ -1,18 +1,23 @@ import { expect, test } from "bun:test"; import type { ApiCustomerV5, TrackResponseV3 } from "@autumn/shared"; -import { ErrCode } from "@autumn/shared"; +import { ApiVersion, ErrCode, ResetInterval } from "@autumn/shared"; +import { setCustomerUsageLimit } from "@tests/integration/balances/utils/usage-limit-utils/customerUsageLimitUtils.js"; import { expectBalanceCorrect } from "@tests/integration/utils/expectBalanceCorrect.js"; +import { expectUsageLimitCorrect } from "@tests/integration/utils/expectUsageLimitCorrect.js"; import { TestFeature } from "@tests/setup/v2Features.js"; import { expectAutumnError } from "@tests/utils/expectUtils/expectErrUtils.js"; import { items } from "@tests/utils/fixtures/items.js"; import { products } from "@tests/utils/fixtures/products.js"; import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; import chalk from "chalk"; +import { AutumnInt } from "@/external/autumn/autumnCli.js"; // custom/internal-model: input_cost=5 $/M, output_cost=15 $/M, markup=0% // in=5000/out=2500 -> 0.0625; in=10000/out=5000 -> 0.125 +const autumnV2_3 = new AutumnInt({ version: ApiVersion.V2_3 }); + // ═══════════════════════════════════════════════════════════════════ // TRACK-TOKENS-LIM-1: default behavior caps deduction at zero balance // ═══════════════════════════════════════════════════════════════════ @@ -252,3 +257,63 @@ test.concurrent( }); }, ); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-LIM-6: a usage limit counts AI credits consumed, not calls +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-lim-6: usage limit counts AI credits consumed, not call count")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, + }); + const freeProd = products.base({ id: "free", items: [aiCreditsItem] }); + + const { customerId } = await initScenario({ + customerId: "track-tokens-lim-6", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + // Cap AI-credit spend at 100 credits/day — well above one call's cost. + await setCustomerUsageLimit({ + autumn: autumnV2_3, + customerId, + featureId: TestFeature.AiCredits, + limit: 100, + interval: ResetInterval.Day, + }); + + // One token track costs 0.125 credits. + await autumnV2_3.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 10000, + output_tokens: 5000, + }); + + const customer = await autumnV2_3.customers.get(customerId); + // The balance is deducted by the dollar cost... + expectBalanceCorrect({ + customer, + featureId: TestFeature.AiCredits, + remaining: 999.875, + usage: 0.125, + }); + // ...and the cap counter must reflect that SAME spend. Regression: before + // the usage-window dimension fix this counted 1 (the call), not 0.125. + expectUsageLimitCorrect({ + customer, + featureId: TestFeature.AiCredits, + usage: 0.125, + limit: 100, + interval: ResetInterval.Day, + }); + }, +); diff --git a/server/tests/unit/usage-windows/fullSubjectToUsageWindowLimits.test.ts b/server/tests/unit/usage-windows/fullSubjectToUsageWindowLimits.test.ts index 30dc545da..d479f24f5 100644 --- a/server/tests/unit/usage-windows/fullSubjectToUsageWindowLimits.test.ts +++ b/server/tests/unit/usage-windows/fullSubjectToUsageWindowLimits.test.ts @@ -25,6 +25,11 @@ const creditsFeature = { internal_id: "icredits", type: FeatureType.CreditSystem, } as Feature; +const aiCreditsFeature = { + id: "ai_credits", + internal_id: "iai_credits", + type: FeatureType.AiCreditSystem, +} as Feature; // Credit system whose schema contains action1, for the membership-anchor path. const creditsContainingAction1 = { id: "credits", @@ -210,6 +215,25 @@ describe("fullSubjectToUsageWindowLimits", () => { }); }); + test("an ai-credit-system feature resolves to the balance dimension", () => { + const limits = fullSubjectToUsageWindowLimits({ + fullSubject: buildSubject({ + usageLimits: [ + { feature_id: "ai_credits", limit: 100, interval: ResetInterval.Day }, + ], + }), + featureIds: ["ai_credits"], + features: [aiCreditsFeature], + now: NOW, + }); + + expect(limits).toHaveLength(1); + expect(limits[0]).toMatchObject({ + dimension_type: "balance", + dimension_feature_id: null, + }); + }); + test("the window interval comes from the entry, independent of the entitlement's reset interval", () => { const limits = fullSubjectToUsageWindowLimits({ fullSubject: buildSubject({ diff --git a/shared/utils/usageWindowUtils/convertUsageWindow/getUsageWindowDimension.ts b/shared/utils/usageWindowUtils/convertUsageWindow/getUsageWindowDimension.ts index e1ebba4d5..f546ac62f 100644 --- a/shared/utils/usageWindowUtils/convertUsageWindow/getUsageWindowDimension.ts +++ b/shared/utils/usageWindowUtils/convertUsageWindow/getUsageWindowDimension.ts @@ -1,11 +1,11 @@ import type { UsageWindowDimension } from "../../../models/cusProductModels/cusEntModels/usageWindowModels.js"; -import { FeatureType } from "../../../models/featureModels/featureEnums.js"; import type { Feature } from "../../../models/featureModels/featureModels.js"; +import { isAnyCreditSystem } from "../../featureUtils/classifyFeature/isAnyCreditSystem.js"; /** * Which dimension a usage limit on `feature` counts against: - * - a credit-system feature caps the credit POOL (`balance` dimension, - * counted in credits drained); + * - any credit-system feature (classic or AI token-based) caps the credit + * POOL (`balance` dimension, counted in credits drained); * - any other feature caps that feature's own usage (`metered_feature` * dimension, counted in tracked units). */ @@ -17,7 +17,7 @@ export const getUsageWindowDimension = ({ dimensionType: UsageWindowDimension; dimensionFeatureId: string | null; } => { - const isCreditSystem = feature.type === FeatureType.CreditSystem; + const isCreditSystem = isAnyCreditSystem(feature.type); return { dimensionType: isCreditSystem ? "balance" : "metered_feature",