pi-tokenburn 0.3.1 → 0.5.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/README.md CHANGED
@@ -10,9 +10,10 @@ Totals match the `tokenburn` CLI exactly.
10
10
 
11
11
  - **Status bar:** `🔥 1.20M ($52.79) [today]`, turns red with `OVER BUDGET` past your limit.
12
12
  - **Report chart:** `/tokenburn` shows the summary plus a vertical column chart (y-axis, budget line, current period marked, over-budget columns in red, peak/avg footer). Daily (14 days), weekly (8 weeks) or monthly (6 months).
13
- - **Views:** `/tokenburn day|week|month|year|all`. Status bar period: `/tokenburn status <window>` or `/tokenburn cycle`.
13
+ - **Status line:** `/tokenburn day|week|month|year|all` (or `cycle`) changes the period shown at the bottom. Charts live under `/tokenburn chart <view>`.
14
14
  - **Window panel (optional):** Today / Week / Month / Total below the editor, with in/out tokens and cost.
15
15
  - **Budgets** per day, week and month, in tokens and/or dollars.
16
+ - **Live:** the bottom line updates after every assistant message, not only when a turn ends, at no measurable cost (see Performance).
16
17
  - **Fast:** per-file cache keyed on mtime+size. First scan of ~400 MB of sessions takes under a second; refreshes after that take milliseconds. Refresh never blocks the agent loop.
17
18
 
18
19
  ## Install
@@ -27,22 +28,22 @@ pi install git:github.com/sorinirimies/pi-tokenburn
27
28
 
28
29
  | Command | What it does |
29
30
  |---|---|
30
- | `/tokenburn` · `day` · `today` | Report + **daily** chart (last 14 days) |
31
- | `/tokenburn week` | Report + **weekly** chart (last 8 weeks) |
32
- | `/tokenburn month` | Report + **monthly** chart (last 6 months) |
33
- | `/tokenburn year` | Report + **yearly** chart (last 5 years) |
34
- | `/tokenburn all` | Report + whole history (monthly columns up to 18 months, yearly beyond) |
35
- | `/tokenburn report <view>` · `chart <view>` | Same as above, explicit form |
36
- | `/tokenburn status <today\|week\|month\|year\|total>` | Set the status-bar period |
37
- | `/tokenburn cycle` | Rotate the status-bar period: today → week → month → year → total |
31
+ | `/tokenburn day\|week\|month\|year\|all` | **Switch the bottom status line**: today, this week, this month, this year, or the all-time total |
32
+ | `/tokenburn chart <view>` | Report + **chart**. Views: `day` (14 days), `week` (8 weeks), `month` (6 months), `year` (5 years), `all` (whole history; monthly columns up to 18 months, yearly beyond) |
33
+ | `/tokenburn` | Report + daily chart (same as `/tokenburn chart day`) |
34
+ | `/tokenburn cycle` | Rotate the status line: today → week → month → year → total |
35
+ | `/tokenburn status <today\|week\|month\|year\|total>` | Explicit form of the first row |
38
36
  | `/tokenburn window` | Toggle the panel below the editor |
39
37
  | `/tokenburn budget <day\|week\|month> <tokens>` | Token budget, e.g. `budget day 5000000` |
40
38
  | `/tokenburn budget <day\|week\|month> $<cost>` | Dollar budget, e.g. `budget week $50` |
41
39
  | `/tokenburn cache` | Toggle counting cache read/write tokens |
40
+ | `/tokenburn live` | Toggle refreshing after each assistant message (on by default) |
41
+ | `/tokenburn disable` (`off`) | Turn tokenburn off: status line and panel cleared, **nothing runs in the background**. `chart` still works on demand |
42
+ | `/tokenburn enable` (`on`) | Turn it back on. Both are remembered across sessions |
42
43
 
43
- Tab-completion works for all of the above. Unknown input prints usage instead of failing.
44
+ Tab-completion works for all of these. Unknown input prints usage instead of failing. `report <view>` is kept as an alias of `chart <view>`.
44
45
 
45
- Chart example (`/tokenburn`, with a 600M day budget):
46
+ Chart example (`/tokenburn chart day`, with a 600M day budget):
46
47
 
47
48
  ```
48
49
  Daily tokens · last 14 days ┄ budget 600.00M
@@ -66,6 +67,9 @@ Weeks start on Monday, in local time.
66
67
 
67
68
  ```json
68
69
  {
70
+ "enabled": true,
71
+ "liveRefresh": true,
72
+ "refreshMs": 750,
69
73
  "includeCache": true,
70
74
  "chartColor": true,
71
75
  "showStatus": true,
@@ -77,10 +81,27 @@ Weeks start on Monday, in local time.
77
81
 
78
82
  - `includeCache` (default `true`): count cache read/write tokens in totals, like the `tokenburn` CLI. Cache reads dominate (often 95%+ of tokens) but are billed far cheaper; set `false` to budget on input + output only.
79
83
 
84
+ - `enabled` (default `true`): master switch, same as `/tokenburn enable|disable`.
85
+ - `liveRefresh` (default `true`): refresh after each assistant message. `false` = only when a turn ends. `refreshMs` (default `750`, range 250 to 60000) is the coalescing delay.
80
86
  - `chartColor` (default `true`): color the chart using your pi theme (current period green, over-budget red, budget line yellow). Set `false` for plain text.
81
87
 
82
88
  Environment overrides for the session directory: `TOKENBURN_PI_SESSIONS`, `PI_CODING_AGENT_SESSION_DIR`, `PI_CODING_AGENT_DIR`.
83
89
 
90
+ ## Performance
91
+
92
+ Refreshing is designed to cost nothing you can notice:
93
+
94
+ - **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.
95
+ - **Coalesced:** a burst of messages triggers one refresh (trailing delay, `refreshMs`), never two at once.
96
+ - **Never blocks:** refreshes are fire-and-forget and errors are swallowed; the timer is unref'd and cancelled at shutdown.
97
+ - **Off means off:** `/tokenburn disable` (or `"enabled": false`) reads no files and starts no timers.
98
+
99
+ ## Security notes
100
+
101
+ - **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**.
102
+ - **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.
103
+ - **Hand-edited config is sanitised.** Unknown or non-numeric budget values are dropped.
104
+
84
105
  ## Development
85
106
 
86
107
  ```bash
@@ -15,7 +15,7 @@
15
15
  */
16
16
 
17
17
  import { readdirSync, statSync } from "node:fs";
18
- import { open } from "node:fs/promises";
18
+ import { open, type FileHandle } from "node:fs/promises";
19
19
  import { homedir } from "node:os";
20
20
  import { join } from "node:path";
21
21
  import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
@@ -52,6 +52,12 @@ export const STATUS_WINDOWS = ["today", "week", "month", "year", "total"] as con
52
52
  export type Window = (typeof STATUS_WINDOWS)[number];
53
53
 
54
54
  export interface TokenBurnConfig {
55
+ /** Master switch. When false nothing runs in the background and the status line / panel are cleared. */
56
+ enabled: boolean;
57
+ /** Refresh after each assistant message (coalesced), not only when a turn ends. */
58
+ liveRefresh: boolean;
59
+ /** Trailing delay for live refreshes, in ms (250 to 60000). */
60
+ refreshMs: number;
55
61
  /** Count cache read/write tokens in totals (matches the tokenburn CLI). */
56
62
  includeCache: boolean;
57
63
  /** Color the report chart (set false if your terminal garbles ANSI in notifications). */
@@ -70,6 +76,9 @@ export interface TokenBurnConfig {
70
76
  }
71
77
 
72
78
  const DEFAULT_CONFIG: TokenBurnConfig = {
79
+ enabled: true,
80
+ liveRefresh: true,
81
+ refreshMs: 750,
73
82
  includeCache: true,
74
83
  chartColor: true,
75
84
  showStatus: true,
@@ -101,6 +110,9 @@ export async function loadConfig(): Promise<TokenBurnConfig> {
101
110
  const raw = await readFile(join(getConfigDir(), "tokenburn.json"), "utf8");
102
111
  const parsed = JSON.parse(raw);
103
112
  return {
113
+ enabled: typeof parsed.enabled === "boolean" ? parsed.enabled : DEFAULT_CONFIG.enabled,
114
+ liveRefresh: typeof parsed.liveRefresh === "boolean" ? parsed.liveRefresh : DEFAULT_CONFIG.liveRefresh,
115
+ refreshMs: sanitizeRefreshMs(parsed.refreshMs),
104
116
  includeCache: typeof parsed.includeCache === "boolean" ? parsed.includeCache : DEFAULT_CONFIG.includeCache,
105
117
  chartColor: typeof parsed.chartColor === "boolean" ? parsed.chartColor : DEFAULT_CONFIG.chartColor,
106
118
  showStatus: typeof parsed.showStatus === "boolean" ? parsed.showStatus : DEFAULT_CONFIG.showStatus,
@@ -108,13 +120,32 @@ export async function loadConfig(): Promise<TokenBurnConfig> {
108
120
  statusWindow: (STATUS_WINDOWS as readonly string[]).includes(parsed.statusWindow)
109
121
  ? parsed.statusWindow
110
122
  : DEFAULT_CONFIG.statusWindow,
111
- budgets: typeof parsed.budgets === "object" && parsed.budgets !== null ? { ...parsed.budgets } : {},
123
+ budgets: sanitizeBudgets(parsed.budgets),
112
124
  };
113
125
  } catch {
114
126
  return freshConfig();
115
127
  }
116
128
  }
117
129
 
130
+ /** Live-refresh delay: a finite number clamped to 250..60000 ms, else the default. */
131
+ export function sanitizeRefreshMs(value: unknown): number {
132
+ if (typeof value !== "number" || !Number.isFinite(value)) return DEFAULT_CONFIG.refreshMs;
133
+ return Math.min(60_000, Math.max(250, Math.round(value)));
134
+ }
135
+
136
+ const BUDGET_KEYS = ["dayTokens", "weekTokens", "monthTokens", "dayCost", "weekCost", "monthCost"] as const;
137
+
138
+ /** Keep only known budget keys holding positive, finite numbers (the file is hand-editable). */
139
+ export function sanitizeBudgets(value: unknown): TokenBurnConfig["budgets"] {
140
+ const out: TokenBurnConfig["budgets"] = {};
141
+ if (typeof value !== "object" || value === null || Array.isArray(value)) return out;
142
+ for (const key of BUDGET_KEYS) {
143
+ const n = (value as Record<string, unknown>)[key];
144
+ if (typeof n === "number" && Number.isFinite(n) && n > 0) out[key] = n;
145
+ }
146
+ return out;
147
+ }
148
+
118
149
  /** A config that shares no mutable state with DEFAULT_CONFIG. */
119
150
  export function freshConfig(): TokenBurnConfig {
120
151
  return { ...DEFAULT_CONFIG, budgets: {} };
@@ -145,44 +176,157 @@ interface Turn {
145
176
  cost: number;
146
177
  }
147
178
 
148
- /** Parsed turns per session file, reused while the file is unchanged (mtime + size). */
149
- const FILE_CACHE = new Map<string, { sig: string; turns: Turn[] }>();
179
+ /** What was read from disk, for tests and tuning: refreshes must cost O(new bytes). */
180
+ export const stats = { bytesRead: 0, fullParses: 0, incrementalParses: 0 };
150
181
 
151
- async function parseFile(file: string): Promise<Turn[]> {
152
- const turns: Turn[] = [];
182
+ export function resetStats(): void {
183
+ stats.bytesRead = 0;
184
+ stats.fullParses = 0;
185
+ stats.incrementalParses = 0;
186
+ }
187
+
188
+ /**
189
+ * Per-session-file parse state. Pi session logs are append-only JSONL, so after the first
190
+ * parse only the bytes added since then are read. The state is dropped (full reparse) when
191
+ * the file shrank, was rewritten in place, or its head / last-parsed bytes no longer match.
192
+ */
193
+ interface FileEntry {
194
+ size: number;
195
+ mtimeMs: number;
196
+ /** End of the last COMPLETE (newline-terminated) line that was parsed. */
197
+ offset: number;
198
+ turns: Turn[];
199
+ /** Turn from an unterminated final line; counted now, re-parsed once the line completes. */
200
+ tail: Turn[];
201
+ /** First bytes of the file, to notice a rewrite. */
202
+ head: Buffer;
203
+ /** The bytes just before `offset`, same purpose. */
204
+ guard: Buffer;
205
+ }
206
+
207
+ const FILE_CACHE = new Map<string, FileEntry>();
208
+ const GUARD_BYTES = 64;
209
+ const HEAD_BYTES = 256;
210
+ const CHUNK_BYTES = 1 << 20;
211
+
212
+ /** A Turn from one JSONL line, or undefined when it carries no usage. */
213
+ function turnFromLine(line: Buffer): Turn | undefined {
214
+ if (!line.includes('"usage"')) return undefined; // most lines: skip without a JSON parse
153
215
  try {
154
- const fd = await open(file, "r");
155
- try {
156
- for await (const line of fd.readLines({ encoding: "utf8" })) {
157
- if (!line.includes('"usage"')) continue;
158
- try {
159
- const obj = JSON.parse(line);
160
- const usage = obj.message?.usage;
161
- if (!obj.timestamp || !usage) continue;
162
- const ts = new Date(obj.timestamp).getTime();
163
- if (!Number.isFinite(ts)) continue;
164
- turns.push({
165
- ts,
166
- input: Number(usage.input) || 0,
167
- output: Number(usage.output) || 0,
168
- cacheRead: Number(usage.cacheRead) || 0,
169
- cacheWrite: Number(usage.cacheWrite) || 0,
170
- cost: Number(usage.cost?.total) || 0,
171
- });
172
- } catch {}
173
- }
174
- } finally {
175
- await fd.close();
216
+ const obj = JSON.parse(line.toString("utf8"));
217
+ const usage = obj.message?.usage;
218
+ if (!obj.timestamp || !usage) return undefined;
219
+ const ts = new Date(obj.timestamp).getTime();
220
+ if (!Number.isFinite(ts)) return undefined;
221
+ return {
222
+ ts,
223
+ input: Number(usage.input) || 0,
224
+ output: Number(usage.output) || 0,
225
+ cacheRead: Number(usage.cacheRead) || 0,
226
+ cacheWrite: Number(usage.cacheWrite) || 0,
227
+ cost: Number(usage.cost?.total) || 0,
228
+ };
229
+ } catch {
230
+ return undefined;
231
+ }
232
+ }
233
+
234
+ async function readAt(fd: FileHandle, buf: Buffer, position: number): Promise<number> {
235
+ const { bytesRead } = await fd.read(buf, 0, buf.length, position);
236
+ stats.bytesRead += bytesRead;
237
+ return bytesRead;
238
+ }
239
+
240
+ /** Parse complete lines in [start, end). Returns their turns, where they stop, and the unterminated rest. */
241
+ async function scan(
242
+ fd: FileHandle,
243
+ start: number,
244
+ end: number,
245
+ chunkSize: number,
246
+ ): Promise<{ turns: Turn[]; consumed: number; rest: Buffer }> {
247
+ const turns: Turn[] = [];
248
+ const buf = Buffer.allocUnsafe(Math.min(chunkSize, Math.max(1, end - start)));
249
+ let pos = start;
250
+ let consumed = start;
251
+ let carry: Buffer = Buffer.alloc(0);
252
+ while (pos < end) {
253
+ const want = Math.min(buf.length, end - pos);
254
+ const n = await readAt(fd, want === buf.length ? buf : buf.subarray(0, want), pos);
255
+ if (n === 0) break;
256
+ const base = pos - carry.length;
257
+ const data = carry.length > 0 ? Buffer.concat([carry, buf.subarray(0, n)]) : buf.subarray(0, n);
258
+ let lineStart = 0;
259
+ for (let nl = data.indexOf(10, lineStart); nl !== -1; nl = data.indexOf(10, lineStart)) {
260
+ const t = turnFromLine(data.subarray(lineStart, nl));
261
+ if (t) turns.push(t);
262
+ lineStart = nl + 1;
176
263
  }
177
- } catch {}
178
- return turns;
264
+ if (lineStart > 0) consumed = base + lineStart;
265
+ carry = Buffer.from(data.subarray(lineStart)); // copy: `buf` is reused
266
+ pos += n;
267
+ }
268
+ return { turns, consumed, rest: carry };
179
269
  }
180
270
 
181
- /** Cache key for a file: mtime + size. `undefined` when it cannot be stat'ed. */
182
- export function fileSignature(file: string): string | undefined {
271
+ async function bytesAt(fd: FileHandle, position: number, length: number): Promise<Buffer> {
272
+ if (length <= 0) return Buffer.alloc(0);
273
+ const buf = Buffer.alloc(length);
274
+ const n = await readAt(fd, buf, position);
275
+ return buf.subarray(0, n);
276
+ }
277
+
278
+ /** Is the prefix we parsed earlier still what is on disk? */
279
+ async function stillAppendOnly(fd: FileHandle, prev: FileEntry): Promise<boolean> {
280
+ const head = await bytesAt(fd, 0, prev.head.length);
281
+ if (!head.equals(prev.head)) return false;
282
+ const guard = await bytesAt(fd, prev.offset - prev.guard.length, prev.guard.length);
283
+ return guard.equals(prev.guard);
284
+ }
285
+
286
+ async function refreshEntry(
287
+ file: string,
288
+ st: { size: number; mtimeMs: number },
289
+ prev: FileEntry | undefined,
290
+ chunkSize: number,
291
+ ): Promise<FileEntry | undefined> {
292
+ let fd: FileHandle;
293
+ try {
294
+ fd = await open(file, "r");
295
+ } catch {
296
+ return undefined;
297
+ }
298
+ try {
299
+ const grew = prev !== undefined && st.size > prev.size;
300
+ const incremental = grew && (await stillAppendOnly(fd, prev!));
301
+ const start = incremental ? prev!.offset : 0;
302
+ const turns = incremental ? prev!.turns : [];
303
+ if (incremental) stats.incrementalParses++;
304
+ else stats.fullParses++;
305
+
306
+ const r = await scan(fd, start, st.size, chunkSize);
307
+ for (const t of r.turns) turns.push(t);
308
+ const tailTurn = r.rest.length > 0 ? turnFromLine(r.rest) : undefined;
309
+ return {
310
+ size: st.size,
311
+ mtimeMs: st.mtimeMs,
312
+ offset: r.consumed,
313
+ turns,
314
+ tail: tailTurn ? [tailTurn] : [],
315
+ head: incremental ? prev!.head : await bytesAt(fd, 0, Math.min(HEAD_BYTES, st.size)),
316
+ guard: await bytesAt(fd, Math.max(0, r.consumed - GUARD_BYTES), Math.min(GUARD_BYTES, r.consumed)),
317
+ };
318
+ } catch {
319
+ return undefined;
320
+ } finally {
321
+ await fd.close();
322
+ }
323
+ }
324
+
325
+ /** Size and mtime of a file, or undefined when it cannot be stat'ed. */
326
+ export function fileStat(file: string): { size: number; mtimeMs: number } | undefined {
183
327
  try {
184
328
  const st = statSync(file);
185
- return `${st.mtimeMs}:${st.size}`;
329
+ return { size: st.size, mtimeMs: st.mtimeMs };
186
330
  } catch {
187
331
  return undefined;
188
332
  }
@@ -217,26 +361,32 @@ const emptyStats = (): PeriodStats => ({
217
361
  sessions: 0,
218
362
  });
219
363
 
220
- /** Refresh the per-file cache and return every session file's parsed turns. */
221
- async function loadAll(sessionsRoot: string): Promise<Turn[][]> {
364
+ /** Bring the per-file cache up to date (reading only what changed) and return every file's turns. */
365
+ async function loadAll(sessionsRoot: string, chunkSize = CHUNK_BYTES): Promise<Turn[][]> {
222
366
  const files = listSessionFiles(sessionsRoot);
223
367
  const live = new Set(files);
224
368
  for (const key of FILE_CACHE.keys()) if (!live.has(key)) FILE_CACHE.delete(key);
225
369
 
226
370
  const out: Turn[][] = [];
227
371
  for (const file of files) {
228
- const sig = fileSignature(file);
229
- if (sig === undefined) continue; // vanished between listing and stat
372
+ const st = fileStat(file);
373
+ if (st === undefined) continue; // vanished between listing and stat
230
374
  let entry = FILE_CACHE.get(file);
231
- if (!entry || entry.sig !== sig) {
232
- entry = { sig, turns: await parseFile(file) };
233
- FILE_CACHE.set(file, entry);
375
+ if (!entry || entry.size !== st.size || entry.mtimeMs !== st.mtimeMs) {
376
+ const fresh = await refreshEntry(file, st, entry, chunkSize);
377
+ if (fresh) FILE_CACHE.set(file, fresh);
378
+ else FILE_CACHE.delete(file);
379
+ entry = fresh;
234
380
  }
235
- out.push(entry.turns);
381
+ if (entry) out.push(entry.tail.length > 0 ? entry.turns.concat(entry.tail) : entry.turns);
236
382
  }
237
383
  return out;
238
384
  }
239
385
 
386
+ /** Test hooks: parse with a tiny chunk size (chunk boundaries) / refresh one file directly. */
387
+ export const __loadAllForTests = loadAll;
388
+ export const __refreshEntryForTests = refreshEntry;
389
+
240
390
  const turnTokens = (t: Turn, includeCache: boolean) =>
241
391
  t.input + t.output + (includeCache ? t.cacheRead + t.cacheWrite : 0);
242
392
 
@@ -505,6 +655,80 @@ export function isOverBudget(p: PeriodStats, b: { tokens?: number; cost?: number
505
655
  return (b.tokens !== undefined && p.tokens > b.tokens) || (b.cost !== undefined && p.cost > b.cost);
506
656
  }
507
657
 
658
+ // ---------------------------------------------------------------------------
659
+ // Live refresh scheduling
660
+ // ---------------------------------------------------------------------------
661
+
662
+ export interface Timers {
663
+ set: (fn: () => void, ms: number) => unknown;
664
+ clear: (handle: unknown) => void;
665
+ }
666
+
667
+ const realTimers: Timers = {
668
+ set: (fn, ms) => setTimeout(fn, ms),
669
+ clear: (h) => clearTimeout(h as ReturnType<typeof setTimeout>),
670
+ };
671
+
672
+ /**
673
+ * Runs `run` at most once per `delayMs` window however often `schedule()` is called
674
+ * (trailing edge), and never two runs at once: a request that arrives mid-run triggers one
675
+ * more run afterwards. A burst of 50 messages costs one refresh. Failures are swallowed;
676
+ * the timer never keeps the process alive.
677
+ */
678
+ export function createCoalescer(
679
+ run: () => Promise<void>,
680
+ delayMs: number | (() => number),
681
+ timers: Timers = realTimers,
682
+ ) {
683
+ let timer: unknown;
684
+ let running = false;
685
+ let dirty = false;
686
+
687
+ const execute = async () => {
688
+ if (running) {
689
+ dirty = true;
690
+ return;
691
+ }
692
+ running = true;
693
+ try {
694
+ await run();
695
+ } catch {
696
+ /* a failed refresh must never surface */
697
+ } finally {
698
+ running = false;
699
+ }
700
+ if (dirty) {
701
+ dirty = false;
702
+ schedule();
703
+ }
704
+ };
705
+
706
+ const schedule = () => {
707
+ if (timer !== undefined) return; // one is already pending: coalesce
708
+ const ms = typeof delayMs === "function" ? delayMs() : delayMs;
709
+ timer = timers.set(() => {
710
+ timer = undefined;
711
+ void execute();
712
+ }, ms);
713
+ (timer as { unref?: () => void } | undefined)?.unref?.();
714
+ };
715
+
716
+ const cancel = () => {
717
+ if (timer !== undefined) timers.clear(timer);
718
+ timer = undefined;
719
+ dirty = false;
720
+ };
721
+
722
+ /** Run now (dropping any pending delayed run). */
723
+ const flush = () => {
724
+ if (timer !== undefined) timers.clear(timer);
725
+ timer = undefined;
726
+ return execute();
727
+ };
728
+
729
+ return { schedule, cancel, flush, isPending: () => timer !== undefined };
730
+ }
731
+
508
732
  /** Map chart roles to pi theme colors. */
509
733
  export function themePaint(theme: { fg: (c: any, t: string) => string }): Paint {
510
734
  const color: Record<string, string> = {
@@ -524,6 +748,12 @@ export default function tokenBurnExtension(pi: ExtensionAPI) {
524
748
 
525
749
  const updateUI = async (ctx: ExtensionContext) => {
526
750
  if (!ctx.hasUI) return;
751
+ if (!config.enabled) {
752
+ // Disabled: show nothing and read nothing.
753
+ ctx.ui.setStatus("tokenburn", "");
754
+ ctx.ui.setWidget("tokenburn-panel", []);
755
+ return;
756
+ }
527
757
  const report = await collectUsage(getSessionsDir(), new Date(), config.includeCache);
528
758
  const theme = ctx.ui.theme;
529
759
 
@@ -569,64 +799,121 @@ export default function tokenBurnExtension(pi: ExtensionAPI) {
569
799
  }
570
800
  };
571
801
 
802
+ // The latest context, used by the (delayed) live refresh. Dropped at shutdown so a reloaded
803
+ // runtime never touches a stale one.
804
+ let latestCtx: ExtensionContext | undefined;
805
+ const refresher = createCoalescer(
806
+ async () => {
807
+ if (latestCtx && config.enabled) await updateUI(latestCtx);
808
+ },
809
+ () => config.refreshMs,
810
+ );
811
+
572
812
  pi.on("session_start", async (_event, ctx) => {
813
+ latestCtx = ctx;
573
814
  config = await loadConfig();
574
- await updateUI(ctx);
815
+ // A failing refresh must never break session start.
816
+ await updateUI(ctx).catch(() => {});
817
+ });
818
+
819
+ // Live: after each assistant message, coalesced. Never blocks, never throws into pi, and
820
+ // must return undefined (a message_end handler may otherwise replace the message).
821
+ pi.on("message_end", (event, ctx) => {
822
+ if (!config.enabled || !config.liveRefresh) return undefined;
823
+ if ((event.message as { role?: string } | undefined)?.role !== "assistant") return undefined;
824
+ latestCtx = ctx;
825
+ refresher.schedule();
826
+ return undefined;
575
827
  });
576
828
 
577
- // Fire-and-forget: never block the agent loop on a stats refresh.
829
+ // Final, exact refresh when a turn ends (fire-and-forget).
578
830
  pi.on("turn_end", (_event, ctx) => {
579
- void updateUI(ctx).catch(() => {});
831
+ if (!config.enabled) return;
832
+ latestCtx = ctx;
833
+ void refresher.flush();
834
+ });
835
+
836
+ pi.on("session_shutdown", () => {
837
+ refresher.cancel();
838
+ latestCtx = undefined;
580
839
  });
581
840
 
582
841
  pi.registerCommand("tokenburn", {
583
- description: "TokenBurn: /tokenburn [day|week|month|year|all] chart, status <window>, cycle, window, budget, cache",
842
+ description: "TokenBurn: /tokenburn [day|week|month|year|all] sets the status bar; /tokenburn chart <day|week|month|year|all> shows the chart",
584
843
  getArgumentCompletions: (prefix: string) => completions(prefix),
585
844
  handler: async (args, ctx) => {
586
845
  const cmd = parseCommand(args);
846
+ latestCtx = ctx;
587
847
  const persist = async () => {
588
848
  await saveConfig(config);
589
849
  await updateUI(ctx);
590
850
  };
851
+ // Settings changed while disabled are kept, but nothing shows until it is enabled.
852
+ const note = (msg: string) => ctx.ui.notify(config.enabled ? msg : `${msg} (tokenburn is disabled; /tokenburn enable to show it)`, "info");
591
853
 
592
854
  switch (cmd.kind) {
593
855
  case "window":
594
856
  config.showWidget = !config.showWidget;
595
857
  await persist();
596
- ctx.ui.notify(`tokenburn widget: ${config.showWidget ? "enabled" : "hidden"}`, "info");
858
+ note(`tokenburn widget: ${config.showWidget ? "enabled" : "hidden"}`);
597
859
  return;
598
860
 
599
861
  case "status":
600
862
  config.statusWindow = cmd.window;
601
863
  await persist();
602
- ctx.ui.notify(`tokenburn status window: ${cmd.window}`, "info");
864
+ note(`tokenburn status window: ${cmd.window}`);
603
865
  return;
604
866
 
605
867
  case "cycle": {
606
868
  const next = STATUS_WINDOWS[(STATUS_WINDOWS.indexOf(config.statusWindow) + 1) % STATUS_WINDOWS.length];
607
869
  config.statusWindow = next;
608
870
  await persist();
609
- ctx.ui.notify(`tokenburn status window: ${next}`, "info");
871
+ note(`tokenburn status window: ${next}`);
610
872
  return;
611
873
  }
612
874
 
613
875
  case "cache":
614
876
  config.includeCache = !config.includeCache;
615
877
  await persist();
616
- ctx.ui.notify(`tokenburn: cache tokens ${config.includeCache ? "included" : "excluded"}`, "info");
878
+ note(`tokenburn: cache tokens ${config.includeCache ? "included" : "excluded"}`);
617
879
  return;
618
880
 
619
881
  case "budget": {
620
882
  const key = `${cmd.period}${cmd.isCost ? "Cost" : "Tokens"}` as keyof TokenBurnConfig["budgets"];
621
883
  config.budgets[key] = cmd.amount;
622
884
  await persist();
623
- ctx.ui.notify(
624
- `tokenburn ${cmd.period} budget: ${cmd.isCost ? formatCost(cmd.amount) : formatNum(cmd.amount) + " tokens"}`,
625
- "info",
626
- );
885
+ note(`tokenburn ${cmd.period} budget: ${cmd.isCost ? formatCost(cmd.amount) : formatNum(cmd.amount) + " tokens"}`);
627
886
  return;
628
887
  }
629
888
 
889
+ case "enable":
890
+ if (config.enabled) {
891
+ ctx.ui.notify("tokenburn is already enabled", "info");
892
+ return;
893
+ }
894
+ config.enabled = true;
895
+ await persist();
896
+ ctx.ui.notify("tokenburn enabled", "info");
897
+ return;
898
+
899
+ case "disable":
900
+ if (!config.enabled) {
901
+ ctx.ui.notify("tokenburn is already disabled", "info");
902
+ return;
903
+ }
904
+ config.enabled = false;
905
+ refresher.cancel();
906
+ await persist(); // clears the status line and the panel
907
+ ctx.ui.notify("tokenburn disabled: nothing runs in the background. /tokenburn enable turns it back on (charts still work on demand)", "info");
908
+ return;
909
+
910
+ case "live":
911
+ config.liveRefresh = !config.liveRefresh;
912
+ if (!config.liveRefresh) refresher.cancel();
913
+ await persist();
914
+ note(`tokenburn live refresh ${config.liveRefresh ? "on (after each assistant message)" : "off (only when a turn ends)"}`);
915
+ return;
916
+
630
917
  case "invalid":
631
918
  ctx.ui.notify(cmd.message, "warning");
632
919
  return;
@@ -666,7 +953,7 @@ export default function tokenBurnExtension(pi: ExtensionAPI) {
666
953
  rule,
667
954
  ...chart,
668
955
  rule,
669
- "/tokenburn day|week|month|year|all · status <window> · cycle · window · budget · cache",
956
+ "/tokenburn chart <day|week|month|year|all> · /tokenburn <day|week|month|year|all> sets the status bar · cycle · window · budget · cache",
670
957
  ].join("\n"),
671
958
  "info",
672
959
  );
@@ -691,28 +978,44 @@ export type Command =
691
978
  | { kind: "cycle" }
692
979
  | { kind: "window" }
693
980
  | { kind: "cache" }
981
+ | { kind: "enable" }
982
+ | { kind: "disable" }
983
+ | { kind: "live" }
694
984
  | { kind: "budget"; period: (typeof BUDGET_PERIODS)[number]; amount: number; isCost: boolean }
695
985
  | { kind: "invalid"; message: string };
696
986
 
697
987
  export const USAGE =
698
- "Usage: /tokenburn [day|week|month|year|all] · status <today|week|month|year|total> · cycle · window · budget <day|week|month> <tokens|$cost> · cache";
988
+ "Usage: /tokenburn [day|week|month|year|all] (status-bar period) · chart [day|week|month|year|all] · status <today|week|month|year|total> · cycle · window · budget <day|week|month> <tokens|$cost> · cache · live · enable · disable";
989
+
990
+ /** Words that pick the status-bar period: `day` shows today, `all` shows the all-time total. */
991
+ const WORD_TO_WINDOW: Record<string, Window> = {
992
+ day: "today",
993
+ today: "today",
994
+ week: "week",
995
+ month: "month",
996
+ year: "year",
997
+ all: "total",
998
+ total: "total",
999
+ };
699
1000
 
1001
+ /** Views the chart can show. */
700
1002
  const asView = (w: string | undefined): ReportView | undefined =>
701
- w === "today" ? "day" : (REPORT_VIEWS as readonly string[]).includes(w ?? "") ? (w as ReportView) : undefined;
1003
+ w === "today" ? "day" : w === "total" ? "all" : (REPORT_VIEWS as readonly string[]).includes(w ?? "") ? (w as ReportView) : undefined;
702
1004
 
703
1005
  export function parseCommand(args: string | undefined): Command {
704
1006
  const [cmd, p1, p2] = (args ?? "").trim().toLowerCase().split(/\s+/).filter(Boolean);
705
1007
  if (!cmd) return { kind: "report", view: "day" };
706
1008
 
707
- const view = asView(cmd);
708
- if (view) return { kind: "report", view };
1009
+ // /tokenburn day|week|month|year|all → what the bottom status line shows
1010
+ const window = WORD_TO_WINDOW[cmd];
1011
+ if (window) return { kind: "status", window };
709
1012
 
710
1013
  switch (cmd) {
711
- case "report":
712
- case "chart": {
1014
+ case "chart":
1015
+ case "report": {
713
1016
  if (p1 === undefined) return { kind: "report", view: "day" };
714
1017
  const v = asView(p1);
715
- return v ? { kind: "report", view: v } : { kind: "invalid", message: `Unknown view "${p1}". ${USAGE}` };
1018
+ return v ? { kind: "report", view: v } : { kind: "invalid", message: `Unknown chart "${p1}". ${USAGE}` };
716
1019
  }
717
1020
  case "status":
718
1021
  return (STATUS_WINDOWS as readonly string[]).includes(p1 ?? "")
@@ -725,6 +1028,14 @@ export function parseCommand(args: string | undefined): Command {
725
1028
  return { kind: "window" };
726
1029
  case "cache":
727
1030
  return { kind: "cache" };
1031
+ case "enable":
1032
+ case "on":
1033
+ return { kind: "enable" };
1034
+ case "disable":
1035
+ case "off":
1036
+ return { kind: "disable" };
1037
+ case "live":
1038
+ return { kind: "live" };
728
1039
  case "budget": {
729
1040
  const isCost = (p2 ?? "").startsWith("$");
730
1041
  const amount = Number((p2 ?? "").replace(/^\$/, ""));
@@ -745,17 +1056,29 @@ export interface CompletionItem {
745
1056
  description?: string;
746
1057
  }
747
1058
 
1059
+ const CHART_SPAN: Record<ReportView, string> = {
1060
+ day: "last 14 days",
1061
+ week: "last 8 weeks",
1062
+ month: "last 6 months",
1063
+ year: "last 5 years",
1064
+ all: "whole history",
1065
+ };
1066
+
748
1067
  const TOP_LEVEL: CompletionItem[] = [
749
- { value: "day", label: "day", description: "Report + daily chart (last 14 days)" },
750
- { value: "week", label: "week", description: "Report + weekly chart (last 8 weeks)" },
751
- { value: "month", label: "month", description: "Report + monthly chart (last 6 months)" },
752
- { value: "year", label: "year", description: "Report + yearly chart (last 5 years)" },
753
- { value: "all", label: "all", description: "Report + whole history" },
754
- { value: "status ", label: "status", description: "Set the status-bar period" },
1068
+ { value: "day", label: "day", description: "Status bar: today" },
1069
+ { value: "week", label: "week", description: "Status bar: this week" },
1070
+ { value: "month", label: "month", description: "Status bar: this month" },
1071
+ { value: "year", label: "year", description: "Status bar: this year" },
1072
+ { value: "all", label: "all", description: "Status bar: all-time total" },
1073
+ { value: "chart ", label: "chart", description: "Report + chart: day, week, month, year or all" },
1074
+ { value: "status ", label: "status", description: "Set the status-bar period explicitly" },
755
1075
  { value: "cycle", label: "cycle", description: "Rotate the status-bar period" },
756
1076
  { value: "window", label: "window", description: "Toggle the panel below the editor" },
757
1077
  { value: "budget ", label: "budget", description: "Set a day/week/month budget" },
758
1078
  { value: "cache", label: "cache", description: "Toggle counting cache tokens" },
1079
+ { value: "live", label: "live", description: "Toggle refreshing after each message" },
1080
+ { value: "enable", label: "enable", description: "Turn tokenburn on (status line, panel, refresh)" },
1081
+ { value: "disable", label: "disable", description: "Turn tokenburn off: nothing runs in the background" },
759
1082
  ];
760
1083
 
761
1084
  /** `prefix` is the whole argument text typed so far (pi replaces it with `value`). */
@@ -771,7 +1094,7 @@ export function completions(prefix: string): CompletionItem[] | null {
771
1094
  pool = sub("budget", BUDGET_PERIODS, (w) => `Set the ${w} budget (tokens or $cost)`, " ");
772
1095
  else if (/^(report|chart) /.test(text)) {
773
1096
  const head = text.split(" ")[0];
774
- pool = sub(head, REPORT_VIEWS, (w) => `${head} for ${w}`);
1097
+ pool = sub(head, REPORT_VIEWS, (w) => `Chart: ${CHART_SPAN[w as ReportView]}`);
775
1098
  } else pool = TOP_LEVEL;
776
1099
 
777
1100
  const out = pool.filter((i) => i.value.startsWith(text));
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "pi-tokenburn",
3
- "version": "0.3.1",
3
+ "version": "0.5.0",
4
4
  "description": "Pi extension for live token usage monitoring and budget tracking (Day / Week / Month / Total). 0 prompt tokens.",
5
5
  "keywords": ["pi-package", "pi-extension", "tokenburn", "tokens", "analytics", "budget", "stats"],
6
6
  "license": "MIT",