Repository Public
1import { afterAll, beforeAll, describe, expect, it } from 'vitest';2import { claimRun, executeRun, LeaseLostError, type ExecuteOptions, type RunLogic } from '../src/index.js';3import { createTestDb, getRun, newMission, 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 execOpts = (logic: RunLogic, over: Partial<ExecuteOptions> = {}): ExecuteOptions => ({15 logic, leaseMs: 400, heartbeatMs: 100, inlineSleepMs: 100, dailyCapUsd: async () => null, ...over,16});17function deferred<T = void>(): { promise: Promise<T>; resolve: (v: T) => void } {18 let resolve!: (v: T) => void;19 const promise = new Promise<T>((r) => (resolve = r));20 return { promise, resolve };21}2223describe('leases and fencing', () => {24 it('two workers racing to claim: exactly one wins, and nothing is claimable twice', async () => {25 const { rootRunId } = await newMission(db);26 const claims = await Promise.all(Array.from({ length: 8 }, (_, i) => claimRun(`racer-${i}`, { kernel: db.k, leaseMs: 60_000 })));27 const won = claims.filter((c) => c !== null);28 expect(won).toHaveLength(1);29 expect(won[0]?.run.id).toBe(rootRunId);30 expect(won[0]?.token).toBe('1');31 expect(await claimRun('late', { kernel: db.k })).toBeNull();32 // Clean up for the following tests (which use workers that would otherwise wait for this lease).33 await db.k.pool.query(`UPDATE runs SET status = 'cancelled' WHERE id = $1`, [rootRunId]);34 });3536 it('lease expiry: another worker takes over, and the zombie’s commit is refused by fencing', async () => {37 const { missionId, rootRunId } = await newMission(db);38 const release = deferred();39 const zombieInExec = deferred();40 let execs = 0;41 const logic: RunLogic = async (ctx) => {42 const r = await ctx.step({ kind: 'model_call', name: 'm', effect: 'model_call', request: {} }, async () => {43 const n = ++execs;44 if (n === 1) {45 zombieInExec.resolve();46 await release.promise; // the zombie is stuck here, well past its lease47 }48 return { result: { by: n === 1 ? 'zombie' : 'successor' }, usage: { usd: 0.01 } };49 });50 return done(r);51 };52 // Zombie: its heartbeat never fires (interval ≫ lease), so its 400 ms lease expires while exec hangs.53 const claimed = await claimRun('zombie', { kernel: db.k, leaseMs: 400 });54 expect(claimed?.run.id).toBe(rootRunId);55 const zombie = executeRun(db.k, claimed!, execOpts(logic, { heartbeatMs: 3_600_000 }));56 await zombieInExec.promise;57 const w = startWorker(db, logic);58 const run = await waitStatus(db, rootRunId, ['done']);59 expect(run.result?.payload).toEqual({ by: 'successor' });60 release.resolve(); // the zombie wakes up and tries to commit its result61 expect(await zombie.done).toBe('lease_lost');62 await w.stop();63 const after = await getRun(db, rootRunId);64 expect(after.status).toBe('done');65 expect(after.result?.payload).toEqual({ by: 'successor' });66 expect((await steps(db, rootRunId))[0]?.result).toEqual({ by: 'successor' });67 expect(after.usage.usd).toBeCloseTo(0.01); // accounted once, not twice68 const m = await rows<{ spent_usd: string }>(db, 'SELECT spent_usd FROM missions WHERE id = $1', [missionId]);69 expect(Number(m[0]?.spent_usd)).toBeCloseTo(0.01);70 expect(execs).toBe(2);71 expect(await rows(db, `SELECT 1 FROM events WHERE run_id = $1 AND kind = 'step.completed'`, [rootRunId])).toHaveLength(1);72 });7374 it('a zombie cannot write intents, emits or waits either (LeaseLostError)', async () => {75 const { rootRunId } = await newMission(db);76 const go = deferred();77 const errors: string[] = [];78 const logic: RunLogic = async (ctx) => {79 await go.promise;80 for (const op of [81 () => ctx.step({ kind: 'calc', name: 'a', effect: 'pure', request: {} }, async () => ({ result: {} })),82 () => ctx.emit('system.notice', {}),83 () => ctx.wait('question', { text: '?' }),84 () => ctx.sleep(5000),85 ]) {86 try {87 await op();88 errors.push('no error');89 } catch (err) {90 errors.push(err instanceof LeaseLostError ? 'LeaseLostError' : String(err));91 }92 }93 return done();94 };95 const claimed = await claimRun('zombie2', { kernel: db.k, leaseMs: 60_000 });96 const zombie = executeRun(db.k, claimed!, execOpts(logic, { leaseMs: 60_000, heartbeatMs: 3_600_000 }));97 await db.k.pool.query('UPDATE runs SET fencing_token = fencing_token + 1 WHERE id = $1', [rootRunId]); // someone else claimed98 go.resolve();99 expect(await zombie.done).toBe('lease_lost');100 expect(errors).toEqual(['LeaseLostError', 'LeaseLostError', 'LeaseLostError', 'LeaseLostError']);101 expect(await steps(db, rootRunId)).toHaveLength(0);102 expect(await rows(db, 'SELECT 1 FROM waits WHERE run_id = $1', [rootRunId])).toHaveLength(0);103 await db.k.pool.query(`UPDATE runs SET status = 'cancelled' WHERE id = $1`, [rootRunId]);104 });105106 it('heartbeat failure mid-step aborts ctx.signal; the step stays started and is resumed', async () => {107 const { rootRunId } = await newMission(db);108 let aborted: unknown = null;109 let execs = 0;110 const logic: RunLogic = async (ctx) => {111 const r = await ctx.step({ kind: 'tool_call', name: 'slow', effect: 'read', request: {} }, async (_idem, signal) => {112 if (++execs === 1) {113 await new Promise<void>((resolve) => signal.addEventListener('abort', () => resolve(), { once: true }));114 aborted = signal.reason;115 throw new Error('aborted');116 }117 return { result: { n: execs } };118 });119 return done(r);120 };121 const claimed = await claimRun('hb', { kernel: db.k, leaseMs: 400 });122 const ex = executeRun(db.k, claimed!, execOpts(logic));123 await waitFor(async () => execs === 1, 5000, 'exec started');124 await sleep(250); // a few successful heartbeats keep the lease alive125 expect((await getRun(db, rootRunId)).status).toBe('running');126 await db.k.pool.query('UPDATE runs SET fencing_token = fencing_token + 1, lease_expires_at = now() WHERE id = $1', [rootRunId]);127 expect(await ex.done).toBe('lease_lost');128 expect(aborted).toBeInstanceOf(LeaseLostError);129 expect((await steps(db, rootRunId))[0]?.status).toBe('started'); // no failure recorded by the zombie130 const w = startWorker(db, logic);131 const run = await waitStatus(db, rootRunId, ['done']);132 await w.stop();133 expect(run.result?.payload).toEqual({ n: 2 });134 });135136 it('a worker stalled past its lease may record finished work but may not start a new step', async () => {137 const { rootRunId } = await newMission(db);138 const count = { one: 0, two: 0 };139 const logic: RunLogic = async (ctx) => {140 await ctx.step({ kind: 'tool_call', name: 'one', effect: 'non_idempotent', request: {} }, async () => {141 if (++count.one === 1) await sleep(500); // outlives the 200 ms lease; nobody else claims meanwhile142 return { result: {} };143 });144 await ctx.step({ kind: 'tool_call', name: 'two', effect: 'non_idempotent', request: {} }, async () => ({ result: { n: ++count.two } }));145 return done();146 };147 const claimed = await claimRun('stalled', { kernel: db.k, leaseMs: 200 });148 const stalled = executeRun(db.k, claimed!, execOpts(logic, { leaseMs: 200, heartbeatMs: 3_600_000 }));149 expect(await stalled.done).toBe('lease_lost');150 expect((await steps(db, rootRunId)).map((s) => `${s.name}:${s.status}`)).toEqual(['one:completed']); // result kept, no new intent151 expect(count).toEqual({ one: 1, two: 0 });152 const w = startWorker(db, logic);153 await waitStatus(db, rootRunId, ['done']);154 await w.stop();155 expect(count).toEqual({ one: 1, two: 1 });156 });157158 it('heartbeats keep a long step alive beyond the lease duration', async () => {159 const { rootRunId } = await newMission(db);160 let execs = 0;161 const logic: RunLogic = async (ctx) => {162 await ctx.step({ kind: 'tool_call', name: 'long', effect: 'read', request: {} }, async () => {163 execs++;164 await sleep(1500); // ≈ 4 leases165 return { result: {} };166 });167 return done();168 };169 const a = startWorker(db, logic, { leaseMs: 400, heartbeatMs: 100 });170 const b = startWorker(db, logic, { leaseMs: 400, heartbeatMs: 100 });171 const run = await waitStatus(db, rootRunId, ['done']);172 await Promise.all([a.stop(), b.stop()]);173 expect(execs).toBe(1);174 expect(run.fencing_token).toBe('1');175 });176177 it('graceful stop leaves the run running and resumable; an in-flight step still commits', async () => {178 const { rootRunId } = await newMission(db);179 const inSecond = deferred();180 const count = { first: 0, second: 0, third: 0 };181 let stopping = false;182 const logic: RunLogic = async (ctx) => {183 await ctx.step({ kind: 'tool_call', name: 'first', effect: 'non_idempotent', request: {} }, async () => ({ result: { n: ++count.first } }));184 await ctx.step({ kind: 'tool_call', name: 'second', effect: 'non_idempotent', request: {} }, async () => {185 count.second++;186 inSecond.resolve();187 while (!stopping) await sleep(10);188 await sleep(100); // ignores the abort signal and finishes: its result must not be lost189 return { result: { ok: true } };190 });191 await ctx.step({ kind: 'calc', name: 'third', effect: 'pure', request: {} }, async () => ({ result: { n: ++count.third } }));192 return done();193 };194 const w1 = startWorker(db, logic);195 await inSecond.promise;196 stopping = true;197 await w1.stop();198 const mid = await getRun(db, rootRunId);199 expect(mid.status).toBe('running'); // never marked failed200 expect(mid.lease_owner).toBeNull();201 const st = await steps(db, rootRunId);202 expect(st.map((s) => s.status)).toEqual(['completed', 'completed']); // third never started203 const w2 = startWorker(db, logic);204 await waitStatus(db, rootRunId, ['done']);205 await w2.stop();206 expect(count).toEqual({ first: 1, second: 1, third: 1 });207 });208209 it('the worker loop records its liveness in `workers`', async () => {210 const w = startWorker(db, async () => done(), { workerId: 'liveness-test', version: '9.9.9' });211 const r = await waitFor(async () => (await rows<{ version: string; fresh: boolean }>(db, `SELECT version, last_seen > now() - interval '5 seconds' AS fresh FROM workers WHERE id = 'liveness-test'`))[0], 5000, 'workers row');212 await w.stop();213 expect(r).toEqual({ version: '9.9.9', fresh: true });214 });215});216