From a0676fb1d1e14301e6b07a8eb3e7c76bd34d427d Mon Sep 17 00:00:00 2001 From: Dotta Date: Thu, 24 Sep 2026 17:07:47 -0700 Subject: [PATCH] fix(heartbeat): claim issue ownership with queued runs Co-Authored-By: Paperclip --- ...heartbeat-stale-queue-invalidation.test.ts | 32 +++++++++++++++++ server/src/services/heartbeat.ts | 35 +++++++++++++++++-- 2 files changed, 64 insertions(+), 3 deletions(-) diff --git a/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts b/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts index fc71329a30a..4d49def6677 100644 --- a/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts +++ b/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts @@ -311,6 +311,38 @@ describeEmbeddedPostgres("heartbeat stale queued-run invalidation", () => { }); } + it.each([true, false])("serializes assignment claims by issue without serializing unrelated work (same issue: %s)", async (sameIssue) => { + const { companyId, agentId } = await seedCompanyAndAgent({ maxConcurrentRuns: 2 }); + const firstIssueId = randomUUID(); + const secondIssueId = sameIssue ? firstIssueId : randomUUID(); + for (const id of new Set([firstIssueId, secondIssueId])) { + await db.insert(issues).values({ id, companyId, title: "Concurrent assignment", status: "todo", assigneeAgentId: agentId }); + } + await seedQueuedRun({ companyId, agentId, issueId: firstIssueId, wakeReason: "issue_assigned", contextExtras: { source: "issue.create" } }); + await seedQueuedRun({ companyId, agentId, issueId: secondIssueId, wakeReason: "issue_assigned", contextExtras: { source: "issue.assigned_todo_liveness_dispatch" } }); + let release!: () => void; + const gate = new Promise((resolve) => { release = resolve; }); + mockAdapterExecute.mockImplementation(async () => { + await gate; + return { exitCode: 0, signal: null, timedOut: false, errorMessage: null, summary: "Assignment finished.", provider: "test", model: "test-model" }; + }); + try { + await heartbeat.resumeQueuedRuns(); + expect(await waitForCondition(() => Promise.resolve(mockAdapterExecute.mock.calls.length > 0))).toBe(true); + const runs = await db.select().from(heartbeatRuns); + const running = runs.filter((run) => run.status === "running"); + expect(running).toHaveLength(sameIssue ? 1 : 2); + expect(runs.filter((run) => run.status === "queued")).toHaveLength(sameIssue ? 1 : 0); + for (const run of running) { + const [issue] = await db.select().from(issues).where(eq(issues.id, (run.contextSnapshot as { issueId: string }).issueId)); + expect(issue.executionRunId).toBe(run.id); + } + } finally { + release(); + await heartbeat.drainActiveRunExecutions(); + } + }); + it("skips generic timer wakes with no actionable assigned work before adapter execution", async () => { const { agentId } = await seedCompanyAndAgent({ heartbeatConfig: { diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index dca86e4abca..60181278c4c 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -17639,9 +17639,38 @@ export function heartbeatService( agentNameKey: normalizeAgentNameKey(agent.name), }); } - return tx.update(heartbeatRuns).set(claimValues).where(and( - eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.status, "queued"), - )).returning().then((rows) => rows[0] ?? null); + return tx.transaction(async (claimTx) => { + // Several queued assignments can enter the same admission batch + // before executeRun registers an in-memory owner. Claim the issue + // and run together, so only its exact owner reaches provider setup. + const issueOwner = issueId && context.wakeReason !== "source_scoped_recovery_action" + ? await claimTx.select({ + assigneeAgentId: issues.assigneeAgentId, + executionRunId: issues.executionRunId, + checkoutRunId: issues.checkoutRunId, + }).from(issues).where(and( + eq(issues.id, issueId), eq(issues.companyId, run.companyId), + )).for("update").then((rows) => rows[0] ?? null) + : null; + if (issueOwner?.assigneeAgentId === run.agentId && + ((issueOwner.executionRunId && issueOwner.executionRunId !== run.id) || + (run.scheduledRetryReason === "native_safe_replacement" && + issueOwner.checkoutRunId && issueOwner.checkoutRunId !== run.id))) { + return null; + } + const claimedRun = await claimTx.update(heartbeatRuns).set(claimValues).where(and( + eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.status, "queued"), + )).returning().then((rows) => rows[0] ?? null); + if (claimedRun && issueId && issueOwner?.assigneeAgentId === run.agentId) { + await claimTx.update(issues).set({ + executionRunId: claimedRun.id, + executionAgentNameKey: normalizeAgentNameKey(agent.name), + executionLockedAt: claimedAt, + updatedAt: claimedAt, + }).where(and(eq(issues.id, issueId), eq(issues.companyId, run.companyId))); + } + return claimedRun; + }); }); if (!claimed) return null;