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 <noreply@paperclip.ing>
This commit is contained in:
DottaandPaperclip committed 2026-07-29 19:58:31 +00:00
1 parent 8c353c0c44
commit 026fc3c4e0
4 files changed
+177 -64

No files matched your search

@@ -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);
+1 -1
View File
@@ -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);
+28
View File
@@ -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,
+64 -61
View File
@@ -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<string, unknown>;
},
) => Promise<IssueHeartbeatRunCancellationResult | null>;
finalizeCancelledHeartbeatRun?: (runId: string) => Promise<void>;
};
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<number | null>`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) => {