chore: improve types and atmn
This commit is contained in:
@@ -349,6 +349,32 @@ function normalizeFeatureForCompare(f: Feature): Record<string, unknown> {
|
||||
}));
|
||||
}
|
||||
|
||||
if (f.type === "ai_credit_system") {
|
||||
const ai = f as Extract<Feature, { type: "ai_credit_system" }>;
|
||||
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;
|
||||
}
|
||||
|
||||
|
||||
@@ -60,6 +60,8 @@ export const featureTransformer = createTransformer<RawApiFeature, Feature>({
|
||||
...BASE_COMPUTE,
|
||||
type: () => "ai_credit_system" as const,
|
||||
modelMarkups: (api) => mapModelMarkups(api),
|
||||
defaultMarkup: (api) => api.default_markup ?? undefined,
|
||||
providerMarkups: (api) => api.provider_markups ?? undefined,
|
||||
},
|
||||
},
|
||||
|
||||
|
||||
@@ -16,6 +16,8 @@ export interface ApiFeatureParams {
|
||||
input_cost?: number;
|
||||
output_cost?: number;
|
||||
}>;
|
||||
default_markup?: number;
|
||||
provider_markups?: Record<string, { markup: number }>;
|
||||
}
|
||||
|
||||
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;
|
||||
|
||||
@@ -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(`});`);
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -14,7 +14,7 @@ export const buildAiCreditCostProperty = ({
|
||||
featureDeductions: FeatureDeduction[];
|
||||
entries: Array<{ featureId: string; amount: number }>;
|
||||
}): Record<string, number> | undefined => {
|
||||
const aiDeduction = featureDeductions.find((d) => d.tokenUsage);
|
||||
const aiDeduction = featureDeductions.find((d) => d.tokens);
|
||||
if (!aiDeduction) return;
|
||||
|
||||
const creditCost: Record<string, number> = {};
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
];
|
||||
|
||||
|
||||
@@ -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<CreditCostLookup> => {
|
||||
const costMap = new Map<string, number>();
|
||||
|
||||
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);
|
||||
}),
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<ModelPricingData> => {
|
||||
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}`,
|
||||
|
||||
@@ -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<ApiCustomerV3>(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,
|
||||
});
|
||||
},
|
||||
);
|
||||
|
||||
@@ -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<ApiCustomerV3>(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<ApiCustomerV3>(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,
|
||||
});
|
||||
},
|
||||
);
|
||||
@@ -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<string, unknown>[],
|
||||
};
|
||||
|
||||
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: "" },
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -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,
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
|
||||
70
server/tests/unit/features/get-credit-cost.test.ts
Normal file
70
server/tests/unit/features/get-credit-cost.test.ts
Normal file
@@ -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"),
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user