chore(vendor): import pi 1.3.0
Imported as a linear vendor commit so forge rebases do not replay the upstream root tree. git-subtree-dir: third_party/pi git-subtree-split: 71023aec1546d6d677a65669c5665c0927ce0dd0
This commit is contained in:
229
third_party/pi/packages/codemode/test/worker-lifetime.test.ts
vendored
Normal file
229
third_party/pi/packages/codemode/test/worker-lifetime.test.ts
vendored
Normal file
@@ -0,0 +1,229 @@
|
||||
import { execFile, spawnSync } from "node:child_process";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import { promisify } from "node:util";
|
||||
import { Worker } from "node:worker_threads";
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { type CodemodeResult, CodemodeSandbox } from "../src/index.ts";
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("worker lifetime", () => {
|
||||
// #10748: reject admission before creating a worker, not by terminating an excess worker.
|
||||
it("does not start a worker when admission is denied", async () => {
|
||||
const onWorkerStarted = vi.fn();
|
||||
const acquireWorker = vi.fn(() => {
|
||||
throw new Error("worker capacity exhausted");
|
||||
});
|
||||
const sandbox = new CodemodeSandbox({ acquireWorker, onWorkerStarted });
|
||||
try {
|
||||
expect(await sandbox.execute("while (true) {}")).toMatchObject({
|
||||
ok: false,
|
||||
error: { kind: "sandbox", message: "worker capacity exhausted" },
|
||||
});
|
||||
expect(acquireWorker).toHaveBeenCalledTimes(1);
|
||||
expect(onWorkerStarted).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
await sandbox.close();
|
||||
}
|
||||
});
|
||||
|
||||
// #10748: a failed Worker constructor must return its reserved capacity.
|
||||
it("releases admission when worker construction fails", async () => {
|
||||
const releaseWorker = vi.fn();
|
||||
const acquireWorker = vi.fn(() => releaseWorker);
|
||||
const sandbox = new CodemodeSandbox({
|
||||
acquireWorker,
|
||||
workerUrl: new URL("https://invalid.example/worker.js"),
|
||||
});
|
||||
try {
|
||||
expect(await sandbox.execute("return 1")).toMatchObject({
|
||||
ok: false,
|
||||
error: { kind: "sandbox", message: expect.stringContaining("Failed to start worker:") },
|
||||
});
|
||||
expect(acquireWorker).toHaveBeenCalledTimes(1);
|
||||
expect(releaseWorker).toHaveBeenCalledTimes(1);
|
||||
} finally {
|
||||
await sandbox.close();
|
||||
}
|
||||
});
|
||||
|
||||
// #10748: loading failures must not reserve worker capacity.
|
||||
it("does not acquire admission when wasm loading fails", async () => {
|
||||
const acquireWorker = vi.fn(() => vi.fn());
|
||||
const sandbox = new CodemodeSandbox({
|
||||
acquireWorker,
|
||||
wasm: Promise.reject(new Error("wasm unavailable")),
|
||||
});
|
||||
try {
|
||||
expect(await sandbox.execute("return 1")).toMatchObject({
|
||||
ok: false,
|
||||
error: { kind: "sandbox", message: "Failed to load QuickJS: wasm unavailable" },
|
||||
});
|
||||
expect(acquireWorker).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
await sandbox.close();
|
||||
}
|
||||
});
|
||||
|
||||
// #10748: Bun can leave terminate() pending while a non-interruptible built-in runs.
|
||||
it.each(["abort", "timeout", "close"] as const)("settles %s without waiting for worker exit", async (how) => {
|
||||
vi.useFakeTimers();
|
||||
const terminate = Worker.prototype.terminate;
|
||||
let worker: Worker | undefined;
|
||||
let releaseTermination: (code: number) => void = () => {};
|
||||
const termination = new Promise<number>((resolve) => {
|
||||
releaseTermination = resolve;
|
||||
});
|
||||
const terminateSpy = vi.spyOn(Worker.prototype, "terminate").mockImplementation(function (this: Worker) {
|
||||
worker = this;
|
||||
// Like the real Node/Bun implementation, termination refs the worker while waiting for exit.
|
||||
this.ref();
|
||||
return termination;
|
||||
});
|
||||
const refSpy = vi.spyOn(Worker.prototype, "ref");
|
||||
const unrefSpy = vi.spyOn(Worker.prototype, "unref");
|
||||
let ready: () => void = () => {};
|
||||
const called = new Promise<void>((resolve) => {
|
||||
ready = resolve;
|
||||
});
|
||||
let toolSignal: AbortSignal | undefined;
|
||||
let exited: Promise<void> | undefined;
|
||||
let didExit = false;
|
||||
const sandbox = new CodemodeSandbox({
|
||||
timeoutMs: 100,
|
||||
onWorkerStarted: (promise) => {
|
||||
exited = promise;
|
||||
void promise.then(() => {
|
||||
didExit = true;
|
||||
});
|
||||
},
|
||||
tools: [
|
||||
{
|
||||
name: "hang",
|
||||
execute: (_args, { signal }) => {
|
||||
toolSignal = signal;
|
||||
ready();
|
||||
return new Promise(() => {});
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
const controller = new AbortController();
|
||||
let result: CodemodeResult | undefined;
|
||||
let didClose = false;
|
||||
const pending = sandbox.execute('text("partial"); await tools.hang()', { signal: controller.signal });
|
||||
void pending.then((value) => {
|
||||
result = value;
|
||||
});
|
||||
try {
|
||||
await called;
|
||||
if (how === "abort") controller.abort(new Error("cancelled"));
|
||||
if (how === "timeout") await vi.advanceTimersByTimeAsync(100);
|
||||
if (how === "close") {
|
||||
void sandbox.close().then(() => {
|
||||
didClose = true;
|
||||
});
|
||||
}
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
expect(result).toMatchObject({
|
||||
ok: false,
|
||||
error: { kind: how === "timeout" ? "timeout" : "aborted" },
|
||||
output: [{ type: "text", text: "partial" }],
|
||||
calls: [{ name: "hang", status: "cancelled" }],
|
||||
});
|
||||
expect(toolSignal?.aborted).toBe(true);
|
||||
expect(terminateSpy).toHaveBeenCalledTimes(1);
|
||||
expect(unrefSpy).toHaveBeenCalledTimes(1);
|
||||
expect(refSpy).toHaveBeenCalledTimes(1);
|
||||
expect(unrefSpy.mock.invocationCallOrder[0]).toBeGreaterThan(refSpy.mock.invocationCallOrder[0]);
|
||||
expect(exited).toBeDefined();
|
||||
expect(didExit).toBe(false);
|
||||
if (how === "close") expect(didClose).toBe(true);
|
||||
// Neither close() nor late worker messages may replace the decided outcome or wait for exit.
|
||||
await sandbox.close();
|
||||
worker?.emit("message", { type: "done", ok: true, value: "42", writes: "[]" });
|
||||
expect(await pending).toBe(result);
|
||||
expect(terminateSpy).toHaveBeenCalledTimes(1);
|
||||
} finally {
|
||||
vi.useRealTimers();
|
||||
terminateSpy.mockRestore();
|
||||
if (worker) releaseTermination(await terminate.call(worker));
|
||||
await sandbox.close();
|
||||
}
|
||||
await exited;
|
||||
expect(didExit).toBe(true);
|
||||
});
|
||||
|
||||
// #10748: exit accounting applies to successful executions and host-controlled runs too.
|
||||
it("reports the actual exit of execute() and start() workers", async () => {
|
||||
const exits: Promise<void>[] = [];
|
||||
const releaseWorker = vi.fn();
|
||||
const acquireWorker = vi.fn(() => releaseWorker);
|
||||
const sandbox = new CodemodeSandbox({ acquireWorker, onWorkerStarted: (exited) => void exits.push(exited) });
|
||||
try {
|
||||
expect(await sandbox.execute("return 1")).toMatchObject({ ok: true, value: 1 });
|
||||
const run = sandbox.start("return 2", { startedAt: 0, seed: 0, onCall: () => {} });
|
||||
expect(await run.result).toMatchObject({ ok: true, value: 2 });
|
||||
expect(run.finished).toBe(true);
|
||||
expect(exits).toHaveLength(2);
|
||||
await Promise.all(exits);
|
||||
expect(acquireWorker).toHaveBeenCalledTimes(2);
|
||||
expect(releaseWorker).toHaveBeenCalledTimes(2);
|
||||
} finally {
|
||||
await sandbox.close();
|
||||
}
|
||||
});
|
||||
|
||||
// #10748: a failing lifetime hook must not leak the worker it was handed.
|
||||
it("requests termination when onWorkerStarted throws", async () => {
|
||||
let exited: Promise<void> | undefined;
|
||||
const releaseWorker = vi.fn();
|
||||
const sandbox = new CodemodeSandbox({
|
||||
acquireWorker: () => releaseWorker,
|
||||
onWorkerStarted: (promise) => {
|
||||
exited = promise;
|
||||
throw new Error("lifetime hook failed");
|
||||
},
|
||||
});
|
||||
try {
|
||||
expect(await sandbox.execute("while (true) {}")).toMatchObject({
|
||||
ok: false,
|
||||
error: { kind: "sandbox", message: "lifetime hook failed" },
|
||||
});
|
||||
expect(exited).toBeDefined();
|
||||
await exited;
|
||||
expect(releaseWorker).toHaveBeenCalledTimes(1);
|
||||
} finally {
|
||||
await sandbox.close();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
const runtimes = [
|
||||
{ name: "Node", command: process.execPath, available: true },
|
||||
{ name: "Bun", command: "bun", available: spawnSync("bun", ["--version"], { timeout: 5000 }).status === 0 },
|
||||
];
|
||||
|
||||
describe.each(runtimes)("$name process lifetime", ({ name, command, available }) => {
|
||||
// #10748: verify both prompt cancellation and, on Node, prompt process exit with a real worker.
|
||||
// Bun 1.4.0's pending terminate() keeps the process alive even after unref(), independently of codemode.
|
||||
it.skipIf(!available).each(["abort", "timeout", "close"] as const)(
|
||||
name === "Node" ? "exits promptly after %s with a real worker" : "settles %s promptly with a real worker",
|
||||
async (how) => {
|
||||
const fixture = fileURLToPath(new URL("./fixtures/worker-process-exit.ts", import.meta.url));
|
||||
const { stdout } = await promisify(execFile)(command, [fixture, how], { timeout: 20000 });
|
||||
const lines = stdout.trim().split(/\r?\n/);
|
||||
expect(lines).toHaveLength(2);
|
||||
const result = JSON.parse(lines[0]) as { type: string; kind: string; elapsedMs: number };
|
||||
const exit = JSON.parse(lines[1]) as { type: string; elapsedMs: number };
|
||||
expect(result).toMatchObject({ type: "result", kind: how === "timeout" ? "timeout" : "aborted" });
|
||||
expect(exit.type).toBe("exit");
|
||||
expect(result.elapsedMs).toBeLessThan(1000);
|
||||
if (name === "Node") expect(exit.elapsedMs - result.elapsedMs).toBeLessThan(1000);
|
||||
},
|
||||
25000,
|
||||
);
|
||||
});
|
||||
Reference in New Issue
Block a user