@monotykamary/pi-tps 1.0.0

package/extensions/pi-tps/__tests__/rehydration.test.ts

@@ -0,0 +1,266 @@
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
+import { createTestFixture, activateExtension, tick } from './helpers';
+
+describe('pi-tps extension — rehydration', () => {
+  let fixture: ReturnType<typeof createTestFixture>;
+
+  beforeEach(async () => {
+    fixture = createTestFixture();
+    await activateExtension(fixture);
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  /** Helper to make a structured TPS entry */
+  function makeTpsEntry(
+    data: {
+      provider?: string;
+      modelId?: string;
+      input?: number;
+      output?: number;
+      cacheRead?: number;
+      cacheWrite?: number;
+      total?: number;
+      ttftMs?: number;
+      totalMs?: number;
+      generationMs?: number;
+      stallMs?: number;
+      stallCount?: number;
+      messageCount?: number;
+      tps?: number;
+    } = {}
+  ) {
+    return {
+      type: 'custom' as const,
+      customType: 'tps',
+      data: {
+        model: { provider: data.provider ?? 'openai', modelId: data.modelId ?? 'gpt-4' },
+        tokens: {
+          input: data.input ?? 10,
+          output: data.output ?? 20,
+          cacheRead: data.cacheRead ?? 0,
+          cacheWrite: data.cacheWrite ?? 0,
+          total: data.total ?? 30,
+        },
+        timing: {
+          ttftMs: data.ttftMs ?? 1000,
+          totalMs: data.totalMs ?? 3000,
+          generationMs: data.generationMs ?? 2000,
+          stallMs: data.stallMs ?? 0,
+          stallCount: data.stallCount ?? 0,
+          messageCount: data.messageCount ?? 1,
+        },
+        tps: data.tps ?? 10.0,
+        timestamp: Date.now(),
+      },
+    };
+  }
+
+  it('should restore notification on session resume from structured telemetry', async () => {
+    const { handlers, notifySpy, mockEntries } = fixture;
+
+    mockEntries.push(
+      makeTpsEntry({
+        input: 50,
+        output: 100,
+        total: 150,
+        ttftMs: 1200,
+        totalMs: 5000,
+        generationMs: 4000,
+        tps: 25.0,
+      })
+    );
+
+    handlers['session_start']?.({ reason: 'resume' }, fixture.mockCtx);
+    await tick();
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const msg = notifySpy.mock.calls[0][0] as string;
+    expect(msg).toMatch(/TPS 25\.0 tok\/s/);
+    expect(msg).toMatch(/TTFT 1\.2s/);
+    expect(msg).toMatch(/out 100/);
+    expect(msg).toMatch(/in 50/);
+    expect(notifySpy).toHaveBeenCalledWith(msg, 'info');
+  });
+
+  it('should rehydrate legacy entries when no structured telemetry exists', async () => {
+    const { handlers, notifySpy, mockEntries } = fixture;
+
+    mockEntries.push({
+      type: 'custom',
+      customType: 'tps',
+      data: {
+        message: 'TPS 37.9 tok/s · TTFT 1s · 27s · out 998 · in 917',
+        timestamp: Date.now() - 500,
+      },
+    });
+
+    handlers['session_start']?.({ reason: 'resume' }, fixture.mockCtx);
+    await tick();
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const msg = notifySpy.mock.calls[0][0];
+    expect(msg).toContain('TPS 37.9');
+  });
+
+  it('should prefer structured telemetry over legacy entries', async () => {
+    const { handlers, notifySpy, mockEntries } = fixture;
+
+    mockEntries.push(
+      {
+        type: 'custom',
+        customType: 'tps',
+        data: {
+          message: 'TPS 37.9 tok/s · TTFT 1s · 27s · out 998 · in 917',
+          timestamp: Date.now() - 1000,
+        },
+      },
+      makeTpsEntry({
+        input: 273,
+        output: 51,
+        total: 324,
+        ttftMs: 1000,
+        totalMs: 3800,
+        generationMs: 2400,
+        stallMs: 1400,
+        stallCount: 1,
+        tps: 18.0,
+      })
+    );
+
+    handlers['session_start']?.({ reason: 'resume' }, fixture.mockCtx);
+    await tick();
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const msg = notifySpy.mock.calls[0][0];
+    expect(msg).toContain('TPS 18.0');
+    expect(msg).toContain('TTFT 1.0s');
+    expect(msg).toContain('stall 1.4s×1');
+    expect(msg).not.toContain('TPS 37.9');
+  });
+
+  it('should restore notification on session startup (continuing previous session)', async () => {
+    const { handlers, notifySpy, mockEntries } = fixture;
+
+    mockEntries.push(makeTpsEntry({ tps: 10.0 }));
+
+    handlers['session_start']?.({ reason: 'startup' }, fixture.mockCtx);
+    await tick();
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    expect(notifySpy.mock.calls[0][0]).toMatch(/TPS 10\.0 tok\/s/);
+  });
+
+  it('should restore notification on session reload', async () => {
+    const { handlers, notifySpy, mockEntries } = fixture;
+
+    mockEntries.push(makeTpsEntry({ tps: 10.0 }));
+
+    handlers['session_start']?.({ reason: 'reload' }, fixture.mockCtx);
+    await tick();
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+  });
+
+  it('should restore notification on tree navigation', async () => {
+    const { handlers, notifySpy, mockEntries } = fixture;
+
+    mockEntries.push(makeTpsEntry({ tps: 10.0 }));
+
+    handlers['session_tree']?.({ newLeafId: 'abc123', oldLeafId: 'def456' }, fixture.mockCtx);
+    await tick();
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+  });
+
+  it('should rehydrate most recent entry, preferring structured over legacy', async () => {
+    const { handlers, notifySpy, mockEntries } = fixture;
+
+    mockEntries.push(
+      {
+        type: 'custom',
+        customType: 'tps',
+        data: {
+          message: 'TPS 10.0 tok/s · TTFT 1s · 5s · out 100 · in 50',
+          timestamp: Date.now() - 3000,
+        },
+      },
+      makeTpsEntry({
+        provider: 'a',
+        modelId: 'a-1',
+        input: 5,
+        output: 10,
+        total: 15,
+        ttftMs: 5000,
+        totalMs: 10000,
+        generationMs: 8000,
+        tps: 1.2,
+      }),
+      {
+        type: 'custom',
+        customType: 'tps',
+        data: {
+          message: 'TPS 15.0 tok/s · TTFT 2s · 10s · out 200 · in 100',
+          timestamp: Date.now() - 1000,
+        },
+      },
+      makeTpsEntry({
+        provider: 'b',
+        modelId: 'b-1',
+        input: 50,
+        output: 500,
+        total: 550,
+        ttftMs: 2000,
+        totalMs: 8000,
+        generationMs: 6000,
+        stallMs: 500,
+        stallCount: 1,
+        messageCount: 2,
+        tps: 83.3,
+      })
+    );
+
+    handlers['session_start']?.({ reason: 'resume' }, fixture.mockCtx);
+    await tick();
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const msg = notifySpy.mock.calls[0][0];
+    expect(msg).toContain('TPS 83.3');
+    expect(msg).toContain('stall');
+  });
+
+  it('should rehydrate most recent legacy entry when no structured entry follows', async () => {
+    const { handlers, notifySpy, mockEntries } = fixture;
+
+    mockEntries.push(
+      makeTpsEntry({
+        provider: 'a',
+        modelId: 'a-1',
+        input: 5,
+        output: 10,
+        total: 15,
+        ttftMs: 5000,
+        totalMs: 10000,
+        generationMs: 8000,
+        tps: 1.2,
+      }),
+      {
+        type: 'custom',
+        customType: 'tps',
+        data: {
+          message: 'TPS 37.9 tok/s · TTFT 1s · 27s · out 998 · in 917',
+          timestamp: Date.now() - 1000,
+        },
+      }
+    );
+
+    handlers['session_start']?.({ reason: 'resume' }, fixture.mockCtx);
+    await tick();
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const msg = notifySpy.mock.calls[0][0];
+    expect(msg).toContain('TPS 37.9');
+  });
+});

package/extensions/pi-tps/__tests__/session-export.test.ts

@@ -0,0 +1,204 @@
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import type { ExtensionCommandContext } from '@earendil-works/pi-coding-agent';
+import { unlinkSync, existsSync } from 'fs';
+import { createTestFixture, activateExtension } from './helpers';
+
+vi.mock('child_process', () => ({ execSync: vi.fn() }));
+
+describe('pi-tps extension — session-export command', () => {
+  let fixture: ReturnType<typeof createTestFixture>;
+
+  const branchEntries = [
+    {
+      type: 'custom',
+      customType: 'tps',
+      data: { tps: 10 },
+      id: '1',
+      parentId: null,
+      timestamp: '2026-01-01T00:00:00Z',
+    },
+    {
+      type: 'message',
+      role: 'user',
+      content: 'hello',
+      id: '2',
+      parentId: '1',
+      timestamp: '2026-01-01T00:00:01Z',
+    },
+    {
+      type: 'model_change',
+      provider: 'openai',
+      modelId: 'gpt-4',
+      id: '3',
+      parentId: '2',
+      timestamp: '2026-01-01T00:00:02Z',
+    },
+  ];
+
+  const allEntries = [
+    ...branchEntries,
+    {
+      type: 'custom',
+      customType: 'tps',
+      data: { tps: 20 },
+      id: '4',
+      parentId: '3',
+      timestamp: '2026-01-01T00:00:03Z',
+    },
+    {
+      type: 'message',
+      role: 'assistant',
+      content: 'hi',
+      id: '5',
+      parentId: '4',
+      timestamp: '2026-01-01T00:00:04Z',
+    },
+  ];
+
+  beforeEach(async () => {
+    fixture = createTestFixture();
+    await activateExtension(fixture);
+  });
+
+  afterEach(() => {
+    for (const call of fixture.notifySpy.mock.calls) {
+      const msg = call[0] as string;
+      if (typeof msg === 'string' && msg.includes('→ ')) {
+        const filepath = msg.split('→ ')[1];
+        if (filepath && existsSync(filepath)) {
+          try {
+            unlinkSync(filepath);
+          } catch {
+            /* ignore */
+          }
+        }
+      }
+    }
+    vi.restoreAllMocks();
+  });
+
+  it('should export all entry types from current branch by default', async () => {
+    const exportCtx = {
+      ...fixture.mockCtx,
+      sessionManager: {
+        getBranch: vi.fn().mockReturnValue(branchEntries),
+        getSessionId: vi.fn().mockReturnValue('test-session-id'),
+      },
+    } as ExtensionCommandContext;
+
+    await fixture.commands['session-export'].handler('', exportCtx);
+
+    expect(fixture.notifySpy).toHaveBeenCalledOnce();
+    const msg = fixture.notifySpy.mock.calls[0][0] as string;
+    expect(msg).toContain('Exported 3 entries');
+    expect(msg).toContain('pi-session-branch-');
+    expect(msg).toContain('/pi-sessions/');
+  });
+
+  it('should export all entries from full session with --full flag', async () => {
+    const exportCtx = {
+      ...fixture.mockCtx,
+      sessionManager: {
+        getEntries: vi.fn().mockReturnValue(allEntries),
+        getSessionId: vi.fn().mockReturnValue('test-session-id'),
+      },
+    } as ExtensionCommandContext;
+
+    await fixture.commands['session-export'].handler('--full', exportCtx);
+
+    expect(fixture.notifySpy).toHaveBeenCalledOnce();
+    const msg = fixture.notifySpy.mock.calls[0][0] as string;
+    expect(msg).toContain('Exported 5 entries');
+    expect(msg).toContain('pi-session-full-');
+  });
+
+  it('should include all entry types — messages, custom, model_change', async () => {
+    const exportCtx = {
+      ...fixture.mockCtx,
+      sessionManager: {
+        getBranch: vi.fn().mockReturnValue(branchEntries),
+        getSessionId: vi.fn().mockReturnValue('test-session-id'),
+      },
+    } as ExtensionCommandContext;
+
+    await fixture.commands['session-export'].handler('', exportCtx);
+
+    const msg = fixture.notifySpy.mock.calls[0][0] as string;
+    const filepath = msg.split('→ ')[1];
+    const fs = await import('fs');
+    const content = fs.readFileSync(filepath, 'utf8');
+    const lines = content
+      .trim()
+      .split('\n')
+      .map((l: string) => JSON.parse(l));
+
+    const types = lines.map((l: any) => l.type);
+    expect(types).toContain('custom');
+    expect(types).toContain('message');
+    expect(types).toContain('model_change');
+    expect(lines).toHaveLength(3);
+  });
+
+  it('should preserve parentIds without re-chaining (full session export)', async () => {
+    const exportCtx = {
+      ...fixture.mockCtx,
+      sessionManager: {
+        getEntries: vi.fn().mockReturnValue(allEntries),
+        getSessionId: vi.fn().mockReturnValue('test-session-id'),
+      },
+    } as ExtensionCommandContext;
+
+    await fixture.commands['session-export'].handler('--full', exportCtx);
+
+    const msg = fixture.notifySpy.mock.calls[0][0] as string;
+    const filepath = msg.split('→ ')[1];
+    const fs = await import('fs');
+    const content = fs.readFileSync(filepath, 'utf8');
+    const lines = content
+      .trim()
+      .split('\n')
+      .map((l: string) => JSON.parse(l));
+
+    // Verify the parentId chain is preserved as-is
+    expect(lines.find((l: any) => l.id === '1').parentId).toBeNull();
+    expect(lines.find((l: any) => l.id === '2').parentId).toBe('1');
+    expect(lines.find((l: any) => l.id === '3').parentId).toBe('2');
+    expect(lines.find((l: any) => l.id === '4').parentId).toBe('3');
+    expect(lines.find((l: any) => l.id === '5').parentId).toBe('4');
+  });
+
+  it('should show warning when no entries found', async () => {
+    const exportCtx = {
+      ...fixture.mockCtx,
+      sessionManager: {
+        getBranch: vi.fn().mockReturnValue([]),
+        getSessionId: vi.fn().mockReturnValue('test-session-id'),
+      },
+    } as ExtensionCommandContext;
+
+    await fixture.commands['session-export'].handler('', exportCtx);
+
+    expect(fixture.notifySpy).toHaveBeenCalledOnce();
+    const msg = fixture.notifySpy.mock.calls[0][0] as string;
+    expect(msg).toContain('No entries found');
+    expect(msg).toContain('current-branch');
+    expect(fixture.notifySpy).toHaveBeenCalledWith(msg, 'warning');
+  });
+
+  it('should show warning with --full when no entries found', async () => {
+    const exportCtx = {
+      ...fixture.mockCtx,
+      sessionManager: {
+        getEntries: vi.fn().mockReturnValue([]),
+        getSessionId: vi.fn().mockReturnValue('test-session-id'),
+      },
+    } as ExtensionCommandContext;
+
+    await fixture.commands['session-export'].handler('--full', exportCtx);
+
+    expect(fixture.notifySpy).toHaveBeenCalledOnce();
+    const msg = fixture.notifySpy.mock.calls[0][0] as string;
+    expect(msg).toContain('No entries found');
+    expect(msg).toContain('all-entries');
+  });
+});

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

@@ -0,0 +1,209 @@
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
+import type { AssistantMessage } from '@earendil-works/pi-ai';
+import { createTestFixture, activateExtension, tick } from './helpers';
+import type { TurnStartEvent, TurnEndEvent, MessageUpdateEvent } from './helpers';
+
+describe('pi-tps extension — stall detection', () => {
+  let fixture: ReturnType<typeof createTestFixture>;
+
+  beforeEach(async () => {
+    fixture = createTestFixture();
+    await activateExtension(fixture);
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it('should detect stalls between message_update events', async () => {
+    const { handlers, notifySpy, appendEntrySpy } = fixture;
+
+    const assistantMessage: AssistantMessage = {
+      role: 'assistant',
+      content: [{ type: 'text', text: 'Response with a stall' }],
+      api: 'openai-completions',
+      provider: 'deepseek',
+      model: 'deepseek-v4',
+      usage: {
+        input: 50,
+        output: 100,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 150,
+        cost: { input: 0.001, output: 0.002, cacheRead: 0, cacheWrite: 0, total: 0.003 },
+      },
+      stopReason: 'stop',
+      timestamp: Date.now(),
+    };
+
+    const updateEvent: MessageUpdateEvent = {
+      type: 'message_update',
+      message: assistantMessage,
+      assistantMessageEvent: { type: 'text_delta', delta: 't' },
+    };
+
+    handlers['turn_start']?.({ type: 'turn_start', turnIndex: 0, timestamp: Date.now() });
+    await tick(200); // TTFT
+    handlers['message_start']?.({ type: 'message_start', message: assistantMessage });
+
+    // Normal streaming
+    handlers['message_update']?.(updateEvent);
+    await tick(100);
+    handlers['message_update']?.(updateEvent);
+
+    // STALL: 800ms gap (> 500ms threshold)
+    await tick(800);
+    handlers['message_update']?.(updateEvent);
+
+    // Normal streaming resumes
+    await tick(100);
+    handlers['message_update']?.(updateEvent);
+
+    // Another stall: 600ms gap
+    await tick(600);
+    handlers['message_update']?.(updateEvent);
+
+    await tick(50);
+    handlers['message_end']?.({ type: 'message_end', message: assistantMessage });
+    handlers['turn_end']?.(
+      { type: 'turn_end', turnIndex: 0, message: assistantMessage, toolResults: [] },
+      fixture.mockCtx
+    );
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const notification = notifySpy.mock.calls[0][0] as string;
+
+    expect(notification).toMatch(/stall \d+\.\ds×2/);
+    expect(notification).toContain('TPS');
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.timing.stallCount).toBe(2);
+    expect(data.timing.stallMs).toBeGreaterThanOrEqual(1300);
+  });
+
+  it('should not flag short gaps as stalls', async () => {
+    const { handlers, notifySpy, appendEntrySpy } = fixture;
+
+    const assistantMessage: AssistantMessage = {
+      role: 'assistant',
+      content: [{ type: 'text', text: 'Smooth' }],
+      api: 'openai-completions',
+      provider: 'openai',
+      model: 'gpt-4',
+      usage: {
+        input: 10,
+        output: 20,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 30,
+        cost: { input: 0.001, output: 0.002, cacheRead: 0, cacheWrite: 0, total: 0.003 },
+      },
+      stopReason: 'stop',
+      timestamp: Date.now(),
+    };
+
+    const updateEvent: MessageUpdateEvent = {
+      type: 'message_update',
+      message: assistantMessage,
+      assistantMessageEvent: { type: 'text_delta', delta: 't' },
+    };
+
+    handlers['turn_start']?.({ type: 'turn_start', turnIndex: 0, timestamp: Date.now() });
+    await tick(100);
+    handlers['message_start']?.({ type: 'message_start', message: assistantMessage });
+
+    // Short gaps (< 500ms)
+    handlers['message_update']?.(updateEvent);
+    await tick(200);
+    handlers['message_update']?.(updateEvent);
+    await tick(300);
+    handlers['message_update']?.(updateEvent);
+    await tick(400); // borderline but < 500
+    handlers['message_update']?.(updateEvent);
+
+    await tick(50);
+    handlers['message_end']?.({ type: 'message_end', message: assistantMessage });
+    handlers['turn_end']?.(
+      { type: 'turn_end', turnIndex: 0, message: assistantMessage, toolResults: [] },
+      fixture.mockCtx
+    );
+
+    const notification = notifySpy.mock.calls[0][0] as string;
+    expect(notification).not.toContain('stall');
+
+    const [, data] = appendEntrySpy.mock.calls[0];
+    expect(data.timing.stallCount).toBe(0);
+    expect(data.timing.stallMs).toBe(0);
+  });
+
+  it('should not produce inflated TPS when stall occurs before first stream update', async () => {
+    // This is the bug scenario: a large gap after TTFT causes the stall
+    // detector to fire on the first streaming update, but that means
+    // firstStreamUpdateMs is set AFTER the stall, making streamMs tiny.
+    const { handlers, notifySpy, appendEntrySpy } = fixture;
+
+    const assistantMessage: AssistantMessage = {
+      role: 'assistant',
+      content: [{ type: 'text', text: 'Stall before stream' }],
+      api: 'openai-completions',
+      provider: 'deepseek',
+      model: 'deepseek-v4',
+      usage: {
+        input: 50,
+        output: 100,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 150,
+        cost: { input: 0.001, output: 0.002, cacheRead: 0, cacheWrite: 0, total: 0.003 },
+      },
+      stopReason: 'stop',
+      timestamp: Date.now(),
+    };
+
+    const updateEvent: MessageUpdateEvent = {
+      type: 'message_update',
+      message: assistantMessage,
+      assistantMessageEvent: { type: 'text_delta', delta: 't' },
+    };
+
+    handlers['turn_start']?.({ type: 'turn_start', turnIndex: 0, timestamp: Date.now() });
+    await tick(100);
+    handlers['message_start']?.({ type: 'message_start', message: assistantMessage });
+    await tick(200);
+    handlers['message_update']?.(updateEvent); // TTFT
+
+    // Large gap (stall before any stream update)
+    await tick(3000);
+    handlers['message_update']?.(updateEvent); // First stream update (after stall)
+
+    // Short burst of updates
+    await tick(50);
+    handlers['message_update']?.(updateEvent);
+    await tick(50);
+    handlers['message_update']?.(updateEvent);
+    await tick(50);
+    handlers['message_update']?.(updateEvent);
+    await tick(50);
+    handlers['message_update']?.(updateEvent);
+
+    await tick(50);
+    handlers['message_end']?.({ type: 'message_end', message: assistantMessage });
+    handlers['turn_end']?.(
+      { type: 'turn_end', turnIndex: 0, message: assistantMessage, toolResults: [] },
+      fixture.mockCtx
+    );
+
+    expect(notifySpy).toHaveBeenCalledOnce();
+    const [, data] = appendEntrySpy.mock.calls[0];
+
+    // Verify stall was detected
+    expect(data.timing.stallMs).toBeGreaterThanOrEqual(2000);
+    expect(data.timing.stallCount).toBeGreaterThanOrEqual(1);
+
+    // TPS must be sane (not in the thousands)
+    // With stall guard: stallMs >= streamMs → fallback branch
+    // Fallback: effectiveGenMs = generationMs - stallMs
+    expect(data.tps).not.toBeNull();
+    expect(data.tps!).toBeLessThan(500); // sane, not thousands
+  });
+});