pi-tokenburn 0.4.0 โ†’ 0.5.1

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.
Files changed (3) hide show
  1. package/README.md +101 -17
  2. package/extensions/tokenburn.ts +339 -52
  3. package/package.json +1 -1
package/README.md CHANGED
@@ -1,11 +1,18 @@
1
1
  # pi-tokenburn
2
2
 
3
+ [![npm](https://img.shields.io/npm/v/pi-tokenburn.svg)](https://www.npmjs.com/package/pi-tokenburn)
4
+ [![npm downloads](https://img.shields.io/npm/dm/pi-tokenburn.svg)](https://www.npmjs.com/package/pi-tokenburn)
5
+ [![CI](https://github.com/sorinirimies/pi-tokenburn/actions/workflows/ci.yml/badge.svg)](https://github.com/sorinirimies/pi-tokenburn/actions/workflows/ci.yml)
6
+ [![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)
7
+
3
8
  Token usage and cost monitor with budgets, directly in your [Pi](https://github.com/earendil-works/pi-coding-agent) session.
4
9
 
5
10
  TypeScript port of the pi collector from [tokenburn](https://github.com/sorinirimies/tokenburn).
6
11
  Reads Pi's local session logs (`~/.pi/agent/sessions/**/*.jsonl`). **No network, zero prompt tokens.**
7
12
  Totals match the `tokenburn` CLI exactly.
8
13
 
14
+ <img src="examples/vhs/generated/overview.gif" alt="pi-tokenburn: the live footer counter, a report with a daily chart, switching the footer period and a weekly chart" width="900">
15
+
9
16
  ## Features
10
17
 
11
18
  - **Status bar:** `๐Ÿ”ฅ 1.20M ($52.79) [today]`, turns red with `OVER BUDGET` past your limit.
@@ -13,6 +20,7 @@ Totals match the `tokenburn` CLI exactly.
13
20
  - **Status line:** `/tokenburn day|week|month|year|all` (or `cycle`) changes the period shown at the bottom. Charts live under `/tokenburn chart <view>`.
14
21
  - **Window panel (optional):** Today / Week / Month / Total below the editor, with in/out tokens and cost.
15
22
  - **Budgets** per day, week and month, in tokens and/or dollars.
23
+ - **Live:** the bottom line updates after every assistant message, not only when a turn ends, at no measurable cost (see Performance).
16
24
  - **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
25
 
18
26
  ## Install
@@ -36,23 +44,57 @@ pi install git:github.com/sorinirimies/pi-tokenburn
36
44
  | `/tokenburn budget <day\|week\|month> <tokens>` | Token budget, e.g. `budget day 5000000` |
37
45
  | `/tokenburn budget <day\|week\|month> $<cost>` | Dollar budget, e.g. `budget week $50` |
38
46
  | `/tokenburn cache` | Toggle counting cache read/write tokens |
47
+ | `/tokenburn live` | Toggle refreshing after each assistant message (on by default) |
48
+ | `/tokenburn disable` (`off`) | Turn tokenburn off: status line and panel cleared, **nothing runs in the background**. `chart` still works on demand |
49
+ | `/tokenburn enable` (`on`) | Turn it back on. Both are remembered across sessions |
50
+
51
+ ### Previews
52
+
53
+ Recorded from a real pi with only this plugin loaded, on synthetic data.
54
+
55
+ **Switch what the bottom line shows** (`/tokenburn day|week|month|year|all`, or `cycle`):
56
+
57
+ ![Status line: week, month, year, all, day, then cycle](examples/vhs/generated/status.gif)
58
+
59
+ **Charts** (`/tokenburn chart day|week|month|year|all`): the current period is green, with a peak and average footer:
60
+
61
+ ![Charts: day, week, month, year and all-time](examples/vhs/generated/charts.gif)
62
+
63
+ **Budgets** (`/tokenburn budget day 150000`, `budget week $2.5`): the footer turns red past the limit, and the chart draws the budget line with over-budget columns in red:
64
+
65
+ ![Budgets: OVER BUDGET in the footer and red columns in the chart](examples/vhs/generated/budget.gif)
66
+
67
+ **Live refresh:** the footer updates after every assistant message, not only when a turn ends (here a local mock model answers; pi writes the turn to its session log and the counter moves by itself):
68
+
69
+ ![Live refresh: the counter and cost tick up as the reply lands](examples/vhs/generated/live-refresh.gif)
70
+
71
+ **Enable / disable:** off means off (nothing runs in the background), and charts still work on demand:
72
+
73
+ ![Disable clears the footer, a chart still works, enable brings it back](examples/vhs/generated/enable-disable.gif)
74
+
75
+ **Tab-completion** uses pi's own menu, with a description for every option:
76
+
77
+ ![Autocomplete for /tokenburn and its chart sub-menu](examples/vhs/generated/completion.gif)
39
78
 
40
79
  Tab-completion works for all of these. Unknown input prints usage instead of failing. `report <view>` is kept as an alias of `chart <view>`.
41
80
 
42
- Chart example (`/tokenburn chart day`, with a 600M day budget):
81
+ Chart example (`/tokenburn chart day`, with a 150k day budget):
43
82
 
44
83
  ```
45
- Daily tokens ยท last 14 days โ”„ budget 600.00M
46
- 1.13B โ”‚ โ–ˆโ–ˆ
47
- โ”‚ โ–ˆโ–ˆ
48
- โ”‚โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ–ˆโ–ˆโ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ–ˆโ–ˆโ”„โ–†โ–†โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„
49
- 565.16M โ”‚ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–‡โ–‡
50
- โ”‚ โ–…โ–… โ–ˆโ–ˆ โ–‚โ–‚ โ–ˆโ–ˆ โ–โ– โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–„โ–„ โ–‡โ–‡
51
- โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
52
- 26 27 28 29 30 01 02 03 04 05 06 07 08 09
53
- Sa Su Mo Tu We Th Fr Sa Su Mo Tu We Th Fr
54
- โ–ฒ
55
- ฮฃ 3.83B ($1264.30) ยท avg 273.80M/day ยท peak 1.13B (09-29 Tue)
84
+ Daily tokens ยท last 14 days โ”„ budget 150.0k
85
+ 291.9k โ”‚ โ–ˆโ–ˆ โ–„โ–„
86
+ โ”‚ โ–ˆโ–ˆ โ–ˆโ–ˆ
87
+ โ”‚ โ–ƒโ–ƒ โ–…โ–… โ–ˆโ–ˆ โ–ˆโ–ˆ
88
+ โ”‚โ”„โ”„โ”„โ”„โ”„โ”„โ–ˆโ–ˆโ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ”„โ–ˆโ–ˆโ”„โ”„โ”„โ”„โ”„โ”„โ”„โ–†โ–†โ”„โ”„โ”„โ”„โ–ˆโ–ˆโ”„โ–ˆโ–ˆโ”„โ–…โ–…โ”„
89
+ 146.0k โ”‚ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ
90
+ โ”‚ โ–ˆโ–ˆ โ–‚โ–‚ โ–‡โ–‡ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–‡โ–‡ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ
91
+ โ”‚โ–โ– โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ
92
+ โ”‚โ–ˆโ–ˆ โ–‡โ–‡ โ–ˆโ–ˆ โ–…โ–… โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ƒโ–ƒ โ–โ– โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ โ–ˆโ–ˆ
93
+ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
94
+ 26 27 28 29 30 01 02 03 04 05 06 07 08 09
95
+ Sa Su Mo Tu We Th Fr Sa Su Mo Tu We Th Fr
96
+ โ–ฒ
97
+ ฮฃ 1.74M ($4.81) ยท avg 124.4k/day ยท peak 291.9k (10-07 Wed)
56
98
  ```
57
99
 
58
100
  Weeks start on Monday, in local time.
@@ -63,6 +105,9 @@ Weeks start on Monday, in local time.
63
105
 
64
106
  ```json
65
107
  {
108
+ "enabled": true,
109
+ "liveRefresh": true,
110
+ "refreshMs": 750,
66
111
  "includeCache": true,
67
112
  "chartColor": true,
68
113
  "showStatus": true,
@@ -74,10 +119,32 @@ Weeks start on Monday, in local time.
74
119
 
75
120
  - `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.
76
121
 
122
+ - `enabled` (default `true`): master switch, same as `/tokenburn enable|disable`.
123
+ - `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.
77
124
  - `chartColor` (default `true`): color the chart using your pi theme (current period green, over-budget red, budget line yellow). Set `false` for plain text.
78
125
 
79
126
  Environment overrides for the session directory: `TOKENBURN_PI_SESSIONS`, `PI_CODING_AGENT_SESSION_DIR`, `PI_CODING_AGENT_DIR`.
80
127
 
128
+ ## Performance
129
+
130
+ Refreshing is designed to cost nothing you can notice:
131
+
132
+ - **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.
133
+ - **Coalesced:** a burst of messages triggers one refresh (trailing delay, `refreshMs`), never two at once.
134
+ - **Never blocks:** refreshes are fire-and-forget and errors are swallowed; the timer is unref'd and cancelled at shutdown.
135
+ - **Off means off:** `/tokenburn disable` (or `"enabled": false`) reads no files and starts no timers.
136
+
137
+ ## Demo recordings
138
+
139
+ The GIFs above live in [`examples/vhs/generated/`](examples/vhs/generated) and are stored with **Git LFS** (`git lfs install` once). They are recorded with [VHS](https://github.com/charmbracelet/vhs) from a **real pi** that loads only this extension, on **synthetic** data (`examples/vhs/fixture.sh`: 900 days of made-up sessions in `/tmp` with deliberately small numbers, and a local mock model for the live-refresh demo), never from real logs, paths or credentials:
140
+
141
+ ```sh
142
+ just vhs-all # every tape (examples/vhs/*.tape): needs vhs, ttyd, ffmpeg, pi, bun, python3
143
+ just vhs-tape charts # one tape
144
+ just vhs-list # list the tapes
145
+ just demo # try it yourself in a real pi on the same synthetic data
146
+ ```
147
+
81
148
  ## Security notes
82
149
 
83
150
  - **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,15 +169,16 @@ CI scripts are [nushell](https://www.nushell.sh) (`scripts/`), the same ones loc
102
169
 
103
170
  | Workflow | When | What |
104
171
  |---|---|---|
105
- | **CI** | push / PR | quality gate on Linux; tests on macOS, Windows and extreme time zones |
106
- | **Nightly Dependency Update** | every night (GitHub 02:00 UTC, Gitea 02:30) | `bun update` within ranges, verify on all platforms, commit `chore(deps)`; then **publish a patch** when a *runtime* dependency was upgraded or `feat`/`fix`/`perf` commits are waiting since the last tag |
107
- | **Release** | tag `vX.Y.Z` | validates the tag against `package.json`, runs the gate, `npm publish` (idempotent), creates the release |
172
+ | **CI** | push / PR | quality gate on Linux; tests on macOS and Windows, and in extreme time zones |
173
+ | **Auto-merge library updates** | CI finished on a Dependabot PR | **patch and minor** updates (GitHub Actions) are merged automatically, but only **after CI is green** on that exact commit; a **major** update waits for you. Then it starts the nightly workflow so the update ships |
174
+ | **Nightly Dependency Update** | every night (GitHub 02:00 UTC, Gitea 02:30) and after each auto-merge | `bun update` within ranges, verify on all platforms, commit `chore(deps)`; then **build, tag and publish a new patch** whenever a library was upgraded or merged, or `feat`/`fix`/`perf` commits are waiting since the last tag |
175
+ | **Release** | tag `vX.Y.Z` | validates the tag against `package.json`, runs the gate, `npm publish` (idempotent, with provenance), creates the release |
108
176
 
109
- Dev-tooling-only bumps are committed but never released on their own. A downgrade blocks the whole update.
177
+ So library updates need no human: they are merged once CI passes, built, versioned and published as a patch. A downgrade, a failing check, or a major update stops the chain and waits for review.
110
178
 
111
179
  Manual release: `just bump patch` (or `minor` / `major` / `X.Y.Z`), then `just release-push`.
112
180
 
113
- **Secrets** (repo settings): `NPM_TOKEN` โ€” npm *granular access token* with publish rights and "bypass 2FA" (required for CI publishing). Optional: `GH_PAT` (lets a tag push trigger the release itself), `GITEA_TOKEN` (Gitea).
181
+ **Secrets** (repo settings): `NPM_TOKEN` is an npm *granular access token* with publish rights and "bypass 2FA" (required for CI publishing). Optional: `GH_PAT` (lets a tag push trigger the release itself), `GITEA_TOKEN` (Gitea).
114
182
 
115
183
  Commits follow [Conventional Commits](https://www.conventionalcommits.org); the changelog is generated by git-cliff.
116
184
 
@@ -118,4 +186,20 @@ Commits follow [Conventional Commits](https://www.conventionalcommits.org); the
118
186
 
119
187
  Pi sessions only. For Zed, Claude Code, Codex, Gemini and more, use the Rust [tokenburn](https://github.com/sorinirimies/tokenburn).
120
188
 
189
+ ## Related plugins
190
+
191
+ Three small [pi](https://github.com/earendil-works/pi-coding-agent) extensions that save tokens without adding anything to the system prompt:
192
+
193
+ | Plugin | What it does |
194
+ |---|---|
195
+ | [**pi-edit-first**](https://github.com/sorinirimies/pi-edit-first) | Blocks whole-file `write` rewrites and hand-written project manifests, steering the agent to targeted `edit` calls and scaffolders |
196
+ | [**pi-read-guard**](https://github.com/sorinirimies/pi-read-guard) | Blocks full reads of large files, steering the agent to `offset`/`limit` or search |
197
+ | [**pi-tokenburn**](https://github.com/sorinirimies/pi-tokenburn) | Live token and cost counter in pi's footer, with charts and budgets |
198
+
199
+ ```bash
200
+ pi install npm:pi-edit-first
201
+ pi install npm:pi-read-guard
202
+ pi install npm:pi-tokenburn
203
+ ```
204
+
121
205
  MIT
package/extensions/tokenburn.ts CHANGED
@@ -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,
@@ -115,6 +127,12 @@ export async function loadConfig(): Promise<TokenBurnConfig> {
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
+
118
136
  const BUDGET_KEYS = ["dayTokens", "weekTokens", "monthTokens", "dayCost", "weekCost", "monthCost"] as const;
119
137
 
120
138
  /** Keep only known budget keys holding positive, finite numbers (the file is hand-editable). */
@@ -158,44 +176,157 @@ interface Turn {
158
176
  cost: number;
159
177
  }
160
178
 
161
- /** Parsed turns per session file, reused while the file is unchanged (mtime + size). */
162
- 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 };
163
181
 
164
- async function parseFile(file: string): Promise<Turn[]> {
165
- 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
166
215
  try {
167
- const fd = await open(file, "r");
168
- try {
169
- for await (const line of fd.readLines({ encoding: "utf8" })) {
170
- if (!line.includes('"usage"')) continue;
171
- try {
172
- const obj = JSON.parse(line);
173
- const usage = obj.message?.usage;
174
- if (!obj.timestamp || !usage) continue;
175
- const ts = new Date(obj.timestamp).getTime();
176
- if (!Number.isFinite(ts)) continue;
177
- turns.push({
178
- ts,
179
- input: Number(usage.input) || 0,
180
- output: Number(usage.output) || 0,
181
- cacheRead: Number(usage.cacheRead) || 0,
182
- cacheWrite: Number(usage.cacheWrite) || 0,
183
- cost: Number(usage.cost?.total) || 0,
184
- });
185
- } catch {}
186
- }
187
- } finally {
188
- 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;
189
263
  }
190
- } catch {}
191
- 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 };
269
+ }
270
+
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);
192
276
  }
193
277
 
194
- /** Cache key for a file: mtime + size. `undefined` when it cannot be stat'ed. */
195
- export function fileSignature(file: string): string | undefined {
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 {
196
327
  try {
197
328
  const st = statSync(file);
198
- return `${st.mtimeMs}:${st.size}`;
329
+ return { size: st.size, mtimeMs: st.mtimeMs };
199
330
  } catch {
200
331
  return undefined;
201
332
  }
@@ -230,26 +361,32 @@ const emptyStats = (): PeriodStats => ({
230
361
  sessions: 0,
231
362
  });
232
363
 
233
- /** Refresh the per-file cache and return every session file's parsed turns. */
234
- 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[][]> {
235
366
  const files = listSessionFiles(sessionsRoot);
236
367
  const live = new Set(files);
237
368
  for (const key of FILE_CACHE.keys()) if (!live.has(key)) FILE_CACHE.delete(key);
238
369
 
239
370
  const out: Turn[][] = [];
240
371
  for (const file of files) {
241
- const sig = fileSignature(file);
242
- if (sig === undefined) continue; // vanished between listing and stat
372
+ const st = fileStat(file);
373
+ if (st === undefined) continue; // vanished between listing and stat
243
374
  let entry = FILE_CACHE.get(file);
244
- if (!entry || entry.sig !== sig) {
245
- entry = { sig, turns: await parseFile(file) };
246
- 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;
247
380
  }
248
- out.push(entry.turns);
381
+ if (entry) out.push(entry.tail.length > 0 ? entry.turns.concat(entry.tail) : entry.turns);
249
382
  }
250
383
  return out;
251
384
  }
252
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
+
253
390
  const turnTokens = (t: Turn, includeCache: boolean) =>
254
391
  t.input + t.output + (includeCache ? t.cacheRead + t.cacheWrite : 0);
255
392
 
@@ -518,6 +655,80 @@ export function isOverBudget(p: PeriodStats, b: { tokens?: number; cost?: number
518
655
  return (b.tokens !== undefined && p.tokens > b.tokens) || (b.cost !== undefined && p.cost > b.cost);
519
656
  }
520
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
+
521
732
  /** Map chart roles to pi theme colors. */
522
733
  export function themePaint(theme: { fg: (c: any, t: string) => string }): Paint {
523
734
  const color: Record<string, string> = {
@@ -537,6 +748,12 @@ export default function tokenBurnExtension(pi: ExtensionAPI) {
537
748
 
538
749
  const updateUI = async (ctx: ExtensionContext) => {
539
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
+ }
540
757
  const report = await collectUsage(getSessionsDir(), new Date(), config.includeCache);
541
758
  const theme = ctx.ui.theme;
542
759
 
@@ -582,15 +799,43 @@ export default function tokenBurnExtension(pi: ExtensionAPI) {
582
799
  }
583
800
  };
584
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
+
585
812
  pi.on("session_start", async (_event, ctx) => {
813
+ latestCtx = ctx;
586
814
  config = await loadConfig();
587
815
  // A failing refresh must never break session start.
588
816
  await updateUI(ctx).catch(() => {});
589
817
  });
590
818
 
591
- // Fire-and-forget: never block the agent loop on a stats refresh.
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;
827
+ });
828
+
829
+ // Final, exact refresh when a turn ends (fire-and-forget).
592
830
  pi.on("turn_end", (_event, ctx) => {
593
- 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;
594
839
  });
595
840
 
596
841
  pi.registerCommand("tokenburn", {
@@ -598,49 +843,77 @@ export default function tokenBurnExtension(pi: ExtensionAPI) {
598
843
  getArgumentCompletions: (prefix: string) => completions(prefix),
599
844
  handler: async (args, ctx) => {
600
845
  const cmd = parseCommand(args);
846
+ latestCtx = ctx;
601
847
  const persist = async () => {
602
848
  await saveConfig(config);
603
849
  await updateUI(ctx);
604
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");
605
853
 
606
854
  switch (cmd.kind) {
607
855
  case "window":
608
856
  config.showWidget = !config.showWidget;
609
857
  await persist();
610
- ctx.ui.notify(`tokenburn widget: ${config.showWidget ? "enabled" : "hidden"}`, "info");
858
+ note(`tokenburn widget: ${config.showWidget ? "enabled" : "hidden"}`);
611
859
  return;
612
860
 
613
861
  case "status":
614
862
  config.statusWindow = cmd.window;
615
863
  await persist();
616
- ctx.ui.notify(`tokenburn status window: ${cmd.window}`, "info");
864
+ note(`tokenburn status window: ${cmd.window}`);
617
865
  return;
618
866
 
619
867
  case "cycle": {
620
868
  const next = STATUS_WINDOWS[(STATUS_WINDOWS.indexOf(config.statusWindow) + 1) % STATUS_WINDOWS.length];
621
869
  config.statusWindow = next;
622
870
  await persist();
623
- ctx.ui.notify(`tokenburn status window: ${next}`, "info");
871
+ note(`tokenburn status window: ${next}`);
624
872
  return;
625
873
  }
626
874
 
627
875
  case "cache":
628
876
  config.includeCache = !config.includeCache;
629
877
  await persist();
630
- ctx.ui.notify(`tokenburn: cache tokens ${config.includeCache ? "included" : "excluded"}`, "info");
878
+ note(`tokenburn: cache tokens ${config.includeCache ? "included" : "excluded"}`);
631
879
  return;
632
880
 
633
881
  case "budget": {
634
882
  const key = `${cmd.period}${cmd.isCost ? "Cost" : "Tokens"}` as keyof TokenBurnConfig["budgets"];
635
883
  config.budgets[key] = cmd.amount;
636
884
  await persist();
637
- ctx.ui.notify(
638
- `tokenburn ${cmd.period} budget: ${cmd.isCost ? formatCost(cmd.amount) : formatNum(cmd.amount) + " tokens"}`,
639
- "info",
640
- );
885
+ note(`tokenburn ${cmd.period} budget: ${cmd.isCost ? formatCost(cmd.amount) : formatNum(cmd.amount) + " tokens"}`);
641
886
  return;
642
887
  }
643
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
+
644
917
  case "invalid":
645
918
  ctx.ui.notify(cmd.message, "warning");
646
919
  return;
@@ -680,7 +953,7 @@ export default function tokenBurnExtension(pi: ExtensionAPI) {
680
953
  rule,
681
954
  ...chart,
682
955
  rule,
683
- "/tokenburn chart <day|week|month|year|all> ยท /tokenburn <day|week|month|year|all> sets the status bar ยท cycle ยท window ยท budget ยท cache",
956
+ "chart <view> ยท <view> sets the status bar ยท cycle ยท window ยท budget ยท cache ยท live ยท on|off",
684
957
  ].join("\n"),
685
958
  "info",
686
959
  );
@@ -705,11 +978,14 @@ export type Command =
705
978
  | { kind: "cycle" }
706
979
  | { kind: "window" }
707
980
  | { kind: "cache" }
981
+ | { kind: "enable" }
982
+ | { kind: "disable" }
983
+ | { kind: "live" }
708
984
  | { kind: "budget"; period: (typeof BUDGET_PERIODS)[number]; amount: number; isCost: boolean }
709
985
  | { kind: "invalid"; message: string };
710
986
 
711
987
  export const USAGE =
712
- "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";
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";
713
989
 
714
990
  /** Words that pick the status-bar period: `day` shows today, `all` shows the all-time total. */
715
991
  const WORD_TO_WINDOW: Record<string, Window> = {
@@ -752,6 +1028,14 @@ export function parseCommand(args: string | undefined): Command {
752
1028
  return { kind: "window" };
753
1029
  case "cache":
754
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" };
755
1039
  case "budget": {
756
1040
  const isCost = (p2 ?? "").startsWith("$");
757
1041
  const amount = Number((p2 ?? "").replace(/^\$/, ""));
@@ -792,6 +1076,9 @@ const TOP_LEVEL: CompletionItem[] = [
792
1076
  { value: "window", label: "window", description: "Toggle the panel below the editor" },
793
1077
  { value: "budget ", label: "budget", description: "Set a day/week/month budget" },
794
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" },
795
1082
  ];
796
1083
 
797
1084
  /** `prefix` is the whole argument text typed so far (pi replaces it with `value`). */
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "pi-tokenburn",
3
- "version": "0.4.0",
3
+ "version": "0.5.1",
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",