diff --git a/server/src/__tests__/native-cancellation-request.integration.test.ts b/server/src/__tests__/native-cancellation-request.integration.test.ts index 15d4d554471..9ada636e8f2 100644 --- a/server/src/__tests__/native-cancellation-request.integration.test.ts +++ b/server/src/__tests__/native-cancellation-request.integration.test.ts @@ -1,21 +1,26 @@ import { randomUUID } from "node:crypto"; import { afterAll, beforeAll, describe, expect, it } from "vitest"; import { eq, sql } from "drizzle-orm"; -import { agents, companies, createDb, heartbeatRuns } from "@paperclipai/db"; +import { agents, companies, createDb, heartbeatRuns, issues, nativeRunFinalizations } from "@paperclipai/db"; import { startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js"; import { cancellationIntentId, claimCancellationRequest, startupCancellationFence } from "../services/native-runtime/native-cancellation-request.js"; describe("atomic caller cancellation request ownership", () => { let database: Awaited>; let db: ReturnType; - beforeAll(async () => { database = await startEmbeddedPostgresTestDatabase("caller-stop-"); db = createDb(database.connectionString); }, 90_000); + let cancelNativeSession: typeof import("../services/native-runtime/native-session-executor.js").cancelNativeSession; + beforeAll(async () => { + ({ cancelNativeSession } = await import("../services/native-runtime/native-session-executor.js")); + database = await startEmbeddedPostgresTestDatabase("caller-stop-"); + db = createDb(database.connectionString); + }, 90_000); afterAll(async () => { await database?.cleanup(); }); async function fixture() { const companyId = randomUUID(), agentId = randomUUID(), runId = randomUUID(); await db.insert(companies).values({ id: companyId, name: "Stop claims", issuePrefix: randomUUID().slice(0, 8) }); await db.insert(agents).values({ id: agentId, companyId, name: "Stop target" }); await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, invocationSource: "on_demand", status: "running", runtimeMode: "native" }); - return { companyId, runId }; + return { companyId, runId, agentId }; } const defaultFence = (runId: string) => db.update(heartbeatRuns).set({ resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || jsonb_build_object('startupCancellation', ${startupCancellationFence(new Date().toISOString())})` }).where(eq(heartbeatRuns.id, runId)); it("allows exactly one of two concurrent identities and idempotent same-ID retries", async () => { @@ -51,4 +56,49 @@ describe("atomic caller cancellation request ownership", () => { await db.update(heartbeatRuns).set({ resultJson: { ...run.resultJson, nativeCancellation: { intentId: cancellationIntentId(id) } } }).where(eq(heartbeatRuns.id, runId)); await expect(claimCancellationRequest(db, runId, companyId, id, "board-user")).rejects.toMatchObject({ status: 409 }); }); + async function retryFixture(phase = "retryable_failure") { + const f = await fixture(), issueId = randomUUID(); + await db.insert(issues).values({ id: issueId, companyId: f.companyId, title: "Native retry", status: "in_progress", assigneeAgentId: f.agentId }); + await db.update(heartbeatRuns).set({ nativeIssueId: issueId, status: "failed" }).where(eq(heartbeatRuns.id, f.runId)); + await db.insert(nativeRunFinalizations).values({ runId: f.runId, companyId: f.companyId, issueId, phase, nextAttemptAt: new Date(Date.now() + 30_000) }); + return { ...f, issueId }; + } + it("admits a failed retry and recovers the same caller after audited native cancellation disables it", async () => { + const f = await retryFixture(), id = randomUUID(); + await claimCancellationRequest(db, f.runId, f.companyId, id, "board-user"); + // No provider is attached: exercise the actual durable cancellation path. + const outcome = await cancelNativeSession(f.runId, "Stop", { db, scope: "run", cancellationRequestId: id }); + expect(outcome.auditId).toBeTruthy(); + const [coordinator] = await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)); + expect(coordinator).toMatchObject({ phase: "terminal_failure", failureCode: "native_retry_cancelled", nextAttemptAt: null }); + const retried = await claimCancellationRequest(db, f.runId, f.companyId, id, "board-user"); + expect(retried.resultJson?.nativeCancellation).toMatchObject({ intentId: cancellationIntentId(id), dispatchState: "acknowledged" }); + await expect(claimCancellationRequest(db, f.runId, f.companyId, randomUUID(), "board-user")).rejects.toMatchObject({ status: 409 }); + await expect(claimCancellationRequest(db, f.runId, f.companyId, id, "other-actor")).rejects.toMatchObject({ status: 409 }); + await expect(cancelNativeSession(f.runId, "Stop", { db, scope: "run", cancellationRequestId: id })).resolves.toMatchObject({ auditId: outcome.auditId }); + }); + it.each(["observed", "terminal_failure", "committed"])("does not claim a failed run whose coordinator advanced to %s", async phase => { + const f = await retryFixture(phase); + await expect(claimCancellationRequest(db, f.runId, f.companyId, randomUUID(), "board-user")).rejects.toMatchObject({ status: 409 }); + const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + expect(run.resultJson?.startupCancellation).toBeUndefined(); + }); + it("rejects a coordinator-only concurrent claim, then observes its committed phase without installing a Stop fence", async () => { + const f = await retryFixture(), id = randomUUID(); + let release!: () => void, locked!: () => void; + const acquired = new Promise(resolve => { locked = resolve; }); + const hold = new Promise(resolve => { release = resolve; }); + const advance = db.transaction(async tx => { + await tx.update(nativeRunFinalizations).set({ phase: "observed", nextAttemptAt: null }).where(eq(nativeRunFinalizations.runId, f.runId)); + locked(); await hold; + }); + try { + await acquired; + await expect(claimCancellationRequest(db, f.runId, f.companyId, id, "board-user")).rejects.toMatchObject({ status: 409, message: "Native retry cancellation is busy; retry the same request" }); + } finally { release(); await advance; } + await expect(claimCancellationRequest(db, f.runId, f.companyId, id, "board-user")).rejects.toMatchObject({ status: 409 }); + const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + expect(run.resultJson?.startupCancellation).toBeUndefined(); + }); + }); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index ba39981ae78..1e165d578ff 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -29146,20 +29146,23 @@ export function heartbeatService( if (options.cancellationRequestId) { run = await claimCancellationRequest(db, runId, run.companyId, options.cancellationRequestId, options.cancellationRequestedByUserId ?? null); } + // The caller claim checked retry eligibility under both durable row locks. + // Its own native intent may already have disabled the retry; re-reading only + // retryable_failure would strand that same-ID recovery. const pendingNativeRetry = - run.runtimeMode === "native" && run.status === "failed" - ? await db - .select({ runId: nativeRunFinalizations.runId }) - .from(nativeRunFinalizations) - .where( - and( - eq(nativeRunFinalizations.runId, run.id), - eq(nativeRunFinalizations.companyId, run.companyId), - eq(nativeRunFinalizations.phase, "retryable_failure"), - ), - ) - .then((rows) => rows.length > 0) - : false; + run.runtimeMode === "native" && run.status === "failed" && ( + Boolean(options.cancellationRequestId) || await db + .select({ runId: nativeRunFinalizations.runId }) + .from(nativeRunFinalizations) + .where( + and( + eq(nativeRunFinalizations.runId, run.id), + eq(nativeRunFinalizations.companyId, run.companyId), + eq(nativeRunFinalizations.phase, "retryable_failure"), + ), + ) + .then((rows) => rows.length > 0) + ); if ( !pendingNativeRetry && !CANCELLABLE_HEARTBEAT_RUN_STATUSES.includes( diff --git a/server/src/services/native-runtime/native-cancellation-request.test.ts b/server/src/services/native-runtime/native-cancellation-request.test.ts index 352f1f80ea3..0706ec5b57b 100644 --- a/server/src/services/native-runtime/native-cancellation-request.test.ts +++ b/server/src/services/native-runtime/native-cancellation-request.test.ts @@ -1,5 +1,6 @@ -import { describe, expect, it } from "vitest"; -import { assertCancellationRequest, cancellationIntentId, cancellationRequestId } from "./native-cancellation-request.js"; +import { describe, expect, it, vi } from "vitest"; +import { assertCancellationRequest, cancellationIntentId, cancellationRequestId, claimCancellationRequest } from "./native-cancellation-request.js"; +import { heartbeatRuns, type Db } from "@paperclipai/db"; const id = "11111111-1111-4111-8111-111111111111"; describe("caller-bound native cancellation", () => { it.each([null, "", " x", {}, [], 1, `${id} `, "not-a-uuid"])("rejects malformed request %j", value => { @@ -21,3 +22,57 @@ describe("caller-bound native cancellation", () => { expect(() => assertCancellationRequest({ startupCancellation, nativeCancellation: { intentId: cancellationIntentId(id) } }, id)).not.toThrow(); }); }); + + +describe("failed native retry caller admission", () => { + const companyId = "company", runId = "run", issueId = "issue"; + const claimed = { cancellationRequestId: id, requestedBy: { type: "board", userId: "board-user" } }; + const native = { schema: "paperclip.native-cancellation.v1", intentId: cancellationIntentId(id), + companyId, runId, issueId, scope: "run", dispatchState: "pending", dispatched: false, intentAuditId: "intent-audit" }; + function fixture(phase: string | null, resultJson: Record = {}, status = "failed", failureCode: string | null = null) { + const run = { id: runId, companyId, nativeIssueId: issueId, runtimeMode: "native", status, resultJson }; + const locks: unknown[][] = []; + const update = vi.fn((values: { resultJson: Record }) => { run.resultJson = values.resultJson; + return { where: () => ({ returning: async () => [run] }) }; }); + const tx = { select: () => ({ from: (table: unknown) => ({ where: () => ({ + for: (...args: unknown[]) => { locks.push(args); return { limit: async () => table === heartbeatRuns ? [run] + : phase ? [{ phase, failureCode }] : [] }; }, + }) }) }), update: () => ({ set: update }) }; + const db = { transaction: async (fn: (value: typeof tx) => Promise) => fn(tx) } as unknown as Db; + return { db, run, locks, update }; + } + it("admits a pending retry under the coordinator lock and preserves the exact caller on retry", async () => { + const f = fixture("retryable_failure"); + await claimCancellationRequest(f.db, runId, companyId, id, "board-user"); + const first = structuredClone(f.run.resultJson); + await claimCancellationRequest(f.db, runId, companyId, id, "board-user"); + expect(f.run.resultJson).toEqual(first); + expect(f.update).toHaveBeenCalledTimes(1); + expect(f.locks).toContainEqual(["update", { noWait: true }]); + }); + it.each([null, "observed", "executing", "terminal_failure", "committed"])("rejects missing or advanced retry phase %s without writing", async phase => { + const f = fixture(phase); + await expect(claimCancellationRequest(f.db, runId, companyId, id, "board-user")).rejects.toMatchObject({ status: 409 }); + expect(f.update).not.toHaveBeenCalled(); + }); + it.each(["pending", "acknowledged"])("recovers its own %s intent after native retry cancellation", async dispatchState => { + const f = fixture("terminal_failure", { startupCancellation: claimed, + nativeCancellation: { ...native, dispatchState, ...(dispatchState === "acknowledged" ? { acknowledgementAuditId: "ack-audit" } : {}) } }, "failed", "native_retry_cancelled"); + await expect(claimCancellationRequest(f.db, runId, companyId, id, "board-user")).resolves.toMatchObject({ status: "failed" }); + expect(f.update).not.toHaveBeenCalled(); + await expect(claimCancellationRequest(f.db, runId, companyId, id, "another-actor")).rejects.toMatchObject({ status: 409 }); + }); + it.each([{}, { ...native, intentId: "foreign" }, { ...native, issueId: "foreign" }, + { ...native, intentAuditId: null }, { ...native, dispatchState: "acknowledged" }])("rejects unproven or foreign cancellation recovery %j", nativeCancellation => { + const f = fixture("terminal_failure", { startupCancellation: claimed, nativeCancellation }, "failed", "native_retry_cancelled"); + return expect(claimCancellationRequest(f.db, runId, companyId, id, "board-user")).rejects.toMatchObject({ status: 409 }); + }); + it.each(["succeeded", "cancelled", "timed_out"])("does not admit terminal %s merely because a retry row exists", status => { + const f = fixture("retryable_failure", {}, status); + return expect(claimCancellationRequest(f.db, runId, companyId, id, "board-user")).rejects.toMatchObject({ status: 409 }); + }); + it.each([{ code: "55P03" }, { cause: { code: "55P03" } }])("fails closed on coordinator contention without database details", async error => { + const db = { transaction: async () => { throw error; } } as unknown as Db; + await expect(claimCancellationRequest(db, runId, companyId, id, "board-user")).rejects.toMatchObject({ status: 409, message: "Native retry cancellation is busy; retry the same request" }); + }); +}); diff --git a/server/src/services/native-runtime/native-cancellation-request.ts b/server/src/services/native-runtime/native-cancellation-request.ts index ecdf8d4ec59..54d11f190c5 100644 --- a/server/src/services/native-runtime/native-cancellation-request.ts +++ b/server/src/services/native-runtime/native-cancellation-request.ts @@ -1,5 +1,5 @@ import { and, eq, sql } from "drizzle-orm"; -import { heartbeatRuns, type Db } from "@paperclipai/db"; +import { heartbeatRuns, nativeRunFinalizations, type Db } from "@paperclipai/db"; import { badRequest, conflict } from "../../errors.js"; const UUID = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i; @@ -34,10 +34,30 @@ export async function claimCancellationRequest(db: Db, runId: string, companyId: if (!run || run.runtimeMode !== "native") throw conflict("Correlated Stop requires a native run"); assertCancellationRequest(run.resultJson, requestId, true); const result = record(run.resultJson); + let pendingRetry = false; + if (run.status === "failed" && run.nativeIssueId) { + // Eligibility and the caller fence must share a transaction. Coordinator + // writers need not lock the run. NOWAIT avoids a lock-order deadlock with + // execution's coordinator -> run claim; a busy claim can be retried. + const coordinator = await tx.select().from(nativeRunFinalizations) + .where(and(eq(nativeRunFinalizations.runId, runId), eq(nativeRunFinalizations.companyId, companyId), + eq(nativeRunFinalizations.issueId, run.nativeIssueId))) + .for("update", { noWait: true }).limit(1).then(rows => rows[0]); + const native = record(result.nativeCancellation); + const ownIntent = record(result.startupCancellation).cancellationRequestId === requestId + && native.schema === "paperclip.native-cancellation.v1" && native.intentId === cancellationIntentId(requestId) + && native.runId === runId && native.companyId === companyId && native.issueId === run.nativeIssueId + && native.scope === "run" && ["pending", "acknowledged"].includes(String(native.dispatchState)) + && typeof native.dispatched === "boolean" && typeof native.intentAuditId === "string" && !!native.intentAuditId + && (native.dispatchState === "pending" || (typeof native.acknowledgementAuditId === "string" + && !!native.acknowledgementAuditId && native.acknowledgementAuditId !== native.intentAuditId)); + pendingRetry = coordinator?.phase === "retryable_failure" + || (coordinator?.phase === "terminal_failure" && coordinator.failureCode === "native_retry_cancelled" && ownIntent); + } if (record(result.startupCancellation).cancellationRequestId === requestId) { const actor = record(record(result.startupCancellation).requestedBy); if (actor.type !== "board" || actor.userId !== userId) throw conflict("Cancellation request belongs to another actor"); - if (!["queued", "running", "scheduled_retry"].includes(run.status)) { + if (!pendingRetry && !["queued", "running", "scheduled_retry"].includes(run.status)) { const native = record(result.nativeCancellation); if (run.status !== "cancelled" || native.schema !== "paperclip.native-cancellation.v1" || native.intentId !== cancellationIntentId(requestId) || native.runId !== runId || native.companyId !== companyId @@ -48,12 +68,20 @@ export async function claimCancellationRequest(db: Db, runId: string, companyId: } return run; } - if (!["queued", "running", "scheduled_retry"].includes(run.status)) throw conflict("Correlated Stop run is already terminal"); + if (!pendingRetry && !["queued", "running", "scheduled_retry"].includes(run.status)) throw conflict("Correlated Stop run is already terminal"); const [claimed] = await tx.update(heartbeatRuns).set({ resultJson: { ...result, startupCancellation: { requestedAt: new Date().toISOString(), beforeNativeSelection: false, cancellationRequestId: requestId, requestedBy: { type: "board", userId } }, } }).where(and(eq(heartbeatRuns.id, runId), eq(heartbeatRuns.companyId, companyId))).returning(); if (!claimed) throw conflict("Correlated Stop run changed"); return claimed; + }).catch((error: unknown) => { + // postgres-js can be wrapped by Drizzle; retain no raw database detail in + // the public conflict while allowing the exact request to retry later. + const value = record(error); + if (value.code === "55P03" || record(value.cause).code === "55P03") { + throw conflict("Native retry cancellation is busy; retry the same request"); + } + throw error; }); }