mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-10-26 10:15:21 +01:00
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:
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);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
Reference in new issue
Block a user