From 026fc3c4e05d09fe18f8a745495dbb7202d54ae6 Mon Sep 17 00:00:00 2001 From: Dotta Date: Wed, 29 Jul 2026 19:58:31 +0000 Subject: [PATCH] fix(watchdogs): make reassignment cancellation atomic Commit self-review run cancellation with the assignee update, then terminate the cancelled process after the transaction succeeds so rollback preserves the former owner's active review. Co-Authored-By: Paperclip --- .../task-watchdogs-scheduler.test.ts | 86 +++++++++++- server/src/routes/issues.ts | 2 +- server/src/services/heartbeat.ts | 28 ++++ server/src/services/issues.ts | 125 +++++++++--------- 4 files changed, 177 insertions(+), 64 deletions(-) diff --git a/server/src/__tests__/task-watchdogs-scheduler.test.ts b/server/src/__tests__/task-watchdogs-scheduler.test.ts index 62c043ac12f..c9212dc0ee5 100644 --- a/server/src/__tests__/task-watchdogs-scheduler.test.ts +++ b/server/src/__tests__/task-watchdogs-scheduler.test.ts @@ -1,5 +1,5 @@ import { randomUUID } from "node:crypto"; -import { and, eq } from "drizzle-orm"; +import { and, eq, sql } from "drizzle-orm"; import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; import { activityLog, @@ -582,7 +582,7 @@ describeEmbeddedPostgres("task watchdog scheduler", () => { await wakeStarted; const heartbeat = heartbeatService(db); const reassignment = issueService(db, { - cancelHeartbeatRun: heartbeat.cancelRun, + finalizeCancelledHeartbeatRun: heartbeat.finalizeCancelledRun, }).update(sourceId, { assigneeAgentId: currentAssigneeId }); releaseWake(); await Promise.all([firstReconcile, reassignment]); @@ -616,6 +616,88 @@ describeEmbeddedPostgres("task watchdog scheduler", () => { expect(wakes[0]?.opts?.idempotencyKey).toContain(currentAssigneeId); }); + it("rolls back self-review cancellation when reassignment fails", async () => { + const companyId = await seedCompany(); + const formerAssigneeId = await seedAgent(companyId, { name: "Rollback Former Assignee" }); + const currentAssigneeId = await seedAgent(companyId, { name: "Rollback Current Assignee" }); + const sourceId = await seedIssue(companyId, { + identifier: "WDOG-SELF-REASSIGN-ROLLBACK", + status: "done", + assigneeAgentId: formerAssigneeId, + }); + const watchdog = await seedWatchdog(companyId, sourceId, formerAssigneeId, { mode: "self" }); + const [wakeupRequest] = await db + .insert(agentWakeupRequests) + .values({ + companyId, + agentId: formerAssigneeId, + source: "automation", + triggerDetail: "system", + reason: "task_watchdog_self_review", + status: "running", + claimedAt: new Date(), + }) + .returning(); + const [run] = await db + .insert(heartbeatRuns) + .values({ + companyId, + agentId: formerAssigneeId, + invocationSource: "automation", + triggerDetail: "system", + status: "running", + startedAt: new Date(), + wakeupRequestId: wakeupRequest!.id, + contextSnapshot: { + issueId: sourceId, + taskWatchdogSelfReview: { watchdogId: watchdog!.id }, + }, + }) + .returning(); + await db + .update(agentWakeupRequests) + .set({ runId: run!.id }) + .where(eq(agentWakeupRequests.id, wakeupRequest!.id)); + + await db.execute(sql` + create or replace function fail_watchdog_reassignment_test() + returns trigger as $$ + begin + raise exception 'forced reassignment failure'; + end; + $$ language plpgsql + `); + await db.execute(sql` + create trigger fail_watchdog_reassignment_test_trigger + before update on issues + for each row execute function fail_watchdog_reassignment_test() + `); + const finalizedRunIds: string[] = []; + try { + await expect(issueService(db, { + finalizeCancelledHeartbeatRun: async (runId) => { + finalizedRunIds.push(runId); + }, + }).update(sourceId, { assigneeAgentId: currentAssigneeId })) + .rejects.toThrow("Failed query: update \"issues\""); + } finally { + await db.execute(sql`drop trigger if exists fail_watchdog_reassignment_test_trigger on issues`); + await db.execute(sql`drop function if exists fail_watchdog_reassignment_test()`); + } + + const [unchangedIssue] = await db.select().from(issues).where(eq(issues.id, sourceId)); + const [unchangedRun] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run!.id)); + const [unchangedWake] = await db + .select() + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.id, wakeupRequest!.id)); + expect(unchangedIssue?.assigneeAgentId).toBe(formerAssigneeId); + expect(unchangedRun?.status).toBe("running"); + expect(unchangedWake?.status).toBe("running"); + expect(finalizedRunIds).toEqual([]); + expect(await db.select().from(heartbeatRunEvents).where(eq(heartbeatRunEvents.runId, run!.id))).toHaveLength(0); + }); + it("retries a self review when the wake was not queued", async () => { const companyId = await seedCompany(); const agentId = await seedAgent(companyId); diff --git a/server/src/routes/issues.ts b/server/src/routes/issues.ts index 3f563d5eff8..2c4917b8594 100644 --- a/server/src/routes/issues.ts +++ b/server/src/routes/issues.ts @@ -2593,7 +2593,7 @@ export function issueRoutes( pluginWorkerManager: opts.pluginWorkerManager, }); const svc = issueService(db, { - cancelHeartbeatRun: heartbeat.cancelRun, + finalizeCancelledHeartbeatRun: heartbeat.finalizeCancelledRun, }); const enqueueRecoveryActionWakeup = opts.recoveryActionEnqueueWakeup ?? heartbeat.wakeup; const feedback = feedbackService(db); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 8e226222f4e..55f7718053a 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -16984,6 +16984,32 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) return cancelled; } + async function finalizeCancelledRunInternal(runId: string) { + const run = await getRun(runId); + if (!run || run.status !== "cancelled") return; + + const running = runningProcesses.get(run.id); + try { + if (running) { + await terminateHeartbeatRunProcess({ + pid: running.child.pid ?? run.processPid, + processGroupId: running.processGroupId ?? run.processGroupId, + graceMs: Math.max(1, running.graceSec) * 1000, + }); + } else if (run.processPid || run.processGroupId) { + await terminateHeartbeatRunProcess({ + pid: run.processPid, + processGroupId: run.processGroupId, + }); + } + } finally { + runningProcesses.delete(run.id); + } + + await finalizeAgentStatus(run.agentId, "cancelled"); + await startNextQueuedRunForAgent(run.agentId); + } + async function cancelActiveForAgentInternal(agentId: string, reason = "Cancelled due to agent pause", errorCode = "cancelled") { const agent = await getAgent(agentId); const runs = await db @@ -17213,6 +17239,8 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) getRun, + finalizeCancelledRun: finalizeCancelledRunInternal, + decorateActiveRunStatus: decorateHeartbeatRunRuntimeStatus, recordRuntimeProgress: recordCurrentHeartbeatRunRuntimeProgress, sweepExpiredRuntimeStatuses: sweepExpiredHeartbeatRunRuntimeStatuses, diff --git a/server/src/services/issues.ts b/server/src/services/issues.ts index 742cdf895de..63792d1a634 100644 --- a/server/src/services/issues.ts +++ b/server/src/services/issues.ts @@ -14,6 +14,7 @@ import { documentRevisions, documents, goals, + heartbeatRunEvents, heartbeatRuns, routineRuns, executionWorkspaces, @@ -119,23 +120,8 @@ export const ISSUE_BLOCKER_DIAGNOSTICS_MAX_BLOCKERS = 100; export const ISSUE_WAKE_DIAGNOSTICS_MAX_WAKE_REQUESTS = 50; export const ISSUE_WAKE_DIAGNOSTICS_MAX_ACTIVITY_RECORDS = 50; -type IssueHeartbeatRunCancellationResult = { - id: string; - agentId: string; - status: string; - wakeupRequestId: string | null; -}; - type IssueServiceOptions = { - cancelHeartbeatRun?: ( - runId: string, - reason: string, - options: { - errorCode: string; - eventMessage: string; - eventPayload: Record; - }, - ) => Promise; + finalizeCancelledHeartbeatRun?: (runId: string) => Promise; }; export const ISSUE_WAKE_DIAGNOSTICS_LOOKBACK_DAYS = 14; export const ISSUE_SUBTREE_DIAGNOSTICS_MAX_DEPTH = 8; @@ -6736,6 +6722,8 @@ export function issueService(db: Db, options: IssueServiceOptions = {}) { patch.executionLockedAt = null; } + const cancelledSelfReviewRunIds: string[] = []; + const runUpdate = async (tx: any) => { // Self-watchdog dispatch takes this row lock across wake creation and // the reviewed-fingerprint claim. Taking it before changing the issue @@ -6776,7 +6764,8 @@ export function issueService(db: Db, options: IssueServiceOptions = {}) { )); const hasRunningReview = cancellableRuns.some((run: { status: string }) => run.status === "running"); - if (hasRunningReview && !options.cancelHeartbeatRun) { + const canFinalizeCancellationAfterCommit = dbOrTx === db && Boolean(options.finalizeCancelledHeartbeatRun); + if (hasRunningReview && !canFinalizeCancellationAfterCommit) { throw conflict("Cannot reassign an issue while its self-review run is active", { issueId: existing.id, watchdogId: lockedSelfWatchdog.id, @@ -6786,49 +6775,32 @@ export function issueService(db: Db, options: IssueServiceOptions = {}) { }); } - if (options.cancelHeartbeatRun) { - for (const run of cancellableRuns) { - const cancelled = await options.cancelHeartbeatRun(run.id, cancellationReason, { - errorCode: "issue_reassigned", - eventMessage: "self-review run cancelled for issue reassignment", - eventPayload: { - source: "self_watchdog_reassignment", - issueId: existing.id, - watchdogId: lockedSelfWatchdog.id, - previousAssigneeAgentId: existing.assigneeAgentId, - currentAssigneeAgentId: nextAssigneeAgentId, - }, - }); - if (cancelled?.status === "cancelled") { - cancelledSelfReviewRuns.push({ - id: cancelled.id, - agentId: cancelled.agentId, - wakeupRequestId: cancelled.wakeupRequestId, - }); - } - } - } else { - const directlyCancelledRuns = await tx - .update(heartbeatRuns) - .set({ - status: "cancelled", - finishedAt: cancelledAt, - error: cancellationReason, - errorCode: "issue_reassigned", - updatedAt: cancelledAt, - }) - .where(and( - eq(heartbeatRuns.companyId, existing.companyId), - inArray(heartbeatRuns.status, ["queued", "scheduled_retry"]), - sql`${heartbeatRuns.contextSnapshot} -> 'taskWatchdogSelfReview' ->> 'watchdogId' = ${lockedSelfWatchdog.id}`, - )) - .returning({ - id: heartbeatRuns.id, - agentId: heartbeatRuns.agentId, - wakeupRequestId: heartbeatRuns.wakeupRequestId, - }); - cancelledSelfReviewRuns.push(...directlyCancelledRuns); - } + const directlyCancelledRuns = await tx + .update(heartbeatRuns) + .set({ + status: "cancelled", + finishedAt: cancelledAt, + error: cancellationReason, + errorCode: "issue_reassigned", + updatedAt: cancelledAt, + }) + .where(and( + eq(heartbeatRuns.companyId, existing.companyId), + inArray( + heartbeatRuns.status, + canFinalizeCancellationAfterCommit + ? ["queued", "running", "scheduled_retry"] + : ["queued", "scheduled_retry"], + ), + sql`${heartbeatRuns.contextSnapshot} -> 'taskWatchdogSelfReview' ->> 'watchdogId' = ${lockedSelfWatchdog.id}`, + )) + .returning({ + id: heartbeatRuns.id, + agentId: heartbeatRuns.agentId, + wakeupRequestId: heartbeatRuns.wakeupRequestId, + }); + cancelledSelfReviewRuns.push(...directlyCancelledRuns); + cancelledSelfReviewRunIds.push(...directlyCancelledRuns.map((run: { id: string }) => run.id)); const remainingCancellableRuns = await tx .select({ id: heartbeatRuns.id, status: heartbeatRuns.status }) @@ -6860,6 +6832,29 @@ export function issueService(db: Db, options: IssueServiceOptions = {}) { }) .where(inArray(agentWakeupRequests.id, wakeupRequestIds)); } + for (const cancelledRun of cancelledSelfReviewRuns) { + const [eventSeq] = await tx + .select({ maxSeq: sql`max(${heartbeatRunEvents.seq})` }) + .from(heartbeatRunEvents) + .where(eq(heartbeatRunEvents.runId, cancelledRun.id)); + await tx.insert(heartbeatRunEvents).values({ + companyId: existing.companyId, + runId: cancelledRun.id, + agentId: cancelledRun.agentId, + seq: Number(eventSeq?.maxSeq ?? 0) + 1, + eventType: "lifecycle", + stream: "system", + level: "warn", + message: "self-review run cancelled for issue reassignment", + payload: { + source: "self_watchdog_reassignment", + issueId: existing.id, + watchdogId: lockedSelfWatchdog.id, + previousAssigneeAgentId: existing.assigneeAgentId, + currentAssigneeAgentId: nextAssigneeAgentId, + }, + }); + } await tx .update(agentWakeupRequests) .set({ @@ -7044,7 +7039,15 @@ export function issueService(db: Db, options: IssueServiceOptions = {}) { return enriched; }; - return dbOrTx === db ? db.transaction(runUpdate) : runUpdate(dbOrTx); + if (dbOrTx !== db) return runUpdate(dbOrTx); + + const updated = await db.transaction(runUpdate); + if (updated && options.finalizeCancelledHeartbeatRun) { + for (const runId of cancelledSelfReviewRunIds) { + await options.finalizeCancelledHeartbeatRun(runId); + } + } + return updated; }, clearExecutionWorkspaceEnvironmentSelection: async (companyId: string, environmentId: string) => {