admin-valentin/agent-ecosystem
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: add tests for cost calculation

Browse source
This commit is contained in:
KaustubhPatange 2026-02-22 14:49:11 +05:30
parent a039fdd573
commit 2bb95db596
4 changed files with 4725 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.gitignore vendored
View file

@ -47,5 +47,3 @@ temp/
eslint-fix/
remotion/*
resources/pricing.json

4195
resources/pricing.json Normal file
View file

File diff suppressed because it is too large Load diff

1
src/main/types/chunks.ts
View file

@ -461,6 +461,7 @@ export const EMPTY_METRICS: SessionMetrics = {
cacheReadTokens: 0,
cacheCreationTokens: 0,
messageCount: 0,
costUsd: 0,
};
// =============================================================================

529
test/main/utils/costCalculation.test.ts Normal file
View file

@ -0,0 +1,529 @@
/**
* Tests for cost calculation in jsonl.ts
*/
import { describe, it, expect, beforeEach, vi } from 'vitest';
import * as fs from 'fs';
import { calculateMetrics } from '@main/utils/jsonl';
import type { ParsedMessage } from '@main/types';
// Mock fs module
vi.mock('fs');
describe('Cost Calculation', () => {
// Sample pricing data matching Claude models
const mockPricingData = {
'claude-3-5-sonnet-20241022': {
input_cost_per_token: 0.000003,
output_cost_per_token: 0.000015,
cache_creation_input_token_cost: 0.00000375,
cache_read_input_token_cost: 0.0000003,
input_cost_per_token_above_200k_tokens: 0.000006,
output_cost_per_token_above_200k_tokens: 0.00003,
cache_creation_input_token_cost_above_200k_tokens: 0.0000075,
cache_read_input_token_cost_above_200k_tokens: 0.0000006,
},
'claude-3-opus-20240229': {
input_cost_per_token: 0.000015,
output_cost_per_token: 0.000075,
cache_creation_input_token_cost: 0.00001875,
cache_read_input_token_cost: 0.0000015,
},
};
beforeEach(() => {
// Reset modules to clear pricing cache
vi.resetModules();
// Mock fs.readFileSync to return our test pricing data
vi.mocked(fs.readFileSync).mockReturnValue(JSON.stringify(mockPricingData));
vi.mocked(fs.existsSync).mockReturnValue(true);
});
describe('Basic Cost Calculation', () => {
it('should calculate cost for simple token usage', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 1000,
output_tokens: 500,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
// Expected: (1000 * 0.000003) + (500 * 0.000015) = 0.003 + 0.0075 = 0.0105
expect(metrics.costUsd).toBeCloseTo(0.0105, 6);
});
it('should calculate cost with cache tokens', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 1000,
output_tokens: 500,
cache_creation_input_tokens: 200,
cache_read_input_tokens: 300,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
// Input: 1000 * 0.000003 = 0.003
// Output: 500 * 0.000015 = 0.0075
// Cache creation: 200 * 0.00000375 = 0.00075
// Cache read: 300 * 0.0000003 = 0.00009
// Total: 0.01134
expect(metrics.costUsd).toBeCloseTo(0.01134, 6);
});
it('should return 0 cost when no model is specified', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
usage: {
input_tokens: 1000,
output_tokens: 500,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
expect(metrics.costUsd).toBe(0);
});
it('should return 0 cost when model pricing not found', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'unknown-model',
usage: {
input_tokens: 1000,
output_tokens: 500,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
expect(metrics.costUsd).toBe(0);
});
});
describe('Tiered Pricing', () => {
it('should use base rates for tokens below 200k threshold', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 100_000,
output_tokens: 50_000,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
// Input: 100000 * 0.000003 = 0.3
// Output: 50000 * 0.000015 = 0.75
// Total: 1.05
expect(metrics.costUsd).toBeCloseTo(1.05, 6);
});
it('should use tiered rates for input tokens above 200k threshold', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 250_000,
output_tokens: 1_000,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
// Input: (200000 * 0.000003) + (50000 * 0.000006) = 0.6 + 0.3 = 0.9
// Output: 1000 * 0.000015 = 0.015
// Total: 0.915
expect(metrics.costUsd).toBeCloseTo(0.915, 6);
});
it('should use tiered rates for output tokens above 200k threshold', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 1_000,
output_tokens: 250_000,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
// Input: 1000 * 0.000003 = 0.003
// Output: (200000 * 0.000015) + (50000 * 0.00003) = 3.0 + 1.5 = 4.5
// Total: 4.503
expect(metrics.costUsd).toBeCloseTo(4.503, 6);
});
it('should use tiered rates for cache tokens above 200k threshold', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 1_000,
output_tokens: 1_000,
cache_creation_input_tokens: 250_000,
cache_read_input_tokens: 250_000,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
// Input: 1000 * 0.000003 = 0.003
// Output: 1000 * 0.000015 = 0.015
// Cache creation: (200000 * 0.00000375) + (50000 * 0.0000075) = 0.75 + 0.375 = 1.125
// Cache read: (200000 * 0.0000003) + (50000 * 0.0000006) = 0.06 + 0.03 = 0.09
// Total: 1.233
expect(metrics.costUsd).toBeCloseTo(1.233, 6);
});
it('should handle model without tiered pricing', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-opus-20240229',
usage: {
input_tokens: 250_000,
output_tokens: 250_000,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
// No tiered rates, so use base rates even above 200k
// Input: 250000 * 0.000015 = 3.75
// Output: 250000 * 0.000075 = 18.75
// Total: 22.5
expect(metrics.costUsd).toBeCloseTo(22.5, 6);
});
});
describe('Multiple Messages', () => {
it('should aggregate costs across multiple messages', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 1000,
output_tokens: 500,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
{
type: 'assistant',
uuid: 'msg-2',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 2000,
output_tokens: 1000,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
// Message 1: (1000 * 0.000003) + (500 * 0.000015) = 0.0105
// Message 2: (2000 * 0.000003) + (1000 * 0.000015) = 0.021
// Total: 0.0315
expect(metrics.costUsd).toBeCloseTo(0.0315, 6);
});
it('should use first model found when calculating aggregated cost', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 1000,
output_tokens: 500,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
{
type: 'assistant',
uuid: 'msg-2',
timestamp: new Date(),
content: [],
model: 'claude-3-opus-20240229', // Different model
usage: {
input_tokens: 1000,
output_tokens: 500,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
// Uses first model (sonnet) pricing for all tokens
// Total tokens: 2000 input, 1000 output
// Cost: (2000 * 0.000003) + (1000 * 0.000015) = 0.006 + 0.015 = 0.021
expect(metrics.costUsd).toBeCloseTo(0.021, 6);
});
});
describe('Edge Cases', () => {
it('should handle zero tokens', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 0,
output_tokens: 0,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
expect(metrics.costUsd).toBe(0);
});
it('should handle messages without usage data', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
expect(metrics.costUsd).toBe(0);
});
it('should handle empty messages array', () => {
const messages: ParsedMessage[] = [];
const metrics = calculateMetrics(messages);
expect(metrics.costUsd).toBe(0);
});
it('should handle pricing data load failure gracefully', async () => {
// Suppress expected console.error for this test
const consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
// Reset modules to clear the pricing cache
vi.resetModules();
// Mock fs to throw error BEFORE importing calculateMetrics
vi.mocked(fs.readFileSync).mockImplementation(() => {
throw new Error('File not found');
});
// Re-import calculateMetrics to get fresh instance with cleared cache
const { calculateMetrics: freshCalculateMetrics } = await import('@main/utils/jsonl');
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 1000,
output_tokens: 500,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = freshCalculateMetrics(messages);
expect(metrics.costUsd).toBe(0);
// Verify that console.error was called (error was logged)
expect(consoleErrorSpy).toHaveBeenCalled();
// Restore console.error
consoleErrorSpy.mockRestore();
});
});
describe('Model Name Lookup', () => {
it('should find model with exact match', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 1000,
output_tokens: 500,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
expect(metrics.costUsd).toBeGreaterThan(0);
});
it('should find model with case-insensitive match', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'CLAUDE-3-5-SONNET-20241022',
usage: {
input_tokens: 1000,
output_tokens: 500,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
expect(metrics.costUsd).toBeGreaterThan(0);
});
});
describe('Integration with Other Metrics', () => {
it('should include cost alongside other session metrics', () => {
const messages: ParsedMessage[] = [
{
type: 'assistant',
uuid: 'msg-1',
timestamp: new Date(),
content: [],
model: 'claude-3-5-sonnet-20241022',
usage: {
input_tokens: 1000,
output_tokens: 500,
},
toolCalls: [],
toolResults: [],
isSidechain: false,
},
];
const metrics = calculateMetrics(messages);
// Check that all expected metrics are present
expect(metrics).toHaveProperty('totalTokens');
expect(metrics).toHaveProperty('inputTokens');
expect(metrics).toHaveProperty('outputTokens');
expect(metrics).toHaveProperty('costUsd');
expect(metrics.totalTokens).toBe(1500);
expect(metrics.inputTokens).toBe(1000);
expect(metrics.outputTokens).toBe(500);
expect(metrics.costUsd).toBeCloseTo(0.0105, 6);
});
});
});