@monotykamary/pi-tps 1.2.0 → 1.3.1

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,334 @@
+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/);
+  });
+
+  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

@@ -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,15 +125,35 @@ 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) => {
+    on: vi.fn((event: string, handler: any) => {
       handlers[event] = handler;
       return mockPi as ExtensionAPI;
     }),
     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,