d3ro-voice/apps/desktop/tests/main/services/VoiceConversationService.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { IPC_CHANNELS } from '@d3ro/core/ipc-channels'
import { ErrorCode } from '@d3ro/core/errors'
import {
  getVoiceConversationService,
  resetVoiceConversationServiceForTests,
} from '../../../src/main/services/VoiceConversationService'

type ChatOptions = {
  signal?: AbortSignal
  maxTokens?: number
  timeoutMs?: number
  keepAlive?: string
}

const mocks = vi.hoisted(() => ({
  audioStart: vi.fn(async () => undefined),
  audioStop: vi.fn(async () => undefined),
  chatStream: vi.fn(),
  premiumCancel: vi.fn(),
  speakSentences: vi.fn(async () => undefined),
  ttsStop: vi.fn(),
  send: vi.fn(),
}))

vi.mock('../../../src/main/services/LoggerService', () => ({
  getLogger: () => ({ info: vi.fn(), warn: vi.fn(), error: vi.fn() }),
}))

vi.mock('../../../src/main/services/AudioCaptureService', () => ({
  getAudioCaptureService: () => ({
    start: mocks.audioStart,
    stop: mocks.audioStop,
    on: vi.fn(),
    off: vi.fn(),
  }),
}))

vi.mock('../../../src/main/services/LocalSTTService', () => ({
  getLocalSTTService: () => ({ initialize: vi.fn(async () => undefined) }),
}))

vi.mock('../../../src/main/services/PremiumLLMService', () => ({
  getPremiumLLMService: () => ({
    cancelGeneration: mocks.premiumCancel,
    isAvailable: () => false,
  }),
}))

vi.mock('../../../src/main/services/LocalLLMService', () => ({
  getLocalLLMService: () => ({
    isAvailable: () => true,
    chatStream: mocks.chatStream,
  }),
}))

vi.mock('../../../src/main/services/TTSPlaybackService', () => ({
  getTTSPlaybackService: () => ({
    speakSentences: mocks.speakSentences,
    stop: mocks.ttsStop,
  }),
}))

vi.mock('../../../src/main/services/SoundEffectService', () => ({
  getSoundEffectService: () => ({ play: vi.fn() }),
}))

vi.mock('../../../src/main/services/ConfigService', () => ({
  configGet: () => undefined,
}))

vi.mock('../../../src/main/services/LicenseService', () => ({
  getLicenseService: () => ({
    canUse: () => ({ allowed: true }),
    promptUpgrade: vi.fn(),
  }),
}))

vi.mock('../../../src/main/windows/WindowManager', () => ({
  getMainWindow: () => ({
    isDestroyed: () => false,
    webContents: { send: mocks.send },
  }),
}))

function waitForAbort(signal: AbortSignal): Promise<void> {
  if (signal.aborted) return Promise.resolve()
  return new Promise((resolve) => signal.addEventListener('abort', () => resolve(), { once: true }))
}

async function* waitUntilAborted(options?: ChatOptions): AsyncGenerator<string, string> {
  if (!options?.signal) throw new Error('Expected caller-owned abort signal')
  await waitForAbort(options.signal)
  return ''
}

describe('VoiceConversationService local response ownership', () => {
  beforeEach(() => {
    resetVoiceConversationServiceForTests()
    vi.clearAllMocks()
    mocks.chatStream.mockImplementation((_messages: unknown, options?: ChatOptions) => waitUntilAborted(options))
  })

  afterEach(() => {
    resetVoiceConversationServiceForTests()
  })

  it('rejects a second send without starting a second local chat or adding duplicate history', async () => {
    const service = getVoiceConversationService()
    await service.startSession()

    const firstResponse = service.sendTextMessage('first')
    await vi.waitFor(() => expect(mocks.chatStream).toHaveBeenCalledTimes(1))

    await expect(service.sendTextMessage('second')).rejects.toMatchObject({
      code: ErrorCode.ConversationLLMFailed,
    })
    expect(mocks.chatStream).toHaveBeenCalledTimes(1)
    expect(service.getHistory().map((message) => message.content)).toEqual(['first'])

    service.stopSession()
    await firstResponse
  })

  it.each([
    ['stopSession', (service: ReturnType<typeof getVoiceConversationService>) => service.stopSession()],
    ['cancelResponse', (service: ReturnType<typeof getVoiceConversationService>) => service.cancelResponse()],
  ])('aborts the local stream signal when %s is called', async (_action, cancel) => {
    const service = getVoiceConversationService()
    await service.startSession()

    const response = service.sendTextMessage('cancel me')
    await vi.waitFor(() => expect(mocks.chatStream).toHaveBeenCalledTimes(1))
    const options = mocks.chatStream.mock.calls[0][1] as ChatOptions

    cancel(service)

    expect(options.signal?.aborted).toBe(true)
    await response
  })

  it('sets bounded local chat options for voice responses', async () => {
    const service = getVoiceConversationService()
    await service.startSession()

    const response = service.sendTextMessage('options')
    await vi.waitFor(() => expect(mocks.chatStream).toHaveBeenCalledTimes(1))

    expect(mocks.chatStream).toHaveBeenCalledWith(
      expect.any(Array),
      expect.objectContaining({
        signal: expect.any(AbortSignal),
        maxTokens: 512,
        timeoutMs: 60_000,
        keepAlive: '2m',
      }),
    )

    service.stopSession()
    await response
  })

  it('does not publish late assistant output or start TTS after cancellation', async () => {
    let releaseLateToken: (() => void) | undefined
    const lateToken = new Promise<void>((resolve) => {
      releaseLateToken = resolve
    })
    mocks.chatStream.mockImplementation(() => (async function* (): AsyncGenerator<string, string> {
      await lateToken
      yield 'late answer.'
      return 'late answer.'
    })())

    const service = getVoiceConversationService()
    await service.startSession()
    const response = service.sendTextMessage('cancel before response')
    await vi.waitFor(() => expect(mocks.chatStream).toHaveBeenCalledTimes(1))

    service.cancelResponse()
    mocks.send.mockClear()
    releaseLateToken?.()
    await response

    const lateConversationEvents = mocks.send.mock.calls.filter(([channel]) => (
      channel === IPC_CHANNELS.VOICE_CONVERSATION.ASSISTANT_DELTA
      || channel === IPC_CHANNELS.VOICE_CONVERSATION.ASSISTANT_MESSAGE
      || channel === IPC_CHANNELS.VOICE_CONVERSATION.TTS_STARTED
      || channel === IPC_CHANNELS.VOICE_CONVERSATION.TTS_FINISHED
    ))
    expect(lateConversationEvents).toEqual([])
    expect(mocks.speakSentences).not.toHaveBeenCalled()
    expect(service.getHistory().map((message) => message.content)).toEqual(['cancel before response'])
  })
})