@monotykamary/pi-ledger 0.3.0 → 0.4.0

package/extensions/pi-ledger/__tests__/helpers.ts

@@ -201,7 +201,14 @@ export function createTestFixture(): TestFixture {
 
   const mockPi: Partial<ExtensionAPI> = {
     on: vi.fn((event: string, handler: (...args: unknown[]) => unknown) => {
-      handlers[event] = handler;
+      const previous = handlers[event];
+      handlers[event] = previous
+        ? (...args: unknown[]) => {
+            const previousResult = previous(...args);
+            const result = handler(...args);
+            return result === undefined ? previousResult : result;
+          }
+        : handler;
       return mockPi as ExtensionAPI;
     }) as unknown as ExtensionAPI['on'],
     appendEntry: appendEntrySpy,

package/extensions/pi-ledger/__tests__/in-process-session-observer.test.ts

@@ -0,0 +1,555 @@
+import { afterEach, describe, expect, it, vi } from 'vitest';
+import {
+  installInProcessAgentSessionObserver,
+  type AgentSessionClassLike,
+  type InProcessAgentSessionObserverCallbacks,
+  type InProcessAgentSessionObserverHandle,
+} from '../in-process-session-observer';
+
+type FakeSessionEvent = Record<string, unknown>;
+type FakeListener = (event: FakeSessionEvent) => void;
+
+function assistantEnd(overrides: Record<string, unknown> = {}): FakeSessionEvent {
+  return {
+    type: 'message_end',
+    message: {
+      role: 'assistant',
+      content: [{ type: 'text', text: 'private response text' }],
+      provider: 'openai',
+      model: 'gpt-child',
+      responseId: 'response-1',
+      usage: {
+        input: 120,
+        output: 30,
+        cacheRead: 10,
+        cacheWrite: 5,
+        reasoning: 4,
+        totalTokens: 165,
+        cost: {
+          input: 0.1,
+          output: 0.2,
+          cacheRead: 0.01,
+          cacheWrite: 0.02,
+          total: 0.33,
+        },
+      },
+      stopReason: 'toolUse',
+      timestamp: 1234,
+      ...overrides,
+    },
+  };
+}
+
+function toolStart(
+  toolCallId: string,
+  toolName = 'bash',
+  args: unknown = { command: 'private command' }
+): FakeSessionEvent {
+  return { type: 'tool_execution_start', toolCallId, toolName, args };
+}
+
+function toolEnd(toolCallId: string, toolName = 'bash'): FakeSessionEvent {
+  return {
+    type: 'tool_execution_end',
+    toolCallId,
+    toolName,
+    result: { content: [{ type: 'text', text: 'private tool output' }] },
+    isError: false,
+  };
+}
+
+function createFakeAgentSessionClass() {
+  return class FakeAgentSession {
+    readonly listeners = new Set<FakeListener>();
+    readonly prompts: string[] = [];
+    subscribeCalls = 0;
+    unsubscribeCalls = 0;
+    promptFailure: unknown;
+    subscribeFailure: unknown;
+    unsubscribeFailure: unknown;
+    promptEvent?: FakeSessionEvent;
+    lastPromptResult?: Promise<void>;
+    onUnsubscribe?: () => void;
+
+    constructor(
+      readonly sessionId: string,
+      readonly sessionFile: string | undefined = `/sessions/${sessionId}.jsonl`,
+      readonly model: { provider: string; id: string } | undefined = {
+        provider: 'openai',
+        id: 'session-model',
+      }
+    ) {}
+
+    prompt(text: string): Promise<void> {
+      this.prompts.push(text);
+      if (this.promptEvent) this.emit(this.promptEvent);
+      const result = this.promptFailure
+        ? Promise.reject(this.promptFailure)
+        : Promise.resolve(undefined);
+      this.lastPromptResult = result;
+      return result;
+    }
+
+    subscribe(listener: FakeListener): () => void {
+      this.subscribeCalls += 1;
+      if (this.subscribeFailure) throw this.subscribeFailure;
+      this.listeners.add(listener);
+      return () => {
+        this.unsubscribeCalls += 1;
+        this.listeners.delete(listener);
+        this.onUnsubscribe?.();
+        if (this.unsubscribeFailure) throw this.unsubscribeFailure;
+      };
+    }
+
+    emit(event: FakeSessionEvent): void {
+      for (const listener of [...this.listeners]) listener(event);
+    }
+
+    dispose(): void {
+      this.listeners.clear();
+    }
+  };
+}
+
+const handles: InProcessAgentSessionObserverHandle[] = [];
+
+function installFor(
+  AgentSession: AgentSessionClassLike,
+  callbacks: InProcessAgentSessionObserverCallbacks = {},
+  now: () => number = Date.now,
+  rootSessionIds?: Iterable<string>
+): InProcessAgentSessionObserverHandle {
+  const handle = installInProcessAgentSessionObserver({
+    ...callbacks,
+    rootSessionIds,
+    dependencies: { AgentSession, now },
+  });
+  handles.push(handle);
+  return handle;
+}
+
+afterEach(() => {
+  for (const handle of handles.splice(0).reverse()) handle.uninstall();
+});
+
+describe('in-process AgentSession observer', () => {
+  it('attaches one public listener before the first prompt and never wraps session instances', async () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    const usage = vi.fn();
+    installFor(FakeAgentSession, { onAssistantUsage: usage });
+    const session = new FakeAgentSession('child-1');
+    session.promptEvent = assistantEnd();
+
+    await session.prompt('private first prompt');
+    await session.prompt('private second prompt');
+
+    expect(session.subscribeCalls).toBe(1);
+    expect(session.listeners).toHaveLength(1);
+    expect(session.prompts).toEqual(['private first prompt', 'private second prompt']);
+    expect(Object.hasOwn(session, 'prompt')).toBe(false);
+    expect(usage).toHaveBeenCalledTimes(2);
+  });
+
+  it('excludes tracked root IDs and can observe the same instance after it is untracked', async () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    const usage = vi.fn();
+    const handle = installFor(FakeAgentSession, { onAssistantUsage: usage }, Date.now, [
+      'root-session',
+    ]);
+    const root = new FakeAgentSession('root-session');
+    const child = new FakeAgentSession('child-session');
+
+    await root.prompt('root prompt');
+    await child.prompt('child prompt');
+    root.emit(assistantEnd());
+    child.emit(assistantEnd());
+
+    expect(root.subscribeCalls).toBe(0);
+    expect(child.subscribeCalls).toBe(1);
+    expect(usage).toHaveBeenCalledTimes(1);
+    expect(usage.mock.calls[0]![0].sessionId).toBe('child-session');
+
+    handle.removeRootSessionId('root-session');
+    await root.prompt('now observable');
+    root.emit(assistantEnd({ responseId: 'root-after-untrack' }));
+
+    expect(root.subscribeCalls).toBe(1);
+    expect(usage).toHaveBeenCalledTimes(2);
+    expect(usage.mock.calls[1]![0].responseId).toBe('root-after-untrack');
+  });
+
+  it('reference-counts root exclusions across observer handles', async () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    const firstUsage = vi.fn();
+    const secondUsage = vi.fn();
+    const first = installFor(FakeAgentSession, { onAssistantUsage: firstUsage }, Date.now, [
+      'shared-root',
+    ]);
+    const second = installFor(FakeAgentSession, { onAssistantUsage: secondUsage }, Date.now, [
+      'shared-root',
+    ]);
+    const session = new FakeAgentSession('shared-root');
+
+    await session.prompt('excluded by both');
+    first.removeRootSessionId('shared-root');
+    await session.prompt('still excluded by second');
+
+    expect(session.subscribeCalls).toBe(0);
+
+    second.removeRootSessionId('shared-root');
+    await session.prompt('now observable');
+    session.emit(assistantEnd());
+
+    expect(session.subscribeCalls).toBe(1);
+    expect(firstUsage).toHaveBeenCalledTimes(1);
+    expect(secondUsage).toHaveBeenCalledTimes(1);
+  });
+
+  it('keeps concurrent child sessions and their tool clocks isolated', async () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    let now = 0;
+    const usage = vi.fn();
+    const intervals = vi.fn();
+    installFor(FakeAgentSession, { onAssistantUsage: usage, onToolInterval: intervals }, () => now);
+    const first = new FakeAgentSession('child-a');
+    const second = new FakeAgentSession('child-b');
+    await Promise.all([first.prompt('a'), second.prompt('b')]);
+
+    first.emit(assistantEnd({ responseId: 'response-a' }));
+    second.emit(assistantEnd({ responseId: 'response-b' }));
+    now = 10;
+    first.emit(toolStart('a-tool'));
+    now = 15;
+    second.emit(toolStart('b-tool'));
+    now = 25;
+    first.emit(toolEnd('a-tool'));
+    now = 40;
+    second.emit(toolEnd('b-tool'));
+
+    expect(usage.mock.calls.map((call) => call[0].sessionId)).toEqual(['child-a', 'child-b']);
+    expect(intervals).toHaveBeenCalledTimes(2);
+    expect(intervals.mock.calls[0]![0]).toMatchObject({
+      sessionId: 'child-a',
+      responseId: 'response-a',
+      startedAt: 10,
+      endedAt: 25,
+      durationMs: 15,
+    });
+    expect(intervals.mock.calls[1]![0]).toMatchObject({
+      sessionId: 'child-b',
+      responseId: 'response-b',
+      startedAt: 15,
+      endedAt: 40,
+      durationMs: 25,
+    });
+  });
+
+  it('normalizes assistant usage and parallel tools into text-free union intervals', async () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    let now = 100;
+    const usage = vi.fn();
+    const intervals = vi.fn();
+    installFor(FakeAgentSession, { onAssistantUsage: usage, onToolInterval: intervals }, () => now);
+    const session = new FakeAgentSession('child-usage', '/tmp/child.jsonl');
+    await session.prompt('secret prompt text');
+
+    session.emit(assistantEnd());
+    session.emit(toolStart('call-a', 'bash'));
+    now = 120;
+    session.emit(toolStart('call-b', 'read'));
+    now = 150;
+    session.emit(toolEnd('call-a', 'bash'));
+    expect(intervals).not.toHaveBeenCalled();
+    now = 180;
+    session.emit(toolEnd('call-b', 'read'));
+
+    expect(usage).toHaveBeenCalledWith({
+      type: 'assistant_usage',
+      sessionId: 'child-usage',
+      sessionFile: '/tmp/child.jsonl',
+      model: { provider: 'openai', modelId: 'gpt-child' },
+      responseId: 'response-1',
+      timestamp: 1234,
+      stopReason: 'toolUse',
+      usage: {
+        input: 120,
+        output: 30,
+        cacheRead: 10,
+        cacheWrite: 5,
+        reasoning: 4,
+        totalTokens: 165,
+        cost: {
+          input: 0.1,
+          output: 0.2,
+          cacheRead: 0.01,
+          cacheWrite: 0.02,
+          total: 0.33,
+        },
+      },
+    });
+    expect(intervals).toHaveBeenCalledWith({
+      type: 'tool_interval',
+      sessionId: 'child-usage',
+      sessionFile: '/tmp/child.jsonl',
+      model: { provider: 'openai', modelId: 'gpt-child' },
+      responseId: 'response-1',
+      startedAt: 100,
+      endedAt: 180,
+      durationMs: 80,
+      toolCalls: [
+        { toolCallId: 'call-a', toolName: 'bash' },
+        { toolCallId: 'call-b', toolName: 'read' },
+      ],
+    });
+
+    const serialized = JSON.stringify({ usage: usage.mock.calls, intervals: intervals.mock.calls });
+    expect(serialized).not.toContain('private response text');
+    expect(serialized).not.toContain('private command');
+    expect(serialized).not.toContain('private tool output');
+    expect(serialized).not.toContain('secret prompt text');
+  });
+
+  it('reference-counts installs and restores the exact prompt descriptor after the last uninstall', async () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    const original = Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')!;
+    const firstUsage = vi.fn();
+    const secondUsage = vi.fn();
+    const first = installFor(FakeAgentSession, { onAssistantUsage: firstUsage });
+    const wrapped = Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')!;
+    const second = installFor(FakeAgentSession, { onAssistantUsage: secondUsage });
+
+    expect(Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')!.value).toBe(
+      wrapped.value
+    );
+    const session = new FakeAgentSession('child-refcount');
+    await session.prompt('one');
+    expect(session.subscribeCalls).toBe(1);
+
+    first.uninstall();
+    session.emit(assistantEnd());
+    expect(firstUsage).not.toHaveBeenCalled();
+    expect(secondUsage).toHaveBeenCalledTimes(1);
+    expect(Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')!.value).toBe(
+      wrapped.value
+    );
+
+    second.uninstall();
+    second.uninstall();
+    const restored = Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')!;
+    expect(restored).toEqual(original);
+    expect(restored.value).toBe(original.value);
+    expect(session.listeners).toHaveLength(0);
+    expect(session.unsubscribeCalls).toBe(1);
+
+    await session.prompt('without observer');
+    expect(session.subscribeCalls).toBe(1);
+    const reinstalled = installFor(FakeAgentSession, { onAssistantUsage: firstUsage });
+    await session.prompt('after reinstall');
+    expect(reinstalled.installed).toBe(true);
+    expect(session.subscribeCalls).toBe(2);
+    expect(session.listeners).toHaveLength(1);
+  });
+
+  it('keeps the wrapper installed when the last uninstall reentrantly installs a subscriber', async () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    const original = Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')!;
+    const first = installFor(FakeAgentSession);
+    const wrapped = Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')!;
+    const session = new FakeAgentSession('child-reentrant');
+    await session.prompt('attach first listener');
+
+    const usage = vi.fn();
+    let replacement: InProcessAgentSessionObserverHandle | undefined;
+    session.onUnsubscribe = () => {
+      session.onUnsubscribe = undefined;
+      replacement = installFor(FakeAgentSession, { onAssistantUsage: usage });
+    };
+
+    first.uninstall();
+
+    expect(replacement?.installed).toBe(true);
+    expect(Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')!.value).toBe(
+      wrapped.value
+    );
+    await session.prompt('attach replacement listener');
+    session.emit(assistantEnd());
+    expect(usage).toHaveBeenCalledTimes(1);
+
+    replacement!.uninstall();
+    expect(Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')).toEqual(original);
+  });
+
+  it('shares one prototype wrapper across uncached module evaluations', async () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    vi.resetModules();
+    const firstModule = await import('../in-process-session-observer.js');
+    const first = firstModule.installInProcessAgentSessionObserver({
+      dependencies: { AgentSession: FakeAgentSession, now: Date.now },
+    });
+    handles.push(first);
+    const wrapper = Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')!.value;
+
+    vi.resetModules();
+    const secondModule = await import('../in-process-session-observer.js');
+    const second = secondModule.installInProcessAgentSessionObserver({
+      dependencies: { AgentSession: FakeAgentSession, now: Date.now },
+    });
+    handles.push(second);
+
+    expect(Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')!.value).toBe(
+      wrapper
+    );
+    const session = new FakeAgentSession('child-reload');
+    await session.prompt('one listener');
+    expect(session.subscribeCalls).toBe(1);
+  });
+
+  it('reattaches after an unsubscribe callback removes its listener and then throws', async () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    const errors = vi.fn();
+    const first = installFor(FakeAgentSession, { onError: errors });
+    const session = new FakeAgentSession('child-throwing-unsubscribe');
+    await session.prompt('attach first listener');
+    const unsubscribeError = new Error('unsubscribe failed after removal');
+    session.unsubscribeFailure = unsubscribeError;
+
+    first.uninstall();
+
+    expect(session.listeners).toHaveLength(0);
+    expect(errors).toHaveBeenCalledWith({ phase: 'uninstall', error: unsubscribeError });
+    session.unsubscribeFailure = undefined;
+    const usage = vi.fn();
+    installFor(FakeAgentSession, { onAssistantUsage: usage });
+    await session.prompt('reattach listener');
+    session.emit(assistantEnd());
+
+    expect(session.subscribeCalls).toBe(2);
+    expect(session.listeners).toHaveLength(1);
+    expect(usage).toHaveBeenCalledTimes(1);
+  });
+
+  it('preserves prompt rejection identity and retries a failed subscription without blocking prompts', async () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    const errors = vi.fn();
+    installFor(FakeAgentSession, { onError: errors });
+    const session = new FakeAgentSession('child-errors');
+    const subscribeError = new Error('subscribe failed');
+    session.subscribeFailure = subscribeError;
+
+    const firstResult = session.prompt('still delivered');
+    expect(firstResult).toBe(session.lastPromptResult);
+    await expect(firstResult).resolves.toBeUndefined();
+    expect(session.subscribeCalls).toBe(1);
+    expect(errors).toHaveBeenCalledWith({ phase: 'attach', error: subscribeError });
+
+    session.subscribeFailure = undefined;
+    const promptError = new Error('prompt failed');
+    session.promptFailure = promptError;
+    const secondResult = session.prompt('rejected prompt');
+    expect(secondResult).toBe(session.lastPromptResult);
+    await expect(secondResult).rejects.toBe(promptError);
+    expect(session.subscribeCalls).toBe(2);
+
+    await expect(session.prompt('rejected again')).rejects.toBe(promptError);
+    expect(session.subscribeCalls).toBe(2);
+  });
+
+  it('contains callback failures and tolerates session disposal', async () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    const callbackError = new Error('consumer failed');
+    const errors = vi.fn();
+    const usage = vi.fn(() => {
+      throw callbackError;
+    });
+    const handle = installFor(FakeAgentSession, { onAssistantUsage: usage, onError: errors });
+    const session = new FakeAgentSession('child-dispose');
+    await session.prompt('attach');
+
+    expect(() => session.emit(assistantEnd())).not.toThrow();
+    expect(errors).toHaveBeenCalledWith({ phase: 'callback', error: callbackError });
+    session.dispose();
+    expect(session.listeners).toHaveLength(0);
+    expect(() => session.emit(assistantEnd())).not.toThrow();
+    expect(usage).toHaveBeenCalledTimes(1);
+    expect(() => handle.uninstall()).not.toThrow();
+  });
+
+  it('rejects a malformed process-global patch state without touching the prototype', () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    const seed = installFor(FakeAgentSession);
+    seed.uninstall();
+
+    const registry = Reflect.get(
+      globalThis,
+      Symbol.for('@monotykamary/pi-ledger/in-process-agent-session-observer/v1')
+    ) as { patches: WeakMap<object, unknown> };
+    const malformed = Object.create(null) as object;
+    Object.defineProperty(malformed, 'brand', {
+      get() {
+        throw new Error('poisoned patch brand');
+      },
+    });
+    registry.patches.set(FakeAgentSession.prototype, malformed);
+
+    try {
+      const before = Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')!;
+      const errors = vi.fn();
+      const handle = installFor(FakeAgentSession, { onError: errors });
+
+      expect(Object.isFrozen(registry)).toBe(true);
+      expect(handle).toMatchObject({ installed: false, incompatibility: 'prototype-conflict' });
+      expect(Object.getOwnPropertyDescriptor(FakeAgentSession.prototype, 'prompt')).toEqual(before);
+      expect(errors.mock.calls[0]![0].phase).toBe('install');
+    } finally {
+      registry.patches.delete(FakeAgentSession.prototype);
+    }
+  });
+
+  it('degrades without patching when the public prototype contract is incompatible', async () => {
+    class NonWritableSession {
+      prompt(): Promise<void> {
+        return Promise.resolve();
+      }
+      subscribe(): () => void {
+        return () => undefined;
+      }
+    }
+    const descriptor = Object.getOwnPropertyDescriptor(NonWritableSession.prototype, 'prompt')!;
+    Object.defineProperty(NonWritableSession.prototype, 'prompt', {
+      ...descriptor,
+      writable: false,
+    });
+    const before = Object.getOwnPropertyDescriptor(NonWritableSession.prototype, 'prompt')!;
+    const errors = vi.fn();
+
+    const handle = installFor(NonWritableSession, { onError: errors });
+
+    expect(handle).toMatchObject({ installed: false, incompatibility: 'prompt-descriptor' });
+    expect(Object.getOwnPropertyDescriptor(NonWritableSession.prototype, 'prompt')).toEqual(before);
+    await expect(new NonWritableSession().prompt()).resolves.toBeUndefined();
+    expect(errors.mock.calls[0]![0].phase).toBe('install');
+  });
+
+  it('keeps observing safe fields when optional metadata getters fail', async () => {
+    const FakeAgentSession = createFakeAgentSessionClass();
+    const metadataError = new Error('session file unavailable');
+    const errors = vi.fn();
+    const usage = vi.fn();
+    installFor(FakeAgentSession, { onAssistantUsage: usage, onError: errors });
+    const session = new FakeAgentSession('child-metadata');
+    Object.defineProperty(session, 'sessionFile', {
+      get() {
+        throw metadataError;
+      },
+    });
+
+    await session.prompt('attach');
+    expect(() => session.emit(assistantEnd())).not.toThrow();
+
+    expect(usage).toHaveBeenCalledWith(
+      expect.objectContaining({ sessionId: 'child-metadata', responseId: 'response-1' })
+    );
+    expect(errors).toHaveBeenCalledWith({ phase: 'metadata', error: metadataError });
+  });
+});