extract stream parser module with tests, add mcp tool call logging

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-07 13:49:26 +00:00
parent 86ddcbfead
commit c38d5b0ea1
4 changed files with 550 additions and 133 deletions
@@ -0,0 +1,289 @@
+import { describe, test, expect } from 'bun:test';
+import { processLine, createParseState, parseStream } from './stream-parser';
+import type { PiEvent } from '../../api/pi/types';
+import type { StreamParserCallbacks, ParseState } from './stream-parser';
+
+function makeCallbacks(): { events: PiEvent[]; sessionIds: string[]; callbacks: StreamParserCallbacks } {
+  const events: PiEvent[] = [];
+  const sessionIds: string[] = [];
+  return {
+    events,
+    sessionIds,
+    callbacks: {
+      onEvent: (e) => events.push(e),
+      onSessionId: (id) => sessionIds.push(id),
+    },
+  };
+}
+
+describe('processLine', () => {
+  test('skips empty lines', () => {
+    const state = createParseState();
+    const { events, callbacks } = makeCallbacks();
+    expect(processLine('', state, callbacks)).toBe(false);
+    expect(processLine('   ', state, callbacks)).toBe(false);
+    expect(events).toHaveLength(0);
+  });
+
+  test('skips malformed JSON', () => {
+    const state = createParseState();
+    const { events, callbacks } = makeCallbacks();
+    expect(processLine('not json', state, callbacks)).toBe(false);
+    expect(events).toHaveLength(0);
+  });
+
+  test('handles stream_event text delta', () => {
+    const state = createParseState();
+    const { events, callbacks } = makeCallbacks();
+    const line = JSON.stringify({
+      type: 'stream_event',
+      event: { type: 'content_block_delta', delta: { type: 'text_delta', text: 'hello' } },
+    });
+    processLine(line, state, callbacks);
+    expect(events).toEqual([{ type: 'delta', text: 'hello' }]);
+    expect(state.textBuffer).toBe('hello');
+  });
+
+  test('accumulates text buffer across deltas', () => {
+    const state = createParseState();
+    const { events, callbacks } = makeCallbacks();
+    const mkDelta = (text: string) =>
+      JSON.stringify({
+        type: 'stream_event',
+        event: { type: 'content_block_delta', delta: { type: 'text_delta', text } },
+      });
+    processLine(mkDelta('hello '), state, callbacks);
+    processLine(mkDelta('world'), state, callbacks);
+    expect(state.textBuffer).toBe('hello world');
+    expect(events).toHaveLength(2);
+  });
+
+  test('handles assistant text block — clears text buffer', () => {
+    const state = createParseState();
+    state.textBuffer = 'partial';
+    const { events, callbacks } = makeCallbacks();
+    const line = JSON.stringify({
+      type: 'assistant',
+      message: { content: [{ type: 'text', text: 'full response' }] },
+    });
+    processLine(line, state, callbacks);
+    expect(events).toEqual([{ type: 'text', text: 'full response' }]);
+    expect(state.textBuffer).toBe('');
+  });
+
+  test('handles assistant tool_use block — flushes text buffer first', () => {
+    const state = createParseState();
+    state.textBuffer = 'thinking...';
+    const { events, callbacks } = makeCallbacks();
+    const line = JSON.stringify({
+      type: 'assistant',
+      message: {
+        content: [{ type: 'tool_use', id: 'call_1', name: 'sqlite', input: { query: 'SELECT 1' } }],
+      },
+    });
+    processLine(line, state, callbacks);
+    expect(events).toEqual([
+      { type: 'text', text: 'thinking...' },
+      { type: 'tool:start', toolCallId: 'call_1', toolName: 'sqlite', toolInput: { query: 'SELECT 1' } },
+    ]);
+    expect(state.textBuffer).toBe('');
+  });
+
+  test('handles assistant tool_use without prior text buffer', () => {
+    const state = createParseState();
+    const { events, callbacks } = makeCallbacks();
+    const line = JSON.stringify({
+      type: 'assistant',
+      message: {
+        content: [{ type: 'tool_use', id: 'call_1', name: 'email_db', input: {} }],
+      },
+    });
+    processLine(line, state, callbacks);
+    expect(events).toEqual([
+      { type: 'tool:start', toolCallId: 'call_1', toolName: 'email_db', toolInput: {} },
+    ]);
+  });
+
+  test('handles user tool_result with string content', () => {
+    const state = createParseState();
+    const { events, callbacks } = makeCallbacks();
+    const line = JSON.stringify({
+      type: 'user',
+      message: {
+        content: [{ type: 'tool_result', tool_use_id: 'call_1', content: 'result text', is_error: false }],
+      },
+    });
+    processLine(line, state, callbacks);
+    expect(events).toEqual([
+      { type: 'tool:result', toolCallId: 'call_1', output: 'result text', isError: false },
+    ]);
+  });
+
+  test('handles user tool_result with array content', () => {
+    const state = createParseState();
+    const { events, callbacks } = makeCallbacks();
+    const line = JSON.stringify({
+      type: 'user',
+      message: {
+        content: [
+          {
+            type: 'tool_result',
+            tool_use_id: 'call_2',
+            content: [
+              { type: 'text', text: 'line 1' },
+              { type: 'image', data: 'ignored' },
+              { type: 'text', text: 'line 2' },
+            ],
+            is_error: true,
+          },
+        ],
+      },
+    });
+    processLine(line, state, callbacks);
+    expect(events).toEqual([
+      { type: 'tool:result', toolCallId: 'call_2', output: 'line 1\nline 2', isError: true },
+    ]);
+  });
+
+  test('handles system init — captures session id', () => {
+    const state = createParseState();
+    const { sessionIds, callbacks } = makeCallbacks();
+    const line = JSON.stringify({ type: 'system', subtype: 'init', session_id: 'sess_abc123' });
+    processLine(line, state, callbacks);
+    expect(sessionIds).toEqual(['sess_abc123']);
+  });
+
+  test('handles result — sets gotResult, emits cost', () => {
+    const state = createParseState();
+    state.textBuffer = 'trailing';
+    const { events, sessionIds, callbacks } = makeCallbacks();
+    const line = JSON.stringify({
+      type: 'result',
+      is_error: false,
+      result: 'done',
+      session_id: 'sess_xyz',
+      usage: { input_tokens: 100, output_tokens: 50 },
+      total_cost_usd: 0.003,
+    });
+    processLine(line, state, callbacks);
+    expect(state.gotResult).toBe(true);
+    expect(sessionIds).toEqual(['sess_xyz']);
+    expect(events).toEqual([
+      { type: 'text', text: 'trailing' },
+      { type: 'result', cost: { inputTokens: 100, outputTokens: 50, totalUSD: 0.003 } },
+    ]);
+  });
+
+  test('handles error result — flushes buffer and emits error', () => {
+    const state = createParseState();
+    state.textBuffer = 'partial';
+    const { events, callbacks } = makeCallbacks();
+    const line = JSON.stringify({ type: 'result', is_error: true, result: 'something broke' });
+    processLine(line, state, callbacks);
+    expect(state.gotResult).toBe(true);
+    expect(events).toEqual([
+      { type: 'text', text: 'partial' },
+      { type: 'error', message: 'something broke' },
+    ]);
+  });
+
+  test('handles result with missing usage — defaults to zero', () => {
+    const state = createParseState();
+    const { events, callbacks } = makeCallbacks();
+    const line = JSON.stringify({ type: 'result', is_error: false });
+    processLine(line, state, callbacks);
+    expect(events).toEqual([
+      { type: 'result', cost: { inputTokens: 0, outputTokens: 0, totalUSD: 0 } },
+    ]);
+  });
+
+  test('ignores unknown message types', () => {
+    const state = createParseState();
+    const { events, callbacks } = makeCallbacks();
+    expect(processLine(JSON.stringify({ type: 'unknown_thing' }), state, callbacks)).toBe(false);
+    expect(events).toHaveLength(0);
+  });
+});
+
+describe('parseStream', () => {
+  function makeStream(lines: string[]): ReadableStream<Uint8Array> {
+    const encoder = new TextEncoder();
+    const data = encoder.encode(lines.join('\n') + '\n');
+    return new ReadableStream({
+      start(controller) {
+        controller.enqueue(data);
+        controller.close();
+      },
+    });
+  }
+
+  test('parses a full conversation stream', async () => {
+    const lines = [
+      JSON.stringify({ type: 'system', subtype: 'init', session_id: 'sess_1' }),
+      JSON.stringify({
+        type: 'stream_event',
+        event: { type: 'content_block_delta', delta: { type: 'text_delta', text: 'Hello' } },
+      }),
+      JSON.stringify({
+        type: 'assistant',
+        message: { content: [{ type: 'text', text: 'Hello there' }] },
+      }),
+      JSON.stringify({
+        type: 'result',
+        is_error: false,
+        session_id: 'sess_1',
+        usage: { input_tokens: 10, output_tokens: 5 },
+        total_cost_usd: 0.001,
+      }),
+    ];
+
+    const { events, sessionIds, callbacks } = makeCallbacks();
+    const state = await parseStream(makeStream(lines), callbacks);
+
+    expect(state.gotResult).toBe(true);
+    expect(sessionIds).toEqual(['sess_1', 'sess_1']);
+    expect(events.map((e) => e.type)).toEqual(['delta', 'text', 'result']);
+  });
+
+  test('handles chunked delivery (split mid-line)', async () => {
+    const fullLine = JSON.stringify({
+      type: 'stream_event',
+      event: { type: 'content_block_delta', delta: { type: 'text_delta', text: 'chunked' } },
+    });
+    const resultLine = JSON.stringify({ type: 'result', is_error: false });
+
+    const encoder = new TextEncoder();
+    const mid = Math.floor(fullLine.length / 2);
+    const chunk1 = encoder.encode(fullLine.slice(0, mid));
+    const chunk2 = encoder.encode(fullLine.slice(mid) + '\n' + resultLine + '\n');
+
+    const stream = new ReadableStream({
+      start(controller) {
+        controller.enqueue(chunk1);
+        controller.enqueue(chunk2);
+        controller.close();
+      },
+    });
+
+    const { events, callbacks } = makeCallbacks();
+    const state = await parseStream(stream, callbacks);
+
+    expect(state.gotResult).toBe(true);
+    expect(events[0]).toEqual({ type: 'delta', text: 'chunked' });
+  });
+
+  test('handles trailing data without newline', async () => {
+    const line = JSON.stringify({ type: 'result', is_error: false });
+    const encoder = new TextEncoder();
+    const stream = new ReadableStream({
+      start(controller) {
+        controller.enqueue(encoder.encode(line)); // no trailing newline
+        controller.close();
+      },
+    });
+
+    const { events, callbacks } = makeCallbacks();
+    const state = await parseStream(stream, callbacks);
+    expect(state.gotResult).toBe(true);
+  });
+});