repositories / will
will
owned by admin
src/agent/state.ts
Rawimport { 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,
}
}