diff --git a/server/src/__tests__/task-watchdogs-scheduler.test.ts b/server/src/__tests__/task-watchdogs-scheduler.test.ts index b71d34988a6..6b5be593354 100644 --- a/server/src/__tests__/task-watchdogs-scheduler.test.ts +++ b/server/src/__tests__/task-watchdogs-scheduler.test.ts @@ -450,11 +450,40 @@ describeEmbeddedPostgres("task watchdog scheduler", () => { }); const firstWakes: string[] = []; const firstService = taskWatchdogService(db, { - enqueueWakeup: async (agentId) => { + enqueueWakeup: async (agentId, opts) => { firstWakes.push(agentId); + const [wakeupRequest] = await db + .insert(agentWakeupRequests) + .values({ + companyId, + agentId, + source: opts?.source ?? "automation", + triggerDetail: opts?.triggerDetail, + reason: opts?.reason, + payload: opts?.payload, + status: "queued", + idempotencyKey: opts?.idempotencyKey, + }) + .returning(); + const [run] = await db + .insert(heartbeatRuns) + .values({ + companyId, + agentId, + invocationSource: opts?.source ?? "automation", + triggerDetail: opts?.triggerDetail, + status: "queued", + wakeupRequestId: wakeupRequest!.id, + contextSnapshot: opts?.contextSnapshot, + }) + .returning(); + await db + .update(agentWakeupRequests) + .set({ runId: run!.id }) + .where(eq(agentWakeupRequests.id, wakeupRequest!.id)); signalWakeStarted(); await wakeCanFinish; - return { id: randomUUID() }; + return { id: run!.id }; }, }); @@ -470,6 +499,19 @@ describeEmbeddedPostgres("task watchdog scheduler", () => { watchdogAgentId: currentAssigneeId, lastReviewedFingerprint: null, }); + const [cancelledRun] = await db + .select() + .from(heartbeatRuns) + .where(eq(heartbeatRuns.agentId, formerAssigneeId)); + expect(cancelledRun).toMatchObject({ + status: "cancelled", + errorCode: "issue_reassigned", + }); + const [cancelledWake] = await db + .select() + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.agentId, formerAssigneeId)); + expect(cancelledWake).toMatchObject({ status: "cancelled" }); const { service, wakes } = createService(); const second = await service.reconcileTaskWatchdogs({ companyId }); diff --git a/server/src/services/issues.ts b/server/src/services/issues.ts index ab5a0763497..89124bfd3d2 100644 --- a/server/src/services/issues.ts +++ b/server/src/services/issues.ts @@ -6761,6 +6761,74 @@ export function issueService(db: Db) { .then((rows: Array) => rows[0] ?? null); if (!updated) return null; if (lockedSelfWatchdog) { + const cancelledAt = new Date(); + const cancellationReason = "Cancelled because the self-review issue was reassigned before the run started"; + const cancelledRuns = 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, + }); + const wakeupRequestIds = cancelledRuns + .map((run: { wakeupRequestId: string | null }) => run.wakeupRequestId) + .filter((requestId: string | null): requestId is string => Boolean(requestId)); + if (wakeupRequestIds.length > 0) { + await tx + .update(agentWakeupRequests) + .set({ + status: "cancelled", + finishedAt: cancelledAt, + error: cancellationReason, + updatedAt: cancelledAt, + }) + .where(inArray(agentWakeupRequests.id, wakeupRequestIds)); + } + await tx + .update(agentWakeupRequests) + .set({ + status: "cancelled", + finishedAt: cancelledAt, + error: cancellationReason, + updatedAt: cancelledAt, + }) + .where(and( + eq(agentWakeupRequests.companyId, existing.companyId), + eq(agentWakeupRequests.status, "deferred_issue_execution"), + sql`${agentWakeupRequests.payload} -> '_paperclipWakeContext' -> 'taskWatchdogSelfReview' ->> 'watchdogId' = ${lockedSelfWatchdog.id}`, + )); + for (const cancelledRun of cancelledRuns) { + await logActivity(tx as unknown as Db, { + companyId: existing.companyId, + actorType: actorAgentId ? "agent" : actorUserId ? "user" : "system", + actorId: actorAgentId ?? actorUserId ?? "issue_service", + agentId: cancelledRun.agentId, + runId: null, + action: "heartbeat.cancelled", + entityType: "heartbeat_run", + entityId: cancelledRun.id, + issueId: existing.id, + details: { + source: "self_watchdog_reassignment", + issueId: existing.id, + watchdogId: lockedSelfWatchdog.id, + previousAssigneeAgentId: existing.assigneeAgentId, + currentAssigneeAgentId: nextAssigneeAgentId, + }, + }); + } await tx .update(issueWatchdogs) .set({