From 9296c74e8dfaf2aa386ffe2ca85b203a2f1619ed Mon Sep 17 00:00:00 2001 From: Dotta <34892728+cryppadotta@users.noreply.github.com> Date: Tue, 10 Nov 2026 17:43:46 -0600 Subject: [PATCH] fix(native): let a restart claim take back its lapsed lease at budget admission (#20650) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Thinking Path > - Paperclip is the open source app people use to manage AI agents for work > - After a control-plane restart, restart recovery claims an interrupted native run (`resume_dead_runner`) and grants a 20-minute lease. Budget admission then re-checks that the claim still owns the run before the resumed run may start provider work. > - A host sleep, or the restart-chaos fault TM2, can make that stored lease look expired while its owner is alive. #20625 lets the owner take back such a lease when the row still shows exactly its own claim. That rule did not cover budget admission. > - So a claim whose lease lapsed before dispatch was refused with `Native recovery no longer owns provider dispatch`. The run failed with `adapter_failed`, recovery held it with `native_continuation_requires_reconciliation`, and the user had to click Retry. > - This pull request applies the #20625 rule in budget admission: the claim takes its lease back through the existing re-acquire condition. > - The benefit is that a run resumed after a restart survives a time jump before dispatch with no Retry. Every fence and every budget check stays the same. ## Linked Issues or Issue Description Refs #20625, #20474 **What happened?** Restart-chaos wave L7 (seed 9107, local) on a branch with #20625, task codex B5 TM2. A TM3 restart claimed the run (`resume_dead_runner`, generation 2). TM2 then aged every database time by one hour before the resume reached dispatch. Budget admission (`lockRunBudgetRecoveryOwnerInTransaction`) required a live lease and threw `Native recovery no longer owns provider dispatch`. The run failed with `adapter_failed`, recovery stopped with `native_continuation_requires_reconciliation` / `preserve_without_replay_v1`, and the task needed Retry. **Expected behavior** The claim's own lapsed lease is taken back, and the resume is admitted when the budget allows it. **Steps to reproduce** `node cli/node_modules/tsx/dist/cli.mjs tests/restart-chaos/main.ts --env local --adapters codex,claude --workloads A1,B5,C1 --faults TM1,TM2,TM3 --time-age-hours 1 --concurrency 3 --tasks 6 --seed 9107` (timing-dependent). ## What Changed - `native-execution-lease.ts`: two new helpers. Both use the #20625 own-claim rule. Neither takes a lock out of order. - `readOwnNativeRecoveryClaimLease` (read-only) returns `live`, `lapsed_own_claim` or `lost`. It compares the lease against the database clock. A lapsed lease counts as `lapsed_own_claim` only when all of these hold: - the same lease owner; - this process's controller boot, pid and process start; - a non-null `recovery_state`; - `nativeExecutionLeaseHeldCondition({ remote: false, controllerGeneration })`, which applies the #20625 own-claim rule and every #20474 untouched check; - no non-local environment lease on the run. - `extendOwnNativeRecoveryClaimInTransaction` extends such a lease by the local execution TTL, with one conditional UPDATE on the same predicate. - `budget-reservations.ts`: - `lockRunBudgetRecoveryOwnerInTransaction` still refuses another owner, and refuses a lapsed lease unless it is the run's own untouched claim. It no longer locks the run row, so the documented coordinator → issue → run order holds. - The lease is extended only at the end of an admitted reservation (`reserveRunBudget`) or an admitted dispatch recheck (`recheckRunBudgetReservationInTransaction`). At those points the caller already holds the run row lock. A refused admission, for example one over budget, leaves the lease lapsed. - `heartbeat-native-restart-suspension.test.ts`: the test "still settles a recovery whose claim merely expired" encoded "expired means not the owner". #20474 and #20625 changed that rule for an untouched own claim. The test now has two cases: - an expired claim that another actor recorded on (`awaiting_evidence`) still settles with no provider session; - an untouched expired claim is dispatched. - `doc/plans/2026-11-09-explicit-native-lease-claim-tokens.md`: Step 1 lists the new user of the rule. ## Verification - New `server/src/__tests__/native-recovery-dispatch-lease-reacquire.test.ts`: - `reserveRunBudget` takes back the claim's own lapsed lease and admits the resume. This fails before the change. - An over-budget resume is still refused, and its lease stays lapsed. - Fenced variants each throw `Native recovery no longer owns provider dispatch` and leave the row unchanged: a graceful-shutdown suspension, an appended failure settlement, `awaiting_evidence`, `blocked`, a claim from an earlier generation, another boot, a remote sandbox lease, and another lease owner. - At the dispatch gate (`recheckRunBudgetReservationInTransaction`), an admitted own claim gets its lease extended, and a fenced claim is refused with the row unchanged. - Related suites pass (308 tests): `native-execution-lease-reacquire.test.ts`, `closed-native-receipt-accounting.test.ts`, `heartbeat-native-predispatch-retry.test.ts`, `run-dispatch/adapters/postgres.test.ts`, `heartbeat-native-restart-suspension.test.ts` and `heartbeat-restart-claim-lapse.test.ts`. - `cd server && npx tsc --noEmit` passes. - Live: one L7 wave (seed 9107) ran on the first commit of this branch (`w-20261110T212307-s9107`). It logged 0 `Native recovery no longer owns provider dispatch` refusals, against 1 in the previous wave. Codex B5 TM2 got past dispatch, but it still needed Retry for another reason: after a `runner.drain` timeout, the scheduled same-run resume did not send the continuation of #20645. That case is tracked separately. The other five tasks passed. ## Risks - This widens when a restart claim can take back its own lapsed lease, and only at budget admission. It is local only. It needs the same owner, boot and generation, and the row must show exactly that claim. A claim that any other actor touched still fails as before. - Budget hard-stop behaviour is unchanged. The lease is extended only after budget admission passes. - No schema change and no migration. Rows saved by earlier releases carry the same history shape. ## Model Used - Claude Opus 5.5 (`claude-opus-5-5`), Anthropic, via Claude Code with tool use (shell, file edits, test runs). Extended reasoning. ## Checklist - [x] I have included a thinking path that traces from project context to this change - [x] I have specified the model used (with version and capability details) - [x] I have checked ROADMAP.md and confirmed this PR does not duplicate planned core work - [x] I have searched GitHub for duplicate or related PRs and linked them above - [x] I have either (a) linked existing issues with `Fixes: #` / `Closes #` / `Refs #` OR (b) described the issue in-PR following the relevant issue template - [x] I have not referenced internal/instance-local Paperclip issues or links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip` URLs) - [x] My branch name describes the change (e.g. `docs/...`, `fix/...`) and contains no internal Paperclip ticket id or instance-derived details - [x] I have run tests locally and they pass - [x] I have added or updated tests where applicable - [x] I have updated relevant documentation to reflect my changes - [x] I have considered and documented any risks above - [ ] All Paperclip CI gates are green - [ ] Greptile is 5/5 with no open P2s, recommendations, or follow-ups - [x] I will address all Greptile and reviewer comments before requesting merge 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude --- ...1-09-explicit-native-lease-claim-tokens.md | 6 +- ...eartbeat-native-restart-suspension.test.ts | 21 ++- ...-recovery-dispatch-lease-reacquire.test.ts | 130 ++++++++++++++++++ server/src/services/budget-reservations.ts | 43 +++--- .../native-runtime/native-execution-lease.ts | 60 +++++++- 5 files changed, 238 insertions(+), 22 deletions(-) create mode 100644 server/src/__tests__/native-recovery-dispatch-lease-reacquire.test.ts 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;