diff --git a/server/src/__tests__/heartbeat-lock-release-on-reassignment.test.ts b/server/src/__tests__/heartbeat-lock-release-on-reassignment.test.ts new file mode 100644 index 00000000..167c4c18 --- /dev/null +++ b/server/src/__tests__/heartbeat-lock-release-on-reassignment.test.ts @@ -0,0 +1,271 @@ +import { randomUUID } from "node:crypto"; +import { and, eq } from "drizzle-orm"; +import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; +import { + agents, + agentWakeupRequests, + companies, + createDb, + heartbeatRunEvents, + heartbeatRuns, + issues, +} from "@paperclipai/db"; +import { + getEmbeddedPostgresTestSupport, + startEmbeddedPostgresTestDatabase, +} from "./helpers/embedded-postgres.js"; +import { heartbeatService } from "../services/heartbeat.ts"; + +const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport(); +const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip; + +if (!embeddedPostgresSupport.supported) { + console.warn( + `Skipping embedded Postgres heartbeat lock-release-on-reassignment tests on this host: ${embeddedPostgresSupport.reason ?? "unsupported environment"}`, + ); +} + +describeEmbeddedPostgres("heartbeat lock release on cross-agent reassignment", () => { + let db!: ReturnType; + let tempDb: Awaited> | null = null; + + beforeAll(async () => { + tempDb = await startEmbeddedPostgresTestDatabase("heartbeat-lock-release-on-reassignment-"); + db = createDb(tempDb.connectionString); + }, 60_000); + + afterEach(async () => { + await db.delete(heartbeatRunEvents); + await db.delete(heartbeatRuns); + await db.delete(agentWakeupRequests); + await db.delete(issues); + await db.delete(agents); + await db.delete(companies); + }); + + afterAll(async () => { + await tempDb?.cleanup(); + }); + + async function seedCrossAgentScenario(opts: { holderStatus: "queued" | "running" }) { + const companyId = randomUUID(); + const coderAgentId = randomUUID(); + const reviewerAgentId = randomUUID(); + const issueId = randomUUID(); + const holderRunId = randomUUID(); + const wakeupRequestId = randomUUID(); + const issuePrefix = `T${companyId.replace(/-/g, "").slice(0, 6).toUpperCase()}`; + + await db.insert(companies).values({ + id: companyId, + name: "Paperclip", + issuePrefix, + requireBoardApprovalForNewAgents: false, + }); + + await db.insert(agents).values([ + { + id: coderAgentId, + companyId, + name: "Coder", + role: "engineer", + status: "idle", + adapterType: "process", + adapterConfig: {}, + runtimeConfig: {}, + permissions: {}, + }, + { + id: reviewerAgentId, + companyId, + name: "Reviewer", + role: "engineer", + status: "idle", + adapterType: "process", + adapterConfig: {}, + runtimeConfig: {}, + permissions: {}, + }, + ]); + + await db.insert(agentWakeupRequests).values({ + id: wakeupRequestId, + companyId, + agentId: coderAgentId, + source: "assignment", + status: "queued", + }); + + await db.insert(heartbeatRuns).values({ + id: holderRunId, + companyId, + agentId: coderAgentId, + invocationSource: "assignment", + triggerDetail: "system", + status: opts.holderStatus, + wakeupRequestId, + contextSnapshot: { issueId, taskId: issueId, wakeReason: "issue_assigned" }, + }); + + await db.insert(issues).values({ + id: issueId, + companyId, + title: "Cross-agent reassignment race", + status: "in_review", + priority: "medium", + assigneeAgentId: reviewerAgentId, + executionRunId: holderRunId, + executionAgentNameKey: "coder", + executionLockedAt: new Date(), + issueNumber: 1, + identifier: `${issuePrefix}-1`, + }); + + return { + companyId, + coderAgentId, + reviewerAgentId, + issueId, + holderRunId, + wakeupRequestId, + }; + } + + it("defers a cross-agent wake while the holder is still running and leaves the holder alone", async () => { + const { coderAgentId, reviewerAgentId, issueId, holderRunId, wakeupRequestId } = + await seedCrossAgentScenario({ holderStatus: "running" }); + + const heartbeat = heartbeatService(db); + const followupRun = await heartbeat.wakeup(reviewerAgentId, { + source: "automation", + triggerDetail: "system", + reason: "issue_assigned", + payload: { issueId }, + contextSnapshot: { issueId, taskId: issueId, wakeReason: "issue_assigned" }, + requestedByActorType: "user", + requestedByActorId: "local-board", + }); + + expect(followupRun).toBeNull(); + + const holder = await db + .select({ + status: heartbeatRuns.status, + errorCode: heartbeatRuns.errorCode, + agentId: heartbeatRuns.agentId, + finishedAt: heartbeatRuns.finishedAt, + }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, holderRunId)) + .then((rows) => rows[0] ?? null); + + expect(holder?.status).toBe("running"); + expect(holder?.errorCode).toBeNull(); + expect(holder?.finishedAt).toBeNull(); + expect(holder?.agentId).toBe(coderAgentId); + + const heldWakeup = await db + .select({ status: agentWakeupRequests.status }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.id, wakeupRequestId)) + .then((rows) => rows[0] ?? null); + + expect(heldWakeup?.status).toBe("queued"); + + const deferred = await db + .select({ status: agentWakeupRequests.status, agentId: agentWakeupRequests.agentId }) + .from(agentWakeupRequests) + .where( + and( + eq(agentWakeupRequests.agentId, reviewerAgentId), + eq(agentWakeupRequests.status, "deferred_issue_execution"), + ), + ) + .then((rows) => rows[0] ?? null); + + expect(deferred).not.toBeNull(); + + const issue = await db + .select({ executionRunId: issues.executionRunId }) + .from(issues) + .where(eq(issues.id, issueId)) + .then((rows) => rows[0] ?? null); + + expect(issue?.executionRunId).toBe(holderRunId); + }); + + // Race-guard regression: the cancel UPDATE for the queued holder is pinned + // to the exact non-running status that was read just above it. If a worker + // races in and flips the holder from `queued` → `running` between that + // SELECT and the cancel UPDATE, the status predicate in the WHERE clause + // must guarantee zero rows are clobbered. We simulate the race by + // pre-running the same UPDATE shape against a row that is already + // `running` (the snapshot we would have read was `queued`); the row must + // remain untouched, no wake-request cascade fires, and the lock stays + // owned by the freshly-claimed running holder. + it("guards the cancel UPDATE WHERE clause against a concurrent claim flip to running", async () => { + const { coderAgentId, issueId, holderRunId, wakeupRequestId } = await seedCrossAgentScenario({ + holderStatus: "queued", + }); + + const snapshotStatus = "queued" as const; + + // Concurrent worker claims the queued run after the SELECT but before + // the cancel UPDATE. In production this is a separate transaction + // flipping the row from queued → running. + await db + .update(heartbeatRuns) + .set({ status: "running", startedAt: new Date(), updatedAt: new Date() }) + .where(eq(heartbeatRuns.id, holderRunId)); + + const cancelled = await db + .update(heartbeatRuns) + .set({ + status: "cancelled", + finishedAt: new Date(), + error: "Execution lock released after issue reassigned to a different agent", + errorCode: "lock_released_on_reassignment", + updatedAt: new Date(), + }) + .where( + and( + eq(heartbeatRuns.id, holderRunId), + eq(heartbeatRuns.status, snapshotStatus), + ), + ) + .returning({ id: heartbeatRuns.id }); + + expect(cancelled).toHaveLength(0); + + const holder = await db + .select({ + status: heartbeatRuns.status, + errorCode: heartbeatRuns.errorCode, + agentId: heartbeatRuns.agentId, + finishedAt: heartbeatRuns.finishedAt, + }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, holderRunId)) + .then((rows) => rows[0] ?? null); + + expect(holder?.status).toBe("running"); + expect(holder?.errorCode).toBeNull(); + expect(holder?.finishedAt).toBeNull(); + expect(holder?.agentId).toBe(coderAgentId); + + const heldWakeup = await db + .select({ status: agentWakeupRequests.status }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.id, wakeupRequestId)) + .then((rows) => rows[0] ?? null); + expect(heldWakeup?.status).toBe("queued"); + + const issue = await db + .select({ executionRunId: issues.executionRunId }) + .from(issues) + .where(eq(issues.id, issueId)) + .then((rows) => rows[0] ?? null); + + expect(issue?.executionRunId).toBe(holderRunId); + }); +}); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 41edce2e..56180470 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -10419,6 +10419,59 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) activeExecutionRun = null; } + // A queued/scheduled run holding the lock for an agent that is + // no longer the issue's assignee is stale by design — the issue + // has been re-routed (e.g. blocked → in_review with a different + // assignee). Cancel it and release the lock; otherwise the new + // assignee's wake gets parked in `deferred_issue_execution` + // forever, because the original queued holder will never run + // (the issue's status / target now belongs to someone else). + // + // Race guard: pin the cancel UPDATE to the exact non-running + // status we read above. A worker could transition the holder + // from `queued` → `running` between the SELECT and this UPDATE; + // the status predicate ensures we never clobber a freshly- + // claimed running run. If zero rows matched, leave + // `activeExecutionRun` populated so the defer path runs + // normally against the now-running holder. + if ( + activeExecutionRun && + activeExecutionRun.status !== "running" && + issue.assigneeAgentId && + activeExecutionRun.agentId !== issue.assigneeAgentId + ) { + const cancelled = await tx + .update(heartbeatRuns) + .set({ + status: "cancelled", + finishedAt: new Date(), + error: "Execution lock released after issue reassigned to a different agent", + errorCode: "lock_released_on_reassignment", + updatedAt: new Date(), + }) + .where( + and( + eq(heartbeatRuns.id, activeExecutionRun.id), + eq(heartbeatRuns.status, activeExecutionRun.status), + ), + ) + .returning({ id: heartbeatRuns.id }); + if (cancelled.length > 0) { + if (activeExecutionRun.wakeupRequestId) { + await tx + .update(agentWakeupRequests) + .set({ + status: "cancelled", + finishedAt: new Date(), + error: "Execution lock released after issue reassigned to a different agent", + updatedAt: new Date(), + }) + .where(eq(agentWakeupRequests.id, activeExecutionRun.wakeupRequestId)); + } + activeExecutionRun = null; + } + } + if (!activeExecutionRun && issue.executionRunId) { await tx .update(issues)