AgentHub/server/lib/token-usage.js

/**
 * @file Shared helpers for normalizing Claude transcript `usage` records into
 * per-bucket token tallies. Used by BOTH ingestion paths — the live server-side
 * parser (`server/lib/transcript-cache.js`) and the history importer
 * (`scripts/import-history.js`) — so the two stay in lockstep.
 *
 * A "bucket" is the unit cost is computed against: tokens are grouped by
 * (model, speed, inference_geo, service_tier) because those four dimensions
 * change the per-token RATE (fast mode, US data residency, Batch API). The
 * dimensions are normalized to the small set of values that actually move
 * price; anything unknown collapses to the standard/global default so old
 * transcripts (which lack `speed` / `inference_geo` / `cache_creation`
 * breakdown / `server_tool_use`) price exactly as they did before.
 *
 * A request whose whole prompt (input + cache read + cache write) is above
 * LONG_CONTEXT_TOKENS is ALSO counted into the bucket's `long*` fields — a
 * subset of the totals, like `cacheWrite1h` — so cost can price that part at
 * the model's long-context rates.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const { LONG_CONTEXT_TOKENS } = require("./pricing-constants");

// Separator for composite bucket keys — U+0001 (SOH) cannot occur in a model id.
const BUCKET_SEP = String.fromCharCode(1);

/** Pricing-relevant speed. Anything other than the fast research-preview tier is standard. */
function normalizeSpeed(usage) {
  return usage && usage.speed === "fast" ? "fast" : "standard";
}

/**
 * Pricing-relevant inference geography. Only US-pinned routing carries the 1.1x
 * data-residency premium; "global", "not_available", and absent all map to the
 * standard "global" rate.
 */
function normalizeGeo(usage) {
  return usage && usage.inference_geo === "us" ? "us" : "global";
}

/** Pricing-relevant service tier. Only "batch" changes the rate (50% off). */
function normalizeTier(usage) {
  return usage && usage.service_tier === "batch" ? "batch" : "standard";
}

/** Composite bucket key — stable string usable as an object property. */
function bucketKey(model, speed, geo, tier) {
  return [model, speed, geo, tier].join(BUCKET_SEP);
}

/** A zeroed bucket carrying its four pricing dimensions. */
function emptyBucket(model, speed, geo, tier) {
  return {
    model,
    speed,
    geo,
    tier,
    input: 0,
    output: 0,
    cacheRead: 0,
    cacheWrite: 0, // TOTAL ephemeral cache-creation tokens (5m + 1h)
    cacheWrite1h: 0, // subset of cacheWrite that is the 1h tier; 5m = cacheWrite - cacheWrite1h
    webSearch: 0, // server_tool_use.web_search_requests (billed per 1k)
    webFetch: 0, // server_tool_use.web_fetch_requests (free; tracked for visibility)
    codeExec: 0, // server_tool_use.code_execution_requests (time-billed; estimated)
    // Subsets of input/output/cacheRead/cacheWrite/cacheWrite1h spent by
    // requests whose prompt was above LONG_CONTEXT_TOKENS.
    longInput: 0,
    longOutput: 0,
    longCacheRead: 0,
    longCacheWrite: 0,
    longCacheWrite1h: 0,
  };
}

/**
 * Pull the numeric token / request fields out of a single `usage` record.
 * Tolerant of the older shape: when `cache_creation` breakdown is absent the
 * whole cache-write amount is treated as 5m (cacheWrite1h = 0), and a missing
 * `server_tool_use` yields zero tool requests.
 */
function extractUsageFields(usage) {
  if (!usage || typeof usage !== "object") {
    return {
      input: 0,
      output: 0,
      cacheRead: 0,
      cacheWrite: 0,
      cacheWrite1h: 0,
      webSearch: 0,
      webFetch: 0,
      codeExec: 0,
      longInput: 0,
      longOutput: 0,
      longCacheRead: 0,
      longCacheWrite: 0,
      longCacheWrite1h: 0,
    };
  }
  const cc =
    usage.cache_creation && typeof usage.cache_creation === "object" ? usage.cache_creation : null;
  const ephem5m = cc ? cc.ephemeral_5m_input_tokens || 0 : 0;
  const ephem1h = cc ? cc.ephemeral_1h_input_tokens || 0 : 0;
  // Prefer the explicit total; fall back to the breakdown sum when only that is present.
  const cacheWrite =
    usage.cache_creation_input_tokens != null
      ? usage.cache_creation_input_tokens || 0
      : ephem5m + ephem1h;
  // Never let the 1h subset exceed the recorded total (guards malformed records).
  const cacheWrite1h = Math.min(ephem1h, cacheWrite);
  const stu =
    usage.server_tool_use && typeof usage.server_tool_use === "object"
      ? usage.server_tool_use
      : null;
  const input = usage.input_tokens || 0;
  const output = usage.output_tokens || 0;
  const cacheRead = usage.cache_read_input_tokens || 0;
  const long = input + cacheRead + cacheWrite > LONG_CONTEXT_TOKENS;
  return {
    input,
    output,
    cacheRead,
    cacheWrite,
    cacheWrite1h,
    webSearch: stu ? stu.web_search_requests || 0 : 0,
    webFetch: stu ? stu.web_fetch_requests || 0 : 0,
    codeExec: stu ? stu.code_execution_requests || 0 : 0,
    longInput: long ? input : 0,
    longOutput: long ? output : 0,
    longCacheRead: long ? cacheRead : 0,
    longCacheWrite: long ? cacheWrite : 0,
    longCacheWrite1h: long ? cacheWrite1h : 0,
  };
}

/**
 * A bucket's long-context subsets in `long_*_tokens` column order, or null
 * when it has none (the caller then skips the extra write).
 */
function longArgs(b) {
  const args = [
    b.longInput,
    b.longOutput,
    b.longCacheRead,
    b.longCacheWrite,
    b.longCacheWrite1h,
  ].map((n) => n || 0);
  return args.some((n) => n > 0) ? args : null;
}

/** Add the numeric fields of `src` into `target` in place. */
function accumulateBucket(target, src) {
  target.input += src.input || 0;
  target.output += src.output || 0;
  target.cacheRead += src.cacheRead || 0;
  target.cacheWrite += src.cacheWrite || 0;
  target.cacheWrite1h += src.cacheWrite1h || 0;
  target.webSearch += src.webSearch || 0;
  target.webFetch += src.webFetch || 0;
  target.codeExec += src.codeExec || 0;
  target.longInput = (target.longInput || 0) + (src.longInput || 0);
  target.longOutput = (target.longOutput || 0) + (src.longOutput || 0);
  target.longCacheRead = (target.longCacheRead || 0) + (src.longCacheRead || 0);
  target.longCacheWrite = (target.longCacheWrite || 0) + (src.longCacheWrite || 0);
  target.longCacheWrite1h = (target.longCacheWrite1h || 0) + (src.longCacheWrite1h || 0);
  return target;
}

module.exports = {
  BUCKET_SEP,
  normalizeSpeed,
  normalizeGeo,
  normalizeTier,
  bucketKey,
  emptyBucket,
  extractUsageFields,
  accumulateBucket,
  longArgs,
};