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

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'

const config = vi.hoisted(() => ({
  suggestionEnabled: true,
  suggestionModelId: 'gemma4:e4b',
  llmModelId: 'gemma4:e4b',
  inputExcludedApps: [],
  suggestionTriggerDelayMs: 600,
  suggestionMinPrefixChars: 8,
  suggestionMaxRequestsPerMinute: 6,
  suggestionDailyBudget: 500,
  suggestionRequestTimeoutMs: 8000,
  inputLearnTypedText: false,
  inputTelemetryEnabled: false,
  suggestionOverlayInteractive: true
}))
const localLlm = vi.hoisted(() => ({
  isAvailable: vi.fn(() => true),
  streamGenerate: vi.fn()
}))

vi.mock('../../../src/main/services/ConfigService', () => ({
  configGet: vi.fn((key: keyof typeof config) => config[key]),
  configSet: vi.fn()
}))
vi.mock('../../../src/main/services/LoggerService', () => ({
  getLogger: () => ({ info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() })
}))
vi.mock('../../../src/main/services/LocalLLMService', () => ({
  getLocalLLMService: () => localLlm
}))
vi.mock('../../../src/main/services/InputTelemetryService', () => ({
  getInputTelemetryService: () => ({ listPhrases: vi.fn(() => []), recordExternalText: vi.fn() })
}))
vi.mock('../../../src/main/services/PersonalGraphService', () => ({
  getPersonalGraphService: () => ({ retrieveContext: vi.fn(() => ({ continuations: [], related: [] })) })
}))
vi.mock('../../../src/main/services/TextInsertService', () => ({
  getTextInsertService: () => ({ insertText: vi.fn() })
}))
vi.mock('../../../src/main/db', () => ({ getDatabase: vi.fn() }))
vi.mock('../../../src/main/db/schema', () => ({ suggestions: {} }))

const PREFIX = '오늘 회의에서 논의한 내용을 정리해서'
const context = {
  prefix: PREFIX,
  fullText: PREFIX,
  caretOffset: PREFIX.length,
  anchor: null,
  anchorKind: null,
  isPassword: false,
  isEditable: true,
  isComposing: false,
  hasSelection: false,
  available: true,
  appName: 'notepad.exe',
  windowTitle: 'notes',
  idleMs: 1000,
  capturedAt: Date.now(),
  editedSinceFocus: true,
  typedRecently: true
}

interface InternalSuggestionService {
  _abort: AbortController | null
  _consecutiveFailures: number
  _cooldownUntil: number
  _inFlight: boolean
  _lastRequestAt: number
  _generate(prefix: string, currentContext: typeof context, maxCandidates: number, maxChars: number): Promise<void>
  _abortStaleGeneration(currentPrefix: string): void
  _watchdog(): void
  _generationToken: number
}

function waitForAbort(signal: AbortSignal | undefined): AsyncGenerator<string> {
  return (async function* () {
    await new Promise<void>((_resolve, reject) => {
      if (!signal) {
        reject(new Error('missing abort signal'))
        return
      }
      signal.addEventListener('abort', () => reject(new Error('intentional abort')), { once: true })
    })
    yield 'unreachable'
  })()
}

beforeEach(() => {
  Object.assign(config, {
    suggestionEnabled: true,
    suggestionModelId: 'gemma4:e4b',
    llmModelId: 'gemma4:e4b',
    inputExcludedApps: [],
    suggestionTriggerDelayMs: 600,
    suggestionMinPrefixChars: 8,
    suggestionMaxRequestsPerMinute: 6,
    suggestionDailyBudget: 500,
    suggestionRequestTimeoutMs: 8000,
    inputLearnTypedText: false,
    inputTelemetryEnabled: false,
    suggestionOverlayInteractive: true
  })
  localLlm.isAvailable.mockReset()
  localLlm.isAvailable.mockReturnValue(true)
  localLlm.streamGenerate.mockReset()
})

afterEach(async () => {
  const { resetSuggestionServiceForTests } = await import('../../../src/main/services/SuggestionService')
  resetSuggestionServiceForTests()
  config.suggestionEnabled = true
  vi.useRealTimers()
  vi.clearAllMocks()
})

describe('SuggestionService warm-up', () => {
  it('동시 warm-up 호출을 하나의 1토큰 요청으로 합치고 10분간 유지한다', async () => {
    localLlm.streamGenerate.mockImplementation(async function* () {
      yield 'ok'
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const warmingStates: boolean[] = []
    service.on('state-changed', (state) => warmingStates.push(state.warmingUp))

    await Promise.all([service.warmUp(), service.warmUp()])

    expect(localLlm.streamGenerate).toHaveBeenCalledTimes(1)
    expect(localLlm.streamGenerate).toHaveBeenCalledWith(
      'hi',
      expect.objectContaining({ maxTokens: 1, keepAlive: '10m' })
    )
    expect(warmingStates).toEqual([true, false])
  })

  it('콜드 모델(워밍업 전)에서는 요청 결정이어도 생성 대신 워밍업을 트리거한다', async () => {
    localLlm.streamGenerate.mockImplementation(async function* () {
      yield 'ok'
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService & {
      _warmUpPromise: Promise<void> | null
    }
    const generateSpy = vi.spyOn(internal, '_generate')

    service.handleTypingContext({ ...context, idleMs: 1000 })

    expect(generateSpy).not.toHaveBeenCalled()
    expect(localLlm.streamGenerate).toHaveBeenCalledWith('hi', expect.objectContaining({ maxTokens: 1 }))
    const state = service.getState()
    expect(state.warmingUp).toBe(true)
    expect(state.lastSkipReason).toBe('model-unavailable')

    await internal._warmUpPromise
  })

  it('dismiss(stale)는 진행 중인 워밍업을 취소하지 않는다', async () => {
    let resolveChunk: (() => void) | null = null
    localLlm.streamGenerate.mockImplementation((_text: string, options: { signal?: AbortSignal }) => {
      return (async function* () {
        await new Promise<void>((resolve, reject) => {
          resolveChunk = resolve
          options.signal?.addEventListener('abort', () => reject(new Error('aborted')), { once: true })
        })
        yield 'ok'
      })()
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService & {
      _warmUpAbort: AbortController | null
    }

    const warmUpPromise = service.warmUp()
    await Promise.resolve()
    expect(internal._warmUpAbort?.signal.aborted).toBe(false)

    service.dismiss('stale')

    expect(internal._warmUpAbort?.signal.aborted).toBe(false)
    resolveChunk?.()
    await warmUpPromise

    expect(localLlm.streamGenerate).toHaveBeenCalledTimes(1)
  })

  it('비활성화하면 가용성 재시도 타이머와 warm-up을 취소한다', async () => {
    vi.useFakeTimers()
    localLlm.isAvailable.mockReturnValue(false)
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const warmUp = service.warmUp()

    config.suggestionEnabled = false
    service.applyConfig()
    await warmUp

    expect(vi.getTimerCount()).toBe(0)
    expect(localLlm.streamGenerate).not.toHaveBeenCalled()
  })

  it('dismiss 취소는 실패 쿨다운을 올리거나 재귀 생성하지 않는다', async () => {
    localLlm.streamGenerate.mockImplementation((_text: string, options: { signal?: AbortSignal }) =>
      waitForAbort(options.signal)
    )
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService
    const cleared: string[] = []
    service.on('cleared', ({ reason }) => cleared.push(reason))
    const generation = internal._generate(PREFIX, context, 3, 240)

    await Promise.resolve()
    service.dismiss('dismissed')
    await generation

    expect(internal._consecutiveFailures).toBe(0)
    expect(internal._cooldownUntil).toBe(0)
    expect(localLlm.streamGenerate).toHaveBeenCalledTimes(1)
    expect(cleared).toEqual(['dismissed'])
  })

  it('생성 중 동일 접두의 주기 스냅샷은 요청을 취소하지 않는다', async () => {
    let signal: AbortSignal | undefined
    localLlm.streamGenerate.mockImplementation((_text: string, options: { signal?: AbortSignal }) => {
      signal = options.signal
      return waitForAbort(options.signal)
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService
    const generation = internal._generate(PREFIX, context, 3, 240)

    await Promise.resolve()
    service.handleTypingContext(context)

    expect(signal?.aborted).toBe(false)
    expect(internal._inFlight).toBe(true)
    service.dismiss('dismissed')
    await generation
  })

  it.each([
    ['선택', { hasSelection: true }, 'selection-active'],
    ['포커스 이탈', { isEditable: false }, 'not-editable']
  ])('생성 중 %s은 후보가 없어도 취소하고 정확한 cleared 사유를 낸다', async (_case, update, reason) => {
    let signal: AbortSignal | undefined
    localLlm.streamGenerate.mockImplementation((_text: string, options: { signal?: AbortSignal }) => {
      signal = options.signal
      return waitForAbort(options.signal)
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const cleared: string[] = []
    service.on('cleared', ({ reason: clearedReason }) => cleared.push(clearedReason))
    const generation = (service as unknown as InternalSuggestionService)._generate(PREFIX, context, 3, 240)

    await Promise.resolve()
    expect(service.getState()).toMatchObject({ candidates: [], generating: true })
    service.handleTypingContext({ ...context, ...update })
    await generation

    expect(signal?.aborted).toBe(true)
    expect(cleared).toEqual([reason])
    expect(service.getState()).toMatchObject({ candidates: [], generating: false, partialText: null })
  })

  it('후보 게시 updated 상태에는 생성 플래그와 부분 텍스트가 남지 않는다', async () => {
    localLlm.streamGenerate.mockImplementation(async function* () {
      yield '다음 단계도 확인하겠습니다.'
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const states: Array<{ candidates: unknown[]; generating: boolean; partialText: string | null }> = []
    service.on('updated', (state) => states.push(state))

    await (service as unknown as InternalSuggestionService)._generate(PREFIX, context, 3, 240)

    const published = states.find((state) => state.candidates.length > 0)
    // partialText 는 후보 공개와 함께 비워지지만, generating 은 계속 true 로 남는다 —
    // 첫 후보 공개 뒤 채우기 루프가 백그라운드에서 나머지(최대 12개)를 마저
    // 청하는 중이라는 신호다(설계).
    expect(published).toMatchObject({ generating: true, partialText: null })
  })

  it('stale 취소는 실패 쿨다운을 올리거나 즉시 재시작하지 않고, 보여줄 후보가 없으면 오버레이를 지운다', async () => {
    localLlm.streamGenerate.mockImplementation((_text: string, options: { signal?: AbortSignal }) =>
      waitForAbort(options.signal)
    )
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService
    const cleared: string[] = []
    service.on('cleared', ({ reason }) => cleared.push(reason))
    const generation = internal._generate(PREFIX, context, 3, 240)

    await Promise.resolve()
    internal._abortStaleGeneration('완전히 다른 문맥으로 바뀐 입력입니다')
    await generation

    expect(internal._consecutiveFailures).toBe(0)
    expect(internal._cooldownUntil).toBe(0)
    expect(localLlm.streamGenerate).toHaveBeenCalledTimes(1)
    // 스피너만 뜬 채로 남지 않도록, 후보가 없는 stale 취소는 오버레이를 닫는다.
    expect(cleared).toEqual(['stale'])
  })

  it('stale 취소는 속도 제한 예산을 환급해 다음 요청이 rate-limited 되지 않는다', async () => {
    localLlm.streamGenerate.mockImplementation((_text: string, options: { signal?: AbortSignal }) =>
      waitForAbort(options.signal)
    )
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService & {
      _minuteCount: number
      _dayCount: number
      _prevRequestAt: number
      _lastSkipReason: string | null
    }

    // 분당 카운터 창을 먼저 굳힌다 — 그렇지 않으면 이 인스턴스의 첫 getState() 호출
    // (_generate 내부의 emit('updated', ...) 에서 일어난다) 이 창을 "지금" 으로
    // 다시 잡으며 방금 늘린 카운트를 0으로 되돌린다 (실제 흐름에선 handleTypingContext
    // 가 먼저 창을 굳혀 두므로 일어나지 않는 순서 문제).
    service.getState()

    const generation = internal._generate(PREFIX, context, 3, 240)
    await Promise.resolve()

    expect(internal._minuteCount).toBe(1)
    const requestedAt = internal._lastRequestAt
    expect(requestedAt).toBeGreaterThan(0)

    internal._abortStaleGeneration('완전히 다른 문맥으로 바뀐 입력입니다')
    await generation

    // 취소된 요청이 쓴 예산이 되돌아간다.
    expect(internal._minuteCount).toBe(0)
    expect(internal._lastRequestAt).toBe(internal._prevRequestAt)
    expect(internal._lastRequestAt).toBeLessThan(requestedAt)

    // 되돌아간 예산으로 바로 다음 요청은 rate-limited 로 막히지 않는다.
    service.handleTypingContext({ ...context, idleMs: 1000 })
    expect(internal._lastSkipReason).not.toBe('rate-limited')
  })

  it('pageNext는 다음 페이지 첫 항목으로, 후보가 없는 페이지면 그대로 둔다', async () => {
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as { _candidates: Array<{ text: string; rank: number }> }
    internal._candidates = Array.from({ length: 5 }, (_v, i) => ({ text: `후보${i}`, rank: i }))

    expect(service.getState().activeIndex).toBe(0)
    service.pageNext()
    expect(service.getState().activeIndex).toBe(3)
    service.pageNext()
    // 5개뿐이라 다음 페이지가 없다 — 그대로 둔다.
    expect(service.getState().activeIndex).toBe(3)
    service.pagePrev()
    expect(service.getState().activeIndex).toBe(0)
    service.pagePrev()
    expect(service.getState().activeIndex).toBe(0)
  })

  it('next/previous는 페이지와 무관하게 전체 후보를 가로질러 순환한다', async () => {
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as { _candidates: Array<{ text: string; rank: number }> }
    internal._candidates = Array.from({ length: 4 }, (_v, i) => ({ text: `후보${i}`, rank: i }))

    service.next()
    service.next()
    service.next()
    expect(service.getState().activeIndex).toBe(3)
    service.next()
    expect(service.getState().activeIndex).toBe(0)
    service.previous()
    expect(service.getState().activeIndex).toBe(3)
  })

  it('watchdog는 실제 요청 signal을 abort하고 세대를 무효화한다', async () => {
    let signal: AbortSignal | undefined
    localLlm.streamGenerate.mockImplementation((_text: string, options: { signal?: AbortSignal }) => {
      signal = options.signal
      return waitForAbort(options.signal)
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService
    const cleared: string[] = []
    service.on('cleared', ({ reason }) => cleared.push(reason))
    const generation = internal._generate(PREFIX, context, 3, 240)

    await Promise.resolve()
    internal._lastRequestAt = Date.now() - 13001
    internal._watchdog()
    await generation

    expect(signal?.aborted).toBe(true)
    expect(internal._inFlight).toBe(false)
    expect(internal._consecutiveFailures).toBe(0)
    expect(cleared).toEqual([])
  })
})

describe('SuggestionService 세션 채우기 (최대 12개, 순차 · 페이지)', () => {
  it('첫 요청은 1개만 청하고, 성공하면 채우기 루프가 하나씩 최대 12개까지 채운다', async () => {
    // 접두/확장 중복 판정 때문에 숫자 접미사(1, 10, 11…)는 서로를 중복으로 오판한다
    // ("이어지는 문장 1" 이 "이어지는 문장 10" 의 접두이므로) — 서로소인 글자를 쓴다.
    const LETTERS = 'ABCDEFGHIJKL'
    let call = 0
    localLlm.streamGenerate.mockImplementation(async function* () {
      const letter = LETTERS[call % LETTERS.length]
      call += 1
      yield `이어지는 문장 ${letter}`
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService

    await internal._generate(PREFIX, context, 3, 240)
    // 세션의 첫 요청은 정확히 1개만 청한다 (한꺼번에 여러 개를 청하면 느리다 — 사용자 요청).
    expect(localLlm.streamGenerate).toHaveBeenCalledWith(
      expect.stringContaining('1개'),
      expect.anything()
    )

    await vi.waitFor(() => {
      expect(service.getState().candidates).toHaveLength(12)
    })
    expect(service.getState().generating).toBe(false)
    expect(service.getState().targetTotal).toBe(12)
    expect(localLlm.streamGenerate).toHaveBeenCalledTimes(12)
  })

  it('연속 2번 새 후보가 없으면(전부 중복) 채우기를 멈춘다', async () => {
    localLlm.streamGenerate.mockImplementation(async function* () {
      yield '같은 문장입니다.'
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService

    await internal._generate(PREFIX, context, 3, 240)
    await vi.waitFor(() => {
      expect(service.getState().generating).toBe(false)
    })

    expect(service.getState().candidates).toHaveLength(1)
    // 첫 요청 1번 + 중복으로 끝난 채우기 시도 2번 = 3번.
    expect(localLlm.streamGenerate).toHaveBeenCalledTimes(3)
  })

  it('중복(완전 일치·접두/확장)은 건너뛰고 고유한 후보만 덧붙인다', async () => {
    const sequence = ['같은 문장입니다.', '같은 문장입니다.', '다른 문장입니다.']
    let call = 0
    localLlm.streamGenerate.mockImplementation(async function* () {
      const text = sequence[Math.min(call, sequence.length - 1)]
      call += 1
      yield text
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService

    await internal._generate(PREFIX, context, 3, 240)
    await vi.waitFor(() => {
      expect(service.getState().generating).toBe(false)
    })

    expect(service.getState().candidates.map((c) => c.text)).toEqual([
      '같은 문장입니다.',
      '다른 문장입니다.'
    ])
  })

  it('dismiss는 진행 중인 채우기 요청을 취소하고 루프를 멈춘다', async () => {
    let fillSignal: AbortSignal | undefined
    let firstServed = false
    localLlm.streamGenerate.mockImplementation((_text: string, options: { signal?: AbortSignal }) => {
      if (!firstServed) {
        firstServed = true
        return (async function* () {
          yield '첫 후보입니다.'
        })()
      }
      fillSignal = options.signal
      return waitForAbort(options.signal)
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService

    await internal._generate(PREFIX, context, 3, 240)
    await vi.waitFor(() => expect(fillSignal).toBeDefined())

    expect(fillSignal?.aborted).toBe(false)
    service.dismiss('dismissed')

    expect(fillSignal?.aborted).toBe(true)
    expect(service.getState().candidates).toHaveLength(0)
  })

  it('세션의 첫 요청만 분당/일일 예산을 쓴다 — 채우기 요청은 쓰지 않는다', async () => {
    const LETTERS = 'ABCDEFGHIJKL'
    let call = 0
    localLlm.streamGenerate.mockImplementation(async function* () {
      const letter = LETTERS[call % LETTERS.length]
      call += 1
      yield `문장 ${letter}`
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService & {
      _minuteCount: number
      _dayCount: number
    }
    // 분당 카운터 창을 먼저 굳힌다 (다른 테스트와 같은 이유 — 위 주석 참조).
    service.getState()

    await internal._generate(PREFIX, context, 3, 240)
    await vi.waitFor(() => expect(service.getState().candidates).toHaveLength(12))

    expect(internal._minuteCount).toBe(1)
    expect(internal._dayCount).toBe(1)
  })

  it('세션이 떠 있는 동안 접두가 자라면(다음 문장 시작) 즉시 세션을 끝낸다', async () => {
    localLlm.streamGenerate.mockImplementation(async function* () {
      yield '첫 후보입니다.'
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService
    const cleared: string[] = []
    service.on('cleared', ({ reason }) => cleared.push(reason))

    await internal._generate(PREFIX, context, 3, 240)
    await vi.waitFor(() => expect(service.getState().visible).toBe(true))

    service.handleTypingContext({ ...context, prefix: `${PREFIX} 추가로 입력했습니다`, idleMs: 1000 })

    expect(service.getState().visible).toBe(false)
    expect(cleared).toContain('stale')
  })

  it('세션이 떠 있는 동안 마지막 글자만 IME 조합으로 바뀌면 세션을 유지한다', async () => {
    localLlm.streamGenerate.mockImplementation(async function* () {
      yield '첫 후보입니다.'
    })
    const { getSuggestionService } = await import('../../../src/main/services/SuggestionService')
    const service = getSuggestionService()
    const internal = service as unknown as InternalSuggestionService
    const cleared: string[] = []
    service.on('cleared', ({ reason }) => cleared.push(reason))

    await internal._generate(PREFIX, context, 3, 240)
    await vi.waitFor(() => expect(service.getState().visible).toBe(true))

    const mutatedLastChar = `${PREFIX.slice(0, -1)}요`
    service.handleTypingContext({ ...context, prefix: mutatedLastChar, idleMs: 1000 })

    expect(service.getState().visible).toBe(true)
    expect(cleared).not.toContain('stale')
  })
})