boocode/apps/coder/src/services/backends/__tests__/claude-sdk-map.test.ts

import { describe, it, expect } from 'vitest';
import type { SDKMessage } from '@anthropic-ai/claude-agent-sdk';
import { mapSdkMessage, createClaudeSdkMapState } from '../claude-sdk-map.js';
import type { AgentEvent } from '../../agent-backend.js';

/**
 * Pure mapper for Claude-SDK messages → AgentEvents (claude-sdk-sessionstore #9 Part 2).
 * Verifies the partial-stream → live-delta mapping, tool assembly across blocks, and
 * the final-assistant dedup, with no live `claude` binary involved.
 *
 * Messages are cast through `unknown` to `SDKMessage`: the real SDK shapes carry many
 * fields (uuid, parent_tool_use_id, …) irrelevant to the mapper, which reads only the
 * `type`/`event`/`message.content` it discriminates on. The cast keeps the fixtures
 * minimal while the production code path sees the full real types (the backend's
 * typecheck against the real SDK is the type-safety proof).
 */
function msg(m: unknown): SDKMessage {
  return m as SDKMessage;
}

/** A partial-stream message wrapping one BetaRawMessageStreamEvent. */
function streamEvent(event: unknown): SDKMessage {
  return msg({ type: 'stream_event', event, parent_tool_use_id: null, uuid: 'u', session_id: 's' });
}

describe('mapSdkMessage — partial stream deltas', () => {
  it('maps a text_delta to a text event', () => {
    const state = createClaudeSdkMapState();
    const out = mapSdkMessage(
      streamEvent({ type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'Hello' } }),
      state,
    );
    expect(out).toEqual<AgentEvent[]>([{ type: 'text', text: 'Hello' }]);
  });

  it('maps a thinking_delta to a reasoning event', () => {
    const state = createClaudeSdkMapState();
    const out = mapSdkMessage(
      streamEvent({
        type: 'content_block_delta',
        index: 0,
        delta: { type: 'thinking_delta', thinking: 'pondering', estimated_tokens: null },
      }),
      state,
    );
    expect(out).toEqual<AgentEvent[]>([{ type: 'reasoning', text: 'pondering' }]);
  });

  it('drops empty text/thinking deltas', () => {
    const state = createClaudeSdkMapState();
    expect(
      mapSdkMessage(streamEvent({ type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: '' } }), state),
    ).toEqual([]);
    expect(
      mapSdkMessage(
        streamEvent({ type: 'content_block_delta', index: 0, delta: { type: 'thinking_delta', thinking: '', estimated_tokens: null } }),
        state,
      ),
    ).toEqual([]);
  });

  it('ignores message framing + signature/citation deltas', () => {
    const state = createClaudeSdkMapState();
    expect(mapSdkMessage(streamEvent({ type: 'message_start', message: {} }), state)).toEqual([]);
    expect(mapSdkMessage(streamEvent({ type: 'message_stop' }), state)).toEqual([]);
    expect(
      mapSdkMessage(streamEvent({ type: 'content_block_delta', index: 0, delta: { type: 'signature_delta', signature: 'x' } }), state),
    ).toEqual([]);
  });
});

describe('mapSdkMessage — tool assembly across blocks', () => {
  it('opens a tool_call on content_block_start, buffers input_json_delta, emits tool_update with parsed input on stop', () => {
    const state = createClaudeSdkMapState();

    const started = mapSdkMessage(
      streamEvent({
        type: 'content_block_start',
        index: 1,
        content_block: { type: 'tool_use', id: 'tool-1', name: 'view_file', input: {} },
      }),
      state,
    );
    expect(started).toEqual<AgentEvent[]>([
      { type: 'tool_call', toolCall: { toolCallId: 'tool-1', title: 'view_file', kind: null, status: 'in_progress', rawInput: {}, rawOutput: undefined } },
    ]);

    // args stream in fragments under the same block index
    expect(
      mapSdkMessage(streamEvent({ type: 'content_block_delta', index: 1, delta: { type: 'input_json_delta', partial_json: '{"path":' } }), state),
    ).toEqual([]);
    expect(
      mapSdkMessage(streamEvent({ type: 'content_block_delta', index: 1, delta: { type: 'input_json_delta', partial_json: '"a.ts"}' } }), state),
    ).toEqual([]);

    const stopped = mapSdkMessage(streamEvent({ type: 'content_block_stop', index: 1 }), state);
    expect(stopped).toHaveLength(1);
    const ev = stopped[0]!;
    expect(ev.type).toBe('tool_update');
    if (ev.type === 'tool_update') {
      expect(ev.toolCall.toolCallId).toBe('tool-1');
      expect(ev.toolCall.title).toBe('view_file');
      expect(ev.toolCall.rawInput).toEqual({ path: 'a.ts' });
    }
  });

  it('content_block_stop for a non-tool block (no tracked index) emits nothing', () => {
    const state = createClaudeSdkMapState();
    // text block was streamed at index 0 but never tracked as a tool
    mapSdkMessage(streamEvent({ type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'hi' } }), state);
    expect(mapSdkMessage(streamEvent({ type: 'content_block_stop', index: 0 }), state)).toEqual([]);
  });

  it('falls back to the prior input when the buffered tool JSON is invalid', () => {
    const state = createClaudeSdkMapState();
    mapSdkMessage(
      streamEvent({ type: 'content_block_start', index: 2, content_block: { type: 'tool_use', id: 't2', name: 'grep', input: { q: 'seed' } } }),
      state,
    );
    mapSdkMessage(streamEvent({ type: 'content_block_delta', index: 2, delta: { type: 'input_json_delta', partial_json: '{not json' } }), state);
    const stopped = mapSdkMessage(streamEvent({ type: 'content_block_stop', index: 2 }), state);
    const ev = stopped[0]!;
    if (ev.type === 'tool_update') {
      expect(ev.toolCall.rawInput).toEqual({ q: 'seed' });
    } else {
      throw new Error('expected tool_update');
    }
  });
});

describe('mapSdkMessage — final assistant message', () => {
  function assistant(content: unknown[]): SDKMessage {
    return msg({ type: 'assistant', message: { content }, parent_tool_use_id: null, uuid: 'u', session_id: 's' });
  }

  it('dedups text/thinking (already streamed) and emits a completed tool_update per tool_use block', () => {
    const state = createClaudeSdkMapState();
    const out = mapSdkMessage(
      assistant([
        { type: 'text', text: 'final answer', citations: null },
        { type: 'thinking', thinking: 'reasoned', signature: 'sig' },
        { type: 'tool_use', id: 'tool-9', name: 'find_files', input: { glob: '**/*.ts' } },
      ]),
      state,
    );
    expect(out).toEqual<AgentEvent[]>([
      {
        type: 'tool_update',
        toolCall: { toolCallId: 'tool-9', title: 'find_files', kind: null, status: 'completed', rawInput: { glob: '**/*.ts' }, rawOutput: undefined },
      },
    ]);
  });

  it('preserves a title from a prior partial tool_call snapshot', () => {
    const state = createClaudeSdkMapState();
    mapSdkMessage(
      streamEvent({ type: 'content_block_start', index: 0, content_block: { type: 'tool_use', id: 'tool-x', name: 'view_file', input: {} } }),
      state,
    );
    const out = mapSdkMessage(assistant([{ type: 'tool_use', id: 'tool-x', name: 'view_file', input: { path: 'z' } }]), state);
    const ev = out[0]!;
    if (ev.type === 'tool_update') {
      expect(ev.toolCall.status).toBe('completed');
      expect(ev.toolCall.title).toBe('view_file');
      expect(ev.toolCall.rawInput).toEqual({ path: 'z' });
    } else {
      throw new Error('expected tool_update');
    }
  });
});

describe('mapSdkMessage — non-content messages', () => {
  it('returns [] for system/init, status, result, and other variants', () => {
    const state = createClaudeSdkMapState();
    expect(mapSdkMessage(msg({ type: 'system', subtype: 'init', session_id: 's', uuid: 'u' }), state)).toEqual([]);
    expect(mapSdkMessage(msg({ type: 'system', subtype: 'status', status: null, session_id: 's', uuid: 'u' }), state)).toEqual([]);
    expect(
      mapSdkMessage(msg({ type: 'result', subtype: 'success', result: 'done', session_id: 's', uuid: 'u' }), state),
    ).toEqual([]);
  });
});

describe('mapSdkMessage — user tool results', () => {
  /** A `user` message carrying tool_result blocks (the SDK feeds tool output back here). */
  function userMsg(content: unknown): SDKMessage {
    return msg({ type: 'user', message: { role: 'user', content }, parent_tool_use_id: null, uuid: 'u', session_id: 's' });
  }

  it('maps a string tool_result to a completed tool_update carrying the output', () => {
    const state = createClaudeSdkMapState();
    const out = mapSdkMessage(userMsg([{ type: 'tool_result', tool_use_id: 't1', content: 'done' }]), state);
    expect(out).toEqual<AgentEvent[]>([
      {
        type: 'tool_update',
        toolCall: { toolCallId: 't1', title: 't1', kind: null, status: 'completed', rawInput: undefined, rawOutput: 'done' },
      },
    ]);
  });

  it('marks an is_error result failed', () => {
    const state = createClaudeSdkMapState();
    const out = mapSdkMessage(userMsg([{ type: 'tool_result', tool_use_id: 't1', content: 'boom', is_error: true }]), state);
    const ev = out[0]!;
    if (ev.type !== 'tool_update') throw new Error('expected tool_update');
    expect(ev.toolCall.status).toBe('failed');
    expect(ev.toolCall.rawOutput).toBe('boom');
  });

  it('flattens array text blocks (skipping non-text) and reuses a prior snapshot title', () => {
    const state = createClaudeSdkMapState();
    mapSdkMessage(
      streamEvent({ type: 'content_block_start', index: 1, content_block: { type: 'tool_use', id: 't2', name: 'view_file', input: {} } }),
      state,
    );
    const out = mapSdkMessage(
      userMsg([
        {
          type: 'tool_result',
          tool_use_id: 't2',
          content: [
            { type: 'text', text: 'line1' },
            { type: 'image', source: {} },
            { type: 'text', text: 'line2' },
          ],
        },
      ]),
      state,
    );
    const ev = out[0]!;
    if (ev.type !== 'tool_update') throw new Error('expected tool_update');
    expect(ev.toolCall.toolCallId).toBe('t2');
    expect(ev.toolCall.title).toBe('view_file');
    expect(ev.toolCall.status).toBe('completed');
    expect(ev.toolCall.rawOutput).toBe('line1\nline2');
  });

  it('surfaces a result for an unknown tool_use_id with the id as the title', () => {
    const state = createClaudeSdkMapState();
    const out = mapSdkMessage(userMsg([{ type: 'tool_result', tool_use_id: 'orphan-id', content: 'x' }]), state);
    expect(out[0]).toMatchObject({
      type: 'tool_update',
      toolCall: { toolCallId: 'orphan-id', title: 'orphan-id', kind: null, status: 'completed' },
    });
  });

  it('ignores non-tool_result blocks and non-array content', () => {
    const state = createClaudeSdkMapState();
    expect(mapSdkMessage(userMsg([{ type: 'text', text: 'hi' }]), state)).toEqual([]);
    expect(mapSdkMessage(userMsg('plain string'), state)).toEqual([]);
  });
});