diff --git a/api/server/controllers/agents/__tests__/resume.spec.js b/api/server/controllers/agents/__tests__/resume.spec.js index a52181feb0..ce50537eb8 100644 --- a/api/server/controllers/agents/__tests__/resume.spec.js +++ b/api/server/controllers/agents/__tests__/resume.spec.js @@ -50,6 +50,10 @@ const mockGenerationJobManager = { getJobStore: jest.fn(() => mockJobStore), getResumeState: jest.fn(), setContentParts: jest.fn(), + /** Resume moves ownership: the new owner records its own seal capability + * and rebuilds armed interrupts from the durable queue. */ + updateMetadata: jest.fn().mockResolvedValue(undefined), + rearmQueuedPreempts: jest.fn().mockResolvedValue(0), emitChunk: jest.fn(), emitDone: jest.fn(), emitError: jest.fn(), diff --git a/packages/api/src/agents/steering/__tests__/runtime.spec.ts b/packages/api/src/agents/steering/__tests__/runtime.spec.ts index f4188b4f9e..a51242fd6e 100644 --- a/packages/api/src/agents/steering/__tests__/runtime.spec.ts +++ b/packages/api/src/agents/steering/__tests__/runtime.spec.ts @@ -409,6 +409,64 @@ describe('createSteerPreemptBoundaryHook', () => { expect(GenerationJobManager.isPreemptRequested(streamId)).toBe(false); }); + /** + * A cancel whose cross-replica clear was lost leaves a stale arm. If the + * next boundary drains a DIFFERENT steer, clearing only the drained id + * would leave the stale one level-triggered — it would immediately seal + * the continuation meant to answer the steer just injected, landing on an + * empty boundary as `preempt_incomplete`. + */ + it('a nonempty drain also clears stale arms held since the snapshot', async () => { + const streamId = `preempt-stale-snapshot-${Date.now()}`; + const job = await GenerationJobManager.createJob(streamId, 'user-1'); + + /** Stale: armed, but its steer never reaches the queue (cancelled). */ + await GenerationJobManager.requestPreempt(streamId, 'steer-cancelled', job.createdAt); + /** Live: queued and armed, and this is what the boundary will drain. */ + await GenerationJobManager.steering.enqueue(streamId, { + ...buildSteer('steer-live', 'interrupt me'), + preempt: true, + }); + await GenerationJobManager.requestPreempt(streamId, 'steer-live', job.createdAt); + expect(GenerationJobManager.isPreemptRequested(streamId)).toBe(true); + + const hook = createSteerPreemptBoundaryHook({ + streamId, + jobCreatedAt: job.createdAt, + applySteer: jest.fn(), + }); + const output: SteerDrainOutput = await hook(boundaryInput(), abortSignal); + + expect(output.injectedMessages).toHaveLength(1); + /** Both the drained id and the stale snapshot id are spent. */ + expect(GenerationJobManager.isPreemptRequested(streamId)).toBe(false); + }); + + /** An arm that lands AFTER the snapshot is backed by a live queue item. */ + it('a nonempty drain spares an arm that landed after the snapshot', async () => { + const streamId = `preempt-post-snapshot-${Date.now()}`; + const job = await GenerationJobManager.createJob(streamId, 'user-1'); + await GenerationJobManager.steering.enqueue(streamId, { + ...buildSteer('steer-first', 'first'), + preempt: true, + }); + await GenerationJobManager.requestPreempt(streamId, 'steer-first', job.createdAt); + + const hook = createSteerPreemptBoundaryHook({ + streamId, + jobCreatedAt: job.createdAt, + applySteer: async () => { + /** Arrives mid-drain, after the snapshot was taken. */ + await GenerationJobManager.requestPreempt(streamId, 'steer-later', job.createdAt); + }, + }); + await hook(boundaryInput(), abortSignal); + + expect(GenerationJobManager.getArmedPreemptIds(streamId, job.createdAt)).toEqual([ + 'steer-later', + ]); + }); + it('clears the request even when applySteer throws mid-drain', async () => { const streamId = `preempt-clears-on-error-${Date.now()}`; const job = await GenerationJobManager.createJob(streamId, 'user-1'); diff --git a/packages/api/src/agents/steering/runtime.ts b/packages/api/src/agents/steering/runtime.ts index 853a3a0ca0..aafde3b84b 100644 --- a/packages/api/src/agents/steering/runtime.ts +++ b/packages/api/src/agents/steering/runtime.ts @@ -139,14 +139,28 @@ async function drainAndBuildInjections(opts: SteerDrainHookOptions): Promise item.steerId), - jobCreatedAt, - ); + /** + * Everything armed at snapshot time PLUS everything just drained. The + * boundary has spent its seal, so a snapshotted id that did NOT come back + * from the drain is stale — its steer left the queue by another route + * (typically a cancel whose cross-replica clear was lost). Clearing only + * the drained ids would leave that one level-triggered, sealing the very + * continuation meant to answer the steer we just injected and landing on + * an empty boundary as `preempt_incomplete`. + * + * Ids armed AFTER the snapshot are deliberately excluded: their queue + * items are live and uninjected, and disarming them would strand an + * interrupt the client was already told about. + * + * Not awaited: this runs on the OWNER, where the local disarm is + * synchronous and already effective — the publish only informs other + * replicas, and blocking a boundary drain on it would delay injection. + */ + const spent = new Set(armedBeforeDrain); + for (const item of steers) { + spent.add(item.steerId); + } + void GenerationJobManager.noteSteersRemoved(streamId, [...spent], jobCreatedAt); } return injectedMessages; }