Repository Public
1/** RunContext implementation: the step protocol (CONTRACT §2) over the journal. */2import { randomUUID } from 'node:crypto';3import type { DelegationContract, EventKind, Trust, Usage } from '@aos/shared';4import { canonicalize, canonicalJson, idemKeyFor, sha256Hex } from './canonical.js';5import { query, type Kernel, type Tx } from './db.js';6import {7 BudgetExhaustedError, isControlError, KernelError, KernelStorageError, LeaseLostError, NonDeterminismError,8 recordError, ReconcileAborted, RunCancelled, StepError, Suspend, type RecordedError,9} from './errors.js';10import { TransientSender } from './events.js';11import { Gate } from './gate.js';12import { addUsage, budgetStop, completeStep, failStep, fencedTx, insertIntent, normalizeUsage, zeroUsage, type Completed, type StepMeta, type StepRow } from './journal.js';13import type { ClaimedRun } from './lease.js';14import type { RunRow } from './runs.js';15import type { KernelHooks, RunContext, StepBoundary, StepExec, StepSpec } from './types.js';16import { Emitter } from './context-emit.js';17import { askReconcile, durableWait } from './context-waits.js';1819const QUIET_KINDS = new Set(['now', 'uuid', 'owner_messages']);20const sleepMs = (ms: number): Promise<void> => new Promise((r) => setTimeout(r, ms));2122export interface CtxDeps {23 hooks?: KernelHooks;24 /** Daily cap resolver, called before each new model_call. */25 dailyCapUsd: () => Promise<number | null>;26 /** Timers shorter than this are waited in-process instead of suspending the run. */27 inlineSleepMs: number;28 /** Called once when the run is suspended (stop the heartbeat). */29 onSuspend: () => void;30}3132export class RunCtx implements RunContext {33 readonly runId: string;34 readonly missionId: string;35 readonly parentRunId: string | null;36 readonly contract: DelegationContract;37 readonly workspaceAbs: string;38 readonly signal: AbortSignal;39 readonly gate = new Gate();40 readonly token: string;41 readonly inlineSleepMs: number;42 /** 'suspended' ⇒ the lease has been released by a durable wait; every further call throws Suspend. */43 state: 'open' | 'suspended' = 'open';44 /**45 * Set when ctx.signal is aborted. No new step starts afterwards and every ctx call rethrows it, so a46 * logic that swallows one of these in a broad `catch` cannot carry on: RunCancelled | LeaseLostError |47 * WorkerStopping | KernelStorageError | NonDeterminismError | ReconcileAborted. The executor settles48 * the run from this reason, whatever the logic returns or throws.49 */50 abortReason: KernelError | null = null;5152 private readonly ac = new AbortController();53 private cursor = 0; // next seq, allocated synchronously in call order54 private readonly emitter: Emitter;55 private readonly lastRecordedSeq: number;56 private acc: Usage = zeroUsage();57 private clockMs: number;58 private ownerCursor = 0;59 private readonly startedMs: number;60 private readonly budget: RunRow['budget'];61 private readonly transient: TransientSender;6263 constructor(64 readonly k: Kernel,65 claimed: ClaimedRun,66 readonly journal: Map<number, StepRow>,67 private readonly deps: CtxDeps,68 ) {69 const r = claimed.run;70 this.runId = r.id;71 this.missionId = r.mission_id;72 this.parentRunId = r.parent_run_id;73 this.contract = r.contract as DelegationContract;74 this.workspaceAbs = r.workspace_abs;75 this.signal = this.ac.signal;76 this.token = claimed.token;77 this.inlineSleepMs = deps.inlineSleepMs;78 this.emitter = new Emitter(this, r.next_emit);79 this.budget = r.budget;80 this.startedMs = (r.started_at ?? new Date()).getTime();81 this.clockMs = this.startedMs;82 this.transient = new TransientSender(k, r.mission_id, r.id);83 let last = -1;84 for (const s of journal.values()) if ((s.status === 'completed' || s.status === 'failed') && s.seq > last) last = s.seq;85 this.lastRecordedSeq = last;86 }8788 // ---- state ---------------------------------------------------------------------------------89 get replaying(): boolean {90 return this.cursor <= this.lastRecordedSeq;91 }92 /** Seqs allocated so far (the executor checks the journal has nothing beyond it when logic returns). */93 get allocated(): number {94 return this.cursor;95 }96 abort(reason: KernelError): void {97 if (this.abortReason && !(reason instanceof LeaseLostError)) return;98 if (this.abortReason instanceof LeaseLostError) return;99 this.abortReason = reason;100 if (!this.ac.signal.aborted) this.ac.abort(reason);101 }102 assertOpen(): void {103 if (this.state === 'suspended') throw new Suspend();104 if (this.abortReason) throw this.abortReason;105 }106 markSuspended(): void {107 this.state = 'suspended';108 this.deps.onSuspend();109 }110 /** Fenced transaction; a refused write aborts the context (zombie), a seen cancel request aborts it too. */111 async fenced<T>(fn: (tx: Tx, run: RunRow) => Promise<T>, live = false): Promise<T> {112 let cancel = false;113 try {114 const out = await fencedTx(this.k, this.runId, this.token, (tx, run) => {115 cancel = run.cancel_requested;116 return fn(tx, run);117 }, live);118 if (cancel) this.abort(new RunCancelled());119 return out;120 } catch (err) {121 if (err instanceof LeaseLostError) this.abort(err);122 throw err;123 }124 }125 usage(): Usage {126 return { ...this.acc };127 }128 /**129 * Deterministic under replay: usage is the sum over the steps returned so far, and the wall clock is130 * the journal's (completed_at of the latest returned step) minus runs.started_at — it includes time131 * spent waiting, queued or crashed (nothing is excluded) but not a step still in flight.132 */133 budgetLeft(): { usd: number; turns: number; wallSec: number } {134 return {135 usd: this.budget.usd - this.acc.usd,136 turns: (this.budget.turns ?? 200) - this.acc.turns,137 wallSec: (this.budget.wallSec ?? 6 * 3600) - (this.clockMs - this.startedMs) / 1000,138 };139 }140 account(row: StepRow): void {141 if (row.usage) this.acc = addUsage(this.acc, row.usage);142 if (row.completed_at && row.completed_at.getTime() > this.clockMs) this.clockMs = row.completed_at.getTime();143 }144 applyCompleted(row: StepRow, c: Completed): unknown {145 row.status = 'completed';146 row.result = c.result;147 row.usage = c.usage;148 row.completed_at = c.completedAt;149 this.account(row);150 return canonicalize(c.result);151 }152 checkSame(row: StepRow, meta: StepMeta, hash: string): void {153 if (row.kind !== meta.kind || row.name !== meta.name || row.request_hash !== hash) {154 const what = row.request_hash !== hash && row.kind === meta.kind && row.name === meta.name ? 'request hash differs' : `journal has ${row.kind}/${row.name}`;155 throw new NonDeterminismError(this.runId, row.seq, `logic called ${meta.kind}/${meta.name}, ${what}`);156 }157 }158 rebuildError(row: StepRow): Error {159 const e: RecordedError = row.error ?? { name: 'Error', message: 'step failed' };160 if (e.name === 'BudgetExhaustedError') return new BudgetExhaustedError((e.reason as 'run_budget') ?? 'run_budget', e.message);161 return new StepError(row.seq, row.name, e);162 }163 private async hook(seq: number, meta: StepMeta, boundary: StepBoundary): Promise<void> {164 await this.deps.hooks?.boundary?.({ runId: this.runId, seq, kind: meta.kind, name: meta.name, boundary });165 }166 /** The result is in memory and the effect has happened: retry the commit before giving up. */167 private async retryStorage<T>(fn: () => Promise<T>): Promise<T> {168 for (let attempt = 0; ; attempt++) {169 try {170 return await fn();171 } catch (err) {172 if (!(err instanceof KernelStorageError) || attempt >= 3) throw err;173 await sleepMs(100 * 4 ** attempt);174 }175 }176 }177178 // ---- steps ---------------------------------------------------------------------------------179 step<Req, Res>(spec: StepSpec<Req, Res>, exec: StepExec<Res>): Promise<Res> {180 const seq = this.cursor++; // synchronous: parallel callers get their seq in call order181 return this.sticky(this.runStep(seq, spec, exec));182 }183 /** Errors that end this execution of the run are made sticky (see abortReason). */184 private sticky<T>(p: Promise<T>): Promise<T> {185 return p.catch((err: unknown) => {186 if (err instanceof KernelStorageError || err instanceof NonDeterminismError || err instanceof ReconcileAborted) this.abort(err);187 throw err;188 });189 }190191 private async runStep<Req, Res>(seq: number, spec: StepSpec<Req, Res>, exec: StepExec<Res>): Promise<Res> {192 const meta: StepMeta = { kind: spec.kind, name: spec.name, effect: spec.effect, quiet: QUIET_KINDS.has(spec.kind) };193 const requestJson = canonicalJson(spec.request);194 const hash = sha256Hex(requestJson);195 let force = false;196 for (;;) {197 this.assertOpen();198 await this.gate.shared();199 let out: { v: Res } | 'ask';200 try {201 this.assertOpen();202 out = await this.attempt(seq, spec, meta, requestJson, hash, exec, force);203 } finally {204 this.gate.releaseShared();205 }206 if (out !== 'ask') return out.v;207 const row = this.journal.get(seq) as StepRow;208 const decision = await askReconcile(this, seq, spec, meta, row); // may throw Suspend209 if (decision.t === 'value') return decision.v as Res;210 force = true; // owner: "it did not run — run it"211 }212 }213214 private async attempt<Req, Res>(215 seq: number, spec: StepSpec<Req, Res>, meta: StepMeta, requestJson: string, hash: string, exec: StepExec<Res>, force: boolean,216 ): Promise<{ v: Res } | 'ask'> {217 let row = this.journal.get(seq);218 if (row) {219 this.checkSame(row, meta, hash);220 if (row.status === 'completed') {221 this.account(row);222 return { v: canonicalize(row.result) as Res };223 }224 if (row.status === 'failed') {225 this.account(row);226 throw this.rebuildError(row);227 }228 if (row.status === 'needs_reconcile') return 'ask';229 // 'started': crash between intent and result. pure/read/model_call/idempotent_write re-execute230 // with the same idem_key; non_idempotent is never retried blindly.231 if (spec.effect === 'non_idempotent' && !force) {232 let probe: Awaited<ReturnType<NonNullable<typeof spec.reconcile>>> = { unknown: true };233 if (spec.reconcile) {234 try {235 probe = await spec.reconcile(spec.request, row.idem_key);236 } catch (err) {237 this.k.log.warn({ runId: this.runId, seq, err: (err as Error).message }, 'reconcile probe threw; asking the owner');238 }239 }240 if (probe && 'done' in probe && probe.done === true) return { v: (await this.commit(seq, row, meta, probe.result, undefined)) as Res };241 if (!probe || !('done' in probe)) return 'ask';242 }243 } else {244 await this.hook(seq, meta, 'before_intent');245 row = await this.intent(seq, meta, requestJson, hash);246 if (this.abortReason) throw this.abortReason; // cancel seen in the intent transaction: do not execute247 }248 await this.hook(seq, meta, 'after_intent');249 let out: Awaited<ReturnType<StepExec<Res>>>;250 try {251 out = await exec(row.idem_key, this.signal);252 } catch (err) {253 // Aborted (cancel / lost lease / shutdown): the step stays 'started' and is handled by effect254 // class on the next claim. Otherwise the failure is journaled and rethrown as StepError.255 if (this.abortReason) throw this.abortReason;256 if (isControlError(err)) throw err;257 return this.fail(seq, row, meta, err);258 }259 if (this.abortReason instanceof LeaseLostError) throw this.abortReason;260 await this.hook(seq, meta, 'after_exec');261 const v = (await this.commit(seq, row, meta, out?.result, out?.usage)) as Res;262 await this.hook(seq, meta, 'after_commit');263 if (this.abortReason) throw this.abortReason; // cancel seen in the commit transaction, or shutdown264 return { v };265 }266267 private async intent(seq: number, meta: StepMeta, requestJson: string, hash: string): Promise<StepRow> {268 const idemKey = idemKeyFor(this.runId, seq);269 const row: StepRow = { seq, ...meta, status: 'started', request_hash: hash, result: null, error: null, usage: null, idem_key: idemKey, completed_at: null };270 const cap = meta.effect === 'model_call' ? await this.deps.dailyCapUsd() : null;271 const refused = await this.fenced(async (tx, run) => {272 const stop = meta.effect === 'model_call' ? await budgetStop(tx, run, cap) : null;273 if (!stop) {274 await insertIntent(tx, run, seq, meta, hash, requestJson, idemKey);275 return null;276 }277 // Journaled as a failed step so that replay throws at the same seq even if the cap moved since.278 const error: RecordedError = { name: 'BudgetExhaustedError', message: `Budget exhausted (${stop})`, reason: stop };279 const at = await failStep(tx, run, seq, meta, error, { hash, requestJson, idemKey });280 return { error, at };281 }, true);282 this.journal.set(seq, row);283 if (refused) {284 row.status = 'failed';285 row.error = refused.error;286 row.completed_at = refused.at;287 this.account(row);288 throw this.rebuildError(row);289 }290 return row;291 }292293 private async commit(seq: number, row: StepRow, meta: StepMeta, result: unknown, usage: Partial<Usage> | undefined): Promise<unknown> {294 let json: string;295 try {296 json = canonicalJson(result);297 } catch (err) {298 return this.fail(seq, row, meta, err); // unserialisable result299 }300 const u = normalizeUsage(meta, usage);301 const c = await this.retryStorage(() => this.fenced((tx, run) => completeStep(tx, run, seq, meta, json, u)));302 return this.applyCompleted(row, c);303 }304305 private async fail(seq: number, row: StepRow, meta: StepMeta, err: unknown): Promise<never> {306 const error = recordError(err);307 const at = await this.retryStorage(() => this.fenced((tx, run) => failStep(tx, run, seq, meta, error)));308 row.status = 'failed';309 row.error = error;310 row.completed_at = at;311 this.account(row);312 throw this.rebuildError(row);313 }314315 wait<T = unknown>(kind: 'approval' | 'question', payload: Record<string, unknown>, opts?: { timeoutSec?: number }): Promise<T | { expired: true }> {316 const seq = this.cursor++;317 return this.sticky(durableWait(this, seq, { mode: 'wait', kind, payload, timeoutSec: opts?.timeoutSec ?? null })) as Promise<T | { expired: true }>;318 }319 sleep(ms: number): Promise<void> {320 const seq = this.cursor++;321 return this.sticky(durableWait(this, seq, { mode: 'timer', ms: Math.max(0, Math.floor(ms)) })).then(() => undefined);322 }323 async now(): Promise<Date> {324 const r = await this.step({ kind: 'now', name: 'now', effect: 'pure', request: {} }, async () => ({ result: { iso: new Date().toISOString() } }));325 return new Date(r.iso);326 }327 async uuid(): Promise<string> {328 const r = await this.step({ kind: 'uuid', name: 'uuid', effect: 'pure', request: {} }, async () => ({ result: { uuid: randomUUID() } }));329 return r.uuid;330 }331332 /** Unread owner.message events, journaled: the cursor (`afterId`) is rebuilt from earlier results on replay. */333 async ownerMessages(): Promise<{ id: number; text: string }[]> {334 const afterId = this.ownerCursor;335 const r = await this.step({ kind: 'owner_messages', name: 'owner_messages', effect: 'read', request: { afterId } }, async () => {336 const rows = await query<{ id: unknown; text: string | null }>(337 this.k,338 `SELECT id, payload->>'text' AS text FROM events WHERE run_id = $1 AND kind = 'owner.message' AND id > $2 ORDER BY id`,339 [this.runId, afterId],340 );341 return { result: { messages: rows.map((m) => ({ id: Number(m.id), text: m.text ?? '' })) } };342 });343 for (const m of r.messages) if (m.id > this.ownerCursor) this.ownerCursor = m.id;344 return r.messages;345 }346347 emit(kind: EventKind, payload: Record<string, unknown>, trust: Trust = 'agent'): Promise<void> {348 if (kind === 'model.delta') {349 this.emitTransient(kind, payload);350 return Promise.resolve();351 }352 return this.sticky(this.emitter.emit(kind, payload, trust));353 }354 emitTransient(kind: 'model.delta', payload: Record<string, unknown>): void {355 if (this.state === 'open' && !this.abortReason) this.transient.send(kind, payload);356 }357 /** Resolves when no step, emit or delta of this context is still in flight. */358 async settle(): Promise<void> {359 await this.emitter.idle();360 await this.gate.idle();361 await this.emitter.idle();362 await this.transient.idle();363 }364}365