318 lines
12 KiB
TypeScript
318 lines
12 KiB
TypeScript
import { describe, it, expect, vi, afterEach } from 'vitest';
|
|
import {
|
|
BudgetExceededError,
|
|
BudgetPricingUnavailableError,
|
|
CostTracker,
|
|
DEFAULT_MODEL_PRICING,
|
|
type ModelPricing,
|
|
} from '../src/cost-tracker.js';
|
|
|
|
describe('CostTracker', () => {
|
|
const pricing: Record<string, ModelPricing> = {
|
|
'claude-sonnet': { inputPer1k: 0.003, outputPer1k: 0.015 },
|
|
'claude-haiku': { inputPer1k: 0.00025, outputPer1k: 0.00125 },
|
|
};
|
|
|
|
it('tracks token usage per model', () => {
|
|
const tracker = new CostTracker(pricing);
|
|
tracker.addUsage('claude-sonnet', 1000, 500);
|
|
tracker.addUsage('claude-sonnet', 2000, 300);
|
|
const stats = tracker.getStats();
|
|
expect(stats.totalInputTokens).toBe(3000);
|
|
expect(stats.totalOutputTokens).toBe(800);
|
|
});
|
|
|
|
it('estimates cost', () => {
|
|
const tracker = new CostTracker(pricing);
|
|
tracker.addUsage('claude-sonnet', 1000, 1000);
|
|
const stats = tracker.getStats();
|
|
expect(stats.estimatedCost).toBeCloseTo(0.018, 4);
|
|
});
|
|
|
|
it('handles unknown models with fallback Sonnet pricing', () => {
|
|
const tracker = new CostTracker(pricing);
|
|
tracker.addUsage('unknown-model', 1000, 500);
|
|
const stats = tracker.getStats();
|
|
expect(stats.totalInputTokens).toBe(1000);
|
|
// Fallback: Sonnet pricing ($0.003/1K in, $0.015/1K out)
|
|
// 1K input = $0.003, 0.5K output = $0.0075 -> total $0.0105
|
|
expect(stats.estimatedCost).toBeCloseTo(0.0105, 4);
|
|
});
|
|
|
|
it('formats summary', () => {
|
|
const tracker = new CostTracker(pricing);
|
|
tracker.addUsage('claude-sonnet', 1000, 500);
|
|
const summary = tracker.formatSummary();
|
|
expect(summary).toContain('1000');
|
|
expect(summary).toContain('500');
|
|
expect(summary).toContain('$');
|
|
});
|
|
|
|
describe('current model pricing', () => {
|
|
afterEach(() => vi.restoreAllMocks());
|
|
|
|
it('prices a known Opus-class model at Opus rates ($15/$75 per 1M)', () => {
|
|
expect(DEFAULT_MODEL_PRICING['claude-opus-4-8']).toEqual({ inputPer1k: 0.015, outputPer1k: 0.075 });
|
|
const tracker = new CostTracker();
|
|
tracker.addUsage('claude-opus-4-8', 1000, 1000);
|
|
// 1K in * $0.015 + 1K out * $0.075 = $0.09
|
|
expect(tracker.getStats().estimatedCost).toBeCloseTo(0.09, 4);
|
|
});
|
|
|
|
it('includes current Sonnet and Haiku ids', () => {
|
|
expect(DEFAULT_MODEL_PRICING['claude-sonnet-5']).toBeDefined();
|
|
expect(DEFAULT_MODEL_PRICING['claude-haiku-4-5']).toBeDefined();
|
|
});
|
|
|
|
it('uses provider rates for the live Gemini and Codex acceptance models', () => {
|
|
expect(DEFAULT_MODEL_PRICING['google/gemini-2.5-flash'])
|
|
.toEqual({ inputPer1k: 0.0003, outputPer1k: 0.0025 });
|
|
expect(DEFAULT_MODEL_PRICING['openrouter/openai/gpt-5.3-codex'])
|
|
.toEqual({ inputPer1k: 0.00175, outputPer1k: 0.014 });
|
|
|
|
const tracker = new CostTracker();
|
|
tracker.addUsage('google/gemini-2.5-flash', 1000, 1000);
|
|
tracker.addUsage('openrouter/openai/gpt-5.3-codex', 1000, 1000);
|
|
expect(tracker.getStats().estimatedCost).toBeCloseTo(0.01855, 6);
|
|
});
|
|
|
|
it('warns once and uses family-aware fallback for an unknown Opus id', () => {
|
|
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
|
|
const tracker = new CostTracker();
|
|
const unknownOpus = `claude-opus-9-9-${Math.random().toString(36).slice(2)}`;
|
|
tracker.addUsage(unknownOpus, 1000, 1000);
|
|
tracker.addUsage(unknownOpus, 1000, 1000);
|
|
const stats = tracker.getStats();
|
|
// Opus fallback (not Sonnet): 2 * ($0.015 + $0.075) = $0.18, not $0.036.
|
|
expect(stats.estimatedCost).toBeCloseTo(0.18, 4);
|
|
// Loud warning fired, and only once for the same unknown model.
|
|
expect(warn).toHaveBeenCalledTimes(1);
|
|
expect(warn.mock.calls[0][0]).toContain(unknownOpus);
|
|
});
|
|
|
|
it('treats unlisted Ollama models as local and free without a cloud-pricing warning', () => {
|
|
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
|
|
const tracker = new CostTracker();
|
|
tracker.addUsage('ollama/minimax-m2.7:cloud', 1000, 1000);
|
|
|
|
expect(tracker.getStats().estimatedCost).toBe(0);
|
|
expect(warn).not.toHaveBeenCalled();
|
|
});
|
|
});
|
|
|
|
describe('getDailyTotal', () => {
|
|
it('adds persisted carryover to in-process usage without double-seeding', () => {
|
|
vi.useFakeTimers();
|
|
vi.setSystemTime(new Date('2026-07-21T12:00:00.000Z'));
|
|
const tracker = new CostTracker();
|
|
try {
|
|
tracker.initializeDailyCarryover('2026-07-21', 6);
|
|
tracker.initializeDailyCarryover('2026-07-21', 8);
|
|
tracker.addUsage('claude-sonnet-4-6', 1000, 1000);
|
|
|
|
expect(tracker.hasDailyCarryover('2026-07-21')).toBe(true);
|
|
expect(tracker.getDailyTotal()).toBeCloseTo(6.018, 6);
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('resets carryover and in-process usage at the next UTC day', () => {
|
|
vi.useFakeTimers();
|
|
const tracker = new CostTracker();
|
|
try {
|
|
vi.setSystemTime(new Date('2026-07-20T23:59:00.000Z'));
|
|
tracker.initializeDailyCarryover('2026-07-20', 6);
|
|
tracker.addUsage('claude-sonnet-4-6', 1000, 1000);
|
|
|
|
vi.setSystemTime(new Date('2026-07-21T00:01:00.000Z'));
|
|
expect(tracker.getDailyTotal()).toBe(0);
|
|
expect(tracker.hasDailyCarryover('2026-07-21')).toBe(false);
|
|
|
|
tracker.initializeDailyCarryover('2026-07-21', 2);
|
|
tracker.addUsage('claude-sonnet-4-6', 2000, 1000);
|
|
expect(tracker.getDailyTotal()).toBeCloseTo(2.021, 6);
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
});
|
|
});
|
|
|
|
describe('hard daily spend reservations', () => {
|
|
const pricing: Record<string, ModelPricing> = {
|
|
paid: { inputPer1k: 1, outputPer1k: 1 },
|
|
};
|
|
|
|
it('atomically prevents concurrent reservations from sharing the same capacity', async () => {
|
|
const tracker = new CostTracker(pricing);
|
|
tracker.setBudget(1, 'hard');
|
|
|
|
const attempts = await Promise.allSettled([
|
|
Promise.resolve().then(() => tracker.reserveModelSpend({
|
|
model: 'paid', inputTokens: 400, maxOutputTokens: 200,
|
|
})),
|
|
Promise.resolve().then(() => tracker.reserveModelSpend({
|
|
model: 'paid', inputTokens: 400, maxOutputTokens: 200,
|
|
})),
|
|
]);
|
|
|
|
expect(attempts.filter(result => result.status === 'fulfilled')).toHaveLength(1);
|
|
expect(attempts.find(result => result.status === 'rejected')).toMatchObject({
|
|
reason: expect.any(BudgetExceededError),
|
|
});
|
|
expect(tracker.getReservedDailyTotal()).toBeCloseTo(0.6, 6);
|
|
});
|
|
|
|
it('reconciles a conservative reservation to actual usage exactly once', () => {
|
|
const tracker = new CostTracker(pricing);
|
|
tracker.setBudget(1, 'hard');
|
|
const reservation = tracker.reserveModelSpend({
|
|
model: 'paid', inputTokens: 400, maxOutputTokens: 400, workspaceId: 'workspace-a',
|
|
});
|
|
|
|
tracker.reconcileModelSpend(reservation, { inputTokens: 100, outputTokens: 100 });
|
|
tracker.reconcileModelSpend(reservation, { inputTokens: 900, outputTokens: 900 });
|
|
|
|
expect(tracker.getReservedDailyTotal()).toBe(0);
|
|
expect(tracker.getDailyTotal()).toBeCloseTo(0.2, 6);
|
|
expect(tracker.getWorkspaceCost('workspace-a')).toBeCloseTo(0.2, 6);
|
|
expect(() => tracker.reserveModelSpend({
|
|
model: 'paid', inputTokens: 400, maxOutputTokens: 400,
|
|
})).not.toThrow();
|
|
});
|
|
|
|
it('retains the conservative reservation after ambiguous provider failure', () => {
|
|
const tracker = new CostTracker(pricing);
|
|
tracker.setBudget(1, 'hard');
|
|
const reservation = tracker.reserveModelSpend({
|
|
model: 'paid', inputTokens: 500, maxOutputTokens: 500,
|
|
});
|
|
|
|
tracker.commitReservedModelSpend(reservation);
|
|
|
|
expect(tracker.getReservedDailyTotal()).toBe(0);
|
|
expect(tracker.getDailyTotal()).toBeCloseTo(1, 6);
|
|
expect(() => tracker.reserveModelSpend({
|
|
model: 'paid', inputTokens: 1, maxOutputTokens: 1,
|
|
})).toThrow(BudgetExceededError);
|
|
});
|
|
|
|
it('allows explicitly verified free execution after the paid cap is exhausted', () => {
|
|
const tracker = new CostTracker(pricing);
|
|
tracker.setBudget(1, 'hard');
|
|
const paid = tracker.reserveModelSpend({
|
|
model: 'paid', inputTokens: 500, maxOutputTokens: 500,
|
|
});
|
|
tracker.commitReservedModelSpend(paid);
|
|
|
|
const local = tracker.reserveModelSpend({
|
|
model: 'unpriced-local-model',
|
|
inputTokens: 10_000,
|
|
maxOutputTokens: 10_000,
|
|
billingClass: 'free',
|
|
});
|
|
tracker.reconcileModelSpend(local, { inputTokens: 10_000, outputTokens: 10_000 });
|
|
|
|
expect(tracker.getDailyTotal()).toBeCloseTo(1, 6);
|
|
});
|
|
|
|
it('normalizes zero to disabled and rejects negative budgets', () => {
|
|
const tracker = new CostTracker(pricing);
|
|
|
|
tracker.setBudget(0, 'hard');
|
|
expect(tracker.getBudget()).toEqual({ dailyBudgetUsd: null, mode: 'hard' });
|
|
expect(() => tracker.setBudget(-1, 'hard')).toThrow(/non-negative finite/i);
|
|
});
|
|
|
|
it('preserves an explicit priced classification for ollama-prefixed routes', () => {
|
|
const tracker = new CostTracker({
|
|
'ollama/remote-paid': { inputPer1k: 1, outputPer1k: 1 },
|
|
});
|
|
tracker.setBudget(1, 'hard');
|
|
const reservation = tracker.reserveModelSpend({
|
|
model: 'ollama/remote-paid', inputTokens: 300, maxOutputTokens: 300,
|
|
billingClass: 'priced',
|
|
});
|
|
|
|
tracker.reconcileModelSpend(reservation, { inputTokens: 100, outputTokens: 100 });
|
|
|
|
expect(tracker.getDailyTotal()).toBeCloseTo(0.2, 6);
|
|
});
|
|
|
|
it('commits the reservation when provider usage is non-finite', () => {
|
|
const tracker = new CostTracker(pricing);
|
|
tracker.setBudget(1, 'hard');
|
|
const reservation = tracker.reserveModelSpend({
|
|
model: 'paid', inputTokens: 400, maxOutputTokens: 400,
|
|
});
|
|
|
|
tracker.reconcileModelSpend(reservation, { inputTokens: Number.NaN, outputTokens: 0 });
|
|
|
|
expect(tracker.getDailyTotal()).toBeCloseTo(0.8, 6);
|
|
expect(() => tracker.reserveModelSpend({
|
|
model: 'paid', inputTokens: 200, maxOutputTokens: 1,
|
|
})).toThrow(BudgetExceededError);
|
|
});
|
|
|
|
it('rejects non-finite direct usage before it can poison the ledger', () => {
|
|
const tracker = new CostTracker(pricing);
|
|
|
|
expect(() => tracker.addUsage('paid', Number.NaN, 0)).toThrow(/finite/i);
|
|
expect(tracker.getDailyTotal()).toBe(0);
|
|
});
|
|
|
|
it('fails closed when hard mode lacks trusted pricing for a paid route', () => {
|
|
const tracker = new CostTracker();
|
|
tracker.setBudget(1, 'hard');
|
|
|
|
expect(() => tracker.reserveModelSpend({
|
|
model: 'openrouter/auto', inputTokens: 1, maxOutputTokens: 1,
|
|
billingClass: 'priced',
|
|
})).toThrow(BudgetPricingUnavailableError);
|
|
});
|
|
|
|
it('prices provider-wrapped model IDs only through a trusted catalog suffix', () => {
|
|
const tracker = new CostTracker();
|
|
tracker.setBudget(1, 'hard');
|
|
|
|
const reservation = tracker.reserveModelSpend({
|
|
model: 'anthropic/claude-sonnet-4-6',
|
|
inputTokens: 1_000,
|
|
maxOutputTokens: 1_000,
|
|
billingClass: 'priced',
|
|
});
|
|
tracker.reconcileModelSpend(reservation, { inputTokens: 1_000, outputTokens: 1_000 });
|
|
|
|
expect(tracker.getDailyTotal()).toBeCloseTo(0.018, 6);
|
|
expect(() => tracker.reserveModelSpend({
|
|
model: 'anthropic/untrusted-custom-model',
|
|
inputTokens: 1,
|
|
maxOutputTokens: 1,
|
|
billingClass: 'priced',
|
|
})).toThrow(BudgetPricingUnavailableError);
|
|
});
|
|
|
|
it('honors explicit paid versus free billing for Ollama-routed reservations', () => {
|
|
const tracker = new CostTracker();
|
|
|
|
const paid = tracker.reserveModelSpend({
|
|
model: 'ollama/minimax-m2.7:cloud',
|
|
inputTokens: 1_000,
|
|
maxOutputTokens: 1_000,
|
|
billingClass: 'priced',
|
|
});
|
|
tracker.reconcileModelSpend(paid, { inputTokens: 1_000, outputTokens: 1_000 });
|
|
const local = tracker.reserveModelSpend({
|
|
model: 'ollama/qwen2.5:1.5b',
|
|
inputTokens: 1_000,
|
|
maxOutputTokens: 1_000,
|
|
billingClass: 'free',
|
|
});
|
|
tracker.reconcileModelSpend(local, { inputTokens: 1_000, outputTokens: 1_000 });
|
|
|
|
expect(tracker.getDailyTotal()).toBeCloseTo(0.018, 6);
|
|
});
|
|
|
|
});
|