@monotykamary/pi-tps 1.2.0 → 1.3.0

package/README.md

@@ -20,7 +20,7 @@ _Originally from [badlogic/pi-mono](https://github.com/badlogic/pi-mono/blob/mai
 ## Quick start
 
 ```bash
-pi install https://github.com/monotykamary/pi-tps
+pi install npm:@monotykamary/pi-tps
 ```
 
 ## What's included
@@ -43,6 +43,12 @@ pi install https://github.com/monotykamary/pi-tps
 
 ## Install
 
+```bash
+pi install npm:@monotykamary/pi-tps
+```
+
+Or install from GitHub:
+
 ```bash
 pi install https://github.com/monotykamary/pi-tps
 ```
@@ -63,24 +69,35 @@ Then `/reload` in pi.
 ## Output format
 
 ```
-TPS 42.5 tok/s · TTFT 1.2s · 29.7s · in 567 · out 1.2K · stall 4.3s×1
+TPS 42.5 tok/s · TTFT 1.2s · 29.7s · in 567 · out 1.2K · stall 4.3s×1 · $3.60/M
 ```
 
-| Field   | Description                                                         |
-| ------- | ------------------------------------------------------------------- |
-| `TPS`   | Tokens per second (generation speed, excludes TTFT & stalls)        |
-| `TTFT`  | Time to first token (seconds, 1 decimal)                            |
-| `s`     | Total wall-clock time from request to completion                    |
-| `in`    | Input tokens (human-readable: K/M/B)                                |
-| `out`   | Output tokens (human-readable: K/M/B)                               |
-| `stall` | Accumulated stall time × stall count (shown only when stalls exist) |
+| Field   | Description                                                                              |
+| ------- | ---------------------------------------------------------------------------------------- |
+| `TPS`   | Tokens per second (generation speed, excludes TTFT & stalls)                             |
+| `TTFT`  | Time to first token (seconds, 1 decimal)                                                 |
+| `s`     | Total wall-clock time from request to completion                                         |
+| `in`    | Input tokens (human-readable: K/M/B)                                                     |
+| `out`   | Output tokens (human-readable: K/M/B)                                                    |
+| `stall` | Accumulated stall time × stall count (shown only when stalls exist)                      |
+| `$/M`   | Blended $/million-tokens rate for this turn (shown only when a cost source is available) |
 
 When TPS can't be determined (burst delivery, too few chunks), the field shows `—`:
 
 ```
-TPS — · TTFT 0.8s · 1.3s · in 291 · out 46
+TPS — · TTFT 0.8s · 1.3s · in 291 · out 46 · $2.59/M
 ```
 
+### $/M-tokens rate
+
+The trailing `$/M` segment is a blended per-turn rate: `effectiveCost / (totalTokens / 1_000_000)`, shown to two decimals. Only one cost source contributes per turn (no double-counting):
+
+- When [pi-neuralwatt-provider](https://github.com/monotykamary/pi-neuralwatt-provider) reports the **energy-billed cost** for the turn (what you're actually invoiced), that wins outright.
+- Otherwise the **list-price compute cost** from `message.usage.cost.total` (derived from the model's pricing table via pi-ai) is used.
+- When neither is available (no pricing, zero tokens), the segment is omitted.
+
+The cross-extension handoff is event-based: pi-neuralwatt-provider emits `neuralwatt:turn-energy` with `{ costUsd, energyJoules, turnIndex }` in its own `turn_end`; pi-tps subscribes at load and stashes `costUsd` keyed by `turnIndex`. pi dispatches `turn_end` handlers sequentially, so if the neuralwatt provider was registered before pi-tps the value is already cached when pi-tps's `turn_end` runs; if not, pi-tps falls back to the list-price rate for that one turn (the neuralwatt widget still shows billing separately). No latency is added — pi-tps reads the cache synchronously, never awaits.
+
 Human-readable scaling (for token counts):
 
 - `< 1K`: raw integer (`567`)
@@ -186,29 +203,30 @@ pi.events.on('tps:telemetry', (data) => {
 
 The event payload:
 
-| Field                 | Type             | Description                                         |
-| --------------------- | ---------------- | --------------------------------------------------- |
-| `tps`                 | `number \| null` | Tokens per second, or null when unidentifiable      |
-| `model.provider`      | `string`         | Provider name (e.g. `openai`)                       |
-| `model.modelId`       | `string`         | Model identifier (e.g. `gpt-4`)                     |
-| `tokens.input`        | `number`         | Input tokens (summed across all assistant messages) |
-| `tokens.output`       | `number`         | Output tokens generated by the LLM                  |
-| `tokens.cacheRead`    | `number`         | Cache-read tokens (provider-dependent)              |
-| `tokens.cacheWrite`   | `number`         | Cache-write tokens (provider-dependent)             |
-| `tokens.total`        | `number`         | Total tokens (input + output + cache)               |
-| `timing.ttftMs`       | `number \| null` | Time to first token in milliseconds                 |
-| `timing.totalMs`      | `number`         | Total wall-clock time from request to completion    |
-| `timing.generationMs` | `number`         | Streaming wall clock (message_start → message_end)  |
-| `timing.streamMs`     | `number \| null` | Inter-update span: first → last streaming update    |
-| `timing.stallMs`      | `number`         | Accumulated inference stall time in ms              |
-| `timing.stallCount`   | `number`         | Number of discrete stall events                     |
-| `timing.messageCount` | `number`         | Assistant messages in this turn                     |
-| `cost.input`          | `number \| null` | Input token cost                                    |
-| `cost.output`         | `number \| null` | Output token cost                                   |
-| `cost.cacheRead`      | `number \| null` | Cache-read token cost                               |
-| `cost.cacheWrite`     | `number \| null` | Cache-write token cost                              |
-| `cost.total`          | `number \| null` | Total cost for this turn                            |
-| `timestamp`           | `number`         | Unix timestamp (ms) when telemetry was computed     |
+| Field                 | Type             | Description                                                                                                                                                           |
+| --------------------- | ---------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `tps`                 | `number \| null` | Tokens per second, or null when unidentifiable                                                                                                                        |
+| `model.provider`      | `string`         | Provider name (e.g. `openai`)                                                                                                                                         |
+| `model.modelId`       | `string`         | Model identifier (e.g. `gpt-4`)                                                                                                                                       |
+| `tokens.input`        | `number`         | Input tokens (summed across all assistant messages)                                                                                                                   |
+| `tokens.output`       | `number`         | Output tokens generated by the LLM                                                                                                                                    |
+| `tokens.cacheRead`    | `number`         | Cache-read tokens (provider-dependent)                                                                                                                                |
+| `tokens.cacheWrite`   | `number`         | Cache-write tokens (provider-dependent)                                                                                                                               |
+| `tokens.total`        | `number`         | Total tokens (input + output + cache)                                                                                                                                 |
+| `timing.ttftMs`       | `number \| null` | Time to first token in milliseconds                                                                                                                                   |
+| `timing.totalMs`      | `number`         | Total wall-clock time from request to completion                                                                                                                      |
+| `timing.generationMs` | `number`         | Streaming wall clock (message_start → message_end)                                                                                                                    |
+| `timing.streamMs`     | `number \| null` | Inter-update span: first → last streaming update                                                                                                                      |
+| `timing.stallMs`      | `number`         | Accumulated inference stall time in ms                                                                                                                                |
+| `timing.stallCount`   | `number`         | Number of discrete stall events                                                                                                                                       |
+| `timing.messageCount` | `number`         | Assistant messages in this turn                                                                                                                                       |
+| `cost.input`          | `number \| null` | Input token cost                                                                                                                                                      |
+| `cost.output`         | `number \| null` | Output token cost                                                                                                                                                     |
+| `cost.cacheRead`      | `number \| null` | Cache-read token cost                                                                                                                                                 |
+| `cost.cacheWrite`     | `number \| null` | Cache-write token cost                                                                                                                                                |
+| `cost.total`          | `number \| null` | Total cost for this turn                                                                                                                                              |
+| `timestamp`           | `number`         | Unix timestamp (ms) when telemetry was computed                                                                                                                       |
+| `rateUsdPerMTokens`   | `number \| null` | Blended $/M-tokens rate for this turn (Neuralwatt billed cost when available, else list-price compute cost). `null` when no cost source is usable or tokens are zero. |
 
 When `cost` is unavailable (provider doesn't report it), the entire `cost` object is `null`.
 

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

@@ -0,0 +1,252 @@
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
+import type { AssistantMessage } from '@earendil-works/pi-ai';
+import { createTestFixture, activateExtension, tick, makeAssistantMessage } from './helpers';
+
+/**
+ * Coverage for the blended $/M-tokens notification field:
+ *   rateUsdPerMTokens = effectiveCost / (tokens.total / 1_000_000)
+ *
+ * effectiveCost is:
+ *   - the Neuralwatt billed cost when stashed via the `neuralwatt:turn-energy`
+ *     event (energy-based, what the user actually pays)
+ *   - otherwise the list-price compute cost from message.usage.cost.total
+ *
+ * Only one source contributes per turn (no double-counting): when both are
+ * present the billed cost wins outright.
+ */
+
+function makeMessageWithCost(opts: {
+  input: number;
+  output: number;
+  costTotal: number;
+  provider?: string;
+  model?: string;
+}): AssistantMessage {
+  const { input, output, costTotal, provider = 'openai', model = 'gpt-4' } = opts;
+  return {
+    role: 'assistant',
+    content: [{ type: 'text', text: 'Hello' }],
+    api: 'openai-completions',
+    provider,
+    model,
+    usage: {
+      input,
+      output,
+      cacheRead: 0,
+      cacheWrite: 0,
+      totalTokens: input + output,
+      cost: {
+        input: 0,
+        output: 0,
+        cacheRead: 0,
+        cacheWrite: 0,
+        total: costTotal,
+      },
+    },
+    stopReason: 'stop',
+    timestamp: Date.now(),
+  };
+}
+
+/** Drive a minimal turn that yields primary-branch null TPS (burst) but valid telemetry. */
+async function runBurstTurn(
+  fixture: ReturnType<typeof createTestFixture>,
+  message: AssistantMessage,
+  turnIndex = 0
+) {
+  const { handlers, mockCtx } = fixture;
+  handlers['turn_start']?.({ type: 'turn_start', turnIndex, timestamp: Date.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: 't' },
+  });
+  handlers['message_end']?.({ type: 'message_end', message });
+  handlers['turn_end']?.({ type: 'turn_end', turnIndex, message, toolResults: [] }, mockCtx);
+}
+
+describe('pi-tps extension — blended $/M-tokens rate', () => {
+  let fixture: ReturnType<typeof createTestFixture>;
+
+  beforeEach(async () => {
+    fixture = createTestFixture();
+    await activateExtension(fixture);
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it('shows list-price $/M rate in the banner from message.usage.cost.total', async () => {
+    // 1000 in + 1000 out = 2000 tokens; cost.total = $0.008 → $4.00/M
+    const message = makeMessageWithCost({ input: 1000, output: 1000, costTotal: 0.008 });
+
+    await runBurstTurn(fixture, message);
+
+    const { notifySpy, appendEntrySpy } = fixture;
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const banner = notifySpy.mock.calls[0][0] as string;
+    expect(banner).toContain('$4.00/M');
+    expect(banner).not.toMatch(/\$.*\/M.*\$.*\/M/); // exactly one rate segment
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.rateUsdPerMTokens).toBe(4.0);
+  });
+
+  it('uses Neuralwatt billed cost over list-price when the energy event fires first', async () => {
+    // Same token volume, but billed cost differs from list price.
+    // 2000 tokens; list cost.total = $0.008 ($4.00/M); billed = $0.006 ($3.00/M)
+    const message = makeMessageWithCost({
+      input: 1000,
+      output: 1000,
+      costTotal: 0.008,
+      provider: 'neuralwatt',
+      model: 'moonshotai/Kimi-K2.5',
+    });
+
+    // 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,
+    });
+
+    await runBurstTurn(fixture, message, 0);
+
+    const { notifySpy, appendEntrySpy } = fixture;
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const banner = notifySpy.mock.calls[0][0] as string;
+    // Billed ($3.00/M) wins over list-price ($4.00/M)
+    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);
+    // The cost block still carries the list-price compute cost separately.
+    expect(data.cost).toEqual(expect.objectContaining({ total: 0.008 }));
+  });
+
+  it('null rate when totalTokens is zero (degenerate)', async () => {
+    const message = makeMessageWithCost({ input: 0, output: 1000, costTotal: 0.008 });
+    message.usage.totalTokens = 0;
+
+    await runBurstTurn(fixture, message);
+
+    const { appendEntrySpy, notifySpy } = fixture;
+    const banner = notifySpy.mock.calls[0][0] as string;
+    expect(banner).not.toMatch(/\$.*\/M/);
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.rateUsdPerMTokens).toBeNull();
+  });
+
+  it('null rate when cost is unavailable', async () => {
+    // No usage.cost at all → effectiveCost null → rate null.
+    const message = makeAssistantMessage({ output: 500, input: 500 });
+    // Strip the cost block that makeAssistantMessage adds.
+    (message.usage as any).cost = null;
+    // Allow the isAssistantMessage guard to still pass: it only checks input/output.
+    await runBurstTurn(fixture, message);
+
+    const { appendEntrySpy } = fixture;
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.rateUsdPerMTokens).toBeNull();
+    expect(data.cost).toBeNull();
+  });
+
+  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.
+    const message = makeMessageWithCost({
+      input: 500,
+      output: 500,
+      costTotal: 0.004, // $4.00/M for 1000 tokens
+      provider: 'neuralwatt',
+      model: 'moonshotai/Kimi-K2.5',
+    });
+
+    await runBurstTurn(fixture, message);
+
+    const { notifySpy } = fixture;
+    const banner = notifySpy.mock.calls[0][0] as string;
+    expect(banner).toContain('$4.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
+    const message = makeMessageWithCost({ input: 1000, output: 1000, costTotal: 0.008 });
+
+    await runBurstTurn(fixture, message);
+
+    const { notifySpy } = fixture;
+    const banner = notifySpy.mock.calls[0][0] as string;
+    // Falls back to list-price since no valid turnIndex-keyed entry landed.
+    expect(banner).toContain('$4.00/M');
+  });
+
+  it('rehydrates a rate segment from structured telemetry on session resume', async () => {
+    const { handlers, notifySpy, mockEntries } = fixture;
+
+    mockEntries.push({
+      type: 'custom',
+      customType: 'tps',
+      data: {
+        model: { provider: 'openai', modelId: 'gpt-4' },
+        tokens: { input: 1000, output: 1000, cacheRead: 0, cacheWrite: 0, total: 2000 },
+        timing: {
+          ttftMs: 1000,
+          totalMs: 3000,
+          generationMs: 2000,
+          stallMs: 0,
+          stallCount: 0,
+          messageCount: 1,
+        },
+        tps: 10.0,
+        rateUsdPerMTokens: 4.0,
+        timestamp: Date.now(),
+      },
+    });
+
+    handlers['session_start']?.({ reason: 'resume' }, fixture.mockCtx);
+    await tick();
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const banner = notifySpy.mock.calls[0][0] as string;
+    expect(banner).toContain('$4.00/M');
+  });
+
+  it('rehydrates older telemetry without a rate field without crashing (omits the segment)', async () => {
+    const { handlers, notifySpy, mockEntries } = fixture;
+
+    mockEntries.push({
+      type: 'custom',
+      customType: 'tps',
+      data: {
+        model: { provider: 'openai', modelId: 'gpt-4' },
+        tokens: { input: 50, output: 100, cacheRead: 0, cacheWrite: 0, total: 150 },
+        timing: {
+          ttftMs: 1000,
+          totalMs: 3000,
+          generationMs: 2000,
+          stallMs: 0,
+          stallCount: 0,
+          messageCount: 1,
+        },
+        tps: 10.0,
+        // rateUsdPerMTokens intentionally absent (pre-feature entry)
+        timestamp: Date.now(),
+      },
+    });
+
+    handlers['session_start']?.({ reason: 'resume' }, fixture.mockCtx);
+    await tick();
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const banner = notifySpy.mock.calls[0][0] as string;
+    expect(banner).not.toMatch(/\$.*\/M/);
+  });
+});

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

@@ -77,6 +77,8 @@ export interface TestFixture {
   notifySpy: ReturnType<typeof vi.fn>;
   appendEntrySpy: ReturnType<typeof vi.fn>;
   eventsEmitSpy: ReturnType<typeof vi.fn>;
+  /** Emit an event on the mock event bus, dispatching to registered listeners (and the emit spy). */
+  emitEvent: (event: string, payload: unknown) => void;
   registerCommandSpy: ReturnType<typeof vi.fn>;
   mockEntries: Array<{ type?: string; role?: string; customType?: string; data?: unknown }>;
   mockCtx: ExtensionContext;
@@ -123,6 +125,19 @@ export function createTestFixture(): TestFixture {
   } as any as ExtensionContext;
 
   const eventsEmitSpy = vi.fn();
+  // Real event bus: listeners registered via events.on are dispatched on emit.
+  // eventsEmitSpy still captures every emit for assertions (e.g. 'tps:telemetry').
+  const eventListeners = new Map<string, ((payload: unknown) => void)[]>();
+  const dispatchEvent = (event: string, payload: unknown) => {
+    eventsEmitSpy(event, payload);
+    for (const listener of eventListeners.get(event) ?? []) {
+      try {
+        listener(payload);
+      } catch {
+        // listener errors shouldn't break emit dispatch in tests
+      }
+    }
+  };
 
   const mockPi: Partial<ExtensionAPI> = {
     on: vi.fn((event: string, handler: (...args: unknown[]) => void) => {
@@ -131,7 +146,14 @@ export function createTestFixture(): TestFixture {
     }),
     appendEntry: appendEntrySpy,
     registerCommand: registerCommandSpy,
-    events: { emit: eventsEmitSpy, on: vi.fn() },
+    events: {
+      emit: dispatchEvent,
+      on: vi.fn((event: string, listener: (payload: unknown) => void) => {
+        const list = eventListeners.get(event) ?? [];
+        list.push(listener);
+        eventListeners.set(event, list);
+      }),
+    } as any,
   };
 
   return {
@@ -141,6 +163,7 @@ export function createTestFixture(): TestFixture {
     notifySpy,
     appendEntrySpy,
     eventsEmitSpy,
+    emitEvent: dispatchEvent,
     registerCommandSpy,
     mockEntries,
     mockCtx,

package/extensions/pi-tps/index.ts

@@ -72,6 +72,11 @@ interface ToolExecutionStartEvent {
 /** Minimum gap between token updates to count as a stall (ms) */
 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 ─────────────────────────────────────────────────────────────
 
 /** Structured telemetry persisted per turn in the session JSONL */
@@ -96,6 +101,15 @@ interface TurnTelemetry {
     cacheWrite: number;
     total: number;
   } | null;
+  /**
+   * Blended $/M-tokens rate shown in the notification banner: the per-turn
+   * cost divided by (totalTokens / 1_000_000). Cost source is the Neuralwatt
+   * billed cost when available (energy-based), otherwise the list-price
+   * compute cost from `cost.total`. null when neither is usable or tokens
+   * are zero. Purely derived — recomputed in composeDisplayString, not a
+   * duplicate of the cost block.
+   */
+  rateUsdPerMTokens: number | null;
   timestamp: number;
 }
 
@@ -122,6 +136,20 @@ interface TurnTiming {
 
 // ─── Helpers ────────────────────────────────────────────────────────────────
 
+/**
+ * Compute the blended $/M-tokens rate: costUsd / (totalTokens / 1_000_000).
+ * The single calculation backing the banner's $/M-tokens field. Returns null
+ * when the inputs can't produce a meaningful rate (zero/negative tokens,
+ * non-finite or negative cost).
+ */
+function computeRateUsdPerM(costUsd: number | null, totalTokens: number): number | null {
+  if (costUsd === null || !Number.isFinite(costUsd) || costUsd < 0) return null;
+  if (!Number.isFinite(totalTokens) || totalTokens <= 0) return null;
+  const rate = costUsd / (totalTokens / 1_000_000);
+  if (!Number.isFinite(rate) || rate < 0) return null;
+  return Math.round(rate * 100) / 100;
+}
+
 function isAssistantMessage(message: unknown): message is AssistantMessage {
   if (!message || typeof message !== 'object') return false;
   const msg = message as Record<string, unknown>;
@@ -236,6 +264,9 @@ function composeDisplayString(t: TurnTelemetry): string {
     const stallStr = formatDuration(t.timing.stallMs / 1000);
     parts.push(`stall ${stallStr}×${t.timing.stallCount}`);
   }
+  if (t.rateUsdPerMTokens != null) {
+    parts.push(`$${t.rateUsdPerMTokens.toFixed(2)}/M`);
+  }
   return parts.join(' · ');
 }
 
@@ -243,7 +274,11 @@ function composeDisplayString(t: TurnTelemetry): string {
  * Build structured TurnTelemetry from accumulated turn timing.
  * Returns null if the turn had no meaningful LLM output.
  */
-function buildTelemetry(timing: TurnTiming, turnEndMs: number): TurnTelemetry | null {
+function buildTelemetry(
+  timing: TurnTiming,
+  turnEndMs: number,
+  billedCost: number | null = null
+): TurnTelemetry | null {
   let input = 0;
   let output = 0;
   let cacheRead = 0;
@@ -415,6 +450,13 @@ function buildTelemetry(timing: TurnTiming, turnEndMs: number): TurnTelemetry |
     isPrimaryBranch = false;
   }
 
+  // Single blended $/M-tokens rate for the banner. Cost source: Neuralwatt's
+  // 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);
+  const rateUsdPerMTokens = computeRateUsdPerM(effectiveCost, totalTokens);
+
   return {
     model,
     tokens: { input, output, cacheRead, cacheWrite, total: totalTokens },
@@ -438,6 +480,7 @@ function buildTelemetry(timing: TurnTiming, turnEndMs: number): TurnTelemetry |
           total: costTotal,
         }
       : null,
+    rateUsdPerMTokens,
     timestamp: Date.now(),
   };
 }
@@ -455,6 +498,34 @@ 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.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);
+    }
+  });
+
   // ── Rehydration ─────────────────────────────────────────────────────────
 
   /**
@@ -614,15 +685,28 @@ export default function tpsExtension(pi: ExtensionAPI) {
 
   // ── Persist telemetry ───────────────────────────────────────────────────
 
-  // Calculate, display, and persist telemetry at the end of each LLM turn
-  pi.on('turn_end', (_event: TurnEndEvent, ctx: ExtensionContext) => {
+  // 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.
+  pi.on('turn_end', (event: TurnEndEvent, ctx: ExtensionContext) => {
     if (!currentTiming) return;
 
     const timing = currentTiming;
     currentTiming = null;
 
     const turnEndMs = performance.now();
-    const telemetry = buildTelemetry(timing, turnEndMs);
+
+    // 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;
+    neuralwattBilledCostByTurn.delete(event.turnIndex);
+
+    const telemetry = buildTelemetry(timing, turnEndMs, billedCost);
     if (!telemetry) return;
 
     // ── Dynamic TPS cap ────────────────────────────────────────────────

package/package.json

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