Repository Public
1/** SQL of the step protocol. Every function here runs inside a fenced transaction (see fencedTx). */2import type { EffectClass, RunStatus, Usage } from '@aos/shared';3import { num, query, withTx, type Kernel, type Tx } from './db.js';4import { KernelError, LeaseLostError, type BudgetReason, type RecordedError } from './errors.js';5import { announceRunStatus, RUN_COLS, type RunRow } from './runs.js';67export interface StepRow {8 seq: number;9 kind: string;10 name: string;11 effect: EffectClass;12 status: 'started' | 'completed' | 'failed' | 'needs_reconcile';13 request_hash: string;14 result: unknown;15 error: RecordedError | null;16 usage: Partial<Usage> | null;17 idem_key: string;18 completed_at: Date | null;19}20export interface StepMeta {21 kind: string;22 name: string;23 effect: EffectClass;24 /** Kernel-internal bookkeeping steps (now, uuid, owner_messages) do not emit step.* events. */25 quiet?: boolean;26}2728/** Bodies larger than this are not copied into step.* events (consumers fetch the step instead). */29const EVENT_BODY_MAX = 16_384;30const ZERO: Usage = { usd: 0, inputTokens: 0, outputTokens: 0, cacheReadTokens: 0, cacheWriteTokens: 0, turns: 0, toolCalls: 0 };31const USAGE_KEYS = Object.keys(ZERO) as (keyof Usage)[];3233export function zeroUsage(): Usage {34 return { ...ZERO };35}36export function addUsage(a: Partial<Usage> | null | undefined, b: Partial<Usage> | null | undefined): Usage {37 const out = zeroUsage();38 for (const key of USAGE_KEYS) out[key] = num(a?.[key]) + num(b?.[key]);39 return out;40}41/** Sanitises exec-reported usage; counts one turn per model_call and one tool call per tool_call step. */42export function normalizeUsage(meta: StepMeta, u: Partial<Usage> | undefined): Usage | null {43 const out = zeroUsage();44 for (const key of USAGE_KEYS) {45 const v = Number(u?.[key]);46 out[key] = Number.isFinite(v) && v > 0 ? v : 0;47 }48 if (meta.effect === 'model_call' && out.turns === 0) out.turns = 1;49 if (meta.kind === 'tool_call' && out.toolCalls === 0) out.toolCalls = 1;50 return USAGE_KEYS.some((key) => out[key] !== 0) ? out : null;51}5253/**54 * One transaction guarded by the fencing token: the run row is locked and must still be 'running'55 * under OUR token, otherwise nothing is written and LeaseLostError is thrown.56 * `live` (used for intents, i.e. before starting NEW work) additionally requires an unexpired lease:57 * a worker that stalled past its lease may still record the result of work it already did (nobody58 * has claimed the run, the token proves it), but it may not start anything new.59 */60export async function fencedTx<T>(k: Kernel, runId: string, token: string, fn: (tx: Tx, run: RunRow) => Promise<T>, live = false): Promise<T> {61 return withTx(k, async (tx) => {62 const run = await tx.one<RunRow>(63 `SELECT ${RUN_COLS} FROM runs WHERE id = $1 AND fencing_token = $2::bigint AND status = 'running'${live ? ' AND lease_expires_at > now()' : ''} FOR UPDATE`,64 [runId, token],65 );66 if (!run) throw new LeaseLostError(runId, live ? 'fencing check failed or lease expired' : undefined);67 return fn(tx, run);68 });69}7071export async function loadJournal(k: Kernel, runId: string): Promise<Map<number, StepRow>> {72 const rows = await query<StepRow>(73 k,74 `SELECT seq, kind, name, effect, status, request_hash, result, error, usage, idem_key, completed_at FROM steps WHERE run_id = $1 ORDER BY seq`,75 [runId],76 );77 return new Map(rows.map((r) => [r.seq, r]));78}7980function body(json: string, key: 'request' | 'result'): Record<string, unknown> {81 return json.length <= EVENT_BODY_MAX ? { [key]: JSON.parse(json) as unknown, truncated: false } : { truncated: true };82}8384/** Intent record: the step row is inserted 'started' (+ step.started) before anything executes. */85export async function insertIntent(tx: Tx, run: RunRow, seq: number, meta: StepMeta, hash: string, requestJson: string, idemKey: string): Promise<void> {86 await tx.q(87 `INSERT INTO steps (run_id, seq, kind, name, effect, status, request_hash, request, idem_key) VALUES ($1,$2,$3,$4,$5,'started',$6,$7::jsonb,$8)`,88 [run.id, seq, meta.kind, meta.name, meta.effect, hash, requestJson, idemKey],89 );90 await tx.q('UPDATE runs SET next_seq = GREATEST(next_seq, $2) WHERE id = $1', [run.id, seq + 1]);91 if (!meta.quiet) {92 tx.emit({ missionId: run.mission_id, runId: run.id, seq, kind: 'step.started', payload: { kind: meta.kind, name: meta.name, effect: meta.effect, ...body(requestJson, 'request') } });93 }94}9596export interface Completed {97 /** true ⇒ the row was already completed (an earlier commit whose acknowledgement was lost). */98 already: boolean;99 result: unknown;100 usage: Partial<Usage> | null;101 completedAt: Date;102 cancelRequested: boolean;103}104105/**106 * Result record. ONE transaction: step → completed, usage added to runs.usage, missions.spent_usd and107 * spend_days, next_seq bumped, step.completed + run.usage emitted. Idempotent: if the step is already108 * completed nothing is accounted twice.109 */110export async function completeStep(tx: Tx, run: RunRow, seq: number, meta: StepMeta, resultJson: string, usage: Usage | null): Promise<Completed> {111 const up = await tx.q<{ completed_at: Date }>(112 `UPDATE steps SET status = 'completed', result = $3::jsonb, usage = $4::jsonb, completed_at = now()113 WHERE run_id = $1 AND seq = $2 AND status IN ('started','needs_reconcile') RETURNING completed_at`,114 [run.id, seq, resultJson, usage ? JSON.stringify(usage) : null],115 );116 if (up.length === 0) {117 const ex = await tx.one<StepRow>('SELECT status, result, usage, completed_at FROM steps WHERE run_id = $1 AND seq = $2', [run.id, seq]);118 if (ex?.status === 'completed' && ex.completed_at) {119 return { already: true, result: ex.result, usage: ex.usage, completedAt: ex.completed_at, cancelRequested: run.cancel_requested };120 }121 throw new KernelError(`step ${seq} of run ${run.id} cannot be completed from status ${ex?.status ?? 'missing'}`);122 }123 await tx.q('UPDATE runs SET next_seq = GREATEST(next_seq, $2), usage = $3::jsonb WHERE id = $1', [124 run.id,125 seq + 1,126 JSON.stringify(addUsage(run.usage, usage)),127 ]);128 if (!meta.quiet) {129 tx.emit({ missionId: run.mission_id, runId: run.id, seq, kind: 'step.completed', payload: { kind: meta.kind, name: meta.name, effect: meta.effect, usage, ...body(resultJson, 'result') } });130 }131 if (usage) {132 let spent: number | null = null;133 if (usage.usd > 0) {134 const m = await tx.one<{ spent_usd: unknown }>('UPDATE missions SET spent_usd = spent_usd + $2 WHERE id = $1 RETURNING spent_usd', [run.mission_id, usage.usd]);135 spent = m ? num(m.spent_usd) : null;136 await tx.q(`INSERT INTO spend_days (day, usd) VALUES (current_date, $1) ON CONFLICT (day) DO UPDATE SET usd = spend_days.usd + EXCLUDED.usd`, [usage.usd]);137 }138 tx.emit({ missionId: run.mission_id, runId: run.id, seq, kind: 'run.usage', payload: { usage: addUsage(run.usage, usage), budget: run.budget, missionSpentUsd: spent } });139 }140 return { already: false, result: JSON.parse(resultJson) as unknown, usage, completedAt: (up[0] as { completed_at: Date }).completed_at, cancelRequested: run.cancel_requested };141}142143/** Failure record (exec threw, or the kernel refused the step: budget). Inserts the row if there was no intent. */144export async function failStep(145 tx: Tx,146 run: RunRow,147 seq: number,148 meta: StepMeta,149 error: RecordedError,150 intent?: { hash: string; requestJson: string; idemKey: string },151): Promise<Date> {152 let rows: { completed_at: Date }[];153 if (intent) {154 rows = await tx.q(155 `INSERT INTO steps (run_id, seq, kind, name, effect, status, request_hash, request, idem_key, error, completed_at)156 VALUES ($1,$2,$3,$4,$5,'failed',$6,$7::jsonb,$8,$9::jsonb, now()) RETURNING completed_at`,157 [run.id, seq, meta.kind, meta.name, meta.effect, intent.hash, intent.requestJson, intent.idemKey, JSON.stringify(error)],158 );159 } else {160 rows = await tx.q(161 `UPDATE steps SET status = 'failed', error = $3::jsonb, completed_at = now()162 WHERE run_id = $1 AND seq = $2 AND status IN ('started','needs_reconcile') RETURNING completed_at`,163 [run.id, seq, JSON.stringify(error)],164 );165 if (rows.length === 0) {166 const ex = await tx.one<{ status: string; completed_at: Date | null }>('SELECT status, completed_at FROM steps WHERE run_id = $1 AND seq = $2', [run.id, seq]);167 if (ex?.status === 'failed' && ex.completed_at) return ex.completed_at;168 throw new KernelError(`step ${seq} of run ${run.id} cannot be failed from status ${ex?.status ?? 'missing'}`);169 }170 }171 await tx.q('UPDATE runs SET next_seq = GREATEST(next_seq, $2) WHERE id = $1', [run.id, seq + 1]);172 if (!meta.quiet) {173 tx.emit({ missionId: run.mission_id, runId: run.id, seq, kind: 'step.failed', payload: { kind: meta.kind, name: meta.name, effect: meta.effect, error } });174 }175 return (rows[0] as { completed_at: Date }).completed_at;176}177178/** Hard stops checked before a new model_call: run budget (usd, tokens), mission budget, daily cap. */179export async function budgetStop(tx: Tx, run: RunRow, dailyCapUsd: number | null): Promise<BudgetReason | null> {180 const eps = 1e-9;181 const u = run.usage;182 if (num(u.usd) >= run.budget.usd - eps) return 'run_budget';183 if (run.budget.tokens && num(u.inputTokens) + num(u.outputTokens) >= run.budget.tokens) return 'run_budget';184 const m = await tx.one<{ over: boolean }>('SELECT (spent_usd >= budget_usd) AS over FROM missions WHERE id = $1', [run.mission_id]);185 if (m?.over) return 'mission_budget';186 if (dailyCapUsd !== null) {187 const d = await tx.one<{ usd: unknown }>('SELECT usd FROM spend_days WHERE day = current_date');188 if (d && num(d.usd) >= dailyCapUsd - eps) return 'daily_cap';189 }190 return null;191}192193/** Suspension: the run leaves 'running' and the lease is released, in the caller's fenced transaction. */194export async function suspendRun(tx: Tx, run: RunRow, to: Extract<RunStatus, 'waiting' | 'queued'>, wakeAt: Date | null): Promise<void> {195 await tx.q(`UPDATE runs SET status = $2, wake_at = $3, lease_owner = NULL, lease_expires_at = NULL WHERE id = $1`, [run.id, to, wakeAt]);196 await announceRunStatus(tx, run, to, wakeAt ? { wakeAt: wakeAt.toISOString() } : {});197}198