From fbd12f3c19759e1bcf68f3433f6abff549c8e77e Mon Sep 17 00:00:00 2001 From: ozymandiashh <234437643+ozymandiashh@users.noreply.github.com> Date: Tue, 23 Jun 2026 01:44:27 +0300 Subject: [PATCH] feat(pricing): add user price overrides for models (#390) --- src/config.ts | 2 + src/main.ts | 108 +++++++++++++++++++++++++++++- src/models.ts | 106 +++++++++++++++++++++++++++++ src/usage-aggregator.ts | 11 ++- tests/cli-price-override.test.ts | 85 +++++++++++++++++++++++ tests/models.test.ts | 111 ++++++++++++++++++++++++++++++- 6 files changed, 417 insertions(+), 6 deletions(-) create mode 100644 tests/cli-price-override.test.ts diff --git a/src/config.ts b/src/config.ts index a1e0abab..65620755 100644 --- a/src/config.ts +++ b/src/config.ts @@ -30,6 +30,8 @@ export type CodeburnConfig = { plan?: Plan plans?: PlanConfigMap modelAliases?: Record + // Rates are stored as USD per 1,000,000 tokens; models.ts converts them to per-token ModelCosts. + priceOverrides?: Record // Extra Claude config directories to aggregate usage across (e.g. work / // personal accounts). Honored by getClaudeConfigDirs() below the // CLAUDE_CONFIG_DIRS/CLAUDE_CONFIG_DIR env vars. Lets the macOS menubar (a diff --git a/src/main.ts b/src/main.ts index 9f0df2ea..fb0047de 100644 --- a/src/main.ts +++ b/src/main.ts @@ -2,7 +2,7 @@ import { isAbsolute } from 'path' import { Command } from 'commander' import { installMenubarApp } from './menubar-installer.js' import { exportCsv, exportJson, type PeriodExport } from './export.js' -import { loadPricing, setModelAliases, setLocalModelSavings, setProxyPaths, normalizeProxyPath } from './models.js' +import { loadPricing, setModelAliases, setPriceOverrides, setLocalModelSavings, setProxyPaths, normalizeProxyPath } from './models.js' import { parseAllSessions, filterProjectsByName, filterProjectsByDateRange, clearSessionCache } from './parser.js' import { allProviderNames } from './providers/index.js' import { convertCost } from './currency.js' @@ -29,7 +29,7 @@ import { runAgyStatusLineHook, uninstallAntigravityStatusLineHook, } from './antigravity-statusline.js' -import { clearPlan, readConfig, readPlan, readPlans, saveConfig, savePlan, getConfigFilePath, type Plan, type PlanId, type PlanProvider } from './config.js' +import { clearPlan, readConfig, readPlan, readPlans, saveConfig, savePlan, getConfigFilePath, type CodeburnConfig, type Plan, type PlanId, type PlanProvider } from './config.js' import { clampResetDay, getPlanUsageOrNull, getPlanUsages, type PlanUsage } from './plan-usage.js' import { getPresetPlan, isPlanId, isPlanProvider, PLAN_IDS, PLAN_PROVIDERS, planDisplayName } from './plans.js' import { createRequire } from 'node:module' @@ -59,6 +59,30 @@ function parseInteger(value: string): number { return parseInt(value, 10) } +type PriceOverrideConfig = NonNullable[string] + +type PriceOverrideOptions = { + input?: number + output?: number + cacheRead?: number + cacheCreation?: number + remove?: string + list?: boolean +} + +function invalidUsdPerMillionRate(option: string, value: number | undefined): string | null { + if (value === undefined) return null + if (Number.isFinite(value) && value >= 0) return null + return `Invalid ${option}: expected a finite number >= 0 (USD per 1,000,000 tokens).` +} + +function formatPriceOverrideParts(rates: PriceOverrideConfig): string { + const parts = [`input ${rates.input}`, `output ${rates.output}`] + if (typeof rates.cacheRead === 'number') parts.push(`cache read ${rates.cacheRead}`) + if (typeof rates.cacheCreation === 'number') parts.push(`cache creation ${rates.cacheCreation}`) + return parts.join(', ') +} + type JsonPlanSummary = { id: PlanId provider: PlanProvider @@ -176,6 +200,7 @@ program.hook('preAction', async (thisCommand) => { } const config = await readConfig() setModelAliases(config.modelAliases ?? {}) + setPriceOverrides(config.priceOverrides ?? {}) setLocalModelSavings(config.localModelSavings ?? {}) setProxyPaths(config.proxyPaths ?? []) if (thisCommand.opts<{ verbose?: boolean }>().verbose) { @@ -906,6 +931,85 @@ program console.log(` Config: ${getConfigFilePath()}\n`) }) +program + .command('price-override [model]') + .description('Override or add local model pricing. Rates are USD per 1,000,000 tokens (e.g. --input 0.27).') + .option('--input ', 'Input token price in USD per 1,000,000 tokens', parseNumber) + .option('--output ', 'Output token price in USD per 1,000,000 tokens', parseNumber) + .option('--cache-read ', 'Cache-read token price in USD per 1,000,000 tokens', parseNumber) + .option('--cache-creation ', 'Cache-creation token price in USD per 1,000,000 tokens', parseNumber) + .option('--remove ', 'Remove a price override') + .option('--list', 'List configured price overrides') + .action(async (model?: string, opts?: PriceOverrideOptions) => { + const config = await readConfig() + const overrides = new Map(Object.entries(config.priceOverrides ?? {})) + + if (opts?.list || (!model && !opts?.remove)) { + const entries = [...overrides.entries()] + if (entries.length === 0) { + console.log('\n No price overrides configured.') + console.log(' Rates use USD per 1,000,000 tokens.') + console.log(` Config: ${getConfigFilePath()}`) + console.log(' Add one with: codeburn price-override --input --output \n') + } else { + console.log('\n Price overrides (USD per 1,000,000 tokens):') + for (const [name, rates] of entries) { + console.log(` ${name}: ${formatPriceOverrideParts(rates)}`) + } + console.log(` Config: ${getConfigFilePath()}\n`) + } + return + } + + if (opts?.remove) { + if (!overrides.has(opts.remove)) { + console.error(`\n Price override not found: ${opts.remove}\n`) + process.exitCode = 1 + return + } + overrides.delete(opts.remove) + config.priceOverrides = overrides.size > 0 ? Object.fromEntries(overrides) : undefined + await saveConfig(config) + console.log(`\n Removed price override: ${opts.remove}\n`) + return + } + + const input = opts?.input + const output = opts?.output + const cacheRead = opts?.cacheRead + const cacheCreation = opts?.cacheCreation + if (!model || input === undefined || output === undefined) { + console.error('\n Usage: codeburn price-override --input --output [--cache-read ] [--cache-creation ]\n') + process.exitCode = 1 + return + } + + const invalidRate = [ + invalidUsdPerMillionRate('--input', input), + invalidUsdPerMillionRate('--output', output), + invalidUsdPerMillionRate('--cache-read', cacheRead), + invalidUsdPerMillionRate('--cache-creation', cacheCreation), + ].find((message): message is string => message !== null) + if (invalidRate) { + console.error(`\n ${invalidRate}\n`) + process.exitCode = 1 + return + } + + const override: PriceOverrideConfig = { + input, + output, + ...(cacheRead !== undefined ? { cacheRead } : {}), + ...(cacheCreation !== undefined ? { cacheCreation } : {}), + } + overrides.set(model, override) + config.priceOverrides = Object.fromEntries(overrides) + await saveConfig(config) + console.log(`\n Price override saved: ${model}: ${formatPriceOverrideParts(override)}`) + console.log(' Unit: USD per 1,000,000 tokens') + console.log(` Config: ${getConfigFilePath()}\n`) + }) + program .command('model-savings [local] [baseline]') .description('Track a local model as "savings" rather than cost. Maps a local-model name to a paid baseline so the dashboard can show what the same tokens would have cost on the baseline (e.g. codeburn model-savings "llama3.1:8b" gpt-4o). The local call itself still costs $0 — actual cost is left untouched.') diff --git a/src/models.ts b/src/models.ts index 4d04c1ad..0e6a5930 100644 --- a/src/models.ts +++ b/src/models.ts @@ -14,6 +14,13 @@ export type ModelCosts = { fastMultiplier: number } +type PriceOverrideRates = { + input: number + output: number + cacheRead?: number + cacheCreation?: number +} + type LiteLLMEntry = { input_cost_per_token?: number output_cost_per_token?: number @@ -332,6 +339,10 @@ const BUILTIN_ALIASES: Record = { } let userAliases: Record = {} +let userPriceOverrides: Map = new Map() +let userPriceOverridesConfig: Record = {} +let sortedPriceOverrideKeys: string[] | null = null +let lowercasePriceOverrideIndex: Map | null = null // Called once during CLI startup after config is loaded. // User aliases take precedence over built-ins. @@ -339,6 +350,76 @@ export function setModelAliases(aliases: Record): void { userAliases = aliases } +function priceOverrideRatePerToken(usdPerMillion: number | undefined): number | null { + if (typeof usdPerMillion !== 'number') return null + return safePerTokenRate(usdPerMillion / 1_000_000) +} + +// Called once during CLI startup after config is loaded. +// Config/CLI rates are USD per 1,000,000 tokens; ModelCosts stores USD/token. +export function setPriceOverrides(overrides: Record): void { + const next = new Map() + const nextConfig: Record = {} + for (const [model, rates] of Object.entries(overrides)) { + if (!model || !rates || typeof rates !== 'object') continue + nextConfig[model] = { ...rates } + const input = priceOverrideRatePerToken(rates.input) + const output = priceOverrideRatePerToken(rates.output) + if (input === null || output === null) continue + next.set(model, buildCosts( + input, + output, + priceOverrideRatePerToken(rates.cacheCreation), + priceOverrideRatePerToken(rates.cacheRead), + undefined, + )) + } + userPriceOverrides = next + userPriceOverridesConfig = nextConfig + sortedPriceOverrideKeys = null + lowercasePriceOverrideIndex = null +} + +function getSortedPriceOverrideKeys(): string[] { + if (sortedPriceOverrideKeys === null) { + sortedPriceOverrideKeys = Array.from(userPriceOverrides.keys()).sort((a, b) => b.length - a.length) + } + return sortedPriceOverrideKeys +} + +function getLowercasePriceOverrideIndex(): Map { + if (lowercasePriceOverrideIndex === null) { + lowercasePriceOverrideIndex = new Map() + for (const [key, costs] of userPriceOverrides) { + const lk = key.toLowerCase() + if (!lowercasePriceOverrideIndex.has(lk)) lowercasePriceOverrideIndex.set(lk, costs) + } + } + return lowercasePriceOverrideIndex +} + +function getPriceOverrideExact(...keys: string[]): ModelCosts | null { + for (const key of keys) { + const costs = userPriceOverrides.get(key) + if (costs) return costs + } + return null +} + +function getPriceOverridePrefix(canonical: string): ModelCosts | null { + for (const key of getSortedPriceOverrideKeys()) { + if (canonical.startsWith(key + '-') || canonical === key) { + return userPriceOverrides.get(key)! + } + } + return null +} + +function getPriceOverrideCaseInsensitive(canonical: string, withPrefix: string): ModelCosts | null { + const lowerIndex = getLowercasePriceOverrideIndex() + return lowerIndex.get(canonical.toLowerCase()) ?? lowerIndex.get(withPrefix.toLowerCase()) ?? null +} + // Local-model savings config. Kept separate from userAliases: a `modelAliases` // entry rewrites a model's identity for actual cost; a `localModelSavings` // entry keeps the model cost at $0 and reports the *avoided* spend against a @@ -402,6 +483,22 @@ export function getLocalModelSavingsConfigHash(): string { return parts.join('\u0002') } +export function getPriceOverridesConfigHash(): string { + const keys = Object.keys(userPriceOverridesConfig).sort() + if (keys.length === 0) return '' + const parts = keys.map(k => { + const rates = userPriceOverridesConfig[k] + return [ + k, + rates.input, + rates.output, + rates.cacheRead ?? '', + rates.cacheCreation ?? '', + ].join('\u0001') + }) + return parts.join('\u0002') +} + // Absolute directory prefixes whose sessions are routed through a // subscription-backed proxy (config `proxyPaths`). Stored already-normalized so // the per-project match is a cheap compare. Set during preAction. See @@ -472,6 +569,9 @@ export function getModelCosts(model: string): ModelCosts | null { const canonicalName = getCanonicalName(model) const canonical = resolveAlias(canonicalName) + const override = getPriceOverrideExact(model, withPrefix, canonicalName, canonical) + if (override) return override + // An explicit alias for a bare (un-prefixed) model name is authoritative: it // must win over a coincidental stripped reseller key of the same name. LiteLLM // ships `snowflake/claude-4-opus` ($5), which the bundler strips to a bare @@ -485,6 +585,9 @@ export function getModelCosts(model: string): ModelCosts | null { if (pricingCache.has(canonical)) return pricingCache.get(canonical)! + const prefixOverride = getPriceOverridePrefix(canonical) + if (prefixOverride) return prefixOverride + // Iterate keys longest-first so a model id like `gpt-5-mini` matches the // `gpt-5-mini` entry rather than collapsing to the shorter `gpt-5` entry // due to dictionary insertion order. @@ -494,6 +597,9 @@ export function getModelCosts(model: string): ModelCosts | null { } } + const caseInsensitiveOverride = getPriceOverrideCaseInsensitive(canonical, withPrefix) + if (caseInsensitiveOverride) return caseInsensitiveOverride + // Case-insensitive fallback: gap-filled keys from OpenRouter are lowercase // slugs (e.g. `minimax-m3`), but sessions report `MiniMax-M3`. Only consulted // after the exact/canonical/prefix attempts, so it never changes a match that diff --git a/src/usage-aggregator.ts b/src/usage-aggregator.ts index ad15d0ed..42240122 100644 --- a/src/usage-aggregator.ts +++ b/src/usage-aggregator.ts @@ -2,7 +2,7 @@ import { homedir } from 'node:os' import { CATEGORY_LABELS, type ProjectSummary, type TaskCategory, type DateRange } from './types.js' import { type PeriodData, type ProviderCost, type BreakdownArrays, type MenubarPayload, buildMenubarPayload } from './menubar-json.js' import { parseAllSessions, filterProjectsByName, filterProjectsByDays } from './parser.js' -import { getLocalModelSavingsConfigHash, getShortModelName } from './models.js' +import { getLocalModelSavingsConfigHash, getPriceOverridesConfigHash, getShortModelName } from './models.js' import { getAllProviders } from './providers/index.js' import { aggregateProjectsIntoDays, buildPeriodDataFromDays } from './day-aggregator.js' import { aggregateModelEfficiency } from './model-efficiency.js' @@ -53,12 +53,19 @@ export function buildPeriodData(label: string, projects: ProjectSummary[]): Peri } } +export function getDailyCacheConfigHash(): string { + const savingsHash = getLocalModelSavingsConfigHash() + const overridesHash = getPriceOverridesConfigHash() + if (!overridesHash) return savingsHash + return `localModelSavings=${savingsHash}\u0002priceOverrides=${overridesHash}` +} + async function hydrateCache(): Promise { try { return await ensureCacheHydrated( (range) => parseAllSessions(range, 'all'), aggregateProjectsIntoDays, - getLocalModelSavingsConfigHash(), + getDailyCacheConfigHash(), ) } catch (err) { // Previously swallowed silently, which turned any backfill failure into an diff --git a/tests/cli-price-override.test.ts b/tests/cli-price-override.test.ts new file mode 100644 index 00000000..cdbea134 --- /dev/null +++ b/tests/cli-price-override.test.ts @@ -0,0 +1,85 @@ +import { mkdtemp, readFile, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { spawnSync } from 'node:child_process' + +import { describe, it, expect } from 'vitest' + +const CLI_TIMEOUT_MS = 10_000 + +function runCli(args: string[], home: string) { + return spawnSync(process.execPath, ['--import', 'tsx', 'src/cli.ts', ...args], { + cwd: process.cwd(), + env: { + ...process.env, + HOME: home, + USERPROFILE: home, + HOMEPATH: home, + HOMEDRIVE: '', + }, + encoding: 'utf-8', + }) +} + +function readConfig(home: string): Promise> { + return readFile(join(home, '.config', 'codeburn', 'config.json'), 'utf-8') + .then(raw => JSON.parse(raw) as Record) +} + +describe('codeburn price-override command', () => { + it('saves, lists, and removes a model price override', async () => { + const home = await mkdtemp(join(tmpdir(), 'codeburn-cli-price-override-')) + try { + const set = runCli([ + 'price-override', + 'unpriced-provider/test-model', + '--input', + '0.27', + '--output', + '1.10', + '--cache-read', + '0.03', + '--cache-creation', + '0.42', + ], home) + expect(set.status).toBe(0) + expect(set.stdout).toContain('unpriced-provider/test-model') + expect(set.stdout).toContain('USD per 1,000,000 tokens') + + const saved = await readConfig(home) + expect(saved.priceOverrides).toEqual({ + 'unpriced-provider/test-model': { + input: 0.27, + output: 1.1, + cacheRead: 0.03, + cacheCreation: 0.42, + }, + }) + + const list = runCli(['price-override', '--list'], home) + expect(list.status).toBe(0) + expect(list.stdout).toContain('unpriced-provider/test-model') + expect(list.stdout).toContain('input 0.27') + expect(list.stdout).toContain('output 1.1') + + const remove = runCli(['price-override', '--remove', 'unpriced-provider/test-model'], home) + expect(remove.status).toBe(0) + + const after = await readConfig(home) + expect(after.priceOverrides).toBeUndefined() + } finally { + await rm(home, { recursive: true, force: true }) + } + }, CLI_TIMEOUT_MS) + + it('rejects an invalid rate', async () => { + const home = await mkdtemp(join(tmpdir(), 'codeburn-cli-price-override-')) + try { + const result = runCli(['price-override', 'bad-rate-model', '--input', 'abc', '--output', '1'], home) + expect(result.status).toBe(1) + expect(result.stderr).toContain('Invalid --input') + } finally { + await rm(home, { recursive: true, force: true }) + } + }, CLI_TIMEOUT_MS) +}) diff --git a/tests/models.test.ts b/tests/models.test.ts index 403a5550..8621ab3d 100644 --- a/tests/models.test.ts +++ b/tests/models.test.ts @@ -3,13 +3,28 @@ import { tmpdir } from 'os' import { join } from 'path' import { describe, it, expect, beforeAll, afterEach } from 'vitest' -import { getModelCosts, getShortModelName, calculateCost, loadPricing, setModelAliases } from '../src/models.js' +import { + getModelCosts, + getShortModelName, + calculateCost, + loadPricing, + setModelAliases, + setPriceOverrides, + setLocalModelSavings, + getLocalModelSavingsConfigHash, + getPriceOverridesConfigHash, +} from '../src/models.js' +import { getDailyCacheConfigHash } from '../src/usage-aggregator.js' beforeAll(async () => { await loadPricing() }) -afterEach(() => setModelAliases({})) +afterEach(() => { + setModelAliases({}) + setPriceOverrides({}) + setLocalModelSavings({}) +}) describe('getModelCosts', () => { it('does not match short canonical against longer pricing key', () => { @@ -229,6 +244,98 @@ describe('user aliases via setModelAliases', () => { }) }) +describe('user price overrides', () => { + it('prices a model missing from the pricing snapshot', () => { + const model = 'zz-price-override-missing-model-390' + expect(getModelCosts(model)).toBeNull() + + setPriceOverrides({ + [model]: { input: 1.25, output: 2.5 }, + }) + + const costs = getModelCosts(model) + expect(costs).not.toBeNull() + expect(costs!.inputCostPerToken).toBe(1.25e-6) + expect(costs!.outputCostPerToken).toBe(2.5e-6) + expect(calculateCost(model, 1_000_000, 1_000_000, 0, 0, 0)).toBe(3.75) + }) + + it('wins over snapshot pricing and configured aliases', () => { + setModelAliases({ + 'price-override-aliased-model': 'claude-opus-4-6', + 'price-override-canonical-source': 'price-override-canonical-target', + }) + setPriceOverrides({ + 'gpt-4o': { input: 7, output: 8 }, + 'claude-opus-4-6': { input: 4, output: 5 }, + 'price-override-aliased-model': { input: 2, output: 3 }, + 'price-override-canonical-target': { input: 6, output: 7 }, + }) + + expect(getModelCosts('gpt-4o')!.inputCostPerToken).toBe(7e-6) + expect(getModelCosts('price-override-aliased-model')!.inputCostPerToken).toBe(2e-6) + expect(getModelCosts('price-override-canonical-source')!.inputCostPerToken).toBe(6e-6) + }) + + it('converts USD per 1,000,000 tokens to per-token ModelCosts exactly', () => { + const model = 'price-override-unit-conversion' + setPriceOverrides({ + [model]: { input: 1, output: 0 }, + }) + + expect(getModelCosts(model)!.inputCostPerToken).toBe(1e-6) + expect(calculateCost(model, 1_000_000, 0, 0, 0, 0)).toBe(1) + }) + + it('defaults cache rates from input pricing when omitted', () => { + const model = 'price-override-cache-defaults' + setPriceOverrides({ + [model]: { input: 10, output: 20 }, + }) + + const costs = getModelCosts(model) + expect(costs).not.toBeNull() + expect(costs!.cacheWriteCostPerToken).toBeCloseTo(12.5e-6, 12) + expect(costs!.cacheReadCostPerToken).toBeCloseTo(1e-6, 12) + }) + + it('wins for case-insensitive and prefix matches without shadowing a more-specific exact snapshot entry', () => { + const miniSnapshot = getModelCosts('gpt-5-mini') + expect(miniSnapshot).not.toBeNull() + + setPriceOverrides({ + 'gpt-5': { input: 91, output: 92 }, + }) + + expect(getModelCosts('GPT-5')!.inputCostPerToken).toBe(91e-6) + expect(getModelCosts('gpt-5-foo')!.inputCostPerToken).toBe(91e-6) + + const mini = getModelCosts('gpt-5-mini') + expect(mini).not.toBeNull() + expect(mini!.inputCostPerToken).toBe(miniSnapshot!.inputCostPerToken) + expect(mini!.outputCostPerToken).toBe(miniSnapshot!.outputCostPerToken) + }) + + it('includes price overrides in the daily cache config hash without changing the empty-override hash', () => { + setLocalModelSavings({ local: 'gpt-4o' }) + setPriceOverrides({}) + + const savingsOnly = getLocalModelSavingsConfigHash() + expect(getPriceOverridesConfigHash()).toBe('') + expect(getDailyCacheConfigHash()).toBe(savingsOnly) + + setPriceOverrides({ 'price-hash-model': { input: 1, output: 2 } }) + const firstCombined = getDailyCacheConfigHash() + + setPriceOverrides({ 'price-hash-model': { input: 3, output: 2 } }) + const secondCombined = getDailyCacheConfigHash() + + expect(firstCombined).not.toBe(savingsOnly) + expect(secondCombined).not.toBe(savingsOnly) + expect(secondCombined).not.toBe(firstCombined) + }) +}) + describe('calculateCost - OMP names produce non-zero cost', () => { it('calculates cost for anthropic--claude-4.6-opus', () => { expect(calculateCost('anthropic--claude-4.6-opus', 1000, 200, 0, 0, 0)).toBeGreaterThan(0)