diff --git a/server/src/services/heartbeat/retries.test.ts b/server/src/services/heartbeat/retries.test.ts index 7176303017..2ea61f9c7e 100644 --- a/server/src/services/heartbeat/retries.test.ts +++ b/server/src/services/heartbeat/retries.test.ts @@ -1,8 +1,9 @@ +import { recoveryService } from "../recovery/service.js"; import { canRetryComputerAdmissionWait } from "../cancelled-native-startup.js"; import { ComputerStopPendingError } from "../../modules/computers/index.js"; import { randomUUID } from "node:crypto"; import type { Db } from "@paperclipai/db"; -import { computers, environments, environmentLeases, agents, agentWakeupRequests, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, heartbeatRunEvents, heartbeatRuns, issues } from "@paperclipai/db"; +import { activityLog, computers, environments, environmentLeases, agents, agentWakeupRequests, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, heartbeatRunEvents, heartbeatRuns, issues } from "@paperclipai/db"; import { and, eq } from "drizzle-orm"; import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest"; import { HttpError } from "../../errors.js"; @@ -189,6 +190,7 @@ describePostgres("heartbeat retry module database wiring", () => { afterAll(async () => { await database?.cleanup(); }); afterEach(async () => { + await db.delete(activityLog); await db.delete(environmentLeases); await db.delete(computers); await db.delete(environments); @@ -268,6 +270,65 @@ describePostgres("heartbeat retry module database wiring", () => { expect(queued.status).toBe("queued"); }); + it.each([false, true])("preserves a crashed computer wait through stale-lock recovery before scheduling one successor (cleanup receipt %s)", async settled => { + const f = await computerWaitFixture(settled); + const liveDeadline = new Date(Date.now() + 60_000); + const [waiting] = await db.update(heartbeatRuns).set({ controllerLeaseExpiresAt: liveDeadline }) + .where(eq(heartbeatRuns.id, f.run.id)).returning(); + // Crash after expected cancellation, before the successor transaction. A + // live old controller deadline intentionally delays preparation recovery. + expect(await canRetryComputerAdmissionWait(db, waiting)).toBe(settled); + const sweep = await recoveryService(db, { enqueueWakeup: vi.fn() }).sweepStaleIssueLocks(); + expect(sweep.issueIds).not.toContain(f.issue.id); + expect((await db.select().from(issues).where(eq(issues.id, f.issue.id)))[0]).toMatchObject({ + executionRunId: waiting.id, checkoutRunId: waiting.id, + }); + expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, waiting.id)))[0].controllerLeaseExpiresAt) + .toEqual(liveDeadline); + if (!settled) { + await boundRetries().retries.promoteDueScheduledRetries(new Date()); + expect(await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, waiting.id))).toHaveLength(0); + await db.update(heartbeatRuns).set({ controllerLeaseExpiresAt: new Date(Date.now() - 1) }) + .where(eq(heartbeatRuns.id, waiting.id)); + } + await boundRetries().retries.promoteDueScheduledRetries(new Date()); + await boundRetries().retries.promoteDueScheduledRetries(new Date()); + const successors = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, waiting.id)); + expect(successors).toHaveLength(1); + expect(successors[0]).toMatchObject({ status: "scheduled_retry", scheduledRetryReason: "computer_admission_wait", + contextSnapshot: { preserved: "original context", executionRetryAccounting: { failureRetries: 0 } } }); + expect((await db.select().from(issues).where(eq(issues.id, f.issue.id)))[0].executionRunId).toBe(successors[0].id); + await boundRetries().retries.promoteDueScheduledRetries(new Date(Date.now() + 31_000)); + expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, successors[0].id)))[0].status).toBe("queued"); + }); + + it.each(["malformed", "foreign-computer", "foreign-task", "foreign-company", "provider-event", "lease", "acknowledged-stop", "done", "cancelled"])( + "does not preserve stale task ownership from invalid or completed computer wait: %s", async kind => { + const f = await computerWaitFixture(false); + const result = structuredClone(f.run.resultJson!); + if (kind === "malformed") delete result.computerAdmission; + if (kind === "foreign-computer") (result.computerAdmission as Record).computerId = randomUUID(); + if (kind === "acknowledged-stop") { + result.computerAdmissionRetryOutcome = "aborted"; + result.executionCancellation = { state: "acknowledged", acknowledgedAt: new Date().toISOString() }; + } + await db.update(heartbeatRuns).set({ resultJson: result, + ...(kind === "foreign-task" ? { contextSnapshot: { ...f.run.contextSnapshot, issueId: randomUUID() } } : {}), + }).where(eq(heartbeatRuns.id, f.run.id)); + if (kind === "provider-event") await db.insert(heartbeatRunEvents).values({ companyId: f.companyId, + runId: f.run.id, agentId: f.agent.id, seq: 1, eventType: "adapter.invoke", stream: "system" }); + if (kind === "lease") await db.insert(environmentLeases).values({ companyId: f.companyId, + environmentId: f.environment.id, heartbeatRunId: f.run.id, provider: "boat", status: "released", releasedAt: now }); + if (kind === "foreign-company") { + const other = await fixture(); + await db.update(issues).set({ companyId: other.companyId, assigneeAgentId: other.agent.id }).where(eq(issues.id, f.issue.id)); + } + if (kind === "done" || kind === "cancelled") await db.update(issues).set({ status: kind }).where(eq(issues.id, f.issue.id)); + expect((await recoveryService(db, { enqueueWakeup: vi.fn() }).sweepStaleIssueLocks()).issueIds).toContain(f.issue.id); + expect((await db.select().from(issues).where(eq(issues.id, f.issue.id)))[0]).toMatchObject({ executionRunId: null, checkoutRunId: null }); + }, + ); + it("rotates proof-ineligible waits so later recoverable waits are not starved", async () => { const f = await computerWaitFixture(); await db.update(heartbeatRuns).set({ resultJson: { ...f.run.resultJson, computerAdmissionRetryCheckedAt: now.toISOString() } }).where(eq(heartbeatRuns.id, f.run.id)); diff --git a/server/src/services/recovery/service.ts b/server/src/services/recovery/service.ts index 0a45ab1619..daaeeee7c4 100644 --- a/server/src/services/recovery/service.ts +++ b/server/src/services/recovery/service.ts @@ -7,6 +7,7 @@ import { externalConversationStateSql } from "../slack-conversation-state.js"; import { isIssueReviewPathRecoveryRun } from "./review-path-recovery.js"; import { escalateExhaustedIssueReviewPathRecovery } from "./review-path-recovery-escalation.js"; import { executionRetryAccounting } from "../execution-recovery-attempt.js"; +import { isComputerAdmissionWaitBeforeProvider } from "../cancelled-native-startup.js"; import { isExplicitContinuationRetryClaim } from "../explicit-continuation-retry-claim.js"; import { decideLegacyContinuation, legacyDispositionEpisode, legacyDispositionFingerprint, @@ -6217,6 +6218,15 @@ export function recoveryService( // newer terminal status. Do not test claim ownership using the old status. const owner = originalOwner ? { ...originalOwner, status: runStatusById.get(originalOwner.id) ?? originalOwner.status } : undefined; + if (owner && owner.companyId === issue.companyId && owner.contextSnapshot?.issueId === issue.id && + issue.status !== "done" && issue.status !== "cancelled" && + await isComputerAdmissionWaitBeforeProvider(db, owner)) { + // Expected cancellation retains the task claim while preparation is + // unwinding, including a crash before its retry successor is committed. + // Retry admission separately waits for cleanup or old-controller expiry; + // that policy (or Stop) owns release/transfer of this verified claim. + continue; + } if (owner && isExplicitContinuationRetryClaim(issue, owner)) { // A terminal row can still own cleanup and a pending bounded retry. // Re-enter the same policy rather than erasing its exact-owner proof.