fix(cli): recover Amp usage from tokens.total and ledger-less threads

Jannchie · Jannchie · commit 67c72ef411ab · 2026-07-07T21:24:35.000+09:00
Two dropped-usage gaps vs ccusage:
- Ledger events reading only tokens.input/output ignored tokens.total, so a
  total-only event ({tokens:{total:N}}) computed to zero and was skipped. Fold an
  explicit total that exceeds the itemized parts into billable output.
- A thread with no usageLedger returned []; the current Amp schema stores usage on
  each assistant message. Fall back to per-message usage (ccusage parse_message_usage).
diff --git a/packages/cli/src/adapters/amp.ts b/packages/cli/src/adapters/amp.ts
@@ -50,8 +50,15 @@ async function parseAmpSessionFile(
   const rawEvents = Array.isArray(ledger.events)
     ? ledger.events.filter(isPlainObject) as Record<string, unknown>[]
     : []
+  const state = new SessionParserState(filePath, options, event => baseAmpEvent(event))
+  state.sessionId = sessionId
+
+  // Prefer the usage ledger. A thread with no ledger (older / other Amp schema)
+  // still carries usage on each assistant message, so fall back to that instead of
+  // dropping the whole thread — mirrors ccusage read_thread_file -> parse_message_usage.
   if (rawEvents.length === 0) {
-    return []
+    collectAmpMessageUsage(state, messages, sessionId, threadId, options)
+    return state.events.filter(event => matchesBackfillFilters(event, options))
   }
 
   // Amp ledger entries can be written out of order; sort by timestamp so the
@@ -62,9 +69,6 @@ async function parseAmpSessionFile(
     return ta.localeCompare(tb)
   })
 
-  const state = new SessionParserState(filePath, options, event => baseAmpEvent(event))
-  state.sessionId = sessionId
-
   const firstTs = timestampFrom(stringField(events[0], 'timestamp')) || new Date().toISOString()
   state.ensureSessionStarted(firstTs, 0)
 
@@ -80,7 +84,10 @@ async function parseAmpSessionFile(
     const cache = cacheTokensFor(messages, toMessageId)
     const cacheRead = cache.cacheReadInputTokens
     const cacheWrite = cache.cacheCreationInputTokens
-    const totalTokens = input + output + cacheRead + cacheWrite
+    // Fold an explicit tokens.total that exceeds the itemized parts into billable
+    // output (ccusage apply_total_token_fallback), so total-only ledger events —
+    // {tokens:{total:N}} with no input/output — are counted instead of dropped.
+    const { billableOutput, totalTokens } = foldAmpTotal(input, output, cacheRead, cacheWrite, numberField(tokens, 'total') || 0)
     if (totalTokens <= 0) {
       continue
     }
@@ -91,7 +98,7 @@ async function parseAmpSessionFile(
 
     const metrics: Partial<MetricBag> = {
       tokensInput: (input + cacheRead + cacheWrite) || undefined,
-      tokensOutput: output || undefined,
+      tokensOutput: billableOutput || undefined,
       tokensCachedInput: (cacheRead + cacheWrite) || undefined,
       tokensCacheReadInput: cacheRead || undefined,
       tokensCacheCreationInput: cacheWrite || undefined,
@@ -142,6 +149,109 @@ async function parseAmpSessionFile(
 
 // ── Amp-specific helpers ──
 
+// ccusage apply_total_token_fallback: when an explicit grand total exceeds the sum
+// of the itemized token parts, attribute the shortfall to billable output (and the
+// grand total). It never reduces the parts sum. Returns codetime's billable output
+// (reasoning/extra folded in) and the reconciled grand total.
+function foldAmpTotal(
+  input: number,
+  output: number,
+  cacheRead: number,
+  cacheWrite: number,
+  explicitTotal: number,
+): { billableOutput: number, totalTokens: number } {
+  const partsSum = input + output + cacheRead + cacheWrite
+  const missing = Math.max(0, explicitTotal - partsSum)
+  return { billableOutput: output + missing, totalTokens: partsSum + missing }
+}
+
+// Fallback usage path for Amp threads with no usage ledger: read each assistant
+// message's own `usage` object (model/timestamp/inputTokens/outputTokens/
+// cacheCreationInputTokens/cacheReadInputTokens/totalTokens). Mirrors ccusage
+// parse_message_usage.
+function collectAmpMessageUsage(
+  state: SessionParserState,
+  messages: Record<string, unknown>[],
+  sessionId: string,
+  threadId: string,
+  _options: Record<string, unknown> & { _: string[] },
+): void {
+  const rows: Array<{
+    ts: string
+    model: string | undefined
+    input: number
+    billableOutput: number
+    cacheRead: number
+    cacheWrite: number
+    totalTokens: number
+  }> = []
+  for (const message of messages) {
+    if (stringField(message, 'role') !== 'assistant') {
+      continue
+    }
+    const usage = objectField(message, 'usage')
+    if (Object.keys(usage).length === 0) {
+      continue
+    }
+    const ts = timestampFrom(stringField(usage, 'timestamp')) || timestampFrom(stringField(message, 'timestamp'))
+    if (!ts) {
+      continue
+    }
+    const model = stringField(usage, 'model') || stringField(message, 'model') || undefined
+    const input = numberField(usage, 'inputTokens') || 0
+    const output = numberField(usage, 'outputTokens') || 0
+    const cacheWrite = numberField(usage, 'cacheCreationInputTokens') || 0
+    const cacheRead = numberField(usage, 'cacheReadInputTokens') || 0
+    const { billableOutput, totalTokens } = foldAmpTotal(input, output, cacheRead, cacheWrite, numberField(usage, 'totalTokens') || 0)
+    if (totalTokens <= 0) {
+      continue
+    }
+    rows.push({ ts, model, input, billableOutput, cacheRead, cacheWrite, totalTokens })
+  }
+  if (rows.length === 0) {
+    return
+  }
+  rows.sort((a, b) => a.ts.localeCompare(b.ts))
+  state.ensureSessionStarted(rows[0].ts, 0)
+  let lastTs = rows[0].ts
+  for (const [index, r] of rows.entries()) {
+    lastTs = r.ts
+    state.push(
+      baseAmpEvent({
+        ts: r.ts,
+        type: 'model.usage',
+        sessionId,
+        model: r.model,
+        confidence: 'exact',
+        metrics: {
+          tokensInput: (r.input + r.cacheRead + r.cacheWrite) || undefined,
+          tokensOutput: r.billableOutput || undefined,
+          tokensCachedInput: (r.cacheRead + r.cacheWrite) || undefined,
+          tokensCacheReadInput: r.cacheRead || undefined,
+          tokensCacheCreationInput: r.cacheWrite || undefined,
+          tokensTotal: r.totalTokens,
+          modelCalls: 1,
+        },
+        refs: stringRefs({ threadId }),
+      }),
+      index + 1,
+      'messages',
+      'model.usage',
+    )
+  }
+  state.push(
+    baseAmpEvent({
+      ts: lastTs,
+      type: 'session.ended',
+      sessionId,
+      confidence: 'derived',
+    }),
+    rows.length,
+    'messages',
+    'session',
+  )
+}
+
 function baseAmpEvent(
   event: Omit<CanonicalEvent, 'schemaVersion' | 'source' | 'agent' | 'workspaceId'>,
 ): CanonicalEvent {
diff --git a/packages/cli/test/amp.test.ts b/packages/cli/test/amp.test.ts
@@ -0,0 +1,116 @@
+import type { CanonicalEvent } from '@codetime/shared'
+import assert from 'node:assert/strict'
+import { mkdtemp, writeFile } from 'node:fs/promises'
+import { tmpdir } from 'node:os'
+import path from 'node:path'
+// eslint-disable-next-line test/no-import-node-test -- This repo uses node:test as the runner.
+import { test } from 'node:test'
+import { createAmpAdapter } from '../src/adapters/amp.ts'
+
+const adapter = createAmpAdapter()
+
+async function parse(thread: unknown): Promise<CanonicalEvent[]> {
+  const dir = await mkdtemp(path.join(tmpdir(), 'amp-'))
+  const file = path.join(dir, 'T-x.json')
+  await writeFile(file, JSON.stringify(thread), 'utf8')
+  return adapter.parseSessionFile!(file, { _: [] })
+}
+
+function usageEvents(events: CanonicalEvent[]): CanonicalEvent[] {
+  return events.filter(e => e.type === 'model.usage')
+}
+
+test('parity: ccusage amp falls_back_to_total_tokens_when_amp_parts_are_missing', async () => {
+  // A ledger event carrying only tokens.total (no input/output) must still be
+  // counted, with the total folded into billable output. ccusage asserts
+  // output_tokens == 345 (adapter/amp/parser.rs falls_back_to_total_tokens_...).
+  const usages = usageEvents(await parse({
+    id: 'T-x',
+    usageLedger: {
+      events: [
+        { id: 'event-a', timestamp: '2026-01-02T00:00:00.000Z', model: 'gpt-5', tokens: { total: 345 } },
+      ],
+    },
+  }))
+
+  assert.equal(usages.length, 1)
+  assert.equal(usages[0].metrics?.tokensOutput, 345)
+  assert.equal(usages[0].metrics?.tokensTotal, 345)
+})
+
+test('parity: ccusage amp ledger folds the shortfall of an explicit total into output', async () => {
+  // {input:10, output:20, total:100} → ccusage counts 100 total (extra 70 folded
+  // into billable output); the raw parts alone would be only 30.
+  const usages = usageEvents(await parse({
+    id: 'T-x',
+    usageLedger: {
+      events: [
+        { id: 'event-a', timestamp: '2026-01-02T00:00:00.000Z', model: 'gpt-5', tokens: { input: 10, output: 20, total: 100 } },
+      ],
+    },
+  }))
+
+  assert.equal(usages.length, 1)
+  assert.equal(usages[0].metrics?.tokensInput, 10)
+  assert.equal(usages[0].metrics?.tokensOutput, 90) // 20 + (100 - 30) folded
+  assert.equal(usages[0].metrics?.tokensTotal, 100)
+})
+
+test('parity: ccusage amp reads_usage_from_messages_when_ledger_is_missing', async () => {
+  // No usageLedger → fall back to per-assistant-message usage. ccusage asserts
+  // output_tokens == 178 with cache 986/11372 (parser.rs reads_usage_from_messages_...).
+  const usages = usageEvents(await parse({
+    id: 'T-x',
+    messages: [
+      {
+        role: 'assistant',
+        usage: {
+          model: 'claude-haiku-4-5-20251001',
+          inputTokens: 10,
+          outputTokens: 178,
+          cacheCreationInputTokens: 986,
+          cacheReadInputTokens: 11_372,
+          timestamp: '2026-01-19T11:42:10.652Z',
+        },
+      },
+    ],
+  }))
+
+  assert.equal(usages.length, 1)
+  assert.equal(usages[0].model, 'claude-haiku-4-5-20251001')
+  assert.equal(usages[0].metrics?.tokensOutput, 178)
+  // cache-inclusive input: 10 + 11372 (read) + 986 (write).
+  assert.equal(usages[0].metrics?.tokensInput, 10 + 11_372 + 986)
+  assert.equal(usages[0].metrics?.tokensCacheReadInput, 11_372)
+  assert.equal(usages[0].metrics?.tokensCacheCreationInput, 986)
+})
+
+test('parity: ccusage amp empty_usage_ledger_falls_back_to_message_usage', async () => {
+  // An empty/malformed usageLedger must not suppress the message-usage fallback.
+  const usages = usageEvents(await parse({
+    id: 'T-x',
+    usageLedger: {},
+    messages: [
+      { role: 'assistant', usage: { model: 'gpt-5', outputTokens: 178, timestamp: '2026-01-19T11:42:10.652Z' } },
+    ],
+  }))
+
+  assert.equal(usages.length, 1)
+  assert.equal(usages[0].metrics?.tokensOutput, 178)
+})
+
+test('amp still parses a normal ledger event unchanged', async () => {
+  const usages = usageEvents(await parse({
+    id: 'T-x',
+    usageLedger: {
+      events: [
+        { id: 'e', timestamp: '2026-01-02T00:00:00.000Z', model: 'gpt-5', tokens: { input: 100, output: 50 } },
+      ],
+    },
+  }))
+
+  assert.equal(usages.length, 1)
+  assert.equal(usages[0].metrics?.tokensInput, 100)
+  assert.equal(usages[0].metrics?.tokensOutput, 50)
+  assert.equal(usages[0].metrics?.tokensTotal, 150)
+})
