@monotykamary/pi-tps 1.3.2 → 1.3.4

package/.github/workflows/test.yml

@@ -13,17 +13,18 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
+      - uses: pnpm/action-setup@v4
       - name: Setup Node.js
         uses: actions/setup-node@v4
         with:
           node-version: '22'
-          cache: 'npm'
+          cache: 'pnpm'
 
       - name: Install dependencies
-        run: npm ci
+        run: pnpm install --frozen-lockfile
 
       - name: Type check
-        run: npm run typecheck
+        run: pnpm typecheck
 
   test:
     runs-on: ubuntu-latest
@@ -32,20 +33,21 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
+      - uses: pnpm/action-setup@v4
       - name: Setup Node.js
         uses: actions/setup-node@v4
         with:
           node-version: '22'
-          cache: 'npm'
+          cache: 'pnpm'
 
       - name: Install dependencies
-        run: npm ci
+        run: pnpm install --frozen-lockfile
 
       - name: Run tests
-        run: npm test
+        run: pnpm test
 
       - name: Run tests with coverage
-        run: npm run test:coverage
+        run: pnpm test:coverage
 
       - name: Upload coverage
         uses: actions/upload-artifact@v4

package/.pi/fabric/mesh/actors/actors.json

@@ -0,0 +1,4 @@
+{
+  "format": 1,
+  "actors": []
+}

package/extensions/pi-tps/__tests__/cost-rate.test.ts

@@ -52,7 +52,8 @@ function makeMessageWithCost(opts: {
 async function runBurstTurn(
   fixture: ReturnType<typeof createTestFixture>,
   message: AssistantMessage,
-  turnIndex = 0
+  turnIndex = 0,
+  beforeTurnEnd?: () => void
 ) {
   const { handlers, mockCtx } = fixture;
   handlers['turn_start']?.({ type: 'turn_start', turnIndex, timestamp: Date.now() });
@@ -65,6 +66,7 @@ async function runBurstTurn(
     assistantMessageEvent: { type: 'text_delta', delta: 't' },
   });
   handlers['message_end']?.({ type: 'message_end', message });
+  beforeTurnEnd?.();
   handlers['turn_end']?.({ type: 'turn_end', turnIndex, message, toolResults: [] }, mockCtx);
 }
 
@@ -108,15 +110,15 @@ describe('pi-tps extension — blended $/M-tokens rate', () => {
     });
 
     // Simulate the neuralwatt provider's turn_end running BEFORE ours: it emits
-    // the per-turn energy event, which our listener stashes keyed by turnIndex.
-    fixture.emitEvent('neuralwatt:turn-energy', {
-      costUsd: 0.006,
-      energyJoules: 21.6,
-      turnIndex: 0,
+    // the per-turn energy event after streaming but before pi-tps handles turn_end.
+    await runBurstTurn(fixture, message, 0, () => {
+      fixture.emitEvent('neuralwatt:turn-energy', {
+        costUsd: 0.006,
+        energyJoules: 21.6,
+        turnIndex: 0,
+      });
     });
 
-    await runBurstTurn(fixture, message, 0);
-
     const { notifySpy, appendEntrySpy } = fixture;
     expect(notifySpy).toHaveBeenCalledOnce();
     const banner = notifySpy.mock.calls[0][0] as string;
@@ -157,6 +159,102 @@ describe('pi-tps extension — blended $/M-tokens rate', () => {
     expect(data.cost).toBeNull();
   });
 
+  it('omits the rate when pi reports an all-zero cost block for an unpriced model', async () => {
+    const message = makeMessageWithCost({
+      input: 500,
+      output: 500,
+      costTotal: 0,
+      provider: 'makora',
+      model: 'zai-org/GLM-5.2-NVFP4',
+    });
+
+    await runBurstTurn(fixture, message);
+
+    const { appendEntrySpy, notifySpy } = fixture;
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.cost).toBeNull();
+    expect(data.rateUsdPerMTokens).toBeNull();
+    expect(notifySpy.mock.calls[0][0]).not.toMatch(/\$.*\/M/);
+  });
+
+  it('does not reuse a late billed cost after /tree when a turn index repeats', async () => {
+    const neuralwattMessage = makeMessageWithCost({
+      input: 1178,
+      output: 1235,
+      costTotal: 0.0102352,
+      provider: 'neuralwatt',
+      model: 'glm-5.2-short',
+    });
+    neuralwattMessage.usage.cacheRead = 8192;
+    neuralwattMessage.usage.totalTokens = 10605;
+
+    await runBurstTurn(fixture, neuralwattMessage, 2);
+    fixture.emitEvent('neuralwatt:turn-energy', {
+      costUsd: 0.00116,
+      energyJoules: 835.2,
+      turnIndex: 2,
+    });
+
+    fixture.handlers['session_tree']?.(
+      { type: 'session_tree', newLeafId: null, oldLeafId: 'old-leaf' },
+      fixture.mockCtx
+    );
+
+    const unpricedMessage = makeMessageWithCost({
+      input: 8460,
+      output: 717,
+      costTotal: 0,
+      provider: 'makora',
+      model: 'zai-org/GLM-5.2-NVFP4',
+    });
+    unpricedMessage.usage.cacheRead = 1216;
+    unpricedMessage.usage.totalTokens = 10393;
+
+    await runBurstTurn(fixture, unpricedMessage, 2);
+
+    const latestTelemetry = fixture.appendEntrySpy.mock.calls.at(-1)![1];
+    expect(latestTelemetry.model).toEqual({
+      provider: 'makora',
+      modelId: 'zai-org/GLM-5.2-NVFP4',
+    });
+    expect(latestTelemetry.cost).toBeNull();
+    expect(latestTelemetry.rateUsdPerMTokens).toBeNull();
+
+    const latestBanner = fixture.notifySpy.mock.calls.at(-1)![0] as string;
+    expect(latestBanner).not.toContain('$0.11/M');
+    expect(latestBanner).not.toMatch(/\$.*\/M/);
+  });
+
+  it('does not apply an early current-turn event to a previous run with the same index', async () => {
+    const previousMessage = makeMessageWithCost({
+      input: 1000,
+      output: 1000,
+      costTotal: 0.008,
+    });
+    await runBurstTurn(fixture, previousMessage, 0);
+
+    const currentMessage = makeMessageWithCost({
+      input: 1000,
+      output: 1000,
+      costTotal: 0.01,
+      provider: 'neuralwatt',
+      model: 'glm-5.2',
+    });
+
+    await runBurstTurn(fixture, currentMessage, 0, () => {
+      fixture.emitEvent('neuralwatt:turn-energy', {
+        costUsd: 0.006,
+        energyJoules: 21.6,
+        turnIndex: 0,
+      });
+      expect(fixture.appendEntrySpy).toHaveBeenCalledTimes(1);
+    });
+
+    expect(fixture.appendEntrySpy).toHaveBeenCalledTimes(2);
+    const latestTelemetry = fixture.appendEntrySpy.mock.calls.at(-1)![1];
+    expect(latestTelemetry.rateUsdPerMTokens).toBe(3);
+  });
+
   it('falls back to list-price rate when billed-cost event misses (out-of-order load)', async () => {
     // Neuralwatt turn but the energy event never arrives (provider loaded after
     // us). Must not block or crash — falls back to the list-price compute rate.
@@ -224,12 +322,13 @@ describe('pi-tps extension — blended $/M-tokens rate', () => {
       model: 'moonshotai/Kimi-K2.5',
     });
 
-    fixture.emitEvent('neuralwatt:turn-energy', {
-      costUsd: 0.006,
-      energyJoules: 21.6,
-      turnIndex: 0,
+    await runBurstTurn(fixture, message, 0, () => {
+      fixture.emitEvent('neuralwatt:turn-energy', {
+        costUsd: 0.006,
+        energyJoules: 21.6,
+        turnIndex: 0,
+      });
     });
-    await runBurstTurn(fixture, message, 0);
 
     // A stray duplicate event must not append a second corrected entry.
     fixture.emitEvent('neuralwatt:turn-energy', {

package/extensions/pi-tps/__tests__/helpers.ts

@@ -108,7 +108,7 @@ export function createTestFixture(): TestFixture {
     ui: { notify: notifySpy } as any,
     sessionManager: {
       getEntries: vi.fn().mockReturnValue(mockEntries),
-      getBranch: vi.fn(),
+      getBranch: vi.fn().mockReturnValue(mockEntries),
       getSessionId: vi.fn(),
     },
     modelRegistry: undefined as any,

package/extensions/pi-tps/__tests__/rehydration.test.ts

@@ -175,6 +175,22 @@ describe('pi-tps extension — rehydration', () => {
     expect(notifySpy).toHaveBeenCalledOnce();
   });
 
+  it('should restore from the active branch instead of an abandoned branch', async () => {
+    const { handlers, notifySpy, mockEntries, mockCtx } = fixture;
+    const activeEntry = makeTpsEntry({ tps: 10.0 });
+    const abandonedEntry = makeTpsEntry({ tps: 99.0 });
+    mockEntries.push(activeEntry, abandonedEntry);
+    (mockCtx.sessionManager.getBranch as ReturnType<typeof vi.fn>).mockReturnValue([activeEntry]);
+
+    handlers['session_tree']?.({ newLeafId: 'active', oldLeafId: 'abandoned' }, mockCtx);
+    await tick();
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const msg = notifySpy.mock.calls[0][0] as string;
+    expect(msg).toContain('TPS 10.0');
+    expect(msg).not.toContain('TPS 99.0');
+  });
+
   it('should rehydrate most recent entry, preferring structured over legacy', async () => {
     const { handlers, notifySpy, mockEntries } = fixture;
 

package/extensions/pi-tps/index.ts

@@ -74,8 +74,6 @@ const STALL_THRESHOLD_MS = 500;
 
 /** Event name emitted by pi-neuralwatt-provider per turn with energy-billed cost data. */
 const NEURALWATT_ENERGY_EVENT = 'neuralwatt:turn-energy';
-/** Cap the per-turn billed-cost cache (turnIndex → costUsd) to avoid unbounded growth in long sessions. */
-const NEURALWATT_ENERGY_CACHE_MAX = 32;
 
 // ─── Data types ─────────────────────────────────────────────────────────────
 
@@ -115,6 +113,7 @@ interface TurnTelemetry {
 
 /** In-memory state accumulated during one LLM turn */
 interface TurnTiming {
+  turnIndex: number;
   turnStartMs: number;
   lastUpdateMs: number;
   firstTokenMs: number | null;
@@ -489,7 +488,11 @@ function buildTelemetry(
   // billed cost when present (energy-based, what the user actually pays),
   // otherwise the list-price compute cost from message.usage.cost. Never both
   // — billedCost wins outright when present, so there's no double-counting.
-  const effectiveCost = billedCost ?? (hasCost ? costTotal : null);
+  // Pi represents models without pricing with an all-zero cost block. Treat
+  // that as unavailable rather than displaying a misleading $0.00/M rate. An
+  // explicit billed cost remains valid even when it is zero.
+  const listPriceCost = hasCost && Number.isFinite(costTotal) && costTotal > 0 ? costTotal : null;
+  const effectiveCost = billedCost ?? listPriceCost;
   const rateUsdPerMTokens = computeRateUsdPerM(effectiveCost, totalTokens);
 
   return {
@@ -506,15 +509,16 @@ function buildTelemetry(
     },
     tps,
     isPrimaryBranch,
-    cost: hasCost
-      ? {
-          input: costInput,
-          output: costOutput,
-          cacheRead: costCacheRead,
-          cacheWrite: costCacheWrite,
-          total: costTotal,
-        }
-      : null,
+    cost:
+      listPriceCost !== null
+        ? {
+            input: costInput,
+            output: costOutput,
+            cacheRead: costCacheRead,
+            cacheWrite: costCacheWrite,
+            total: costTotal,
+          }
+        : null,
     rateUsdPerMTokens,
     timestamp: Date.now(),
   };
@@ -533,24 +537,12 @@ export default function tpsExtension(pi: ExtensionAPI) {
   // Cached session entries for argument completion (captured on session_start / session_tree)
   let cachedEntries: Array<{ type?: string; customType?: string; data?: unknown }> = [];
 
-  // ── Neuralwatt per-turn billed-cost integration ─────────────────────────
-  // pi-neuralwatt-provider emits NEURALWATT_ENERGY_EVENT in its turn_end handler
-  // (after its tee reader drains) with { costUsd, energyJoules, turnIndex }. We stash
-  // only the one number we need — billed costUsd — keyed by turnIndex, and read it
-  // synchronously in our own turn_end. pi awaits turn_end handlers sequentially in
-  // registration order, so if the neuralwatt provider was registered before us the
-  // event has already fired and the cache hit is immediate; if after us we miss it
-  // for that one turn and fall back to the list-price compute rate (the neuralwatt
-  // widget still shows billing separately). Zero latency, no awaiting, no duplicate
-  // capture — the raw energy/cost records stay solely in the neuralwatt provider's
-  // own session entries. Non-Neuralwatt turns never emit, so this stays empty.
-  const neuralwattBilledCostByTurn = new Map<number, number>();
-
-  // Tracks the most recently committed turn so a late-arriving
-  // `neuralwatt:turn-energy` event (provider loaded AFTER us) can retroactively
-  // correct a turn that we already persisted with the list-price fallback.
-  // Load-order-independent: when the provider loads before us the live cache
-  // hits and `billedApplied` is already true, so this never fires.
+  // pi-neuralwatt-provider emits billed cost during its turn_end handler. Keep
+  // an early event only for the active turn; a late event corrects only the
+  // turn that was just committed. turnIndex resets to zero on every agent run,
+  // so billed costs must never survive beyond either of those scopes.
+  let pendingNeuralwattBilledCost: { turnIndex: number; costUsd: number } | null = null;
+
   let lastCommittedTurn: {
     turnIndex: number;
     telemetry: TurnTelemetry;
@@ -563,43 +555,34 @@ export default function tpsExtension(pi: ExtensionAPI) {
     const p = payload as Record<string, unknown>;
     const turnIndex = typeof p.turnIndex === 'number' ? p.turnIndex : null;
     const costUsd = typeof p.costUsd === 'number' ? p.costUsd : null;
-    if (turnIndex === null || costUsd === null) return; // require turnIndex correlation
-    neuralwattBilledCostByTurn.set(turnIndex, costUsd);
-
-    // Late-arrival correction: the event reached us after our turn_end already
-    // committed this turn on the list-price fallback. Recompute the blended
-    // rate from the billed cost, persist a corrected `tps` entry (which later
-    // resume/export reads in place of the stale list-price one), and refresh
-    // the banner — but only if no newer turn is mid-flight (currentTiming set
-    // means a new turn is streaming; updating the banner then would be
-    // misleading). Idempotent: once `billedApplied` flips true we never
-    // re-correct, so repeated or duplicate events are safe.
-    const committed = lastCommittedTurn;
-    if (
-      committed &&
-      !committed.billedApplied &&
-      committed.turnIndex === turnIndex &&
-      !Number.isNaN(costUsd)
-    ) {
-      const totalTokens = committed.telemetry.tokens.total;
-      const correctedRate = computeRateUsdPerM(costUsd, totalTokens);
-      if (correctedRate !== null && correctedRate !== committed.telemetry.rateUsdPerMTokens) {
-        const corrected = { ...committed.telemetry, rateUsdPerMTokens: correctedRate };
-        committed.telemetry = corrected;
-        committed.billedApplied = true;
-        pi.appendEntry('tps', corrected);
-        pi.events?.emit('tps:telemetry', corrected);
-        cachedEntries.push({ type: 'custom', customType: 'tps' });
-        if (committed.ctx.hasUI && !currentTiming) {
-          committed.ctx.ui.notify(composeDisplayString(corrected), 'info');
-        }
+    if (turnIndex === null || costUsd === null || !Number.isFinite(costUsd) || costUsd < 0) {
+      return;
+    }
+
+    // Provider loaded before pi-tps: turn_end has not reached us yet.
+    if (currentTiming) {
+      if (currentTiming.turnIndex === turnIndex) {
+        pendingNeuralwattBilledCost = { turnIndex, costUsd };
       }
+      return;
     }
-    // Bound the cache (oldest = lowest turnIndex)
-    if (neuralwattBilledCostByTurn.size > NEURALWATT_ENERGY_CACHE_MAX) {
-      let oldest = Infinity;
-      for (const k of neuralwattBilledCostByTurn.keys()) if (k < oldest) oldest = k;
-      if (oldest !== Infinity) neuralwattBilledCostByTurn.delete(oldest);
+
+    // Provider loaded after pi-tps: correct the turn we just committed. Do not
+    // cache this event; the same turnIndex can belong to the next agent run.
+    const committed = lastCommittedTurn;
+    if (!committed || committed.billedApplied || committed.turnIndex !== turnIndex) return;
+
+    committed.billedApplied = true;
+    const correctedRate = computeRateUsdPerM(costUsd, committed.telemetry.tokens.total);
+    if (correctedRate === null || correctedRate === committed.telemetry.rateUsdPerMTokens) return;
+
+    const corrected = { ...committed.telemetry, rateUsdPerMTokens: correctedRate };
+    committed.telemetry = corrected;
+    pi.appendEntry('tps', corrected);
+    pi.events?.emit('tps:telemetry', corrected);
+    cachedEntries.push({ type: 'custom', customType: 'tps' });
+    if (committed.ctx.hasUI) {
+      committed.ctx.ui.notify(composeDisplayString(corrected), 'info');
     }
   });
 
@@ -613,7 +596,7 @@ export default function tpsExtension(pi: ExtensionAPI) {
    */
   function restoreTPSNotification(ctx: ExtensionContext) {
     if (!ctx.hasUI) return;
-    const entries = cachedEntries.length > 0 ? cachedEntries : ctx.sessionManager.getEntries();
+    const entries = ctx.sessionManager.getBranch();
     for (let i = entries.length - 1; i >= 0; i--) {
       const entry = entries[i];
       if (entry.type === 'custom' && entry.customType === 'tps') {
@@ -646,6 +629,8 @@ export default function tpsExtension(pi: ExtensionAPI) {
 
   // Restore notification after /tree navigation (same session, different branch)
   pi.on('session_tree', (_event: SessionTreeEvent, ctx: ExtensionContext) => {
+    pendingNeuralwattBilledCost = null;
+    lastCommittedTurn = null;
     cachedEntries = ctx.sessionManager.getEntries();
     restoreTPSNotification(ctx);
   });
@@ -654,7 +639,10 @@ export default function tpsExtension(pi: ExtensionAPI) {
 
   // Track when a turn starts (request sent to LLM)
   pi.on('turn_start', (_event: TurnStartEvent) => {
+    pendingNeuralwattBilledCost = null;
+    lastCommittedTurn = null;
     currentTiming = {
+      turnIndex: _event.turnIndex,
       turnStartMs: performance.now(),
       turnStartTimestamp: typeof _event.timestamp === 'number' ? _event.timestamp : Date.now(),
       lastUpdateMs: performance.now(),
@@ -764,9 +752,8 @@ export default function tpsExtension(pi: ExtensionAPI) {
   // ── Persist telemetry ───────────────────────────────────────────────────
 
   // Calculate, display, and persist telemetry at the end of each LLM turn.
-  // Synchronous: the Neuralwatt billed cost (if any) was stashed by our
-  // NEURALWATT_ENERGY_EVENT listener, keyed by turnIndex, and is read here
-  // synchronously. No awaiting, no added latency to turn dispatch.
+  // The Neuralwatt listener stores an early billed cost only while this exact
+  // turn is active, so the handoff remains synchronous and adds no latency.
   pi.on('turn_end', (event: TurnEndEvent, ctx: ExtensionContext) => {
     if (!currentTiming) return;
 
@@ -782,8 +769,11 @@ export default function tpsExtension(pi: ExtensionAPI) {
     // energy source is available do we fall back to the compute rate.
     // Prefer the live energy event, fall back to the persisted energy entry,
     // and finally to the list-price compute cost.
-    let billedCost = neuralwattBilledCostByTurn.get(event.turnIndex) ?? null;
-    neuralwattBilledCostByTurn.delete(event.turnIndex);
+    let billedCost =
+      pendingNeuralwattBilledCost?.turnIndex === event.turnIndex
+        ? pendingNeuralwattBilledCost.costUsd
+        : null;
+    pendingNeuralwattBilledCost = null;
 
     if (billedCost === null && timing.turnStartTimestamp) {
       billedCost = findEnergyCostFromSession(ctx, timing.turnStartTimestamp);

package/npm-shrinkwrap.json

@@ -1,12 +1,12 @@
 {
   "name": "@monotykamary/pi-tps",
-  "version": "1.3.2",
+  "version": "1.3.4",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "@monotykamary/pi-tps",
-      "version": "1.3.2",
+      "version": "1.3.4",
       "hasInstallScript": true,
       "license": "MIT",
       "devDependencies": {

package/package.json

@@ -1,19 +1,11 @@
 {
   "name": "@monotykamary/pi-tps",
-  "version": "1.3.2",
+  "version": "1.3.4",
   "description": "Tokens-per-second tracker for pi — see your LLM generation speed after every agent turn",
   "keywords": [
     "pi-package"
   ],
   "license": "MIT",
-  "scripts": {
-    "test": "vitest run",
-    "test:watch": "vitest",
-    "test:coverage": "vitest run --coverage",
-    "typecheck": "tsc --noEmit",
-    "lint:dead": "knip --no-gitignore",
-    "postinstall": "simple-git-hooks 2>/dev/null || true"
-  },
   "pi": {
     "extensions": [
       "./extensions"
@@ -32,8 +24,8 @@
   "devDependencies": {
     "@commitlint/cli": "21.0.1",
     "@commitlint/config-conventional": "21.0.1",
-    "@earendil-works/pi-ai": "0.79.8",
-    "@earendil-works/pi-coding-agent": "0.79.8",
+    "@earendil-works/pi-ai": "0.83.0",
+    "@earendil-works/pi-coding-agent": "0.83.0",
     "@types/node": "25.9.1",
     "@vitest/coverage-v8": "4.1.7",
     "knip": "6.14.1",
@@ -52,5 +44,13 @@
     "fast-xml-builder": "1.2.0",
     "protobufjs": "8.4.0",
     "ws": "8.20.1"
+  },
+  "scripts": {
+    "test": "vitest run",
+    "test:watch": "vitest",
+    "test:coverage": "vitest run --coverage",
+    "typecheck": "tsc --noEmit",
+    "lint:dead": "knip --no-gitignore",
+    "postinstall": "simple-git-hooks 2>/dev/null || true"
   }
-}
+}

package/pnpm-workspace.yaml

@@ -0,0 +1,9 @@
+allowBuilds:
+  '@google/genai': false
+  protobufjs: false
+  simple-git-hooks: false
+minimumReleaseAgeExclude:
+  - '@earendil-works/pi-agent-core@0.83.0'
+  - '@earendil-works/pi-ai@0.83.0'
+  - '@earendil-works/pi-coding-agent@0.83.0'
+  - '@earendil-works/pi-tui@0.83.0'