Repository Public
1import { afterAll, beforeAll, describe, expect, it } from 'vitest';2import { createRun, expireWaits, postOwnerMessage, requestCancel, resolveWait, type RunLogic } from '../src/index.js';3import { createTestDb, getRun, newMission, pendingWait, rows, sleep, startWorker, steps, waitFor, waitStatus, type TestDb } from './helpers.js';45let db: TestDb;6beforeAll(async () => {7 db = await createTestDb();8});9afterAll(async () => {10 await db.cleanup();11});1213const done = (payload: unknown = null) => ({ status: 'done' as const, summary: 'ok', payload, artifacts: [], evidence: [], openQuestions: [] });14const missionStatus = async (id: string): Promise<string> => (await rows<{ status: string }>(db, 'SELECT status FROM missions WHERE id = $1', [id]))[0]?.status ?? '';15const eventKinds = async (runId: string, kind: string): Promise<Record<string, unknown>[]> =>16 (await rows<{ payload: Record<string, unknown> }>(db, 'SELECT payload FROM events WHERE run_id = $1 AND kind = $2 ORDER BY id', [runId, kind])).map((e) => e.payload);1718/** The run is suspended on a timer: queued with a wake_at in the future. */19const sleeping = (runId: string): Promise<unknown> =>20 waitFor(async () => (await rows(db, `SELECT 1 FROM runs WHERE id = $1 AND status = 'queued' AND wake_at IS NOT NULL`, [runId])).length > 0, 10_000, 'run sleeping');2122describe('durable waits', () => {23 it('a question survives worker restarts; the answer is returned on replay', async () => {24 const { missionId, rootRunId } = await newMission(db);25 const count = { a: 0, b: 0 };26 const logic: RunLogic = async (ctx) => {27 await ctx.step({ kind: 'calc', name: 'a', effect: 'pure', request: {} }, async () => ({ result: { n: ++count.a } }));28 const ans = await ctx.wait<{ answer: string }>('question', { text: 'Which colour?', options: ['red', 'blue'] });29 const b = await ctx.step({ kind: 'calc', name: 'b', effect: 'pure', request: { ans } }, async () => ({ result: { n: ++count.b, ans } }));30 return done(b);31 };32 const w1 = startWorker(db, logic);33 const waiting = await waitStatus(db, rootRunId, ['waiting']);34 await w1.stop(); // no worker at all while the owner thinks35 expect(waiting.lease_owner).toBeNull();36 expect(await missionStatus(missionId)).toBe('waiting');37 const w = await pendingWait(db, rootRunId);38 expect(w?.kind).toBe('question');39 expect(w?.payload).toEqual({ text: 'Which colour?', options: ['red', 'blue'] });40 expect(w?.step_seq).toBe(1);41 const res = await resolveWait(w!.id, { answer: 'blue' }, { name: 'owner' }, db.k);42 expect(res).toEqual({ resolved: true, status: 'resolved', runId: rootRunId });43 expect((await getRun(db, rootRunId)).status).toBe('queued');44 expect(await resolveWait(w!.id, { answer: 'red' }, 'owner', db.k)).toMatchObject({ resolved: false, status: 'resolved' }); // idempotent45 await expect(resolveWait('00000000-0000-4000-8000-000000000000', {}, 'owner', db.k)).rejects.toThrow(/not found/);46 const w2 = startWorker(db, logic);47 const run = await waitStatus(db, rootRunId, ['done']);48 await w2.stop();49 expect(run.result?.payload).toEqual({ n: 1, ans: { answer: 'blue' } });50 expect(count).toEqual({ a: 1, b: 1 });51 expect((await eventKinds(rootRunId, 'wait.created'))).toHaveLength(1);52 expect((await eventKinds(rootRunId, 'wait.resolved')).map((p) => p.status)).toEqual(['resolved']);53 expect((await steps(db, rootRunId))[1]).toMatchObject({ kind: 'wait', name: 'question', status: 'completed', result: { answer: 'blue' } });54 });5556 it('a resolution arriving while the run is still running is not lost (no suspension at all)', async () => {57 const { rootRunId } = await newMission(db);58 const logic: RunLogic = async (ctx) => {59 // The wait row is inserted at once; the suspension must wait for the slow sibling step to settle.60 const [ans] = await Promise.all([61 ctx.wait<{ decision: string }>('approval', { tool: 'shell', input: { command: 'ls' } }),62 ctx.step({ kind: 'tool_call', name: 'slow', effect: 'read', request: {} }, async () => {63 await sleep(600);64 return { result: {} };65 }),66 ]);67 return done(ans);68 };69 const w = startWorker(db, logic);70 const wait = await waitFor(() => pendingWait(db, rootRunId), 5000, 'pending wait');71 expect((await getRun(db, rootRunId)).status).toBe('running'); // window between insert and suspend72 await resolveWait(wait.id, { decision: 'approve' }, 'owner', db.k);73 const run = await waitStatus(db, rootRunId, ['done']);74 await w.stop();75 expect(run.result?.payload).toEqual({ decision: 'approve' });76 expect(run.fencing_token).toBe('1'); // one claim: it never suspended77 expect((await eventKinds(rootRunId, 'run.status')).map((p) => p.status)).toEqual(['running', 'done']);78 });7980 it('a crash between inserting the wait and suspending is recovered (run left running with a pending wait)', async () => {81 const { rootRunId } = await newMission(db);82 const logic: RunLogic = async (ctx) => done(await ctx.wait('question', { text: 'q' }));83 // Simulate the crash: the state a worker leaves when killed right after the insert transaction.84 await db.k.pool.query(`UPDATE runs SET status = 'running', fencing_token = 1, lease_owner = 'dead', lease_expires_at = now() - interval '1 second', started_at = now() WHERE id = $1`, [rootRunId]);85 await db.k.pool.query(86 `INSERT INTO steps (run_id, seq, kind, name, effect, status, request_hash, request, idem_key)87 VALUES ($1, 0, 'wait', 'question', 'pure', 'started', encode(sha256(convert_to('{"kind":"question","payload":{"text":"q"}}', 'UTF8')), 'hex'), '{}', 'k')`, [rootRunId]);88 const ins = await rows<{ id: string }>(db, `INSERT INTO waits (run_id, step_seq, mission_id, kind, payload) SELECT id, 0, mission_id, 'question', '{"text":"q"}' FROM runs WHERE id = $1 RETURNING id`, [rootRunId]);89 const w = startWorker(db, logic);90 await waitStatus(db, rootRunId, ['waiting']);91 await resolveWait(ins[0]!.id, { answer: 'yes' }, 'owner', db.k);92 const run = await waitStatus(db, rootRunId, ['done']);93 await w.stop();94 expect(run.result?.payload).toEqual({ answer: 'yes' });95 });9697 it('wait expiry: the janitor expires it, the run resumes and wait() returns {expired:true}', async () => {98 const { rootRunId } = await newMission(db);99 const logic: RunLogic = async (ctx) => done(await ctx.wait('approval', { tool: 'git_push' }, { timeoutSec: 0.4 }));100 const w = startWorker(db, logic);101 await waitStatus(db, rootRunId, ['waiting']);102 expect(await expireWaits(db.k)).toEqual({ expired: 0 }); // not due yet103 await sleep(450);104 expect(await expireWaits(db.k)).toEqual({ expired: 1 });105 const run = await waitStatus(db, rootRunId, ['done']);106 await w.stop();107 expect(run.result?.payload).toEqual({ expired: true });108 expect((await eventKinds(rootRunId, 'wait.resolved')).map((p) => p.status)).toEqual(['expired']);109 expect(await expireWaits(db.k)).toEqual({ expired: 0 });110 });111112 it('the worker loop runs the janitor itself', async () => {113 const { rootRunId } = await newMission(db);114 const w = startWorker(db, async (ctx) => done(await ctx.wait('question', { text: 'q' }, { timeoutSec: 0.3 })), { janitorMs: 100 });115 const run = await waitStatus(db, rootRunId, ['done']);116 await w.stop();117 expect(run.result?.payload).toEqual({ expired: true });118 });119});120121describe('timers', () => {122 it('sleep survives a worker restart and does not re-execute earlier steps', async () => {123 const { missionId, rootRunId } = await newMission(db);124 const count = { a: 0, b: 0 };125 const logic: RunLogic = async (ctx) => {126 await ctx.step({ kind: 'calc', name: 'a', effect: 'pure', request: {} }, async () => ({ result: { n: ++count.a } }));127 await ctx.sleep(900);128 await ctx.step({ kind: 'calc', name: 'b', effect: 'pure', request: {} }, async () => ({ result: { n: ++count.b } }));129 return done();130 };131 const t0 = Date.now();132 const w1 = startWorker(db, logic);133 await sleeping(rootRunId);134 await w1.stop();135 const r = await rows<{ wake_in: string; lease_owner: string | null }>(db, 'SELECT EXTRACT(EPOCH FROM (wake_at - now())) AS wake_in, lease_owner FROM runs WHERE id = $1', [rootRunId]);136 expect(Number(r[0]?.wake_in)).toBeGreaterThan(0.3);137 expect(r[0]?.lease_owner).toBeNull();138 expect(await missionStatus(missionId)).toBe('queued');139 const w2 = startWorker(db, logic);140 await sleep(200);141 expect((await getRun(db, rootRunId)).status).toBe('queued'); // not claimable before wake_at142 await waitStatus(db, rootRunId, ['done']);143 await w2.stop();144 expect(Date.now() - t0).toBeGreaterThanOrEqual(900);145 expect(count).toEqual({ a: 1, b: 1 });146 expect((await steps(db, rootRunId))[1]).toMatchObject({ kind: 'timer', status: 'completed' });147 });148149 it('a short sleep is waited in-process (no suspension)', async () => {150 const { rootRunId } = await newMission(db);151 const t0 = Date.now();152 const w = startWorker(db, async (ctx) => {153 await ctx.sleep(60);154 return done();155 });156 const run = await waitStatus(db, rootRunId, ['done']);157 await w.stop();158 expect(Date.now() - t0).toBeGreaterThanOrEqual(60);159 expect(run.fencing_token).toBe('1');160 });161});162163describe('cancellation', () => {164 it('cancel while waiting: cancelled immediately, wait cancelled, mission mirrored', async () => {165 const { missionId, rootRunId } = await newMission(db);166 const w = startWorker(db, async (ctx) => done(await ctx.wait('question', { text: 'q' })));167 await waitStatus(db, rootRunId, ['waiting']);168 const r = await requestCancel(missionId, db.k);169 expect(r).toEqual({ cancelled: [rootRunId], requested: [] });170 const run = await getRun(db, rootRunId);171 expect(run.status).toBe('cancelled');172 expect(run.result?.status).toBe('cancelled');173 expect(await missionStatus(missionId)).toBe('cancelled');174 expect(await rows(db, `SELECT 1 FROM waits WHERE run_id = $1 AND status = 'cancelled'`, [rootRunId])).toHaveLength(1);175 await sleep(200);176 await w.stop();177 expect((await getRun(db, rootRunId)).status).toBe('cancelled'); // nobody resurrected it178 expect(await requestCancel(missionId, db.k)).toEqual({ cancelled: [], requested: [] });179 });180181 it('cancel while running: ctx.signal is aborted promptly and the run ends cancelled', async () => {182 const { missionId, rootRunId } = await newMission(db);183 let abortedAfterMs = -1;184 let started = 0;185 const w = startWorker(db, async (ctx) => {186 await ctx.step({ kind: 'model_call', name: 'm', effect: 'model_call', request: {} }, async (_idem, signal) => {187 started = Date.now();188 await new Promise<void>((resolve) => signal.addEventListener('abort', () => resolve(), { once: true }));189 abortedAfterMs = Date.now() - started;190 throw new Error('aborted by signal');191 });192 return done();193 }, { heartbeatMs: 60_000, leaseMs: 180_000 }); // only the cancel doorbell can make this fast194 await waitFor(async () => started > 0, 5000, 'exec started');195 const r = await requestCancel(rootRunId, db.k); // by run id196 expect(r).toEqual({ cancelled: [], requested: [rootRunId] });197 const run = await waitStatus(db, rootRunId, ['cancelled'], 5000);198 await w.stop();199 expect(abortedAfterMs).toBeGreaterThanOrEqual(0);200 expect(abortedAfterMs).toBeLessThan(3000);201 expect(run.result?.status).toBe('cancelled');202 expect(await missionStatus(missionId)).toBe('cancelled');203 expect((await steps(db, rootRunId))[0]?.status).toBe('started'); // not recorded as a failure204 });205206 it('cancel is also seen at the next step boundary when the step ignores the signal', async () => {207 const { rootRunId } = await newMission(db);208 let inExec = false;209 let second = 0;210 const w = startWorker(db, async (ctx) => {211 await ctx.step({ kind: 'tool_call', name: 'deaf', effect: 'read', request: {} }, async () => {212 inExec = true;213 await sleep(500);214 return { result: {} };215 });216 await ctx.step({ kind: 'calc', name: 'never', effect: 'pure', request: {} }, async () => ({ result: { n: ++second } }));217 return done();218 }, { heartbeatMs: 60_000, leaseMs: 180_000 });219 await waitFor(async () => inExec, 5000, 'exec started');220 await requestCancel(rootRunId, db.k);221 await waitStatus(db, rootRunId, ['cancelled']);222 await w.stop();223 expect(second).toBe(0);224 });225226 it('cancel a queued run, and cascade to descendants', async () => {227 const { missionId, rootRunId } = await newMission(db);228 const child = await createRun(rootRunId, { objective: 'sub-task', doneCriteria: ['done'] }, { provider: 'openai', model: 'gpt-6-luna', budget: { usd: 1 } }, db.k);229 const grandchild = await createRun(child.runId, { objective: 'sub-sub-task', doneCriteria: ['done'] }, { provider: 'openai', model: 'gpt-6-luna', budget: { usd: 1 } }, db.k);230 await expect(createRun(rootRunId, { objective: 'x', doneCriteria: [] }, { provider: 'openai', model: 'm', budget: { usd: 1 } }, db.k)).rejects.toThrow(); // no done criteria231 const r = await requestCancel(child.runId, db.k); // a sub-tree only232 expect(r.cancelled.sort()).toEqual([child.runId, grandchild.runId].sort());233 expect((await getRun(db, rootRunId)).status).toBe('queued');234 expect(await missionStatus(missionId)).toBe('queued'); // the mission mirrors the ROOT run only235 const all = await requestCancel(missionId, db.k);236 expect(all.cancelled).toEqual([rootRunId]);237 expect(await missionStatus(missionId)).toBe('cancelled');238 await expect(createRun(rootRunId, { objective: 'x', doneCriteria: ['y'] }, { provider: 'openai', model: 'm', budget: { usd: 1 } }, db.k)).rejects.toThrow(/cancelled/);239 await expect(postOwnerMessage(rootRunId, 'too late', db.k)).rejects.toThrow(/cancelled/);240 });241});242243describe('owner messages', () => {244 it('are journaled: each message is delivered once and replays return the recorded batches', async () => {245 const { rootRunId } = await newMission(db);246 await postOwnerMessage(rootRunId, 'first', db.k);247 const batches: string[] = [];248 const logic: RunLogic = async (ctx) => {249 const m1 = await ctx.ownerMessages();250 await ctx.wait('question', { text: 'continue?' });251 const m2 = await ctx.ownerMessages();252 await ctx.sleep(700);253 const m3 = await ctx.ownerMessages();254 batches.push(JSON.stringify([m1.map((m) => m.text), m2.map((m) => m.text), m3.map((m) => m.text)]));255 return done({ m1: m1.map((m) => m.text), m2: m2.map((m) => m.text), m3: m3.map((m) => m.text) });256 };257 const w = startWorker(db, logic);258 await waitStatus(db, rootRunId, ['waiting']);259 await postOwnerMessage(rootRunId, 'second', db.k);260 await postOwnerMessage(rootRunId, 'third', db.k);261 const wait = await pendingWait(db, rootRunId);262 await resolveWait(wait!.id, { answer: 'yes' }, 'owner', db.k);263 await sleeping(rootRunId);264 await postOwnerMessage(rootRunId, 'fourth', db.k);265 const run = await waitStatus(db, rootRunId, ['done']);266 await w.stop();267 expect(run.result?.payload).toEqual({ m1: ['first'], m2: ['second', 'third'], m3: ['fourth'] });268 expect(batches).toHaveLength(1);269 const st = await steps(db, rootRunId);270 expect(st.filter((s) => s.kind === 'owner_messages')).toHaveLength(3);271 expect(await rows(db, `SELECT 1 FROM events WHERE run_id = $1 AND kind = 'owner.message' AND trust = 'owner'`, [rootRunId])).toHaveLength(4);272 });273});274