diff --git a/packages/db/src/events.test.ts b/packages/db/src/events.test.ts index 1ab2d24..e8ca709 100644 --- a/packages/db/src/events.test.ts +++ b/packages/db/src/events.test.ts @@ -1766,6 +1766,42 @@ describe("claimSteering", () => { where: { id: { in: ["steer-1", "steer-2"] }, claimedAt: null }, data: { runId: "run-1", claimedAt: expect.any(Date) }, }); + expect(tx.steeringMessage.findMany).toHaveBeenCalledWith( + expect.objectContaining({ + where: expect.objectContaining({ claimedAt: null }), + }), + ); + }); + + it("returns nothing when steering was already claimed by another caller", async () => { + const tx = { + $queryRaw: vi.fn(), + run: { findFirst: vi.fn().mockResolvedValue({ id: "run-1" }) }, + steeringMessage: { + findMany: vi.fn().mockResolvedValue([ + { + id: "steer-1", + messageId: "message-1", + message: { seq: 5, blocks: [{ kind: "text", text: "Already claimed" }] }, + }, + ]), + updateMany: vi.fn().mockResolvedValue({ count: 0 }), + }, + }; + const prisma = { + $transaction: vi.fn(async (callback: (client: typeof tx) => unknown) => callback(tx)), + } as unknown as PrismaClient; + + await expect( + claimSteering(prisma, { + threadId: "thread-1", + botId: "bot-1", + runId: "run-1", + leaseOwner: "worker-1", + leaseFence: 2, + seenIds: [], + }), + ).resolves.toEqual([]); }); }); diff --git a/packages/db/src/events.ts b/packages/db/src/events.ts index 90a2daf..cdc4a91 100644 --- a/packages/db/src/events.ts +++ b/packages/db/src/events.ts @@ -482,6 +482,7 @@ export async function claimSteering( const steering = await tx.steeringMessage.findMany({ where: { botId: input.botId, + claimedAt: null, id: input.seenIds.length ? { notIn: input.seenIds } : undefined, OR: [{ runId: null }, { runId: input.runId }], message: { @@ -496,10 +497,11 @@ export async function claimSteering( orderBy: [{ message: { seq: "asc" } }, { id: "asc" }], }); if (steering.length === 0) return []; - await tx.steeringMessage.updateMany({ + const claimed = await tx.steeringMessage.updateMany({ where: { id: { in: steering.map((item) => item.id) }, claimedAt: null }, data: { runId: input.runId, claimedAt: new Date() }, }); + if (claimed.count === 0) return []; return steering.map((item) => ({ id: item.id, messageId: item.messageId, diff --git a/packages/testkit/src/executor-lifecycle.test.ts b/packages/testkit/src/executor-lifecycle.test.ts index 35a8292..21a536c 100644 --- a/packages/testkit/src/executor-lifecycle.test.ts +++ b/packages/testkit/src/executor-lifecycle.test.ts @@ -334,6 +334,39 @@ describeIntegration("run executor lifecycle", () => { expect(userMessages[0]!.seq).toBeLessThan(userMessages[1]!.seq); }); + it("does not re-deliver steering already claimed by the active run", async () => { + const seeded = await seedRun("steering-duplicate", "start the analysis", { + status: "running", + leaseOwner: "steering-worker", + leaseFence: 5, + leaseExpiresAt: new Date(Date.now() + 60_000), + startedAt: new Date(), + }); + await handles.prisma.attempt.create({ + data: { runId: seeded.run.id, fence: 5, status: "running" }, + }); + const events = createThreadEvents(handles.prisma); + await events.sendUserMessage({ + spaceId: seeded.me.spaceId, + threadId: seeded.thread.id, + botId: seeded.bot.id, + userId: seeded.me.userId, + blocks: [{ kind: "text", text: "Use the revised data." }], + prompt: "Use the revised data.", + trigger: "follow_up", + }); + const fence = { + threadId: seeded.thread.id, + botId: seeded.bot.id, + runId: seeded.run.id, + leaseOwner: "steering-worker", + leaseFence: 5, + }; + + await expect(events.claimSteering({ ...fence, seenIds: [] })).resolves.toHaveLength(1); + await expect(events.claimSteering({ ...fence, seenIds: [] })).resolves.toEqual([]); + }); + it.each([ { claim: false, fresh: false, outcome: "failed" as const }, { claim: true, fresh: false, outcome: "failed" as const },