@monotykamary/pi-tps 1.0.0

package/extensions/pi-tps/__tests__/stall-reduction.test.ts

@@ -0,0 +1,139 @@
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import type { AssistantMessage } from '@earendil-works/pi-ai';
+import { createTestFixture, activateExtension } from './helpers';
+import type { MessageUpdateEvent } from './helpers';
+
+/**
+ * Tests partial stall reduction in the fallback branch.
+ *
+ * When stalls dominate the effective generation window (effectiveGenMs < 200ms
+ * OR stallMs > 85% of generationMs), the raw TPS would explode because the
+ * denominator is tiny. Partial reduction divides stallMs by 2 before subtracting
+ * it, giving a much larger safe denominator and a saner TPS.
+ */
+describe('pi-tps extension — partial stall reduction', () => {
+  let fixture: ReturnType<typeof createTestFixture>;
+
+  beforeEach(async () => {
+    fixture = createTestFixture();
+    await activateExtension(fixture);
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  /**
+   * Drive a full turn with mocked performance.now() so every
+   * timestamp is deterministic (no real timers needed).
+   */
+  function driveTurn(clocks: {
+    turnStart: number;
+    messageStart: number;
+    firstUpdate: number;
+    streamUpdates: number[];
+    messageEnd: number;
+    turnEnd?: number;
+  }) {
+    const { handlers, 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 assistantMessage: AssistantMessage = {
+      role: 'assistant',
+      content: [{ type: 'text' as const, text: 'x'.repeat(5000) }],
+      api: 'openai-completions',
+      provider: 'openai',
+      model: 'gpt-4',
+      usage: {
+        input: 100,
+        output: 5000,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 5100,
+        cost: { input: 0.001, output: 0.002, cacheRead: 0, cacheWrite: 0, total: 0.003 },
+      },
+      stopReason: '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' },
+      });
+    }
+    handlers['message_end']?.({ type: 'message_end', message: assistantMessage });
+    handlers['turn_end']?.(
+      { type: 'turn_end', turnIndex: 0, message: assistantMessage, toolResults: [] },
+      fixture.mockCtx
+    );
+
+    spy.mockRestore();
+    return { appendEntrySpy };
+  }
+
+  it('partially reduces stalls when effectiveGenMs is tiny', () => {
+    // generationMs ~ 50_300, stallMs ~ 50_120
+    // effectiveGenMs = 50_300 - 50_120 = 180  (< ACTIVE_TIME_THRESHOLD=200)
+    // stallMs (50_120) > effectiveGenMs (180) → partial reduction kicks in
+    // safeGenMs = max(50_300 - 50_120/2, 50) = max(25_240, 50) = 25_240
+    // raw = 5_000 / 25.24 = 198.1
+    // Without partial reduction: 5_000 / 0.18 = 27_778
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 100,
+      firstUpdate: 200,
+      streamUpdates: [250, 300],
+      messageEnd: 50400,
+    });
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.tps).not.toBeNull();
+    expect(data.tps).toBeLessThan(300);
+    expect(data.tps).toBeGreaterThan(50);
+  });
+
+  it('partially reduces stalls when stallMs dominates generationMs', () => {
+    // generationMs = 15_000, stallMs = 13_500
+    // effectiveGenMs = 1_500  (>200 so absolute threshold NOT hit)
+    // BUT stallMs/genMs ratio = 90% > STALL_DOMINANCE_RATIO (85%)
+    // → partial reduction kicks in via ratio branch
+    // safeGenMs = max(15_000 - 13_500/2, 50) = max(8_250, 50) = 8_250
+    // raw = 5_000 / 8.25 = 606
+    // Without partial reduction: 5_000 / 1.5 = 3_333
+    const { appendEntrySpy } = driveTurn({
+      turnStart: 0,
+      messageStart: 50,
+      firstUpdate: 150,
+      streamUpdates: [300, 800],
+      messageEnd: 15050,
+    });
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.tps).not.toBeNull();
+    expect(data.tps).toBeLessThan(700);
+    expect(data.tps).toBeGreaterThan(100);
+  });
+});