d3ro-voice/apps/desktop/tests/main/services/local-llm-service.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { ErrorCode } from '@d3ro/core/errors'
import { initInMemoryConfig, resetInMemoryConfig } from '../../../src/main/services/ConfigService'
import {
  getLocalLLMService,
  resetLocalLLMServiceForTests,
} from '../../../src/main/services/LocalLLMService'

const encoder = new TextEncoder()

function completedStream(lines: string[]): ReadableStream<Uint8Array> {
  return new ReadableStream({
    start(controller) {
      controller.enqueue(encoder.encode(lines.join('\n')))
      controller.close()
    },
  })
}

function completedGenerateResponse(): Response {
  return new Response(JSON.stringify({
    model: 'gemma4:e4b',
    response: 'ok',
    done: true,
  }), { status: 200 })
}

function markAvailable(): void {
  const service = getLocalLLMService() as unknown as { _available: boolean }
  service._available = true
}

describe('LocalLLMService request lifecycle', () => {
  beforeEach(() => {
    initInMemoryConfig()
    resetLocalLLMServiceForTests()
    markAvailable()
  })

  afterEach(() => {
    resetLocalLLMServiceForTests()
    resetInMemoryConfig()
    vi.unstubAllGlobals()
    vi.useRealTimers()
  })

  it('sends the bounded default num_predict to chat requests', async () => {
    const cancel = vi.fn()
    const fetchMock = vi.fn(async (_input: RequestInfo | URL, init?: RequestInit) => {
      const body = JSON.parse(String(init?.body)) as { options: { num_predict: number } }
      expect(body.options.num_predict).toBe(512)
      return new Response(new ReadableStream<Uint8Array>({
        start(controller) {
          controller.enqueue(encoder.encode([
            JSON.stringify({ message: { content: 'hello' }, done: false }),
            JSON.stringify({ message: { content: '' }, done: true }),
          ].join('\n') + '\n'))
        },
        cancel,
      }), { status: 200 })
    })
    vi.stubGlobal('fetch', fetchMock)

    const output: string[] = []
    for await (const token of getLocalLLMService().chatStream([{ role: 'user', content: 'hello' }])) {
      output.push(token)
    }

    expect(output).toEqual(['hello'])
    expect(fetchMock).toHaveBeenCalledTimes(1)
    expect(cancel).toHaveBeenCalledTimes(1)
  })

  it('propagates external cancellation and request deadlines to fetch', async () => {
    const signals: AbortSignal[] = []
    const fetchMock = vi.fn((_input: RequestInfo | URL, init?: RequestInit) => new Promise<Response>((_resolve, reject) => {
      const signal = init?.signal as AbortSignal
      signals.push(signal)
      signal.addEventListener('abort', () => reject(new DOMException('Aborted', 'AbortError')), { once: true })
    }))
    vi.stubGlobal('fetch', fetchMock)

    const external = new AbortController()
    const externalRequest = getLocalLLMService().generate('first', { signal: external.signal })
    const externalExpectation = expect(externalRequest).rejects.toMatchObject({ code: ErrorCode.LLMProcessingCancelled })
    external.abort()
    await externalExpectation
    expect(signals[0].aborted).toBe(true)

    vi.useFakeTimers()
    const timeoutRequest = getLocalLLMService().generate('second', { timeoutMs: 25 })
    const timeoutExpectation = expect(timeoutRequest).rejects.toMatchObject({ code: ErrorCode.LLMProcessingTimeout })
    await vi.advanceTimersByTimeAsync(25)
    await timeoutExpectation
    expect(signals[1].aborted).toBe(true)
  })

  it('keeps other active requests cancellable after one request completes', async () => {
    const signals: AbortSignal[] = []
    const fetchMock = vi.fn((_input: RequestInfo | URL, init?: RequestInit) => {
      const signal = init?.signal as AbortSignal
      signals.push(signal)
      if (signals.length === 1) return Promise.resolve(completedGenerateResponse())
      return new Promise<Response>((_resolve, reject) => {
        signal.addEventListener('abort', () => reject(new DOMException('Aborted', 'AbortError')), { once: true })
      })
    })
    vi.stubGlobal('fetch', fetchMock)

    const first = getLocalLLMService().generate('first')
    const second = getLocalLLMService().generate('second')
    const third = getLocalLLMService().generate('third')
    await expect(first).resolves.toMatchObject({ text: 'ok' })

    const secondExpectation = expect(second).rejects.toMatchObject({ code: ErrorCode.LLMProcessingCancelled })
    const thirdExpectation = expect(third).rejects.toMatchObject({ code: ErrorCode.LLMProcessingCancelled })
    getLocalLLMService().cancelGeneration()

    await secondExpectation
    await thirdExpectation
    expect(signals[1].aborted).toBe(true)
    expect(signals[2].aborted).toBe(true)
  })

  it('rejects a stream that reaches EOF without a done frame', async () => {
    vi.stubGlobal('fetch', vi.fn(async () => new Response(completedStream([
      JSON.stringify({ model: 'gemma4:e4b', response: 'partial', done: false }),
    ]), { status: 200 })))

    const stream = getLocalLLMService().streamGenerate('hello')
    await expect(stream.next()).resolves.toMatchObject({ value: 'partial', done: false })
    await expect(stream.next()).rejects.toMatchObject({ code: ErrorCode.LLMProcessingFailed })
  })

  it('finishes on a done frame without waiting for the server to close the stream', async () => {
    const cancel = vi.fn()
    const body = new ReadableStream<Uint8Array>({
      start(controller) {
        controller.enqueue(encoder.encode([
          JSON.stringify({ model: 'gemma4:e4b', response: 'partial', done: false }),
          JSON.stringify({ model: 'gemma4:e4b', response: '', done: true }),
        ].join('\n') + '\n'))
      },
      cancel,
    })
    vi.stubGlobal('fetch', vi.fn(async () => new Response(body, { status: 200 })))

    const stream = getLocalLLMService().streamGenerate('hello', { timeoutMs: 500 })
    await expect(stream.next()).resolves.toMatchObject({ value: 'partial', done: false })
    await expect(stream.next()).resolves.toMatchObject({ value: expect.objectContaining({ text: 'partial' }), done: true })
    expect(cancel).toHaveBeenCalledTimes(1)
  })

  it('deduplicates concurrent ensureRunning calls and polling startup', async () => {
    const service = getLocalLLMService()
    const privateService = service as unknown as {
      _ensureRunning: () => Promise<'running' | 'starting' | 'not-installed' | 'failed'>
      _checkAvailability: () => Promise<void>
    }
    let release: ((value: 'running') => void) | undefined
    const pending = new Promise<'running'>((resolve) => {
      release = resolve
    })
    const ensure = vi.spyOn(privateService, '_ensureRunning').mockReturnValue(pending)

    const first = service.ensureRunning()
    const second = service.ensureRunning()
    expect(ensure).toHaveBeenCalledTimes(1)
    release?.('running')
    await expect(Promise.all([first, second])).resolves.toEqual(['running', 'running'])

    vi.useFakeTimers()
    const availability = vi.spyOn(privateService, '_checkAvailability').mockResolvedValue()
    service.startPolling()
    service.startPolling()
    expect(availability).toHaveBeenCalledTimes(1)
    await vi.advanceTimersByTimeAsync(5000)
    expect(availability).toHaveBeenCalledTimes(2)
    service.stopPolling()
  })
})