diff --git a/server/src/__tests__/execution-lock-orphan-cleanup.test.ts b/server/src/__tests__/execution-lock-orphan-cleanup.test.ts new file mode 100644 index 00000000..b7428624 --- /dev/null +++ b/server/src/__tests__/execution-lock-orphan-cleanup.test.ts @@ -0,0 +1,461 @@ +import { randomUUID } from "node:crypto"; +import { eq } from "drizzle-orm"; +import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; +import { + activityLog, + 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 execution-lock orphan-cleanup tests on this host: ${embeddedPostgresSupport.reason ?? "unsupported environment"}`, + ); +} + +describeEmbeddedPostgres("execution lock orphan cleanup", () => { + let db!: ReturnType; + let tempDb: Awaited> | null = null; + + beforeAll(async () => { + tempDb = await startEmbeddedPostgresTestDatabase("paperclip-exec-lock-orphan-"); + db = createDb(tempDb.connectionString); + }, 30_000); + + afterEach(async () => { + await db.delete(agentWakeupRequests); + await db.delete(activityLog); + await db.delete(heartbeatRunEvents); + await db.delete(issues); + await db.delete(heartbeatRuns); + await db.delete(agents); + await db.delete(companies); + }); + + afterAll(async () => { + await tempDb?.cleanup(); + }); + + async function seedCompany() { + const companyId = randomUUID(); + await db.insert(companies).values({ + id: companyId, + name: "Paperclip", + issuePrefix: `T${companyId.replace(/-/g, "").slice(0, 6).toUpperCase()}`, + requireBoardApprovalForNewAgents: false, + }); + return companyId; + } + + async function seedAgent(companyId: string, name = "CEO") { + const agentId = randomUUID(); + await db.insert(agents).values({ + id: agentId, + companyId, + name, + role: "engineer", + status: "running", + adapterType: "codex_local", + adapterConfig: {}, + runtimeConfig: {}, + permissions: {}, + }); + return agentId; + } + + async function seedIssue(companyId: string, overrides: Partial = {}) { + const issueId = randomUUID(); + await db.insert(issues).values({ + id: issueId, + companyId, + title: "Fixture issue", + status: "in_progress", + priority: "medium", + ...overrides, + }); + return issueId; + } + + describe("heartbeat run finalization", () => { + it("clears execution_run_id on every issue that references the finalized run, not just the run's contextSnapshot issue", async () => { + // Regression test for the "stale execution lock" bug: + // + // A single heartbeat run can end up stamped onto multiple issues' execution_run_id: + // - the issue it explicitly checked out (run.contextSnapshot.issueId) + // - any *other* issue whose enqueueWakeup hit the "legacy run" fallback + // and reattached this same run to that issue's execution_run_id. + // + // The original releaseIssueExecutionAndPromote implementation only resolved the + // single context issue (or rows[0] when no context issue existed), so all the + // other issues were left with execution_run_id pointing at a finalized run — + // an orphan lock that blocked any future agent from checking them out. + const companyId = await seedCompany(); + const ceoAgentId = await seedAgent(companyId, "CEO"); + + const contextIssueId = await seedIssue(companyId); + const orphanIssueId = await seedIssue(companyId); + + const runId = randomUUID(); + await db.insert(heartbeatRuns).values({ + id: runId, + companyId, + agentId: ceoAgentId, + invocationSource: "assignment", + status: "queued", + contextSnapshot: { issueId: contextIssueId }, + }); + + // Both issues reference the same run — exactly the state produced when the + // CEO explicitly checks out one issue and enqueueWakeup's legacy-run fallback + // stamps the same run.id onto a sibling issue. + await db + .update(issues) + .set({ + executionRunId: runId, + executionAgentNameKey: "ceo", + executionLockedAt: new Date(), + }) + .where(eq(issues.id, contextIssueId)); + await db + .update(issues) + .set({ + executionRunId: runId, + executionAgentNameKey: "ceo", + executionLockedAt: new Date(), + }) + .where(eq(issues.id, orphanIssueId)); + + await heartbeatService(db).cancelRun(runId); + + const [contextAfter] = await db.select().from(issues).where(eq(issues.id, contextIssueId)); + const [orphanAfter] = await db.select().from(issues).where(eq(issues.id, orphanIssueId)); + + expect(contextAfter?.executionRunId).toBeNull(); + expect(contextAfter?.executionAgentNameKey).toBeNull(); + expect(contextAfter?.executionLockedAt).toBeNull(); + + expect(orphanAfter?.executionRunId).toBeNull(); + expect(orphanAfter?.executionAgentNameKey).toBeNull(); + expect(orphanAfter?.executionLockedAt).toBeNull(); + }); + + it("clears the execution lock on every orphan when three or more issues share the finalized run", async () => { + // The production symptom of this bug surfaced with four issues pointing at + // the same run id; two is the minimum to reproduce but higher fan-out is + // what we actually observed in the field. The bulk UPDATE path should be + // O(n) without per-row round-trips, so exercise n>2 explicitly. + const companyId = await seedCompany(); + const ceoAgentId = await seedAgent(companyId, "CEO"); + + const contextIssueId = await seedIssue(companyId); + const orphanAId = await seedIssue(companyId); + const orphanBId = await seedIssue(companyId); + const orphanCId = await seedIssue(companyId); + + const runId = randomUUID(); + await db.insert(heartbeatRuns).values({ + id: runId, + companyId, + agentId: ceoAgentId, + invocationSource: "assignment", + status: "queued", + contextSnapshot: { issueId: contextIssueId }, + }); + + for (const issueId of [contextIssueId, orphanAId, orphanBId, orphanCId]) { + await db + .update(issues) + .set({ + executionRunId: runId, + executionAgentNameKey: "ceo", + executionLockedAt: new Date(), + }) + .where(eq(issues.id, issueId)); + } + + await heartbeatService(db).cancelRun(runId); + + const rows = await db.select().from(issues).where(eq(issues.companyId, companyId)); + expect(rows).toHaveLength(4); + for (const row of rows) { + expect(row.executionRunId).toBeNull(); + expect(row.executionAgentNameKey).toBeNull(); + expect(row.executionLockedAt).toBeNull(); + } + }); + + it("clears orphan locks even when the finalizing run has no contextSnapshot issueId", async () => { + // Not every heartbeat run is tied to a specific issue via contextSnapshot + // (e.g. assignment-less wakeups). releaseIssueExecutionAndPromote must + // still clear every issue whose execution_run_id matches the run in that + // case — the legacy behavior picked rows[0] as the "primary" issue for + // deferred-wake promotion, which we intentionally preserve, but cleanup + // must span all matches. + const companyId = await seedCompany(); + const ceoAgentId = await seedAgent(companyId, "CEO"); + + const orphanAId = await seedIssue(companyId); + const orphanBId = await seedIssue(companyId); + + const runId = randomUUID(); + await db.insert(heartbeatRuns).values({ + id: runId, + companyId, + agentId: ceoAgentId, + invocationSource: "assignment", + status: "queued", + contextSnapshot: {}, + }); + + for (const issueId of [orphanAId, orphanBId]) { + await db + .update(issues) + .set({ + executionRunId: runId, + executionAgentNameKey: "ceo", + executionLockedAt: new Date(), + }) + .where(eq(issues.id, issueId)); + } + + await heartbeatService(db).cancelRun(runId); + + const [orphanAAfter] = await db.select().from(issues).where(eq(issues.id, orphanAId)); + const [orphanBAfter] = await db.select().from(issues).where(eq(issues.id, orphanBId)); + + expect(orphanAAfter?.executionRunId).toBeNull(); + expect(orphanBAfter?.executionRunId).toBeNull(); + }); + + it("never clears execution locks on issues in another company", async () => { + // Defense-in-depth regression: a finalizing run in company A must never + // affect rows in company B even if (pathologically) an issue in B carries + // the same execution_run_id. The `company_id = run.companyId` scope in + // the cleanup query is the only line protecting tenant isolation here. + const companyAId = await seedCompany(); + const companyBId = await seedCompany(); + const agentAId = await seedAgent(companyAId, "CEO-A"); + + const issueAId = await seedIssue(companyAId); + const issueBId = await seedIssue(companyBId); + + const runAId = randomUUID(); + await db.insert(heartbeatRuns).values({ + id: runAId, + companyId: companyAId, + agentId: agentAId, + invocationSource: "assignment", + status: "queued", + contextSnapshot: { issueId: issueAId }, + }); + + // Company B's issue pathologically references company A's run id. The FK + // from issues.execution_run_id → heartbeat_runs.id permits this because + // it doesn't enforce a company-match constraint; in production this + // state is only reachable via a bug, but the scoping predicate should + // make us robust against it. + await db + .update(issues) + .set({ + executionRunId: runAId, + executionAgentNameKey: "ceo-a", + executionLockedAt: new Date(), + }) + .where(eq(issues.id, issueAId)); + await db + .update(issues) + .set({ + executionRunId: runAId, + executionAgentNameKey: "ceo-b", + executionLockedAt: new Date(), + }) + .where(eq(issues.id, issueBId)); + + await heartbeatService(db).cancelRun(runAId); + + const [issueAAfter] = await db.select().from(issues).where(eq(issues.id, issueAId)); + const [issueBAfter] = await db.select().from(issues).where(eq(issues.id, issueBId)); + + expect(issueAAfter?.executionRunId).toBeNull(); + expect(issueAAfter?.companyId).toBe(companyAId); + + expect(issueBAfter?.executionRunId).toBe(runAId); + expect(issueBAfter?.executionAgentNameKey).toBe("ceo-b"); + expect(issueBAfter?.executionLockedAt).not.toBeNull(); + expect(issueBAfter?.companyId).toBe(companyBId); + }); + + it("clears checkout_run_id on every sibling and preserves an in-flight retry's execution_run_id pointer", async () => { + // Exercises the second of the two scoped bulk UPDATEs in + // releaseIssueExecutionAndPromote, and the invariant that motivated + // splitting it from the first one: + // + // - One sibling has BOTH execution_run_id and checkout_run_id pinned at + // the finalizing run (normal "checked out and executing" shape) — both + // columns must be cleared. + // + // - Another sibling is in the codex-transient / process-loss retry shape: + // execution_run_id has already been moved to a *different* retry run, + // while checkout_run_id is left pinned at the original failed run. + // The checkout column must be cleared but the retry's execution_run_id + // pointer must NOT be clobbered — otherwise the retry can never check + // itself out and the bug re-surfaces under a different name. + // + // - A third sibling has checkout_run_id pinned at the finalizing run with + // execution_run_id null — covers the post-status-change shape where the + // issue's checkout pointer outlived its execution lock. + const companyId = await seedCompany(); + const ceoAgentId = await seedAgent(companyId, "CEO"); + + const dualLockIssueId = await seedIssue(companyId); + const retryIssueId = await seedIssue(companyId); + const checkoutOnlyIssueId = await seedIssue(companyId); + + const finalizingRunId = randomUUID(); + const retryRunId = randomUUID(); + await db.insert(heartbeatRuns).values([ + { + id: finalizingRunId, + companyId, + agentId: ceoAgentId, + invocationSource: "assignment", + status: "queued", + contextSnapshot: { issueId: dualLockIssueId }, + }, + { + // Marked `running` rather than `queued` so cancelRun(finalizingRunId) + // does not also reconcile this queued retry into a cancelled state + // (which would then release its own execution_run_id and defeat the + // point of this test). In production, retries are typically already + // running by the time the original failing run finalizes. + id: retryRunId, + companyId, + agentId: ceoAgentId, + invocationSource: "assignment", + status: "running", + contextSnapshot: { issueId: retryIssueId }, + }, + ]); + + await db + .update(issues) + .set({ + executionRunId: finalizingRunId, + executionAgentNameKey: "ceo", + executionLockedAt: new Date(), + checkoutRunId: finalizingRunId, + }) + .where(eq(issues.id, dualLockIssueId)); + await db + .update(issues) + .set({ + executionRunId: retryRunId, + executionAgentNameKey: "ceo", + executionLockedAt: new Date(), + checkoutRunId: finalizingRunId, + }) + .where(eq(issues.id, retryIssueId)); + await db + .update(issues) + .set({ + checkoutRunId: finalizingRunId, + }) + .where(eq(issues.id, checkoutOnlyIssueId)); + + await heartbeatService(db).cancelRun(finalizingRunId); + + const [dualAfter] = await db.select().from(issues).where(eq(issues.id, dualLockIssueId)); + const [retryAfter] = await db.select().from(issues).where(eq(issues.id, retryIssueId)); + const [checkoutOnlyAfter] = await db + .select() + .from(issues) + .where(eq(issues.id, checkoutOnlyIssueId)); + + // Dual-lock sibling: both columns cleared. + expect(dualAfter?.executionRunId).toBeNull(); + expect(dualAfter?.executionAgentNameKey).toBeNull(); + expect(dualAfter?.executionLockedAt).toBeNull(); + expect(dualAfter?.checkoutRunId).toBeNull(); + + // Retry sibling: checkout column cleared, retry's execution pointer preserved. + expect(retryAfter?.checkoutRunId).toBeNull(); + expect(retryAfter?.executionRunId).toBe(retryRunId); + expect(retryAfter?.executionAgentNameKey).toBe("ceo"); + expect(retryAfter?.executionLockedAt).not.toBeNull(); + + // Checkout-only sibling: column cleared, no execution side-effects. + expect(checkoutOnlyAfter?.checkoutRunId).toBeNull(); + expect(checkoutOnlyAfter?.executionRunId).toBeNull(); + }); + + it("does not touch execution locks on issues owned by unrelated runs", async () => { + const companyId = await seedCompany(); + const ceoAgentId = await seedAgent(companyId, "CEO"); + const seAgentId = await seedAgent(companyId, "SE1"); + + const finalizingRunId = randomUUID(); + const unrelatedRunId = randomUUID(); + + const contextIssueId = await seedIssue(companyId); + const unrelatedIssueId = await seedIssue(companyId); + + await db.insert(heartbeatRuns).values([ + { + id: finalizingRunId, + companyId, + agentId: ceoAgentId, + invocationSource: "assignment", + status: "queued", + contextSnapshot: { issueId: contextIssueId }, + }, + { + id: unrelatedRunId, + companyId, + agentId: seAgentId, + invocationSource: "assignment", + status: "running", + contextSnapshot: { issueId: unrelatedIssueId }, + }, + ]); + + await db + .update(issues) + .set({ + executionRunId: finalizingRunId, + executionAgentNameKey: "ceo", + executionLockedAt: new Date(), + }) + .where(eq(issues.id, contextIssueId)); + await db + .update(issues) + .set({ + executionRunId: unrelatedRunId, + executionAgentNameKey: "se1", + executionLockedAt: new Date(), + }) + .where(eq(issues.id, unrelatedIssueId)); + + await heartbeatService(db).cancelRun(finalizingRunId); + + const [contextAfter] = await db.select().from(issues).where(eq(issues.id, contextIssueId)); + const [unrelatedAfter] = await db.select().from(issues).where(eq(issues.id, unrelatedIssueId)); + + expect(contextAfter?.executionRunId).toBeNull(); + expect(unrelatedAfter?.executionRunId).toBe(unrelatedRunId); + expect(unrelatedAfter?.executionAgentNameKey).toBe("se1"); + }); + }); +}); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 24a38df7..39f6794a 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -9524,52 +9524,109 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) const recoveryAgentNameKey = normalizeAgentNameKey(recoveryAgent?.name); const promotionResult = await db.transaction(async (tx) => { - if (contextIssueId) { - await tx.execute( - sql`select id from issues where company_id = ${run.companyId} and id = ${contextIssueId} for update`, - ); - } else { - await tx.execute( - sql`select id from issues where company_id = ${run.companyId} and execution_run_id = ${run.id} for update`, - ); - } + // Lock the context issue (if any) AND every issue that still references this run. + // + // A single run can hold execution locks on multiple issues: the caller's context + // issue (set via svc.checkout) plus any additional issues stamped by + // enqueueWakeup's "legacy run" fallback when the run was the only queued/running + // run matching their contextSnapshot.issueId. Historically this function only + // resolved and cleared the lock on *one* issue (rows[0]), leaving the others + // with an executionRunId pointing at a finalized run. Subsequent checkouts from + // the assigned agent then failed with 409 and the issue stayed blocked forever. + // `order by id` makes row-lock acquisition deterministic across concurrent + // finalizations, which keeps deadlock risk independent of PostgreSQL's plan + // choice when multiple issues match. + await tx.execute( + contextIssueId + ? sql` + select id from issues + where company_id = ${run.companyId} + and ( + id = ${contextIssueId} + or execution_run_id = ${run.id} + or checkout_run_id = ${run.id} + ) + order by id + for update + ` + : sql` + select id from issues + where company_id = ${run.companyId} + and (execution_run_id = ${run.id} or checkout_run_id = ${run.id}) + order by id + for update + `, + ); - let issue = await tx + const candidateIssues = await tx .select() .from(issues) .where( and( eq(issues.companyId, run.companyId), - contextIssueId ? eq(issues.id, contextIssueId) : eq(issues.executionRunId, run.id), + contextIssueId + ? or( + eq(issues.id, contextIssueId), + eq(issues.executionRunId, run.id), + eq(issues.checkoutRunId, run.id), + ) + : or(eq(issues.executionRunId, run.id), eq(issues.checkoutRunId, run.id)), ), ) - .then((rows) => rows[0] ?? null); + .orderBy(asc(issues.id)); + + // Clear orphaned execution-lock columns that still point at this finalizing + // run, across every sibling issue in one statement so it scales with N + // orphans without N round-trips. Rows are already held under FOR UPDATE from + // the lock query above. + // + // The two columns are cleared in separate UPDATEs so we never clobber a + // retry's executionRunId pointer: when a process-loss or codex-transient + // retry is scheduled mid-finalization, it moves `executionRunId` from this + // run to the retry run while leaving `checkoutRunId` pinned at this run. + // Only the checkout column should be released in that case; the execution + // column now belongs to the retry. + const promotionUpdateTimestamp = new Date(); + await tx + .update(issues) + .set({ + executionRunId: null, + executionAgentNameKey: null, + executionLockedAt: null, + updatedAt: promotionUpdateTimestamp, + }) + .where( + and(eq(issues.companyId, run.companyId), eq(issues.executionRunId, run.id)), + ); + // `checkoutRunId` clear is symmetric to #6008's per-issue self-heal, + // extended to all siblings: covers paths where the issue's assignee or + // status changed between checkout and termination, which + // adoptStaleCheckoutRun's narrow WHERE clause cannot reach. + await tx + .update(issues) + .set({ + checkoutRunId: null, + updatedAt: promotionUpdateTimestamp, + }) + .where( + and(eq(issues.companyId, run.companyId), eq(issues.checkoutRunId, run.id)), + ); + + // Deferred-wake promotion is bound to a single primary issue: the run's context + // issue when present, otherwise the first candidate we found (preserves the + // legacy rows[0] selection for runs that were not tied to a specific issue). + let issue = + (contextIssueId + ? candidateIssues.find((candidate) => candidate.id === contextIssueId) + : candidateIssues[0]) ?? null; if (!issue) return null; if (issue.executionRunId && issue.executionRunId !== run.id) return null; - // Clear lock columns that point at the terminating run. checkoutRunId is - // cleared here in addition to executionRunId so the issue self-heals even - // if its assignee or status changed between checkout and termination — - // adoptStaleCheckoutRun's narrow WHERE clause cannot cover those paths. - if (issue.executionRunId === run.id || issue.checkoutRunId === run.id) { - await tx - .update(issues) - .set({ - checkoutRunId: null, - executionRunId: null, - executionAgentNameKey: null, - executionLockedAt: null, - updatedAt: new Date(), - }) - .where( - and( - eq(issues.id, issue.id), - or(eq(issues.executionRunId, run.id), eq(issues.checkoutRunId, run.id)), - ), - ); - } - + // Workspace-validation recovery: if the finalizing run failed workspace + // validation, surface the primary issue for the blocked-recovery comment path. + // Sibling lock cleanup is already done above; only the primary issue carries + // the recovery surface because the comment is attached to a single issue. if ( isWorkspaceValidationFailedRun(run) && (issue.status === "todo" || issue.status === "in_progress") && @@ -9585,6 +9642,7 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) }; } + while (true) { const deferred = await tx .select()