@monotykamary/pi-tps 1.3.0 → 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
@@ -249,4 +319,86 @@ describe('pi-tps extension — blended $/M-tokens rate', () => {
     const banner = notifySpy.mock.calls[0][0] as string;
     expect(banner).not.toMatch(/\$.*\/M/);
   });
+
+  it('falls back to a persisted neuralwatt-energy entry when the live event is missed', async () => {
+    const now = Date.now();
+    const message = makeMessageWithCost({
+      input: 1000,
+      output: 1000,
+      costTotal: 0.008, // list price: $4.00/M
+      provider: 'neuralwatt',
+      model: 'moonshotai/Kimi-K2.5',
+    });
+
+    // Provider appended its energy entry before pi-tps handled turn_end.
+    fixture.mockEntries.push({
+      type: 'custom',
+      customType: 'neuralwatt-energy',
+      data: { energy_joules: 21.6, cost_usd: 0.006 }, // $3.00/M for 2000 tokens
+      timestamp: now,
+    });
+
+    const { handlers, notifySpy, appendEntrySpy } = fixture;
+    handlers['turn_start']?.({ type: 'turn_start', turnIndex: 0, timestamp: now - 100 });
+    await tick(50);
+    handlers['message_start']?.({ type: 'message_start', message });
+    await tick(50);
+    handlers['message_update']?.({
+      type: 'message_update',
+      message,
+      assistantMessageEvent: { type: 'text_delta', delta: 'H' },
+    });
+    handlers['message_end']?.({ type: 'message_end', message });
+    handlers['turn_end']?.(
+      { type: 'turn_end', turnIndex: 0, message, toolResults: [] },
+      fixture.mockCtx
+    );
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const banner = notifySpy.mock.calls[0][0] as string;
+    expect(banner).toContain('$3.00/M');
+    expect(banner).not.toContain('$4.00/M');
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.rateUsdPerMTokens).toBe(3.0);
+  });
+
+  it('ignores stale neuralwatt-energy entries from before this turn', async () => {
+    const now = Date.now();
+    const message = makeMessageWithCost({
+      input: 1000,
+      output: 1000,
+      costTotal: 0.008, // list price: $4.00/M
+      provider: 'neuralwatt',
+      model: 'moonshotai/Kimi-K2.5',
+    });
+
+    fixture.mockEntries.push({
+      type: 'custom',
+      customType: 'neuralwatt-energy',
+      data: { energy_joules: 1, cost_usd: 9999 }, // absurd cost, should be ignored
+      timestamp: now - 1000,
+    });
+
+    const { handlers, notifySpy } = fixture;
+    handlers['turn_start']?.({ type: 'turn_start', turnIndex: 0, timestamp: now });
+    await tick(50);
+    handlers['message_start']?.({ type: 'message_start', message });
+    await tick(50);
+    handlers['message_update']?.({
+      type: 'message_update',
+      message,
+      assistantMessageEvent: { type: 'text_delta', delta: 'H' },
+    });
+    handlers['message_end']?.({ type: 'message_end', message });
+    handlers['turn_end']?.(
+      { type: 'turn_end', turnIndex: 0, message, toolResults: [] },
+      fixture.mockCtx
+    );
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const banner = notifySpy.mock.calls[0][0] as string;
+    expect(banner).toContain('$4.00/M');
+    expect(banner).not.toContain('$9999');
+  });
 });

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

@@ -140,7 +140,7 @@ export function createTestFixture(): TestFixture {
   };
 
   const mockPi: Partial<ExtensionAPI> = {
-    on: vi.fn((event: string, handler: (...args: unknown[]) => void) => {
+    on: vi.fn((event: string, handler: any) => {
       handlers[event] = handler;
       return mockPi as ExtensionAPI;
     }),

package/extensions/pi-tps/index.ts

@@ -132,6 +132,7 @@ interface TurnTiming {
   messageCount: number;
   isToolCall: boolean; // tool_execution_start fired during this turn
   isPrimaryBranch: boolean; // TPS came from primary-branch (reliable) measurement
+  turnStartTimestamp: number; // wall-clock ms at turn_start, for correlating session energy entries
 }
 
 // ─── Helpers ────────────────────────────────────────────────────────────────
@@ -150,6 +151,40 @@ function computeRateUsdPerM(costUsd: number | null, totalTokens: number): number
   return Math.round(rate * 100) / 100;
 }
 
+/**
+ * Fallback: find a Neuralwatt energy entry that was persisted in this session
+ * during the current turn. This covers the case where pi-neuralwatt-provider
+ * loads before pi-tps, so the energy entry exists before our turn_end runs,
+ * but the live event cache missed (e.g. malformed payload, event ordering, or
+ * a race).
+ *
+ * Scans backward from the most recent entry and requires the entry timestamp
+ * to be >= turnStartTimestamp so a previous turn's energy data is not reused.
+ */
+function findEnergyCostFromSession(
+  ctx: ExtensionContext,
+  turnStartTimestamp: number
+): number | null {
+  if (!ctx.sessionManager?.getEntries) return null;
+
+  const entries = ctx.sessionManager.getEntries();
+  for (let i = entries.length - 1; i >= 0; i--) {
+    const e = entries[i];
+    if (e.type !== 'custom' || e.customType !== 'neuralwatt-energy') continue;
+
+    const entryTimestamp = typeof e.timestamp === 'number' ? e.timestamp : 0;
+    if (entryTimestamp < turnStartTimestamp) return null;
+
+    const data = e.data as Record<string, unknown> | null | undefined;
+    if (!data) continue;
+
+    const costUsd = typeof data.cost_usd === 'number' ? data.cost_usd : null;
+    if (costUsd !== null && Number.isFinite(costUsd) && costUsd >= 0) return costUsd;
+  }
+
+  return null;
+}
+
 function isAssistantMessage(message: unknown): message is AssistantMessage {
   if (!message || typeof message !== 'object') return false;
   const msg = message as Record<string, unknown>;
@@ -511,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>;
@@ -518,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;
@@ -579,6 +656,7 @@ export default function tpsExtension(pi: ExtensionAPI) {
   pi.on('turn_start', (_event: TurnStartEvent) => {
     currentTiming = {
       turnStartMs: performance.now(),
+      turnStartTimestamp: typeof _event.timestamp === 'number' ? _event.timestamp : Date.now(),
       lastUpdateMs: performance.now(),
       firstTokenMs: null,
       currentMessageStartMs: null,
@@ -698,14 +776,19 @@ export default function tpsExtension(pi: ExtensionAPI) {
     const turnEndMs = performance.now();
 
     // Pick the effective cost for the blended $/M-tokens rate: Neuralwatt's
-    // billed cost when present (cache hit), otherwise the list-price compute
-    // cost from message.usage.cost. The neuralwatt provider emits its event
-    // in its own turn_end, which runs before or after ours depending on load
-    // order — a cache miss here just falls back to the compute rate for this
-    // one turn (no double-counting: only one source contributes to the rate).
-    const billedCost = neuralwattBilledCostByTurn.get(event.turnIndex) ?? null;
+    // billed cost when present, otherwise the list-price compute cost from
+    // message.usage.cost. The live event is the primary source; if it missed,
+    // we read the persisted neuralwatt-energy session entry. Only if neither
+    // 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);
 
+    if (billedCost === null && timing.turnStartTimestamp) {
+      billedCost = findEnergyCostFromSession(ctx, timing.turnStartTimestamp);
+    }
+
     const telemetry = buildTelemetry(timing, turnEndMs, billedCost);
     if (!telemetry) return;
 
@@ -731,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);