@monotykamary/pi-tps 1.1.0 → 1.1.1

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/index.ts

@@ -303,7 +303,7 @@ 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;
@@ -601,7 +601,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.1.1",
   "description": "Tokens-per-second tracker for pi — see your LLM generation speed after every agent turn",
   "keywords": [
     "pi-package"