diff --git a/server/src/services/cancelled-native-startup.ts b/server/src/services/cancelled-native-startup.ts index a215b0bd66..f7e13b561c 100644 --- a/server/src/services/cancelled-native-startup.ts +++ b/server/src/services/cancelled-native-startup.ts @@ -73,16 +73,12 @@ export async function isCancelledNativeStartup(db: Db, run: Run, coordinator: Co return !execution; } -/** Resource waiting is replayable only before any provider admission, independently of operator Stop. */ -export async function canRetryComputerAdmissionWait(db: Db, run: Run): Promise { +/** Retain task ownership for a typed wait that never admitted provider work. + * This proves no provider effects, not that the preparing controller is gone. */ +export async function isComputerAdmissionWaitBeforeProvider(db: Db, run: Run): Promise { if (!isComputerAdmissionWait(run) || !run.finishedAt || run.runtimeModeResolvedAt || run.nativeIssueId || run.nativeSessionId || run.processPid || run.processGroupId || - run.processStartedAt || run.sessionIdAfter || adapterExecutionControls.has(run.id)) return false; - const settledAt = run.resultJson?.computerAdmissionPreparationSettledAt; - const settled = typeof settledAt === "string" && Number.isFinite(Date.parse(settledAt)); - if (!settled && !(run.executionStage === "preparing" && run.controllerBootId && - run.controllerBootId !== legacyControllerBootId && run.controllerLeaseExpiresAt && - run.controllerLeaseExpiresAt <= new Date())) return false; + run.processStartedAt || run.sessionIdAfter) return false; const admission = run.resultJson!.computerAdmission as Record; const [computer] = await db.select({ id: computers.id }).from(computers).where(and( eq(computers.id, admission.computerId), eq(computers.companyId, run.companyId), @@ -105,3 +101,19 @@ export async function canRetryComputerAdmissionWait(db: Db, run: Run): Promise { + return hasSettledComputerAdmissionPreparation(run) && isComputerAdmissionWaitBeforeProvider(db, run); +} diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 23594b880d..8c2753179a 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -5386,13 +5386,17 @@ export function heartbeatService( executionStage: "settled", controllerLeaseExpiresAt: null, resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || - ${JSON.stringify(aborted - ? { computerAdmissionRetryOutcome: "aborted" } - : { computerAdmissionPreparationSettledAt: new Date().toISOString() })}::jsonb`, + case when ${aborted} or ${heartbeatRuns.resultJson}->>'computerAdmissionRetryOutcome' = 'aborted' + then ${JSON.stringify({ computerAdmissionRetryOutcome: "aborted", + executionCancellation: { state: "acknowledged", acknowledgedAt: new Date().toISOString() } })}::jsonb + else ${JSON.stringify({ computerAdmissionPreparationSettledAt: new Date().toISOString() })}::jsonb end`, }).where(and(eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.status, "cancelled"), eq(heartbeatRuns.controllerBootId, legacyControllerBootId), eq(heartbeatRuns.errorCode, "computer_admission_wait"))).returning(); - if (settledWait && !aborted) await scheduleComputerAdmissionRetry(settledWait); + if (settledWait?.resultJson?.computerAdmissionRetryOutcome === "aborted") { + await releaseIssueExecutionAndPromote(settledWait, { suppressImmediateRecovery: true }); + await finalizeAgentStatus(run.agentId, "cancelled", undefined, { wasFirstHeartbeat: timerClaimWasFirstHeartbeat(run) }); + } else if (settledWait) await scheduleComputerAdmissionRetry(settledWait); } // Terminalization precedes lease and adapter cleanup. Only now is the // owner gone; retry pending input for ordinary completions as well as Stop. diff --git a/server/src/services/heartbeat/run-control.test.ts b/server/src/services/heartbeat/run-control.test.ts index 0234f326f4..7898b9f706 100644 --- a/server/src/services/heartbeat/run-control.test.ts +++ b/server/src/services/heartbeat/run-control.test.ts @@ -1,11 +1,16 @@ +import { createHeartbeatRetries } from "./retries.js"; +import { createRunDispatch } from "../../modules/run-dispatch/index.js"; +import { ComputerStopPendingError } from "../../modules/computers/index.js"; +import { heartbeatService } from "../heartbeat.js"; +import { canRetryComputerAdmissionWait } from "../cancelled-native-startup.js"; import { randomUUID } from "node:crypto"; import type { Db } from "@paperclipai/db"; -import { agents, agentWakeupRequests, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, heartbeatRuns, issues, projects } from "@paperclipai/db"; +import { issueComments, issueRecoveryActions, computers, environments, agents, agentWakeupRequests, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, heartbeatRuns, issues, projects } from "@paperclipai/db"; import { eq, sql } from "drizzle-orm"; import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../__tests__/helpers/embedded-postgres.js"; import { runningProcesses } from "../../adapters/index.js"; -import { adapterExecutionControls } from "../adapter-execution-control.js"; +import { adapterExecutionControls, createAdapterExecutionControl } from "../adapter-execution-control.js"; import { terminateLocalService } from "../local-service-supervisor.js"; import { createHeartbeatRunState } from "./run-state.js"; import { createHeartbeatRunPreparation } from "./run-preparation.js"; @@ -170,7 +175,7 @@ describe.skipIf(!support.supported)("heartbeat run-control database wiring", () afterEach(() => runningProcesses.clear()); async function fixture() { - const [company] = await db.insert(companies).values({ name: "Run control", issuePrefix: "RC" }).returning(); + const [company] = await db.insert(companies).values({ name: "Run control", issuePrefix: "RC", defaultResponsibleUserId: "board" }).returning(); const [agent] = await db.insert(agents).values({ companyId: company.id, name: "Controlled agent", adapterType: "codex_local" }).returning(); return { company, agent }; } @@ -180,6 +185,171 @@ describe.skipIf(!support.supported)("heartbeat run-control database wiring", () return run; } + async function computerWaitFixture(settled = true) { + const f = await fixture(); + const [issue] = await db.insert(issues).values({ companyId: f.company.id, title: "Wait for computer", status: "in_progress", assigneeAgentId: f.agent.id }).returning(); + const [environment] = await db.insert(environments).values({ name: randomUUID(), driver: "computer" }).returning(); + const [computer] = await db.insert(computers).values({ companyId: f.company.id, environmentId: environment.id, providerId: randomUUID(), ledger: {} }).returning(); + const id = randomUUID(), now = new Date(); + const [run] = await db.insert(heartbeatRuns).values({ id, companyId: f.company.id, agentId: f.agent.id, + scopeKind: "issue", issueId: issue.id, status: "cancelled", runtimeMode: "legacy", errorCode: "computer_admission_wait", finishedAt: now, + executionStage: settled ? "settled" : "preparing", controllerBootId: randomUUID(), controllerLeaseExpiresAt: new Date(Date.now() + 60_000), + contextSnapshot: { issueId: issue.id, computerAdmissionDeferredWhileAssignee: true }, + resultJson: { executionRecovery: { kind: "computer_admission_wait", providerWorkStarted: false }, + computerAdmission: { companyId: f.company.id, runId: id, environmentId: environment.id, computerId: computer.id, stopId: "stop_test" }, + cancellation: { source: "control_plane", expected: true, initiator: { type: "system" } }, + ...(settled ? { computerAdmissionPreparationSettledAt: now.toISOString() } : {}) }, + }).returning(); + await db.update(issues).set({ executionRunId: run.id }).where(eq(issues.id, issue.id)); + return { ...f, issue, run }; + } + const operatorStop = { resultJson: { cancelledByActorType: "user", cancelledByUserId: "board" } }; + + function lifecycleRetries(deps = callbacks(db)) { + const runControl = createHeartbeatRunControl(db, deps); + return createHeartbeatRetries(db, { ...deps, + resolveResponsibleUserIdForRunContext: async () => "board", + getAgentInvokability: async () => ({ invokable: true }), + escalatePlanApprovalResumeFailureNeedsAttention: async () => null, + recordPlanApprovalResumeFailureRetry: async () => null, + releaseIssueExecutionAndPromote: runControl.releaseIssueExecutionAndPromote, + getWorktreeExecutionCutoff: async () => null, + applyRunDispatchPostCommitEffects: () => undefined, + runDispatch: createRunDispatch(db), + }); + } + + it("real lifecycle terminalization retains a typed wait's task lock before cleanup and only schedules afterward", async () => { + const f = await computerWaitFixture(false), deps = callbacks(db); + const control = createAdapterExecutionControl(); + adapterExecutionControls.set(f.run.id, control); + deps.activeRunExecutions.add(f.run.id); + const [running] = await db.update(heartbeatRuns).set({ status: "running", resultJson: null, errorCode: null, finishedAt: null }).where(eq(heartbeatRuns.id, f.run.id)).returning(); + const admission = f.run.resultJson!.computerAdmission as { companyId: string; environmentId: string; computerId: string; stopId: string; runId: string }; + try { + const retries = lifecycleRetries(deps); + await retries.finalizeComputerAdmissionDeferral(running, new ComputerStopPendingError(admission), true); + const cancelled = (await deps.getRun(f.run.id))!; + expect(cancelled).toMatchObject({ status: "cancelled", errorCode: "computer_admission_wait" }); + expect((await db.select().from(issues).where(eq(issues.id, f.issue.id)))[0].executionRunId).toBe(f.run.id); + expect(await db.select().from(issueRecoveryActions).where(eq(issueRecoveryActions.sourceIssueId, f.issue.id))).toHaveLength(0); + expect(await retries.scheduleBoundedRetryForRun(cancelled, f.agent, { retryReason: "computer_admission_wait" })).toMatchObject({ outcome: "not_scheduled" }); + deps.activeRunExecutions.delete(f.run.id); adapterExecutionControls.delete(f.run.id); control.finish(); + const [settled] = await db.update(heartbeatRuns).set({ executionStage: "settled", controllerLeaseExpiresAt: null, + resultJson: { ...cancelled.resultJson, computerAdmissionPreparationSettledAt: new Date().toISOString() }, + }).where(eq(heartbeatRuns.id, f.run.id)).returning(); + expect(await lifecycleRetries().scheduleBoundedRetryForRun(settled, f.agent, { retryReason: "computer_admission_wait" })).toMatchObject({ outcome: "scheduled" }); + } finally { adapterExecutionControls.delete(f.run.id); control.finish(); } + }); + + it("a real Stop between scheduler preflight and source lock prevents the successor", async () => { + const f = await computerWaitFixture(); + const heartbeat = heartbeatService(db); + const deps = callbacks(db); + // Session resolution happens after the source's preflight proof and before + // its transaction; issue/task locking must re-read the persisted Stop. + deps.resolveSessionBeforeForWakeup = async () => { + await createHeartbeatRunControl(db, deps).cancelRunInternal(f.run.id, "Stop at preflight", operatorStop); + return null; + }; + expect(await lifecycleRetries(deps).scheduleBoundedRetryForRun(f.run, f.agent, { retryReason: "computer_admission_wait" })).toMatchObject({ outcome: "not_scheduled" }); + expect(await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, f.run.id))).toHaveLength(0); + await heartbeat.promoteDueScheduledRetries(new Date(Date.now() + 60_000)); + expect(await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, f.run.id))).toHaveLength(0); + }); + + it.each([true, false])("Stop durably suppresses a computer wait before scheduling (cleanup settled=%s)", async settled => { + const f = await computerWaitFixture(settled), deps = callbacks(db); + const control = createAdapterExecutionControl(); + if (!settled) { deps.activeRunExecutions.add(f.run.id); adapterExecutionControls.set(f.run.id, control); } + try { + const stopped = await createHeartbeatRunControl(db, deps).cancelRunInternal(f.run.id, "Stop waiting", operatorStop); + expect(stopped).toMatchObject({ status: "cancelled", resultJson: { computerAdmissionRetryOutcome: "aborted", + cancellation: { source: "operator", initiator: { type: "user", id: "board" } } } }); + if (!settled) { + expect(control.controller.signal.aborted).toBe(true); + expect(deps.wakeQueue.releaseIssueExecution).not.toHaveBeenCalled(); + } + // Preparation finishing later cannot erase a real Stop, even when the + // old executor writes its cleanup receipt after the request returns. + await db.update(heartbeatRuns).set({ executionStage: "settled", controllerLeaseExpiresAt: null, + resultJson: sql`${heartbeatRuns.resultJson} || ${JSON.stringify({ computerAdmissionPreparationSettledAt: new Date().toISOString() })}::jsonb`, + }).where(eq(heartbeatRuns.id, f.run.id)); + expect(await canRetryComputerAdmissionWait(db, (await deps.getRun(f.run.id))!)).toBe(false); + expect(await heartbeatService(db).scheduleBoundedRetry(f.run.id, { retryReason: "computer_admission_wait" })).toMatchObject({ outcome: "not_scheduled" }); + expect(await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, f.run.id))).toHaveLength(0); + expect(deps.envOrchestrator.releaseForRun).not.toHaveBeenCalled(); + } finally { adapterExecutionControls.delete(f.run.id); control.finish(); } + }); + + 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); + expect(stopped?.resultJson?.computerAdmissionRetryOutcome).toBe("aborted"); + expect(stopped?.resultJson?.executionCancellation).toBeUndefined(); + 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); + }); + + it("Stop of a settled computer wait retains old queued messages instead of promoting them", async () => { + const f = await computerWaitFixture(); + const [comment] = await db.insert(issueComments).values({ companyId: f.company.id, issueId: f.issue.id, + authorType: "user", authorUserId: "board", body: "Queued before Stop" }).returning(); + const [wake] = await db.insert(agentWakeupRequests).values({ companyId: f.company.id, agentId: f.agent.id, + source: "automation", reason: "issue_execution_deferred", status: "deferred_issue_execution", requestedByActorType: "system", + payload: { issueId: f.issue.id, commentId: comment.id, _paperclipWakeContext: { wakeCommentIds: [comment.id] } }, + }).returning(); + const stopped = await heartbeatService(db).cancelRun(f.run.id, "Stop waiting", operatorStop); + expect(stopped?.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); + expect((await db.select().from(issues).where(eq(issues.id, f.issue.id)))[0].executionRunId).toBeNull(); + }); + + it.each(["scheduled_retry", "queued", "running"])("a stale-parent Stop cancels its exact %s successor through normal cancellation", async status => { + const f = await computerWaitFixture(), deps = callbacks(db); + const scheduled = await heartbeatService(db).scheduleBoundedRetry(f.run.id, { retryReason: "computer_admission_wait", delayMs: 30_000 }); + if (scheduled.outcome !== "scheduled") throw new Error("Expected scheduled successor"); + await db.update(heartbeatRuns).set({ status, ...(status === "running" ? { runtimeModeResolvedAt: new Date(), startedAt: new Date() } : {}) }).where(eq(heartbeatRuns.id, scheduled.run.id)); + if (status === "running") runningProcesses.set(scheduled.run.id, { child: { pid: 424242 } as never, graceSec: 4, processGroupId: 424242 }); + const [unrelated] = await db.insert(heartbeatRuns).values({ companyId: f.company.id, agentId: f.agent.id, status: "running", runtimeModeResolvedAt: new Date() }).returning(); + const controller = createHeartbeatRunControl(db, deps); + await controller.cancelRunInternal(f.run.id, "Stop waiting", operatorStop); + await controller.cancelRunInternal(f.run.id, "Stop again", operatorStop); + expect(await deps.getRun(scheduled.run.id)).toMatchObject({ status: "cancelled" }); + expect(await deps.getRun(unrelated.id)).toMatchObject({ status: "running" }); + expect(terminateLocalService).toHaveBeenCalledTimes(status === "running" ? 1 : 0); + expect(await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, f.run.id))).toHaveLength(1); + }); + + it("retries successor cancellation after a Stop failure without reopening parent admission", async () => { + const f = await computerWaitFixture(), deps = callbacks(db); + const scheduled = await heartbeatService(db).scheduleBoundedRetry(f.run.id, { retryReason: "computer_admission_wait", delayMs: 30_000 }); + if (scheduled.outcome !== "scheduled") throw new Error("Expected successor"); + await db.update(heartbeatRuns).set({ status: "running", runtimeModeResolvedAt: new Date() }).where(eq(heartbeatRuns.id, scheduled.run.id)); + runningProcesses.set(scheduled.run.id, { child: { pid: 424242 } as never, graceSec: 4, processGroupId: 424242 }); + vi.mocked(terminateLocalService).mockRejectedValueOnce(new Error("Stop failed")); + const controller = createHeartbeatRunControl(db, deps); + await expect(controller.cancelRunInternal(f.run.id, "Stop waiting", operatorStop)).rejects.toThrow("Stop failed"); + expect(await deps.getRun(f.run.id)).toMatchObject({ resultJson: { computerAdmissionRetryOutcome: "aborted" } }); + await controller.cancelRunInternal(f.run.id, "Stop again", operatorStop); + expect(await deps.getRun(scheduled.run.id)).toMatchObject({ status: "cancelled" }); + expect(terminateLocalService).toHaveBeenCalledTimes(2); + }); + + it.each(["agent", "reason", "issue"])("does not forward stale-parent Stop across a changed successor %s", async kind => { + const f = await computerWaitFixture(), deps = callbacks(db); + const [otherAgent] = await db.insert(agents).values({ companyId: f.company.id, name: "Reassigned" }).returning(); + const [otherIssue] = await db.insert(issues).values({ companyId: f.company.id, title: "Other task" }).returning(); + const [successor] = await db.insert(heartbeatRuns).values({ companyId: f.company.id, + agentId: kind === "agent" ? otherAgent.id : f.agent.id, scopeKind: "issue", issueId: kind === "issue" ? otherIssue.id : f.issue.id, + retryOfRunId: f.run.id, scheduledRetryReason: kind === "reason" ? "transient_failure" : "computer_admission_wait", status: "queued", + }).returning(); + await createHeartbeatRunControl(db, deps).cancelRunInternal(f.run.id, "Stop waiting", operatorStop); + expect(await deps.getRun(successor.id)).toMatchObject({ status: "queued" }); + }); + it("reports a missing run without entering cancellation", async () => { const deps = callbacks(db); await expect(createHeartbeatRunControl(db, deps).cancelRunInternal(randomUUID())).rejects.toThrow("Heartbeat run not found"); diff --git a/server/src/services/heartbeat/run-control.ts b/server/src/services/heartbeat/run-control.ts index 7006c54325..beebcafc1e 100644 --- a/server/src/services/heartbeat/run-control.ts +++ b/server/src/services/heartbeat/run-control.ts @@ -31,7 +31,7 @@ import { admitExplicitNativeContinuation, undeliveredLegacyUserCommentIds, } from "../explicit-native-continuation.js"; -import { isCancelledNativeStartup } from "../cancelled-native-startup.js"; +import { isCancelledNativeStartup, isComputerAdmissionWaitBeforeProvider, hasSettledComputerAdmissionPreparation } from "../cancelled-native-startup.js"; import { executionBlockerPredicate, getExecutionBlocker, @@ -842,16 +842,78 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont return Math.max(100, Math.min(30_000, Math.trunc(requestedGraceMs))); } + async function cancelComputerAdmissionWait( + run: typeof heartbeatRuns.$inferSelect, reason: string, options: CancelRunOptions, + ) { + const cancellation = requestedRunCancellation(options.resultJson ?? {}, reason); + const stop = async (tx: Db) => { + // Match retry scheduling's lock order, including a scheduler that already + // transferred the task to this wait's exact successor. + if (run.scopeKind === "issue" && run.issueId) await tx.select({ id: issues.id }).from(issues) + .where(and(eq(issues.companyId, run.companyId), eq(issues.id, run.issueId))).for("update"); + const [current] = await tx.select().from(heartbeatRuns).where(and( + eq(heartbeatRuns.companyId, run.companyId), eq(heartbeatRuns.id, run.id), + )).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; + const [stopped] = alreadyStopped ? [current] : await tx.update(heartbeatRuns).set({ + error: reason, updatedAt: new Date(), + resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || ${JSON.stringify({ + computerAdmissionRetryOutcome: "aborted", cancellation, + // This Stop did not cancel provider work; preparation never admitted it. + executionRecovery: { kind: "bootstrap", providerWorkStarted: false }, + })}::jsonb`, + }).where(eq(heartbeatRuns.id, current.id)).returning(); + const [successor] = await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and( + eq(heartbeatRuns.companyId, run.companyId), eq(heartbeatRuns.agentId, run.agentId), + eq(heartbeatRuns.retryOfRunId, run.id), eq(heartbeatRuns.scheduledRetryReason, "computer_admission_wait"), + eq(heartbeatRuns.scopeKind, current.scopeKind), + current.issueId ? eq(heartbeatRuns.issueId, current.issueId) : isNull(heartbeatRuns.issueId), + )).limit(1); + return { stopped, successorId: successor?.id, updated: !alreadyStopped }; + }; + const outcome = await db.transaction(tx => options.budgetEnforcement + ? withCurrentBudgetEnforcement(tx as unknown as Db, options.budgetEnforcement, stop) + : stop(tx as unknown as Db)); + if (!outcome) return getRun(run.id); + adapterExecutionControls.get(run.id)?.controller.abort(new Error(reason)); + if (outcome.successorId) { + // Native cancellation request IDs belong to one run and cannot be lent + // to its successor. The successor uses its existing owned Stop path. + const { cancellationRequestId: _requestId, cancellationRequestedByUserId: _requestUser, ...successorOptions } = options; + await cancelRunInternal(outcome.successorId, reason, successorOptions); + } + if (outcome.updated) { + await appendRunEvent(outcome.stopped, { eventType: "lifecycle", stream: "system", level: "info", + message: "computer admission retry cancelled", payload: { cancellation } }); + } + // 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) || + ${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 }); + await finalizeAgentStatus(run.agentId, "cancelled", undefined, { wasFirstHeartbeat: timerClaimWasFirstHeartbeat(run) }); + } + return getRun(run.id); + } + async function cancelRunInternal( runId: string, reason = "Cancelled by control plane", options: CancelRunOptions = {}, - ) { + ): Promise { let run = await getRun(runId); if (!run) throw notFound("Heartbeat run not found"); if (options.cancellationRequestId) { run = await claimCancellationRequest(db, runId, run.companyId, options.cancellationRequestId, options.cancellationRequestedByUserId ?? null); } + if (run.status === "cancelled" && run.errorCode === "computer_admission_wait") { + return cancelComputerAdmissionWait(run, reason, options); + } // The caller claim checked retry eligibility under both durable row locks. // This is only a cancellation candidate: dispatch rechecks the coordinator // under lock, and the final CAS requires its own unchanged acknowledged diff --git a/server/src/services/legacy-execution-recovery.ts b/server/src/services/legacy-execution-recovery.ts index 7d11a8ec6f..fc874bd90b 100644 --- a/server/src/services/legacy-execution-recovery.ts +++ b/server/src/services/legacy-execution-recovery.ts @@ -1,4 +1,4 @@ -import { canRetryComputerAdmissionWait } from "./cancelled-native-startup.js"; +import { isComputerAdmissionWaitBeforeProvider } from "./cancelled-native-startup.js"; import { hasRequiredWorkspaceRecovery, preserveWorkspaceRestoreRecoveryMetadata, LEGACY_WORKSPACE_RECOVERY_SCHEMA } from "./workspace-restore-recovery-state.js"; import { hasUnrestoredRemoteWorkspace, preserveLegacyWorkspaceRestoreSources, type LegacyWorkspaceRestoreSource } from "./legacy-workspace-restore-recovery.js"; import { isPreDispatchReviewWaitVerified } from "./pre-dispatch-review-wait.js"; @@ -62,7 +62,7 @@ export function legacyExecutionNeedsReconciliation( export async function legacyExecutionNeedsReconciliationWithEvidence(db: Db, run: Run): Promise { return await hasUnrestoredRemoteWorkspace(db, run) || (legacyExecutionNeedsReconciliation(run) && !(await isPreDispatchReviewWaitVerified(db, run)) && - !(await canRetryComputerAdmissionWait(db, run))); + !(await isComputerAdmissionWaitBeforeProvider(db, run))); } /** Persist the failed legacy run, owned lock release and operator decision together. */