@monotykamary/pi-tps 1.3.1 → 1.3.2

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

@@ -175,6 +175,76 @@ 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',
+    });
+
+    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', {
+      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/index.ts

@@ -546,6 +546,18 @@ export default function tpsExtension(pi: ExtensionAPI) {
   // 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.
+  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>;
@@ -553,6 +565,36 @@ export default function tpsExtension(pi: ExtensionAPI) {
     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');
+        }
+      }
+    }
     // Bound the cache (oldest = lowest turnIndex)
     if (neuralwattBilledCostByTurn.size > NEURALWATT_ENERGY_CACHE_MAX) {
       let oldest = Infinity;
@@ -772,6 +814,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.2",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "@monotykamary/pi-tps",
-      "version": "1.3.1",
+      "version": "1.3.2",
       "hasInstallScript": true,
       "license": "MIT",
       "devDependencies": {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@monotykamary/pi-tps",
-  "version": "1.3.1",
+  "version": "1.3.2",
   "description": "Tokens-per-second tracker for pi — see your LLM generation speed after every agent turn",
   "keywords": [
     "pi-package"