pi-tokenburn 0.7.4 → 0.8.2

package/README.md

@@ -17,7 +17,7 @@ Totals match the `tokenburn` CLI exactly.
 
 ## Features
 
-- **Status bar:** `🔥 1.2M ($5.10) [today] · 52k/h · 95% hit · ~$12.40 saved`: tokens, cost, burn rate, cache hit ratio and estimated cache savings for the chosen period (each extra toggleable). Turns red with `OVER BUDGET` past your limit.
+- **Status bar:** `🔥 1.2M ($5.10) [today] · 52k/h · 95% hit · ~$12.40 saved`: tokens, cost, burn rate, cache hit ratio and estimated savings (prompt cache, plus [pi-context-diet](https://github.com/sorinirimies/pi-context-diet) and [pi-squeeze](https://github.com/sorinirimies/pi-squeeze) when installed) for the chosen period (each extra toggleable). Turns red with `OVER BUDGET` past your limit.
 - **Tables:** usage per model, provider, project and session, for any period, with cache hit %, savings, cost and share.
 - **Filter and export:** limit everything to a provider, model or project; write csv or json.
 - **Pricing:** a price table fills in (or overrides) costs the logs don't have.
@@ -61,7 +61,7 @@ pi install git:github.com/sorinirimies/pi-tokenburn
 | `/tokenburn warn <pct>` | Notify once when a budget reaches `pct` percent (default 80) and once when it is exceeded |
 | `/tokenburn rate` | Toggle the burn rate in the status line (`42k/h`: tokens per hour averaged over the status period) |
 | `/tokenburn hit` | Toggle the cache hit ratio (`95% hit`) in the status line |
-| `/tokenburn saved` | Toggle the estimated cache savings (`~$12.40 saved`) in the status line |
+| `/tokenburn saved` | Toggle the estimated savings (`~$12.40 saved`: prompt cache, pi-context-diet, pi-squeeze) in the status line |
 | `/tokenburn pricing` | Where costs come from (session logs vs your price table) and which models are unpriced |
 | `/tokenburn price <model> <in> <out>` | Set a price in $ per million tokens (key `model` or `provider/model`); `price remove <model>` deletes it. `pricing mode <fill\|override>`: use the table only where the log has no cost (default), or always |
 | `/tokenburn window` | Toggle the panel below the editor |
@@ -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/command/completions.ts

@@ -41,7 +41,7 @@ const TOP_LEVEL: CompletionItem[] = [
 	{ value: "price ", label: "price", description: "Set a $/M-token price for a model" },
 	{ value: "rate", label: "rate", description: "Toggle the tokens/hour burn rate in the status line" },
 	{ value: "hit", label: "hit", description: "Toggle the cache hit ratio in the status line" },
-	{ value: "saved", label: "saved", description: "Toggle the estimated cache savings in the status line" },
+	{ value: "saved", label: "saved", description: "Toggle the estimated savings (cache, pi context-diet, pi squeeze) in the status line" },
 	{ value: "cycle", label: "cycle", description: "Rotate the status-bar period" },
 	{ value: "window", label: "window", description: "Toggle the panel below the editor" },
 	{ value: "budget ", label: "budget", description: "Set a day/week/month budget" },

package/extensions/core/aggregate.ts

@@ -1,6 +1,7 @@
 /** Aggregation of parsed turns into period totals and time series. */
+import { DIET_LOG, loadPluginLog, pluginLogPath, SQUEEZE_LOG } from "./plugin-logs.ts";
 import { loadAll, type Turn } from "./parser.ts";
-import { cacheValue, selectTurns, type UsageOptions } from "./select.ts";
+import { cacheValue, hasFilter, selectTurns, type UsageOptions } from "./select.ts";
 import { bucketLabel, bucketStart, type Unit } from "./time.ts";
 import type { PeriodStats, TokenBurnReport } from "./types.ts";
 
@@ -14,6 +15,13 @@ const emptyStats = (): PeriodStats => ({
 	sessions: 0,
 	cacheBase: 0,
 	cacheActual: 0,
+	cacheReadPriced: 0,
+	dietTokens: 0,
+	dietRequests: 0,
+	dietUsd: 0,
+	squeezeTokens: 0,
+	squeezeOutputs: 0,
+	squeezeUsd: 0,
 });
 
 export const HOUR_MS = 3_600_000;
@@ -33,6 +41,8 @@ export async function collectUsage(
 	now: Date = new Date(),
 	includeCache = true,
 	opts?: UsageOptions,
+	/** Where the pi plugins' savings logs are (default: pi's agent dir). */
+	logDir?: string,
 ): Promise<TokenBurnReport> {
 	const startOfHour = now.getTime() - HOUR_MS;
 	const startOfDay = bucketStart(now.getTime(), "day");
@@ -59,6 +69,7 @@ export async function collectUsage(
 		const cv = cacheValue(t);
 		p.cacheBase += cv.base;
 		p.cacheActual += cv.actual;
+		if (cv.base > 0) p.cacheReadPriced += t.cacheRead;
 	};
 
 	for (const turns of await loadSelected(sessionsRoot, opts)) {
@@ -81,6 +92,35 @@ export async function collectUsage(
 		if (maxTs >= startOfHour) report.hour.sessions++;
 	}
 
+	// The pi plugins' logs are not attributed to providers, models or projects: only unfiltered views count them.
+	if (!hasFilter(opts?.filter)) {
+		const windows: [PeriodStats, number][] = [
+			[report.total, -Infinity],
+			[report.year, startOfYear],
+			[report.month, startOfMonth],
+			[report.week, startOfWeek],
+			[report.today, startOfDay],
+			[report.hour, startOfHour],
+		];
+		const [diet, squeeze] = await Promise.all([loadPluginLog(pluginLogPath(DIET_LOG, logDir)), loadPluginLog(pluginLogPath(SQUEEZE_LOG, logDir))]);
+		for (const [p, start] of windows) {
+			for (const r of diet) {
+				if (r.ts < start) continue;
+				p.dietTokens += r.tokens;
+				p.dietRequests++;
+			}
+			for (const r of squeeze) {
+				if (r.ts < start) continue;
+				p.squeezeTokens += r.tokens;
+				p.squeezeOutputs++;
+			}
+			// Priced at the window's own prices: elided context would mostly have been cache reads,
+			// squeezed bash output would have been fresh input.
+			p.dietUsd = p.cacheReadPriced > 0 ? (p.dietTokens * p.cacheActual) / p.cacheReadPriced : 0;
+			p.squeezeUsd = p.cacheReadPriced > 0 ? (p.squeezeTokens * p.cacheBase) / p.cacheReadPriced : 0;
+		}
+	}
+
 	return report;
 }
 

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/config.ts

@@ -23,7 +23,7 @@ export interface TokenBurnConfig {
 	showRate: boolean;
 	/** Show the cache hit ratio of the status window in the status line. */
 	showHit: boolean;
-	/** Show the estimated cache savings of the status window in the status line. */
+	/** Show the estimated savings (prompt cache, pi-context-diet, pi-squeeze) of the status window in the status line. */
 	showSaved: boolean;
 	/** Notify once when a budget reaches this percentage (1 to 100); 0 turns warnings off. */
 	warnPct: number;

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/core/plugin-logs.ts

@@ -0,0 +1,57 @@
+/**
+ * The savings logs of pi plugins, in pi's agent dir: one JSON line per saving
+ * (`{"ts":"2026-01-02T09:15:00.000Z","tokens":18400,...}`). Only `ts` and `tokens` are read.
+ *
+ * - pi-context-diet: `context-diet-savings.jsonl`, one line per request it trimmed;
+ * - pi-squeeze: `squeeze-savings.jsonl`, one line per bash output it shortened.
+ *
+ * A missing or unreadable file means "no savings"; lines that do not parse are skipped.
+ */
+import { readFile, stat } from "node:fs/promises";
+import { join } from "node:path";
+import { getConfigDir } from "./config.ts";
+
+export const DIET_LOG = "context-diet-savings.jsonl";
+export const SQUEEZE_LOG = "squeeze-savings.jsonl";
+
+export interface LogRecord {
+	ts: number;
+	tokens: number;
+}
+
+/** Where a plugin's log is: `file` in pi's agent dir. */
+export const pluginLogPath = (file: string, dir: string = getConfigDir()) => join(dir, file);
+
+export function parsePluginLog(text: string): LogRecord[] {
+	const out: LogRecord[] = [];
+	for (const line of text.split("\n")) {
+		if (!line.trim()) continue;
+		try {
+			const r = JSON.parse(line) as { ts?: unknown; tokens?: unknown };
+			const ts = typeof r.ts === "string" ? Date.parse(r.ts) : NaN;
+			const tokens = r.tokens;
+			if (Number.isFinite(ts) && typeof tokens === "number" && Number.isFinite(tokens) && tokens > 0) out.push({ ts, tokens });
+		} catch {
+			/* a torn or foreign line */
+		}
+	}
+	return out;
+}
+
+// The status line refreshes after every message: re-read a log only when it changed.
+const cache = new Map<string, { mtimeMs: number; size: number; records: LogRecord[] }>();
+
+export async function loadPluginLog(path: string): Promise<LogRecord[]> {
+	try {
+		const s = await stat(path);
+		const hit = cache.get(path);
+		if (hit && hit.mtimeMs === s.mtimeMs && hit.size === s.size) return hit.records;
+		const records = parsePluginLog(await readFile(path, "utf8"));
+		cache.set(path, { mtimeMs: s.mtimeMs, size: s.size, records });
+		return records;
+	} catch {
+		return [];
+	}
+}
+
+export const clearPluginLogCache = () => cache.clear();

package/extensions/core/types.ts

@@ -20,6 +20,16 @@ export interface PeriodStats {
 	/** What cache reads would have cost as plain input vs what they cost (turns with a known input price). */
 	cacheBase: number;
 	cacheActual: number;
+	/** Cache-read tokens of the turns with a known input price (`cacheActual` / this = the cache-read price). */
+	cacheReadPriced: number;
+	/** pi-context-diet: tokens kept out of requests, the requests, and what they would have cost as cache reads. */
+	dietTokens: number;
+	dietRequests: number;
+	dietUsd: number;
+	/** pi-squeeze: tokens cut from bash output, the outputs, and what they would have cost as input. */
+	squeezeTokens: number;
+	squeezeOutputs: number;
+	squeezeUsd: number;
 }
 
 export interface TokenBurnReport {

package/extensions/features/report/report.ts

@@ -27,6 +27,12 @@ export function renderReport(input: {
 		...(models.cacheRead + models.cacheWrite > 0
 			? [`cache: ${Math.round(cacheHitRatio(models) * 100)}% of prompt tokens read from cache${cacheSavings(models) > 0 ? ` · saved ~${formatCost(cacheSavings(models))}` : ""}`]
 			: []),
+		...(report.total.dietTokens > 0
+			? [`pi context-diet: ~${formatNum(report.total.dietTokens)} tokens kept out of ${report.total.dietRequests} requests${report.total.dietUsd > 0 ? ` · saved ~${formatCost(report.total.dietUsd)}` : ""}`]
+			: []),
+		...(report.total.squeezeTokens > 0
+			? [`pi squeeze: ~${formatNum(report.total.squeezeTokens)} tokens cut from ${report.total.squeezeOutputs} bash outputs${report.total.squeezeUsd > 0 ? ` · saved ~${formatCost(report.total.squeezeUsd)}` : ""}`]
+			: []),
 		...forecastLines(report, config, now),
 	];
 	return [

package/extensions/features/status-line/commands.ts

@@ -73,6 +73,6 @@ export const statusLineCommands: HandlersFor<"window" | "status" | "cycle" | "ca
 	async saved(_cmd, { config, persist, note }) {
 		config.showSaved = !config.showSaved;
 		await persist();
-		note(`tokenburn estimated cache savings in the status line: ${onOff(config.showSaved)}`);
+		note(`tokenburn estimated savings (cache, pi context-diet, pi squeeze) in the status line: ${onOff(config.showSaved)}`);
 	},
 };

package/extensions/features/status-line/status.ts

@@ -20,8 +20,8 @@ export function windowRate(stat: PeriodStats, w: Window, firstTs: number | undef
 	return { tokensPerHour: stat.tokens / hours, costPerHour: stat.cost / hours };
 }
 
-/** Estimated dollars saved by cache reads in a window. */
-export const statSavings = (s: PeriodStats): number => Math.max(0, s.cacheBase - s.cacheActual);
+/** Estimated dollars saved in a window: cache reads, plus what pi-context-diet and pi-squeeze kept out. */
+export const statSavings = (s: PeriodStats): number => Math.max(0, s.cacheBase - s.cacheActual) + s.dietUsd + s.squeezeUsd;
 
 /** Status-line extras for the window, in order: burn rate, cache hit ratio, cache savings (each only when enabled and meaningful). */
 export function statusExtras(report: TokenBurnReport, cfg: TokenBurnConfig, now: number): string[] {

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.7.4",
+  "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",