@wierdbytes/pi-facelift 0.6.4 → 0.6.6

package/CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## 0.6.6
+
+- **Fix:** prompt-cache warm-ups no longer inflate `worked` time / deflate `tps` or start a phantom working timer while idle.
+
+## 0.6.5
+
+- **Fix:** pi crashing with "extension ctx is stale" when the session was replaced or reloaded while the working timer was ticking.
+
 ## 0.6.4
 
 - **Fix:** declare `shiki` as a runtime dependency instead of a dev-only one.

package/README.md

@@ -49,7 +49,8 @@ tool output.
   `⏱ worked 45.2s (total: 1m15s) · tps: 756`
 
   - `worked` — model-only time (request dispatch → response end, so
-    time-to-first-token counts; tool execution is excluded).
+    time-to-first-token counts; tool execution and pi's prompt-cache
+    warm-up requests are excluded).
   - `total` — wall-clock time from your prompt to the final response,
     including retries, tool calls, and all other overhead.
   - `tps` — output tokens per second over the `worked` time. The live

package/index.ts

@@ -1180,11 +1180,24 @@ export default function piFaceliftExtension(pi: PiFaceliftApi, deps?: PiFacelift
 	};
 
 	const renderWorkingMessage = (ctx: SessionContextLike): void => {
-		if (ctx.hasUI) {
-			ctx.ui.setWorkingMessage(workingMessageText(workingTracker.elapsedMs(), workingTracker.liveTps()));
+		// The ticker captures `ctx`; after session replacement / reload any
+		// access to it throws (stale ctx). Inside a timer that would be an
+		// uncaughtException and kill pi, so swallow and stop ticking.
+		try {
+			if (ctx.hasUI) {
+				ctx.ui.setWorkingMessage(workingMessageText(workingTracker.elapsedMs(), workingTracker.liveTps()));
+			}
+		} catch {
+			stopWorkingTick();
 		}
 	};
 
+	// Session replaced / reloaded / exiting mid-stream: message_end and
+	// agent_settled may never arrive, so kill the ticker here.
+	pi.on("session_shutdown", () => {
+		stopWorkingTick();
+	});
+
 	// `before_agent_start` fires right after the user submits the prompt,
 	// so it anchors the wall-clock `total` (retries + tools + overhead).
 	pi.on("before_agent_start", () => {
@@ -1198,11 +1211,22 @@ export default function piFaceliftExtension(pi: PiFaceliftApi, deps?: PiFacelift
 		workingTracker.ensureRun();
 	});
 
+	// pi emits `turn_start` (awaited) right before every agent request, so it
+	// arms the tracker. Cache-warmer replays also fire `before_provider_request`
+	// but arrive outside this window (during tools or after settle) and never
+	// get a `message_end` — unguarded they would hold a segment open forever.
+	pi.on("turn_start", () => {
+		if (!faceliftConfig.showWorkingTime) return;
+		workingTracker.armRequest();
+	});
+
 	pi.on("before_provider_request", (_event, ctx) => {
 		if (!faceliftConfig.showWorkingTime) return;
 		// Start the segment when the request is dispatched (not on the first
-		// token), so time-to-first-token counts as working time too.
-		workingTracker.beginSegment();
+		// token), so time-to-first-token counts as working time too. A `false`
+		// means a cache warm-up (not armed) or a repeat payload hook for an
+		// already-open request — leave the ticker alone.
+		if (!workingTracker.beginSegment()) return;
 		stopWorkingTick();
 		renderWorkingMessage(ctx);
 		workingTick = setInterval(() => renderWorkingMessage(ctx), 1000);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@wierdbytes/pi-facelift",
-  "version": "0.6.4",
+  "version": "0.6.6",
   "description": "Cosmetic facelift for built-in pi tool output — open-right rounded frames, syntax-highlighted reads, status-aware bash summaries, find/grep.",
   "type": "module",
   "files": [

package/working-time.test.ts

@@ -4,9 +4,13 @@
  * The tracker measures two metrics: `worked` (model streaming only,
  * tools excluded) and `totalMs` (wall-clock from user prompt to settle,
  * retries/tools/overhead included). The wiring drives it from pi
- * lifecycle events: before_agent_start / agent_start /
+ * lifecycle events: before_agent_start / agent_start / turn_start /
  * before_provider_request / message_end / agent_settled. The durable
  * history line is produced via pi.appendEntry + registerEntryRenderer.
+ *
+ * `before_provider_request` also fires for pi's prompt-cache warm-ups,
+ * which never produce a `message_end`; only requests armed by `turn_start`
+ * may open a segment.
  */
 
 import { beforeEach, describe, expect, it, vi } from "vitest";
@@ -27,6 +31,7 @@ describe("WorkingTimeTracker", () => {
 		const t = new WorkingTimeTracker(() => now);
 		t.beginRun();
 		now = 2000;
+		t.armRequest();
 		t.beginSegment();
 		now = 4500;
 		t.endSegment();
@@ -40,11 +45,13 @@ describe("WorkingTimeTracker", () => {
 		let now = 0;
 		const t = new WorkingTimeTracker(() => now);
 		t.beginRun();
+		t.armRequest();
 		t.beginSegment();
 		now = 1000;
 		t.endSegment();
 		// tool + overhead from 1000..9000 counts toward total, not worked
 		now = 9000;
+		t.armRequest();
 		t.beginSegment();
 		now = 10000;
 		t.endSegment();
@@ -57,6 +64,7 @@ describe("WorkingTimeTracker", () => {
 		let now = 0;
 		const t = new WorkingTimeTracker(() => now);
 		t.beginRun();
+		t.armRequest();
 		t.beginSegment();
 		now = 500;
 		t.endSegment();
@@ -71,6 +79,7 @@ describe("WorkingTimeTracker", () => {
 		let now = 0;
 		const t = new WorkingTimeTracker(() => now);
 		t.beginRun();
+		t.armRequest();
 		t.beginSegment();
 		t.addDelta(4000); // ~1000 estimated tokens at 4 chars/token
 		now = 2000; // 2s of worked time
@@ -81,6 +90,7 @@ describe("WorkingTimeTracker", () => {
 	it("liveTps is 0 before any time elapses", () => {
 		const t = new WorkingTimeTracker(() => 0);
 		t.beginRun();
+		t.armRequest();
 		t.beginSegment();
 		t.addDelta(400);
 		expect(t.liveTps()).toBe(0);
@@ -90,6 +100,7 @@ describe("WorkingTimeTracker", () => {
 		let now = 0;
 		const t = new WorkingTimeTracker(() => now);
 		t.beginRun();
+		t.armRequest();
 		t.beginSegment();
 		now = 2000;
 		t.recordExactTokens(1500);
@@ -104,12 +115,14 @@ describe("WorkingTimeTracker", () => {
 		const t = new WorkingTimeTracker(() => now);
 		t.beginRun();
 		// first model segment: 0..1000
+		t.armRequest();
 		t.beginSegment();
 		now = 1000;
 		t.endSegment();
 		// "tool" runs from 1000..9000 — must NOT count
 		now = 9000;
 		// second model segment: 9000..10000
+		t.armRequest();
 		t.beginSegment();
 		now = 10000;
 		t.endSegment();
@@ -120,6 +133,7 @@ describe("WorkingTimeTracker", () => {
 		let now = 0;
 		const t = new WorkingTimeTracker(() => now);
 		t.beginRun();
+		t.armRequest();
 		t.beginSegment();
 		now = 3000;
 		expect(t.elapsedMs()).toBe(3000);
@@ -132,10 +146,63 @@ describe("WorkingTimeTracker", () => {
 		expect(t.settle()).toBeUndefined();
 	});
 
+	it("ignores an unarmed request (cache warm-up while idle)", () => {
+		let now = 0;
+		const t = new WorkingTimeTracker(() => now);
+		expect(t.beginSegment()).toBe(false);
+		now = 5000;
+		expect(t.isActive()).toBe(false);
+		expect(t.elapsedMs()).toBe(0);
+	});
+
+	it("ignores a cache warm-up during tool execution", () => {
+		let now = 0;
+		const t = new WorkingTimeTracker(() => now);
+		t.beginRun();
+		t.armRequest();
+		expect(t.beginSegment()).toBe(true);
+		now = 1000;
+		t.endSegment();
+		// tool runs 1000..9000; the warmer replays the request at 4000
+		now = 4000;
+		expect(t.beginSegment()).toBe(false);
+		now = 9000;
+		t.armRequest();
+		expect(t.beginSegment()).toBe(true);
+		now = 10000;
+		t.endSegment();
+		expect(t.settle()?.ms).toBe(2000);
+	});
+
+	it("keeps the first start when the request hook fires again mid-segment", () => {
+		let now = 0;
+		const t = new WorkingTimeTracker(() => now);
+		t.beginRun();
+		t.armRequest();
+		expect(t.beginSegment()).toBe(true);
+		now = 1000;
+		expect(t.beginSegment()).toBe(false); // e.g. a provider re-sending the payload
+		now = 3000;
+		t.endSegment();
+		expect(t.settle()?.ms).toBe(3000);
+	});
+
+	it("disarms on settle and on beginRun", () => {
+		const t = new WorkingTimeTracker(() => 0);
+		t.beginRun();
+		t.armRequest();
+		t.settle();
+		expect(t.beginSegment()).toBe(false);
+		t.armRequest();
+		t.beginRun();
+		expect(t.beginSegment()).toBe(false);
+	});
+
 	it("resets state after settle", () => {
 		let now = 0;
 		const t = new WorkingTimeTracker(() => now);
 		t.beginRun();
+		t.armRequest();
 		t.beginSegment();
 		now = 1000;
 		t.endSegment();
@@ -264,11 +331,13 @@ describe("working-time wiring", () => {
 		fire("before_agent_start", {}); // anchors the total clock
 		vi.advanceTimersByTime(500); // pre-loop overhead
 		fire("agent_start", {});
+		fire("turn_start", {});
 		fire("before_provider_request", {});
 		vi.advanceTimersByTime(5000);
 		fire("message_end", asst);
 		// tool gap — counts toward total, NOT worked
 		vi.advanceTimersByTime(3000);
+		fire("turn_start", {});
 		fire("before_provider_request", {});
 		vi.advanceTimersByTime(2000);
 		fire("message_end", asst);
@@ -287,6 +356,7 @@ describe("working-time wiring", () => {
 		const { fire, entries, ctx } = loadExtension();
 		fire("before_agent_start", {});
 		fire("agent_start", {});
+		fire("turn_start", {});
 		fire("before_provider_request", {});
 		// Stream ~4000 chars => ~1000 estimated tokens over 2s of worked time.
 		fire("message_update", {
@@ -311,6 +381,7 @@ describe("working-time wiring", () => {
 		const { fire, entries } = loadExtension();
 		fire("before_agent_start", {});
 		fire("agent_start", {});
+		fire("turn_start", {});
 		fire("before_provider_request", {});
 		vi.advanceTimersByTime(1000);
 		fire("message_end", user); // stray non-assistant end — ignored
@@ -327,6 +398,7 @@ describe("working-time wiring", () => {
 		const { fire, ctx } = loadExtension();
 		fire("before_agent_start", {});
 		fire("agent_start", {});
+		fire("turn_start", {});
 		fire("before_provider_request", {});
 		vi.advanceTimersByTime(2000);
 		const calls = ctx.ui.setWorkingMessage.mock.calls.map((c) => c[0]);
@@ -340,6 +412,64 @@ describe("working-time wiring", () => {
 		expect(last?.[0]).toBeUndefined();
 	});
 
+	it("ignores a cache warm-up during tool execution (no segment, no ticker)", () => {
+		vi.useFakeTimers();
+		const { fire, entries, ctx } = loadExtension();
+		fire("before_agent_start", {});
+		fire("agent_start", {});
+		fire("turn_start", {});
+		fire("before_provider_request", {});
+		vi.advanceTimersByTime(1000);
+		fire("message_end", asst);
+		// tool executes; the cache warmer replays the request mid-tool
+		vi.advanceTimersByTime(2000);
+		const callsBefore = ctx.ui.setWorkingMessage.mock.calls.length;
+		fire("before_provider_request", {});
+		vi.advanceTimersByTime(5000);
+		expect(ctx.ui.setWorkingMessage.mock.calls.length).toBe(callsBefore);
+		fire("turn_start", {});
+		fire("before_provider_request", {});
+		vi.advanceTimersByTime(1000);
+		fire("message_end", asst);
+		fire("agent_settled", {});
+		vi.useRealTimers();
+		expect(entries).toHaveLength(1);
+		const data = entries[0].data as WorkingTimeEntry;
+		expect(data.ms).toBe(2000); // the 7s tool gap stays excluded
+		expect(data.totalMs).toBe(9000);
+	});
+
+	it("ignores a cache warm-up after settle (idle warming): no phantom run or ticker", () => {
+		vi.useFakeTimers();
+		const { fire, entries, ctx } = loadExtension();
+		fire("before_agent_start", {});
+		fire("agent_start", {});
+		fire("turn_start", {});
+		fire("before_provider_request", {});
+		vi.advanceTimersByTime(1000);
+		fire("message_end", asst);
+		fire("agent_settled", {});
+		ctx.ui.setWorkingMessage.mockClear();
+		// idle warm-up minutes later
+		vi.advanceTimersByTime(270_000);
+		fire("before_provider_request", {});
+		vi.advanceTimersByTime(5000);
+		expect(ctx.ui.setWorkingMessage).not.toHaveBeenCalled();
+		// the next real run is measured from its own prompt
+		fire("before_agent_start", {});
+		fire("agent_start", {});
+		fire("turn_start", {});
+		fire("before_provider_request", {});
+		vi.advanceTimersByTime(1500);
+		fire("message_end", asst);
+		fire("agent_settled", {});
+		vi.useRealTimers();
+		expect(entries).toHaveLength(2);
+		const data = entries[1].data as WorkingTimeEntry;
+		expect(data.ms).toBe(1500);
+		expect(data.totalMs).toBe(1500);
+	});
+
 	it("registers an entry renderer that produces a dim worked+total line", () => {
 		const { entryRenderers } = loadExtension();
 		const renderer = entryRenderers.get(WORKING_TIME_ENTRY);

package/working-time.ts

@@ -61,6 +61,7 @@ export interface WorkingTimeEntry {
  * Lifecycle (driven by pi events):
  *   before_agent_start      → beginRun()   (total clock starts here)
  *   agent_start             → ensureRun()  (fallback if before_agent_start skipped)
+ *   turn_start              → armRequest() (an agent request is coming)
  *   before_provider_request → beginSegment() + start the live ticker
  *   message_update  (asst)  → addDelta()     (feed the live tps estimate)
  *   message_end     (asst)  → recordExactTokens() + endSegment() + stop ticker
@@ -68,6 +69,15 @@ export interface WorkingTimeEntry {
  *
  * Retries / auto-compaction fire their own agent_start but pi only emits
  * `agent_settled` when it truly stops, so a run is measured end-to-end.
+ *
+ * `before_provider_request` is not exclusive to agent turns: pi's prompt
+ * cache warmer replays the last request (`maxTokens: 1`) with the same
+ * options, so the hook fires for warm-ups too — mid-run while tools execute,
+ * and after the run settles in `idle` warming mode. Those requests never
+ * produce a `message_end`, so an unguarded segment would stay open and
+ * swallow tool time (or start a phantom run while idle). Segments therefore
+ * open only while armed: from `turn_start` (pi emits one before every agent
+ * request) until the assistant `message_end`.
  */
 export class WorkingTimeTracker {
 	private readonly now: () => number;
@@ -76,6 +86,7 @@ export class WorkingTimeTracker {
 	private runStart: number | undefined;
 	private runEstChars = 0;
 	private runExactTokens = 0;
+	private requestArmed = false;
 
 	constructor(now: () => number = () => Date.now()) {
 		this.now = now;
@@ -88,6 +99,7 @@ export class WorkingTimeTracker {
 		this.runStart = this.now();
 		this.runEstChars = 0;
 		this.runExactTokens = 0;
+		this.requestArmed = false;
 	}
 
 	/** Start a run only if one isn't already in progress (idempotent). */
@@ -95,14 +107,29 @@ export class WorkingTimeTracker {
 		if (this.runStart === undefined) this.beginRun();
 	}
 
-	/** Begin timing an assistant streaming segment. */
-	beginSegment(): void {
+	/**
+	 * Mark that the agent is about to issue a real request (pi `turn_start`).
+	 * Only an armed tracker lets `beginSegment` open a segment.
+	 */
+	armRequest(): void {
 		this.ensureRun();
-		if (this.segmentStart === undefined) this.segmentStart = this.now();
+		this.requestArmed = true;
+	}
+
+	/**
+	 * Begin timing an assistant streaming segment. Ignored unless armed (see
+	 * `armRequest`) or when a segment is already open. Returns `true` only
+	 * when a new segment was opened.
+	 */
+	beginSegment(): boolean {
+		if (!this.requestArmed || this.segmentStart !== undefined) return false;
+		this.segmentStart = this.now();
+		return true;
 	}
 
-	/** Close the current segment, folding its duration into the total. */
+	/** Close the current segment, folding its duration into the total, and disarm. */
 	endSegment(): void {
+		this.requestArmed = false;
 		if (this.segmentStart !== undefined) {
 			this.accumulatorMs += Math.max(0, this.now() - this.segmentStart);
 			this.segmentStart = undefined;