From ec3a3ceb39075a0a022812ee89c2621937060cdb Mon Sep 17 00:00:00 2001 From: cynarAI <155274210+cynarAI@users.noreply.github.com> Date: Tue, 22 Sep 2026 12:00:14 +0200 Subject: [PATCH] fix(adapters): poll run lease during silent runtime work (#67) When a run worker loses its lease while the runtime yields no events (long tool or model work), lease loss was only detected on the next heartbeat renewal (up to 60s). A reclaimed worker could execute in parallel during that window. Add a dedicated lease poll that aborts stale workers within a few seconds and call runtime.abort so pi-runtime stops promptly. Co-authored-by: Cursor Agent Co-authored-by: cynarAI --- .../adapters/src/executor-lease-poll.test.ts | 173 ++++++++++++++++++ packages/adapters/src/executor.ts | 55 ++++-- 2 files changed, 208 insertions(+), 20 deletions(-) create mode 100644 packages/adapters/src/executor-lease-poll.test.ts diff --git a/packages/adapters/src/executor-lease-poll.test.ts b/packages/adapters/src/executor-lease-poll.test.ts new file mode 100644 index 0000000..ccefd2f --- /dev/null +++ b/packages/adapters/src/executor-lease-poll.test.ts @@ -0,0 +1,173 @@ +import type { AdapterContext, AgentRunRequest } from "@rakazo/adapter-kit"; +import { describe, expect, it, vi } from "vitest"; +import type * as AutoReviewModule from "./auto-review.js"; +import type * as ComputerLifecycleModule from "./computer-lifecycle.js"; +import { createRunExecutor } from "./executor.js"; + +vi.mock("./computer-lifecycle.js", async (importOriginal) => ({ + ...(await importOriginal()), + acquireComputerExecutionLease: async () => null, + provisionComputer: async () => ({ id: "computer-1", kind: "desktop" }), +})); + +vi.mock("./auto-review.js", async (importOriginal) => ({ + ...(await importOriginal()), + resolveAutoReviewChecker: () => ({ provider: "scripted", model: "checker" }), + isAutoReviewCheckerConfigured: () => true, + runAutoReviewJudge: vi.fn(), +})); + +/** Runtime that blocks without yielding until the executor abort signal fires. */ +function silentBlockingRuntime() { + let blocked = false; + const runtimeRun = vi.fn(async function* ( + _request: AgentRunRequest, + context?: Partial, + ) { + blocked = true; + await new Promise((_resolve, reject) => { + const signal = context?.signal; + if (!signal) return; + if (signal.aborted) { + reject(Object.assign(new Error("Aborted"), { name: "AbortError" })); + return; + } + const onAbort = () => { + signal.removeEventListener("abort", onAbort); + reject(Object.assign(new Error("Aborted"), { name: "AbortError" })); + }; + signal.addEventListener("abort", onAbort); + }); + yield { type: "done" as const, text: "stopped" }; + }); + const runtimeAbort = vi.fn(async () => undefined); + return { + blocked: () => blocked, + runtime: { + describe: () => ({ capabilities: { scripted: false } }), + run: runtimeRun, + abort: runtimeAbort, + }, + runtimeRun, + runtimeAbort, + }; +} + +function executorFixture() { + const run = { + id: "run-1", + botId: "bot-1", + threadId: "thread-1", + taskId: "task-1", + spaceId: "space-1", + userId: "user-1", + status: "queued", + trigger: "user", + leaseFence: 0, + leaseOwner: null as string | null, + startedAt: null as Date | null, + checkpoint: null, + }; + const prisma = { + run: { + findUnique: vi.fn(async () => run), + findUniqueOrThrow: vi.fn(async () => run), + updateMany: vi.fn(async ({ data }: { data: Record }) => { + Object.assign(run, data); + return { count: 1 }; + }), + }, + bot: { + findUnique: vi.fn(async () => ({ computerId: "computer-1", computerSwitching: false })), + findUniqueOrThrow: vi.fn(async () => ({ + id: run.botId, + name: "Assistant", + title: "Assistant", + description: "Test assistant", + computerId: "computer-1", + computer: { id: "computer-1", scope: "dedicated" }, + instructions: null, + modelProvider: null, + modelId: null, + memoryScope: null, + })), + findMany: vi.fn(async () => []), + }, + attempt: { + create: vi.fn(async () => ({ id: "attempt-1" })), + update: vi.fn(), + updateMany: vi.fn(async () => ({ count: 1 })), + }, + thread: { findUniqueOrThrow: vi.fn(async () => ({ id: run.threadId, groupId: null })) }, + message: { findMany: vi.fn(async () => []), findUnique: vi.fn(async () => null) }, + task: { + findUniqueOrThrow: vi.fn(async () => ({ id: run.taskId, prompt: "silent tool work" })), + }, + connection: { findMany: vi.fn(async () => []) }, + spaceModelPreference: { findFirst: vi.fn(async () => null) }, + userModelCredential: { findFirst: vi.fn(async () => null) }, + deploymentSettings: { + findUnique: vi.fn(async () => ({ + defaultModelProvider: "scripted", + defaultModelId: "scripted", + })), + }, + taughtSkill: { findMany: vi.fn(async () => []) }, + agentSecret: { findMany: vi.fn(async () => []) }, + agentSkill: { findMany: vi.fn(async () => []) }, + scratchpadItem: { findMany: vi.fn(async () => []) }, + actionApprovalRule: { findMany: vi.fn(async () => []) }, + actionAutoReviewPreference: { findUnique: vi.fn(async () => ({ enabled: false })) }, + externalEffect: { + findMany: vi.fn(async () => []), + findUnique: vi.fn(async () => null), + }, + }; + const { runtime, runtimeRun, runtimeAbort, blocked } = silentBlockingRuntime(); + const executor = createRunExecutor({ + prisma, + runtime, + connector: { + discoverTools: async () => [], + resolveCall: async () => undefined, + execute: async function* () {}, + }, + sandbox: { describe: () => ({ capabilities: { graphical: false } }) }, + memory: { + describe: () => ({ capabilities: {} }), + read: async () => ({ documents: [] }), + search: async () => [], + commit: vi.fn(async () => ({ revision: "rev-1" })), + exportMarkdown: async function* () {}, + }, + memoryProviders: { resolve: async () => null }, + events: { + append: vi.fn(async () => undefined), + pauseRunForInput: vi.fn(async () => true), + finalizeRun: vi.fn(async () => ({ continuationRunId: null })), + }, + jobs: { enqueue: vi.fn(async () => undefined) }, + secrets: [], + } as unknown as Parameters[0]); + + return { run, prisma, executor, runtimeRun, runtimeAbort, blocked }; +} + +describe("executor lease poll", () => { + it("stops a stale worker during silent runtime work after lease reclaim", async () => { + const { run, executor, runtimeRun, runtimeAbort, blocked } = executorFixture(); + const workerDone = executor.continueRun(run.id, "worker-a"); + + await vi.waitFor(() => expect(runtimeRun).toHaveBeenCalled()); + expect(blocked()).toBe(true); + + Object.assign(run, { leaseOwner: "worker-b", leaseFence: 2, status: "running" }); + + await vi.waitFor(() => expect(runtimeAbort).toHaveBeenCalledWith("run-1"), { + timeout: 5_000, + }); + await workerDone; + + expect(run.leaseOwner).toBe("worker-b"); + }); +}); diff --git a/packages/adapters/src/executor.ts b/packages/adapters/src/executor.ts index 4164e0e..258c8bd 100644 --- a/packages/adapters/src/executor.ts +++ b/packages/adapters/src/executor.ts @@ -950,6 +950,39 @@ export function createRunExecutor(deps: ExecutorDeps) { let screenRelease: { computer: ComputerRef; context: AdapterContext } | undefined; let runAbortController: AbortController | null = null; let detachShutdown: (() => void) | undefined; + const invalidateLostLease = () => { + if (!leaseValid) return; + leaseValid = false; + runAbortController?.abort(); + void deps.runtime.abort(runId).catch(() => undefined); + }; + const verifyLeaseOwnership = async (): Promise => { + if (!leaseValid) return false; + const still = await deps.prisma.run.findUnique({ + where: { id: runId }, + select: { status: true, leaseOwner: true, leaseFence: true, checkpoint: true }, + }); + if ( + !still || + still.status === "cancelled" || + still.leaseOwner !== workerId || + still.leaseFence !== fence + ) { + invalidateLostLease(); + return false; + } + const releasedHold = takeoverCheckpointOf(still.checkpoint); + if (heldForTakeover && releasedHold) { + await requeueComputerRun(deps, runId, workerId, fence, releasedHold, false); + invalidateLostLease(); + return false; + } + return true; + }; + const leasePoll = setInterval(() => { + void verifyLeaseOwnership().catch(() => invalidateLostLease()); + }, 2_000); + leasePoll.unref?.(); const heartbeat = setInterval(() => { void Promise.all([ renewRunLease(deps, runId, workerId, fence), @@ -3596,26 +3629,7 @@ export function createRunExecutor(deps: ExecutorDeps) { const now = Date.now(); if (now - lastLeaseCheckAt >= 1_000) { lastLeaseCheckAt = now; - const still = await deps.prisma.run.findUnique({ - where: { id: runId }, - select: { status: true, leaseOwner: true, leaseFence: true, checkpoint: true }, - }); - if ( - !still || - still.status === "cancelled" || - still.leaseOwner !== workerId || - still.leaseFence !== fence - ) { - leaseValid = false; - return; - } - const releasedHold = takeoverCheckpointOf(still.checkpoint); - if (heldForTakeover && releasedHold) { - await requeueComputerRun(deps, runId, workerId, fence, releasedHold, false); - leaseValid = false; - runAbortController?.abort(); - return; - } + if (!(await verifyLeaseOwnership())) return; } if (event.type === "text") { @@ -4187,6 +4201,7 @@ export function createRunExecutor(deps: ExecutorDeps) { } } finally { detachShutdown?.(); + clearInterval(leasePoll); clearInterval(heartbeat); if (!retainComputerLease) { if (screenRelease) {