orca/src/main/agent-awake-service.test.ts

import { beforeEach, describe, expect, it, vi } from 'vitest'
import { AgentAwakeService, AGENT_AWAKE_STATUS_STALE_AFTER_MS } from './agent-awake-service'
import type { AgentAwakeStatus } from './agent-awake-service'

vi.mock('electron', () => ({
  powerMonitor: {
    on: vi.fn(),
    off: vi.fn()
  },
  powerSaveBlocker: {
    start: vi.fn(),
    stop: vi.fn(),
    isStarted: vi.fn()
  }
}))

function workingStatus(overrides: Partial<AgentAwakeStatus> = {}): AgentAwakeStatus {
  return {
    state: 'working',
    receivedAt: 1_000,
    observedInCurrentRuntime: true,
    ...overrides
  }
}

function createBlocker() {
  const startedIds = new Set<number>()
  let nextId = 1
  return {
    start: vi.fn(() => {
      const id = nextId++
      startedIds.add(id)
      return id
    }),
    stop: vi.fn((id: number) => {
      startedIds.delete(id)
    }),
    isStarted: vi.fn((id: number) => startedIds.has(id)),
    startedIds
  }
}

function createMacosAssertion() {
  return {
    start: vi.fn(),
    stop: vi.fn(),
    dispose: vi.fn()
  }
}

function createLinuxAssertion() {
  return {
    start: vi.fn(),
    stop: vi.fn(),
    dispose: vi.fn()
  }
}

function createPowerMonitor() {
  const listeners = new Set<() => void>()
  return {
    on: vi.fn((_event: 'resume', listener: () => void) => {
      listeners.add(listener)
    }),
    off: vi.fn((_event: 'resume', listener: () => void) => {
      listeners.delete(listener)
    }),
    emitResume: () => {
      for (const listener of listeners) {
        listener()
      }
    }
  }
}

function createService(
  now: () => number,
  blocker = createBlocker(),
  macosAssertion = createMacosAssertion(),
  linuxAssertion = createLinuxAssertion(),
  powerMonitor: ReturnType<typeof createPowerMonitor> | null = null
): AgentAwakeService {
  return new AgentAwakeService({
    blocker,
    linuxAssertion,
    macosAssertion,
    now,
    powerMonitor,
    logger: {
      debug: vi.fn(),
      warn: vi.fn()
    }
  })
}

describe('AgentAwakeService', () => {
  beforeEach(() => {
    vi.useRealTimers()
  })

  it('does not start when disabled even with a running status', () => {
    const blocker = createBlocker()
    const service = createService(() => 1_000, blocker)

    service.setStatuses([workingStatus()])

    expect(blocker.start).not.toHaveBeenCalled()
  })

  it('starts Electron and platform assertions when enabled with a fresh working status', () => {
    const blocker = createBlocker()
    const macosAssertion = createMacosAssertion()
    const linuxAssertion = createLinuxAssertion()
    const service = createService(() => 1_000, blocker, macosAssertion, linuxAssertion)

    service.setEnabled(true)
    service.setStatuses([workingStatus()])

    expect(blocker.start).toHaveBeenCalledTimes(1)
    expect(blocker.start).toHaveBeenCalledWith('prevent-display-sleep')
    expect(macosAssertion.start).toHaveBeenCalledTimes(1)
    expect(linuxAssertion.start).toHaveBeenCalledTimes(1)
  })

  it('stays awake in On mode without a working agent', () => {
    const blocker = createBlocker()
    const service = createService(() => 1_000, blocker)

    service.setMode('on')

    expect(blocker.start).toHaveBeenCalledWith('prevent-display-sleep')
    expect(service.getStatus()).toEqual({ mode: 'on', active: true })
  })

  it('publishes Auto activity changes to status subscribers', () => {
    const service = createService(() => 1_000)
    const listener = vi.fn()
    service.subscribe(listener)

    service.setMode('auto')
    service.setStatuses([workingStatus()])
    service.setStatuses([workingStatus(), workingStatus()])
    service.setStatuses([workingStatus()])
    service.setStatuses([])

    expect(listener).toHaveBeenNthCalledWith(1, {
      mode: 'auto',
      active: false
    })
    expect(listener).toHaveBeenNthCalledWith(2, {
      mode: 'auto',
      active: true
    })
    expect(listener).toHaveBeenNthCalledWith(3, {
      mode: 'auto',
      active: false
    })
    expect(listener).toHaveBeenCalledTimes(3)
  })

  it('starts and stops from settings flips around an already-running status', () => {
    const blocker = createBlocker()
    const macosAssertion = createMacosAssertion()
    const linuxAssertion = createLinuxAssertion()
    const service = createService(() => 1_000, blocker, macosAssertion, linuxAssertion)

    service.setStatuses([workingStatus()])
    service.setEnabled(true)
    service.setEnabled(false)

    expect(blocker.start).toHaveBeenCalledTimes(1)
    expect(blocker.stop).toHaveBeenCalledWith(1)
    expect(macosAssertion.start).toHaveBeenCalledTimes(1)
    expect(macosAssertion.stop).toHaveBeenCalled()
    expect(linuxAssertion.start).toHaveBeenCalledTimes(1)
    expect(linuxAssertion.stop).toHaveBeenCalled()
  })

  it('ignores startup-hydrated working statuses that were not observed in this runtime', () => {
    const blocker = createBlocker()
    const service = createService(() => 1_000, blocker)

    service.setEnabled(true)
    service.setStatuses([workingStatus({ observedInCurrentRuntime: false })])

    expect(blocker.start).not.toHaveBeenCalled()
  })

  it('does not start for blocked, waiting, or done statuses', () => {
    const blocker = createBlocker()
    const service = createService(() => 1_000, blocker)

    service.setEnabled(true)
    service.setStatuses([
      workingStatus({ state: 'blocked' }),
      workingStatus({ state: 'waiting' }),
      workingStatus({ state: 'done' })
    ])

    expect(blocker.start).not.toHaveBeenCalled()
  })

  it('does not start a second blocker when one working status replaces another', () => {
    const blocker = createBlocker()
    const service = createService(() => 1_000, blocker)

    service.setEnabled(true)
    service.setStatuses([workingStatus({ receivedAt: 1_000 })])
    service.setStatuses([workingStatus({ receivedAt: 1_100 })])

    expect(blocker.start).toHaveBeenCalledTimes(1)
  })

  it('stops when the last running status is dropped', () => {
    const blocker = createBlocker()
    const macosAssertion = createMacosAssertion()
    const linuxAssertion = createLinuxAssertion()
    const service = createService(() => 1_000, blocker, macosAssertion, linuxAssertion)

    service.setEnabled(true)
    service.setStatuses([workingStatus()])
    service.setStatuses([])

    expect(blocker.stop).toHaveBeenCalledWith(1)
    expect(macosAssertion.stop).toHaveBeenCalledWith('status-change')
    expect(linuxAssertion.stop).toHaveBeenCalledWith('status-change')
  })

  it('does not start for a stale working status', () => {
    const blocker = createBlocker()
    const service = createService(() => AGENT_AWAKE_STATUS_STALE_AFTER_MS + 1_001, blocker)

    service.setEnabled(true)
    service.setStatuses([workingStatus({ receivedAt: 1_000 })])

    expect(blocker.start).not.toHaveBeenCalled()
  })

  it('stops when the only running status becomes stale without another event', () => {
    vi.useFakeTimers()
    let now = 1_000
    const blocker = createBlocker()
    const macosAssertion = createMacosAssertion()
    const linuxAssertion = createLinuxAssertion()
    const service = createService(() => now, blocker, macosAssertion, linuxAssertion)

    service.setEnabled(true)
    service.setStatuses([workingStatus({ receivedAt: 1_000 })])
    now = 1_000 + AGENT_AWAKE_STATUS_STALE_AFTER_MS + 1
    vi.advanceTimersByTime(AGENT_AWAKE_STATUS_STALE_AFTER_MS)

    expect(blocker.stop).toHaveBeenCalledWith(1)
    expect(macosAssertion.stop).toHaveBeenCalledWith('stale-expiry')
    expect(linuxAssertion.stop).toHaveBeenCalledWith('stale-expiry')
    service.dispose()
  })

  it('reschedules stale expiry for a newer running status', () => {
    vi.useFakeTimers()
    let now = 1_000
    const blocker = createBlocker()
    const macosAssertion = createMacosAssertion()
    const linuxAssertion = createLinuxAssertion()
    const service = createService(() => now, blocker, macosAssertion, linuxAssertion)

    service.setEnabled(true)
    service.setStatuses([workingStatus({ receivedAt: 1_000 })])
    now = 2_000
    service.setStatuses([workingStatus({ receivedAt: 2_000 })])
    now = 1_000 + AGENT_AWAKE_STATUS_STALE_AFTER_MS + 1
    vi.advanceTimersByTime(AGENT_AWAKE_STATUS_STALE_AFTER_MS)

    expect(blocker.stop).not.toHaveBeenCalled()
    now = 2_000 + AGENT_AWAKE_STATUS_STALE_AFTER_MS + 1
    vi.advanceTimersByTime(1_000)

    expect(blocker.stop).toHaveBeenCalledWith(1)
    service.dispose()
  })

  it('keeps the blocker id when stop fails and Electron reports it is still started', () => {
    const blocker = createBlocker()
    blocker.stop.mockImplementation(() => {
      throw new Error('stop failed')
    })
    const service = createService(() => 1_000, blocker)

    service.setEnabled(true)
    service.setStatuses([workingStatus()])
    service.setStatuses([])
    service.setStatuses([])

    expect(blocker.stop).toHaveBeenCalledTimes(2)
    expect(blocker.stop).toHaveBeenCalledWith(1)
  })

  it('disposes by clearing timers and stopping an active blocker once', () => {
    vi.useFakeTimers()
    const blocker = createBlocker()
    const macosAssertion = createMacosAssertion()
    const linuxAssertion = createLinuxAssertion()
    const service = createService(() => 1_000, blocker, macosAssertion, linuxAssertion)

    service.setEnabled(true)
    service.setStatuses([workingStatus()])
    service.dispose()
    vi.advanceTimersByTime(AGENT_AWAKE_STATUS_STALE_AFTER_MS)

    expect(blocker.stop).toHaveBeenCalledTimes(1)
    expect(blocker.stop).toHaveBeenCalledWith(1)
    expect(macosAssertion.dispose).toHaveBeenCalledTimes(1)
    expect(linuxAssertion.dispose).toHaveBeenCalledTimes(1)
  })

  it('reconciles assertions on power resume while work is still eligible', () => {
    const blocker = createBlocker()
    const macosAssertion = createMacosAssertion()
    const linuxAssertion = createLinuxAssertion()
    const monitor = createPowerMonitor()
    const service = createService(() => 1_000, blocker, macosAssertion, linuxAssertion, monitor)

    service.setEnabled(true)
    service.setStatuses([workingStatus()])
    blocker.startedIds.clear()
    monitor.emitResume()

    expect(blocker.start).toHaveBeenCalledTimes(2)
    expect(macosAssertion.start).toHaveBeenCalledTimes(2)
    expect(linuxAssertion.start).toHaveBeenCalledTimes(2)
  })

  it('unsubscribes the resume listener on dispose', () => {
    const blocker = createBlocker()
    const macosAssertion = createMacosAssertion()
    const linuxAssertion = createLinuxAssertion()
    const monitor = createPowerMonitor()
    const service = createService(() => 1_000, blocker, macosAssertion, linuxAssertion, monitor)

    service.dispose()

    expect(monitor.on).toHaveBeenCalledTimes(1)
    expect(monitor.off).toHaveBeenCalledTimes(1)
  })
})