boocode/apps/coder/src/services/__tests__/stream-json-parser.test.ts

import { describe, it, expect } from 'vitest';
import {
  makeStreamJsonParser,
  makeStreamJsonState,
  parseStreamJsonLine,
  type AgentEventList,
} from '../stream-json-parser.js';
import type { AgentEvent } from '../agent-backend.js';
import type { AcpToolSnapshot } from '../acp-tool-snapshot.js';

// Helpers to JSON-encode the representative Claude-Code stream-json lines.
const sys = (sessionId: string) =>
  JSON.stringify({ type: 'system', subtype: 'init', session_id: sessionId, tools: ['read', 'edit'] });

const streamEvent = (event: unknown) => JSON.stringify({ type: 'stream_event', event });

const textDelta = (index: number, text: string) =>
  streamEvent({ type: 'content_block_delta', index, delta: { type: 'text_delta', text } });

const thinkingDelta = (index: number, thinking: string) =>
  streamEvent({ type: 'content_block_delta', index, delta: { type: 'thinking_delta', thinking } });

const toolStart = (index: number, id: string, name: string) =>
  streamEvent({ type: 'content_block_start', index, content_block: { type: 'tool_use', id, name } });

const inputJsonDelta = (index: number, partial: string) =>
  streamEvent({ type: 'content_block_delta', index, delta: { type: 'input_json_delta', partial_json: partial } });

const blockStop = (index: number) => streamEvent({ type: 'content_block_stop', index });

const resultLine = (input: number, output: number, sessionId?: string) =>
  JSON.stringify({ type: 'result', subtype: 'success', session_id: sessionId, usage: { input_tokens: input, output_tokens: output } });

describe('parseStreamJsonLine (pure per-line mapping)', () => {
  it('captures session_id from the system init line and emits no events', () => {
    const state = makeStreamJsonState();
    const events = parseStreamJsonLine(sys('sess-abc'), state);
    expect(events).toEqual([]);
    expect(state.sessionId).toBe('sess-abc');
  });

  it('maps a text_delta stream_event → a text event', () => {
    const state = makeStreamJsonState();
    expect(parseStreamJsonLine(textDelta(0, 'Hello'), state)).toEqual([{ type: 'text', text: 'Hello' }]);
  });

  it('maps a thinking_delta stream_event → a reasoning event', () => {
    const state = makeStreamJsonState();
    expect(parseStreamJsonLine(thinkingDelta(0, 'pondering'), state)).toEqual([
      { type: 'reasoning', text: 'pondering' },
    ]);
  });

  it('tolerates a garbage / non-JSON line (returns [], no throw)', () => {
    const state = makeStreamJsonState();
    expect(parseStreamJsonLine('not json at all {{{', state)).toEqual([]);
    expect(parseStreamJsonLine('', state)).toEqual([]);
    expect(parseStreamJsonLine('   ', state)).toEqual([]);
    // A truncated/partial JSON object also yields [] rather than throwing.
    expect(parseStreamJsonLine('{"type":"stream_event","eve', state)).toEqual([]);
  });

  it('ignores unknown top-level line types and the user (tool-result) line', () => {
    const state = makeStreamJsonState();
    expect(parseStreamJsonLine(JSON.stringify({ type: 'user', message: {} }), state)).toEqual([]);
    expect(parseStreamJsonLine(JSON.stringify({ type: 'whatever' }), state)).toEqual([]);
  });

  it('assembles a tool call across input_json_delta chunks (split across lines)', () => {
    const state = makeStreamJsonState();
    // start → tool_call (running, empty args)
    const start = parseStreamJsonLine(toolStart(1, 'toolu_1', 'edit_file'), state);
    expect(start).toHaveLength(1);
    expect(start[0]!.type).toBe('tool_call');
    const startSnap = (start[0] as { type: 'tool_call'; toolCall: AcpToolSnapshot }).toolCall;
    expect(startSnap.toolCallId).toBe('toolu_1');
    expect(startSnap.title).toBe('edit_file');
    expect(startSnap.status).toBe('in_progress');
    expect(startSnap.rawInput).toEqual({});

    // args streamed in fragments — no events until stop
    expect(parseStreamJsonLine(inputJsonDelta(1, '{"path":"a'), state)).toEqual([]);
    expect(parseStreamJsonLine(inputJsonDelta(1, '.ts","content":'), state)).toEqual([]);
    expect(parseStreamJsonLine(inputJsonDelta(1, '"hi"}'), state)).toEqual([]);

    // stop → tool_update with the parsed, fully-assembled input
    const stop = parseStreamJsonLine(blockStop(1), state);
    expect(stop).toHaveLength(1);
    expect(stop[0]!.type).toBe('tool_update');
    const stopSnap = (stop[0] as { type: 'tool_update'; toolCall: AcpToolSnapshot }).toolCall;
    expect(stopSnap.toolCallId).toBe('toolu_1');
    expect(stopSnap.status).toBe('completed');
    expect(stopSnap.rawInput).toEqual({ path: 'a.ts', content: 'hi' });
  });

  it('falls back to {_raw} when accumulated tool args are not valid JSON', () => {
    const state = makeStreamJsonState();
    parseStreamJsonLine(toolStart(0, 'toolu_x', 'run'), state);
    parseStreamJsonLine(inputJsonDelta(0, '{"broken'), state);
    const stop = parseStreamJsonLine(blockStop(0), state);
    const snap = (stop[0] as { type: 'tool_update'; toolCall: AcpToolSnapshot }).toolCall;
    expect(snap.rawInput).toEqual({ _raw: '{"broken' });
  });

  it('captures usage from message_delta and result lines', () => {
    const state = makeStreamJsonState();
    parseStreamJsonLine(streamEvent({ type: 'message_delta', usage: { output_tokens: 42 } }), state);
    expect(state.usage.outputTokens).toBe(42);
    parseStreamJsonLine(resultLine(100, 250, 'sess-z'), state);
    expect(state.usage.inputTokens).toBe(100);
    expect(state.usage.outputTokens).toBe(250);
    expect(state.sessionId).toBe('sess-z');
  });

  it('maps a terminal assistant message (fallback) → text + reasoning + tool events', () => {
    const state = makeStreamJsonState();
    const line = JSON.stringify({
      type: 'assistant',
      session_id: 'sess-asst',
      message: {
        content: [
          { type: 'thinking', thinking: 'let me think' },
          { type: 'text', text: 'Here is the answer' },
          { type: 'tool_use', id: 'toolu_9', name: 'view_file', input: { path: 'x.ts' } },
        ],
        usage: { input_tokens: 5, output_tokens: 7 },
      },
    });
    const events = parseStreamJsonLine(line, state);
    expect(events).toEqual([
      { type: 'reasoning', text: 'let me think' },
      { type: 'text', text: 'Here is the answer' },
      {
        type: 'tool_update',
        toolCall: { toolCallId: 'toolu_9', title: 'view_file', kind: null, status: 'completed', rawInput: { path: 'x.ts' } },
      },
    ]);
    expect(state.usage).toEqual({ inputTokens: 5, outputTokens: 7 });
    expect(state.sessionId).toBe('sess-asst');
  });
});

describe('makeStreamJsonParser (stateful wrapper over a full turn)', () => {
  it('streams a representative turn: init → text → thinking → tool → result', () => {
    const parser = makeStreamJsonParser();
    const all: AgentEvent[] = [];
    const feed = (line: string): AgentEventList => {
      const evs = parser.push(line);
      all.push(...evs);
      return evs;
    };

    feed(sys('sess-1'));
    feed(textDelta(0, 'Reading '));
    feed(textDelta(0, 'the file. '));
    feed(thinkingDelta(0, 'I should edit it'));
    feed(toolStart(1, 'toolu_a', 'edit_file'));
    feed(inputJsonDelta(1, '{"path":'));
    feed(inputJsonDelta(1, '"main.ts"}'));
    feed(blockStop(1));
    feed(textDelta(0, 'Done.'));
    feed(resultLine(120, 80, 'sess-1'));

    expect(all).toEqual([
      { type: 'text', text: 'Reading ' },
      { type: 'text', text: 'the file. ' },
      { type: 'reasoning', text: 'I should edit it' },
      {
        type: 'tool_call',
        toolCall: { toolCallId: 'toolu_a', title: 'edit_file', kind: null, status: 'in_progress', rawInput: {} },
      },
      {
        type: 'tool_update',
        toolCall: { toolCallId: 'toolu_a', title: 'edit_file', kind: null, status: 'completed', rawInput: { path: 'main.ts' } },
      },
      { type: 'text', text: 'Done.' },
    ]);

    expect(parser.usage()).toEqual({ inputTokens: 120, outputTokens: 80 });
    expect(parser.sessionId()).toBe('sess-1');
  });

  it('a garbage line interleaved mid-turn does not derail subsequent parsing', () => {
    const parser = makeStreamJsonParser();
    expect(parser.push(textDelta(0, 'a'))).toEqual([{ type: 'text', text: 'a' }]);
    expect(parser.push('>>> not json <<<')).toEqual([]);
    expect(parser.push(textDelta(0, 'b'))).toEqual([{ type: 'text', text: 'b' }]);
  });
});