@monotykamary/pi-tps 1.1.1 → 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/__tests__/volume-gate.test.ts

@@ -0,0 +1,372 @@
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import type { AssistantMessage } from '@earendil-works/pi-ai';
+import { createTestFixture, activateExtension } from './helpers';
+
+describe('pi-tps extension — volume-based TPS gate', () => {
+  let fixture: ReturnType<typeof createTestFixture>;
+
+  beforeEach(async () => {
+    fixture = createTestFixture();
+    await activateExtension(fixture);
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  /**
+   * Drive a full turn with mocked performance.now() timestamps.
+   * Customizable output token count for volume-based testing.
+   *
+   * Timestamp mapping (one performance.now() per handler call):
+   *   [0] turn_start   → turnStartMs, lastUpdateMs
+   *   [1] turn_start   → (second call — same init)
+   *   [2] message_start → currentMessageStartMs, lastUpdateMs reset
+   *   [3] message_update (TTFT) → firstTokenMs, lastUpdateMs
+   *   [4..4+n-1] message_update (stream) → each streamUpdate timestamp
+   *   [4+n] message_end → generation time end
+   *   [5+n] turn_end → total time
+   *
+   * streamMs = last streamUpdate - first streamUpdate (not total window!)
+   * updateCount = streamUpdates.length (post-TTFT events)
+   */
+  function driveTurn(clocks: {
+    turnStart: number;
+    messageStart: number;
+    firstUpdate: number;
+    streamUpdates: number[];
+    messageEnd: number;
+    turnEnd?: number;
+    output: number;
+    input?: number;
+    isToolCall?: boolean;
+  }) {
+    const { handlers, notifySpy, appendEntrySpy } = fixture;
+
+    const timestamps = [
+      clocks.turnStart,
+      clocks.turnStart,
+      clocks.messageStart,
+      clocks.firstUpdate,
+      ...clocks.streamUpdates,
+      clocks.messageEnd,
+      clocks.turnEnd ?? clocks.messageEnd,
+    ];
+
+    let callIdx = 0;
+    const spy = vi.spyOn(performance, 'now').mockImplementation(() => {
+      return timestamps[Math.min(callIdx++, timestamps.length - 1)];
+    });
+
+    const input = clocks.input ?? 50;
+    const assistantMessage: AssistantMessage = {
+      role: 'assistant',
+      content: [{ type: 'text', text: 'Response' }],
+      api: 'openai-completions',
+      provider: 'openai',
+      model: 'gpt-4',
+      usage: {
+        input,
+        output: clocks.output,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: input + clocks.output,
+        cost: {
+          input: 0.001,
+          output: 0.002,
+          cacheRead: 0,
+          cacheWrite: 0,
+          total: 0.003,
+        },
+      },
+      stopReason: clocks.isToolCall ? 'toolUse' : 'stop',
+      timestamp: Date.now(),
+    };
+
+    handlers['turn_start']?.({ type: 'turn_start', turnIndex: 0, timestamp: Date.now() });
+    handlers['message_start']?.({ type: 'message_start', message: assistantMessage });
+    handlers['message_update']?.({
+      type: 'message_update',
+      message: assistantMessage,
+      assistantMessageEvent: { type: 'text_delta', delta: 't' },
+    });
+    for (const _ts of clocks.streamUpdates) {
+      handlers['message_update']?.({
+        type: 'message_update',
+        message: assistantMessage,
+        assistantMessageEvent: { type: 'text_delta', delta: 't' },
+      });
+    }
+
+    if (clocks.isToolCall) {
+      handlers['tool_execution_start']?.({
+        type: 'tool_execution_start',
+        toolCallId: 'call_123',
+        toolName: 'bash',
+        args: { command: 'ls' },
+      });
+    }
+
+    handlers['message_end']?.({ type: 'message_end', message: assistantMessage });
+    handlers['turn_end']?.(
+      { type: 'turn_end', turnIndex: 0, message: assistantMessage, toolResults: [] },
+      fixture.mockCtx
+    );
+
+    spy.mockRestore();
+    return { notifySpy, appendEntrySpy };
+  }
+
+  // ── Primary branch: volume gate ────────────────────────────────────────
+
+  it('should null TPS when primary-branch TPS exceeds plausibility ceiling', () => {
+    // streamUpdates: [400, 450, 500, 550, 600] → streamMs = 200ms
+    // 3000 tokens / 0.2s = 15,000 TPS — exceeds 10,000 ceiling
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 200,
+      firstUpdate: 200.123,
+      streamUpdates: [400, 450, 500, 550, 600],
+      messageEnd: 700,
+      output: 3000,
+    });
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.tps).toBeNull();
+    expect(data.isPrimaryBranch).toBe(false);
+  });
+
+  it('should preserve primary-branch TPS when it is within plausibility', () => {
+    // streamUpdates: [400, 500, 600, 700, 800] → streamMs = 400ms
+    // 1000 tokens / 0.4s = 2,500 TPS — well within 10,000 ceiling
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 200,
+      firstUpdate: 200.123,
+      streamUpdates: [400, 500, 600, 700, 800],
+      messageEnd: 900,
+      output: 1000,
+    });
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.tps).not.toBeNull();
+    expect(data.tps).toBeGreaterThanOrEqual(2000);
+    expect(data.tps).toBeLessThanOrEqual(3000);
+    expect(data.isPrimaryBranch).toBe(true);
+  });
+
+  it('should preserve primary-branch TPS for high volume with long enough window', () => {
+    // streamUpdates: [500, 600, 700, 800, 1000] → streamMs = 500ms
+    // 3000 tokens / 0.5s = 6,000 TPS — within ceiling
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 200,
+      firstUpdate: 200.123,
+      streamUpdates: [500, 600, 700, 800, 1000],
+      messageEnd: 1100,
+      output: 3000,
+    });
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.tps).not.toBeNull();
+    expect(data.tps).toBeGreaterThanOrEqual(5500);
+    expect(data.tps).toBeLessThanOrEqual(6500);
+    expect(data.isPrimaryBranch).toBe(true);
+  });
+
+  it('should null TPS for very large token volume over minimum effective span', () => {
+    // streamUpdates: [400, 450, 500, 550, 600] → streamMs = 200ms
+    // 5000 tokens / 0.2s = 25,000 TPS — well beyond ceiling
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 200,
+      firstUpdate: 200.123,
+      streamUpdates: [400, 450, 500, 550, 600],
+      messageEnd: 700,
+      output: 5000,
+    });
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.tps).toBeNull();
+    expect(data.isPrimaryBranch).toBe(false);
+  });
+
+  // ── Fallback branch: volume gate ───────────────────────────────────────
+
+  it('should null TPS when fallback-branch TPS exceeds plausibility ceiling', () => {
+    // 2 updates (fallback branch), generationMs = 200ms
+    // 3000 tokens / 0.2s = 15,000 TPS — exceeds ceiling
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 50,
+      firstUpdate: 50.1,
+      streamUpdates: [50.15, 50.3],
+      messageEnd: 250,
+      output: 3000,
+    });
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.tps).toBeNull();
+  });
+
+  it('should preserve fallback-branch TPS when it is within plausibility', () => {
+    // 2 updates (fallback branch), generationMs = 200ms
+    // 400 tokens / 0.2s = 2,000 TPS — within ceiling
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 50,
+      firstUpdate: 50.1,
+      streamUpdates: [50.15, 50.3],
+      messageEnd: 250,
+      output: 400,
+    });
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.tps).not.toBeNull();
+    expect(data.tps).toBeGreaterThanOrEqual(1500);
+    expect(data.tps).toBeLessThanOrEqual(2500);
+  });
+
+  // ── Boundary: exactly at ceiling ──────────────────────────────────────
+
+  it('should preserve TPS when exactly at the plausibility ceiling (not exceeded)', () => {
+    // streamUpdates: [400, 450, 500, 550, 600] → streamMs = 200ms
+    // 2000 tokens / 0.2s = 10,000 TPS — exactly at ceiling, not > ceiling
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 200,
+      firstUpdate: 200.123,
+      streamUpdates: [400, 450, 500, 550, 600],
+      messageEnd: 700,
+      output: 2000,
+    });
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.tps).not.toBeNull();
+    expect(data.tps).toBeGreaterThanOrEqual(9900);
+    expect(data.tps).toBeLessThanOrEqual(10100);
+    expect(data.isPrimaryBranch).toBe(true);
+  });
+
+  it('should null TPS just above the plausibility ceiling', () => {
+    // streamUpdates: [400, 450, 500, 550, 600] → streamMs = 200ms
+    // 2100 tokens / 0.2s = 10,500 TPS — just above threshold
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 200,
+      firstUpdate: 200.123,
+      streamUpdates: [400, 450, 500, 550, 600],
+      messageEnd: 700,
+      output: 2100,
+    });
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.tps).toBeNull();
+    expect(data.isPrimaryBranch).toBe(false);
+  });
+
+  // ── Notification display ──────────────────────────────────────────────
+
+  it('should show TPS dash when volume gate nulls TPS', () => {
+    const { notifySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 200,
+      firstUpdate: 200.123,
+      streamUpdates: [400, 450, 500, 550, 600],
+      messageEnd: 700,
+      output: 3000,
+    });
+
+    const notification = notifySpy.mock.calls[0][0] as string;
+    expect(notification).toContain('TPS —');
+  });
+
+  // ── Interaction with dynamic TPS cap ──────────────────────────────────
+
+  it('should not let volume-gated turns set the dynamic TPS cap', () => {
+    // Turn 1: 3000 tokens / 0.2s = 15,000 TPS — volume gates to null.
+    // The cap condition requires isPrimaryBranch && tps !== null,
+    // both of which are false after the volume gate. So the cap
+    // should NOT be set from this turn.
+    driveTurn({
+      turnStart: 0,
+      messageStart: 200,
+      firstUpdate: 200.123,
+      streamUpdates: [400, 450, 500, 550, 600],
+      messageEnd: 700,
+      output: 3000,
+    });
+
+    const [, data1] = fixture.appendEntrySpy.mock.calls[0];
+    expect(data1.tps).toBeNull();
+    expect(data1.isPrimaryBranch).toBe(false);
+
+    // Turn 2: reliable non-tool-call streaming at ~50 TPS (20 tokens / 0.4s)
+    // This should set the cap at ~50 TPS, not at 15,000.
+    // If the volume-gated turn had set the cap at 15,000, the cap
+    // would be 15,000 and a subsequent tool call would be allowed
+    // up to 15,000 — which is inflated.
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 200,
+      firstUpdate: 200.123,
+      streamUpdates: [400, 500, 600, 700, 800],
+      messageEnd: 900,
+      output: 20,
+    });
+
+    const [, data2] = appendEntrySpy.mock.calls[1];
+    expect(data2.tps).not.toBeNull();
+    expect(data2.tps).toBeGreaterThanOrEqual(40);
+    expect(data2.tps).toBeLessThanOrEqual(60);
+    expect(data2.isPrimaryBranch).toBe(true);
+  });
+
+  // ── Volume gate doesn't affect normal token counts ────────────────────
+
+  it('should not affect TPS for normal token counts even at short effective span', () => {
+    // streamUpdates: [400, 450, 500, 550, 600] → streamMs = 200ms
+    // 20 tokens / 0.2s = 100 TPS — well within ceiling
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 200,
+      firstUpdate: 200.123,
+      streamUpdates: [400, 450, 500, 550, 600],
+      messageEnd: 700,
+      output: 20,
+    });
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.tps).not.toBeNull();
+    expect(data.tps).toBeGreaterThanOrEqual(80);
+    expect(data.tps).toBeLessThanOrEqual(120);
+    expect(data.isPrimaryBranch).toBe(true);
+  });
+
+  // ── High-output `read`-style burst ────────────────────────────────────
+
+  it('should null TPS when provider dumps 1000+ tokens in a burst that passes timing gates', () => {
+    // Real-world scenario: a provider spits out 1500 tokens in a single
+    // fast burst. The timing gates pass (5+ updates with ≥1ms gaps,
+    // 200ms+ effective span), but the rate is inflated because the
+    // generation window is too short relative to the volume.
+    // streamUpdates: [400, 450, 500, 550, 600] → streamMs = 200ms
+    // 1500 tokens / 0.2s = 7,500 TPS — below 10,000 ceiling, passes gate
+    // This is actually plausible for a very fast provider, so it should
+    // NOT be nulled. Test that the gate is not overly aggressive.
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 200,
+      firstUpdate: 200.123,
+      streamUpdates: [400, 450, 500, 550, 600],
+      messageEnd: 700,
+      output: 1500,
+    });
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.tps).not.toBeNull();
+    // 1500 / 0.2 = 7,500 TPS — plausible, within ceiling
+    expect(data.tps).toBeGreaterThanOrEqual(7000);
+    expect(data.tps).toBeLessThanOrEqual(8000);
+  });
+});

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;
@@ -308,6 +343,19 @@ function buildTelemetry(timing: TurnTiming, turnEndMs: number): TurnTelemetry |
   const STALL_REDUCTION_DENOM = 2;
   const STALL_DOMINANCE_RATIO = 0.85;
 
+  // Maximum plausible generation speed (tokens/second). Beyond this, the
+  // measured TPS is almost certainly a measurement artifact — the effective
+  // generation window is too short relative to the token volume to
+  // distinguish genuine inference from a buffer-flush dispatch of
+  // pre-generated tokens. At 10_000 TPS this is 5× the fastest known
+  // commercial inference (Cerebras ~2_000 tok/s). The gate is phrased as
+  // "extraordinary claims require extraordinary evidence": for X output
+  // tokens, the minimum reliable measurement window is X / MAX_PLAUSIBLE_TPS
+  // seconds. Below that, the volume of tokens cannot be reliably timed, and
+  // TPS is set to null. This is mathematically equivalent to nulling out
+  // TPS when computed TPS > MAX_PLAUSIBLE_TPS.
+  const MAX_PLAUSIBLE_TPS = 10_000;
+
   const streamMs =
     timing.updateCount > 0 && timing.firstStreamUpdateMs !== null
       ? timing.lastStreamUpdateMs - timing.firstStreamUpdateMs
@@ -385,6 +433,30 @@ function buildTelemetry(timing: TurnTiming, turnEndMs: number): TurnTelemetry |
     tps = null;
   }
 
+  // Volume-based sanity gate: extraordinary TPS claims require
+  // proportionally longer measurement windows. When a provider emits a
+  // large volume of tokens in a short time window, the measured rate is
+  // dominated by dispatch/buffer-flush timing rather than actual inference
+  // speed. For X output tokens, the minimum reliable generation window is
+  // X / MAX_PLAUSIBLE_TPS seconds — below that, the volume of tokens cannot
+  // be reliably timed, and TPS is set to null.
+  //
+  // Mathematically equivalent to: computed TPS > MAX_PLAUSIBLE_TPS → null.
+  // But phrased as a measurement-reliability principle: the evidence
+  // (effective window duration) must be proportional to the claim (token
+  // volume ÷ inferred rate).
+  if (tps !== null && tps > MAX_PLAUSIBLE_TPS) {
+    tps = null;
+    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 },
@@ -408,6 +480,7 @@ function buildTelemetry(timing: TurnTiming, turnEndMs: number): TurnTelemetry |
           total: costTotal,
         }
       : null,
+    rateUsdPerMTokens,
     timestamp: Date.now(),
   };
 }
@@ -425,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 ─────────────────────────────────────────────────────────
 
   /**
@@ -584,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.1.1",
+  "version": "1.3.0",
   "description": "Tokens-per-second tracker for pi — see your LLM generation speed after every agent turn",
   "keywords": [
     "pi-package"