deepseek-harness/packages/llm/tests/assembler.spec.ts

import { describe, expect, it } from 'vitest'
import { BlockAssembler, CallId, type StreamChunk } from '@deepseek-ai/dsh-llm'

describe('BlockAssembler', () => {
  it('assembles interleaved text, reasoning, and tool-call deltas', () => {
    const chunks: StreamChunk[] = [
      { type: 'block-start', index: 0, blockType: 'reasoning' },
      { type: 'reasoning-delta', index: 0, text: 'thinking…' },
      { type: 'block-end', index: 0, block: { type: 'reasoning', text: 'thinking…' } },
      { type: 'block-start', index: 1, blockType: 'text' },
      { type: 'text-delta', index: 1, text: 'Hello' },
      { type: 'text-delta', index: 1, text: ' world' },
      { type: 'block-start', index: 2, blockType: 'tool-call' },
      { type: 'tool-call-delta', index: 2, id: CallId('call-1'), name: 'echo', argumentsDelta: '{"text":' },
      { type: 'tool-call-delta', index: 2, id: CallId('call-1'), argumentsDelta: '"hi"}' },
      { type: 'usage', usage: { inputTokens: 10, outputTokens: 5 } },
      { type: 'finish', reason: { kind: 'tool-calls' } },
    ]
    const assembler = new BlockAssembler()
    for (const chunk of chunks) assembler.push(chunk)

    expect(assembler.blocks()).toEqual([
      { type: 'reasoning', text: 'thinking…' },
      { type: 'text', text: 'Hello world' },
      { type: 'tool-call', id: CallId('call-1'), name: 'echo', arguments: '{"text":"hi"}' },
    ])
    expect(assembler.usage).toEqual({ inputTokens: 10, outputTokens: 5 })
    expect(assembler.finish).toEqual({ kind: 'tool-calls' })
    expect(assembler.message().role).toBe('assistant')
  })

  it('returns the completed block from push() on block-end', () => {
    const assembler = new BlockAssembler()
    expect(assembler.push({ type: 'block-start', index: 0, blockType: 'text' })).toBeUndefined()
    expect(assembler.push({ type: 'text-delta', index: 0, text: 'hi' })).toBeUndefined()
    const block = assembler.push({ type: 'block-end', index: 0, block: { type: 'text', text: 'hi' } })
    expect(block).toEqual({ type: 'text', text: 'hi' })
  })

  it('tolerates deltas without explicit block-start/end', () => {
    const assembler = new BlockAssembler()
    assembler.push({ type: 'text-delta', index: 0, text: 'implicit' })
    expect(assembler.blocks()).toEqual([{ type: 'text', text: 'implicit' }])
    expect(assembler.finish).toEqual({ kind: 'stop' })
  })

  it('returns undefined usage when no usage chunk was received', () => {
    const assembler = new BlockAssembler()
    assembler.push({ type: 'text-delta', index: 0, text: 'no usage' })
    expect(assembler.usage).toBeUndefined()
  })

  it('reuses an existing partial when ensure() is called with a tracked index', () => {
    const assembler = new BlockAssembler()
    // block-start creates the partial; block-end calls ensure() on the same index
    assembler.push({ type: 'block-start', index: 0, blockType: 'text' })
    // push a delta first to guarantee the partial exists
    assembler.push({ type: 'text-delta', index: 0, text: 'hi' })
    // block-end's ensure() must find the existing partial (the second branch path)
    const block = assembler.push({ type: 'block-end', index: 0, block: { type: 'text', text: 'hi' } })
    expect(block).toEqual({ type: 'text', text: 'hi' })
  })

  it('throws from assemble() when a partial has an unhandled blockType', () => {
    const assembler = new BlockAssembler()
    // Directly push a block-end for an image block whose block-start never
    // called ensure — but the image block-type flows through normally.
    // What we really need is a partial whose blockType is not text/reasoning/tool-call.
    // We can achieve this via a block-start for 'image' followed by blocks().
    assembler.push({ type: 'block-start', index: 0, blockType: 'image' } as unknown as StreamChunk)
    expect(() => assembler.blocks()).toThrow('cannot assemble incomplete block of type "image"')
  })

  it('mustGet throws when an index is missing from the partials map (invariant violation)', () => {
    const assembler = new BlockAssembler()
    // Force the invariant violation: manually corrupt the data structures.
    /* eslint-disable */
    const hack = assembler as any
    hack.order.push(99)
    /* eslint-enable */
    expect(() => assembler.blocks()).toThrow('BlockAssembler invariant violated')
  })

  it('assembles open blocks at end of stream via flushRemaining', () => {
    const assembler = new BlockAssembler()
    assembler.push({ type: 'text-delta', index: 0, text: 'open' })
    assembler.push({ type: 'reasoning-delta', index: 1, text: 'thinking' })

    // flushReady returns nothing because index 0 is incomplete and blocking
    const ready = assembler.flushReady()
    expect(ready).toEqual([])

    // flushRemaining assembles everything still open
    const remaining = assembler.flushRemaining()
    expect(remaining).toEqual([
      { type: 'text', text: 'open' },
      { type: 'reasoning', text: 'thinking' },
    ])

    // blocks() now matches the flushed view
    expect(assembler.blocks()).toEqual(remaining)
  })

  it('result() omits usage key when no usage was received', () => {
    const assembler = new BlockAssembler()
    assembler.push({ type: 'text-delta', index: 0, text: 'msg' })
    const result = assembler.result()
    expect(result.message).toBeDefined()
    expect(result.finish).toEqual({ kind: 'stop' })
    // usage should NOT be present on the object at all
    expect('usage' in result).toBe(false)
  })

  it('ignores duplicate block-start for the same index', () => {
    const assembler = new BlockAssembler()
    assembler.push({ type: 'block-start', index: 0, blockType: 'text' })
    assembler.push({ type: 'text-delta', index: 0, text: 'one' })
    // duplicate block-start — should be no-op (false branch of has check)
    assembler.push({ type: 'block-start', index: 0, blockType: 'text' })
    assembler.push({ type: 'text-delta', index: 0, text: ' two' })
    assembler.push({ type: 'block-end', index: 0, block: { type: 'text', text: 'one two' } })
    expect(assembler.blocks()).toEqual([{ type: 'text', text: 'one two' }])
  })

  it('ignores tool-call-delta stragglers after block-end', () => {
    const assembler = new BlockAssembler()
    assembler.push({ type: 'block-start', index: 0, blockType: 'tool-call' })
    assembler.push({ type: 'tool-call-delta', index: 0, id: CallId('c1'), name: 'echo', argumentsDelta: '{}' })
    assembler.push({ type: 'block-end', index: 0, block: { type: 'tool-call', id: CallId('c1'), name: 'echo', arguments: '{}' } })
    // straggler after block-end — partial.block is set, so early return
    assembler.push({ type: 'tool-call-delta', index: 0, id: CallId('c1'), name: 'evil', argumentsDelta: 'oops' })
    expect(assembler.blocks()).toEqual([{ type: 'tool-call', id: CallId('c1'), name: 'echo', arguments: '{}' }])
  })

  it('assembles tool-call with generated id fallback when no id provided', () => {
    const assembler = new BlockAssembler()
    assembler.push({ type: 'tool-call-delta', index: 0, argumentsDelta: '{}' } as StreamChunk)
    // No id and no name provided — uses fallback id `call-{index}` and empty name
    const blocks = assembler.blocks()
    expect(blocks).toEqual([
      { type: 'tool-call', id: CallId('call-0'), name: '', arguments: '{}' },
    ])
  })

  it('includes usage in result() when usage was received', () => {
    const assembler = new BlockAssembler()
    assembler.push({ type: 'text-delta', index: 0, text: 'msg' })
    assembler.push({ type: 'usage', usage: { inputTokens: 5, outputTokens: 3 } })
    const result = assembler.result()
    expect(result.usage).toEqual({ inputTokens: 5, outputTokens: 3 })
    expect('usage' in result).toBe(true)
  })
})

describe('assertNever', () => {
  it('throws with diagnostics when a value escapes a closed union at runtime', async () => {
    const { assertNever } = await import('@deepseek-ai/dsh-llm')
    expect(() => assertNever({ type: 'rogue' } as never, 'test-context'))
      .toThrow('unreachable variant in test-context: {"type":"rogue"}')
    expect(() => assertNever(undefined as never)).toThrow('unreachable variant: undefined')
  })

  it('BlockAssembler.push rejects chunks outside the closed StreamChunk union', () => {
    const assembler = new BlockAssembler()
    expect(() => assembler.push({ type: 'rogue-chunk' } as unknown as StreamChunk))
      .toThrow('unreachable variant in BlockAssembler.push')
  })
})

describe('BlockAssembler regressions (property-test findings)', () => {
  it('first block-end wins: a duplicate block-end for a closed index is ignored', () => {
    // Found by fast-check (the property-testing RFC): two block-ends at the same index made the
    // streamed prefix (first block) disagree with final blocks() (second
    // block). The first close must win — same straggler rule as post-close
    // deltas — so streaming and one-shot assembly stay identical.
    const chunks: StreamChunk[] = [
      { type: 'block-end', index: 0, block: { type: 'reasoning', text: 'first' } },
      { type: 'block-end', index: 0, block: { type: 'text', text: 'second' } },
    ]
    const streaming = new BlockAssembler()
    const flushed = []
    for (const chunk of chunks) {
      streaming.push(chunk)
      flushed.push(...streaming.flushReady())
    }
    flushed.push(...streaming.flushRemaining())

    const oneShot = new BlockAssembler()
    for (const chunk of chunks) oneShot.push(chunk)

    expect(flushed).toEqual([{ type: 'reasoning', text: 'first' }])
    expect(oneShot.blocks()).toEqual([{ type: 'reasoning', text: 'first' }])
    expect(flushed).toEqual(oneShot.blocks())
  })

  it('push returns undefined for a duplicate block-end (it closed nothing)', () => {
    const a = new BlockAssembler()
    expect(a.push({ type: 'block-end', index: 0, block: { type: 'text', text: 'x' } }))
      .toEqual({ type: 'text', text: 'x' })
    expect(a.push({ type: 'block-end', index: 0, block: { type: 'text', text: 'y' } }))
      .toBeUndefined()
  })
})