@monotykamary/pi-tps 1.1.0 → 1.2.0

package/extensions/pi-tps/__tests__/dynamic-tps-cap.test.ts

@@ -128,15 +128,15 @@ describe('pi-tps extension — dynamic TPS cap', () => {
       messageEnd: 900,
     });
 
-    // Turn 2: tool call with fallback TPS (2 updates, 100ms generationMs)
-    // Without cap: 20 tokens / 0.055s ≈ 363 TPS (inflated)
-    // With cap: min(363, 50) = 50 TPS
+    // Turn 2: tool call with fallback TPS (2 updates, 250ms generationMs)
+    // Without cap: 20 tokens / 0.25s ≈ 80 TPS (feasible but from short window)
+    // With cap: min(80, 50) = 50 TPS
     const { appendEntrySpy, notifySpy } = driveTurn({
       turnStart: 0,
       messageStart: 100,
       firstUpdate: 100.1,
       streamUpdates: [100.15, 100.3],
-      messageEnd: 200,
+      messageEnd: 350,
       isToolCall: true,
     });
 
@@ -147,20 +147,20 @@ describe('pi-tps extension — dynamic TPS cap', () => {
     expect(data.tps).toBeGreaterThan(0);
   });
 
-  // ── Tool calls do not set the cap ────────────────────────────────────────
+  // ── Tool calls do not set the cap from fallback ─────────────────────────────
 
-  it('should not let tool-call turns set the cap', () => {
-    // Turn 1: tool call → no cap exists yet, TPS is null
+  it('should not let fallback-branch tool-call turns set the cap', () => {
+    // Turn 1: tool call with fallback TPS — should NOT set the cap
     const { appendEntrySpy: spy1 } = driveTurn({
       turnStart: 0,
       messageStart: 100,
       firstUpdate: 100.1,
       streamUpdates: [100.15, 100.3],
-      messageEnd: 200,
+      messageEnd: 350,
       isToolCall: true,
     });
     const [, data1] = spy1.mock.calls[0];
-    // No cap → tool call TPS is null
+    // No cap → fallback tool call TPS is null
     expect(data1.tps).toBeNull();
 
     // Turn 2: reliable streaming response at ~50 TPS → sets the cap
@@ -175,13 +175,13 @@ describe('pi-tps extension — dynamic TPS cap', () => {
     expect(data2.tps).toBeGreaterThanOrEqual(40);
     expect(data2.tps).toBeLessThanOrEqual(60);
 
-    // Turn 3: another tool call — should now be clamped to 50
+    // Turn 3: another fallback tool call — should now be clamped to 50
     const { appendEntrySpy: spy3 } = driveTurn({
       turnStart: 0,
       messageStart: 100,
       firstUpdate: 100.1,
       streamUpdates: [100.15, 100.3],
-      messageEnd: 200,
+      messageEnd: 350,
       isToolCall: true,
     });
     const [, data3] = spy3.mock.calls[2];
@@ -189,6 +189,40 @@ describe('pi-tps extension — dynamic TPS cap', () => {
     expect(data3.tps).toBeLessThanOrEqual(55);
   });
 
+  // ── Primary-branch tool calls (reasoning) set the cap ──────────────────────
+
+  it('should let primary-branch tool-call turns set the cap (e.g. reasoning before tool call)', () => {
+    // Turn 1: tool call with PRIMARY-branch TPS (reasoning + tool call, enough updates/time)
+    // 20 tokens / 0.4s = 50 TPS from primary branch, isToolCall = true
+    const { appendEntrySpy: spy1 } = driveTurn({
+      turnStart: 0,
+      messageStart: 200,
+      firstUpdate: 200.123,
+      streamUpdates: [400, 500, 600, 700, 800],
+      messageEnd: 900,
+      isToolCall: true,
+    });
+    const [, data1] = spy1.mock.calls[0];
+    // Primary branch + isToolCall → TPS is still computed (not null/capped)
+    expect(data1.tps).toBeGreaterThanOrEqual(40);
+    expect(data1.tps).toBeLessThanOrEqual(60);
+    expect(data1.isPrimaryBranch).toBe(true);
+
+    // Turn 2: fallback tool call — should be clamped to the cap from turn 1
+    const { appendEntrySpy: spy2 } = driveTurn({
+      turnStart: 0,
+      messageStart: 100,
+      firstUpdate: 100.1,
+      streamUpdates: [100.15, 100.3],
+      messageEnd: 350,
+      isToolCall: true,
+    });
+    const [, data2] = spy2.mock.calls[1];
+    expect(data2.tps).not.toBeNull();
+    // Clamped to ~50 cap set by the primary-branch tool call in turn 1
+    expect(data2.tps).toBeLessThanOrEqual(55);
+  });
+
   // ── Cold start: no cap yet ────────────────────────────────────────────────
 
   it('should show null TPS for tool calls when no cap exists yet', () => {
@@ -197,7 +231,7 @@ describe('pi-tps extension — dynamic TPS cap', () => {
       messageStart: 100,
       firstUpdate: 100.1,
       streamUpdates: [100.15, 100.3],
-      messageEnd: 200,
+      messageEnd: 350,
       isToolCall: true,
     });
 
@@ -228,7 +262,7 @@ describe('pi-tps extension — dynamic TPS cap', () => {
       messageStart: 100,
       firstUpdate: 100.1,
       streamUpdates: [100.15, 100.3],
-      messageEnd: 200,
+      messageEnd: 350,
       isToolCall: false,
     });
 
@@ -344,7 +378,7 @@ describe('pi-tps extension — dynamic TPS cap', () => {
       messageStart: 100,
       firstUpdate: 100.1,
       streamUpdates: [100.15, 100.3],
-      messageEnd: 200,
+      messageEnd: 350,
       isToolCall: true,
     });
 

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

@@ -303,11 +303,24 @@ function buildTelemetry(timing: TurnTiming, turnEndMs: number): TurnTelemetry |
   const MIN_STREAM_MS = 1;
   const MIN_STREAM_UPDATES = 5;
   const MIN_INTER_CHUNK_MS = 1;
-  const MIN_GENERATION_MS = 50;
+  const MIN_GENERATION_MS = 200;
   const ACTIVE_TIME_THRESHOLD_MS = 200;
   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 +398,23 @@ 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;
+  }
+
   return {
     model,
     tokens: { input, output, cacheRead, cacheWrite, total: totalTokens },
@@ -601,7 +631,7 @@ export default function tpsExtension(pi: ExtensionAPI) {
     // short outputs over tiny time windows.
     const modelKey = `${telemetry.model.provider}:${telemetry.model.modelId}`;
 
-    if (telemetry.isPrimaryBranch && !timing.isToolCall && telemetry.tps !== null) {
+    if (telemetry.isPrimaryBranch && telemetry.tps !== null) {
       const currentCap = tpsCaps.get(modelKey);
       if (currentCap === undefined || telemetry.tps > currentCap) {
         tpsCaps.set(modelKey, telemetry.tps);

package/npm-shrinkwrap.json

@@ -1,12 +1,12 @@
 {
   "name": "pi-tps",
-  "version": "1.1.0",
+  "version": "1.1.1",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "pi-tps",
-      "version": "1.1.0",
+      "version": "1.1.1",
       "hasInstallScript": true,
       "license": "MIT",
       "devDependencies": {

package/package.json

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