src/lib/message-segments.test.ts

import { describe, expect, it } from 'vitest';
import type { ChatMessage, MessageSegment } from './contracts';
import { appendToolSegment, normalizeMessageSegments, segmentsForMessage, updateRoundSegments } from './message-segments';

describe('ordered message segments', () => {
  it('preserves the first observed order across reasoning, text, tools, and later text', () => {
    let segments: MessageSegment[] = [];
    segments = updateRoundSegments(segments, 'round-0', 'Checking inputs', '');
    segments = updateRoundSegments(segments, 'round-0', 'Checking inputs', 'I will calculate it.');
    segments = appendToolSegment(segments, 'call-1', 'round-0');
    segments = updateRoundSegments(segments, 'round-1', 'The result is valid', 'The answer is 437.');
    expect(segments.map((segment) => segment.kind)).toEqual([
      'reasoning', 'text', 'tool', 'reasoning', 'text',
    ]);
    expect(segments[2]).toMatchObject({ roundId: 'round-0' });
  });

  it('removes streamed textual tool markup when the final visible text is empty', () => {
    const streamed = updateRoundSegments([], 'round-0', '', '<tool_call>');
    expect(updateRoundSegments(streamed, 'round-0', '', '')).toEqual([]);
  });

  it('normalizes legacy artifacts next to their producing tool without duplicating them', () => {
    const message: ChatMessage = {
      id: 'assistant',
      role: 'assistant',
      content: 'Before tool',
      reasoning: 'Think',
      timestamp: '10:00',
      tools: [{ id: 'artifact-1', name: 'bar_chart', input: '{}', state: 'complete' }],
      artifacts: [{
        id: 'artifact-1',
        title: 'Chart',
        source: '<main>Chart</main>',
        sandboxedSource: '<main>Chart</main>',
        runtimeToken: 'runtime',
        createdAt: 1,
      }],
    };
    const normalized = normalizeMessageSegments(message);
    expect(normalized.tools?.[0]?.artifactId).toBe('artifact-1');
    expect(segmentsForMessage(normalized).map((segment) => segment.kind)).toEqual([
      'reasoning', 'text', 'tool',
    ]);
  });
});