extract stream parser module with tests, add mcp tool call logging
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -13,8 +13,8 @@
|
|||||||
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
|
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
|
||||||
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
|
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
|
||||||
import { z } from 'zod';
|
import { z } from 'zod';
|
||||||
import { readdirSync, existsSync, readFileSync } from 'node:fs';
|
import { readdirSync, existsSync, readFileSync, appendFileSync, mkdirSync } from 'node:fs';
|
||||||
import { join } from 'node:path';
|
import { join, dirname } from 'node:path';
|
||||||
|
|
||||||
// ── Types ──
|
// ── Types ──
|
||||||
|
|
||||||
@@ -186,6 +186,31 @@ function buildZodSchema(inputs: Record<string, ToolParam>): z.ZodRawShape {
|
|||||||
return shape;
|
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 ──
|
// ── Main ──
|
||||||
|
|
||||||
const rawDirs = process.env.PI_TOOLS_DIRS ?? '';
|
const rawDirs = process.env.PI_TOOLS_DIRS ?? '';
|
||||||
@@ -218,11 +243,14 @@ for (const { entryFile, meta } of tools) {
|
|||||||
inputSchema: z.object(schema),
|
inputSchema: z.object(schema),
|
||||||
},
|
},
|
||||||
async (params) => {
|
async (params) => {
|
||||||
|
const start = performance.now();
|
||||||
let executeFn: Function | undefined;
|
let executeFn: Function | undefined;
|
||||||
try {
|
try {
|
||||||
const mod = await import(capturedEntry);
|
const mod = await import(capturedEntry);
|
||||||
executeFn = mod.execute ?? mod.default?.execute;
|
executeFn = mod.execute ?? mod.default?.execute;
|
||||||
} catch (err) {
|
} catch (err) {
|
||||||
|
const ms = Math.round(performance.now() - start);
|
||||||
|
logToolCall(toolName, ms, false, `load failed: ${String(err)}`);
|
||||||
return {
|
return {
|
||||||
content: [{ type: 'text' as const, text: `Failed to load ${toolName}: ${String(err)}` }],
|
content: [{ type: 'text' as const, text: `Failed to load ${toolName}: ${String(err)}` }],
|
||||||
isError: true,
|
isError: true,
|
||||||
@@ -230,14 +258,27 @@ for (const { entryFile, meta } of tools) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if (typeof executeFn !== 'function') {
|
if (typeof executeFn !== 'function') {
|
||||||
|
const ms = Math.round(performance.now() - start);
|
||||||
|
logToolCall(toolName, ms, false, 'no execute function');
|
||||||
return {
|
return {
|
||||||
content: [{ type: 'text' as const, text: `${toolName}/index.ts must export an "execute" function` }],
|
content: [{ type: 'text' as const, text: `${toolName}/index.ts must export an "execute" function` }],
|
||||||
isError: true,
|
isError: true,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
const result = await executeFn('mcp', params, undefined, undefined);
|
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;
|
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 { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs';
|
||||||
import { join, resolve } from 'node:path';
|
import { join, resolve } from 'node:path';
|
||||||
import type { Subprocess } from 'bun';
|
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 type { ClaudeSpawnParams, ClaudeSpawnStreamingParams, ClaudeCodeResult } from '../protocol';
|
||||||
import { setClaudeSession, clearClaudeSession, getClaudeSession } from './state';
|
import { setClaudeSession, clearClaudeSession, getClaudeSession } from './state';
|
||||||
|
import { parseStream } from './stream-parser';
|
||||||
import { getProxySecret } from './proxy';
|
import { getProxySecret } from './proxy';
|
||||||
import { DATA_PATH, getHomeDir, getGlobalToolsDir, getUserToolsDir } from '../../data-path';
|
import { DATA_PATH, getHomeDir, getGlobalToolsDir, getUserToolsDir } from '../../data-path';
|
||||||
import { generateContainerContext } from '../../generate-container-context';
|
import { generateContainerContext } from '../../generate-container-context';
|
||||||
@@ -36,6 +37,7 @@ function generateMcpConfig(email: string): string {
|
|||||||
env: {
|
env: {
|
||||||
PI_TOOLS_DIRS: toolsDirs,
|
PI_TOOLS_DIRS: toolsDirs,
|
||||||
OFFICER_EMAIL_DB: join(DATA_PATH, email, 'emails.db'),
|
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
|
// Process NDJSON stream
|
||||||
try {
|
try {
|
||||||
const stdout = proc.stdout as ReadableStream<Uint8Array>;
|
const stdout = proc.stdout as ReadableStream<Uint8Array>;
|
||||||
const reader = stdout.getReader();
|
const callbacks = {
|
||||||
const decoder = new TextDecoder();
|
onEvent,
|
||||||
let buffer = '';
|
onSessionId: (sessionId: string) => setClaudeSession(sessionKey, sessionId),
|
||||||
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;
|
const state = await parseStream(stdout, callbacks);
|
||||||
if (sessionId) {
|
|
||||||
setClaudeSession(sessionKey, sessionId);
|
|
||||||
}
|
|
||||||
|
|
||||||
onEvent({ type: 'result', cost });
|
|
||||||
}
|
|
||||||
} catch {
|
|
||||||
// Skip malformed JSON lines
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
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);
|
|
||||||
}
|
|
||||||
|
|
||||||
clearTimeout(timeout);
|
clearTimeout(timeout);
|
||||||
|
|
||||||
if (!gotResult) {
|
if (!state.gotResult) {
|
||||||
const exitCode = await proc.exited;
|
const exitCode = await proc.exited;
|
||||||
const stderr = await new Response(proc.stderr as ReadableStream<Uint8Array>).text();
|
const stderr = await new Response(proc.stderr as ReadableStream<Uint8Array>).text();
|
||||||
if (textBuffer) {
|
if (state.textBuffer) {
|
||||||
onEvent({ type: 'text', text: textBuffer });
|
onEvent({ type: 'text', text: state.textBuffer });
|
||||||
}
|
}
|
||||||
if (exitCode !== 0) {
|
if (exitCode !== 0) {
|
||||||
onEvent({ type: 'error', message: `Claude Code exited with code ${exitCode}: ${stderr.trim().slice(0, 200)}` });
|
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 };
|
||||||
Reference in New Issue
Block a user