diff --git a/server/src/modules/wake-queue/adapters/postgres.ts b/server/src/modules/wake-queue/adapters/postgres.ts index d2624b2512..e38e6127ea 100644 --- a/server/src/modules/wake-queue/adapters/postgres.ts +++ b/server/src/modules/wake-queue/adapters/postgres.ts @@ -1019,7 +1019,7 @@ export function createPostgresWakeQueueAdapter(db: Db, deps: WakeQueuePostgresAd async withIssueExecutionLock(input, fn): Promise { return db.transaction(async (rawTx) => { const tx = rawTx as unknown as Db; - const run = await tx + let run = await tx .select() .from(heartbeatRuns) .where(and(eq(heartbeatRuns.id, input.runId), eq(heartbeatRuns.companyId, input.companyId))) @@ -1027,7 +1027,7 @@ export function createPostgresWakeQueueAdapter(db: Db, deps: WakeQueuePostgresAd if (!run) { throw new Error(`wake-queue: run ${input.runId} was not found while releasing issue execution`); } - const runSnapshot = toRunSnapshot(run); + let runSnapshot = toRunSnapshot(run); const contextIssueId = readNonEmptyString(parseObject(run.contextSnapshot).issueId); // Lock the context issue (if any) and every issue that still references this @@ -1055,6 +1055,22 @@ export function createPostgresWakeQueueAdapter(db: Db, deps: WakeQueuePostgresAd `, ); + // Stop and retry scheduling serialize on task then run. A computer + // wait may have been stopped while this release waited for the task. + // Never drain old input from the stale pre-lock suppression snapshot. + if (run.errorCode === "computer_admission_wait") { + const [current] = await tx.select().from(heartbeatRuns).where(and( + eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.companyId, input.companyId), + )).for("update"); + if (!current) throw new Error("wake-queue: computer admission source disappeared"); + run = current; + runSnapshot = toRunSnapshot(run); + if (run.resultJson?.computerAdmissionRetryOutcome === "aborted" && + parseObject(run.resultJson?.executionCancellation).state !== "acknowledged") { + return { outcome: { kind: "released" }, postCommitEffects: [], run: runSnapshot }; + } + } + const candidateIssues = await tx .select() .from(issues) diff --git a/server/src/services/heartbeat/retries.ts b/server/src/services/heartbeat/retries.ts index ee5d00245a..8cba29f396 100644 --- a/server/src/services/heartbeat/retries.ts +++ b/server/src/services/heartbeat/retries.ts @@ -1426,16 +1426,20 @@ export function createHeartbeatRetries(db: Db, dependencies: HeartbeatRetryDepen retry.errorCode === "continuation_user_authorization_missing"; if (retry.outcome !== "scheduled") { const reason = "reason" in retry ? retry.reason : "The retry limit was reached."; - await db.update(heartbeatRuns).set({ + const [suppressed] = await db.update(heartbeatRuns).set({ ...(authorizationLost ? { status: "failed", errorCode: "computer_admission_retry_unavailable" } : {}), error: `Computer admission retry was not scheduled: ${reason}`, resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || ${JSON.stringify({ computerAdmissionRetryOutcome: "suppressed" })}::jsonb`, - }).where(and(eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.status, "cancelled"), - eq(heartbeatRuns.errorCode, COMPUTER_ADMISSION_WAIT_RETRY_REASON))); - await appendRunEvent(run, { eventType: "lifecycle", stream: "system", level: authorizationLost ? "warn" : "info", + }).where(and(eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.companyId, run.companyId), + eq(heartbeatRuns.status, "cancelled"), eq(heartbeatRuns.errorCode, COMPUTER_ADMISSION_WAIT_RETRY_REASON), + sql`${heartbeatRuns.resultJson}->>'computerAdmissionRetryOutcome' is null`)).returning(); + // Stop or another terminal owner won. Its acknowledgement and release + // must not be replaced by this scheduler's stale suppression outcome. + if (!suppressed) return; + await appendRunEvent(suppressed, { eventType: "lifecycle", stream: "system", level: authorizationLost ? "warn" : "info", message: `Computer admission retry was not scheduled: ${reason}`, payload: { retryScheduled: false } }); - await releaseIssueExecutionAndPromote(run); + await releaseIssueExecutionAndPromote(suppressed); } await finalizeAgentStatus(run.agentId, authorizationLost ? "failed" : "cancelled", null, { wasFirstHeartbeat: timerClaimWasFirstHeartbeat(run) }); diff --git a/server/src/services/heartbeat/run-control.test.ts b/server/src/services/heartbeat/run-control.test.ts index a06b7501b8..bac6f5a93f 100644 --- a/server/src/services/heartbeat/run-control.test.ts +++ b/server/src/services/heartbeat/run-control.test.ts @@ -1,7 +1,9 @@ +import { createReleaseIssueExecution } from "../../modules/wake-queue/application/use-cases.js"; +import { createPostgresWakeQueueAdapter } from "../../modules/wake-queue/adapters/postgres.js"; import { admitExplicitNativeContinuation } from "../explicit-native-continuation.js"; import { legacyExecutionNeedsReconciliationWithEvidence, terminalizeLegacyExecution } from "../legacy-execution-recovery.js"; import { createPostgresRunDispatchAdapter } from "../../modules/run-dispatch/adapters/postgres.js"; -import { createHeartbeatRetries } from "./retries.js"; +import { createHeartbeatRetries, type HeartbeatRetryDependencies } from "./retries.js"; import { createRunDispatch } from "../../modules/run-dispatch/index.js"; import { ComputerStopPendingError } from "../../modules/computers/index.js"; import { heartbeatService } from "../heartbeat.js"; @@ -302,7 +304,7 @@ describe.skipIf(!support.supported)("heartbeat run-control database wiring", () const operatorStop = { resultJson: { cancelledByActorType: "user", cancelledByUserId: "board" } }; - function lifecycleRetries(deps = callbacks(db)) { + function lifecycleRetries(deps = callbacks(db), overrides: Partial = {}) { const runControl = createHeartbeatRunControl(db, deps); return createHeartbeatRetries(db, { ...deps, resolveResponsibleUserIdForRunContext: async () => "board", @@ -313,6 +315,7 @@ describe.skipIf(!support.supported)("heartbeat run-control database wiring", () getWorktreeExecutionCutoff: async () => null, applyRunDispatchPostCommitEffects: () => undefined, runDispatch: createRunDispatch(db), + ...overrides, }); } @@ -379,6 +382,74 @@ describe.skipIf(!support.supported)("heartbeat run-control database wiring", () } finally { adapterExecutionControls.delete(f.run.id); control.finish(); } }); + it.each(["stop-first", "suppression-first"])("the full retry wrapper preserves Stop and old queued input (%s)", async ordering => { + const f = await computerWaitFixture(), deps = callbacks(db); + deps.wakeQueue.releaseIssueExecution.mockImplementation(createReleaseIssueExecution({ + issueLock: createPostgresWakeQueueAdapter(db, { resolveResponsibleUserId: async () => "board", + getRoutineEnv: async () => ({ routineId: null, env: null, responsibleUserId: "board" }), + resolveSessionBeforeForWakeup: async () => null }), + recovery: { escalateStrandedAssignedIssue: async () => {}, escalateStrandedRecoveryIssueInPlace: async () => {} }, + })); + const control = createHeartbeatRunControl(db, deps); + const [comment] = await db.insert(issueComments).values({ companyId: f.company.id, issueId: f.issue.id, + authorType: "user", authorUserId: "board", body: "Input queued before Stop" }).returning(); + const [wake] = await db.insert(agentWakeupRequests).values({ companyId: f.company.id, agentId: f.agent.id, + source: "automation", reason: "issue_commented", status: "deferred_issue_execution", + requestedByActorType: "user", requestedByActorId: "board", payload: { issueId: f.issue.id, commentId: comment.id, + _paperclipWakeContext: { issueId: f.issue.id, wakeCommentId: comment.id, wakeCommentIds: [comment.id] } }, + }).returning(); + let finishStop!: () => void, stopRecorded!: () => void; + const stopGate = new Promise(resolve => { finishStop = resolve; }); + const recorded = new Promise(resolve => { stopRecorded = resolve; }); + let stopping: ReturnType | undefined; + const append = deps.appendRunEvent.getMockImplementation()!; + deps.appendRunEvent.mockImplementation(async (run, event) => { + await append(run, event); + if (event.message === "computer admission retry cancelled") { + stopRecorded(); await stopGate; + } + if (ordering === "suppression-first" && event.message?.startsWith("Computer admission retry was not scheduled:")) { + stopping = control.cancelRunInternal(f.run.id, "Stop waiting", operatorStop); + await recorded; + } + }); + if (ordering === "stop-first") deps.getAgent.mockImplementation(async () => { + stopping = control.cancelRunInternal(f.run.id, "Stop waiting", operatorStop); + await recorded; + return f.agent; + }); + try { + await lifecycleRetries(deps, { getAgentInvokability: async () => ordering === "stop-first" ? { invokable: true } + : { invokable: false, reason: "paused", message: "Paused", details: {}, invalidOrgChain: false } }) + .scheduleComputerAdmissionRetry(f.run); + const [current] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.run.id)); + expect(current.resultJson?.computerAdmissionRetryOutcome).toBe("aborted"); + if (ordering === "stop-first") { + expect(deps.wakeQueue.releaseIssueExecution).not.toHaveBeenCalled(); + } + expect((await db.select().from(issues).where(eq(issues.id, f.issue.id)))[0].executionRunId).toBe(f.run.id); + expect((await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, wake.id)))[0].status).toBe("deferred_issue_execution"); + } finally { finishStop(); await stopping; } + expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.run.id)))[0].resultJson?.executionCancellation) + .toMatchObject({ state: "acknowledged" }); + expect((await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, wake.id)))[0]) + .toMatchObject({ status: "deferred_issue_execution", runId: null }); + expect(await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.companyId, f.company.id))).toHaveLength(1); + }); + + it("does not release a Stop whose acknowledgement CAS loses terminal ownership", async () => { + const f = await computerWaitFixture(), deps = callbacks(db); + const append = deps.appendRunEvent.getMockImplementation()!; + deps.appendRunEvent.mockImplementation(async (run, event) => { + await append(run, event); + if (event.message === "computer admission retry cancelled") await db.update(heartbeatRuns) + .set({ status: "failed" }).where(eq(heartbeatRuns.id, f.run.id)); + }); + await createHeartbeatRunControl(db, deps).cancelRunInternal(f.run.id, "Stop waiting", operatorStop); + expect(deps.wakeQueue.releaseIssueExecution).not.toHaveBeenCalled(); + expect(deps.finalizeAgentStatus).not.toHaveBeenCalled(); + }); + it("Stop persists suppression but leaves a remote controller's live cleanup unacknowledged", async () => { const f = await computerWaitFixture(false), deps = callbacks(db); const stopped = await createHeartbeatRunControl(db, deps).cancelRunInternal(f.run.id, "Stop remote preparation", operatorStop); diff --git a/server/src/services/heartbeat/run-control.ts b/server/src/services/heartbeat/run-control.ts index 1aeb7943c8..4586f29bb8 100644 --- a/server/src/services/heartbeat/run-control.ts +++ b/server/src/services/heartbeat/run-control.ts @@ -858,7 +858,13 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont )).for("update"); if (!current || current.status !== "cancelled" || current.errorCode !== "computer_admission_wait") return null; const alreadyStopped = current.resultJson?.computerAdmissionRetryOutcome === "aborted"; - if (!alreadyStopped && !(await isComputerAdmissionWaitBeforeProvider(tx, current))) return null; + // A scheduler may have suppressed its successor but not released the + // task yet. Stop can still claim this exact verified pre-provider wait; + // removing only that disposition for proof does not authorize a retry. + const beforeSuppression = current.resultJson?.computerAdmissionRetryOutcome === "suppressed" + ? { ...current, resultJson: { ...current.resultJson, computerAdmissionRetryOutcome: undefined } } + : current; + if (!alreadyStopped && !(await isComputerAdmissionWaitBeforeProvider(tx, beforeSuppression))) return null; const [stopped] = alreadyStopped ? [current] : await tx.update(heartbeatRuns).set({ error: reason, updatedAt: new Date(), resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || ${JSON.stringify({ @@ -893,11 +899,12 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont // The executor's finally block owns this acknowledgement while cleanup is // live. Until then retain the task lock and do not drain queued messages. if (!activeRunExecutions.has(run.id) && hasSettledComputerAdmissionPreparation(outcome.stopped)) { - await db.update(heartbeatRuns).set({ resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || + const [acknowledged] = await db.update(heartbeatRuns).set({ resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || ${JSON.stringify({ executionCancellation: { state: "acknowledged", acknowledgedAt: new Date().toISOString() } })}::jsonb`, - }).where(and(eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.status, "cancelled"), - sql`${heartbeatRuns.resultJson}->>'computerAdmissionRetryOutcome' = 'aborted'`)); - await releaseIssueExecutionAndPromote(outcome.stopped, { suppressImmediateRecovery: true }); + }).where(and(eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.companyId, run.companyId), eq(heartbeatRuns.status, "cancelled"), + sql`${heartbeatRuns.resultJson}->>'computerAdmissionRetryOutcome' = 'aborted'`)).returning(); + if (!acknowledged) return getRun(run.id); + await releaseIssueExecutionAndPromote(acknowledged, { suppressImmediateRecovery: true }); await finalizeAgentStatus(run.agentId, "cancelled", undefined, { wasFirstHeartbeat: timerClaimWasFirstHeartbeat(run) }); } return getRun(run.id);