fix: preserve Stop through computer admission retry races

server/src/modules/wake-queue/adapters/postgres.ts

@@ -1019,7 +1019,7 @@ export function createPostgresWakeQueueAdapter(db: Db, deps: WakeQueuePostgresAd
     async withIssueExecutionLock(input, fn): Promise<ReleaseTransactionResult & { run: RunSnapshot }> {
       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)

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) });

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<HeartbeatRetryDependencies> = {}) {
     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<void>(resolve => { finishStop = resolve; });
+    const recorded = new Promise<void>(resolve => { stopRecorded = resolve; });
+    let stopping: ReturnType<typeof control.cancelRunInternal> | 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);

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);