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
@@ -13,8 +13,8 @@
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
 import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
 import { z } from 'zod';
-import { readdirSync, existsSync, readFileSync } from 'node:fs';
-import { join } from 'node:path';
+import { readdirSync, existsSync, readFileSync, appendFileSync, mkdirSync } from 'node:fs';
+import { join, dirname } from 'node:path';

 // ── Types ──

@@ -186,6 +186,31 @@ function buildZodSchema(inputs: Record<string, ToolParam>): z.ZodRawShape {
  return shape;
 }

+// ── Logging ──
+
+const LOG_FILE = process.env.MCP_TOOLS_LOG ?? '';
+
+function logToolCall(toolName: string, durationMs: number, success: boolean, error?: string): void {
+  const ts = new Date().toISOString();
+  const status = success ? 'ok' : 'error';
+  const line = error
+    ? `${ts}\t${toolName}\t${status}\t${durationMs}ms\t${error}\n`
+    : `${ts}\t${toolName}\t${status}\t${durationMs}ms\n`;
+
+  // Always log to stderr for process-level visibility
+  console.error(`[mcp-tools] ${toolName} ${status} (${durationMs}ms)${error ? ': ' + error : ''}`);
+
+  // Write to log file if configured
+  if (LOG_FILE) {
+    try {
+      mkdirSync(dirname(LOG_FILE), { recursive: true });
+      appendFileSync(LOG_FILE, line);
+    } catch {
+      // Don't fail tool calls over logging
+    }
+  }
+}
+
 // ── Main ──

 const rawDirs = process.env.PI_TOOLS_DIRS ?? '';
@@ -218,11 +243,14 @@ for (const { entryFile, meta } of tools) {
      inputSchema: z.object(schema),
    },
    async (params) => {
+      const start = performance.now();
      let executeFn: Function | undefined;
      try {
        const mod = await import(capturedEntry);
        executeFn = mod.execute ?? mod.default?.execute;
      } catch (err) {
+        const ms = Math.round(performance.now() - start);
+        logToolCall(toolName, ms, false, `load failed: ${String(err)}`);
        return {
          content: [{ type: 'text' as const, text: `Failed to load ${toolName}: ${String(err)}` }],
          isError: true,
@@ -230,14 +258,27 @@ for (const { entryFile, meta } of tools) {
      }

      if (typeof executeFn !== 'function') {
+        const ms = Math.round(performance.now() - start);
+        logToolCall(toolName, ms, false, 'no execute function');
        return {
          content: [{ type: 'text' as const, text: `${toolName}/index.ts must export an "execute" function` }],
          isError: true,
        };
      }

-      const result = await executeFn('mcp', params, undefined, undefined);
-      return result;
+      try {
+        const result = await executeFn('mcp', params, undefined, undefined);
+        const ms = Math.round(performance.now() - start);
+        logToolCall(toolName, ms, !result.isError, result.isError ? 'tool returned error' : undefined);
+        return result;
+      } catch (err) {
+        const ms = Math.round(performance.now() - start);
+        logToolCall(toolName, ms, false, String(err));
+        return {
+          content: [{ type: 'text' as const, text: `${toolName} threw: ${String(err)}` }],
+          isError: true,
+        };
+      }
    },
  );

@@ -1,9 +1,10 @@
 import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs';
 import { join, resolve } from 'node:path';
 import type { Subprocess } from 'bun';
-import type { PiEvent, MessageCost } from '../../api/pi/types';
+import type { PiEvent } from '../../api/pi/types';
 import type { ClaudeSpawnParams, ClaudeSpawnStreamingParams, ClaudeCodeResult } from '../protocol';
 import { setClaudeSession, clearClaudeSession, getClaudeSession } from './state';
+import { parseStream } from './stream-parser';
 import { getProxySecret } from './proxy';
 import { DATA_PATH, getHomeDir, getGlobalToolsDir, getUserToolsDir } from '../../data-path';
 import { generateContainerContext } from '../../generate-container-context';
@@ -36,6 +37,7 @@ function generateMcpConfig(email: string): string {
        env: {
          PI_TOOLS_DIRS: toolsDirs,
          OFFICER_EMAIL_DB: join(DATA_PATH, email, 'emails.db'),
+          MCP_TOOLS_LOG: join(DATA_PATH, email, 'logs', 'mcp-tools.log'),
        },
      },
    },
@@ -263,141 +265,20 @@ export async function spawnClaudeStreaming(
  // Process NDJSON stream
  try {
    const stdout = proc.stdout as ReadableStream<Uint8Array>;
-    const reader = stdout.getReader();
-    const decoder = new TextDecoder();
-    let buffer = '';
-    let textBuffer = '';
-    let gotResult = false;
-
-    const processLine = (line: string) => {
-      if (!line.trim()) return;
-      try {
-        const msg = JSON.parse(line) as Record<string, unknown>;
-        const type = msg.type as string;
-
-        if (type === 'stream_event') {
-          const event = msg.event as Record<string, unknown> | undefined;
-          if (event?.type === 'content_block_delta') {
-            const delta = event.delta as Record<string, unknown> | undefined;
-            if (delta?.type === 'text_delta' && typeof delta.text === 'string') {
-              textBuffer += delta.text;
-              onEvent({ type: 'delta', text: delta.text });
-            }
-          }
-        } else if (type === 'assistant') {
-          const message = msg.message as Record<string, unknown> | undefined;
-          const content = message?.content as Array<Record<string, unknown>> | undefined;
-          if (Array.isArray(content)) {
-            for (const block of content) {
-              if (block.type === 'text' && typeof block.text === 'string') {
-                onEvent({ type: 'text', text: block.text });
-                textBuffer = '';
-              } else if (block.type === 'tool_use') {
-                if (textBuffer) {
-                  onEvent({ type: 'text', text: textBuffer });
-                  textBuffer = '';
-                }
-                onEvent({
-                  type: 'tool:start',
-                  toolCallId: (block.id as string) ?? '',
-                  toolName: (block.name as string) ?? 'unknown',
-                  toolInput: (block.input as Record<string, unknown>) ?? {},
-                });
-              }
-            }
-          }
-        } else if (type === 'user') {
-          const message = msg.message as Record<string, unknown> | undefined;
-          const content = message?.content as Array<Record<string, unknown>> | undefined;
-          if (Array.isArray(content)) {
-            for (const block of content) {
-              if (block.type === 'tool_result') {
-                let output = '';
-                if (typeof block.content === 'string') {
-                  output = block.content;
-                } else if (Array.isArray(block.content)) {
-                  output = (block.content as Array<Record<string, unknown>>)
-                    .filter((c) => c.type === 'text')
-                    .map((c) => c.text as string)
-                    .join('\n');
-                }
-                onEvent({
-                  type: 'tool:result',
-                  toolCallId: (block.tool_use_id as string) ?? '',
-                  output,
-                  isError: (block.is_error as boolean) ?? false,
-                });
-              }
-            }
-          }
-        } else if (type === 'system' && msg.subtype === 'init') {
-          const sessionId = msg.session_id as string | undefined;
-          if (sessionId) {
-            setClaudeSession(sessionKey, sessionId);
-          }
-        } else if (type === 'result') {
-          gotResult = true;
-          clearTimeout(timeout);
-
-          const isError = (msg.is_error as boolean) ?? false;
-          const resultText = (msg.result as string) ?? '';
-
-          if (isError) {
-            if (textBuffer) {
-              onEvent({ type: 'text', text: textBuffer });
-              textBuffer = '';
-            }
-            onEvent({ type: 'error', message: resultText || 'Claude Code returned an error' });
-            return;
-          }
-
-          if (textBuffer) {
-            onEvent({ type: 'text', text: textBuffer });
-            textBuffer = '';
-          }
-
-          const usage = msg.usage as Record<string, number> | undefined;
-          const cost: MessageCost = {
-            inputTokens: usage?.input_tokens ?? 0,
-            outputTokens: usage?.output_tokens ?? 0,
-            totalUSD: (msg.total_cost_usd as number) ?? 0,
-          };
-
-          const sessionId = msg.session_id as string | undefined;
-          if (sessionId) {
-            setClaudeSession(sessionKey, sessionId);
-          }
-
-          onEvent({ type: 'result', cost });
-        }
-      } catch {
-        // Skip malformed JSON lines
-      }
+    const callbacks = {
+      onEvent,
+      onSessionId: (sessionId: string) => setClaudeSession(sessionKey, sessionId),
    };

-    while (true) {
-      const { done, value } = await reader.read();
-      if (done) break;
-      const chunk = decoder.decode(value, { stream: true });
-      buffer += chunk;
-      const lines = buffer.split('\n');
-      buffer = lines.pop()!;
-      for (const line of lines) {
-        processLine(line);
-      }
-    }
-
-    if (buffer.trim()) {
-      processLine(buffer);
-    }
+    const state = await parseStream(stdout, callbacks);

    clearTimeout(timeout);

-    if (!gotResult) {
+    if (!state.gotResult) {
      const exitCode = await proc.exited;
      const stderr = await new Response(proc.stderr as ReadableStream<Uint8Array>).text();
-      if (textBuffer) {
-        onEvent({ type: 'text', text: textBuffer });
+      if (state.textBuffer) {
+        onEvent({ type: 'text', text: state.textBuffer });
      }
      if (exitCode !== 0) {
        onEvent({ type: 'error', message: `Claude Code exited with code ${exitCode}: ${stderr.trim().slice(0, 200)}` });
@@ -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);
+  });
+});
@@ -0,0 +1,206 @@
+/**
+ * Claude Code NDJSON Stream Parser
+ *
+ * Parses the stream-json output format from Claude Code CLI into PiEvents.
+ * Handles: stream_event (deltas), assistant (text/tool_use), user (tool_result),
+ * system init, and result messages.
+ */
+
+import type { PiEvent, MessageCost } from '../../api/pi/types';
+
+type SessionCallback = (sessionId: string) => void;
+
+type StreamParserCallbacks = {
+  onEvent: (event: PiEvent) => void;
+  onSessionId: SessionCallback;
+};
+
+type ParseState = {
+  textBuffer: string;
+  gotResult: boolean;
+};
+
+function flushTextBuffer(state: ParseState, onEvent: (event: PiEvent) => void): void {
+  if (state.textBuffer) {
+    onEvent({ type: 'text', text: state.textBuffer });
+    state.textBuffer = '';
+  }
+}
+
+function handleStreamEvent(msg: Record<string, unknown>, state: ParseState, onEvent: (event: PiEvent) => void): void {
+  const event = msg.event as Record<string, unknown> | undefined;
+  if (event?.type === 'content_block_delta') {
+    const delta = event.delta as Record<string, unknown> | undefined;
+    if (delta?.type === 'text_delta' && typeof delta.text === 'string') {
+      state.textBuffer += delta.text;
+      onEvent({ type: 'delta', text: delta.text });
+    }
+  }
+}
+
+function handleAssistant(msg: Record<string, unknown>, state: ParseState, onEvent: (event: PiEvent) => void): void {
+  const message = msg.message as Record<string, unknown> | undefined;
+  const content = message?.content as Array<Record<string, unknown>> | undefined;
+  if (!Array.isArray(content)) return;
+
+  for (const block of content) {
+    if (block.type === 'text' && typeof block.text === 'string') {
+      onEvent({ type: 'text', text: block.text });
+      state.textBuffer = '';
+    } else if (block.type === 'tool_use') {
+      flushTextBuffer(state, onEvent);
+      onEvent({
+        type: 'tool:start',
+        toolCallId: (block.id as string) ?? '',
+        toolName: (block.name as string) ?? 'unknown',
+        toolInput: (block.input as Record<string, unknown>) ?? {},
+      });
+    }
+  }
+}
+
+function handleUser(msg: Record<string, unknown>, onEvent: (event: PiEvent) => void): void {
+  const message = msg.message as Record<string, unknown> | undefined;
+  const content = message?.content as Array<Record<string, unknown>> | undefined;
+  if (!Array.isArray(content)) return;
+
+  for (const block of content) {
+    if (block.type === 'tool_result') {
+      let output = '';
+      if (typeof block.content === 'string') {
+        output = block.content;
+      } else if (Array.isArray(block.content)) {
+        output = (block.content as Array<Record<string, unknown>>)
+          .filter((c) => c.type === 'text')
+          .map((c) => c.text as string)
+          .join('\n');
+      }
+      onEvent({
+        type: 'tool:result',
+        toolCallId: (block.tool_use_id as string) ?? '',
+        output,
+        isError: (block.is_error as boolean) ?? false,
+      });
+    }
+  }
+}
+
+function handleResult(
+  msg: Record<string, unknown>,
+  state: ParseState,
+  callbacks: StreamParserCallbacks,
+): void {
+  state.gotResult = true;
+
+  const isError = (msg.is_error as boolean) ?? false;
+  const resultText = (msg.result as string) ?? '';
+
+  if (isError) {
+    flushTextBuffer(state, callbacks.onEvent);
+    callbacks.onEvent({ type: 'error', message: resultText || 'Claude Code returned an error' });
+    return;
+  }
+
+  flushTextBuffer(state, callbacks.onEvent);
+
+  const usage = msg.usage as Record<string, number> | undefined;
+  const cost: MessageCost = {
+    inputTokens: usage?.input_tokens ?? 0,
+    outputTokens: usage?.output_tokens ?? 0,
+    totalUSD: (msg.total_cost_usd as number) ?? 0,
+  };
+
+  const sessionId = msg.session_id as string | undefined;
+  if (sessionId) {
+    callbacks.onSessionId(sessionId);
+  }
+
+  callbacks.onEvent({ type: 'result', cost });
+}
+
+/**
+ * Process a single NDJSON line from Claude Code's stream output.
+ * Returns false if the line was skipped (empty or malformed), true otherwise.
+ */
+export function processLine(
+  line: string,
+  state: ParseState,
+  callbacks: StreamParserCallbacks,
+): boolean {
+  if (!line.trim()) return false;
+
+  let msg: Record<string, unknown>;
+  try {
+    msg = JSON.parse(line) as Record<string, unknown>;
+  } catch {
+    return false;
+  }
+
+  const type = msg.type as string;
+
+  switch (type) {
+    case 'stream_event':
+      handleStreamEvent(msg, state, callbacks.onEvent);
+      break;
+    case 'assistant':
+      handleAssistant(msg, state, callbacks.onEvent);
+      break;
+    case 'user':
+      handleUser(msg, callbacks.onEvent);
+      break;
+    case 'system':
+      if (msg.subtype === 'init') {
+        const sessionId = msg.session_id as string | undefined;
+        if (sessionId) callbacks.onSessionId(sessionId);
+      }
+      break;
+    case 'result':
+      handleResult(msg, state, callbacks);
+      break;
+    default:
+      return false;
+  }
+
+  return true;
+}
+
+/**
+ * Create a fresh parse state for a new stream.
+ */
+export function createParseState(): ParseState {
+  return { textBuffer: '', gotResult: false };
+}
+
+/**
+ * Read an NDJSON stream and process each line.
+ * Returns the final parse state.
+ */
+export async function parseStream(
+  stdout: ReadableStream<Uint8Array>,
+  callbacks: StreamParserCallbacks,
+): Promise<ParseState> {
+  const reader = stdout.getReader();
+  const decoder = new TextDecoder();
+  const state = createParseState();
+  let buffer = '';
+
+  while (true) {
+    const { done, value } = await reader.read();
+    if (done) break;
+    const chunk = decoder.decode(value, { stream: true });
+    buffer += chunk;
+    const lines = buffer.split('\n');
+    buffer = lines.pop()!;
+    for (const line of lines) {
+      processLine(line, state, callbacks);
+    }
+  }
+
+  if (buffer.trim()) {
+    processLine(buffer, state, callbacks);
+  }
+
+  return state;
+}
+
+export type { ParseState, StreamParserCallbacks };