pi-tokenburn 0.8.0 → 0.8.2

package/README.md

@@ -183,6 +183,7 @@ Environment overrides for the session directory: `TOKENBURN_PI_SESSIONS`, `PI_CO
 Refreshing is designed to cost nothing you can notice:
 
 - **Incremental:** pi's session logs are append-only, so after the first parse only the bytes added since then are read (with a cheap integrity check that falls back to a full parse if a file was rewritten or truncated). On a 35 MB active session a refresh after one new message takes **~1.2 ms**, down from ~18 ms, and the cost no longer grows with the session.
+- **Fast start:** pi waits for nothing: the first refresh runs in the background. The parse state is saved in `tokenburn-cache.json` (pi's agent dir, at most once a minute and at shutdown), so a new pi run reads only what was appended since; on a 2.4 GB session history the status line appears in **~40 ms** instead of ~4 s, and pi starts ~5 s sooner. Delete the file any time: the next start parses everything once.
 - **Coalesced:** a burst of messages triggers one refresh (trailing delay, `refreshMs`), never two at once.
 - **Never blocks:** refreshes are fire-and-forget and errors are swallowed; the timer is unref'd and cancelled at shutdown.
 - **Off means off:** `/tokenburn disable` (or `"enabled": false`) reads no files and starts no timers.
@@ -200,7 +201,7 @@ just demo             # try it yourself in a real pi on the same synthetic data
 
 ## Security notes
 
-- **Local and read-only.** It reads pi's session logs and its own `tokenburn.json`. Nothing leaves your machine: no network, no telemetry, no shell commands, **zero runtime dependencies**.
+- **Local.** It reads pi's session logs, its own `tokenburn.json` and the savings logs of pi-context-diet / pi-squeeze, and writes only `tokenburn.json` and its parse cache `tokenburn-cache.json`. Nothing leaves your machine: no network, no telemetry, no shell commands, **zero runtime dependencies**.
 - **Never breaks pi.** Stats refreshes are fire-and-forget and swallow their own errors, so a failure here cannot interrupt a turn or a session start.
 - **Hand-edited config is sanitised.** Unknown or non-numeric budget values are dropped.
 
@@ -224,7 +225,7 @@ One folder per feature under `extensions/features/`, each with its logic and its
 ```
 extensions/
   tokenburn.ts            wiring only: pi events + the /tokenburn command
-  core/                   types, config, session-log parser (incremental cache), aggregation, filters + prices, paint
+  core/                   types, config, session-log parser (incremental cache, saved between runs), aggregation, filters + prices, paint
   command/                parse + completions (pure), handler table, the Host handlers receive
   features/
     status-line/          footer text, optional panel, burn rate / cache hit / savings

package/extensions/core/cache-file.ts

@@ -0,0 +1,143 @@
+/**
+ * The session-log parse cache, saved between pi runs in `tokenburn-cache.json` (pi's agent dir).
+ *
+ * Without it every pi start re-reads the whole session history (gigabytes for a long-used pi)
+ * before the status line can show anything. With it, a start reads only what was appended since.
+ * The parser still checks size, mtime and the head / last-parsed bytes of every file, so a stale
+ * or foreign cache file only costs a re-parse, never wrong numbers. Any read or write failure
+ * is ignored: the cache is an optimization.
+ */
+import { readFile, rename, writeFile } from "node:fs/promises";
+import { join } from "node:path";
+import { getConfigDir } from "./config.ts";
+import { cacheGeneration, exportCache, importCache, type FileEntry, type Turn } from "./parser.ts";
+
+export const CACHE_FILE = "tokenburn-cache.json";
+const VERSION = 1;
+/** Rewrite the file at most this often while pi runs; the last state is written at shutdown. */
+export const SAVE_EVERY_MS = 60_000;
+
+export const cacheFilePath = (dir: string = getConfigDir()) => join(dir, CACHE_FILE);
+
+/** A turn without its session (the file it belongs to); strings are indexes into `strings`. */
+type Row = [ts: number, input: number, output: number, cacheRead: number, cacheWrite: number, cost: number, costInput: number, costCacheRead: number, provider: number, model: number, project: number];
+
+interface SavedEntry {
+	size: number;
+	mtimeMs: number;
+	offset: number;
+	cwd?: string;
+	head: string;
+	guard: string;
+	turns: Row[];
+	tail: Row[];
+}
+
+interface Saved {
+	version: number;
+	strings: string[];
+	files: Record<string, SavedEntry>;
+}
+
+export function encodeCache(entries: [string, FileEntry][]): Saved {
+	const strings: string[] = [];
+	const index = new Map<string, number>();
+	const id = (s: string) => {
+		let i = index.get(s);
+		if (i === undefined) {
+			i = strings.push(s) - 1;
+			index.set(s, i);
+		}
+		return i;
+	};
+	const row = (t: Turn): Row => [t.ts, t.input, t.output, t.cacheRead, t.cacheWrite, t.cost, t.costInput, t.costCacheRead, id(t.provider), id(t.model), id(t.project)];
+	const files: Record<string, SavedEntry> = {};
+	for (const [file, e] of entries) {
+		files[file] = {
+			size: e.size,
+			mtimeMs: e.mtimeMs,
+			offset: e.offset,
+			...(e.cwd === undefined ? {} : { cwd: e.cwd }),
+			head: e.head.toString("base64"),
+			guard: e.guard.toString("base64"),
+			turns: e.turns.map(row),
+			tail: e.tail.map(row),
+		};
+	}
+	return { version: VERSION, strings, files };
+}
+
+const isNum = (v: unknown): v is number => typeof v === "number" && Number.isFinite(v);
+
+/** The entries of a saved cache; anything malformed is dropped (that file is simply re-parsed). */
+export function decodeCache(data: unknown): [string, FileEntry][] {
+	const saved = data as Partial<Saved> | null;
+	if (!saved || saved.version !== VERSION || !Array.isArray(saved.strings) || typeof saved.files !== "object" || saved.files === null) return [];
+	const strings = saved.strings;
+	const str = (i: unknown) => (Number.isInteger(i) && typeof strings[i as number] === "string" ? strings[i as number] : undefined);
+	const turn = (r: unknown, session: string): Turn | undefined => {
+		if (!Array.isArray(r) || r.length !== 11 || !r.slice(0, 8).every(isNum)) return undefined;
+		const [provider, model, project] = [str(r[8]), str(r[9]), str(r[10])];
+		if (provider === undefined || model === undefined || project === undefined) return undefined;
+		const [ts, input, output, cacheRead, cacheWrite, cost, costInput, costCacheRead] = r as number[];
+		return { ts: ts!, input: input!, output: output!, cacheRead: cacheRead!, cacheWrite: cacheWrite!, cost: cost!, costInput: costInput!, costCacheRead: costCacheRead!, provider, model, project, session };
+	};
+	const turns = (rows: unknown, session: string): Turn[] | undefined => {
+		if (!Array.isArray(rows)) return undefined;
+		const out: Turn[] = [];
+		for (const r of rows) {
+			const t = turn(r, session);
+			if (!t) return undefined;
+			out.push(t);
+		}
+		return out;
+	};
+	const out: [string, FileEntry][] = [];
+	for (const [file, e] of Object.entries(saved.files as Record<string, Partial<SavedEntry>>)) {
+		if (!e || !isNum(e.size) || !isNum(e.mtimeMs) || !isNum(e.offset) || typeof e.head !== "string" || typeof e.guard !== "string") continue;
+		if (e.cwd !== undefined && typeof e.cwd !== "string") continue;
+		const [t, tail] = [turns(e.turns, file), turns(e.tail, file)];
+		if (!t || !tail) continue;
+		out.push([file, { size: e.size, mtimeMs: e.mtimeMs, offset: e.offset, cwd: e.cwd, head: Buffer.from(e.head, "base64"), guard: Buffer.from(e.guard, "base64"), turns: t, tail }]);
+	}
+	return out;
+}
+
+let savedGeneration = -1;
+let lastSave = -Infinity;
+
+/** Seed the parser from the cache file (a no-op when it is missing or broken, or the parser is warm). */
+export async function restoreCache(path: string = cacheFilePath()): Promise<void> {
+	try {
+		importCache(decodeCache(JSON.parse(await readFile(path, "utf8"))));
+		savedGeneration = cacheGeneration();
+	} catch {
+		/* no cache yet, or a broken one: the first refresh parses everything */
+	}
+}
+
+/**
+ * Write the parse state when it changed, at most every SAVE_EVERY_MS unless `force` (shutdown).
+ * Written to a temp file and renamed, so a crash never leaves a torn cache. Returns whether it wrote.
+ */
+export async function persistCache(path: string = cacheFilePath(), opts: { force?: boolean; now?: number } = {}): Promise<boolean> {
+	const now = opts.now ?? Date.now();
+	const gen = cacheGeneration();
+	if (gen === savedGeneration || (!opts.force && now - lastSave < SAVE_EVERY_MS)) return false;
+	try {
+		const tmp = `${path}.${process.pid}.tmp`;
+		await writeFile(tmp, JSON.stringify(encodeCache(exportCache())));
+		await rename(tmp, path);
+		savedGeneration = gen;
+		lastSave = now;
+		return true;
+	} catch {
+		return false;
+	}
+}
+
+/** Test hook: forget what was saved and when. */
+export function resetCacheFileState(): void {
+	savedGeneration = -1;
+	lastSave = -Infinity;
+}

package/extensions/core/parser.ts

@@ -230,11 +230,25 @@ function listSessionFiles(root: string): string[] {
 	return files;
 }
 
-/** Bring the per-file cache up to date (reading only what changed) and return every file's turns. */
-export async function loadAll(sessionsRoot: string, chunkSize = CHUNK_BYTES): Promise<Turn[][]> {
+/** Bring the per-file cache up to date (reading only what changed) and return every file's turns.
+ * Calls run one at a time: an incremental parse appends to the cached turns, so two at once
+ * (the background startup refresh and a /tokenburn command) would count the new turns twice. */
+export function loadAll(sessionsRoot: string, chunkSize = CHUNK_BYTES): Promise<Turn[][]> {
+	const run = queue.then(() => loadAllNow(sessionsRoot, chunkSize));
+	queue = Promise.allSettled([run]);
+	return run;
+}
+
+let queue: Promise<unknown> = Promise.resolve();
+
+/** Bumped whenever a cache entry changes, so the cache file is only rewritten when needed. */
+let generation = 0;
+export const cacheGeneration = () => generation;
+
+async function loadAllNow(sessionsRoot: string, chunkSize: number): Promise<Turn[][]> {
 	const files = listSessionFiles(sessionsRoot);
 	const live = new Set(files);
-	for (const key of FILE_CACHE.keys()) if (!live.has(key)) FILE_CACHE.delete(key);
+	for (const key of FILE_CACHE.keys()) if (!live.has(key) && FILE_CACHE.delete(key)) generation++;
 
 	const out: Turn[][] = [];
 	for (const file of files) {
@@ -245,6 +259,7 @@ export async function loadAll(sessionsRoot: string, chunkSize = CHUNK_BYTES): Pr
 			const fresh = await refreshEntry(file, st, entry, chunkSize);
 			if (fresh) FILE_CACHE.set(file, fresh);
 			else FILE_CACHE.delete(file);
+			generation++;
 			entry = fresh;
 		}
 		if (entry) out.push(entry.tail.length > 0 ? entry.turns.concat(entry.tail) : entry.turns);
@@ -259,3 +274,14 @@ export const __refreshEntryForTests = refreshEntry;
 export function clearCache(): void {
 	FILE_CACHE.clear();
 }
+
+/** The parse state of every session file, to save it. */
+export const exportCache = (): [string, FileEntry][] => [...FILE_CACHE];
+
+/** Seed a fresh runtime with saved parse state. A warm cache is newer than any file: kept. */
+export function importCache(entries: Iterable<[string, FileEntry]>): void {
+	if (FILE_CACHE.size > 0) return;
+	for (const [file, entry] of entries) FILE_CACHE.set(file, entry);
+}
+
+export type { FileEntry };

package/extensions/tokenburn.ts

@@ -22,6 +22,7 @@ import type { Host } from "./command/host.ts";
 import { completions } from "./command/completions.ts";
 import { parseCommand } from "./command/parse.ts";
 import { collectUsage } from "./core/aggregate.ts";
+import { persistCache, restoreCache } from "./core/cache-file.ts";
 import { freshConfig, getSessionsDir, loadConfig, saveConfig, type TokenBurnConfig } from "./core/config.ts";
 import type { UsageFilter, UsageOptions } from "./core/select.ts";
 import { budgetWarnings } from "./features/budgets/budgets.ts";
@@ -49,6 +50,10 @@ export * from "./features/breakdown/table.ts";
 export * from "./features/pricing/report.ts";
 export * from "./features/export/export.ts";
 
+/** The background refresh started by the latest session_start (tests await it). */
+let startup: Promise<void> = Promise.resolve();
+export const startupRefresh = () => startup;
+
 export default function tokenBurnExtension(pi: ExtensionAPI) {
 	let config: TokenBurnConfig = freshConfig();
 	/** Budget warnings already shown (id = budget + period bucket + level): each shows once. */
@@ -84,6 +89,7 @@ export default function tokenBurnExtension(pi: ExtensionAPI) {
 
 		if (config.showWidget) ctx.ui.setWidget("tokenburn-panel", renderWidget(report, config, theme), { placement: "belowEditor" });
 		else ctx.ui.setWidget("tokenburn-panel", []);
+		void persistCache();
 	};
 
 	// The latest context, used by the (delayed) live refresh. Dropped at shutdown so a reloaded
@@ -99,8 +105,10 @@ export default function tokenBurnExtension(pi: ExtensionAPI) {
 	pi.on("session_start", async (_event, ctx) => {
 		latestCtx = ctx;
 		config = await loadConfig();
-		// A failing refresh must never break session start.
-		await updateUI(ctx).catch(() => {});
+		if (config.enabled) await restoreCache();
+		// pi waits for this handler, and the first scan of a long session history takes seconds:
+		// draw in the background. A failing refresh must never break session start.
+		startup = updateUI(ctx).catch(() => {});
 	});
 
 	// Live: after each assistant message, coalesced. Never blocks, never throws into pi, and
@@ -120,9 +128,10 @@ export default function tokenBurnExtension(pi: ExtensionAPI) {
 		void refresher.flush();
 	});
 
-	pi.on("session_shutdown", () => {
+	pi.on("session_shutdown", async () => {
 		refresher.cancel();
 		latestCtx = undefined;
+		await persistCache(undefined, { force: true });
 	});
 
 	pi.registerCommand("tokenburn", {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-tokenburn",
-  "version": "0.8.0",
+  "version": "0.8.2",
   "description": "Pi extension for live token usage monitoring and budget tracking (hour / day / week / month / year / all). 0 prompt tokens.",
   "keywords": ["pi-package", "pi-extension", "tokenburn", "tokens", "analytics", "budget", "stats"],
   "license": "MIT",