Luigit
repositories / will

will

owned by admin

src/agent/state.ts

Raw
import { readFile, rename, writeFile } from 'node:fs/promises'
import { dirname, join } from 'node:path'
import type { BudgetState } from '../shared/budget.ts'
import { type BudgetConfig, emptyBudgetState, rollover } from '../shared/budget.ts'
import { type Clock, systemClock } from '../shared/time.ts'
import type { HeartbeatState } from '../shared/wake.ts'

/**
 * Durable agent state on the agent volume: telegram offset, heartbeat claim,
 * budget accounting, tracked pushes. Written atomically after every change.
 */

export interface AgentStateData {
  telegramOffset?: number | undefined
  botUserId?: number | undefined
  botUsername?: string | undefined
  heartbeat: HeartbeatState
  budget: BudgetState
  pushedShas: string[]
  handledCiRuns: string[]
  storageNoticeState: 'ok' | 'warn' | 'stop'
  deployEventCursor?: string | undefined
}

export function emptyAgentState(clock: Clock = systemClock): AgentStateData {
  return {
    heartbeat: {},
    budget: emptyBudgetState(clock),
    pushedShas: [],
    handledCiRuns: [],
    storageNoticeState: 'ok',
  }
}

export class AgentState {
  #data: AgentStateData
  readonly #path: string

  private constructor(path: string, data: AgentStateData) {
    this.#path = path
    this.#data = data
  }

  static async load(path: string, clock: Clock = systemClock): Promise<AgentState> {
    let data: AgentStateData
    try {
      data = JSON.parse(await readFile(path, 'utf8')) as AgentStateData
    } catch {
      data = emptyAgentState(clock)
    }
    data.budget = rollover(data.budget ?? emptyBudgetState(clock), clock)
    data.heartbeat ??= {}
    data.pushedShas ??= []
    data.handledCiRuns ??= []
    if (data.storageNoticeState === undefined) data.storageNoticeState = 'ok'
    return new AgentState(path, data)
  }

  get data(): Readonly<AgentStateData> {
    return this.#data
  }

  update(mutate: (draft: AgentStateData) => void): Promise<void> {
    mutate(this.#data)
    return this.#persist()
  }

  async #persist(): Promise<void> {
    const tmp = join(dirname(this.#path), `.state-${process.pid}.tmp`)
    await writeFile(tmp, JSON.stringify(this.#data), 'utf8')
    await rename(tmp, this.#path)
  }
}

export function budgetConfigFromEnv(
  env: Record<string, string | undefined>,
  fallback: BudgetConfig,
): BudgetConfig {
  const cap = env.WILL_BUDGET_DAILY_CAP
  const fraction = env.WILL_BUDGET_RECOVERY_FRACTION
  return {
    dailyCapCalls:
      cap !== undefined && cap !== '' ? Number.parseInt(cap, 10) : fallback.dailyCapCalls,
    recoveryReserveFraction:
      fraction !== undefined && fraction !== ''
        ? Number.parseFloat(fraction)
        : fallback.recoveryReserveFraction,
  }
}