Replace hardcoded subtraction chain with registry-driven decomposition

packages/js/src/__tests__/calcPrice.test.ts

@@ -41,18 +41,12 @@ describe('Core Price Calculation Function', () => {
 
       const result = calcPrice(usage, modelPrice)
 
-      const uncachedInputTokens = 1000 - 200 - 100
-      const uncachedInputPrice = (uncachedInputTokens * 1.0) / 1_000_000
-      const cacheWritePrice = (200 * 0.5) / 1_000_000
-      const cacheReadPrice = (100 * 0.1) / 1_000_000
-      const inputPrice = uncachedInputPrice + cacheWritePrice + cacheReadPrice
-      const outputPrice = (500 * 2.0) / 1_000_000
-      const totalPrice = inputPrice + outputPrice
-      expect(result).toMatchObject({
-        input_price: inputPrice,
-        output_price: outputPrice,
-        total_price: totalPrice,
-      })
+      // input leaf: 1000 - 200 - 100 = 700 uncached text, 100 cache_read, 200 cache_write
+      // input: (700*1.0 + 100*0.1 + 200*0.5) / 1_000_000 = 0.00081
+      // output: (500*2.0) / 1_000_000 = 0.001
+      expect(result.input_price).toBeCloseTo(0.00081, 10)
+      expect(result.output_price).toBeCloseTo(0.001, 10)
+      expect(result.total_price).toBeCloseTo(0.00181, 10)
     })
 
     it('should handle audio tokens correctly', () => {

packages/js/src/__tests__/decompose.test.ts

@@ -0,0 +1,150 @@
+import { describe, expect, it } from 'vitest'
+
+import type { Usage } from '../types'
+
+import { computeLeafValues, isDescendantOrSelf } from '../decompose'
+import { getUnit, TOKENS_FAMILY } from '../units'
+
+describe('Containment', () => {
+  it('self is descendant-or-self', () => {
+    const unit = getUnit('input_mtok')
+    expect(isDescendantOrSelf(unit, unit)).toBe(true)
+  })
+
+  it('child is descendant', () => {
+    expect(isDescendantOrSelf(getUnit('input_mtok'), getUnit('cache_read_mtok'))).toBe(true)
+  })
+
+  it('parent is NOT descendant of child', () => {
+    expect(isDescendantOrSelf(getUnit('cache_read_mtok'), getUnit('input_mtok'))).toBe(false)
+  })
+
+  it('grandchild is descendant', () => {
+    expect(isDescendantOrSelf(getUnit('input_mtok'), getUnit('cache_read_audio_mtok'))).toBe(true)
+  })
+
+  it('sibling not descendant', () => {
+    expect(isDescendantOrSelf(getUnit('cache_read_mtok'), getUnit('cache_write_mtok'))).toBe(false)
+  })
+
+  it('different direction not descendant', () => {
+    expect(isDescendantOrSelf(getUnit('input_mtok'), getUnit('output_mtok'))).toBe(false)
+  })
+})
+
+describe('Leaf Values', () => {
+  const family = TOKENS_FAMILY
+
+  it('simple text model', () => {
+    const priced = new Set(['input_mtok', 'output_mtok'])
+    const usage = { input_tokens: 1000, output_tokens: 500 }
+    expect(computeLeafValues(priced, usage, family)).toEqual({ input_mtok: 1000, output_mtok: 500 })
+  })
+
+  it('with cache', () => {
+    const priced = new Set(['cache_read_mtok', 'cache_write_mtok', 'input_mtok', 'output_mtok'])
+    const usage = { cache_read_tokens: 200, cache_write_tokens: 100, input_tokens: 1000, output_tokens: 500 }
+    expect(computeLeafValues(priced, usage, family)).toEqual({
+      cache_read_mtok: 200,
+      cache_write_mtok: 100,
+      input_mtok: 700,
+      output_mtok: 500,
+    })
+  })
+
+  it('with audio', () => {
+    const priced = new Set(['input_audio_mtok', 'input_mtok', 'output_mtok'])
+    const usage = { input_audio_tokens: 300, input_tokens: 1000, output_tokens: 500 }
+    expect(computeLeafValues(priced, usage, family)).toEqual({
+      input_audio_mtok: 300,
+      input_mtok: 700,
+      output_mtok: 500,
+    })
+  })
+
+  it('lattice: cache_read_audio carved from both', () => {
+    const priced = new Set(['cache_read_audio_mtok', 'cache_read_mtok', 'input_audio_mtok', 'input_mtok'])
+    const usage = {
+      cache_audio_read_tokens: 50,
+      cache_read_tokens: 200,
+      input_audio_tokens: 300,
+      input_tokens: 1000,
+    }
+    expect(computeLeafValues(priced, usage, family)).toEqual({
+      cache_read_audio_mtok: 50,
+      cache_read_mtok: 150,
+      input_audio_mtok: 250,
+      input_mtok: 550,
+    })
+  })
+
+  it('unpriced audio stays in catch-all', () => {
+    const priced = new Set(['cache_read_mtok', 'input_mtok', 'output_mtok'])
+    const usage = {
+      cache_read_tokens: 200,
+      input_audio_tokens: 300,
+      input_tokens: 1000,
+      output_tokens: 500,
+    }
+    expect(computeLeafValues(priced, usage, family)).toEqual({
+      cache_read_mtok: 200,
+      input_mtok: 800,
+      output_mtok: 500,
+    })
+  })
+
+  it('unpriced cache stays in catch-all', () => {
+    const priced = new Set(['input_mtok', 'output_mtok'])
+    const usage = { cache_read_tokens: 200, input_tokens: 1000, output_tokens: 500 }
+    expect(computeLeafValues(priced, usage, family)).toEqual({
+      input_mtok: 1000,
+      output_mtok: 500,
+    })
+  })
+
+  it('negative leaf raises error', () => {
+    const priced = new Set(['cache_read_mtok', 'input_mtok'])
+    const usage = { cache_read_tokens: 200, input_tokens: 100 }
+    expect(() => computeLeafValues(priced, usage, family)).toThrow(/Negative leaf value.*input_mtok/)
+  })
+
+  it('full 7-unit model', () => {
+    const priced = new Set([
+      'cache_read_audio_mtok',
+      'cache_read_mtok',
+      'cache_write_mtok',
+      'input_audio_mtok',
+      'input_mtok',
+      'output_audio_mtok',
+      'output_mtok',
+    ])
+    const usage = {
+      cache_audio_read_tokens: 50,
+      cache_read_tokens: 200,
+      cache_write_tokens: 100,
+      input_audio_tokens: 300,
+      input_tokens: 1000,
+      output_audio_tokens: 150,
+      output_tokens: 800,
+    }
+    expect(computeLeafValues(priced, usage, family)).toEqual({
+      cache_read_audio_mtok: 50,
+      cache_read_mtok: 150,
+      cache_write_mtok: 100,
+      input_audio_mtok: 250,
+      input_mtok: 450,
+      output_audio_mtok: 150,
+      output_mtok: 650,
+    })
+  })
+
+  it('accepts Usage object (optional fields)', () => {
+    const priced = new Set(['cache_read_mtok', 'input_mtok', 'output_mtok'])
+    const usage: Usage = { cache_read_tokens: 200, input_tokens: 1000, output_tokens: 500 }
+    expect(computeLeafValues(priced, usage as Record<string, unknown>, family)).toEqual({
+      cache_read_mtok: 200,
+      input_mtok: 800,
+      output_mtok: 500,
+    })
+  })
+})

packages/js/src/__tests__/units.test.ts

@@ -0,0 +1,64 @@
+import { describe, expect, it } from 'vitest'
+
+import { FIELD_TO_UNIT, getFamily, getUnit, TOKENS_FAMILY } from '../units'
+
+describe('Unit Registry', () => {
+  it('should load the tokens family', () => {
+    const family = getFamily('tokens')
+    expect(family.id).toBe('tokens')
+    expect(family.per).toBe(1_000_000)
+  })
+
+  it('should have 20 token units', () => {
+    const family = getFamily('tokens')
+    expect(Object.keys(family.units)).toHaveLength(20)
+  })
+
+  it('should look up a unit by ID', () => {
+    const unit = getUnit('input_mtok')
+    expect(unit.familyId).toBe('tokens')
+    expect(unit.usageKey).toBe('input_tokens')
+    expect(unit.dimensions).toEqual({ direction: 'input' })
+  })
+
+  it('should throw on unknown unit', () => {
+    expect(() => getUnit('nonexistent')).toThrow()
+  })
+
+  it('should throw on unknown family', () => {
+    expect(() => getFamily('nonexistent')).toThrow()
+  })
+
+  it('should have all 7 currently-used units', () => {
+    for (const unitId of [
+      'input_mtok',
+      'output_mtok',
+      'cache_read_mtok',
+      'cache_write_mtok',
+      'input_audio_mtok',
+      'cache_read_audio_mtok',
+      'output_audio_mtok',
+    ]) {
+      expect(TOKENS_FAMILY.units[unitId]).toBeDefined()
+    }
+  })
+
+  it('should map cache_audio_read_mtok field to cache_read_audio_mtok unit', () => {
+    expect(FIELD_TO_UNIT.cache_audio_read_mtok).toBe('cache_read_audio_mtok')
+  })
+
+  it('should have correct usage_key for cache_read_audio_mtok', () => {
+    const unit = getUnit('cache_read_audio_mtok')
+    expect(unit.usageKey).toBe('cache_audio_read_tokens')
+  })
+
+  it('should validate all unit dimensions against family dimensions', () => {
+    const family = getFamily('tokens')
+    for (const unit of Object.values(family.units)) {
+      for (const [dimKey, dimVal] of Object.entries(unit.dimensions)) {
+        expect(family.dimensions[dimKey]).toBeDefined()
+        expect(family.dimensions[dimKey]).toContain(dimVal)
+      }
+    }
+  })
+})

packages/js/src/data.ts

@@ -9432,7 +9432,9 @@ export const data: Provider[] = [
           starts_with: 'gpt-4o-audio-preview',
         },
         context_window: 128000,
+        price_comments: 'input_mtok set equal to input_audio_mtok (audio-only input model, catch-all required by ancestor coverage rule)',
         prices: {
+          input_mtok: 2.5,
           output_mtok: 10,
           input_audio_mtok: 2.5,
         },
@@ -9483,7 +9485,9 @@ export const data: Provider[] = [
         match: {
           starts_with: 'gpt-4o-mini-audio',
         },
+        price_comments: 'input_mtok set equal to input_audio_mtok (audio-only input model, catch-all required by ancestor coverage rule)',
         prices: {
+          input_mtok: 0.15,
           output_mtok: 0.6,
           input_audio_mtok: 0.15,
         },
@@ -9518,8 +9522,11 @@ export const data: Provider[] = [
         match: {
           equals: 'gpt-4o-mini-tts',
         },
+        price_comments:
+          'output_mtok set equal to output_audio_mtok (audio-only output model, catch-all required by ancestor coverage rule)',
         prices: {
           input_mtok: 0.6,
+          output_mtok: 12,
           output_audio_mtok: 12,
         },
       },

packages/js/src/decompose.ts

@@ -0,0 +1,65 @@
+import type { UnitDef, UnitFamily } from './units'
+
+/**
+ * True if candidate's dimensions are a (non-strict) superset of ancestor's.
+ */
+export function isDescendantOrSelf(ancestor: UnitDef, candidate: UnitDef): boolean {
+  if (ancestor.familyId !== candidate.familyId) return false
+  return Object.entries(ancestor.dimensions).every(([k, v]) => candidate.dimensions[k] === v)
+}
+
+/**
+ * Get a usage value by key from a plain object. Returns 0 for missing/undefined/null.
+ */
+function getUsageValue(usage: Record<string, unknown>, key: string): number {
+  const val = usage[key]
+  return typeof val === 'number' ? val : 0
+}
+
+/**
+ * Compute leaf values for each priced unit via Mobius inversion on the containment poset.
+ *
+ * Only priced units participate. Unpriced units' tokens stay in the nearest
+ * priced ancestor's catch-all. Throws on negative leaf values.
+ *
+ * Precondition: if a unit is priced, all its ancestors must also be priced
+ * (ancestor coverage rule). Violating this produces silently incorrect results.
+ */
+export function computeLeafValues(pricedUnitIds: Set<string>, usage: Record<string, unknown>, family: UnitFamily): Record<string, number> {
+  const result: Record<string, number> = {}
+
+  for (const unitId of pricedUnitIds) {
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    const unit = family.units[unitId]!
+    const targetDepth = Object.keys(unit.dimensions).length
+
+    let leafValue = 0
+    for (const otherId of pricedUnitIds) {
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      const other = family.units[otherId]!
+      if (!isDescendantOrSelf(unit, other)) continue
+      const depthDiff = Object.keys(other.dimensions).length - targetDepth
+      const coefficient = (-1) ** depthDiff
+      leafValue += coefficient * getUsageValue(usage, other.usageKey)
+    }
+
+    if (leafValue < 0) {
+      const involved = [...pricedUnitIds]
+        .filter((oid) => {
+          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+          const o = family.units[oid]!
+          return isDescendantOrSelf(unit, o) && getUsageValue(usage, o.usageKey) !== 0
+        })
+        .map((oid) => {
+          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+          const o = family.units[oid]!
+          return `${o.usageKey}=${String(getUsageValue(usage, o.usageKey))}`
+        })
+      throw new Error(`Negative leaf value (${String(leafValue)}) for ${unitId}: inconsistent usage values: ${involved.join(', ')}`)
+    }
+
+    result[unitId] = leafValue
+  }
+
+  return result
+}