voyage/lib/stats/token-usage.mjs

// lib/stats/token-usage.mjs
// SKAL-2 — pure token-usage parser + cache-aware cost derivation for Voyage
// observability. Captures MAIN-CONTEXT token/cost from the Claude Code
// transcript (transcript_path). Each assistant record carries message.usage
// with a per-REQUEST snapshot.
//
// Verified 2026-06-26 against a real local transcript (216 lines): 70
// assistant records collapse to 31 distinct requestIds — records duplicate
// per requestId (streamed snapshots), so we dedup by requestId keeping the
// LAST occurrence (GH #28197 streaming-placeholder mitigation), then sum
// across requests. The main transcript carried zero isSidechain:true records.
//
// v1 scope = MAIN-CONTEXT only. Sub-agent (swarm) turns live in separate
// agent-*.jsonl siblings and are a documented v2 follow-on — so the
// main-transcript sum UNDER-counts total Voyage cost. Every record is stamped
// scope:'main-context' so no reader mistakes cost_usd for the session total.
//
// The four core functions (parseTranscriptUsage, deriveCost, buildRecord,
// upsertSessionRecord) are PURE (no I/O). captureTokenUsage is the impure
// shell that wires them to the filesystem (read transcript → upsert jsonl).
//
// Zero npm dependencies. Node stdlib only.

import { readFileSync, existsSync, writeFileSync, renameSync, statSync } from 'node:fs';
import { join, dirname } from 'node:path';

// Per-Mtok USD prices, resolved 2026-06-26 via the claude-api skill reference:
// base input/output from the model table; cache rates from the prompt-caching
// doc multipliers (cache_read 0.1x, write_5m 1.25x, write_1h 2.0x of input).
export const PRICE_TABLE = Object.freeze({
  'claude-opus-4-8': Object.freeze({
    input: 5.0,
    output: 25.0,
    cache_read: 0.5,
    cache_write_5m: 6.25,
    cache_write_1h: 10.0,
  }),
});

// Date the PRICE_TABLE values were resolved/verified. Bump when prices change.
export const PRICE_TABLE_VERSION = '2026-06-26';

function num(v) {
  return typeof v === 'number' && Number.isFinite(v) ? v : 0;
}

/**
 * Single pass over transcript JSONL text. Keeps only main-chain assistant
 * records (type==='assistant', isSidechain!==true), dedups by requestId
 * keeping the LAST occurrence (per-request usage snapshots duplicate across
 * streamed records), and sums the four usage token fields. Malformed lines
 * are skipped.
 *
 * @returns {{tokens_input:number, tokens_output:number,
 *            tokens_cache_creation:number, tokens_cache_read:number}}
 */
export function parseTranscriptUsage(text) {
  const lines = (text || '').split('\n');
  // requestId -> usage (last occurrence wins). Records without a requestId
  // get a unique sentinel key so each is still counted exactly once.
  const byRequest = new Map();
  let anon = 0;
  for (const line of lines) {
    const trimmed = line.trim();
    if (trimmed === '') continue;
    let obj;
    try { obj = JSON.parse(trimmed); }
    catch { continue; }
    if (!obj || obj.type !== 'assistant') continue;
    if (obj.isSidechain === true) continue;
    const usage = obj.message && obj.message.usage;
    if (!usage || typeof usage !== 'object') continue;
    const key = typeof obj.requestId === 'string' && obj.requestId
      ? obj.requestId
      : `__anon_${anon++}`;
    byRequest.set(key, usage); // last occurrence wins
  }
  const totals = {
    tokens_input: 0,
    tokens_output: 0,
    tokens_cache_creation: 0,
    tokens_cache_read: 0,
  };
  for (const usage of byRequest.values()) {
    totals.tokens_input += num(usage.input_tokens);
    totals.tokens_output += num(usage.output_tokens);
    totals.tokens_cache_creation += num(usage.cache_creation_input_tokens);
    totals.tokens_cache_read += num(usage.cache_read_input_tokens);
  }
  return totals;
}

/**
 * Cache-aware USD cost. Lumped cache_creation is priced at the 5m write rate
 * (Claude Code's default cache TTL; verified the dominant case on a real
 * transcript — the 1h split is a documented v2 refinement). When the model is
 * absent from the table we REFUSE to estimate: {cost_usd:null, is_estimate:true}.
 *
 * @returns {{cost_usd:number|null, is_estimate:boolean}}
 */
export function deriveCost(totals, model, priceTable = PRICE_TABLE) {
  const price = priceTable && priceTable[model];
  if (!price) return { cost_usd: null, is_estimate: true };
  const t = totals || {};
  const cost =
    (num(t.tokens_input) * price.input +
      num(t.tokens_output) * price.output +
      num(t.tokens_cache_creation) * price.cache_write_5m +
      num(t.tokens_cache_read) * price.cache_read) / 1_000_000;
  return { cost_usd: cost, is_estimate: false };
}

/**
 * Build the flat numeric record. Stamps scope:'main-context' and
 * price_table_version. session_id is written for upsert keying but MUST be
 * stripped at export (CWE-212, enforced by the field allowlist in Step 2).
 *
 * @param {{sessionId:string, model:string, totals:object,
 *          priceTable?:object, now:string}} args
 */
export function buildRecord({ sessionId, model, totals, priceTable = PRICE_TABLE, now }) {
  const t = totals || {};
  const { cost_usd, is_estimate } = deriveCost(t, model, priceTable);
  return {
    ts: now,
    session_id: sessionId,
    scope: 'main-context',
    model,
    tokens_input: num(t.tokens_input),
    tokens_output: num(t.tokens_output),
    tokens_cache_creation: num(t.tokens_cache_creation),
    tokens_cache_read: num(t.tokens_cache_read),
    cost_usd,
    is_estimate,
    price_table_version: PRICE_TABLE_VERSION,
  };
}

/**
 * Pure read-modify-write. Parses existing JSONL lines, REPLACES the line whose
 * session_id matches the record (dropping any further same-session dupes),
 * else APPENDS. Returns the new file text (one line per session, trailing
 * newline). One-line-per-session is what keeps cross-session sums from
 * double-counting (vs append-and-grow).
 */
export function upsertSessionRecord(existingText, record) {
  const lines = (existingText || '').split('\n');
  const out = [];
  let replaced = false;
  for (const line of lines) {
    const trimmed = line.trim();
    if (trimmed === '') continue;
    let obj;
    try { obj = JSON.parse(trimmed); }
    catch { out.push(trimmed); continue; } // preserve unparseable lines verbatim
    if (obj && obj.session_id === record.session_id) {
      if (!replaced) { out.push(JSON.stringify(record)); replaced = true; }
      // else: duplicate same-session line — drop it
    } else {
      out.push(trimmed);
    }
  }
  if (!replaced) out.push(JSON.stringify(record));
  return out.join('\n') + '\n';
}

/**
 * Last main-chain (non-sidechain) assistant model in the transcript. Used to
 * pick the price-table key. Returns null when no model is found (→ deriveCost
 * refuses to estimate). Pure.
 */
export function lastMainChainModel(text) {
  const lines = (text || '').split('\n');
  let model = null;
  for (const line of lines) {
    const trimmed = line.trim();
    if (trimmed === '') continue;
    let obj;
    try { obj = JSON.parse(trimmed); }
    catch { continue; }
    if (!obj || obj.type !== 'assistant' || obj.isSidechain === true) continue;
    const m = obj.message && obj.message.model;
    if (typeof m === 'string' && m) model = m;
  }
  return model;
}

/**
 * Impure capture shell (Step 4). Reads the transcript, derives main-context
 * token totals + cost, and UPSERTS a one-line-per-session record into
 * {dataDir}/token-usage-stats.jsonl via an atomic temp+rename write
 * (EXDEV mitigation: tmp lives in the same dir as the target).
 *
 * Returns the written record, or null when skipped (no path / no dataDir /
 * transcript not a readable regular file). Throws propagate to the caller —
 * the Stop hook wraps this in try/catch so capture stays fail-open.
 *
 * @param {{transcriptPath:string, sessionId:string, dataDir:string,
 *          now?:string}} args
 */
export function captureTokenUsage({ transcriptPath, sessionId, dataDir, now }) {
  if (!transcriptPath || !dataDir) return null;
  let st;
  try { st = statSync(transcriptPath); }
  catch { return null; }
  if (!st.isFile()) return null;

  const text = readFileSync(transcriptPath, 'utf-8');
  const totals = parseTranscriptUsage(text);
  const model = lastMainChainModel(text);
  const record = buildRecord({
    sessionId,
    model,
    totals,
    now: now || new Date().toISOString(),
  });

  const outPath = join(dataDir, 'token-usage-stats.jsonl');
  const existing = existsSync(outPath) ? readFileSync(outPath, 'utf-8') : '';
  const updated = upsertSessionRecord(existing, record);

  const tmpPath = join(dirname(outPath), '.token-usage-stats.jsonl.tmp');
  writeFileSync(tmpPath, updated);
  renameSync(tmpPath, outPath);
  return record;
}