diff --git a/src/server/runner.ts b/src/server/runner.ts index c9de2d66..c61e775a 100644 --- a/src/server/runner.ts +++ b/src/server/runner.ts @@ -5,7 +5,10 @@ import { research, type Config } from './research.js'; import { DEFAULT_AGENT_RUN_TIMEOUT_MS } from './platform-config.js'; export class Runner { private timer?: ReturnType; - private active = new Map(); + private active = new Map< + string, + { claim: Claim; controller: AbortController } + >(); constructor( private store: Store, private config: Config, @@ -26,20 +29,18 @@ export class Runner { stop() { clearInterval(this.timer); this.timer = undefined; - for (const task of this.store.tasks()) { - if (this.active.has(task.id) && task.lease) - this.store.interrupt( - { ...task, lease: task.lease }, - 'Server stopped during this run. Review completed effects before retrying.', - ); - } + for (const { claim } of this.active.values()) + this.store.interrupt( + claim, + 'Server stopped during this run. Review completed effects before retrying.', + ); this.abortAll(); } abort(id: string) { - this.active.get(id)?.abort(new Error('Run stopped.')); + this.active.get(id)?.controller.abort(new Error('Run stopped.')); } abortAll() { - for (const controller of this.active.values()) + for (const { controller } of this.active.values()) controller.abort(new Error('Run stopped because settings changed.')); } async tick() { @@ -62,7 +63,7 @@ export class Runner { ); if (!claim) return; const controller = new AbortController(); - this.active.set(claim.id, controller); + this.active.set(claim.id, { claim, controller }); const ownershipCheck = setInterval(() => { try { if (!this.store.owns(claim)) diff --git a/tests/runner.test.ts b/tests/runner.test.ts index 3de211f4..4d3d3002 100644 --- a/tests/runner.test.ts +++ b/tests/runner.test.ts @@ -1,4 +1,7 @@ import { afterEach, expect, it, vi } from 'vitest'; +import { mkdtempSync, rmSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; import { Store } from '../src/server/store.js'; import { Runner } from '../src/server/runner.js'; import { research, type Config } from '../src/server/research.js'; @@ -135,6 +138,72 @@ it('holds active work for review on graceful shutdown', async () => { expect(store.claim()).toBeTruthy(); store.close(); }); +it('does not interrupt a replacement claim when an old runner stops', async () => { + const dir = mkdtempSync(join(tmpdir(), 'opendots-runner-')); + const store = new Store(join(dir, 'test.sqlite')); + const replacementStore = new Store(join(dir, 'test.sqlite')); + const result = { text: 'Finished', sources: [], sample: false }; + let finishOld!: () => void; + let finishReplacement!: () => void; + const executeOld = vi.fn( + () => + new Promise( + (resolve) => (finishOld = () => resolve(result)), + ), + ); + const executeReplacement = vi.fn( + () => + new Promise( + (resolve) => (finishReplacement = () => resolve(result)), + ), + ); + const oldRunner = new Runner(store, config, executeOld); + const replacementRunner = new Runner( + replacementStore, + config, + executeReplacement, + ); + const task = store.createTask('Keep working'); + const oldRun = oldRunner.tick(); + let replacementRun: Promise | undefined; + try { + expect(executeOld).toHaveBeenCalledOnce(); + const oldLease = store.task(task.id)!.lease; + replacementStore.action(task.id, 'pause'); + replacementStore.action(task.id, 'run'); + replacementRun = replacementRunner.tick(); + expect(executeReplacement).toHaveBeenCalledOnce(); + const replacementLease = replacementStore.task(task.id)!.lease; + expect(replacementLease).not.toBe(oldLease); + + oldRunner.stop(); + + expect(replacementStore.task(task.id)).toMatchObject({ + status: 'running', + lease: replacementLease, + }); + finishOld(); + await oldRun; + finishReplacement(); + await replacementRun; + expect(replacementStore.task(task.id)?.status).toBe('completed'); + expect(replacementStore.detail(task.id)?.runs[0]).toMatchObject({ + id: replacementLease, + status: 'completed', + result, + }); + } finally { + finishOld?.(); + finishReplacement?.(); + await oldRun; + await replacementRun; + oldRunner.stop(); + replacementRunner.stop(); + store.close(); + replacementStore.close(); + rmSync(dir, { recursive: true, force: true }); + } +}); it('survives a claim failure and runs work on the next tick', async () => { const store = new Store(':memory:'); const runner = new Runner(store, { mode: 'sample', baseUrl: '' });