@monotykamary/pi-tps 1.3.1 → 1.3.3

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.
@@ -175,6 +273,77 @@ describe('pi-tps extension — blended $/M-tokens rate', () => {
     expect(banner).toContain('$4.00/M');
   });
 
+  it('corrects the banner + persisted entry when the energy event arrives late', async () => {
+    // Provider loads AFTER pi-tps: our turn_end runs first on the list-price
+    // fallback, then the provider's turn_end emits the energy event, which
+    // must retroactively correct the rate without a second turn_end.
+    const message = makeMessageWithCost({
+      input: 1000,
+      output: 1000,
+      costTotal: 0.008, // list price: $4.00/M; billed: $3.00/M
+      provider: 'neuralwatt',
+      model: 'moonshotai/Kimi-K2.5',
+    });
+
+    await runBurstTurn(fixture, message, 0);
+
+    const { notifySpy, appendEntrySpy } = fixture;
+    // Initially committed on the list-price fallback.
+    expect(appendEntrySpy.mock.calls[0][0]).toBe('tps');
+    expect(appendEntrySpy.mock.calls[0][1].rateUsdPerMTokens).toBe(4.0);
+    expect(notifySpy.mock.calls[0][0]).toContain('$4.00/M');
+
+    // Provider's turn_end fires the energy event after ours.
+    fixture.emitEvent('neuralwatt:turn-energy', {
+      costUsd: 0.006, // $3.00/M for 2000 tokens
+      energyJoules: 21.6,
+      turnIndex: 0,
+    });
+
+    // A corrected `tps` entry is appended in place of the list-price one.
+    expect(appendEntrySpy.mock.calls).toHaveLength(2);
+    expect(appendEntrySpy.mock.calls[1][0]).toBe('tps');
+    expect(appendEntrySpy.mock.calls[1][1].rateUsdPerMTokens).toBe(3.0);
+
+    // Banner is refreshed to the billed rate.
+    const lastBanner = notifySpy.mock.calls.at(-1)![0] as string;
+    expect(lastBanner).toContain('$3.00/M');
+    expect(lastBanner).not.toContain('$4.00/M');
+  });
+
+  it('does not double-correct a turn already billed at commit time', async () => {
+    // Provider loads BEFORE pi-tps: the event fires first, our turn_end reads
+    // it from the live cache (billedApplied=true), so no late correction.
+    const message = makeMessageWithCost({
+      input: 1000,
+      output: 1000,
+      costTotal: 0.008,
+      provider: 'neuralwatt',
+      model: 'moonshotai/Kimi-K2.5',
+    });
+
+    await runBurstTurn(fixture, message, 0, () => {
+      fixture.emitEvent('neuralwatt:turn-energy', {
+        costUsd: 0.006,
+        energyJoules: 21.6,
+        turnIndex: 0,
+      });
+    });
+
+    // A stray duplicate event must not append a second corrected entry.
+    fixture.emitEvent('neuralwatt:turn-energy', {
+      costUsd: 0.006,
+      energyJoules: 21.6,
+      turnIndex: 0,
+    });
+
+    const { appendEntrySpy, notifySpy } = fixture;
+    expect(appendEntrySpy.mock.calls.filter((c) => c[0] === 'tps')).toHaveLength(1);
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const banner = notifySpy.mock.calls[0][0] as string;
+    expect(banner).toContain('$3.00/M');
+  });
+
   it('ignores neuralwatt:turn-energy payloads lacking a numeric turnIndex', async () => {
     // Defensive: malformed event must not pollute the cache.
     fixture.emitEvent('neuralwatt:turn-energy', { costUsd: 0.006, energyJoules: 21.6 }); // no turnIndex

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,31 +537,52 @@ 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>();
+  // 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;
+    billedApplied: boolean;
+    ctx: ExtensionContext;
+  } | null = null;
 
   pi.events?.on(NEURALWATT_ENERGY_EVENT, (payload: unknown) => {
     if (!payload || typeof payload !== 'object') return;
     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);
-    // 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);
+    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;
+    }
+
+    // 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');
     }
   });
 
@@ -571,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') {
@@ -604,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);
   });
@@ -612,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(),
@@ -722,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;
 
@@ -740,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);
@@ -772,6 +804,17 @@ export default function tpsExtension(pi: ExtensionAPI) {
       }
     }
 
+    // Record this turn so a late-arriving neuralwatt:turn-energy event
+    // (provider loaded after us) can retroactively correct the rate. Billed
+    // cost wins when present, so a turn already billed at commit time is
+    // never re-corrected.
+    lastCommittedTurn = {
+      turnIndex: event.turnIndex,
+      telemetry,
+      billedApplied: billedCost !== null,
+      ctx,
+    };
+
     // Persist structured telemetry to session for export and rehydration
     pi.appendEntry('tps', telemetry);
 

package/npm-shrinkwrap.json

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

package/package.json

@@ -1,19 +1,11 @@
 {
   "name": "@monotykamary/pi-tps",
-  "version": "1.3.1",
+  "version": "1.3.3",
   "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"
@@ -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,4 @@
+allowBuilds:
+  '@google/genai': false
+  protobufjs: false
+  simple-git-hooks: false