diff --git a/src/agents/plugins/claude/__tests__/statusline-installer.test.ts b/src/agents/plugins/claude/__tests__/statusline-installer.test.ts index 14a670185..0620e6cfa 100644 --- a/src/agents/plugins/claude/__tests__/statusline-installer.test.ts +++ b/src/agents/plugins/claude/__tests__/statusline-installer.test.ts @@ -25,8 +25,6 @@ vi.mock('@/utils/security.js', () => ({ })); describe('statusline-installer', () => { - // Derived paths go through path.join in production, so compute expected values the same - // way to get the correct separator on each OS (backslashes on Windows). const CLAUDE_HOME = '/home/testuser/claude'; const SCRIPT_PATH = join(CLAUDE_HOME, 'codemie-budget-status.js'); const LEGACY_SCRIPT_PATH = join(CLAUDE_HOME, 'codemie-statusline.mjs'); @@ -35,6 +33,24 @@ describe('statusline-installer', () => { let fsp: typeof import('fs/promises'); let fsMod: typeof import('fs'); + const realReadFile = async (p: string) => { + const { readFile } = await vi.importActual('fs/promises'); + return readFile(p, 'utf-8'); + }; + + const mockScriptSources = (settings?: string) => + vi.mocked(fsp.readFile).mockImplementation((async (p: string) => { + const path = String(p).split('\\').join('/'); + if (path.includes('plugin/statusline.mjs')) { + return realReadFile('src/agents/plugins/claude/plugin/statusline.mjs'); + } + if (path.includes('plugin/transcript-cost.mjs')) { + return realReadFile('src/agents/plugins/claude/plugin/transcript-cost.mjs'); + } + if (settings !== undefined) return settings; + throw new Error('ENOENT'); + }) as never); + beforeEach(async () => { vi.resetModules(); vi.resetAllMocks(); @@ -48,9 +64,7 @@ describe('statusline-installer', () => { describe('installStatusline', () => { it('deploys the script and reports alreadyConfigured=false when settings.json has no statusLine yet', async () => { - vi.mocked(fsp.readFile) - .mockResolvedValueOnce('#!/usr/bin/env node\n// statusline' as any) // script source - .mockResolvedValueOnce(JSON.stringify({ theme: 'dark' }) as any); // settings.json + mockScriptSources(JSON.stringify({ theme: 'dark' })); vi.mocked(fsMod.existsSync).mockReturnValue(true); vi.mocked(fsp.writeFile).mockResolvedValue(undefined); vi.mocked(fsp.chmod).mockResolvedValue(undefined); @@ -69,9 +83,7 @@ describe('statusline-installer', () => { }); it('reports alreadyConfigured=true (and still refreshes settings) when statusLine already exists', async () => { - vi.mocked(fsp.readFile) - .mockResolvedValueOnce('// script' as any) - .mockResolvedValueOnce(JSON.stringify({ statusLine: { type: 'command', command: 'node "/old.js"' } }) as any); + mockScriptSources(JSON.stringify({ statusLine: { type: 'command', command: 'node "/old.js"' } })); vi.mocked(fsMod.existsSync).mockReturnValue(true); vi.mocked(fsp.writeFile).mockResolvedValue(undefined); vi.mocked(fsp.chmod).mockResolvedValue(undefined); @@ -82,8 +94,75 @@ describe('statusline-installer', () => { expect(result.alreadyConfigured).toBe(true); }); + it('generates a self-contained script with no package imports', async () => { + mockScriptSources(JSON.stringify({})); + vi.mocked(fsMod.existsSync).mockReturnValue(true); + vi.mocked(fsp.writeFile).mockResolvedValue(undefined); + vi.mocked(fsp.chmod).mockResolvedValue(undefined); + + const { installStatusline } = await import('../statusline-installer.js'); + await installStatusline(); + + const scriptWrite = vi.mocked(fsp.writeFile).mock.calls.find(([p]) => p === SCRIPT_PATH); + const script = scriptWrite![1] as string; + expect(script).not.toContain("from './transcript-cost.mjs'"); + expect(script).not.toContain('__CODEMIE_TRANSCRIPT_COST_IMPORT__'); + expect(script).not.toContain("from 'module'"); + expect(script.startsWith('#!')).toBe(true); + }); + + it('prices a transcript with the rates it is given', async () => { + mockScriptSources(JSON.stringify({})); + vi.mocked(fsMod.existsSync).mockReturnValue(true); + vi.mocked(fsp.writeFile).mockResolvedValue(undefined); + vi.mocked(fsp.chmod).mockResolvedValue(undefined); + + const { installStatusline } = await import('../statusline-installer.js'); + await installStatusline(); + const script = vi.mocked(fsp.writeFile).mock.calls.find(([p]) => p === SCRIPT_PATH)![1] as string; + + const dataUrl = `data:text/javascript;base64,${Buffer.from( + script.replace('function transcriptCostUSD', 'export function transcriptCostUSD') + ).toString('base64')}`; + const { transcriptCostUSD } = await import(/* @vite-ignore */ dataUrl); + + const usage = { input_tokens: 1_000_000, output_tokens: 0 }; + const line = JSON.stringify({ requestId: 'r1', message: { id: 'm1', model: 'claude-sonnet-4-5', usage } }); + const prices = { 'claude-sonnet-4-5': { input: 3, output: 15, cacheRead: 0.3, cacheCreation: 3.75, cacheWrite1h: 6 } }; + expect(transcriptCostUSD([line, line].join('\n'), prices)).toBeCloseTo(3, 6); + }); + + it('fails loudly when a source is missing its substitution marker', async () => { + mockScriptSources(JSON.stringify({})); + vi.mocked(fsp.readFile).mockImplementationOnce((async () => '#!/usr/bin/env node') as never); + vi.mocked(fsMod.existsSync).mockReturnValue(true); + vi.mocked(fsp.writeFile).mockResolvedValue(undefined); + + const { installStatusline } = await import('../statusline-installer.js'); + await expect(installStatusline()).rejects.toThrow(/missing its/); + expect(fsp.writeFile).not.toHaveBeenCalled(); + }); + + it('refuses to inline transcript-cost.mjs if it grows an import', async () => { + mockScriptSources(JSON.stringify({})); + vi.mocked(fsp.readFile).mockImplementation((async (p: string) => { + const path = String(p).split('\\').join('/'); + if (path.includes('plugin/transcript-cost.mjs')) return "import x from 'y';\nexport function f() {}"; + if (path.includes('plugin/statusline.mjs')) { + return realReadFile('src/agents/plugins/claude/plugin/statusline.mjs'); + } + return JSON.stringify({}); + }) as never); + vi.mocked(fsMod.existsSync).mockReturnValue(true); + vi.mocked(fsp.writeFile).mockResolvedValue(undefined); + + const { installStatusline } = await import('../statusline-installer.js'); + await expect(installStatusline()).rejects.toThrow(/must have no imports/); + expect(fsp.writeFile).not.toHaveBeenCalled(); + }); + it('creates ~/.claude when it does not exist', async () => { - vi.mocked(fsp.readFile).mockResolvedValueOnce('// script' as any); + mockScriptSources(); vi.mocked(fsMod.existsSync).mockReturnValueOnce(false).mockReturnValueOnce(false); vi.mocked(fsp.mkdir).mockResolvedValue(undefined); vi.mocked(fsp.writeFile).mockResolvedValue(undefined); @@ -96,9 +175,7 @@ describe('statusline-installer', () => { }); it('throws ConfigurationError and does not overwrite malformed settings.json', async () => { - vi.mocked(fsp.readFile) - .mockResolvedValueOnce('// script' as any) - .mockResolvedValueOnce('{ bad json' as any); + mockScriptSources('{ bad json'); vi.mocked(fsMod.existsSync).mockReturnValue(true); vi.mocked(fsp.writeFile).mockResolvedValue(undefined); vi.mocked(fsp.chmod).mockResolvedValue(undefined); diff --git a/src/agents/plugins/claude/plugin/__tests__/statusline.test.ts b/src/agents/plugins/claude/plugin/__tests__/statusline.test.ts index 0ef6df897..f32a0fe37 100644 --- a/src/agents/plugins/claude/plugin/__tests__/statusline.test.ts +++ b/src/agents/plugins/claude/plugin/__tests__/statusline.test.ts @@ -11,7 +11,10 @@ import { resolveBudget, isMainModule, ctxBar, + resolveCost, + resolveGatewayPrices, } from '../statusline.mjs'; +import { transcriptCostUSD, sumTranscriptUsage, priceFromGatewayCost, lookupPrice } from '../transcript-cost.mjs'; const YELLOW = '\x1b[0;33m'; const GREEN = '\x1b[0;32m'; @@ -299,3 +302,319 @@ describe('isMainModule', () => { expect(isMainModule(undefined, url)).toBe(false); }); }); + +describe('transcript cost', () => { + const PRICES = { + 'claude-sonnet-4-5': { input: 3, output: 15, cacheRead: 0.3, cacheCreation: 3.75, cacheWrite1h: 6 }, + 'claude-haiku-4-5': { input: 1, output: 5, cacheRead: 0.1, cacheCreation: 1.25, cacheWrite1h: 2 }, + }; + const row = (id: string, usage: Record, extra: Record = {}) => + JSON.stringify({ type: 'assistant', requestId: `req-${id}`, message: { id, model: 'claude-sonnet-4-5', usage }, ...extra }); + + it('counts one API response once when Claude Code splits it across rows', () => { + const usage = { input_tokens: 1_000_000, output_tokens: 0 }; + const text = [row('msg-1', usage), row('msg-1', usage)].join('\n'); + expect(transcriptCostUSD(text, PRICES)).toBeCloseTo(3, 6); + }); + + it('keeps the most complete row when a partial one shares the key', () => { + const partial = row('msg-1', { input_tokens: 10, output_tokens: 0 }); + const full = row('msg-1', { input_tokens: 1_000_000, output_tokens: 0 }); + expect(transcriptCostUSD([partial, full].join('\n'), PRICES)).toBeCloseTo(3, 6); + expect(transcriptCostUSD([full, partial].join('\n'), PRICES)).toBeCloseTo(3, 6); + }); + + it('sums distinct responses and splits cache writes by TTL', () => { + const text = [ + row('msg-1', { input_tokens: 1_000_000, output_tokens: 0 }), + row('msg-2', { + input_tokens: 0, output_tokens: 0, + cache_creation_input_tokens: 1_000_000, + cache_creation: { ephemeral_1h_input_tokens: 1_000_000 }, + }), + ].join('\n'); + expect(transcriptCostUSD(text, PRICES)).toBeCloseTo(9, 6); + }); + + it('ignores synthetic rows, usage-less rows, and a partially-written trailing line', () => { + const text = [ + row('msg-1', { input_tokens: 1_000_000, output_tokens: 0 }), + JSON.stringify({ type: 'assistant', message: { id: 'x', model: '', usage: { input_tokens: 999_999_999 } } }), + JSON.stringify({ type: 'user', message: { role: 'user' } }), + '{"partial": ', + ].join('\n'); + expect(transcriptCostUSD(text, PRICES)).toBeCloseTo(3, 6); + }); + + it('counts unkeyable rows individually', () => { + const line = JSON.stringify({ message: { model: 'claude-sonnet-4-5', usage: { input_tokens: 1_000_000, output_tokens: 0 } } }); + expect(transcriptCostUSD([line, line].join('\n'), PRICES)).toBeCloseTo(6, 6); + }); + + it('returns null when nothing is priceable, so the caller can fall back', () => { + expect(transcriptCostUSD('', PRICES)).toBeNull(); + expect(transcriptCostUSD(row('msg-1', { input_tokens: 5 }), null)).toBeNull(); + expect(transcriptCostUSD(JSON.stringify({ message: { model: 'not-listed', usage: { input_tokens: 5 } } }), PRICES)).toBeNull(); + }); + + it('groups usage per model', () => { + const text = [ + row('msg-1', { input_tokens: 100, output_tokens: 0 }), + JSON.stringify({ requestId: 'r2', message: { id: 'msg-2', model: 'claude-haiku-4-5', usage: { input_tokens: 200, output_tokens: 0 } } }), + ].join('\n'); + const byModel = sumTranscriptUsage(text); + expect(byModel.get('claude-sonnet-4-5').input).toBe(100); + expect(byModel.get('claude-haiku-4-5').input).toBe(200); + }); +}); + +describe('resolveCost', () => { + const transcript = JSON.stringify({ + requestId: 'r1', + message: { id: 'm1', model: 'claude-sonnet-4-5', usage: { input_tokens: 1_000_000, output_tokens: 0 } }, + }); + + it('prices from the transcript rather than the cost Claude Code reports', async () => { + const cost = await resolveCost( + { transcriptPath: '/fake/t.jsonl', cost: 99 }, + { + readFile: async () => transcript, + gatewayPrices: { 'claude-sonnet-4-5': { input: 3, output: 15, cacheRead: 0.3, cacheCreation: 3.75, cacheWrite1h: 6 } }, + } + ); + expect(cost).toBeCloseTo(3, 6); + }); + + it('falls back to the reported cost when the transcript is unreadable', async () => { + const cost = await resolveCost( + { transcriptPath: '/fake/t.jsonl', cost: 42 }, + { readFile: async () => { throw new Error('ENOENT'); } } + ); + expect(cost).toBe(42); + }); + + it('falls back when the transcript has nothing priceable (never a misleading $0)', async () => { + const cost = await resolveCost( + { transcriptPath: '/fake/t.jsonl', cost: 42 }, + { readFile: async () => '' } + ); + expect(cost).toBe(42); + }); + + it('falls back when no transcript path is provided', async () => { + expect(await resolveCost({ transcriptPath: null, cost: 7 })).toBe(7); + }); +}); + +describe('priceFromGatewayCost', () => { + it('converts per-token gateway rates to per-million and derives the 1h write rate', () => { + const price = priceFromGatewayCost({ + input: 0.0000055, + output: 0.0000275, + cache_read_input_token_cost: 5.5e-7, + cache_creation_input_token_cost: 0.000006875, + }); + expect(price).toEqual({ input: 5.5, output: 27.5, cacheRead: 0.55, cacheCreation: 6.875, cacheWrite1h: 11 }); + }); + + it('derives cache rates from input when the gateway omits them', () => { + const price = priceFromGatewayCost({ input: 0.000003, output: 0.000015 }); + expect(price.cacheRead).toBeCloseTo(0.3, 6); + expect(price.cacheCreation).toBeCloseTo(3.75, 6); + }); + + it('returns null without a usable input rate', () => { + expect(priceFromGatewayCost(undefined)).toBeNull(); + expect(priceFromGatewayCost({ output: 0.00001 })).toBeNull(); + }); +}); + +describe('pricing source', () => { + const line = JSON.stringify({ + requestId: 'r1', + message: { id: 'm1', model: 'claude-sonnet-4-5', usage: { input_tokens: 1_000_000, output_tokens: 0 } }, + }); + + it('prices with the tenant rate the gateway reports', () => { + const gateway = { 'claude-sonnet-4-5': { input: 3.3, output: 16.5, cacheRead: 0.33, cacheCreation: 4.125, cacheWrite1h: 6.6 } }; + expect(transcriptCostUSD(line, gateway)).toBeCloseTo(3.3, 6); + }); + + it('returns null without gateway rates so the caller keeps the reported cost', () => { + expect(transcriptCostUSD(line, null)).toBeNull(); + expect(transcriptCostUSD(line, { 'other-model': { input: 99, output: 99, cacheRead: 0, cacheCreation: 0, cacheWrite1h: 0 } })).toBeNull(); + }); +}); + +describe('resolveGatewayPrices', () => { + const config = JSON.stringify({ + activeProfile: 'p', + profiles: { p: { provider: 'ai-run-sso', codeMieUrl: 'https://cm.test' } }, + }); + const models = [ + { deployment_name: 'claude-opus-5', cost: { input: 0.0000055, output: 0.0000275 } }, + { deployment_name: 'no-cost-model' }, + ]; + + it('fetches, keys by lowercased deployment name, and caches', async () => { + const writeFile = vi.fn().mockResolvedValue(undefined); + const readFile = vi.fn(async (p: string) => + String(p).includes('model-prices-cache') ? Promise.reject(new Error('ENOENT')) : config + ); + const fetchImpl = vi.fn().mockResolvedValue({ ok: true, json: async () => models }); + + const prices = await resolveGatewayPrices({ + readFile, writeFile, fetchImpl, getAuthHeadersImpl: async () => ({ cookie: 'x' }), + }); + + expect(prices!['claude-opus-5'].input).toBeCloseTo(5.5, 6); + expect(prices!['no-cost-model']).toBeUndefined(); + expect(fetchImpl.mock.calls[0][0]).toBe('https://cm.test/code-assistant-api/v1/llm_models?include_all=true'); + expect(writeFile).toHaveBeenCalled(); + }); + + it('serves a fresh cache without hitting the network', async () => { + const fetchImpl = vi.fn(); + const cached = { + schema: 1, + tenant: 'https://cm.test/code-assistant-api', + ts: Date.now(), + value: { 'claude-opus-5': { input: 1 } }, + }; + const prices = await resolveGatewayPrices({ + readFile: async (p: string) => + String(p).includes('model-prices-cache') ? JSON.stringify(cached) : config, + writeFile: vi.fn(), + fetchImpl, + getAuthHeadersImpl: async () => ({ cookie: 'x' }), + }); + + expect(prices).toEqual(cached.value); + expect(fetchImpl).not.toHaveBeenCalled(); + }); + + it('returns null when unauthenticated, on a failed response, or with no config', async () => { + const readFile = vi.fn(async (p: string) => + String(p).includes('model-prices-cache') ? Promise.reject(new Error('ENOENT')) : config + ); + const base = { readFile, writeFile: vi.fn() }; + + expect(await resolveGatewayPrices({ ...base, fetchImpl: vi.fn(), getAuthHeadersImpl: async () => null })).toBeNull(); + expect(await resolveGatewayPrices({ + ...base, + fetchImpl: vi.fn().mockResolvedValue({ ok: false }), + getAuthHeadersImpl: async () => ({ cookie: 'x' }), + })).toBeNull(); + expect(await resolveGatewayPrices({ + readFile: async () => { throw new Error('ENOENT'); }, + writeFile: vi.fn(), + fetchImpl: vi.fn(), + getAuthHeadersImpl: async () => ({ cookie: 'x' }), + })).toBeNull(); + }); + + it('returns null when the network call throws', async () => { + const readFile = vi.fn(async (p: string) => + String(p).includes('model-prices-cache') ? Promise.reject(new Error('ENOENT')) : config + ); + const prices = await resolveGatewayPrices({ + readFile, + writeFile: vi.fn(), + fetchImpl: vi.fn().mockRejectedValue(new Error('offline')), + getAuthHeadersImpl: async () => ({ cookie: 'x' }), + }); + expect(prices).toBeNull(); + }); +}); + +describe('resolveGatewayPrices — provider gating', () => { + const run = (profile: Record) => { + const fetchImpl = vi.fn().mockResolvedValue({ ok: true, json: async () => [] }); + return resolveGatewayPrices({ + readFile: async (p: string) => + String(p).includes('model-prices-cache') + ? Promise.reject(new Error('ENOENT')) + : JSON.stringify({ activeProfile: 'p', profiles: { p: profile } }), + writeFile: vi.fn(), + fetchImpl, + getAuthHeadersImpl: async () => ({ cookie: 'secret' }), + }).then((prices) => ({ prices, fetchImpl })); + }; + + it('never contacts a subscription provider, whose baseUrl is not a CodeMie host', async () => { + for (const provider of ['anthropic-subscription', 'moonshot-subscription']) { + const { prices, fetchImpl } = await run({ + provider, codeMieUrl: 'https://cm.test', baseUrl: 'https://api.anthropic.com', + }); + expect(prices).toBeNull(); + expect(fetchImpl).not.toHaveBeenCalled(); + } + }); + + it('skips a gateway provider with no codeMieUrl', async () => { + const { prices, fetchImpl } = await run({ provider: 'bedrock', baseUrl: 'https://bedrock.aws' }); + expect(prices).toBeNull(); + expect(fetchImpl).not.toHaveBeenCalled(); + }); + + it('builds the endpoint from codeMieUrl, tolerating a trailing slash', async () => { + const { fetchImpl } = await run({ provider: 'ai-run-sso', codeMieUrl: 'https://cm.test/' }); + expect(fetchImpl.mock.calls[0][0]).toBe('https://cm.test/code-assistant-api/v1/llm_models?include_all=true'); + }); +}); + +describe('resolveGatewayPrices — keying and cache identity', () => { + const gatewayProfile = (codeMieUrl: string) => + JSON.stringify({ activeProfile: 'p', profiles: { p: { provider: 'ai-run-sso', codeMieUrl } } }); + + const run = (config: string, cached?: unknown, models: unknown[] = []) => { + const fetchImpl = vi.fn().mockResolvedValue({ ok: true, json: async () => models }); + const writeFile = vi.fn().mockResolvedValue(undefined); + return resolveGatewayPrices({ + readFile: async (p: string) => + String(p).includes('model-prices-cache') + ? (cached === undefined ? Promise.reject(new Error('ENOENT')) : JSON.stringify(cached)) + : config, + writeFile, + fetchImpl, + getAuthHeadersImpl: async () => ({ cookie: 'x' }), + }).then((prices) => ({ prices, fetchImpl, writeFile })); + }; + + it('keys dotted deployment names so transcript lookups match', async () => { + const { prices } = await run(gatewayProfile('https://cm.test'), undefined, [ + { deployment_name: 'us.anthropic.claude-sonnet-4-6', cost: { input: 0.000003, output: 0.000015 } }, + ]); + expect(lookupPrice('us.anthropic.claude-sonnet-4-6', prices)).not.toBeNull(); + }); + + it('keeps a model whose input rate is legitimately zero', async () => { + const { prices } = await run(gatewayProfile('https://cm.test'), undefined, [ + { deployment_name: 'free-model', cost: { input: 0, output: 0 } }, + ]); + expect(prices!['free-model']).toBeDefined(); + expect(prices!['free-model'].input).toBe(0); + }); + + it('does not double-append the api base when codeMieUrl already carries it', async () => { + const { fetchImpl } = await run(gatewayProfile('https://cm.test/code-assistant-api')); + expect(fetchImpl.mock.calls[0][0]).toBe('https://cm.test/code-assistant-api/v1/llm_models?include_all=true'); + }); + + it('ignores a cache entry belonging to a different tenant', async () => { + const stale = { schema: 1, tenant: 'https://other.test/code-assistant-api', ts: Date.now(), value: { m: { input: 1 } } }; + const { prices, fetchImpl } = await run(gatewayProfile('https://cm.test'), stale, [ + { deployment_name: 'mine', cost: { input: 0.000002 } }, + ]); + expect(fetchImpl).toHaveBeenCalled(); + expect(prices!.mine).toBeDefined(); + expect(prices!.m).toBeUndefined(); + }); + + it('stamps the tenant on the cache it writes', async () => { + const { writeFile } = await run(gatewayProfile('https://cm.test'), undefined, [ + { deployment_name: 'm', cost: { input: 0.000002 } }, + ]); + expect(JSON.parse(writeFile.mock.calls[0][1] as string).tenant).toBe('https://cm.test/code-assistant-api'); + }); +}); diff --git a/src/agents/plugins/claude/plugin/statusline.mjs b/src/agents/plugins/claude/plugin/statusline.mjs index 16dfd705e..fd0f93f63 100644 --- a/src/agents/plugins/claude/plugin/statusline.mjs +++ b/src/agents/plugins/claude/plugin/statusline.mjs @@ -10,6 +10,8 @@ import fs from 'fs/promises'; import os from 'os'; import path from 'path'; import { fileURLToPath } from 'url'; +/* __CODEMIE_TRANSCRIPT_COST_IMPORT__ */ +import { transcriptCostUSD, priceFromGatewayCost, normalizeModelKey } from './transcript-cost.mjs'; const HOME = process.env.CODEMIE_HOME || path.join(os.homedir(), '.codemie'); const CACHE_FILE = path.join(HOME, 'budget-cache.json'); @@ -17,6 +19,9 @@ const CONFIG_FILE = path.join(HOME, 'codemie-cli.config.json'); const CREDS_DIR = path.join(HOME, 'credentials'); const CACHE_TTL_MS = 60_000; const CACHE_SCHEMA = 2; // bump when the cache.value shape changes, to discard stale pre-upgrade entries +const PRICES_CACHE_FILE = path.join(HOME, 'model-prices-cache.json'); +const PRICES_CACHE_TTL_MS = 6 * 60 * 60 * 1000; +const PRICES_CACHE_SCHEMA = 1; const ENCRYPTION_KEY = (() => { const id = os.hostname() + os.platform() + os.arch(); @@ -97,9 +102,106 @@ export function extractBasicInfo(ctx) { tokOut: ctx?.context_window?.total_output_tokens ?? null, cost: ctx?.cost?.total_cost_usd ?? null, durationMs: ctx?.cost?.total_duration_ms ?? null, + transcriptPath: ctx?.transcript_path ?? null, }; } +const PROVIDERS = new Set(['ai-run-sso', 'jwt', 'litellm', 'bedrock']); + +/** Mirrors `ensureApiBase` in providers/core/codemie-auth-helpers.ts, which this cannot import. */ +function ensureApiBase(rawUrl) { + const base = String(rawUrl).replace(/\/+$/, ''); + return /\/code-assistant-api(\/|$)/i.test(base) ? base : `${base}/code-assistant-api`; +} + +/** + * Per-model rates from the gateway's `/v1/llm_models`, keyed by {@link normalizeModelKey}. + * Cached on disk per tenant. Returns null for non-gateway providers or when unavailable, so + * callers keep the cost Claude Code reports. Never throws. + */ +export async function resolveGatewayPrices({ + readFile = fs.readFile, + writeFile = fs.writeFile, + fetchImpl = fetch, + getAuthHeadersImpl = getAuthHeaders, +} = {}) { + let profile; + try { + const config = JSON.parse(await readFile(CONFIG_FILE, 'utf8')); + profile = config.profiles?.[config.activeProfile]; + } catch { + return null; + } + + const { codeMieUrl, provider } = profile ?? {}; + if (!codeMieUrl || !PROVIDERS.has(provider)) { + return null; + } + + // Rates are per-tenant, so a cache entry is only valid for the gateway that produced it. + const tenant = ensureApiBase(codeMieUrl).toLowerCase(); + try { + const cache = JSON.parse(await readFile(PRICES_CACHE_FILE, 'utf8')); + if ( + cache.schema === PRICES_CACHE_SCHEMA + && cache.tenant === tenant + && Date.now() - cache.ts < PRICES_CACHE_TTL_MS + ) { + return cache.value; + } + } catch {} + + try { + const headers = await getAuthHeadersImpl(codeMieUrl); + if (!headers) { + return null; + } + const res = await fetchImpl(`${tenant}/v1/llm_models?include_all=true`, { headers }); + if (!res.ok) { + return null; + } + const body = await res.json(); + const models = Array.isArray(body) ? body : (body?.data ?? []); + const prices = {}; + for (const model of models) { + const name = model?.deployment_name || model?.base_name; + const price = name && priceFromGatewayCost(model.cost); + if (price) { + prices[normalizeModelKey(name)] = price; + } + } + if (!Object.keys(prices).length) { + return null; + } + try { + await writeFile( + PRICES_CACHE_FILE, + JSON.stringify({ schema: PRICES_CACHE_SCHEMA, tenant, ts: Date.now(), value: prices }), + 'utf8' + ); + } catch {} + return prices; + } catch { + return null; + } +} + +export async function resolveCost( + { transcriptPath, cost }, + { readFile = fs.readFile, gatewayPrices = null } = {} +) { + if (!transcriptPath) { + return cost; + } + try { + const text = await readFile(transcriptPath, 'utf8'); + const derived = transcriptCostUSD(text, gatewayPrices); + return derived ?? cost; + } catch { + return cost; + } +} + export function formatDuration(ms) { if (typeof ms !== 'number' || Number.isNaN(ms) || ms < 0) return null; const mins = Math.floor(ms / 60000); @@ -256,9 +358,14 @@ export async function main() { } const branchPromise = basic.cwd ? gitBranch(basic.cwd) : Promise.resolve(''); - const [budgetResult, branch] = await Promise.all([resolveBudget(), branchPromise]); - - process.stdout.write(buildStatusLine({ ...basic, branch, ...budgetResult })); + const [budgetResult, branch, gatewayPrices] = await Promise.all([ + resolveBudget(), + branchPromise, + resolveGatewayPrices(), + ]); + const cost = await resolveCost(basic, { gatewayPrices }); + + process.stdout.write(buildStatusLine({ ...basic, branch, cost, ...budgetResult })); } // Compares decoded paths (not raw strings) so this correctly matches even when the diff --git a/src/agents/plugins/claude/plugin/transcript-cost.mjs b/src/agents/plugins/claude/plugin/transcript-cost.mjs new file mode 100644 index 000000000..c953ffefd --- /dev/null +++ b/src/agents/plugins/claude/plugin/transcript-cost.mjs @@ -0,0 +1,119 @@ +export function priceFromGatewayCost(cost) { + if (cost?.input == null) { + return null; + } + const input = cost.input * 1_000_000; + const output = (cost.output ?? 0) * 1_000_000; + const cacheRead = + cost.cache_read_input_token_cost != null + ? cost.cache_read_input_token_cost * 1_000_000 + : input * 0.1; + const cacheCreation = + cost.cache_creation_input_token_cost != null + ? cost.cache_creation_input_token_cost * 1_000_000 + : input * 1.25; + return { input, output, cacheRead, cacheCreation, cacheWrite1h: input * 2 }; +} + +/** Key models identically on both sides of the price map: lowercased, dots folded to dashes. */ +export function normalizeModelKey(model) { + return String(model || "") + .toLowerCase() + .replace(/\./g, "-"); +} + +/** + * Rates for `model` from `gatewayPrices`, or null when the tenant's gateway does not list it + * (caller treats the model as unpriced). + */ +export function lookupPrice(model, gatewayPrices) { + return gatewayPrices?.[normalizeModelKey(model)] ?? null; +} + +/** USD for one model's usage. Rates are per 1,000,000 tokens. */ +function costForUsage(usage, price) { + const tokens1h = Math.min(usage.cacheCreation1h, usage.cacheCreation); + const tokens5m = usage.cacheCreation - tokens1h; + return ( + (usage.input * price.input + + usage.output * price.output + + usage.cacheRead * price.cacheRead + + tokens1h * price.cacheWrite1h + + tokens5m * price.cacheCreation) / + 1_000_000 + ); +} + +/** + * Usage per model, counting each API response once. Claude Code writes one response across + * several JSONL rows (`thinking` + `text`, or `text` + `tool_use`) that each repeat the full + * usage, so rows are keyed by `message.id` + `requestId` and the most complete row wins. + */ +export function sumTranscriptUsage(text) { + const byKey = new Map(); + const unkeyed = []; + for (const line of String(text).split("\n")) { + if (!line) continue; + let row; + try { + row = JSON.parse(line); + } catch { + continue; + } + const u = row?.message?.usage; + if (!u) continue; + const model = row.message.model; + if (!model || model === "") continue; + const usage = { + input: u.input_tokens ?? 0, + output: u.output_tokens ?? 0, + cacheRead: u.cache_read_input_tokens ?? 0, + cacheCreation: u.cache_creation_input_tokens ?? 0, + cacheCreation1h: u.cache_creation?.ephemeral_1h_input_tokens ?? 0, + }; + const weight = + usage.input + usage.output + usage.cacheRead + usage.cacheCreation; + const id = row.message.id; + const reqId = row.requestId; + if (!id && !reqId) { + unkeyed.push({ model, usage }); + continue; + } + const key = `${id ?? ""}::${reqId ?? ""}`; + const current = byKey.get(key); + if (!current || weight > current.weight) { + byKey.set(key, { model, usage, weight }); + } + } + + const byModel = new Map(); + for (const { model, usage } of [...byKey.values(), ...unkeyed]) { + const acc = byModel.get(model); + if (!acc) { + byModel.set(model, { ...usage }); + continue; + } + acc.input += usage.input; + acc.output += usage.output; + acc.cacheRead += usage.cacheRead; + acc.cacheCreation += usage.cacheCreation; + acc.cacheCreation1h += usage.cacheCreation1h; + } + return byModel; +} + +/** + * Session cost in USD, or null when no usage row prices against `gatewayPrices` so the caller + * can fall back rather than render $0. + */ +export function transcriptCostUSD(text, gatewayPrices) { + let total = 0; + let priced = false; + for (const [model, usage] of sumTranscriptUsage(text)) { + const price = lookupPrice(model, gatewayPrices); + if (!price) continue; + priced = true; + total += costForUsage(usage, price); + } + return priced ? total : null; +} diff --git a/src/agents/plugins/claude/statusline-installer.ts b/src/agents/plugins/claude/statusline-installer.ts index df979fd9d..75c1287c1 100644 --- a/src/agents/plugins/claude/statusline-installer.ts +++ b/src/agents/plugins/claude/statusline-installer.ts @@ -20,15 +20,58 @@ export interface InstallStatuslineResult { alreadyConfigured: boolean; } +const TRANSCRIPT_COST_IMPORT_MARKER = '/* __CODEMIE_TRANSCRIPT_COST_IMPORT__ */'; + +function escapeRegExp(literal: string): string { + return literal.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); +} + +/** + * Build the deployed statusline as one self-contained file, inlining `transcript-cost.mjs` + * into `statusline.mjs`. The script runs as `node ` from ~/.claude after this process + * exits, so it cannot import from the package. + * + * @throws ConfigurationError when statusline.mjs has lost its substitution marker. + */ +async function buildStatuslineScript(): Promise { + const pluginDir = join(getDirname(import.meta.url), 'plugin'); + const [statusline, transcriptCost] = await Promise.all([ + readFile(join(pluginDir, 'statusline.mjs'), 'utf-8'), + readFile(join(pluginDir, 'transcript-cost.mjs'), 'utf-8'), + ]); + + if (!statusline.includes(TRANSCRIPT_COST_IMPORT_MARKER)) { + throw new ConfigurationError( + `Statusline source statusline.mjs is missing its ${TRANSCRIPT_COST_IMPORT_MARKER} marker` + ); + } + + // Inlining flattens the module into statusline.mjs's scope, so an import here would be + // dropped and surface only as a ReferenceError in the deployed copy. + if (/^import\s/m.test(transcriptCost)) { + throw new ConfigurationError( + 'Statusline source transcript-cost.mjs must have no imports; it is inlined into the deployed script' + ); + } + + const inlinedCost = transcriptCost.replace(/^export /gm, ''); + + const body = statusline.replace( + new RegExp(`${escapeRegExp(TRANSCRIPT_COST_IMPORT_MARKER)}\\n?import[^\\n]*'\\./transcript-cost\\.mjs';`), + '' + ); + + const lines = body.split('\n'); + const shebang = lines[0].startsWith('#!') ? `${lines.shift()}\n` : ''; + return `${shebang}// GENERATED by codemie install statusline — do not edit. Source: plugin/statusline.mjs\n${inlinedCost}\n${lines.join('\n')}`; +} + export async function installStatusline(): Promise { const claudeHome = resolveHomeDir('.claude'); const scriptPath = join(claudeHome, SCRIPT_FILENAME); const settingsPath = join(claudeHome, 'settings.json'); - const scriptContent = await readFile( - join(getDirname(import.meta.url), 'plugin/statusline.mjs'), - 'utf-8' - ); + const scriptContent = await buildStatuslineScript(); if (!existsSync(claudeHome)) { await mkdir(claudeHome, { recursive: true });