discovery-by-llm/frontend/tests/recommendationStream.test.ts

import { afterEach, describe, expect, it, vi } from 'vitest'
import { streamRecommendations, StreamTransportError } from '../src/lib/recommendationStream'
import type { StreamEvent } from '../src/lib/types'

const encoder = new TextEncoder()

function streamResponse(
  chunks: Uint8Array[],
  options: { close?: boolean; cancel?: () => void } = {},
): Response {
  let index = 0
  const body = new ReadableStream<Uint8Array>({
    pull(controller) {
      const chunk = chunks[index]
      if (chunk) {
        controller.enqueue(chunk)
        index += 1
      }
      if (options.close !== false && index === chunks.length) controller.close()
    },
    cancel() {
      options.cancel?.()
    },
  })
  return new Response(body, {
    headers: { 'Content-Type': 'application/x-ndjson; charset=utf-8' },
  })
}

function stubResponse(response: Response): void {
  vi.stubGlobal(
    'fetch',
    vi.fn(async () => response),
  )
}

afterEach(() => vi.unstubAllGlobals())

describe('streamRecommendations', () => {
  it('decodes fragmented NDJSON across JSON and multi-byte boundaries', async () => {
    const symbol = String.fromCodePoint(0x1f3b5)
    const records = [
      {
        type: 'metadata',
        request_id: 'request-1',
        intent_summary: `Focused ${symbol} listening`,
        candidate_count: 4,
      },
      {
        type: 'track',
        rank: 1,
        track: {
          id: 'track-1',
          uri: 'spotify:track:1',
          title: `Signal ${symbol}`,
          artists: ['Artist'],
          album_name: 'Album',
          album_art_url: null,
          external_url: 'https://open.spotify.com/track/1',
        },
        justification: 'It fits.',
      },
      { type: 'warning', code: 'limited_pool', message: 'The pool was limited.' },
      { type: 'warning', code: 'limited_pool', message: 'The pool stayed limited.' },
      { type: 'done', track_count: 1, total_ms: 12 },
    ]
    const payload = `${records.map((record) => JSON.stringify(record)).join('\n')}\n`
    const bytes = encoder.encode(payload)
    const jsonCut = encoder.encode(payload.slice(0, payload.indexOf('candidate_count') + 5)).length
    const symbolIndex = payload.indexOf(symbol)
    const symbolCut = encoder.encode(payload.slice(0, symbolIndex)).length + 2
    const cuts = [jsonCut, symbolCut, bytes.length - 8].sort((left, right) => left - right)
    const chunks = [
      bytes.slice(0, cuts[0]),
      bytes.slice(cuts[0], cuts[1]),
      bytes.slice(cuts[1], cuts[2]),
      bytes.slice(cuts[2]),
    ]
    stubResponse(streamResponse(chunks))
    const events: StreamEvent[] = []

    await streamRecommendations(
      { schema_version: 1, query: 'focus', history: [], prior_recommendations: [] },
      new AbortController().signal,
      (event) => events.push(event),
    )

    expect(events).toHaveLength(5)
    expect(events[0]).toMatchObject({
      type: 'metadata',
      intent_summary: `Focused ${symbol} listening`,
    })
    expect(events[1]).toMatchObject({ type: 'track', track: { title: `Signal ${symbol}` } })
    expect(events.filter((event) => event.type === 'warning')).toHaveLength(2)
    expect(events.at(-1)).toMatchObject({ type: 'done', track_count: 1 })
  })

  it('treats an error as terminal without waiting for done', async () => {
    let wasCancelled = false
    const payload = [
      JSON.stringify({
        type: 'metadata',
        request_id: 'request-1',
        intent_summary: 'Intent',
        candidate_count: 0,
      }),
      JSON.stringify({ type: 'error', code: 'upstream', message: 'Upstream failed.' }),
      JSON.stringify({ type: 'done', track_count: 0, total_ms: 4 }),
    ].join('\n')
    stubResponse(
      streamResponse([encoder.encode(payload)], {
        close: false,
        cancel: () => {
          wasCancelled = true
          throw new Error('Cleanup failed.')
        },
      }),
    )
    const events: StreamEvent[] = []

    await streamRecommendations(
      { schema_version: 1, query: 'focus', history: [], prior_recommendations: [] },
      new AbortController().signal,
      (event) => events.push(event),
    )

    expect(events.map((event) => event.type)).toEqual(['metadata', 'error'])
    expect(wasCancelled).toBe(true)
  })

  it.each([
    ['an event before metadata', [{ type: 'warning', code: 'early', message: 'No metadata.' }]],
    [
      'duplicate metadata',
      [
        { type: 'metadata', request_id: 'one', intent_summary: 'First', candidate_count: 1 },
        { type: 'metadata', request_id: 'two', intent_summary: 'Second', candidate_count: 1 },
      ],
    ],
  ])('rejects %s and cancels the reader', async (_label, records) => {
    let wasCancelled = false
    const payload = `${records.map((record) => JSON.stringify(record)).join('\n')}\n`
    stubResponse(
      streamResponse([encoder.encode(payload)], {
        close: false,
        cancel: () => {
          wasCancelled = true
        },
      }),
    )

    const result = streamRecommendations(
      { schema_version: 1, query: 'focus', history: [], prior_recommendations: [] },
      new AbortController().signal,
      () => undefined,
    )

    await expect(result).rejects.toMatchObject<Partial<StreamTransportError>>({ kind: 'protocol' })
    expect(wasCancelled).toBe(true)
  })

  it('cancels the reader after a parse failure', async () => {
    let wasCancelled = false
    stubResponse(
      streamResponse([encoder.encode('{not-json}\n')], {
        close: false,
        cancel: () => {
          wasCancelled = true
        },
      }),
    )

    const result = streamRecommendations(
      { schema_version: 1, query: 'focus', history: [], prior_recommendations: [] },
      new AbortController().signal,
      () => undefined,
    )

    await expect(result).rejects.toMatchObject<Partial<StreamTransportError>>({ kind: 'parse' })
    expect(wasCancelled).toBe(true)
  })

  it('reports a transport failure when the body ends before a terminal event', async () => {
    const payload = `${JSON.stringify({
      type: 'metadata',
      request_id: 'request-1',
      intent_summary: 'Intent',
      candidate_count: 0,
    })}\n`
    stubResponse(streamResponse([encoder.encode(payload)]))

    const result = streamRecommendations(
      { schema_version: 1, query: 'focus', history: [], prior_recommendations: [] },
      new AbortController().signal,
      () => undefined,
    )

    await expect(result).rejects.toMatchObject<Partial<StreamTransportError>>({
      kind: 'unexpected_eof',
    })
  })
})