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 <cursoragent@cursor.com> Co-authored-by: cynarAI <cynarAI@users.noreply.github.com>
This commit is contained in:
cynarAI
GitHub
parent
b4a0dbc74d
commit
ec3a3ceb39
173
packages/adapters/src/executor-lease-poll.test.ts
Normal file
173
packages/adapters/src/executor-lease-poll.test.ts
Normal file
@ -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<typeof ComputerLifecycleModule>()),
|
||||
acquireComputerExecutionLease: async () => null,
|
||||
provisionComputer: async () => ({ id: "computer-1", kind: "desktop" }),
|
||||
}));
|
||||
|
||||
vi.mock("./auto-review.js", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof AutoReviewModule>()),
|
||||
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<AdapterContext>,
|
||||
) {
|
||||
blocked = true;
|
||||
await new Promise<void>((_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<string, unknown> }) => {
|
||||
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<typeof createRunExecutor>[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");
|
||||
});
|
||||
});
|
||||
@ -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<boolean> => {
|
||||
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) {
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user