mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
fix(coding-agent): keep codemode cancellation prompt and bounded
Refs #10748
This commit is contained in:
parent
b9ab2a6238
commit
3daa20280f
11 changed files with 495 additions and 25 deletions
|
|
@ -2,6 +2,19 @@
|
|||
|
||||
## [Unreleased]
|
||||
|
||||
### Breaking Changes
|
||||
|
||||
- Worker-backed results and `close()` now settle after requesting shutdown, not after worker exit. Hosts that require actual exit must track `onWorkerStarted(exited)` or retain capacity through `acquireWorker()`'s release callback
|
||||
|
||||
### Added
|
||||
|
||||
- Added `CodemodeSandbox`'s `onWorkerStarted(exited)` callback so hosts can count a worker against concurrency limits until it actually exits, independently of its script result
|
||||
- Added `CodemodeSandbox`'s `acquireWorker()` hook to reserve shared capacity before starting a worker and release it only on actual exit or constructor failure
|
||||
|
||||
### Fixed
|
||||
|
||||
- Fixed aborts, timeouts, and `close()` waiting for worker termination on Bun while a non-interruptible built-in runs; shutdown now requests termination without waiting for exit ([#10748](https://github.com/earendil-works/pi/issues/10748))
|
||||
|
||||
## [1.2.0] - 2026-10-12
|
||||
|
||||
### Breaking Changes
|
||||
|
|
|
|||
|
|
@ -59,6 +59,24 @@ await sandbox.close();
|
|||
|
||||
`memoryLimitBytes` caps the VM's heap. Allocations beyond it fail inside the script as `InternalError: out of memory`.
|
||||
|
||||
Aborts, timeouts, and `close()` request worker termination without waiting for the worker to exit. On Bun, a built-in such as a long `String.prototype.indexOf` may not poll QuickJS's interrupt handler, so its worker can keep computing until the built-in returns. Stopping unrefs the worker after requesting termination, because `terminate()` refs it internally. This does not guarantee computation has stopped. On Bun 1.4.0, a pending `terminate()` can still delay host process exit despite `unref()`; cancellation results and `close()` do not wait for it.
|
||||
|
||||
Hosts that limit concurrency can track actual worker lifetime with `onWorkerStarted(exited)`, called for each worker started by `execute()` or `start()`. Its promise resolves when that worker exits, not when the script result settles:
|
||||
|
||||
```ts
|
||||
const workers = new Set<Promise<void>>();
|
||||
const sandbox = new CodemodeSandbox({
|
||||
onWorkerStarted(exited) {
|
||||
workers.add(exited);
|
||||
void exited.then(() => workers.delete(exited));
|
||||
},
|
||||
});
|
||||
```
|
||||
|
||||
Keep a concurrency slot occupied until `exited` resolves. `close()` does not wait for these promises.
|
||||
|
||||
To enforce an admission limit, pass `acquireWorker()`. It runs synchronously after WASM loading, immediately before each worker is created. Throw to reject the execution without starting a worker; otherwise return a release callback. The sandbox calls it once on actual worker exit or if the worker constructor fails. Share reservation state across sandbox instances, and do not release it when a script result settles.
|
||||
|
||||
## Store
|
||||
|
||||
`store`/`load` let scripts keep values across executions. The sandbox does not persist anything itself: pass the current values as `options.store`, and a successful result reports what the script changed as `result.storeWrites` (`{ set, delete }`). Failed executions report no writes.
|
||||
|
|
@ -167,7 +185,7 @@ Bind the class with `script_name`. Binding a class exported from the agent's own
|
|||
}
|
||||
```
|
||||
|
||||
Then replace the sandbox. Tools, options other than `wasm` and `workerUrl`, and results stay the same:
|
||||
Then replace the sandbox. Tools, options other than the worker-specific `wasm`, `workerUrl`, `acquireWorker`, and `onWorkerStarted`, and results stay the same:
|
||||
|
||||
```ts
|
||||
import type { CodemodeSandboxDurableObject } from "@earendil-works/pi-codemode/cloudflare";
|
||||
|
|
@ -331,8 +349,8 @@ A run has no deadline unless `timeoutMs` is given. Result numbers must run 1, 2,
|
|||
| kind | meaning |
|
||||
| --------- | --------------------------------------------------------------------------- |
|
||||
| `script` | the script threw or failed to parse; `stack` points at `codemode.js:<line>` |
|
||||
| `timeout` | the deadline expired or the interrupt budget ran out; the VM was stopped |
|
||||
| `aborted` | `options.signal` fired or `close()` was called; the VM was stopped |
|
||||
| `timeout` | the deadline expired or the interrupt budget ran out; VM shutdown was requested |
|
||||
| `aborted` | `options.signal` fired or `close()` was called; VM shutdown was requested |
|
||||
| `sandbox` | the VM or its transport failed, for example a wasm trap, a missing worker file, or a failed remote |
|
||||
|
||||
`result.output` holds the text and image items in the order the script produced them, also for failed executions. The host keeps all of it until the script ends, so output is limited to `MAX_OUTPUT_CHARS` (16 Mi) characters of text and base64 image data and `MAX_OUTPUT_ITEMS` (100000) items. Past either limit the script fails with a `RangeError`, even if it catches the error. `result.calls` lists every tool call with `status: "ok" | "error" | "cancelled"`. A call that is still running when the script returns (not awaited) is aborted through the tool's `signal` and reported as `cancelled`.
|
||||
|
|
@ -351,6 +369,6 @@ The worker keeps script execution off the host thread: QuickJS runs synchronousl
|
|||
|
||||
## Runtime notes
|
||||
|
||||
- Works the same on Node and Bun, including memory limits, interrupts, and a catchable `RangeError` for deep recursion (QuickJS's stack guard is enabled; without it the wasm stack overflows and traps).
|
||||
- Node and Bun support memory limits, interrupts, and a catchable `RangeError` for deep recursion (QuickJS's stack guard is enabled; without it the wasm stack overflows and traps). Bun may delay worker exit during non-interruptible built-ins; results do not wait for exit.
|
||||
- Stack traces use QuickJS frames (`at f (codemode.js:2:31)`), prefixed with `Name: message` like V8. The wrapper prefix shares line 1 with the script, so line numbers are exact; column numbers on line 1 are shifted.
|
||||
- QuickJS is an interpreter. Glue code and filtering tool results are fast enough; heavy computation is slower than in V8.
|
||||
|
|
|
|||
|
|
@ -14,8 +14,8 @@
|
|||
* results go back through {@link VmChannel.post}. Replies produced before the launcher has
|
||||
* returned its channel are queued and flushed when it arrives.
|
||||
* 3. The first terminal event wins: `done`, `crash`, a launcher failure, the deadline, or an abort.
|
||||
* `finish()` aborts pending tool calls, stops the channel, and resolves the result once the
|
||||
* channel reports it stopped. Later events are ignored.
|
||||
* `finish()` aborts pending tool calls, asks the channel to stop, and resolves the result once
|
||||
* that request completes. A worker or remote VM may exit later. Later events are ignored.
|
||||
*
|
||||
* The launcher, not this class, is responsible for actually stopping a running VM: a worker is
|
||||
* terminated, an in-thread VM is interrupted and disposed, and a remote VM is told to stop at its
|
||||
|
|
@ -71,7 +71,7 @@ export interface VmEvents {
|
|||
/** The host's handle on a running VM. */
|
||||
export interface VmChannel {
|
||||
post(message: HostToWorkerMessage): void;
|
||||
/** Stops the VM. Called once, after the execution has its result. */
|
||||
/** Requests VM shutdown. Called once; may resolve before a worker or remote VM actually exits. */
|
||||
stop(): Promise<void>;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -14,27 +14,28 @@ function defaultWorkerUrl(): URL {
|
|||
* termination simple: a runaway script, including one that only spins the microtask queue, is
|
||||
* killed with `terminate()` and cannot poison a later run. Before terminating, the host sets a
|
||||
* shared interrupt flag that the VM polls, because Bun's `terminate()` cannot stop a thread that
|
||||
* is spinning in wasm.
|
||||
* is spinning in wasm. Stopping requests termination without waiting for exit: some built-ins do
|
||||
* not poll the flag, so on Bun a worker can outlive the script's result.
|
||||
*/
|
||||
function workerLauncher(
|
||||
wasm: CodemodeWasmModule | Promise<CodemodeWasmModule> | undefined,
|
||||
workerUrl: string | URL,
|
||||
): VmLauncher {
|
||||
function workerLauncher(options: CodemodeSandboxOptions): VmLauncher {
|
||||
return async (data, events) => {
|
||||
let module: CodemodeWasmModule;
|
||||
try {
|
||||
module = await (wasm ?? loadQuickJSWasm());
|
||||
module = await (options.wasm ?? loadQuickJSWasm());
|
||||
} catch (error) {
|
||||
throw new Error(`Failed to load QuickJS: ${errorMessage(error)}`);
|
||||
}
|
||||
const interrupt = new SharedArrayBuffer(4);
|
||||
const workerData: WorkerData = { ...data, wasm: module, interrupt };
|
||||
const releaseWorker = options.acquireWorker?.();
|
||||
let worker: Worker;
|
||||
try {
|
||||
worker = new Worker(workerUrl, { workerData });
|
||||
worker = new Worker(options.workerUrl ?? defaultWorkerUrl(), { workerData });
|
||||
} catch (error) {
|
||||
releaseWorker?.();
|
||||
throw new Error(`Failed to start worker: ${errorMessage(error)}`);
|
||||
}
|
||||
if (releaseWorker) worker.once("exit", releaseWorker);
|
||||
worker.on("message", (message: unknown) => {
|
||||
// Node 24/26 --watch leaks dependency reports onto the worker channel (nodejs/node#65044).
|
||||
if (
|
||||
|
|
@ -55,13 +56,21 @@ function workerLauncher(
|
|||
worker.on("exit", (code) => {
|
||||
events.failure({ kind: "sandbox", message: `Worker exited with code ${code} before the script settled` });
|
||||
});
|
||||
return {
|
||||
post: (message) => worker.postMessage(message),
|
||||
stop: async () => {
|
||||
Atomics.store(new Int32Array(interrupt), 0, 1);
|
||||
await worker.terminate().catch(() => undefined);
|
||||
},
|
||||
const stop = async (): Promise<void> => {
|
||||
Atomics.store(new Int32Array(interrupt), 0, 1);
|
||||
// A non-interruptible built-in can delay Bun's terminate() until it returns (#10748).
|
||||
// Do not hold up the result while that worker exits.
|
||||
void worker.terminate().catch(() => undefined);
|
||||
// terminate() refs the worker internally; unref it after requesting termination.
|
||||
worker.unref();
|
||||
};
|
||||
try {
|
||||
options.onWorkerStarted?.(new Promise<void>((resolve) => worker.once("exit", () => resolve())));
|
||||
} catch (error) {
|
||||
void stop();
|
||||
throw error;
|
||||
}
|
||||
return { post: (message) => worker.postMessage(message), stop };
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -74,6 +83,6 @@ function workerLauncher(
|
|||
*/
|
||||
export class CodemodeSandbox extends CodemodeSandboxBase {
|
||||
constructor(options: CodemodeSandboxOptions = {}) {
|
||||
super(options, workerLauncher(options.wasm, options.workerUrl ?? defaultWorkerUrl()));
|
||||
super(options, workerLauncher(options));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -60,11 +60,11 @@ export type CodemodeErrorKind =
|
|||
/** The script threw or failed to parse. `name` and `stack` come from the script's error. */
|
||||
| "script"
|
||||
/**
|
||||
* The overall deadline expired, or a VM on the calling thread used up its interrupt budget. The VM
|
||||
* was stopped.
|
||||
* The overall deadline expired, or a VM on the calling thread used up its interrupt budget.
|
||||
* The execution ended and the VM was asked to stop; a worker or remote VM may still be computing.
|
||||
*/
|
||||
| "timeout"
|
||||
/** The caller's signal fired or the sandbox was closed. The VM was stopped. */
|
||||
/** The caller's signal fired or the sandbox was closed. The execution ended and the VM was asked to stop. */
|
||||
| "aborted"
|
||||
/**
|
||||
* The VM or its transport failed outside the script's control (for example a wasm trap, a missing
|
||||
|
|
@ -146,6 +146,20 @@ export interface CodemodeSandboxOptions extends CodemodeSandboxBaseOptions {
|
|||
* specifier of an embedded build entrypoint; other hosts usually use a URL.
|
||||
*/
|
||||
workerUrl?: string | URL;
|
||||
/**
|
||||
* Synchronously reserves capacity immediately before creating each worker, after loading wasm.
|
||||
* Throw to deny admission without starting a worker. Return a release callback, called once if
|
||||
* worker construction fails or when the worker actually exits, not when the result settles.
|
||||
* The release callback must not throw. Reservations can be shared across sandbox instances.
|
||||
*/
|
||||
acquireWorker?(): () => void;
|
||||
/**
|
||||
* Called once for each worker started by execute() or start(). `exited` resolves when that worker
|
||||
* actually exits, which may be after the result and close() settle: on Bun a non-interruptible
|
||||
* built-in can delay termination. Use it to keep the worker counted against host concurrency
|
||||
* limits until it exits. A thrown callback error fails the execution and requests worker termination.
|
||||
*/
|
||||
onWorkerStarted?(exited: Promise<void>): void;
|
||||
}
|
||||
|
||||
/** A call the script made in a {@link CodemodeSandboxBase.start} run, answered with `CodemodeRun.deliver()`. */
|
||||
|
|
@ -216,7 +230,7 @@ export interface CodemodeRun {
|
|||
readonly result: Promise<CodemodeResult>;
|
||||
/**
|
||||
* Whether the run has its outcome: the script ended, timed out, or failed, or the run was aborted. Set when that is
|
||||
* decided, before `result` settles (which waits for the VM to stop); calls can no longer be answered.
|
||||
* decided, before `result` settles; calls can no longer be answered. It does not mean the VM has exited.
|
||||
*/
|
||||
readonly finished: boolean;
|
||||
/**
|
||||
|
|
|
|||
51
packages/codemode/test/fixtures/worker-process-exit.ts
vendored
Normal file
51
packages/codemode/test/fixtures/worker-process-exit.ts
vendored
Normal file
|
|
@ -0,0 +1,51 @@
|
|||
import { CodemodeSandbox } from "../../src/runtime/host.ts";
|
||||
|
||||
const how = process.argv[2];
|
||||
if (how !== "abort" && how !== "timeout" && how !== "close") throw new Error(`Unknown stop mode: ${how}`);
|
||||
const started = performance.now();
|
||||
process.on("exit", () => {
|
||||
console.log(JSON.stringify({ type: "exit", elapsedMs: performance.now() - started }));
|
||||
});
|
||||
const controller = new AbortController();
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
let workerExited = false;
|
||||
const sandbox = new CodemodeSandbox({
|
||||
memoryLimitBytes: 256 * 1024 * 1024,
|
||||
onWorkerStarted: (exited) => {
|
||||
void exited.then(() => {
|
||||
workerExited = true;
|
||||
});
|
||||
},
|
||||
tools: [
|
||||
{
|
||||
name: "ready",
|
||||
execute: () => {
|
||||
if (how !== "timeout") {
|
||||
timer = setTimeout(() => {
|
||||
if (how === "abort") controller.abort();
|
||||
else void sandbox.close();
|
||||
}, 100);
|
||||
}
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
try {
|
||||
const result = await sandbox.execute(
|
||||
"await tools.ready(); return 'a'.repeat(300000).indexOf('a'.repeat(10000) + 'b')",
|
||||
{ signal: controller.signal, timeoutMs: how === "timeout" ? 300 : Number.POSITIVE_INFINITY },
|
||||
);
|
||||
if (result.ok) throw new Error("The script completed instead of being cancelled");
|
||||
await sandbox.close();
|
||||
console.log(
|
||||
JSON.stringify({
|
||||
type: "result",
|
||||
kind: result.error.kind,
|
||||
elapsedMs: performance.now() - started,
|
||||
workerExited,
|
||||
}),
|
||||
);
|
||||
} finally {
|
||||
clearTimeout(timer);
|
||||
await sandbox.close();
|
||||
}
|
||||
229
packages/codemode/test/worker-lifetime.test.ts
Normal file
229
packages/codemode/test/worker-lifetime.test.ts
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,
|
||||
);
|
||||
});
|
||||
|
|
@ -6,6 +6,8 @@
|
|||
|
||||
- Pinned internal Pi dependencies to exact release versions to prevent fresh installs from mixing incompatible packages ([#10865](https://github.com/earendil-works/pi/issues/10865))
|
||||
- Fixed resumed sessions dropping tools activated during the session. Resume now restores the session's active tools, with `--tools +name`/`-name` entries applied to them, unless `--tools` with tool names, `--no-tools`, or `--no-builtin-tools` selects tools explicitly ([#10785](https://github.com/earendil-works/pi/issues/10785))
|
||||
- Fixed stopping a codemode script or reaching its deadline waiting for a non-interruptible built-in to return on Bun ([#10748](https://github.com/earendil-works/pi/issues/10748))
|
||||
- Fixed repeated codemode timeouts or aborts accumulating surviving workers; admission now limits live workers to two across sessions, counting cancelled workers until actual exit ([#10748](https://github.com/earendil-works/pi/issues/10748))
|
||||
|
||||
## [1.2.0] - 2026-10-12
|
||||
|
||||
|
|
|
|||
|
|
@ -224,5 +224,7 @@ for (const block of result.output) {
|
|||
## Limits
|
||||
|
||||
- A script's VM has 256 MB of memory. Running out throws `InternalError: out of memory`; filter or aggregate large data instead of accumulating it.
|
||||
- On Bun, stopping a script or reaching its deadline ends the call promptly, but a built-in that does not poll QuickJS's interrupt handler can keep computing in its worker until it returns.
|
||||
- At most two codemode workers may be alive across all sessions in one pi process. A cancelled worker occupies its slot until it exits. At capacity, new scripts fail immediately without starting another worker; retry after an existing worker exits.
|
||||
- A script that waits on a promise that can never settle (no tool call pending) fails immediately, since there are no timers.
|
||||
- Scripts cannot start other `codemode` scripts.
|
||||
|
|
|
|||
|
|
@ -54,6 +54,9 @@ const MAX_CONCURRENT_MODEL_CALLS = 4;
|
|||
* `InternalError: out of memory` inside the script.
|
||||
*/
|
||||
const CODEMODE_MEMORY_LIMIT_BYTES = 256 * 1024 * 1024;
|
||||
// Shared across calls and sessions. Allow one replacement while a cancelled worker is still exiting.
|
||||
const MAX_CONCURRENT_CODEMODE_WORKERS = 2;
|
||||
let activeCodemodeWorkers = 0;
|
||||
const MODEL_TYPES: ReadonlySet<string> = new Set<ModelType>(["chat", "image", "classifier"]);
|
||||
|
||||
function truncateText(text: string, maxChars: number): string {
|
||||
|
|
@ -508,6 +511,17 @@ export async function executeCodemode(
|
|||
memoryLimitBytes: CODEMODE_MEMORY_LIMIT_BYTES,
|
||||
wasm: loadQuickJSWasm(getQuickJSWasmPath()),
|
||||
workerUrl: getCodemodeWorkerSpecifier(),
|
||||
acquireWorker: () => {
|
||||
if (activeCodemodeWorkers >= MAX_CONCURRENT_CODEMODE_WORKERS) {
|
||||
throw new Error(
|
||||
`Codemode worker limit reached (${MAX_CONCURRENT_CODEMODE_WORKERS} workers are still running or shutting down). Retry after they exit.`,
|
||||
);
|
||||
}
|
||||
activeCodemodeWorkers++;
|
||||
return () => {
|
||||
activeCodemodeWorkers--;
|
||||
};
|
||||
},
|
||||
});
|
||||
|
||||
let result: CodemodeResult;
|
||||
|
|
|
|||
118
packages/coding-agent/test/codemode-worker-admission.test.ts
Normal file
118
packages/coding-agent/test/codemode-worker-admission.test.ts
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
import { Worker } from "node:worker_threads";
|
||||
import * as codemode from "@earendil-works/pi-codemode";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { executeCodemode } from "../src/extensions/codemode/execute.ts";
|
||||
|
||||
const terminate = Worker.prototype.terminate;
|
||||
const terminations = new Map<Worker, (code: number) => void>();
|
||||
|
||||
beforeEach(() => {
|
||||
vi.spyOn(Worker.prototype, "terminate").mockImplementation(function (this: Worker) {
|
||||
this.ref();
|
||||
return new Promise<number>((resolve) => {
|
||||
terminations.set(this, resolve);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
vi.restoreAllMocks();
|
||||
await Promise.all(
|
||||
[...terminations].map(async ([worker, resolve]) => {
|
||||
resolve(await terminate.call(worker));
|
||||
}),
|
||||
);
|
||||
terminations.clear();
|
||||
});
|
||||
|
||||
const run = (code: string, signal?: AbortSignal) =>
|
||||
executeCodemode("worker-admission", { code }, signal, undefined, undefined);
|
||||
|
||||
describe("codemode worker admission", () => {
|
||||
// #10748: fresh sandboxes must share capacity and retain it until actual worker exit.
|
||||
it.each(["timeout", "abort", "success"] as const)("bounds surviving workers after %s", async (how) => {
|
||||
const once = vi.spyOn(Worker.prototype, "once");
|
||||
for (let i = 0; i < 2; i++) {
|
||||
const controller = new AbortController();
|
||||
const exitListenersBefore = once.mock.calls.filter(([event]) => event === "exit").length;
|
||||
const pending = run(
|
||||
how === "timeout"
|
||||
? '// @options: {"timeout_ms": 200}\nwhile (true) {}'
|
||||
: how === "abort"
|
||||
? "while (true) {}"
|
||||
: "return 1",
|
||||
controller.signal,
|
||||
);
|
||||
if (how === "abort") {
|
||||
await vi.waitFor(() =>
|
||||
expect(once.mock.calls.filter(([event]) => event === "exit").length).toBeGreaterThan(
|
||||
exitListenersBefore,
|
||||
),
|
||||
);
|
||||
controller.abort();
|
||||
}
|
||||
const result = await pending;
|
||||
if (how === "success") expect(result.isError).toBeUndefined();
|
||||
else {
|
||||
expect(result.isError).toBe(true);
|
||||
expect(result.content).toContainEqual({
|
||||
type: "text",
|
||||
text: expect.stringContaining(how === "timeout" ? "Script timed out:" : "Script aborted:"),
|
||||
});
|
||||
}
|
||||
expect(terminations.size).toBe(i + 1);
|
||||
}
|
||||
|
||||
const blocked = await run("return 3");
|
||||
expect(blocked).toMatchObject({ isError: true, details: { calls: [] } });
|
||||
expect(blocked.content).toContainEqual({
|
||||
type: "text",
|
||||
text: expect.stringContaining("Codemode worker limit reached (2 workers are still running or shutting down)"),
|
||||
});
|
||||
expect(terminations.size).toBe(2);
|
||||
|
||||
const aborted = await run("while (true) {}", AbortSignal.abort());
|
||||
expect(aborted.isError).toBe(true);
|
||||
expect(aborted.content).toContainEqual({ type: "text", text: expect.stringContaining("Script aborted:") });
|
||||
expect(terminations.size).toBe(2);
|
||||
|
||||
// Releasing one actual worker, not its result, permits exactly one new execution.
|
||||
const [worker, resolve] = [...terminations][0];
|
||||
resolve(await terminate.call(worker));
|
||||
terminations.delete(worker);
|
||||
expect((await run("return 42")).isError).toBeUndefined();
|
||||
expect(terminations.size).toBe(2);
|
||||
});
|
||||
|
||||
// #10748: timed-out launches resuming after WASM loading must obey the same limit.
|
||||
it("bounds workers that start after their results already timed out", async () => {
|
||||
const wasm = await codemode.loadQuickJSWasm();
|
||||
const loading = Promise.withResolvers<codemode.CodemodeWasmModule>();
|
||||
vi.spyOn(codemode, "loadQuickJSWasm").mockReturnValue(loading.promise);
|
||||
try {
|
||||
const results = await Promise.all(
|
||||
Array.from({ length: 3 }, () => run('// @options: {"timeout_ms": 20}\nwhile (true) {}')),
|
||||
);
|
||||
expect(results.every((result) => result.isError)).toBe(true);
|
||||
expect(terminations.size).toBe(0);
|
||||
loading.resolve(wasm);
|
||||
const blocked = await run("return 1");
|
||||
expect(blocked.isError).toBe(true);
|
||||
expect(blocked.content).toContainEqual({
|
||||
type: "text",
|
||||
text: expect.stringContaining("Codemode worker limit reached"),
|
||||
});
|
||||
expect(terminations.size).toBe(2);
|
||||
} finally {
|
||||
loading.resolve(wasm);
|
||||
}
|
||||
});
|
||||
|
||||
// #10748: concurrent callers must not race past the shared admission limit.
|
||||
it("admits only two workers from a concurrent burst", async () => {
|
||||
const results = await Promise.all(Array.from({ length: 6 }, () => run("return 1")));
|
||||
expect(results.filter((result) => !result.isError)).toHaveLength(2);
|
||||
expect(results.filter((result) => result.isError)).toHaveLength(4);
|
||||
expect(terminations.size).toBe(2);
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue