From ba88e86f63d0642b4802ca63d5700b2e6426a994 Mon Sep 17 00:00:00 2001 From: Dotta Date: Wed, 29 Apr 2026 14:15:24 -0500 Subject: [PATCH] Cancel queued runs for cancelled issues --- .../heartbeat-retry-scheduling.test.ts | 99 +++++++++++++++++++ .../issue-comment-reopen-routes.test.ts | 31 +++--- server/src/routes/issues.ts | 58 ++++++++--- server/src/services/heartbeat.ts | 90 ++++++++++++++++- 4 files changed, 247 insertions(+), 31 deletions(-) diff --git a/server/src/__tests__/heartbeat-retry-scheduling.test.ts b/server/src/__tests__/heartbeat-retry-scheduling.test.ts index edf94d6eb4e..4f61f867f4c 100644 --- a/server/src/__tests__/heartbeat-retry-scheduling.test.ts +++ b/server/src/__tests__/heartbeat-retry-scheduling.test.ts @@ -62,6 +62,105 @@ describeEmbeddedPostgres("heartbeat bounded retry scheduling", () => { await tempDb?.cleanup(); }); + it("cancels queued issue runs and deferred wakeups when an issue is cancelled", async () => { + const companyId = randomUUID(); + const agentId = randomUUID(); + const issueId = randomUUID(); + const queuedWakeupId = randomUUID(); + const deferredWakeupId = randomUUID(); + const queuedRunId = randomUUID(); + + await db.insert(companies).values({ + id: companyId, + name: "Paperclip", + issuePrefix: "TQUEUE", + requireBoardApprovalForNewAgents: false, + }); + await db.insert(agents).values({ + id: agentId, + companyId, + name: "CodexCoder", + role: "engineer", + status: "active", + adapterType: "codex_local", + adapterConfig: {}, + runtimeConfig: { + heartbeat: { + wakeOnDemand: true, + maxConcurrentRuns: 1, + }, + }, + permissions: {}, + }); + await db.insert(issues).values({ + id: issueId, + companyId, + title: "Cancel queued work", + status: "cancelled", + priority: "medium", + assigneeAgentId: agentId, + cancelledAt: new Date(), + }); + await db.insert(agentWakeupRequests).values([ + { + id: queuedWakeupId, + companyId, + agentId, + source: "automation", + triggerDetail: "system", + reason: "issue_commented", + payload: { issueId }, + status: "queued", + }, + { + id: deferredWakeupId, + companyId, + agentId, + source: "automation", + triggerDetail: "system", + reason: "issue_execution_deferred", + payload: { issueId }, + status: "deferred_issue_execution", + }, + ]); + await db.insert(heartbeatRuns).values({ + id: queuedRunId, + companyId, + agentId, + invocationSource: "automation", + triggerDetail: "system", + status: "queued", + wakeupRequestId: queuedWakeupId, + contextSnapshot: { issueId, wakeCommentId: randomUUID() }, + }); + await db + .update(agentWakeupRequests) + .set({ runId: queuedRunId }) + .where(eq(agentWakeupRequests.id, queuedWakeupId)); + + const result = await heartbeat.cancelIssueRuns({ + companyId, + issueId, + reason: "Cancelled because the issue was marked cancelled", + }); + + expect(result.runs.map((run) => run.id)).toEqual([queuedRunId]); + expect(result.wakeups.map((wakeup) => wakeup.id)).toEqual([deferredWakeupId]); + + const [queuedRun] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, queuedRunId)); + const wakeups = await db + .select() + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.companyId, companyId)); + + expect(queuedRun?.status).toBe("cancelled"); + expect(queuedRun?.errorCode).toBe("cancelled"); + expect(Object.fromEntries(wakeups.map((wakeup) => [wakeup.id, wakeup.status]))).toEqual({ + [queuedWakeupId]: "cancelled", + [deferredWakeupId]: "cancelled", + }); + }); + async function seedRetryFixture(input: { runId: string; companyId: string; diff --git a/server/src/__tests__/issue-comment-reopen-routes.test.ts b/server/src/__tests__/issue-comment-reopen-routes.test.ts index 3cca0371fb6..735afd3fd81 100644 --- a/server/src/__tests__/issue-comment-reopen-routes.test.ts +++ b/server/src/__tests__/issue-comment-reopen-routes.test.ts @@ -24,6 +24,7 @@ const mockHeartbeatService = vi.hoisted(() => ({ getRun: vi.fn(async () => null), getActiveRunForAgent: vi.fn(async () => null), cancelRun: vi.fn(async () => null), + cancelIssueRuns: vi.fn(async () => ({ runs: [], wakeups: [] })), })); const mockAgentService = vi.hoisted(() => ({ @@ -110,6 +111,7 @@ vi.mock("../services/routines.js", () => ({ vi.mock("../services/index.js", () => ({ accessService: () => mockAccessService, agentService: () => mockAgentService, + companyService: () => ({}), documentService: () => ({}), executionWorkspaceService: () => ({}), feedbackService: () => mockFeedbackService, @@ -221,6 +223,7 @@ describe.sequential("issue comment reopen routes", () => { mockHeartbeatService.getRun.mockReset(); mockHeartbeatService.getActiveRunForAgent.mockReset(); mockHeartbeatService.cancelRun.mockReset(); + mockHeartbeatService.cancelIssueRuns.mockReset(); mockAgentService.getById.mockReset(); mockAgentService.list.mockReset(); mockAgentService.resolveByReference.mockReset(); @@ -242,6 +245,7 @@ describe.sequential("issue comment reopen routes", () => { mockHeartbeatService.getRun.mockResolvedValue(null); mockHeartbeatService.getActiveRunForAgent.mockResolvedValue(null); mockHeartbeatService.cancelRun.mockResolvedValue(null); + mockHeartbeatService.cancelIssueRuns.mockResolvedValue({ runs: [], wakeups: [] }); mockLogActivity.mockResolvedValue(undefined); mockFeedbackService.listIssueVotesForUser.mockResolvedValue([]); mockFeedbackService.saveIssueVote.mockResolvedValue({ @@ -967,17 +971,14 @@ describe.sequential("issue comment reopen routes", () => { ...issue, ...patch, })); - mockHeartbeatService.getRun.mockResolvedValue({ - id: "run-1", - companyId: "company-1", - agentId: "22222222-2222-4222-8222-222222222222", - status: "running", - }); - mockHeartbeatService.cancelRun.mockResolvedValue({ - id: "run-1", - companyId: "company-1", - agentId: "22222222-2222-4222-8222-222222222222", - status: "cancelled", + mockHeartbeatService.cancelIssueRuns.mockResolvedValue({ + runs: [{ + id: "run-1", + companyId: "company-1", + agentId: "22222222-2222-4222-8222-222222222222", + status: "cancelled", + }], + wakeups: [], }); const res = await request(await installActor(createApp())) @@ -985,8 +986,11 @@ describe.sequential("issue comment reopen routes", () => { .send({ status: "cancelled" }); expect(res.status).toBe(200); - expect(mockHeartbeatService.getRun).toHaveBeenCalledWith("run-1"); - expect(mockHeartbeatService.cancelRun).toHaveBeenCalledWith("run-1"); + expect(mockHeartbeatService.cancelIssueRuns).toHaveBeenCalledWith({ + companyId: "company-1", + issueId: "11111111-1111-4111-8111-111111111111", + reason: "Cancelled because the issue was marked cancelled", + }); expect(mockLogActivity).toHaveBeenCalledWith( expect.anything(), expect.objectContaining({ @@ -1022,6 +1026,7 @@ describe.sequential("issue comment reopen routes", () => { expect(res.status).toBe(200); expect(mockHeartbeatService.cancelRun).not.toHaveBeenCalled(); + expect(mockHeartbeatService.cancelIssueRuns).not.toHaveBeenCalled(); }); it("writes decision ids into executionState and inserts the decision inside the transaction", async () => { diff --git a/server/src/routes/issues.ts b/server/src/routes/issues.ts index 9ce8447c0e9..f7c5973964d 100644 --- a/server/src/routes/issues.ts +++ b/server/src/routes/issues.ts @@ -2115,10 +2115,6 @@ export function issueRoutes( } } - const runToCancelForCancelledStatus = shouldCancelActiveRunForCancelledStatus - ? await resolveActiveIssueRun(existing) - : null; - if (hiddenAtRaw !== undefined) { updateFields.hiddenAt = hiddenAtRaw ? new Date(hiddenAtRaw) : null; } @@ -2301,26 +2297,50 @@ export function issueRoutes( return; } - let cancelledStatusRunId: string | null = null; - if (runToCancelForCancelledStatus) { + const cancelledStatusRunIds: string[] = []; + const cancelledStatusWakeupIds: string[] = []; + if (shouldCancelActiveRunForCancelledStatus) { try { - const cancelled = await heartbeat.cancelRun(runToCancelForCancelledStatus.id); - if (cancelled) { - cancelledStatusRunId = cancelled.id; + const cancelled = await heartbeat.cancelIssueRuns({ + companyId: issue.companyId, + issueId: issue.id, + reason: "Cancelled because the issue was marked cancelled", + }); + for (const cancelledRun of cancelled.runs) { + cancelledStatusRunIds.push(cancelledRun.id); await logActivity(db, { - companyId: cancelled.companyId, + companyId: cancelledRun.companyId, actorType: actor.actorType, actorId: actor.actorId, agentId: actor.agentId, runId: actor.runId, action: "heartbeat.cancelled", entityType: "heartbeat_run", - entityId: cancelled.id, - details: { agentId: cancelled.agentId, source: "issue_status_cancelled", issueId: existing.id }, + entityId: cancelledRun.id, + details: { agentId: cancelledRun.agentId, source: "issue_status_cancelled", issueId: existing.id }, + }); + } + for (const cancelledWakeup of cancelled.wakeups) { + cancelledStatusWakeupIds.push(cancelledWakeup.id); + await logActivity(db, { + companyId: issue.companyId, + actorType: actor.actorType, + actorId: actor.actorId, + agentId: actor.agentId, + runId: actor.runId, + action: "issue.wakeup_cancelled", + entityType: "agent_wakeup_request", + entityId: cancelledWakeup.id, + details: { + agentId: cancelledWakeup.agentId, + previousReason: cancelledWakeup.reason, + source: "issue_status_cancelled", + issueId: existing.id, + }, }); } } catch (err) { - logger.warn({ err, issueId: existing.id, runId: runToCancelForCancelledStatus.id }, "failed to cancel run for cancelled issue"); + logger.warn({ err, issueId: existing.id }, "failed to cancel runs for cancelled issue"); await logActivity(db, { companyId: existing.companyId, actorType: actor.actorType, @@ -2328,8 +2348,8 @@ export function issueRoutes( agentId: actor.agentId, runId: actor.runId, action: "heartbeat.cancel_failed", - entityType: "heartbeat_run", - entityId: runToCancelForCancelledStatus.id, + entityType: "issue", + entityId: existing.id, details: { source: "issue_status_cancelled", issueId: existing.id }, }); } @@ -2401,7 +2421,13 @@ export function issueRoutes( ...(resumeRequested === true ? { resumeIntent: true, followUpRequested: true } : {}), ...(reopened ? { reopened: true, reopenedFrom: reopenFromStatus } : {}), ...(interruptedRunId ? { interruptedRunId } : {}), - ...(cancelledStatusRunId ? { cancelledStatusRunId } : {}), + ...(cancelledStatusRunIds.length > 0 + ? { + cancelledStatusRunId: cancelledStatusRunIds[0], + cancelledStatusRunIds, + } + : {}), + ...(cancelledStatusWakeupIds.length > 0 ? { cancelledStatusWakeupIds } : {}), _previous: hasFieldChanges ? previous : undefined, ...summarizeIssueReferenceActivityDetails( updateReferenceDiff diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index a79c8f79b08..39ea8722f04 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -7304,7 +7304,11 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) return wakeupIds.length; } - async function cancelRunInternal(runId: string, reason = "Cancelled by control plane") { + async function cancelRunInternal( + runId: string, + reason = "Cancelled by control plane", + opts?: { startNextQueuedRun?: boolean }, + ) { const run = await getRun(runId); if (!run) throw notFound("Heartbeat run not found"); if (!CANCELLABLE_HEARTBEAT_RUN_STATUSES.includes(run.status as (typeof CANCELLABLE_HEARTBEAT_RUN_STATUSES)[number])) return run; @@ -7354,10 +7358,90 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) runningProcesses.delete(run.id); await finalizeAgentStatus(run.agentId, "cancelled"); - await startNextQueuedRunForAgent(run.agentId); + if (opts?.startNextQueuedRun !== false) { + await startNextQueuedRunForAgent(run.agentId); + } return cancelled; } + async function cancelIssueRunsInternal(input: { + companyId: string; + issueId: string; + reason?: string; + }) { + const reason = input.reason ?? "Cancelled because the issue was cancelled"; + const now = new Date(); + + const cancelledWakeups = await db + .update(agentWakeupRequests) + .set({ + status: "cancelled", + finishedAt: now, + error: reason, + updatedAt: now, + }) + .where( + and( + eq(agentWakeupRequests.companyId, input.companyId), + inArray(agentWakeupRequests.status, ["queued", "deferred_issue_execution"]), + isNull(agentWakeupRequests.runId), + sql`${agentWakeupRequests.payload} ->> 'issueId' = ${input.issueId}`, + ), + ) + .returning({ + id: agentWakeupRequests.id, + agentId: agentWakeupRequests.agentId, + reason: agentWakeupRequests.reason, + status: agentWakeupRequests.status, + payload: agentWakeupRequests.payload, + }); + + const issue = await db + .select({ + executionRunId: issues.executionRunId, + }) + .from(issues) + .where(and(eq(issues.companyId, input.companyId), eq(issues.id, input.issueId))) + .then((rows) => rows[0] ?? null); + + const contextIssueId = sql`${heartbeatRuns.contextSnapshot} ->> 'issueId'`; + const runs = await db + .select() + .from(heartbeatRuns) + .where( + and( + eq(heartbeatRuns.companyId, input.companyId), + inArray(heartbeatRuns.status, [...CANCELLABLE_HEARTBEAT_RUN_STATUSES]), + issue?.executionRunId + ? or(eq(heartbeatRuns.id, issue.executionRunId), sql`${contextIssueId} = ${input.issueId}`) + : sql`${contextIssueId} = ${input.issueId}`, + ), + ) + .orderBy( + sql`case when ${heartbeatRuns.status} = 'running' then 0 when ${heartbeatRuns.status} = 'queued' then 1 else 2 end`, + asc(heartbeatRuns.createdAt), + ); + + const cancelledRuns: Array = []; + const affectedAgentIds = new Set(); + for (const run of runs) { + const cancelled = await cancelRunInternal(run.id, reason, { startNextQueuedRun: false }); + affectedAgentIds.add(run.agentId); + if (cancelled?.status === "cancelled") { + cancelledRuns.push(cancelled); + } + } + + for (const agentId of affectedAgentIds) { + await startNextQueuedRunForAgent(agentId); + } + + return { + runs: cancelledRuns, + wakeups: cancelledWakeups, + }; + } + async function cancelActiveForAgentInternal(agentId: string, reason = "Cancelled due to agent pause") { const agent = await getAgent(agentId); const runs = await db @@ -7734,6 +7818,8 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) cancelRun: (runId: string, reason?: string) => cancelRunInternal(runId, reason), + cancelIssueRuns: (input: { companyId: string; issueId: string; reason?: string }) => cancelIssueRunsInternal(input), + cancelActiveForAgent: (agentId: string) => cancelActiveForAgentInternal(agentId), cancelBudgetScopeWork,