From f25838d21c2aef8479da5e62f02cd21146b47f9e Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 10 Jun 2026 13:05:45 +0100 Subject: [PATCH] chore: improve types and atmn --- packages/atmn/src/commands/push/push.ts | 26 ++++ .../src/lib/transforms/apiToSdk/feature.ts | 2 + .../src/lib/transforms/sdkToApi/feature.ts | 32 +++-- .../src/lib/transforms/sdkToCode/feature.ts | 14 +- .../actions/generateAndUpdateAgentRules.ts | 8 +- .../track/utils/buildAiCreditCostProperty.ts | 2 +- .../track/utils/getTokenTrackParams.ts | 12 +- .../utils/deduction/computeCreditCosts.ts | 26 ++-- .../balances/utils/types/featureDeduction.ts | 11 +- .../internal/features/creditSystemUtils.ts | 23 ++-- .../features/utils/getModelPricing.ts | 6 +- .../track/basic/track-tokens-orbs.test.ts | 73 +++++++++-- .../track/basic/track-tokens-replay.test.ts | 120 ++++++++++++++++++ .../generate-and-update-agent-rules.test.ts | 82 ++++++++++++ .../track/handle-track-tokens.test.ts | 9 +- .../unit/features/get-credit-cost.test.ts | 70 ++++++++++ .../unit/features/get-model-pricing.test.ts | 17 +++ 17 files changed, 467 insertions(+), 66 deletions(-) create mode 100644 server/tests/integration/balances/track/basic/track-tokens-replay.test.ts create mode 100644 server/tests/unit/agent/generate-and-update-agent-rules.test.ts create mode 100644 server/tests/unit/features/get-credit-cost.test.ts diff --git a/packages/atmn/src/commands/push/push.ts b/packages/atmn/src/commands/push/push.ts index fe3480669..a3384ce97 100644 --- a/packages/atmn/src/commands/push/push.ts +++ b/packages/atmn/src/commands/push/push.ts @@ -349,6 +349,32 @@ function normalizeFeatureForCompare(f: Feature): Record { })); } + if (f.type === "ai_credit_system") { + const ai = f as Extract; + if (ai.modelMarkups && Object.keys(ai.modelMarkups).length > 0) { + result.modelMarkups = Object.fromEntries( + Object.entries(ai.modelMarkups) + .sort(([a], [b]) => a.localeCompare(b)) + .map(([modelId, entry]) => [ + modelId, + { + markup: entry.markup, + inputCost: entry.inputCost, + outputCost: entry.outputCost, + }, + ]), + ); + } + if (ai.defaultMarkup != null) result.defaultMarkup = ai.defaultMarkup; + if (ai.providerMarkups && Object.keys(ai.providerMarkups).length > 0) { + result.providerMarkups = Object.fromEntries( + Object.entries(ai.providerMarkups).sort(([a], [b]) => + a.localeCompare(b), + ), + ); + } + } + return result; } diff --git a/packages/atmn/src/lib/transforms/apiToSdk/feature.ts b/packages/atmn/src/lib/transforms/apiToSdk/feature.ts index e4ba94ffb..793025736 100644 --- a/packages/atmn/src/lib/transforms/apiToSdk/feature.ts +++ b/packages/atmn/src/lib/transforms/apiToSdk/feature.ts @@ -60,6 +60,8 @@ export const featureTransformer = createTransformer({ ...BASE_COMPUTE, type: () => "ai_credit_system" as const, modelMarkups: (api) => mapModelMarkups(api), + defaultMarkup: (api) => api.default_markup ?? undefined, + providerMarkups: (api) => api.provider_markups ?? undefined, }, }, diff --git a/packages/atmn/src/lib/transforms/sdkToApi/feature.ts b/packages/atmn/src/lib/transforms/sdkToApi/feature.ts index 2c53db6f5..d963639e3 100644 --- a/packages/atmn/src/lib/transforms/sdkToApi/feature.ts +++ b/packages/atmn/src/lib/transforms/sdkToApi/feature.ts @@ -16,6 +16,8 @@ export interface ApiFeatureParams { input_cost?: number; output_cost?: number; }>; + default_markup?: number; + provider_markups?: Record; } export function transformFeatureToApi(feature: Feature): ApiFeatureParams { @@ -44,17 +46,25 @@ export function transformFeatureToApi(feature: Feature): ApiFeatureParams { })); } - if (feature.type === "ai_credit_system" && feature.modelMarkups) { - base.model_markups = Object.fromEntries( - Object.entries(feature.modelMarkups).map(([modelId, entry]) => [ - modelId, - { - markup: entry.markup, - input_cost: entry.inputCost, - output_cost: entry.outputCost, - }, - ]) - ); + if (feature.type === "ai_credit_system") { + if (feature.modelMarkups) { + base.model_markups = Object.fromEntries( + Object.entries(feature.modelMarkups).map(([modelId, entry]) => [ + modelId, + { + markup: entry.markup, + input_cost: entry.inputCost, + output_cost: entry.outputCost, + }, + ]) + ); + } + if (feature.defaultMarkup !== undefined) { + base.default_markup = feature.defaultMarkup; + } + if (feature.providerMarkups) { + base.provider_markups = feature.providerMarkups; + } } return base; diff --git a/packages/atmn/src/lib/transforms/sdkToCode/feature.ts b/packages/atmn/src/lib/transforms/sdkToCode/feature.ts index 54a728f17..eddbd30b1 100644 --- a/packages/atmn/src/lib/transforms/sdkToCode/feature.ts +++ b/packages/atmn/src/lib/transforms/sdkToCode/feature.ts @@ -42,9 +42,17 @@ export function buildFeatureCode(feature: Feature, varNameOverride?: string): st lines.push(`\tcreditSchema: ${formatValue(feature.creditSchema)},`); } - // Add modelMarkups for ai_credit_system features - if (feature.type === "ai_credit_system" && feature.modelMarkups) { - lines.push(`\tmodelMarkups: ${formatValue(feature.modelMarkups)},`); + // Add markup config for ai_credit_system features + if (feature.type === "ai_credit_system") { + if (feature.modelMarkups) { + lines.push(`\tmodelMarkups: ${formatValue(feature.modelMarkups)},`); + } + if (feature.defaultMarkup !== undefined) { + lines.push(`\tdefaultMarkup: ${feature.defaultMarkup},`); + } + if (feature.providerMarkups) { + lines.push(`\tproviderMarkups: ${formatValue(feature.providerMarkups)},`); + } } lines.push(`});`); diff --git a/server/src/internal/agent/rules/actions/generateAndUpdateAgentRules.ts b/server/src/internal/agent/rules/actions/generateAndUpdateAgentRules.ts index 34cd35d26..f3acc2ee4 100644 --- a/server/src/internal/agent/rules/actions/generateAndUpdateAgentRules.ts +++ b/server/src/internal/agent/rules/actions/generateAndUpdateAgentRules.ts @@ -11,13 +11,17 @@ export const generateAndUpdateAgentRules = async ({ endTime?: string; startTime?: string; }) => { - const generated = await generateAgentRules({ ctx, endTime, startTime }); + const [generated, existing] = await Promise.all([ + generateAgentRules({ ctx, endTime, startTime }), + agentRulesRepo.get({ db: ctx.db, orgId: ctx.org.id }), + ]); + // Generation only derives entity/credit rules; never overwrite user-written notes. const rules = await agentRulesRepo.upsert({ db: ctx.db, metadata: generated.metadata, orgId: ctx.org.id, orgSlug: ctx.org.slug, - rules: generated.rules, + rules: { ...generated.rules, notes: existing.notes }, }); return { diff --git a/server/src/internal/balances/track/utils/buildAiCreditCostProperty.ts b/server/src/internal/balances/track/utils/buildAiCreditCostProperty.ts index 66156e522..673225f5d 100644 --- a/server/src/internal/balances/track/utils/buildAiCreditCostProperty.ts +++ b/server/src/internal/balances/track/utils/buildAiCreditCostProperty.ts @@ -14,7 +14,7 @@ export const buildAiCreditCostProperty = ({ featureDeductions: FeatureDeduction[]; entries: Array<{ featureId: string; amount: number }>; }): Record | undefined => { - const aiDeduction = featureDeductions.find((d) => d.tokenUsage); + const aiDeduction = featureDeductions.find((d) => d.tokens); if (!aiDeduction) return; const creditCost: Record = {}; diff --git a/server/src/internal/balances/track/utils/getTokenTrackParams.ts b/server/src/internal/balances/track/utils/getTokenTrackParams.ts index 97a2996c1..16a9aedc5 100644 --- a/server/src/internal/balances/track/utils/getTokenTrackParams.ts +++ b/server/src/internal/balances/track/utils/getTokenTrackParams.ts @@ -139,12 +139,14 @@ export const getTokenTrackParams = async ({ { feature: aiCreditFeature, deduction: 1, - tokenUsage: { - modelName: input.model_id, - inputTokens: input.input_tokens, - outputTokens: input.output_tokens, + tokens: { + usage: { + modelName: input.model_id, + inputTokens: input.input_tokens, + outputTokens: input.output_tokens, + }, + cost, }, - precomputedCreditCost: cost, }, ]; diff --git a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts index 0ffe954a9..442ef83ae 100644 --- a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts +++ b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts @@ -8,8 +8,8 @@ export type CreditCostLookup = (entitlementId: string) => number; /** * Computes the credit cost for each customer entitlement and returns a lookup - * function. Uses precomputedCreditCost when available (token tracking), - * otherwise calls getCreditCost per entitlement (credit system schema lookups). + * function. Token deductions carry their USD cost from the API layer; all other + * costs come from credit system schema ratios. */ export const computeCreditCosts = async ({ cusEnts, @@ -20,33 +20,23 @@ export const computeCreditCosts = async ({ }): Promise => { const costMap = new Map(); - const tokens = deduction.tokenUsage - ? { - input: deduction.tokenUsage.inputTokens, - output: deduction.tokenUsage.outputTokens, - } - : undefined; - await Promise.all( cusEnts.map(async (ce) => { - // Precomputed cost (from /track/tokens) is in the AI credit feature's - // native unit (USD). It applies 1:1 to that feature's own entitlement, - // but parent credit systems still need their schema ratio applied — - // fall through to getCreditCost with amount = precomputed cost. + // A token deduction's cost is in the AI feature's native unit (USD): it + // applies 1:1 to its own entitlement, while parent credit systems apply + // their schema ratio to it via getCreditCost's amount. if ( - deduction.precomputedCreditCost != null && + deduction.tokens && ce.entitlement.feature.id === deduction.feature.id ) { - costMap.set(ce.id, deduction.precomputedCreditCost); + costMap.set(ce.id, deduction.tokens.cost); return; } const creditCost = await getCreditCost({ featureId: deduction.feature.id, creditSystem: ce.entitlement.feature, - amount: deduction.precomputedCreditCost, - modelName: deduction.tokenUsage?.modelName, - tokens, + amount: deduction.tokens?.cost, }); costMap.set(ce.id, creditCost); }), diff --git a/server/src/internal/balances/utils/types/featureDeduction.ts b/server/src/internal/balances/utils/types/featureDeduction.ts index ae6b7755d..8448112a5 100644 --- a/server/src/internal/balances/utils/types/featureDeduction.ts +++ b/server/src/internal/balances/utils/types/featureDeduction.ts @@ -7,13 +7,18 @@ export type TokenUsage = { outputTokens: number; }; +/** Token usage and its USD cost are priced together at the API layer — one cannot exist without the other. */ +export type TokenDeduction = { + usage: TokenUsage; + cost: number; +}; + export type FeatureDeduction = { feature: Feature; deduction: number; targetBalance?: number; - tokenUsage?: TokenUsage; - /** Pre-computed dollar cost; if set, the deduction layer skips its own getCreditCost call. */ - precomputedCreditCost?: number; + /** Present only for track_tokens deductions; standard deductions omit it. */ + tokens?: TokenDeduction; lock?: LockParams; lockReceipt?: LockReceipt; lockReceiptKey?: string; diff --git a/server/src/internal/features/creditSystemUtils.ts b/server/src/internal/features/creditSystemUtils.ts index c195ade41..19bb0b3c1 100644 --- a/server/src/internal/features/creditSystemUtils.ts +++ b/server/src/internal/features/creditSystemUtils.ts @@ -96,17 +96,22 @@ export const getCreditCost = async ({ return amount; } if (isAiCreditSystem(creditSystem.type)) { - if (!tokens || !modelName) { - throw new RecaseError({ - message: "modelName and tokens must be provided for AI credit systems", - code: ErrCode.InvalidRequest, - statusCode: 400, + if (tokens && modelName) { + return await getModelCreditCost({ + modelName, + creditSystem, + ...tokens, }); } - return await getModelCreditCost({ - modelName, - creditSystem, - ...tokens, + // No token context (plain /track values, balance updates, queued replays): + // the feature's own balance is already in USD, so the value maps 1:1. + if (featureId === creditSystem.id) { + return amount; + } + throw new RecaseError({ + message: "modelName and tokens must be provided for AI credit systems", + code: ErrCode.InvalidRequest, + statusCode: 400, }); } // If tracking the credit system feature itself, 1:1 mapping diff --git a/server/src/internal/features/utils/getModelPricing.ts b/server/src/internal/features/utils/getModelPricing.ts index d45d23cd6..0b4cdf0a9 100644 --- a/server/src/internal/features/utils/getModelPricing.ts +++ b/server/src/internal/features/utils/getModelPricing.ts @@ -7,9 +7,13 @@ const CACHE_KEY = "models_dev_pricing"; const STALE_KEY = `${CACHE_KEY}_stale`; const TTL_PRIMARY = 60 * 60 * 3; const TTL_STALE = 60 * 60 * 24 * 3; +// Runs inside the track request path — a hanging models.dev must not hang tracks. +const FETCH_TIMEOUT_MS = 5000; const fetchFromSource = async (): Promise => { - const response = await fetch("https://models.dev/api.json"); + const response = await fetch("https://models.dev/api.json", { + signal: AbortSignal.timeout(FETCH_TIMEOUT_MS), + }); if (!response.ok) { throw new InternalError({ message: `models.dev returned ${response.status}`, diff --git a/server/tests/integration/balances/track/basic/track-tokens-orbs.test.ts b/server/tests/integration/balances/track/basic/track-tokens-orbs.test.ts index 36dcde8b8..77f2a4576 100644 --- a/server/tests/integration/balances/track/basic/track-tokens-orbs.test.ts +++ b/server/tests/integration/balances/track/basic/track-tokens-orbs.test.ts @@ -10,13 +10,14 @@ import { Decimal } from "decimal.js"; // ═══════════════════════════════════════════════════════════════════ // TRACK-TOKENS-ORBS: AI credit system nested inside a parent credit system -// Verifies that a single /track/tokens call deducts USD from the AI credit -// feature AND deducts the ratio-mapped amount from any parent credit -// system whose schema references it. +// +// Parent credit systems are overflow pools (same semantics as classic +// metered → credits deduction order): a token track drains the AI credit +// balance first, and only the overflow is ratio-mapped onto the parent. // ═══════════════════════════════════════════════════════════════════ test.concurrent( - `${chalk.yellowBright("track-tokens-orbs: AI credit system inside parent credit system deducts both balances")}`, + `${chalk.yellowBright("track-tokens-orbs-1: AI balance covers the cost — parent orbs untouched")}`, async () => { const aiCreditsItem = items.free({ featureId: TestFeature.AiCredits, @@ -24,7 +25,7 @@ test.concurrent( }); const orbsItem = items.free({ featureId: TestFeature.Orbs, - includedUsage: 50_000, // 50,000 orbs + includedUsage: 50_000, // orbs schema: 1000 orbs per $1 of AI usage }); const freeProd = products.base({ id: "free", @@ -32,7 +33,7 @@ test.concurrent( }); const { customerId, autumnV1, autumnV2 } = await initScenario({ - customerId: "track-tokens-orbs", + customerId: "track-tokens-orbs-1", setup: [ s.customer({ testClock: false }), s.products({ list: [freeProd] }), @@ -49,9 +50,6 @@ test.concurrent( .div(1_000_000) .toNumber(); // 0.125 - // Orbs schema: 1000 orbs per $1 of AI usage - const expectedOrbsCost = new Decimal(expectedUsdCost).mul(1000).toNumber(); // 125 - const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { customer_id: customerId, feature_id: TestFeature.AiCredits, @@ -71,10 +69,61 @@ test.concurrent( usage: expectedUsdCost, }); - // Parent orbs balance dropped by USD cost × 1000 + // AI balance covered the full cost, so the parent overflow pool is untouched expect(customer.features[TestFeature.Orbs]).toMatchObject({ - balance: new Decimal(50_000).minus(expectedOrbsCost).toNumber(), - usage: expectedOrbsCost, + balance: 50_000, + usage: 0, + }); + }, +); + +test.concurrent( + `${chalk.yellowBright("track-tokens-orbs-2: cost exceeding AI balance overflows into parent orbs at the schema ratio")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 100, // $100 of AI usage + }); + const orbsItem = items.free({ + featureId: TestFeature.Orbs, + includedUsage: 50_000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem, orbsItem], + }); + + const { customerId, autumnV1, autumnV2 } = await initScenario({ + customerId: "track-tokens-orbs-2", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + // (5 * 24M) / 1M = $120 > the $100 AI balance + const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 24_000_000, + output_tokens: 0, + }); + expect(trackRes.value).toBeCloseTo(120, 10); + + const customer = await autumnV1.customers.get(customerId); + + // AI pool fully drained + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: 0, + usage: 100, + }); + + // $20 overflow lands on orbs at 1000 orbs per $1 + expect(customer.features[TestFeature.Orbs]).toMatchObject({ + balance: new Decimal(50_000).minus(20_000).toNumber(), + usage: 20_000, }); }, ); diff --git a/server/tests/integration/balances/track/basic/track-tokens-replay.test.ts b/server/tests/integration/balances/track/basic/track-tokens-replay.test.ts new file mode 100644 index 000000000..23e501696 --- /dev/null +++ b/server/tests/integration/balances/track/basic/track-tokens-replay.test.ts @@ -0,0 +1,120 @@ +import { expect, test } from "bun:test"; + +import type { ApiCustomerV3 } from "@autumn/shared"; +import { ApiVersion, ApiVersionClass } 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 { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; +import { Decimal } from "decimal.js"; +import { runQueuedTrack } from "@/internal/balances/track/runQueuedTrack.js"; + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-REPLAY: queued replay + plain value tracks on AI credit features +// +// When Redis fails open, track_tokens queues only the TrackParams body — the +// token context (FeatureDeduction.tokens) is not serialized. The +// replay worker rebuilds deductions from {feature_id, value}, so the USD value +// must deduct 1:1 from the AI credit balance, exactly like the original token +// track would have. Parent credit systems are overflow pools: untouched while +// the AI balance covers the deduction (same as live track_tokens behavior). +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-replay-1: queued replay body deducts AI credits 1:1")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 100, // $100 of AI usage + }); + const orbsItem = items.free({ + featureId: TestFeature.Orbs, + includedUsage: 50_000, // orbs schema: 1000 orbs per $1 of AI usage + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem, orbsItem], + }); + + const { customerId, autumnV1, ctx } = await initScenario({ + customerId: "track-tokens-replay-1", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + // The USD cost computed by the original track_tokens call; only this + // survives in the queued body. + const usdCost = 0.125; + + await runQueuedTrack({ + ctx: { ...ctx, apiVersion: new ApiVersionClass(ApiVersion.V2_1) }, + body: { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + value: usdCost, + idempotency_key: `replay-${crypto.randomUUID()}`, + }, + apiVersion: ApiVersion.V2_1, + }); + + const customer = await autumnV1.customers.get(customerId); + + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(100).minus(usdCost).toNumber(), + usage: usdCost, + }); + expect(customer.features[TestFeature.Orbs]).toMatchObject({ + balance: 50_000, + usage: 0, + }); + }, +); + +test.concurrent( + `${chalk.yellowBright("track-tokens-replay-2: plain /track with a USD value deducts an AI credit balance 1:1")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 100, + }); + const orbsItem = items.free({ + featureId: TestFeature.Orbs, + includedUsage: 50_000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem, orbsItem], + }); + + const { customerId, autumnV1, autumnV2 } = await initScenario({ + customerId: "track-tokens-replay-2", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + const usdValue = 5; + await autumnV2.post("/track", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + value: usdValue, + }); + + const customer = await autumnV1.customers.get(customerId); + + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(100).minus(usdValue).toNumber(), + usage: usdValue, + }); + expect(customer.features[TestFeature.Orbs]).toMatchObject({ + balance: 50_000, + usage: 0, + }); + }, +); diff --git a/server/tests/unit/agent/generate-and-update-agent-rules.test.ts b/server/tests/unit/agent/generate-and-update-agent-rules.test.ts new file mode 100644 index 000000000..9eff79834 --- /dev/null +++ b/server/tests/unit/agent/generate-and-update-agent-rules.test.ts @@ -0,0 +1,82 @@ +import { beforeEach, describe, expect, mock, test } from "bun:test"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; + +const generatedRules = { + entity_rules: { attach_to_entities: true, entity_feature_id: "deployments" }, + credit_rules: { credit_feature_id: "credits" }, + notes: "", +}; + +const mockState = { + existingNotes: "", + upsertCalls: [] as Record[], +}; + +mock.module( + "@/internal/agent/workflows/generateAgentRules/generateAgentRules.js", + () => ({ + generateAgentRules: async () => ({ + rules: generatedRules, + metadata: { generated_from: "axiom" }, + unconfigured: false, + }), + }), +); + +mock.module("@/internal/agent/rules/repos/index.js", () => ({ + agentRulesRepo: { + get: async () => ({ + entity_rules: { attach_to_entities: false, entity_feature_id: "" }, + credit_rules: { credit_feature_id: "" }, + notes: mockState.existingNotes, + metadata: {}, + org_id: "org_test", + org_slug: "test", + updated_at: null, + }), + upsert: async (args: { rules: typeof generatedRules }) => { + mockState.upsertCalls.push(args); + return { ...args.rules, metadata: {}, org_id: "org_test" }; + }, + }, +})); + +const { generateAndUpdateAgentRules } = await import( + "@/internal/agent/rules/actions/generateAndUpdateAgentRules.js" +); + +const ctx = { + db: {}, + org: { id: "org_test", slug: "test" }, +} as unknown as AutumnContext; + +describe("generateAndUpdateAgentRules", () => { + beforeEach(() => { + mockState.existingNotes = ""; + mockState.upsertCalls = []; + }); + + test("preserves existing user notes when applying generated rules", async () => { + mockState.existingNotes = "Always attach add-ons at the customer level."; + + const result = await generateAndUpdateAgentRules({ ctx }); + + expect(mockState.upsertCalls).toHaveLength(1); + expect(mockState.upsertCalls[0]).toMatchObject({ + rules: { + entity_rules: generatedRules.entity_rules, + credit_rules: generatedRules.credit_rules, + notes: "Always attach add-ons at the customer level.", + }, + }); + expect(result.notes).toBe("Always attach add-ons at the customer level."); + }); + + test("keeps notes empty when none were saved", async () => { + await generateAndUpdateAgentRules({ ctx }); + + expect(mockState.upsertCalls[0]).toMatchObject({ + rules: { notes: "" }, + }); + }); +}); diff --git a/server/tests/unit/balances/track/handle-track-tokens.test.ts b/server/tests/unit/balances/track/handle-track-tokens.test.ts index decca2c78..81f01e3e4 100644 --- a/server/tests/unit/balances/track/handle-track-tokens.test.ts +++ b/server/tests/unit/balances/track/handle-track-tokens.test.ts @@ -19,7 +19,14 @@ const featureDeductions = [ { feature: { id: "ai_credits" }, deduction: 1, - precomputedCreditCost: 3.5, + tokens: { + usage: { + modelName: "openai/gpt-4.1", + inputTokens: 100, + outputTokens: 50, + }, + cost: 3.5, + }, }, ]; diff --git a/server/tests/unit/features/get-credit-cost.test.ts b/server/tests/unit/features/get-credit-cost.test.ts new file mode 100644 index 000000000..2d97d89f1 --- /dev/null +++ b/server/tests/unit/features/get-credit-cost.test.ts @@ -0,0 +1,70 @@ +import { describe, expect, test } from "bun:test"; +import { + ErrCode, + type Feature, + FeatureType, + FeatureUsageType, +} from "@autumn/shared"; +import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; + +// Uses custom/* models so pricing resolves offline (no models.dev fetch). +const CUSTOM_MODEL = "custom/foo"; + +const aiCreditFeature: Feature = { + internal_id: "fe_ai_credits", + org_id: "org_test", + created_at: Date.now(), + env: "sandbox" as Feature["env"], + id: "ai_credits", + name: "AI Credits", + type: FeatureType.AiCreditSystem, + config: { schema: [], usage_type: FeatureUsageType.Single }, + archived: false, + event_names: [], + model_markups: { + [CUSTOM_MODEL]: { markup: 0, input_cost: 1000, output_cost: 2000 }, + }, +}; + +describe("getCreditCost — AI credit system without token context", () => { + test("self feature with no tokens maps 1:1 (plain /track values, queued replays)", async () => { + const cost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + amount: 5.25, + }); + expect(cost).toBe(5.25); + }); + + test("self feature with no tokens defaults to a per-unit cost of 1", async () => { + const cost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + }); + expect(cost).toBe(1); + }); + + test("self feature WITH tokens still prices through the model (not 1:1)", async () => { + const cost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + modelName: CUSTOM_MODEL, + tokens: { input: 1000, output: 500 }, + }); + // (1000 * 1000 + 2000 * 500) / 1_000_000 = 2.0 + expect(cost).toBeCloseTo(2.0, 10); + }); + + test("non-self feature with no tokens throws", async () => { + expect( + getCreditCost({ + featureId: "some_other_feature", + creditSystem: aiCreditFeature, + amount: 5, + }), + ).rejects.toMatchObject({ + code: ErrCode.InvalidRequest, + message: expect.stringContaining("modelName and tokens"), + }); + }); +}); diff --git a/server/tests/unit/features/get-model-pricing.test.ts b/server/tests/unit/features/get-model-pricing.test.ts index 67ac21dfe..9b06aedeb 100644 --- a/server/tests/unit/features/get-model-pricing.test.ts +++ b/server/tests/unit/features/get-model-pricing.test.ts @@ -113,3 +113,20 @@ test("fetch failure with no stale cache throws InternalError", async () => { message: "Failed to fetch models.dev pricing and no cache available", }); }); + +test("fetch carries an abort timeout so a hanging models.dev cannot hang tracks", async () => { + let capturedSignal: AbortSignal | undefined; + globalThis.fetch = Object.assign( + async (_input: unknown, init?: RequestInit) => { + fetchCalls++; + capturedSignal = init?.signal ?? undefined; + return Response.json(pricingData); + }, + { preconnect: realFetch.preconnect }, + ) as typeof fetch; + + await getModelsDevPricing(); + + expect(capturedSignal).toBeInstanceOf(AbortSignal); + expect(capturedSignal?.aborted).toBe(false); +});