model-router/tests/pricing.test.ts

import type { On } from 'claude-code'
import { expect, mock, test, type Engine } from 'claude-code/testing'

import {
  addOverhead,
  addWriterRun,
  actualTotal,
  counterfactual,
  emptyTask,
  parsePrices,
  priceRun,
  saved,
  savingsReport,
  tokenizerFactor,
} from '../hooks/pricing'

const U = (i: number, o: number, w = 0, r = 0) => ({
  input_tokens: i,
  output_tokens: o,
  cache_creation_input_tokens: w,
  cache_read_input_tokens: r,
})
const M = 1_000_000
const round = (n: number) => Math.round(n * 1e6) / 1e6

test('prices: each part of the usage at its own rate', async () => {
  // Sonnet 5.5: in 2, out 10, write 2.5, read 0.1 per MTok
  expect(priceRun('claude-sonnet-5-5', U(M, M, M, M)).actual).toBe(2 + 10 + 2.5 + 0.1)
  // the same tokens at Opus 5.5: 4 + 20 + 5 + 0.2
  expect(round(priceRun('claude-sonnet-5-5', U(M, M, M, M)).opus)).toBe(29.2)
  expect(priceRun('claude-opus-5-5', U(M, 0)).actual).toBe(4)
  expect(priceRun('claude-haiku-5-5', U(M, 0)).actual).toBe(0.1)
  expect(priceRun('claude-haiku-4-5', U(M, 0)).actual).toBe(1)
})

test('prices: the most specific model wins, unknown models are flagged', async () => {
  expect(priceRun('claude-opus-5-5-20260601', U(M, 0)).label).toBe('Opus 5.5')
  expect(priceRun('claude-opus-5', U(M, 0)).label).toBe('Opus 5')
  expect(priceRun('claude-sonnet-5', U(M, 0)).label).toBe('Sonnet 5')
  expect(priceRun('gpt-9', U(M, 0)).known).toBe(false)
})

test('older tokenizers are scaled before pricing as Opus 5.5', async () => {
  expect(tokenizerFactor('claude-haiku-4-5')).toBe(1.3)
  expect(tokenizerFactor('claude-sonnet-5-5')).toBe(1)
  expect(tokenizerFactor('claude-haiku-5-5')).toBe(1)
  expect(round(priceRun('claude-haiku-4-5', U(M, 0)).opus)).toBe(5.2)
})

test('a task below opus: every writer token at opus price is the comparison', async () => {
  let t = emptyTask('clamp')
  t = addWriterRun(t, 'claude-sonnet-5-5', U(M, 0))
  t = addWriterRun(t, 'claude-sonnet-5-5', U(M, 0)) // resumed round
  t = addOverhead(t, 'claude-opus-5-5', U(M / 2, 0)) // verifier
  expect(t.writerActual).toBe(4)
  expect(counterfactual(t)).toBe(8)
  expect(actualTotal(t)).toBe(6)
  expect(saved(t)).toBe(2)
})

test('a chain that ended on opus: the opus run is the comparison, failed rounds are loss', async () => {
  let t = emptyTask('clamp')
  t = addWriterRun(t, 'claude-sonnet-5-5', U(M, 0)) // failed twice
  t = addWriterRun(t, 'claude-opus-5-5', U(M, 0)) // rewrite
  expect(counterfactual(t)).toBe(4)
  expect(actualTotal(t)).toBe(6)
  expect(saved(t)).toBe(-2)
})

test('unpriced models are listed, never counted as free', async () => {
  const t = addWriterRun(emptyTask('x'), 'mystery-1', U(M, M))
  expect(t.unpriced).toEqual(['mystery-1'])
  expect(t.writerActual).toBe(0)
  expect(savingsReport({ T1: t }, 5)).toContain('[no price: mystery-1]')
})

test('price overrides come first; bad JSON falls back to list prices', async () => {
  const table = parsePrices('[{"match":"sonnet-5-5","price":{"in":1,"out":1,"write":1,"read":1}}]')
  expect(priceRun('claude-sonnet-5-5', U(M, 0), table).actual).toBe(1)
  expect(priceRun('claude-opus-5-5', U(M, 0), table).actual).toBe(4)
  expect(priceRun('claude-sonnet-5-5', U(M, 0), parsePrices('{oops')).actual).toBe(2)
})

test('the report: newest first, totals over every task, filter by task', async () => {
  const a = addWriterRun(emptyTask('first'), 'claude-sonnet-5-5', U(M, 0))
  const b = addWriterRun(emptyTask('second'), 'claude-sonnet-5-5', U(2 * M, 0))
  const r = savingsReport({ T1: a, T2: b }, 1)
  expect(r.indexOf('T2')).toBeGreaterThan(-1)
  expect(r).not.toContain('T1 ')
  expect(r).toMatch(/total\s+\$6\.00\s+\$12\.00\s+\$6\.00\s+50%/)
  expect(savingsReport({ T1: a, T2: b }, 9, 'T1')).toContain('first')
  expect(savingsReport({ T1: a, T2: b }, 9, 'T1')).not.toContain('second')
})

// ---------- end to end ----------

const USAGE = { input_tokens: 1000, output_tokens: 10, cache_creation_input_tokens: 0, cache_read_input_tokens: 0 }
const ok = (stdout: string) => ({ value: { exitCode: 0, stdout, stderr: '', isStdoutTruncated: false, isStderrTruncated: false } })
const DEV = '/home/homelab/projects/goodwe/clones/OPT-11-P4-dev'

function world(on: On, models: string[]) {
  mock.store(on)
  const clock = mock.clock(on)
  on('ui.status', async () => ({ value: undefined }))
  on('ui.toast', async () => ({ value: undefined }))
  on('model.complete', async (_$, e) => {
    if (e.system?.startsWith('You route')) return { value: { isAnswered: true as const, text: 'sonnet', usage: USAGE } }
    return { value: { isAnswered: true as const, text: '{"verdict":"PASS"}', usage: USAGE } }
  })
  on('process.run', async (_$, e) => {
    const a = e.argv.join(' ')
    if (a.startsWith('git merge-base')) return ok('base\n')
    if (a.startsWith('git diff')) return ok('diff')
    return ok('')
  })
  on('agent.spawn', async (_$, e) => {
    models.push(String(e.model))
    return { model: e.model === 'sonnet' ? 'claude-sonnet-5-5' : String(e.model), agentId: `a${models.length}` }
  })
  on('turn.complete', async (_$, e) => ({ text: e.answer }))
  on('prompt.submit', async (_$, e) => ({ text: e.text }))
  return clock
}

async function end($: Engine, agentId: string, model: string, input: number) {
  await ($.turn as unknown as { complete: (e: unknown) => Promise<unknown> }).complete({
    answer: 'done',
    durationMs: 1,
    isAborted: false,
    turnId: 't',
    agentId,
    reason: 'answer',
    usage: { model, input_tokens: input, output_tokens: 0, cache_creation_input_tokens: 0, cache_read_input_tokens: 0 },
  })
}

test('end to end: a routed dev agent is priced on the model it ran on, a kept review agent is not counted', async ($, on) => {
  const models: string[] = []
  const clock = world(on, models)
  const dev = await $.agent.spawn({ prompt: `Clone: ${DEV}/\nExpected: x`, description: 'P4 dev' } as never)
  await end($, dev.agentId as string, 'claude-sonnet-5-5', M)
  await clock.advance(1)
  const review = await $.agent.spawn({ prompt: 'Clone: /x/clones/OPT-11-P4-review/', description: 'P4 review', model: 'opus' } as never)
  await end($, review.agentId as string, 'claude-opus-5-5', M)
  await clock.advance(1)

  const r = (await $.command.run({ command: 'router', args: 'savings' } as never)).text ?? ''
  expect(r).toContain('P4 dev')
  expect(r).not.toContain('P4 review')
  // writer $2.00 + classifier and verifier (1000 in + 10 out each, Haiku 4.5 then Opus 5.5)
  expect(r).toMatch(/T1\s+\$2\.01\s+\$4\.00/)
  expect(r).toContain('Sonnet 5.5')
})

test('agents started in parallel each get their own task, and every run is counted', async ($, on) => {
  const models: string[] = []
  const clock = world(on, models)
  const ids = (
    await Promise.all(
      ['A', 'D', 'E', 'F', 'G'].map(c => $.agent.spawn({ prompt: `Scope chunk ${c}: list the open cards`, description: `Scope chunk ${c}` } as never)),
    )
  ).map(r => r.agentId as string)
  await Promise.all(ids.map(id => end($, id, 'claude-sonnet-5-5', M)))
  await clock.advance(1)

  const r = (await $.command.run({ command: 'router', args: 'savings' } as never)).text ?? ''
  for (const c of ['A', 'D', 'E', 'F', 'G']) expect(r).toContain(`Scope chunk ${c}`)
  expect(r).toContain('5 of 5 tasks')
  // five writer runs of $2.00 each, all counted
  expect(r).toMatch(/total\s+\$10\.\d\d\s+\$20\.00/)
})

test('a task whose writer has not finished shows as running and stays out of the total', async () => {
  const running = addOverhead(emptyTask('MAINT-14 dev'), 'claude-haiku-4-5', U(M, 0)) // $1 of classifier, no writer yet
  const done = addWriterRun(emptyTask('done'), 'claude-sonnet-5-5', U(M, 0))
  const r = savingsReport({ T1: done, T2: running }, 5)
  expect(r).toMatch(/T2 .* running /)
  expect(r).toMatch(/total\s+\$2\.00\s+\$4\.00/)
})