diff --git a/doc/plans/2026-11-09-explicit-native-lease-claim-tokens.md b/doc/plans/2026-11-09-explicit-native-lease-claim-tokens.md index 7edd7b5cf08..e37bcbd02d6 100644 --- a/doc/plans/2026-11-09-explicit-native-lease-claim-tokens.md +++ b/doc/plans/2026-11-09-explicit-native-lease-claim-tokens.md @@ -38,7 +38,11 @@ succeeds only when all of these hold: After a result is accepted, the owner may only relinquish a lapsed lease, and never renews it. The rule lives in `nativeExecutionLeaseHeldCondition` and `lockNativeExecutionLeaseForReacquire` (`native-execution-lease.ts`). Renewal, -failure settlement and the final release use it. +failure settlement and the final release use it. Budget admission of a restart +claim (`lockRunBudgetRecoveryOwnerInTransaction`) uses it too, through +`holdOwnNativeRecoveryClaimInTransaction` (2026-11-10). A restart claim whose lease +lapsed in a time jump before dispatch therefore takes the lease back when the row +still shows exactly that claim. The budget checks that follow are unchanged. This works because every consumer that acts on a lapsed executor lease of a running, pre-result run already writes one of those fields under the coordinator diff --git a/server/src/__tests__/heartbeat-native-restart-suspension.test.ts b/server/src/__tests__/heartbeat-native-restart-suspension.test.ts index 9a5b1cf6209..0ce76469c5f 100644 --- a/server/src/__tests__/heartbeat-native-restart-suspension.test.ts +++ b/server/src/__tests__/heartbeat-native-restart-suspension.test.ts @@ -348,13 +348,13 @@ describeEmbeddedPostgres("heartbeat native recovery interrupted by graceful shut expect(resumed).toHaveBeenCalledOnce(); }, 90_000); - it("still settles a recovery whose claim merely expired without a shutdown suspension", async () => { + it("still settles a recovery whose expired claim another actor recorded on, without a shutdown suspension", async () => { const f = await seedRecoverableRun("EXPIRED-NOT-SUSPENDED"); const dying = vi.fn(() => { throw new Error("An expired claim must not open a provider session"); }); const heartbeat = heartbeatService(f.db, { nativeSessionBackendFactory: dying, beforeNativeRuntimeSelection: async () => { - await f.db.update(nativeRunFinalizations).set({ leaseExpiresAt: new Date(0) }) + await f.db.update(nativeRunFinalizations).set({ leaseExpiresAt: new Date(0), recoveryState: "awaiting_evidence" }) .where(eq(nativeRunFinalizations.runId, f.runId)); }, }); @@ -367,6 +367,23 @@ describeEmbeddedPostgres("heartbeat native recovery interrupted by graceful shut expect(run.status).toBe("failed"); }, 60_000); + it("dispatches its own restart claim whose lease lapsed untouched before dispatch (host sleep)", async () => { + const f = await seedRecoverableRun("EXPIRED-UNTOUCHED"); + const reached = vi.fn(() => { throw new Error("provider reached after the untouched claim was taken back"); }); + const heartbeat = heartbeatService(f.db, { + nativeSessionBackendFactory: reached, + beforeNativeRuntimeSelection: async () => { + await f.db.update(nativeRunFinalizations).set({ leaseExpiresAt: new Date(0) }) + .where(eq(nativeRunFinalizations.runId, f.runId)); + }, + }); + try { + const { claims } = await heartbeat.recoverNativeRunsAfterRestart(new Date(), [f.runId]); + expect(claims).toEqual([expect.objectContaining({ runId: f.runId })]); + } finally { await heartbeat.drainActiveRunExecutions(); } + expect(reached).toHaveBeenCalled(); + }, 60_000); + async function shutdownDispositions(db: Awaited>["db"], runId: string) { const events = await db.select().from(heartbeatRunEvents).where(and(eq(heartbeatRunEvents.runId, runId), diff --git a/server/src/__tests__/native-recovery-dispatch-lease-reacquire.test.ts b/server/src/__tests__/native-recovery-dispatch-lease-reacquire.test.ts new file mode 100644 index 00000000000..ace8c7cab0c --- /dev/null +++ b/server/src/__tests__/native-recovery-dispatch-lease-reacquire.test.ts @@ -0,0 +1,130 @@ +import { randomUUID } from "node:crypto"; +import { eq, sql } from "drizzle-orm"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { + agents, budgetPolicies, budgetReservations, companies, costEvents, createDb, environmentLeases, heartbeatRuns, issues, + nativeRunFinalizations, type Db, +} from "@paperclipai/db"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js"; +import { recheckRunBudgetReservationInTransaction, reserveRunBudget } from "../services/budget-reservations.js"; +import { currentNativeControllerIdentity, type NativeControllerIdentity } from "../services/native-runtime/native-restart-recovery.js"; + +const GENERATION = 2; + +const support = await getEmbeddedPostgresTestSupport(); +(support.supported ? describe : describe.skip)("budget admission of a restart claim whose lease lapsed in a time jump", () => { + let database: Awaited>; + let db: Db; + let controller: NativeControllerIdentity; + + beforeAll(async () => { + database = await startEmbeddedPostgresTestDatabase("native-recovery-dispatch-lease-"); + db = createDb(database.connectionString); + controller = await currentNativeControllerIdentity(); + }, 120_000); + afterAll(async () => { await database?.cleanup(); }); + + function claimEntry(overrides: Record = {}) { + return { at: new Date().toISOString(), restartKind: "hard", controllerBootId: controller.bootId, + controllerPid: controller.pid, controllerProcessStartedAt: controller.processStartedAt.toISOString(), + controllerGeneration: GENERATION, disposition: "resume_dead_runner", reason: "dead_runner_with_exact_provider_checkpoint", + stateRootAction: "reuse_exact_root", recoveryRequestId: null, providerAttempt: 1, ...overrides }; + } + + async function seed(options: { + recoveryState?: string; history?: Array>; bootId?: string; remote?: boolean; + overBudget?: boolean; + } = {}) { + const companyId = randomUUID(), agentId = randomUUID(); + await db.insert(companies).values({ id: companyId, name: "Time jump", issuePrefix: `T${companyId.slice(0, 4)}`.toUpperCase() }); + await db.insert(agents).values({ id: agentId, companyId, name: "Codex", role: "engineer", status: "idle", + adapterType: "paperclip_runner" }); + const [issue] = await db.insert(issues).values({ companyId, title: "Resumed task", status: "in_progress", + assigneeAgentId: agentId }).returning(); + const [run] = await db.insert(heartbeatRuns).values({ companyId, agentId, invocationSource: "assignment", + issueId: issue!.id, nativeIssueId: issue!.id, runtimeMode: "native", status: "running", + contextSnapshot: { issueId: issue!.id } }).returning(); + const leaseOwner = `${controller.bootId}:${GENERATION}:${randomUUID()}`; + await db.insert(nativeRunFinalizations).values({ companyId, issueId: issue!.id, runId: run!.id, phase: "observed", + attempt: 1, leaseOwner, leaseExpiresAt: sql`clock_timestamp() - interval '40 minutes'`, + controllerBootId: options.bootId ?? controller.bootId, controllerPid: controller.pid, + controllerProcessStartedAt: controller.processStartedAt, controllerGeneration: GENERATION, + recoveryState: options.recoveryState ?? "resuming_session", recoveryHistory: options.history ?? [claimEntry()] }); + if (options.remote) await db.insert(environmentLeases).values({ companyId, issueId: issue!.id, + heartbeatRunId: run!.id, provider: "daytona" }); + if (options.overBudget) { + await db.insert(budgetPolicies).values({ companyId, scopeType: "company", scopeId: companyId, metric: "billed_cents", + windowKind: "calendar_month_utc", amount: 100, hardStopEnabled: true, notifyEnabled: false }); + await db.insert(costEvents).values({ companyId, agentId, provider: "openai", model: "gpt", costCents: 250, + occurredAt: new Date(), costStatus: "reported" }); + } + const coordinator = async () => (await db.select().from(nativeRunFinalizations) + .where(eq(nativeRunFinalizations.runId, run!.id)))[0]!; + return { companyId, runId: run!.id, leaseOwner, coordinator }; + } + + it("takes back the lease of its own restart claim and admits the resume", async () => { + const f = await seed(); + const reserved = await reserveRunBudget(db, f.companyId, f.runId, null, {}, f.leaseOwner); + expect(reserved).not.toHaveProperty("admissionBlockReason"); + const saved = await f.coordinator(); + expect(saved.leaseOwner).toBe(f.leaseOwner); + expect(saved.recoveryState).toBe("resuming_session"); + expect(saved.leaseExpiresAt!.getTime()).toBeGreaterThan(Date.now() + 15 * 60_000); + }); + + it("still refuses a resume that is over its hard-stop budget", async () => { + const f = await seed({ overBudget: true }); + const refusal = await reserveRunBudget(db, f.companyId, f.runId, null, {}, f.leaseOwner) + .then(result => "admissionBlockReason" in result ? String(result.admissionBlockReason) : null, + (error: Error) => error.message); + expect(refusal).toMatch(/budget/i); + expect((await f.coordinator()).leaseExpiresAt!.getTime()).toBeLessThan(Date.now()); + }); + + it("takes the lease back at the dispatch gate only when the resume is admitted", async () => { + const admitted = await seed(); + const [reservation] = await db.insert(budgetReservations).values({ companyId: admitted.companyId, runId: admitted.runId, + agentId: (await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, admitted.runId)))[0]!.agentId, + amountCents: "0", state: "held", providerStartedAt: new Date() }).returning(); + const gate = async (f: Awaited>, reservationId: string) => db.transaction(async tx => { + const [run] = await tx.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)).for("update"); + return recheckRunBudgetReservationInTransaction(tx as unknown as Db, [], run!, reservationId, f.leaseOwner); + }); + expect(await gate(admitted, reservation!.id)).toBeNull(); + expect((await admitted.coordinator()).leaseExpiresAt!.getTime()).toBeGreaterThan(Date.now() + 15 * 60_000); + + const fenced = await seed({ recoveryState: "awaiting_evidence" }); + const [fencedReservation] = await db.insert(budgetReservations).values({ companyId: fenced.companyId, runId: fenced.runId, + agentId: (await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, fenced.runId)))[0]!.agentId, + amountCents: "0", state: "held", providerStartedAt: new Date() }).returning(); + const before = await fenced.coordinator(); + expect(await gate(fenced, fencedReservation!.id)).toBe("Native recovery no longer owns provider dispatch"); + expect(await fenced.coordinator()).toEqual(before); + }); + + it.each<[string, () => Parameters[0]]>([ + ["a graceful shutdown suspended the claim", () => ({ recoveryState: "awaiting_runner_reattach", history: [ + claimEntry({ disposition: "reattach_existing_runner" }), + { at: new Date().toISOString(), restartKind: "graceful", disposition: "restart_suspended", reason: "SIGTERM" }] })], + ["failure settlement appended its outcome", () => ({ history: [claimEntry(), + { at: new Date().toISOString(), disposition: "retryable_failure", reason: "native_session_lease_lost" }] })], + ["recovery now waits for evidence", () => ({ recoveryState: "awaiting_evidence" })], + ["the run is blocked", () => ({ recoveryState: "blocked" })], + ["the claim belongs to an earlier generation", () => ({ history: [claimEntry({ controllerGeneration: GENERATION - 1 })] })], + ["another boot owns the row", () => ({ bootId: `boot-${randomUUID()}` })], + ["the run executes on a remote sandbox", () => ({ remote: true })], + ])("stays fenced out when %s", async (_name, options) => { + const f = await seed(options()); + const before = await f.coordinator(); + await expect(reserveRunBudget(db, f.companyId, f.runId, null, {}, f.leaseOwner)) + .rejects.toThrow("Native recovery no longer owns provider dispatch"); + expect(await f.coordinator()).toEqual(before); + }); + + it("stays fenced out when another owner holds the claim", async () => { + const f = await seed(); + await expect(reserveRunBudget(db, f.companyId, f.runId, null, {}, `${controller.bootId}:${GENERATION}:${randomUUID()}`)) + .rejects.toThrow("Native recovery no longer owns provider dispatch"); + }); +}); diff --git a/server/src/services/budget-reservations.ts b/server/src/services/budget-reservations.ts index c417a4d0cd4..49c47eb14a1 100644 --- a/server/src/services/budget-reservations.ts +++ b/server/src/services/budget-reservations.ts @@ -9,6 +9,7 @@ import { budgetServiceInTransaction } from "./budgets.js"; import { withBudgetSpan } from "./budget-tracing.js"; import type { ActivityPublication } from "./activity-log.js"; import { readClaimedNativeStopRecovery } from "./native-runtime/native-cancellation-request.js"; +import { extendOwnNativeRecoveryClaimInTransaction, readOwnNativeRecoveryClaimLease } from "./native-runtime/native-execution-lease.js"; /** Native finalization takes its coordinator before issue/run locks. Recovery * admission uses the same order after the company accounting lock, so its @@ -17,9 +18,10 @@ export async function lockRunBudgetRecoveryOwnerInTransaction( tx: Db, companyId: string, runId: string, recoveryLeaseOwner?: string, ) { if (!recoveryLeaseOwner) return; - const [coordinator] = await tx.select({ leaseOwner: nativeRunFinalizations.leaseOwner, leaseExpiresAt: nativeRunFinalizations.leaseExpiresAt }).from(nativeRunFinalizations) + const [coordinator] = await tx.select({ leaseOwner: nativeRunFinalizations.leaseOwner }).from(nativeRunFinalizations) .where(and(eq(nativeRunFinalizations.companyId, companyId), eq(nativeRunFinalizations.runId, runId))).for("update"); - if (!coordinator || coordinator.leaseOwner !== recoveryLeaseOwner || !coordinator.leaseExpiresAt || coordinator.leaseExpiresAt.getTime() <= Date.now()) + if (!coordinator || coordinator.leaseOwner !== recoveryLeaseOwner + || await readOwnNativeRecoveryClaimLease(tx, { companyId, runId, leaseOwner: recoveryLeaseOwner }) === "lost") throw conflict("Native recovery no longer owns provider dispatch"); } @@ -45,14 +47,13 @@ export async function recheckRunBudgetReservationInTransaction( eq(budgetReservations.state, "held"), )).for("update"); if (!reservation) return "Provider dispatch no longer owns a held budget reservation"; - if (recoveryLeaseOwner) { - const [owner] = run.runtimeMode === "native" ? await tx.select({ runId: nativeRunFinalizations.runId }) - .from(nativeRunFinalizations).where(and(eq(nativeRunFinalizations.companyId, run.companyId), - eq(nativeRunFinalizations.runId, run.id), eq(nativeRunFinalizations.leaseOwner, recoveryLeaseOwner), - sql`${nativeRunFinalizations.leaseExpiresAt} > clock_timestamp()`)) : []; - if (!owner) return "Native recovery no longer owns provider dispatch"; - } - if (stopRecovery) return null; + if (recoveryLeaseOwner && (run.runtimeMode !== "native" || await readOwnNativeRecoveryClaimLease(tx, + { companyId: run.companyId, runId: run.id, leaseOwner: recoveryLeaseOwner }) === "lost")) + return "Native recovery no longer owns provider dispatch"; + const keepRecoveryClaim = async () => !recoveryLeaseOwner || await extendOwnNativeRecoveryClaimInTransaction(tx, + { companyId: run.companyId, runId: run.id, leaseOwner: recoveryLeaseOwner }) + ? null : "Native recovery no longer owns provider dispatch"; + if (stopRecovery) return keepRecoveryClaim(); const { block, policies, observations } = await budgetServiceInTransaction(tx, publications) .evaluateAdmission(run.companyId, run.agentId, { projectId: reservation.projectId, runId: run.id, admittedRecovery: Boolean(recoveryLeaseOwner) }); @@ -63,7 +64,7 @@ export async function recheckRunBudgetReservationInTransaction( return "Agent is paused or its lifecycle is not ready for new work"; // An owned recovery reuses the hold reserveRunBudget already admitted; // other runs' holds filling the budget must not refuse it. - if (recoveryLeaseOwner) return null; + if (recoveryLeaseOwner) return keepRecoveryClaim(); for (const [index, policy] of policies.entries()) { const observed = observations[index]; // This operation's hold is already included in exposure. Do not propose it @@ -104,13 +105,18 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string, if (existing) { // The native coordinator fences recovery ownership. Reuse the original // hold only for that live owner, never for an ordinary duplicate dispatch. - const [owner] = recoveryLeaseOwner && run.runtimeMode === "native" && existing.state === "held" - && existing.projectId === projectId ? await tx.select({ runId: nativeRunFinalizations.runId }) - .from(nativeRunFinalizations).where(and(eq(nativeRunFinalizations.companyId, companyId), - eq(nativeRunFinalizations.runId, runId), eq(nativeRunFinalizations.leaseOwner, recoveryLeaseOwner), - sql`${nativeRunFinalizations.leaseExpiresAt} > clock_timestamp()`)) : []; + const owner = recoveryLeaseOwner && run.runtimeMode === "native" && existing.state === "held" + && existing.projectId === projectId + && await readOwnNativeRecoveryClaimLease(tx, { companyId, runId, leaseOwner: recoveryLeaseOwner }) !== "lost"; + const keepRecoveryClaim = async () => { + if (!await extendOwnNativeRecoveryClaimInTransaction(tx, { companyId, runId, leaseOwner: recoveryLeaseOwner! })) + throw conflict("Native recovery no longer owns provider dispatch"); + }; if (owner) { - if (stopRecovery) return { ...existing, reused: true }; + if (stopRecovery) { + await keepRecoveryClaim(); + return { ...existing, reused: true }; + } // Recovery has already made this run queued/running, so its own pending // accounting is excluded. Current overruns, other unfinished runs and // scope pauses still apply before any more paid work starts. @@ -120,6 +126,7 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string, .where(and(eq(agents.companyId, companyId), eq(agents.id, run.agentId))); if (!agent || agent.lifecycleState !== "ready" || ["paused", "terminated", "pending_approval"].includes(agent.status)) return { admissionBlockReason: "Agent is paused and cannot start new work" }; + await keepRecoveryClaim(); return { ...existing, reused: true }; } throw conflict("Provider dispatch has already reserved this run"); @@ -149,6 +156,8 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string, usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({ accountingReceiptReady: false, ledgerScope: { ...(run.usageJson?.ledgerScope as object ?? {}), ...ledgerScope, projectId } })}::jsonb`, }).where(eq(heartbeatRuns.id, runId)); await queueBudgetEvaluation(tx, { companyId, agentId: run.agentId, projectId }); + if (recoveryLeaseOwner && !await extendOwnNativeRecoveryClaimInTransaction(tx, { companyId, runId, leaseOwner: recoveryLeaseOwner })) + throw conflict("Native recovery no longer owns provider dispatch"); return { ...reservation, reused: false }; }, { isolationLevel: "read committed" })); // Commit required block/resume transitions even when dispatch is denied. diff --git a/server/src/services/native-runtime/native-execution-lease.ts b/server/src/services/native-runtime/native-execution-lease.ts index 49571e3b091..b63b7e1398f 100644 --- a/server/src/services/native-runtime/native-execution-lease.ts +++ b/server/src/services/native-runtime/native-execution-lease.ts @@ -1,7 +1,7 @@ import { notifyLiveRunRecovery } from "../live-run-recovery-notifications.js"; import type { Db } from "@paperclipai/db"; -import { heartbeatRuns, livenessDb, nativeRunFinalizations } from "@paperclipai/db"; -import { and, eq, gt, isNull, or, sql, type SQL } from "drizzle-orm"; +import { environmentLeases, heartbeatRuns, livenessDb, nativeRunFinalizations } from "@paperclipai/db"; +import { and, eq, gt, isNotNull, isNull, notExists, or, sql, type SQL } from "drizzle-orm"; import { currentNativeControllerIdentity, type NativeControllerIdentity } from "./native-restart-recovery.js"; import { OPERATOR_ATTESTED_CONTROLLER_HOST } from "./native-controller-host-recovery.js"; import { NATIVE_OWNERSHIP_UNVERIFIED_ERROR_CODE } from "./native-runner-ownership.js"; @@ -126,6 +126,62 @@ export async function lockNativeExecutionLeaseForReacquire(tx: Pick 'local'`))), + )!; +} + +export async function readOwnNativeRecoveryClaimLease( + tx: Db, input: OwnNativeRecoveryClaimInput, +): Promise<"live" | "lapsed_own_claim" | "lost"> { + const [row] = await tx.select({ + live: sql`coalesce(${nativeRunFinalizations.leaseExpiresAt} > clock_timestamp(), false)`, + }).from(nativeRunFinalizations).where(and(eq(nativeRunFinalizations.companyId, input.companyId), + eq(nativeRunFinalizations.runId, input.runId), eq(nativeRunFinalizations.leaseOwner, input.leaseOwner))); + if (!row) return "lost"; + if (row.live) return "live"; + const condition = await ownNativeRecoveryClaimCondition(tx, input); + if (!condition) return "lost"; + const [held] = await tx.select({ runId: nativeRunFinalizations.runId }).from(nativeRunFinalizations).where(condition); + return held ? "lapsed_own_claim" : "lost"; +} + +export async function extendOwnNativeRecoveryClaimInTransaction(tx: Db, input: OwnNativeRecoveryClaimInput): Promise { + const state = await readOwnNativeRecoveryClaimLease(tx, input); + if (state === "live") return true; + if (state === "lost") return false; + const condition = await ownNativeRecoveryClaimCondition(tx, input); + if (!condition) return false; + const [held] = await tx.update(nativeRunFinalizations) + .set({ + leaseExpiresAt: sql`clock_timestamp() + (${NATIVE_SESSION_EXECUTION_LEASE_TTL_MS} * interval '1 millisecond')`, + updatedAt: sql`clock_timestamp()`, + }) + .where(condition) + .returning({ runId: nativeRunFinalizations.runId }); + return Boolean(held); +} + /** Preparation can perform provider I/O only under its current recovery claim. */ export async function assertNativeSessionPreparationAuthority(input: { db: Db; companyId: string; runId: string; leaseOwner: string | undefined;