end

open-sse/utils/usageTracking.js

@@ -141,6 +141,68 @@ export function normalizeUsage(usage) {
   return normalized;
 }
 
+/**
+ * Canonicalize usage into ONE storage/cost convention so token counts and cost
+ * are consistent across providers:
+ *   prompt_tokens               = total input INCLUDING cache read + cache creation
+ *   cached_tokens               = cache-read portion (subset of prompt_tokens)
+ *   cache_creation_input_tokens = cache-write portion (subset of prompt_tokens)
+ *   completion_tokens, reasoning_tokens, total_tokens
+ *
+ * Discriminator: Claude reports cache_read_input_tokens with a prompt that
+ * EXCLUDES cache, so we fold cache into prompt. OpenAI/Gemini report
+ * cached_tokens already counted inside prompt, so we pass through. Idempotent:
+ * once folded the output carries cached_tokens (not cache_read_input_tokens),
+ * so re-running takes the passthrough branch and does not double-add.
+ *
+ * @param {object} usage - a normalizeUsage()-shaped object
+ * @returns {object|null} canonical token object, or null for invalid input
+ */
+export function canonicalizeUsage(usage) {
+  if (!usage || typeof usage !== "object" || Array.isArray(usage)) return null;
+
+  const num = (v) => (Number.isFinite(Number(v)) ? Number(v) : 0);
+  const completion = num(usage.completion_tokens ?? usage.output_tokens);
+  const reasoning = num(usage.reasoning_tokens);
+  // Fall back to the nested prompt_tokens_details.cache_creation_tokens shape
+  // (buildUsage()'s OpenAI-forwarding format) when the top-level field is
+  // absent, so callers that pass a buildUsage() object through don't silently
+  // drop cache_creation.
+  const cacheCreation = num(usage.cache_creation_input_tokens ?? usage.prompt_tokens_details?.cache_creation_tokens);
+
+  let prompt = num(usage.prompt_tokens ?? usage.input_tokens);
+  let cached;
+
+  // Claude path: prompt excludes cache; cache_read_input_tokens and/or
+  // cache_creation_input_tokens are separate. A cache-miss "first write" only
+  // carries cache_creation_input_tokens (no cache_read_input_tokens yet), so
+  // check both fields — otherwise a first-write request falls through to the
+  // OpenAI passthrough branch below and cache_creation never gets folded in.
+  // Guard on the absence of `cached_tokens`: our own canonical output always
+  // sets that key (even to 0), so re-running canonicalizeUsage on an already-
+  // folded result takes the passthrough branch instead of folding again.
+  if (usage.cached_tokens === undefined &&
+      (usage.cache_read_input_tokens !== undefined || usage.cache_creation_input_tokens !== undefined)) {
+    cached = num(usage.cache_read_input_tokens);
+    prompt = prompt + cached + cacheCreation;
+  } else {
+    // OpenAI/Gemini path (or already-canonical input): prompt already includes cached_tokens.
+    cached = num(usage.cached_tokens);
+  }
+
+  const result = {
+    prompt_tokens: prompt,
+    completion_tokens: completion,
+    // Recompute rather than pass through: when the fold branch ran above,
+    // an upstream total_tokens (cache-exclusive) would otherwise be stale.
+    total_tokens: prompt + completion,
+    cached_tokens: cached,
+    cache_creation_input_tokens: cacheCreation,
+  };
+  if (reasoning > 0) result.reasoning_tokens = reasoning;
+  return result;
+}
+
 /**
  * Check if usage has valid token data
  * Valid = has at least one token field with value > 0
@@ -171,6 +233,19 @@ export function hasValidUsage(usage) {
 export function extractUsage(chunk) {
   if (!chunk || typeof chunk !== "object") return null;
 
+  // Claude format (message_start event): carries input_tokens + cache_read +
+  // cache_creation. message_delta later carries only the final output_tokens,
+  // so callers must MERGE (mergeUsage), not overwrite, to keep cache counts.
+  if (chunk.type === "message_start" && chunk.message?.usage && typeof chunk.message.usage === "object") {
+    const u = chunk.message.usage;
+    return normalizeUsage({
+      prompt_tokens: u.input_tokens || 0,
+      completion_tokens: u.output_tokens || 0,
+      cache_read_input_tokens: u.cache_read_input_tokens,
+      cache_creation_input_tokens: u.cache_creation_input_tokens
+    });
+  }
+
   // Claude format (message_delta event)
   if (chunk.type === "message_delta" && chunk.usage && typeof chunk.usage === "object") {
     return normalizeUsage({
@@ -232,6 +307,27 @@ export function extractUsage(chunk) {
   return null;
 }
 
+// Field-wise max-merge of two usage objects. Anthropic splits usage across
+// events: message_start has real input+cache (output is a placeholder 1),
+// message_delta has the real cumulative output (input/cache absent). Max keeps
+// the meaningful value from each without clobbering. Idempotent for other
+// providers that emit a single complete usage object.
+export function mergeUsage(prev, next) {
+  if (!prev) return next || null;
+  if (!next) return prev;
+  const merged = { ...prev };
+  for (const [k, v] of Object.entries(next)) {
+    // typeof NaN === "number" — guard with Number.isFinite so one malformed
+    // chunk can't poison the whole accumulation (Math.max(x, NaN) is NaN).
+    if (typeof v === "number" && Number.isFinite(v)) {
+      merged[k] = Math.max(typeof merged[k] === "number" ? merged[k] : 0, v);
+    } else if (v && typeof v === "object") {
+      merged[k] = v; // nested details objects: take latest
+    }
+  }
+  return merged;
+}
+
 /**
  * Estimate input tokens from request body
  * Calculate total body size for more accurate estimation