@monotykamary/pi-tps 1.3.0 → 1.3.1

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

@@ -249,4 +249,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>;
@@ -579,6 +614,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 +734,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;