138 lines · 4.7 KB · TypeScript
RawDownload
import { createReadStream, statSync } from "node:fs";import readline from "node:readline"; /** Everything the session has spent, for the whole of its life. * *  This is deliberately NOT the context window. Context resets on every *  compaction; spend does not. A session that has compacted four times still *  cost what it cost, and that total is the number worth showing. */export type SessionUsageTotals = {  input: number;  output: number;  cacheRead: number;  cacheWrite: number;  reasoning: number;  total: number;  /** Total cost in USD when the provider reports one; 0 for local models. */  cost: number;  /** Assistant round-trips, i.e. how many times a model was actually called. */  calls: number;  /** Successful compactions, each one a point where the context was discarded. */  compactions: number;}; export function emptyUsageTotals(): SessionUsageTotals {  return {    input: 0,    output: 0,    cacheRead: 0,    cacheWrite: 0,    reasoning: 0,    total: 0,    cost: 0,    calls: 0,    compactions: 0,  };} type CacheEntry = { size: number; mtimeMs: number; totals: SessionUsageTotals }; // Rollouts are append-only, so a file whose size and mtime are unchanged has// unchanged totals. Keyed by path; one entry per session file.const cache = new Map<string, CacheEntry>(); function numeric(source: Record<string, unknown> | null, keys: string[]): number {  if (!source) return 0;  for (const key of keys) {    const value = source[key];    if (typeof value === "number" && Number.isFinite(value)) return value;  }  return 0;} function asRecord(value: unknown): Record<string, unknown> | null {  return value && typeof value === "object" && !Array.isArray(value)    ? (value as Record<string, unknown>)    : null;} /** Fold one rollout line into the running totals. */export function accumulateUsageLine(totals: SessionUsageTotals, line: string): SessionUsageTotals {  // Cheap pre-filter: the vast majority of lines are tool output and user text  // with no usage block at all, and JSON.parse on a multi-GB log is the whole  // cost of this scan.  const hasUsage = line.includes('"usage"');  const hasCompaction = line.includes("compaction");  if (!hasUsage && !hasCompaction) return totals;   let entry: Record<string, unknown> | null = null;  try {    entry = asRecord(JSON.parse(line));  } catch {    return totals;  }  if (!entry) return totals;   if (entry.type === "compaction" || entry.customType === "compaction") {    return { ...totals, compactions: totals.compactions + 1 };  }   const message = asRecord(entry.message);  if (!message || message.role !== "assistant") return totals;  const usage = asRecord(message.usage);  if (!usage) return totals;   const input = numeric(usage, ["input", "input_tokens", "prompt_tokens"]);  const output = numeric(usage, ["output", "output_tokens", "completion_tokens"]);  const cacheRead = numeric(usage, ["cacheRead", "cache_read_input_tokens"]);  const cacheWrite = numeric(usage, ["cacheWrite", "cache_creation_input_tokens"]);  const reasoning = numeric(usage, ["reasoning", "reasoning_tokens"]);  const reported = numeric(usage, ["totalTokens", "total_tokens", "total"]);  const cost = numeric(asRecord(usage.cost), ["total"]);   return {    input: totals.input + input,    output: totals.output + output,    cacheRead: totals.cacheRead + cacheRead,    cacheWrite: totals.cacheWrite + cacheWrite,    reasoning: totals.reasoning + reasoning,    total: totals.total + (reported || input + output),    cost: totals.cost + cost,    calls: totals.calls + 1,    compactions: totals.compactions,  };} /** Walk a rollout and total what it spent. * *  Streams line by line so memory stays flat regardless of file size, and *  memoises on (size, mtime) — an append-only log that has not grown cannot *  have different totals, and the status panel asks for this on every session *  open. */export async function readSessionUsageTotals(filepath: string): Promise<SessionUsageTotals> {  let stat: { size: number; mtimeMs: number };  try {    stat = statSync(filepath);  } catch {    return emptyUsageTotals();  }   const cached = cache.get(filepath);  if (cached && cached.size === stat.size && cached.mtimeMs === stat.mtimeMs) {    return cached.totals;  }   let totals = emptyUsageTotals();  try {    const stream = createReadStream(filepath, { encoding: "utf-8" });    const lines = readline.createInterface({ input: stream, crlfDelay: Infinity });    for await (const line of lines) {      if (line) totals = accumulateUsageLine(totals, line);    }  } catch {    return emptyUsageTotals();  }   cache.set(filepath, { size: stat.size, mtimeMs: stat.mtimeMs, totals });  return totals;}