pi-tokenburn 0.4.0 → 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 +18 -0
- package/extensions/tokenburn.ts +338 -51
- package/package.json +1 -1
package/README.md
CHANGED
|
@@ -13,6 +13,7 @@ Totals match the `tokenburn` CLI exactly.
|
|
|
13
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
|
|
@@ -36,6 +37,9 @@ pi install git:github.com/sorinirimies/pi-tokenburn
|
|
|
36
37
|
| `/tokenburn budget <day\|week\|month> <tokens>` | Token budget, e.g. `budget day 5000000` |
|
|
37
38
|
| `/tokenburn budget <day\|week\|month> $<cost>` | Dollar budget, e.g. `budget week $50` |
|
|
38
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 |
|
|
39
43
|
|
|
40
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>`.
|
|
41
45
|
|
|
@@ -63,6 +67,9 @@ Weeks start on Monday, in local time.
|
|
|
63
67
|
|
|
64
68
|
```json
|
|
65
69
|
{
|
|
70
|
+
"enabled": true,
|
|
71
|
+
"liveRefresh": true,
|
|
72
|
+
"refreshMs": 750,
|
|
66
73
|
"includeCache": true,
|
|
67
74
|
"chartColor": true,
|
|
68
75
|
"showStatus": true,
|
|
@@ -74,10 +81,21 @@ Weeks start on Monday, in local time.
|
|
|
74
81
|
|
|
75
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.
|
|
76
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.
|
|
77
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.
|
|
78
87
|
|
|
79
88
|
Environment overrides for the session directory: `TOKENBURN_PI_SESSIONS`, `PI_CODING_AGENT_SESSION_DIR`, `PI_CODING_AGENT_DIR`.
|
|
80
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
|
+
|
|
81
99
|
## Security notes
|
|
82
100
|
|
|
83
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**.
|
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
|
-
/**
|
|
162
|
-
const
|
|
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
|
-
|
|
165
|
-
|
|
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
|
|
168
|
-
|
|
169
|
-
|
|
170
|
-
|
|
171
|
-
|
|
172
|
-
|
|
173
|
-
|
|
174
|
-
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
187
|
-
|
|
188
|
-
|
|
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
|
-
|
|
191
|
-
|
|
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
|
-
/**
|
|
195
|
-
|
|
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
|
|
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
|
-
/**
|
|
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
|
|
242
|
-
if (
|
|
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.
|
|
245
|
-
|
|
246
|
-
FILE_CACHE.set(file,
|
|
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
|
-
//
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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;
|
|
@@ -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.
|
|
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",
|