feat(usage): track cached tokens + correct input/output/cache cost (#2209)
Normalize every provider to one cache-inclusive convention via canonicalizeUsage() before persist, and price cached + cache_creation as subsets of prompt_tokens in calculateCostFromTokens() to stop double-counting. usageRepo now delegates cost math to a single source. Surface Cached tokens/cost across dashboard (overview, tokens, cost, details). Merge Claude message_start cache with message_delta output so cache counts survive. Compatible LLM nodes now allow multiple API-key connections (key pool). Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
hodtien
84
tests/unit/cached-token-e2e.test.js
Normal file
84
tests/unit/cached-token-e2e.test.js
Normal file
@@ -0,0 +1,84 @@
|
||||
// End-to-end: a cache-bearing request flows through canonicalizeUsage →
|
||||
// saveRequestUsage → getUsageStats, proving cached tokens are persisted,
|
||||
// aggregated, and cost is computed correctly (the bug this branch fixes).
|
||||
import fs from "node:fs";
|
||||
import os from "node:os";
|
||||
import path from "node:path";
|
||||
import { describe, it, expect, beforeAll, afterAll, vi } from "vitest";
|
||||
import { canonicalizeUsage } from "../../open-sse/utils/usageTracking.js";
|
||||
|
||||
const originalDataDir = process.env.DATA_DIR;
|
||||
let tempDir;
|
||||
let db;
|
||||
|
||||
beforeAll(async () => {
|
||||
tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "9router-cached-e2e-"));
|
||||
process.env.DATA_DIR = tempDir;
|
||||
vi.resetModules();
|
||||
db = await import("@/lib/db/index.js");
|
||||
await db.initDb();
|
||||
});
|
||||
|
||||
afterAll(() => {
|
||||
if (tempDir) fs.rmSync(tempDir, { recursive: true, force: true });
|
||||
if (originalDataDir === undefined) delete process.env.DATA_DIR;
|
||||
else process.env.DATA_DIR = originalDataDir;
|
||||
});
|
||||
|
||||
describe("cached-token end-to-end (persist + aggregate + cost)", () => {
|
||||
it("Claude cache usage: canonical prompt is inclusive, cached persisted, cost correct", async () => {
|
||||
// Raw Claude usage (cache-EXCLUSIVE prompt): input 100, cache_read 200, cache_creation 30, output 50
|
||||
const canonical = canonicalizeUsage({
|
||||
prompt_tokens: 100,
|
||||
completion_tokens: 50,
|
||||
cache_read_input_tokens: 200,
|
||||
cache_creation_input_tokens: 30,
|
||||
});
|
||||
expect(canonical.prompt_tokens).toBe(330); // inclusive
|
||||
|
||||
await db.saveRequestUsage({
|
||||
provider: "anthropic",
|
||||
model: "claude-sonnet-4-6",
|
||||
connectionId: "c-cache",
|
||||
tokens: canonical,
|
||||
endpoint: "/v1/messages",
|
||||
status: "ok",
|
||||
});
|
||||
|
||||
const stats = await db.getUsageStats("24h");
|
||||
expect(stats.totalCachedTokens).toBe(200);
|
||||
expect(stats.totalPromptTokens).toBe(330);
|
||||
expect(stats.byProvider.anthropic.cachedTokens).toBe(200);
|
||||
|
||||
// Cost: nonCached=330-200-30=100 @3 + cached 200 @0.30 + creation 30 @3.75 + output 50 @15
|
||||
const expected = (100 * 3 + 200 * 0.3 + 30 * 3.75 + 50 * 15) / 1_000_000;
|
||||
const hist = await db.getUsageHistory({ provider: "anthropic" });
|
||||
expect(hist.length).toBe(1);
|
||||
expect(hist[0].cost).toBeCloseTo(expected, 12);
|
||||
expect(hist[0].tokens.cached_tokens).toBe(200);
|
||||
expect(hist[0].tokens.cache_creation_input_tokens).toBe(30);
|
||||
});
|
||||
|
||||
it("OpenAI cache usage: inclusive prompt passes through, cached counted once", async () => {
|
||||
const canonical = canonicalizeUsage({
|
||||
prompt_tokens: 1000, // already includes cached
|
||||
completion_tokens: 200,
|
||||
cached_tokens: 600,
|
||||
});
|
||||
expect(canonical.prompt_tokens).toBe(1000);
|
||||
expect(canonical.cached_tokens).toBe(600);
|
||||
|
||||
await db.saveRequestUsage({
|
||||
provider: "openai",
|
||||
model: "gpt-4o",
|
||||
connectionId: "c-oai",
|
||||
tokens: canonical,
|
||||
endpoint: "/v1/chat/completions",
|
||||
status: "ok",
|
||||
});
|
||||
|
||||
const hist = await db.getUsageHistory({ provider: "openai" });
|
||||
expect(hist[0].tokens.prompt_tokens).toBe(1000);
|
||||
expect(hist[0].tokens.cached_tokens).toBe(600);
|
||||
});
|
||||
});
|
||||
188
tests/unit/cached-token-usage.test.js
Normal file
188
tests/unit/cached-token-usage.test.js
Normal file
@@ -0,0 +1,188 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { canonicalizeUsage, extractUsage, mergeUsage } from "../../open-sse/utils/usageTracking.js";
|
||||
import { calculateCostFromTokens } from "../../open-sse/providers/pricing.js";
|
||||
import { toOpenAIUsage } from "../../open-sse/translator/concerns/usage.js";
|
||||
|
||||
// Canonical convention (single source of truth for storage + cost):
|
||||
// prompt_tokens = total input INCLUDING cache read + cache creation
|
||||
// cached_tokens = cache-read portion (subset of prompt_tokens)
|
||||
// cache_creation_input_tokens = cache-write portion (subset of prompt_tokens)
|
||||
// completion_tokens = output
|
||||
// Discriminator: Claude reports cache separately (prompt EXCLUDES cache);
|
||||
// OpenAI/Gemini report prompt INCLUDING cached_tokens.
|
||||
describe("canonicalizeUsage", () => {
|
||||
it("folds Claude exclusive cache into an inclusive prompt count", () => {
|
||||
// Claude: input_tokens excludes cache; cache_read + cache_creation are separate
|
||||
const out = canonicalizeUsage({
|
||||
prompt_tokens: 100,
|
||||
completion_tokens: 50,
|
||||
cache_read_input_tokens: 200,
|
||||
cache_creation_input_tokens: 30,
|
||||
});
|
||||
expect(out.prompt_tokens).toBe(330); // 100 + 200 + 30
|
||||
expect(out.completion_tokens).toBe(50);
|
||||
expect(out.cached_tokens).toBe(200);
|
||||
expect(out.cache_creation_input_tokens).toBe(30);
|
||||
});
|
||||
|
||||
it("passes through OpenAI inclusive prompt unchanged", () => {
|
||||
// OpenAI: prompt_tokens already includes cached_tokens (a subset)
|
||||
const out = canonicalizeUsage({
|
||||
prompt_tokens: 330,
|
||||
completion_tokens: 50,
|
||||
cached_tokens: 200,
|
||||
});
|
||||
expect(out.prompt_tokens).toBe(330);
|
||||
expect(out.cached_tokens).toBe(200);
|
||||
expect(out.cache_creation_input_tokens).toBe(0);
|
||||
});
|
||||
|
||||
it("passes through Gemini inclusive prompt (cachedContent already counted)", () => {
|
||||
const out = canonicalizeUsage({
|
||||
prompt_tokens: 500,
|
||||
completion_tokens: 80,
|
||||
cached_tokens: 120,
|
||||
reasoning_tokens: 40,
|
||||
});
|
||||
expect(out.prompt_tokens).toBe(500);
|
||||
expect(out.cached_tokens).toBe(120);
|
||||
expect(out.reasoning_tokens).toBe(40);
|
||||
});
|
||||
|
||||
it("handles no-cache usage", () => {
|
||||
const out = canonicalizeUsage({ prompt_tokens: 100, completion_tokens: 50 });
|
||||
expect(out.prompt_tokens).toBe(100);
|
||||
expect(out.cached_tokens).toBe(0);
|
||||
expect(out.cache_creation_input_tokens).toBe(0);
|
||||
});
|
||||
|
||||
it("is idempotent (running twice yields the same canonical shape)", () => {
|
||||
const once = canonicalizeUsage({
|
||||
prompt_tokens: 100,
|
||||
completion_tokens: 50,
|
||||
cache_read_input_tokens: 200,
|
||||
cache_creation_input_tokens: 30,
|
||||
});
|
||||
const twice = canonicalizeUsage(once);
|
||||
expect(twice.prompt_tokens).toBe(330);
|
||||
expect(twice.cached_tokens).toBe(200);
|
||||
expect(twice.cache_creation_input_tokens).toBe(30);
|
||||
expect(twice.completion_tokens).toBe(50);
|
||||
});
|
||||
|
||||
it("returns null for invalid input", () => {
|
||||
expect(canonicalizeUsage(null)).toBeNull();
|
||||
expect(canonicalizeUsage(undefined)).toBeNull();
|
||||
});
|
||||
|
||||
it("folds a Claude cache-miss first write (cache_creation only, no cache_read yet)", () => {
|
||||
// Cache-miss on first write: upstream emits cache_creation_input_tokens but
|
||||
// no cache_read_input_tokens at all (not even 0). Must still fold into prompt
|
||||
// instead of falling through to the OpenAI passthrough branch.
|
||||
const out = canonicalizeUsage({
|
||||
prompt_tokens: 100,
|
||||
completion_tokens: 20,
|
||||
cache_creation_input_tokens: 500,
|
||||
});
|
||||
expect(out.prompt_tokens).toBe(600); // 100 + 0 (no read) + 500
|
||||
expect(out.cached_tokens).toBe(0);
|
||||
expect(out.cache_creation_input_tokens).toBe(500);
|
||||
});
|
||||
});
|
||||
|
||||
describe("calculateCostFromTokens (canonical inclusive convention)", () => {
|
||||
const pricing = { input: 3, output: 15, cached: 0.3, cache_creation: 3.75 };
|
||||
|
||||
it("prices cached + cache_creation as subsets of an inclusive prompt without double-counting", () => {
|
||||
// prompt=330 includes 200 cached + 30 cache_creation → 100 full-price input
|
||||
const cost = calculateCostFromTokens(
|
||||
{ prompt_tokens: 330, completion_tokens: 50, cached_tokens: 200, cache_creation_input_tokens: 30 },
|
||||
pricing
|
||||
);
|
||||
const expected =
|
||||
(100 * 3 + 200 * 0.3 + 30 * 3.75 + 50 * 15) / 1_000_000;
|
||||
expect(cost).toBeCloseTo(expected, 12);
|
||||
});
|
||||
|
||||
it("does not let cache_creation drive nonCached negative", () => {
|
||||
// pathological: cached + creation exceeds prompt → nonCached clamps at 0
|
||||
const cost = calculateCostFromTokens(
|
||||
{ prompt_tokens: 100, completion_tokens: 0, cached_tokens: 80, cache_creation_input_tokens: 40 },
|
||||
pricing
|
||||
);
|
||||
const expected = (0 * 3 + 80 * 0.3 + 40 * 3.75) / 1_000_000;
|
||||
expect(cost).toBeCloseTo(expected, 12);
|
||||
});
|
||||
|
||||
it("matches plain input pricing when no cache present", () => {
|
||||
const cost = calculateCostFromTokens({ prompt_tokens: 100, completion_tokens: 50 }, pricing);
|
||||
expect(cost).toBeCloseTo((100 * 3 + 50 * 15) / 1_000_000, 12);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Anthropic streaming usage (message_start carries cache, message_delta output-only)", () => {
|
||||
it("extractUsage reads input + cache from message_start", () => {
|
||||
const u = extractUsage({
|
||||
type: "message_start",
|
||||
message: { usage: { input_tokens: 100, output_tokens: 1, cache_read_input_tokens: 200, cache_creation_input_tokens: 30 } },
|
||||
});
|
||||
expect(u.prompt_tokens).toBe(100);
|
||||
expect(u.cache_read_input_tokens).toBe(200);
|
||||
expect(u.cache_creation_input_tokens).toBe(30);
|
||||
});
|
||||
|
||||
it("merges message_start cache with message_delta output without clobbering", () => {
|
||||
// Real Anthropic SSE: cache only in message_start, real output only in message_delta.
|
||||
const start = extractUsage({
|
||||
type: "message_start",
|
||||
message: { usage: { input_tokens: 100, output_tokens: 1, cache_read_input_tokens: 200, cache_creation_input_tokens: 30 } },
|
||||
});
|
||||
const delta = extractUsage({ type: "message_delta", usage: { output_tokens: 50 } });
|
||||
const merged = mergeUsage(start, delta);
|
||||
expect(merged.prompt_tokens).toBe(100);
|
||||
expect(merged.cache_read_input_tokens).toBe(200);
|
||||
expect(merged.cache_creation_input_tokens).toBe(30);
|
||||
expect(merged.completion_tokens).toBe(50);
|
||||
|
||||
// And it canonicalizes to a cache-inclusive prompt for storage/cost.
|
||||
const canon = canonicalizeUsage(merged);
|
||||
expect(canon.prompt_tokens).toBe(330); // 100 + 200 + 30
|
||||
expect(canon.cached_tokens).toBe(200);
|
||||
expect(canon.cache_creation_input_tokens).toBe(30);
|
||||
expect(canon.completion_tokens).toBe(50);
|
||||
});
|
||||
|
||||
it("does not let a NaN field poison the running max-merge", () => {
|
||||
// typeof NaN === "number", so a naive Math.max(prev, NaN) is NaN — one
|
||||
// malformed chunk must not wipe out an already-accumulated good value.
|
||||
const prev = { prompt_tokens: 100, cache_read_input_tokens: 200 };
|
||||
const bad = { prompt_tokens: NaN, completion_tokens: 50 };
|
||||
const merged = mergeUsage(prev, bad);
|
||||
expect(merged.prompt_tokens).toBe(100);
|
||||
expect(merged.cache_read_input_tokens).toBe(200);
|
||||
expect(merged.completion_tokens).toBe(50);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Kiro usage pass-through", () => {
|
||||
it("passes through plain input/output when no cache fields are present", () => {
|
||||
const out = toOpenAIUsage({ inputTokens: 100, outputTokens: 50 }, "kiro");
|
||||
expect(out.prompt_tokens).toBe(100);
|
||||
expect(out.completion_tokens).toBe(50);
|
||||
expect(out.total_tokens).toBe(150);
|
||||
expect(out.prompt_tokens_details).toBeUndefined();
|
||||
});
|
||||
|
||||
it("forward-compat: surfaces cache fields if Kiro event shape grows them", () => {
|
||||
// ponytail: Amazon Q upstream doesn't expose cache today, but if it starts
|
||||
// sending cache_read_input_tokens / cache_creation_input_tokens / cachedTokens,
|
||||
// cost tracking should pick them up automatically without another change.
|
||||
const out = toOpenAIUsage(
|
||||
{ inputTokens: 500, outputTokens: 100, cache_read_input_tokens: 200, cache_creation_input_tokens: 50 },
|
||||
"kiro"
|
||||
);
|
||||
expect(out.prompt_tokens_details).toBeDefined();
|
||||
expect(out.prompt_tokens_details.cached_tokens).toBe(200);
|
||||
expect(out.prompt_tokens_details.cache_creation_tokens).toBe(50);
|
||||
});
|
||||
});
|
||||
@@ -38,14 +38,14 @@ async function setupTestContext(nodeData) {
|
||||
};
|
||||
}
|
||||
|
||||
function makeRequest(provider) {
|
||||
function makeRequest(provider, name = "Test Connection") {
|
||||
return new Request("https://9router.local/api/providers", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
provider,
|
||||
apiKey: "test-key",
|
||||
name: "Test Connection",
|
||||
name,
|
||||
defaultModel: "test-model",
|
||||
}),
|
||||
});
|
||||
@@ -156,8 +156,8 @@ describe("compatible provider connections API", () => {
|
||||
});
|
||||
cleanup = ctx.cleanup;
|
||||
|
||||
const firstResponse = await ctx.POST(makeRequest(ctx.node.id));
|
||||
const secondResponse = await ctx.POST(makeRequest(ctx.node.id));
|
||||
const firstResponse = await ctx.POST(makeRequest(ctx.node.id, "Key A"));
|
||||
const secondResponse = await ctx.POST(makeRequest(ctx.node.id, "Key B"));
|
||||
const storedConnections = await ctx.getProviderConnections({ provider: ctx.node.id });
|
||||
|
||||
expect(firstResponse.status).toBe(201);
|
||||
|
||||
Reference in New Issue
Block a user