From 885a564b3b22dc6eefd55a78a8f6bb757ddfbc69 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 20 Oct 2026 20:49:50 +0200 Subject: [PATCH] fix(coding-agent): queue codemode scripts at the worker limit Scripts wait in arrival order for a free worker instead of failing immediately. The wait counts toward timeout_ms, honours abort, and is capped at 60 seconds. PI_CODEMODE_MAX_WORKERS raises the process-wide limit of two workers for hosts running many sessions in one process. closes #11184 --- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/docs/codemode.md | 4 +- .../docs/environment-variables.md | 1 + .../src/extensions/codemode/execute.ts | 130 ++++++++++++++---- .../test/codemode-worker-admission.test.ts | 113 +++++++++++---- 5 files changed, 192 insertions(+), 57 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 34551bcbc..cbf3ab469 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -19,6 +19,7 @@ - Fixed the terminal staying in mouse-tracking, focus-reporting, Kitty keyboard, and alternate-screen modes when pi exits without its normal terminal teardown, which flooded the shell prompt with escape sequences ([#11170](https://github.com/earendil-works/pi/issues/11170)) - Fixed concurrent extension `select`/`confirm`/`input`/`editor` dialogs and inline `custom()` components in interactive mode: a second dialog replaced the first and the first promise never settled; closing the visible dialog now shows the previous one again ([#11166](https://github.com/earendil-works/pi/issues/11166), [#9909](https://github.com/earendil-works/pi/issues/9909)) - Fixed pasting in the shell hanging after copying text in pi and suspending it with Ctrl+Z on Linux: the background `wl-copy`/`xclip`/`xsel` process that keeps serving the clipboard was suspended along with pi ([#11179](https://github.com/earendil-works/pi/issues/11179)) +- Fixed codemode scripts failing immediately when both codemode workers are busy, for example with parallel codemode calls or in-process subagents. Scripts now wait up to 60 seconds for a free worker, and `PI_CODEMODE_MAX_WORKERS` raises the limit of two workers ([#11184](https://github.com/earendil-works/pi/issues/11184)) ## [1.3.0] - 2026-10-15 diff --git a/packages/coding-agent/docs/codemode.md b/packages/coding-agent/docs/codemode.md index 12b41277d..6dd76cbad 100644 --- a/packages/coding-agent/docs/codemode.md +++ b/packages/coding-agent/docs/codemode.md @@ -13,7 +13,7 @@ A script may start with an options line. It must be the first line; an options l ``` - `max_output_tokens` (default 10000) limits the output. Longer output keeps its start and end, and the full text is written to a temp file whose path is included in the result. A script fails when its output passes 16777216 characters of text and base64 image data or 100000 `text()`, `image()`, and `console` calls; write large data to a file with a tool instead. -- `timeout_ms` is a hard deadline for the whole script. It is unset by default. Image generation can take minutes, so do not set a short deadline for scripts that generate images. +- `timeout_ms` is a hard deadline for the whole script, including any wait for a free worker. It is unset by default. Image generation can take minutes, so do not set a short deadline for scripts that generate images. The result starts with `Script completed` or `Script failed`, the wall time, and the output. Text and image items appear in order, each on its own line. When the output has more than one text item (from `text()` or `return`), each starts with a `==> text N/M <==` line. `console` calls follow in one `` block with one line per call. A failed script keeps its partial output, followed by `Script error:` and the error. Tool calls are real: calls made before a failure are not undone. Calls still running when the script ends are cancelled, and unawaited promises are discarded. @@ -225,6 +225,6 @@ for (const block of result.output) { - 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. +- At most two codemode workers may be alive across all sessions in one pi process; `PI_CODEMODE_MAX_WORKERS` changes the limit. A cancelled worker occupies its slot until it exits. At capacity, new scripts wait for a free worker in arrival order. A script that gets no worker within 60 seconds, or before its `timeout_ms`, fails without running. - 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. diff --git a/packages/coding-agent/docs/environment-variables.md b/packages/coding-agent/docs/environment-variables.md index 8a1975043..828d617dd 100644 --- a/packages/coding-agent/docs/environment-variables.md +++ b/packages/coding-agent/docs/environment-variables.md @@ -84,6 +84,7 @@ These variables are read by Pi itself: | `PI_OFFLINE` | Disable automatic network activity, including model catalog refreshes | | `PI_SKIP_VERSION_CHECK` | Disable the `pi.dev` latest-version request | | `PI_TELEMETRY` | Override install/update telemetry and provider attribution headers: `1`/`true`/`yes` or `0`/`false`/`no` | +| `PI_CODEMODE_MAX_WORKERS` | Maximum codemode workers alive at once across all sessions in the process; default `2`. Each worker's script may use up to 256 MB. Raise it when one process runs several sessions, such as in-process subagents. Further scripts wait for a free worker. See [Codemode](codemode.md) | | `PI_CACHE_RETENTION` | Set to `long` for extended provider prompt caching where supported | | `PI_SHARE_VIEWER_URL` | Override the base URL used by `/share` | | `PI_RADIUS_GATEWAY` | Override the Radius gateway origin used by `/bug` uploads | diff --git a/packages/coding-agent/src/extensions/codemode/execute.ts b/packages/coding-agent/src/extensions/codemode/execute.ts index 2d3bfbec8..b3635999f 100644 --- a/packages/coding-agent/src/extensions/codemode/execute.ts +++ b/packages/coding-agent/src/extensions/codemode/execute.ts @@ -55,8 +55,60 @@ const MAX_CONCURRENT_MODEL_CALLS = 4; */ 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; +const DEFAULT_MAX_CODEMODE_WORKERS = 2; + +/** `PI_CODEMODE_MAX_WORKERS` if it is a positive integer. Read per call so SDK hosts can set it at runtime. */ +function maxCodemodeWorkers(): number { + const value = Number(process.env.PI_CODEMODE_MAX_WORKERS); + return Number.isInteger(value) && value > 0 ? value : DEFAULT_MAX_CODEMODE_WORKERS; +} +/** + * How long a script waits for a free worker before failing. Bounds the wait when every worker is + * held by a script that waits on another script, for example an in-process subagent that uses codemode. + */ +const CODEMODE_WORKER_WAIT_MS = 60_000; let activeCodemodeWorkers = 0; +/** Scripts waiting for a worker, in arrival order. */ +const codemodeWorkerWaiters: (() => void)[] = []; + +/** + * Reserve one worker slot, waiting in arrival order while all are taken. Resolves with why it could + * not reserve one: the signal aborted, or `timeoutMs` passed. + */ +function reserveCodemodeWorker( + signal: AbortSignal | undefined, + timeoutMs: number, +): Promise<"reserved" | "aborted" | "timeout"> { + if (signal?.aborted) return Promise.resolve("aborted"); + return new Promise((resolve) => { + const settle = (outcome: "reserved" | "aborted" | "timeout") => { + clearTimeout(timer); + signal?.removeEventListener("abort", onAbort); + const index = codemodeWorkerWaiters.indexOf(grant); + if (index >= 0) codemodeWorkerWaiters.splice(index, 1); + resolve(outcome); + }; + const grant = () => settle("reserved"); + const onAbort = () => settle("aborted"); + const timer = setTimeout(() => settle("timeout"), timeoutMs); + codemodeWorkerWaiters.push(grant); + signal?.addEventListener("abort", onAbort, { once: true }); + admitCodemodeWorkers(); + }); +} + +/** Give free slots to waiting scripts in arrival order. */ +function admitCodemodeWorkers(): void { + while (codemodeWorkerWaiters.length > 0 && activeCodemodeWorkers < maxCodemodeWorkers()) { + activeCodemodeWorkers++; + codemodeWorkerWaiters.shift()?.(); + } +} + +function releaseCodemodeWorker(): void { + activeCodemodeWorkers--; + admitCodemodeWorkers(); +} const MODEL_TYPES: ReadonlySet = new Set(["chat", "image", "classifier"]); function truncateText(text: string, maxChars: number): string { @@ -499,37 +551,63 @@ export async function executeCodemode( }, })); - const sandbox = new CodemodeSandbox({ - tools: sandboxTools, - globals: [ - ...createDiscoveryGlobals(callable, samples, options), - ...(options.models && ctx - ? createModelGlobals(ctx.modelRegistry, toolCallId, calls, publish, addModelUsage, addGeneratedImages) - : []), - ], - timeoutMs: sourceOptions.timeoutMs ?? Number.POSITIVE_INFINITY, - 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--; - }; - }, - }); + // The wait for a worker counts toward the script's deadline. + const deadline = sourceOptions.timeoutMs ?? Number.POSITIVE_INFINITY; + const waitStartedAt = performance.now(); + const admission = await reserveCodemodeWorker(signal, Math.min(deadline, CODEMODE_WORKER_WAIT_MS)); + const waitedMs = performance.now() - waitStartedAt; + if (admission !== "reserved") { + const waited = ((performance.now() - startedAt) / 1000).toFixed(1); + const reason = signal?.reason; + const error = + admission === "aborted" + ? `Script aborted: ${reason instanceof Error ? reason.message : "Execution aborted"}` + : deadline <= CODEMODE_WORKER_WAIT_MS + ? `Script timed out: Execution timed out after ${deadline} ms while waiting for a free codemode worker` + : `No codemode worker became free within ${CODEMODE_WORKER_WAIT_MS / 1000} seconds. All ${maxCodemodeWorkers()} workers, shared by every session in this process, are held by other scripts. Retry after they finish.`; + return { + content: [ + { type: "text", text: `Script failed\nWall time ${waited} seconds\nOutput:\n` }, + { type: "text", text: `Script error:\n${error}\n\nThe script did not start.` }, + ], + details: snapshot(), + isError: true, + }; + } + // The reservation covers the one worker this script starts. A launch that resumes after the result + // settled finds it released and starts no worker. + let reserved = true; + let sandbox: CodemodeSandbox | undefined; let result: CodemodeResult; try { + sandbox = new CodemodeSandbox({ + tools: sandboxTools, + globals: [ + ...createDiscoveryGlobals(callable, samples, options), + ...(options.models && ctx + ? createModelGlobals(ctx.modelRegistry, toolCallId, calls, publish, addModelUsage, addGeneratedImages) + : []), + ], + timeoutMs: Math.max(1, Math.round(deadline - waitedMs)), + memoryLimitBytes: CODEMODE_MEMORY_LIMIT_BYTES, + wasm: loadQuickJSWasm(getQuickJSWasmPath()), + workerUrl: getCodemodeWorkerSpecifier(), + acquireWorker: () => { + if (!reserved) throw new Error("The script already settled"); + // From here the sandbox releases the slot when the worker actually exits. + reserved = false; + return releaseCodemodeWorker; + }, + }); const store = ctx ? readCodemodeStore(ctx.sessionManager.getBranch()) : {}; result = await sandbox.execute(code, { signal, store }); } finally { - await sandbox.close(); + await sandbox?.close(); + if (reserved) { + reserved = false; + releaseCodemodeWorker(); + } } // Calls still marked running were cut off by the script ending, a timeout, or an abort. for (const call of calls) { diff --git a/packages/coding-agent/test/codemode-worker-admission.test.ts b/packages/coding-agent/test/codemode-worker-admission.test.ts index 5026a3fde..119fc82d2 100644 --- a/packages/coding-agent/test/codemode-worker-admission.test.ts +++ b/packages/coding-agent/test/codemode-worker-admission.test.ts @@ -17,6 +17,7 @@ beforeEach(() => { afterEach(async () => { vi.restoreAllMocks(); + vi.unstubAllEnvs(); await Promise.all( [...terminations].map(async ([worker, resolve]) => { resolve(await terminate.call(worker)); @@ -28,6 +29,15 @@ afterEach(async () => { const run = (code: string, signal?: AbortSignal) => executeCodemode("worker-admission", { code }, signal, undefined, undefined); +/** Let the first held worker actually exit. */ +async function releaseOne(): Promise { + const [worker, resolve] = [...terminations][0]; + terminations.delete(worker); + resolve(await terminate.call(worker)); +} + +const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + 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) => { @@ -63,29 +73,43 @@ describe("codemode worker admission", () => { 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)"), + // #11184: at capacity, a script waits for a worker instead of failing. + let waitingSettled = false; + const waiting = run("return 3").finally(() => { + waitingSettled = true; }); + + const abortWhileWaiting = new AbortController(); + const aborted = run("while (true) {}", abortWhileWaiting.signal); + await sleep(50); + abortWhileWaiting.abort(); + const abortedResult = await aborted; + expect(abortedResult.isError).toBe(true); + expect(abortedResult.content).toContainEqual({ + type: "text", + text: expect.stringMatching(/Script aborted:[\s\S]*The script did not start\./), + }); + + const timedOut = await run('// @options: {"timeout_ms": 50}\nreturn 1'); + expect(timedOut.isError).toBe(true); + expect(timedOut.content).toContainEqual({ + type: "text", + text: expect.stringContaining("Execution timed out after 50 ms while waiting for a free codemode worker"), + }); + + expect(waitingSettled).toBe(false); 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(); + // Releasing one actual worker, not its result, admits exactly the waiting script. + await releaseOne(); + const admitted = await waiting; + expect(admitted.isError).toBeUndefined(); + expect(admitted.content).toContainEqual({ type: "text", text: "3" }); 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 () => { + // #10748: timed-out launches resuming after WASM loading must not start workers. + it("starts no worker for a launch that resumes after its result timed out", async () => { const wasm = await codemode.loadQuickJSWasm(); const loading = Promise.withResolvers(); vi.spyOn(codemode, "loadQuickJSWasm").mockReturnValue(loading.promise); @@ -96,23 +120,54 @@ describe("codemode worker admission", () => { 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); + expect((await run("return 1")).isError).toBeUndefined(); + // Only the last script's worker started. + expect(terminations.size).toBe(1); } 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); + // #10748, #11184: a concurrent burst runs two workers at a time and queues the rest. + it("queues a concurrent burst behind two workers", async () => { + const results: Awaited>[] = []; + const pending = Array.from({ length: 6 }, () => run("return 1").then((result) => results.push(result))); + for (let released = 0; released < 4; released++) { + await vi.waitFor(() => { + expect(terminations.size).toBe(2); + expect(results).toHaveLength(released + 2); + }); + await sleep(50); + expect(results).toHaveLength(released + 2); + await releaseOne(); + } + await Promise.all(pending); + expect(results.every((result) => result.isError === undefined)).toBe(true); expect(terminations.size).toBe(2); }); + + // #11184: in-process hosts with many sessions can raise the limit. + it("reads the limit from PI_CODEMODE_MAX_WORKERS", async () => { + vi.stubEnv("PI_CODEMODE_MAX_WORKERS", "3"); + const results: Awaited>[] = []; + const pending = Array.from({ length: 4 }, () => run("return 1").then((result) => results.push(result))); + await vi.waitFor(() => expect(results).toHaveLength(3)); + await sleep(50); + expect(results).toHaveLength(3); + expect(terminations.size).toBe(3); + await releaseOne(); + await Promise.all(pending); + expect(results.every((result) => result.isError === undefined)).toBe(true); + }); + + it("ignores an invalid PI_CODEMODE_MAX_WORKERS", async () => { + vi.stubEnv("PI_CODEMODE_MAX_WORKERS", "0"); + const results: Awaited>[] = []; + const pending = Array.from({ length: 3 }, () => run("return 1").then((result) => results.push(result))); + await vi.waitFor(() => expect(results).toHaveLength(2)); + await sleep(50); + expect(results).toHaveLength(2); + await releaseOne(); + await Promise.all(pending); + }); });