create-walle 0.9.21 → 0.9.23

package/template/claude-task-manager/lib/session-token-usage.js

@@ -0,0 +1,494 @@
+'use strict';
+
+// Per-session token-usage summary.
+//
+// Produces the numbers shown next to the model-name chip in the top bar:
+//   { total, ctx, ctxWindow, ctxPct, exact, totalLabel, breakdown }
+//
+// Sources, in order of fidelity:
+//   - Claude Code: ~/.claude/projects/<hash>/<id>.jsonl — every assistant entry
+//     carries message.usage (input/output/cache_read/cache_creation tokens). We
+//     sum them for `total` and take the last turn for `ctx`. EXACT.
+//   - Codex: ~/.codex/sessions/.../rollout-*.jsonl — periodic `token_count`
+//     events whose `total_token_usage` is already cumulative and which also
+//     carry `model_context_window`. We read the LAST such event. EXACT.
+//   - Everything else (Cursor/Gemini/Aider/OpenCode/Wall-E): estimate by
+//     tokenizing the parsed message text. Marked `exact:false` so the UI can
+//     render a `~` prefix.
+//
+// Claude files can be tens of MB and we recompute on every turn-finish, so the
+// Claude path scans INCREMENTALLY: JSONL is append-only, so we keep a running
+// total + byte offset per file and only read the bytes appended since last time.
+// Codex events are cumulative, so we only ever read the file tail.
+
+const fs = require('fs');
+
+const DEFAULT_CTX_WINDOW = 200000;
+// Codex: a cumulative token_count event normally sits within the last few KB, so
+// read a small tail first and grow only if it isn't found (huge final tool blob).
+const CODEX_TAIL_BYTES = 64 * 1024;
+const CODEX_TAIL_BYTES_MAX = 4 * 1024 * 1024;
+
+function _contextWindowOverride(value) {
+  const n = Math.round(Number(value) || 0);
+  return Number.isFinite(n) && n > 0 ? n : 0;
+}
+
+// Fallback context windows by model id (Codex carries its own window in-band, so
+// this is only consulted for Claude + the estimate tier). Conservative defaults.
+// `override` is the preferred path for provider/catalog metadata. The static map
+// is only a cold-start safety net for transcripts seen before the registry cache
+// has refreshed.
+function ctxWindowForModel(modelId, override) {
+  const fromRegistry = _contextWindowOverride(override);
+  if (fromRegistry) return fromRegistry;
+  const id = String(modelId || '').toLowerCase();
+  if (!id) return DEFAULT_CTX_WINDOW;
+  if (/\[1m\]|-1m\b|context-1m/.test(id)) return 1000000;
+  // Cold-start defaults only. The live window is read from the model registry
+  // when available, and per-session transcript evidence can still correct it.
+  if (/^claude-(opus-4-8|fable-5|mythos-5)\b/.test(id)) return 1000000;
+  if (/^claude-/.test(id)) return 200000;
+  if (/^gpt-5|^o[34]\b/.test(id)) return 256000;
+  if (/^gpt-4/.test(id)) return 128000;
+  if (/^gemini/.test(id)) return 1000000;
+  if (/^(deepseek|kimi|moonshot|qwen)/.test(id)) return 128000;
+  return DEFAULT_CTX_WINDOW;
+}
+
+// Claude's usable context window is NOT a fixed property of the model id — Claude Code
+// runs a per-session 200K vs 1M ("extended context") mode, so a static model→window map
+// is fundamentally incomplete. Read the window from the session's own authoritative
+// signals instead (the Claude analog of Codex's in-band model_context_window):
+//   • AUTO compact_boundary `preTokens` — Claude Code auto-compacts near the limit, so
+//     its recorded pre-compaction token count is its OWN tight measurement of the window.
+//     This is authoritative: it can correct the static estimate UP *or* DOWN. (Manual
+//     /compact is ignored — the user can compact at any size.)
+//   • Largest single-turn prompt (input+cache_read+cache_creation) — a hard lower bound,
+//     since the API cannot accept a prompt larger than the window. Only raises, never
+//     lowers, the static default (it's a floor, not a measurement of the true window).
+// The evidence is classified to the nearest standard window with a tolerance, so an
+// auto-compact that overshoots a round 1M (e.g. 1,008,316) reads as the 1M tier instead
+// of rounding up to 2M. A fresh session with no evidence falls back to the model map.
+const _CTX_WINDOW_TIERS = [200000, 1000000, 2000000, 5000000, 10000000];
+const _CTX_WINDOW_TOLERANCE = 1.05;
+function _snapToCtxTier(observed) {
+  const obs = Math.max(0, Number(observed) || 0);
+  for (const tier of _CTX_WINDOW_TIERS) { if (obs <= tier * _CTX_WINDOW_TOLERANCE) return tier; }
+  return obs;
+}
+function _effectiveCtxWindow(baseWindow, evidence) {
+  const base = Number(baseWindow) > 0 ? Number(baseWindow) : DEFAULT_CTX_WINDOW;
+  const ev = evidence || {};
+  const autoAt = Math.max(0, Number(ev.autoCompactAt) || 0);
+  const maxPrompt = Math.max(0, Number(ev.maxPrompt) || 0);
+  // Authoritative auto-compact measurement wins outright (up or down).
+  if (autoAt > 0) return _snapToCtxTier(autoAt);
+  // Otherwise a prompt larger than the static estimate proves the estimate is too small.
+  if (maxPrompt > base) return Math.max(_snapToCtxTier(maxPrompt), base);
+  return base;
+}
+
+// Compact human label: 980 -> "980", 1234 -> "1.2k", 1_200_000 -> "1.2M".
+function formatTokens(n) {
+  const v = Math.max(0, Math.round(Number(n) || 0));
+  if (v < 1000) return String(v);
+  if (v < 1_000_000) {
+    const k = v / 1000;
+    return (k < 10 ? k.toFixed(1) : Math.round(k)) + 'k';
+  }
+  const m = v / 1_000_000;
+  return (m < 10 ? m.toFixed(1) : Math.round(m)) + 'M';
+}
+
+function _isClaude(provider) { return String(provider || '').toLowerCase() === 'claude'; }
+function _isCodex(provider) { return String(provider || '').toLowerCase() === 'codex'; }
+
+// Match a brace-balanced JSON object beginning at the `{` that follows `key` at
+// or after `from`. Returns the parsed object or null. Used for Codex events,
+// whose token objects are small.
+function _matchObjectAfter(text, key, from = 0) {
+  const at = text.indexOf(key, from);
+  if (at === -1) return null;
+  let i = text.indexOf('{', at);
+  if (i === -1) return null;
+  const start = i;
+  let depth = 0;
+  let inStr = false;
+  let esc = false;
+  for (; i < text.length; i++) {
+    const ch = text[i];
+    if (inStr) {
+      if (esc) esc = false;
+      else if (ch === '\\') esc = true;
+      else if (ch === '"') inStr = false;
+      continue;
+    }
+    if (ch === '"') inStr = true;
+    else if (ch === '{') depth++;
+    else if (ch === '}') {
+      depth--;
+      if (depth === 0) {
+        try { return JSON.parse(text.slice(start, i + 1)); }
+        catch { return null; }
+      }
+    }
+  }
+  return null;
+}
+
+const _CODEX_TOKEN_COUNT_PAYLOAD_RE = /"payload"\s*:\s*\{\s*"type"\s*:\s*"token_count"/;
+
+function _codexSummaryFromTokenCountLines(text) {
+  if (!text) return null;
+  let end = text.length;
+  while (end > 0) {
+    const start = text.lastIndexOf('\n', end - 1);
+    const line = text.slice(start + 1, end).trim();
+    if (line && _CODEX_TOKEN_COUNT_PAYLOAD_RE.test(line)) {
+      const summary = _codexSummaryFromTokenCountLine(line);
+      if (summary) return summary;
+    }
+    if (start === -1) break;
+    end = start;
+  }
+  return null;
+}
+
+function _codexSummaryFromTokenCountLine(line) {
+  if (!line) return null;
+  try {
+    const obj = JSON.parse(line);
+    const payload = obj && obj.payload;
+    const info = payload && payload.info;
+    if (!payload || payload.type !== 'token_count' || !info) return null;
+    if (!info.total_token_usage || typeof info.total_token_usage !== 'object') return null;
+    return {
+      total: info.total_token_usage,
+      last: info.last_token_usage && typeof info.last_token_usage === 'object' ? info.last_token_usage : null,
+      window: Number(info.model_context_window) || 0,
+    };
+  } catch {
+    return null;
+  }
+}
+
+function _codexCurrentContextTokens(summary) {
+  const last = summary && summary.last;
+  if (last && typeof last === 'object') {
+    // Codex emits a post-compaction token_count where input/output are zero but
+    // total_tokens is the new compacted context size. Treat 0 as valid; do not
+    // fall through to cumulative total_token_usage or the badge pins at 100%.
+    const lastTotal = Number(last.total_tokens);
+    if (Number.isFinite(lastTotal) && lastTotal > 0) return lastTotal;
+    const lastInput = Number(last.input_tokens);
+    if (Number.isFinite(lastInput) && lastInput > 0) return lastInput;
+    const componentTotal = (Number(last.input_tokens) || 0)
+      + (Number(last.output_tokens) || 0)
+      + (Number(last.reasoning_output_tokens) || 0);
+    return Math.max(0, componentTotal);
+  }
+  const total = summary && summary.total;
+  return Number(total && total.input_tokens) || 0;
+}
+
+// --- Claude: sum message.usage across assistant entries ------------------------
+//
+// Extract usage numbers from a single JSONL line WITHOUT JSON.parsing the whole
+// (possibly multi-MB) line: the four counters are uniquely-named keys, so a field
+// regex is both faster and immune to giant tool-result blobs.
+const _CLAUDE_FIELD_RE = {
+  input: /"input_tokens":\s*(\d+)/,
+  output: /"output_tokens":\s*(\d+)/,
+  cacheRead: /"cache_read_input_tokens":\s*(\d+)/,
+  cacheCreate: /"cache_creation_input_tokens":\s*(\d+)/,
+};
+
+function _claudeUsageFromLine(line) {
+  // Only assistant entries carry message.usage; cheap guard before regex work.
+  if (line.indexOf('"usage"') === -1) return null;
+  const num = (re) => { const m = line.match(re); return m ? parseInt(m[1], 10) || 0 : 0; };
+  const u = {
+    input: num(_CLAUDE_FIELD_RE.input),
+    output: num(_CLAUDE_FIELD_RE.output),
+    cacheRead: num(_CLAUDE_FIELD_RE.cacheRead),
+    cacheCreate: num(_CLAUDE_FIELD_RE.cacheCreate),
+  };
+  if (!(u.input || u.output || u.cacheRead || u.cacheCreate)) return null;
+  return u;
+}
+
+// Scan a chunk of appended JSONL text, accumulating into `acc`. Returns the
+// last full assistant turn seen (for ctx), or acc.lastTurn unchanged.
+function _scanClaudeChunk(text, acc) {
+  let start = 0;
+  while (start < text.length) {
+    let end = text.indexOf('\n', start);
+    if (end === -1) end = text.length;
+    const line = text.slice(start, end);
+    start = end + 1;
+    if (!line) continue;
+    // Authoritative window signal: Claude Code's AUTO compact_boundary records the
+    // pre-compaction token count near the context limit. Track the largest. Manual
+    // /compact is ignored (the user can compact at any size, so it says nothing about
+    // the window). These are system lines without usage, so they never double-count.
+    if (line.indexOf('compact_boundary') !== -1) {
+      if (/"trigger"\s*:\s*"auto"/.test(line)) {
+        const pt = line.match(/"preTokens"\s*:\s*(\d+)/);
+        if (pt) { const v = Number(pt[1]) || 0; if (v > (acc.autoCompactAt || 0)) acc.autoCompactAt = v; }
+      }
+      continue;
+    }
+    const u = _claudeUsageFromLine(line);
+    if (!u) continue;
+    acc.input += u.input;
+    acc.output += u.output;
+    acc.cacheRead += u.cacheRead;
+    acc.cacheCreate += u.cacheCreate;
+    // Current context occupancy ~ the most recent turn's prompt + reply.
+    acc.lastTurn = u.input + u.cacheRead + u.cacheCreate + u.output;
+    // High-water prompt size (input side only) — a hard lower bound on the real
+    // context window, used to self-correct a stale static window estimate.
+    const prompt = u.input + u.cacheRead + u.cacheCreate;
+    if (prompt > (acc.maxPrompt || 0)) acc.maxPrompt = prompt;
+  }
+}
+
+// Read the last `readLen` bytes of a file as utf8.
+function _readFileTail(jsonlPath, fileSize, readLen) {
+  const len = Math.min(readLen, fileSize);
+  const fd = fs.openSync(jsonlPath, 'r');
+  try {
+    const buf = Buffer.alloc(len);
+    fs.readSync(fd, buf, 0, len, Math.max(0, fileSize - len));
+    return buf.toString('utf8');
+  } finally { fs.closeSync(fd); }
+}
+
+// Read a small tail first; grow (×8) only when the window didn't contain a
+// token_count event. lastIndexOf+forward-scan means a window that includes the
+// key includes its (EOF-ward) object, so a partial leading line never corrupts.
+function _codexSummaryGrowing(jsonlPath, fileSize) {
+  let readLen = CODEX_TAIL_BYTES;
+  for (;;) {
+    let text;
+    try { text = _readFileTail(jsonlPath, fileSize, readLen); } catch { return null; }
+    const c = _codexSummaryFromTail(text);
+    if (c) return c;
+    if (readLen >= fileSize || readLen >= CODEX_TAIL_BYTES_MAX) return null;
+    readLen = Math.min(readLen * 8, CODEX_TAIL_BYTES_MAX, fileSize);
+  }
+}
+
+// --- Codex: read the last cumulative token_count event -------------------------
+function _codexSummaryFromTail(text) {
+  const structured = _codexSummaryFromTokenCountLines(text);
+  if (structured) return structured;
+
+  // Find the last token_count event; its total_token_usage is cumulative.
+  const lastIdx = text.lastIndexOf('"total_token_usage"');
+  if (lastIdx === -1) return null;
+  const total = _matchObjectAfter(text, '"total_token_usage"', lastIdx);
+  if (!total) return null;
+  const last = _matchObjectAfter(text, '"last_token_usage"', lastIdx);
+  const winM = text.slice(lastIdx).match(/"model_context_window":\s*(\d+)/);
+  const window = winM ? parseInt(winM[1], 10) || 0 : 0;
+  return { total, last, window };
+}
+
+// --- Estimate tier -------------------------------------------------------------
+let _encoder = null;
+let _encoderTried = false;
+function _estimateTokens(str) {
+  const text = String(str || '');
+  if (!text) return 0;
+  if (!_encoderTried) {
+    _encoderTried = true;
+    try { _encoder = require('gpt-tokenizer'); }
+    catch { _encoder = null; }
+  }
+  if (_encoder && typeof _encoder.encode === 'function') {
+    try { return _encoder.encode(text).length; } catch { /* fall through */ }
+  }
+  // Heuristic fallback (~4 chars/token) when the tokenizer dep isn't present.
+  return Math.ceil(text.length / 4);
+}
+
+function estimateFromMessages(messages, model, opts = {}) {
+  const list = Array.isArray(messages) ? messages : [];
+  let total = 0;
+  for (let k = 0; k < list.length; k++) {
+    const m = list[k];
+    // Structured-capture rows (tool calls/results, reasoning, patches) carry
+    // capped derived text; counting them would shift the estimate tier for
+    // sessions imported after the capture upgrade. Estimate from the same
+    // user/assistant prose the pre-capture importer produced.
+    if (m && m.role === 'system' && m.metadata && m.metadata.kind) continue;
+    total += _estimateTokens(m && (m.text || m.content));
+  }
+  const ctxWindow = ctxWindowForModel(model, opts.contextWindow || opts.ctxWindow);
+  return _finalize({ total, ctx: total, ctxWindow, exact: false,
+    breakdown: { input: 0, output: 0, cacheRead: 0, cacheCreate: 0, reasoning: 0, estimated: total } });
+}
+
+// --- Shaping -------------------------------------------------------------------
+function _finalize({ total, ctx, ctxWindow, exact, breakdown, cost }) {
+  const window = ctxWindow > 0 ? ctxWindow : DEFAULT_CTX_WINDOW;
+  const ctxClamped = Math.max(0, Math.min(Number(ctx) || 0, window));
+  const ctxPct = window > 0 ? Math.round((ctxClamped / window) * 100) : 0;
+  return {
+    total: Math.max(0, Math.round(Number(total) || 0)),
+    ctx: Math.max(0, Math.round(Number(ctx) || 0)),
+    ctxWindow: window,
+    ctxPct,
+    exact: !!exact,
+    totalLabel: formatTokens(total),
+    breakdown: breakdown || null,
+    cost: (cost === undefined ? null : cost),
+  };
+}
+
+// --- Per-file cache (incremental for Claude, tail for Codex) -------------------
+// keyed by jsonlPath -> { size, mtimeMs, offset, acc, summary }
+// Bounded LRU: a long-lived primary sees many transcripts (incl. rotated/compacted
+// ones whose paths never recur), so an unbounded Map would leak. Map preserves
+// insertion order; we evict the oldest key and refresh recency on hit.
+const _cache = new Map();
+const _CACHE_MAX = 500;
+
+function _cacheSet(key, val) {
+  if (_cache.has(key)) _cache.delete(key);
+  _cache.set(key, val);
+  if (_cache.size > _CACHE_MAX) {
+    const oldest = _cache.keys().next().value;
+    if (oldest !== undefined) _cache.delete(oldest);
+  }
+}
+
+function _statSafe(p) {
+  try { return fs.statSync(p); } catch { return null; }
+}
+
+// `statHint` ({size, mtimeMs}) lets a caller that already freshly stat'd the file
+// (e.g. _sessionPayload via _activeSessionFileInfo) skip a redundant statSync.
+function _computeFromFile(provider, jsonlPath, model, statHint, opts = {}) {
+  const st = (statHint && Number.isFinite(statHint.size) && Number.isFinite(statHint.mtimeMs))
+    ? statHint
+    : _statSafe(jsonlPath);
+  if (!st || (typeof st.isFile === 'function' && !st.isFile())) return null;
+  const sizeKey = st.size;
+  const mtimeKey = st.mtimeMs;
+  const cached = _cache.get(jsonlPath);
+  if (cached && cached.size === sizeKey && cached.mtimeMs === mtimeKey && cached.summary) {
+    _cacheSet(jsonlPath, cached); // refresh LRU recency
+    return cached.summary;
+  }
+
+  if (_isCodex(provider)) {
+    const c = _codexSummaryGrowing(jsonlPath, st.size);
+    if (!c) return null;
+    const tt = c.total || {};
+    // "total" = genuinely-new tokens this session (exclude cached re-reads, which
+    // are re-counted every turn and would balloon into the billions). Codex's
+    // input_tokens INCLUDES cached_input_tokens, so subtract them out.
+    const cachedIn = Number(tt.cached_input_tokens) || 0;
+    const nonCachedIn = Math.max(0, (Number(tt.input_tokens) || 0) - cachedIn);
+    const total = nonCachedIn + (Number(tt.output_tokens) || 0) + (Number(tt.reasoning_output_tokens) || 0);
+    const ctx = _codexCurrentContextTokens(c);
+    const summary = _finalize({
+      total, ctx, ctxWindow: c.window || ctxWindowForModel(model, opts.contextWindow || opts.ctxWindow), exact: true,
+      breakdown: {
+        input: Number(tt.input_tokens) || 0,
+        output: Number(tt.output_tokens) || 0,
+        cacheRead: Number(tt.cached_input_tokens) || 0,
+        cacheCreate: 0,
+        reasoning: Number(tt.reasoning_output_tokens) || 0,
+      },
+    });
+    _cacheSet(jsonlPath, { size: sizeKey, mtimeMs: mtimeKey, summary });
+    return summary;
+  }
+
+  // Claude (and any other provider whose transcript matches the Claude usage
+  // shape): incremental append-scan.
+  let acc;
+  let offset;
+  if (cached && cached.acc && st.size >= cached.offset) {
+    acc = cached.acc;
+    offset = cached.offset;
+  } else {
+    acc = { input: 0, output: 0, cacheRead: 0, cacheCreate: 0, lastTurn: 0, maxPrompt: 0, autoCompactAt: 0 };
+    offset = 0;
+  }
+  if (st.size > offset) {
+    try {
+      const fd = fs.openSync(jsonlPath, 'r');
+      try {
+        const len = st.size - offset;
+        const buf = Buffer.alloc(len);
+        fs.readSync(fd, buf, 0, len, offset);
+        _scanClaudeChunk(buf.toString('utf8'), acc);
+      } finally { fs.closeSync(fd); }
+    } catch { return null; }
+  }
+  // "total" = genuinely-new tokens (non-cached input + new cache-creation +
+  // output). cache_READ is the context re-read every turn; including it balloons
+  // the count into the billions, so it lives in the breakdown/tooltip only.
+  const total = acc.input + acc.cacheCreate + acc.output;
+  if (acc.input + acc.output + acc.cacheRead + acc.cacheCreate === 0) {
+    // No usage found (e.g. not actually a Claude transcript). Let caller fall
+    // back to estimation rather than showing 0.
+    _cacheSet(jsonlPath, { size: sizeKey, mtimeMs: mtimeKey, acc, offset: st.size, summary: null });
+    return null;
+  }
+  const summary = _finalize({
+    total,
+    ctx: acc.lastTurn || total,
+    ctxWindow: _effectiveCtxWindow(ctxWindowForModel(model, opts.contextWindow || opts.ctxWindow), { maxPrompt: acc.maxPrompt, autoCompactAt: acc.autoCompactAt }),
+    exact: true,
+    breakdown: {
+      input: acc.input, output: acc.output,
+      cacheRead: acc.cacheRead, cacheCreate: acc.cacheCreate, reasoning: 0,
+    },
+  });
+  _cacheSet(jsonlPath, { size: sizeKey, mtimeMs: mtimeKey, acc, offset: st.size, summary });
+  return summary;
+}
+
+/**
+ * Compute the token-usage summary for a session.
+ * @param {object} o
+ * @param {string} o.provider  agent type: 'claude' | 'codex' | 'cursor' | ...
+ * @param {string} [o.jsonlPath]  transcript path (Claude/Codex authoritative)
+ * @param {string} [o.model]  model id, for the ctx-window fallback
+ * @param {Array}  [o.messages]  parsed messages, for the estimate tier
+ * @param {number} [o.contextWindow] provider/catalog context-window override
+ * @param {object} [o.stat]  a fresh {size, mtimeMs} the caller already stat'd,
+ *   to skip a redundant statSync (must be current — used as the cache key + read range)
+ * @returns {object|null} summary, or null when nothing is computable
+ */
+function computeSessionTokenUsage({ provider, jsonlPath, model, messages, stat, contextWindow, ctxWindow } = {}) {
+  const opts = { contextWindow: contextWindow || ctxWindow };
+  if (jsonlPath && (_isClaude(provider) || _isCodex(provider) || !provider)) {
+    const fromFile = _computeFromFile(provider, jsonlPath, model, stat, opts);
+    if (fromFile) return fromFile;
+  }
+  if (Array.isArray(messages) && messages.length) {
+    return estimateFromMessages(messages, model, opts);
+  }
+  return null;
+}
+
+function _resetCacheForTests() { _cache.clear(); }
+function _cacheSizeForTests() { return _cache.size; }
+
+module.exports = {
+  computeSessionTokenUsage,
+  estimateFromMessages,
+  formatTokens,
+  ctxWindowForModel,
+  _effectiveCtxWindow,
+  _resetCacheForTests,
+  _cacheSizeForTests,
+  _CACHE_MAX,
+};