feat(coding-agent): codemode and MCP

This adds codemode and MCP support to pi. Codemode (packages/codemode) runs
model-written JavaScript in a QuickJS wasm VM inside a worker. From there the
script calls pi's tools as async functions, uses a per-session store, and can
read the model catalog and run classifiers. The codemode tool in coding-agent
is compatible with Codex's exec tool, and codemode.mode works like Codex's
tool modes (on or only). Large results are written to a JSON temp file. MCP
(packages/mcp) is a standalone client for stdio and streamable HTTP servers,
with OAuth.

The core only gets general mechanisms; codemode, tool_search and MCP are
built-in extensions that use them. Other extensions can replace these
built-ins. Tools declare an exposure: direct, model-only, codemode, deferred
or hidden. prepareLoadout() lets an orchestrating tool change which tool
declarations the model sees. ctx.executeTool() runs nested calls through
validation and the tool_call/tool_result hooks. Those calls emit events with
parentToolCallId and are recorded as bounded nestedCalls on the parent
result, which compaction and HTML export use. Tools can also return
structuredContent with an outputSchema, and can report errors with isError
instead of throwing. MCP servers come from mcp.json (global, or per project
after the project is trusted) or from pi.registerMcpServer(). They are
managed with /mcp and pi mcp list|login|logout, and their runtimes load on
first use.

Closes #10040
This commit is contained in:
Armin Ronacher 2026-09-29 09:12:50 +02:00
parent b485fa3128
commit 8562bcf66a
143 changed files with 17620 additions and 225 deletions

View file

@ -0,0 +1,361 @@
import { Worker } from "node:worker_threads";
import { toCodemodeIdentifier } from "../identifier.ts";
import type {
CodemodeCall,
CodemodeCallStatus,
CodemodeError,
CodemodeExecuteOptions,
CodemodeOutputItem,
CodemodeResult,
CodemodeSandboxOptions,
CodemodeStoreWrites,
CodemodeTool,
} from "../types.ts";
import { type CodemodeWasmModule, loadQuickJSWasm } from "../wasm.ts";
import {
type HostToWorkerMessage,
isWorkerToHostMessage,
type WorkerData,
type WorkerToHostMessage,
} from "./protocol.ts";
const DEFAULT_TIMEOUT_MS = 300_000;
const IDENTIFIER = /^[A-Za-z_$][A-Za-z0-9_$]*$/;
const RESERVED_GLOBALS: ReadonlySet<string> = new Set([
"tools",
"ALL_TOOLS",
"console",
"text",
"image",
"exit",
"globalThis",
"store",
"load",
]);
function errorMessage(error: unknown): string {
return error instanceof Error ? error.message : String(error);
}
function serializeStore(store: Readonly<Record<string, unknown>> | undefined): Record<string, string> {
const serialized: Record<string, string> = {};
for (const [key, value] of Object.entries(store ?? {})) {
const json = JSON.stringify(value);
if (json !== undefined) serialized[key] = json;
}
return serialized;
}
function parseStoreWrites(json: string): CodemodeStoreWrites {
const writes: CodemodeStoreWrites = { set: {}, delete: [] };
for (const [key, value] of JSON.parse(json) as [string, string?][]) {
if (value === undefined) writes.delete.push(key);
else writes.set[key] = JSON.parse(value);
}
return writes;
}
function defaultWorkerUrl(): URL {
// `.ts` when running from source (tests, tsx), `.js` from the published dist.
return new URL(import.meta.url.endsWith(".ts") ? "./worker.ts" : "./worker.js", import.meta.url);
}
interface PendingCall {
record: CodemodeCall | undefined;
startedAt: number;
controller: AbortController;
}
interface ExecutionOptions {
code: string;
tools: ReadonlyMap<string, CodemodeTool>;
globals: ReadonlyMap<string, CodemodeTool>;
timeoutMs: number;
signal: AbortSignal | undefined;
memoryLimitBytes: number | undefined;
store: Record<string, string>;
wasm: Promise<CodemodeWasmModule>;
workerUrl: URL;
}
/**
* One script run in its own worker and QuickJS VM. A fresh worker per run keeps
* termination simple: a runaway script, including one that only spins the
* microtask queue, is killed with `terminate()` and cannot poison a later run.
*/
class Execution {
readonly promise: Promise<CodemodeResult>;
private resolveResult!: (result: CodemodeResult) => void;
private worker: Worker | undefined;
private readonly interrupt = new SharedArrayBuffer(4);
private readonly tools: ReadonlyMap<string, CodemodeTool>;
private readonly globals: ReadonlyMap<string, CodemodeTool>;
private readonly signal: AbortSignal | undefined;
private readonly timer: NodeJS.Timeout | undefined;
private readonly output: CodemodeOutputItem[] = [];
private readonly calls: CodemodeCall[] = [];
private readonly pending = new Map<number, PendingCall>();
private finished = false;
constructor(options: ExecutionOptions) {
this.promise = new Promise<CodemodeResult>((resolve) => {
this.resolveResult = resolve;
});
this.tools = options.tools;
this.globals = options.globals;
this.signal = options.signal;
if (Number.isFinite(options.timeoutMs)) {
this.timer = setTimeout(() => {
this.finish({ kind: "timeout", message: `Execution timed out after ${options.timeoutMs} ms` });
}, options.timeoutMs);
}
if (options.signal) {
if (options.signal.aborted) {
this.onAbort();
} else {
options.signal.addEventListener("abort", this.onAbort, { once: true });
}
}
options.wasm.then(
(wasm) => this.start(options, wasm),
(error: unknown) => {
this.finish({ kind: "sandbox", message: `Failed to load QuickJS: ${errorMessage(error)}` });
},
);
}
abort(message: string): Promise<CodemodeResult> {
this.finish({ kind: "aborted", message });
return this.promise;
}
private start(options: ExecutionOptions, wasm: CodemodeWasmModule): void {
if (this.finished) return;
const workerData: WorkerData = {
code: options.code,
tools: [...options.tools.values()].map((tool) => ({
name: tool.name,
jsName: toCodemodeIdentifier(tool.name),
description: tool.description ?? "",
})),
globals: [...options.globals.values()].map((global) => ({
name: global.name,
spread: global.spread === true,
})),
wasm,
memoryLimitBytes: options.memoryLimitBytes,
store: options.store,
interrupt: this.interrupt,
};
let worker: Worker;
try {
worker = new Worker(options.workerUrl, { workerData });
} catch (error) {
this.finish({ kind: "sandbox", message: `Failed to start worker: ${errorMessage(error)}` });
return;
}
this.worker = worker;
worker.on("message", (message: unknown) => this.handleMessage(message));
worker.on("error", (error: unknown) => {
this.finish({
kind: "sandbox",
name: error instanceof Error ? error.name : undefined,
message: errorMessage(error),
});
});
worker.on("exit", (code) => {
this.finish({ kind: "sandbox", message: `Worker exited with code ${code} before the script settled` });
});
}
private readonly onAbort = (): void => {
const reason: unknown = this.signal?.reason;
this.finish({ kind: "aborted", message: reason instanceof Error ? reason.message : "Execution aborted" });
};
private post(message: HostToWorkerMessage): void {
this.worker?.postMessage(message);
}
private handleMessage(message: unknown): void {
if (this.finished || !isWorkerToHostMessage(message)) return;
switch (message.type) {
case "output":
this.output.push(message.item);
break;
case "call":
void this.handleCall(message);
break;
case "done":
this.handleDone(message);
break;
case "crash":
this.finish({ kind: "sandbox", message: message.message });
break;
}
}
private handleDone(message: Extract<WorkerToHostMessage, { type: "done" }>): void {
if (!message.ok) {
const parsed = JSON.parse(message.error) as Omit<CodemodeError, "kind">;
this.finish({ kind: "script", ...parsed });
return;
}
this.finish(undefined, message.value === undefined ? undefined : JSON.parse(message.value), message.writes);
}
private async handleCall(message: Extract<WorkerToHostMessage, { type: "call" }>): Promise<void> {
const { id, name } = message;
const isTool = message.target === "tool";
const record: CodemodeCall | undefined = isTool ? { name, status: "cancelled", durationMs: 0 } : undefined;
if (record) this.calls.push(record);
const pending: PendingCall = { record, startedAt: performance.now(), controller: new AbortController() };
this.pending.set(id, pending);
let status: CodemodeCallStatus;
let reply: HostToWorkerMessage;
try {
const tool = (isTool ? this.tools : this.globals).get(name);
if (!tool) throw new Error(`Unknown ${isTool ? "tool" : "global"} "${name}"`);
const args: unknown = message.args === undefined ? undefined : JSON.parse(message.args);
const value = await tool.execute(args, { signal: pending.controller.signal });
reply = { type: "result", id, ok: true, payload: value === undefined ? undefined : JSON.stringify(value) };
status = "ok";
} catch (error) {
reply = { type: "result", id, ok: false, payload: errorMessage(error) };
status = "error";
}
// Already cancelled by finish(): the record keeps "cancelled" and the
// worker is gone or going.
if (!this.pending.delete(id)) return;
if (record) {
record.status = status;
record.durationMs = performance.now() - pending.startedAt;
}
this.post(reply);
}
private finish(error: CodemodeError | undefined, value?: unknown, writes?: string): void {
if (this.finished) return;
this.finished = true;
clearTimeout(this.timer);
this.signal?.removeEventListener("abort", this.onAbort);
const now = performance.now();
for (const pending of this.pending.values()) {
if (pending.record) pending.record.durationMs = now - pending.startedAt;
pending.controller.abort();
}
this.pending.clear();
const result: CodemodeResult = error
? { ok: false, error, output: this.output, calls: this.calls }
: {
ok: true,
value,
output: this.output,
calls: this.calls,
storeWrites: writes === undefined ? { set: {}, delete: [] } : parseStoreWrites(writes),
};
if (!this.worker) {
this.resolveResult(result);
return;
}
Atomics.store(new Int32Array(this.interrupt), 0, 1);
this.worker
.terminate()
.catch(() => undefined)
.then(() => this.resolveResult(result));
}
}
/**
* Runs JavaScript in a QuickJS VM (a separate wasm instance) inside a worker
* thread. The script sees `tools.<name>(args)` for every registered tool, `ALL_TOOLS`,
* the output helpers `text`, `image`, `exit`, and `console.*`, `store`/`load`, and the
* configured globals; nothing else (no timers, `fetch`, `process`, `require`, modules).
*
* Each `execute()` gets its own worker and VM; the sandbox only holds the tool
* table and defaults. `close()` aborts in-flight executions.
*/
export class CodemodeSandbox {
private readonly toolsByName = new Map<string, CodemodeTool>();
private readonly globalsByName = new Map<string, CodemodeTool>();
private readonly timeoutMs: number;
private readonly memoryLimitBytes: number | undefined;
private readonly wasm: CodemodeWasmModule | Promise<CodemodeWasmModule> | undefined;
private readonly workerUrl: URL;
private readonly running = new Set<Execution>();
private closed = false;
constructor(options: CodemodeSandboxOptions = {}) {
this.timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS;
this.memoryLimitBytes = options.memoryLimitBytes;
this.wasm = options.wasm;
this.workerUrl = options.workerUrl ?? defaultWorkerUrl();
for (const tool of options.tools ?? []) this.registerTool(tool);
const namespaces = new Set<string>();
for (const global of options.globals ?? []) {
const parts = global.name.split(".");
if (parts.length > 2 || !parts.every((part) => IDENTIFIER.test(part)) || RESERVED_GLOBALS.has(parts[0])) {
throw new Error(`Invalid global name "${global.name}"`);
}
if (this.globalsByName.has(global.name)) throw new Error(`Global "${global.name}" is already registered`);
if (parts.length === 2) namespaces.add(parts[0]);
this.globalsByName.set(global.name, global);
}
for (const name of namespaces) {
if (this.globalsByName.has(name)) throw new Error(`Global "${name}" conflicts with the namespace "${name}"`);
}
}
/** Throws if a tool with the same name is already registered. */
registerTool(tool: CodemodeTool): void {
if (this.toolsByName.has(tool.name)) throw new Error(`Tool "${tool.name}" is already registered`);
this.toolsByName.set(tool.name, tool);
}
unregisterTool(name: string): boolean {
return this.toolsByName.delete(name);
}
get tools(): CodemodeTool[] {
return [...this.toolsByName.values()];
}
get globals(): CodemodeTool[] {
return [...this.globalsByName.values()];
}
/**
* `code` is an async function body: `return` and top-level `await` work.
* Never rejects for script failures; those come back as `{ ok: false }`.
* The script can use `store(key, value)` and `load(key)` on `options.store`.
*/
execute(code: string, options: CodemodeExecuteOptions = {}): Promise<CodemodeResult> {
if (this.closed) return Promise.reject(new Error("Sandbox is closed"));
const execution = new Execution({
code,
tools: new Map(this.toolsByName),
globals: this.globalsByName,
timeoutMs: options.timeoutMs ?? this.timeoutMs,
signal: options.signal,
memoryLimitBytes: this.memoryLimitBytes,
store: serializeStore(options.store),
wasm: this.wasm === undefined ? loadQuickJSWasm() : Promise.resolve(this.wasm),
workerUrl: this.workerUrl,
});
this.running.add(execution);
return execution.promise.finally(() => this.running.delete(execution));
}
/** Aborts in-flight executions (they resolve with `kind: "aborted"`) and rejects new ones. */
async close(): Promise<void> {
this.closed = true;
await Promise.all([...this.running].map((execution) => execution.abort("Sandbox closed")));
}
}

View file

@ -0,0 +1,324 @@
/**
* JavaScript evaluated inside the QuickJS VM before the script runs.
*
* The VM is its own wasm instance, so nothing here guards a realm boundary.
* The prelude keeps the host bridge in a closure so the script cannot call it
* directly, and builds `tools`, `ALL_TOOLS`, the output helpers (`text`, `image`,
* `exit`, `console`), and globals on top of it. Tool arguments and results cross
* as JSON strings and are parsed on this side.
*
* Output helpers: `text(value)` appends a text item (non-strings are
* JSON-stringified), `image(urlOrItem)` appends an image from a base64 `data:` URL, an
* `{ image_url }` object, or an MCP `ImageContent` block, and `exit()` ends the script
* successfully. `console.*` appends text items like `text()`.
*
* `store(key, value)` and `load(key)` are synchronous: they work on a snapshot of
* JSON text passed in as `storeJson`, and the keys the script wrote are reported
* with a successful "done".
*
* Evaluates to a function `(bridge, toolsJson, globalsJson, storeJson) => { settle, run, stalled }`.
* `toolsJson` lists `{ name, jsName, description }`: `tools[jsName]` and `tools[name]` call the
* tool, and `ALL_TOOLS` lists `{ name: jsName, description }`.
* `stalled()` reports a script that has not finished while no host call is pending: with no timers
* or I/O in the VM, nothing can ever resume it.
* `globalsJson` lists `{ name, spread }`; `a.b` names are grouped into a frozen `a` object.
* `bridge(kind, a, b, c)` with kind "call" or "global" (id, name, argsJson),
* "output" ("text", text) or ("image", data, mimeType), or "done" (ok, valueJsonOrErrorJson, writesJson).
*/
export const MAX_STORE_VALUE_CHARS = 256 * 1024;
export const MAX_STORE_TOTAL_CHARS = 1024 * 1024;
const IMAGE_HELPER_EXPECTS =
"image expects a non-empty image URL string, an object with image_url, or a raw MCP image block";
export const PRELUDE_SOURCE: string = `(function (bridge, toolsJson, globalsJson, storeJson) {
"use strict";
const stringify = JSON.stringify;
const parse = JSON.parse;
const promiseThen = Promise.prototype.then;
const ErrorCtor = Error;
const TypeErrorCtor = TypeError;
const pending = new Map();
let nextId = 1;
let finished = false;
// Thrown by exit() to unwind the script after it already reported success.
const EXIT = Object.freeze({});
function done(ok, payload, writes) {
if (finished) return;
finished = true;
bridge("done", ok, payload, writes);
}
function serialize(value) {
return value === undefined ? undefined : stringify(value);
}
// QuickJS stacks list frames only. Prefix "Name: message" like V8 so the
// text reads the same as a Node error, and drop this prelude's frames.
function errorText(error) {
const head = error.message ? error.name + ": " + error.message : String(error.name);
const frames =
typeof error.stack === "string"
? error.stack.split("\\n").filter((line) => line.trim() && !line.includes("codemode-prelude.js"))
: [];
return [head, ...frames].join("\\n");
}
function format(value) {
if (typeof value === "string") return value;
if (value instanceof ErrorCtor) return errorText(value);
try {
const json = stringify(value);
return json === undefined ? String(value) : json;
} catch {
return String(value);
}
}
function describeError(error) {
if (error instanceof ErrorCtor) {
return stringify({ name: error.name, message: error.message, stack: errorText(error) });
}
return stringify({ message: format(error) });
}
function caller(kind, name, spread) {
return (...args) =>
new Promise((resolve, reject) => {
let json;
try {
json = serialize(spread ? args : args[0]);
} catch (error) {
reject(error);
return;
}
const id = nextId++;
pending.set(id, { resolve, reject });
bridge(kind, id, name, json);
});
}
const tools = Object.create(null);
const allTools = [];
for (const { name, jsName, description } of parse(toolsJson)) {
const fn = caller("call", name);
// The first tool wins when two names normalize to the same identifier.
if (!(jsName in tools)) {
tools[jsName] = fn;
allTools.push(Object.freeze({ name: jsName, description }));
}
if (!(name in tools)) tools[name] = fn;
}
Object.freeze(tools);
Object.freeze(allTools);
const namespaces = new Map();
for (const { name, spread } of parse(globalsJson)) {
const fn = caller("global", name, spread);
const dot = name.indexOf(".");
if (dot === -1) {
Object.defineProperty(globalThis, name, { value: fn, enumerable: true });
continue;
}
const namespace = name.slice(0, dot);
if (!namespaces.has(namespace)) namespaces.set(namespace, Object.create(null));
namespaces.get(namespace)[name.slice(dot + 1)] = fn;
}
for (const [namespace, members] of namespaces) {
Object.defineProperty(globalThis, namespace, { value: Object.freeze(members), enumerable: true });
}
// key -> JSON text. Sizes count key and JSON characters.
const stored = new Map(Object.entries(parse(storeJson)));
const writes = new Map();
let storedChars = 0;
for (const [key, json] of stored) storedChars += key.length + json.length;
function checkKey(name, key) {
if (typeof key !== "string") throw new TypeError(name + "() key must be a string");
}
function store(key, value) {
checkKey("store", key);
const previous = stored.has(key) ? key.length + stored.get(key).length : 0;
if (value === undefined) {
stored.delete(key);
storedChars -= previous;
writes.set(key, undefined);
return;
}
let json;
try {
json = stringify(value);
} catch (error) {
throw new TypeError("store(" + stringify(key) + ") value is not JSON-serializable: " + format(error));
}
if (json === undefined) {
throw new TypeError("store(" + stringify(key) + ") value is not JSON-serializable");
}
if (json.length > ${MAX_STORE_VALUE_CHARS}) {
throw new RangeError("store(" + stringify(key) + ") value exceeds ${MAX_STORE_VALUE_CHARS} characters of JSON");
}
const next = storedChars - previous + key.length + json.length;
if (next > ${MAX_STORE_TOTAL_CHARS}) {
throw new RangeError("store is full: stored values would exceed ${MAX_STORE_TOTAL_CHARS} characters of JSON");
}
stored.set(key, json);
storedChars = next;
writes.set(key, json);
}
function load(key) {
checkKey("load", key);
const json = stored.get(key);
return json === undefined ? undefined : parse(json);
}
function serializeWrites() {
const entries = [];
for (const [key, json] of writes) entries.push(json === undefined ? [key] : [key, json]);
return stringify(entries);
}
Object.defineProperty(globalThis, "store", { value: store, enumerable: true });
Object.defineProperty(globalThis, "load", { value: load, enumerable: true });
// Primitives become their string form, everything else JSON.
function outputText(value) {
if (value === undefined || value === null || typeof value !== "object" && typeof value !== "function") {
return String(value);
}
const json = stringify(value);
return json === undefined ? String(value) : json;
}
function text(value) {
let rendered;
try {
rendered = outputText(value);
} catch (error) {
throw new TypeErrorCtor(error instanceof ErrorCtor ? error.message : String(error));
}
if (!finished) bridge("output", "text", rendered);
}
function imageUrl(value) {
if (typeof value === "string") return value;
if (typeof value !== "object" || value === null || Array.isArray(value)) {
throw new TypeErrorCtor(${JSON.stringify(IMAGE_HELPER_EXPECTS)});
}
if (value.image_url !== undefined) {
if (typeof value.image_url !== "string") throw new TypeErrorCtor(${JSON.stringify(IMAGE_HELPER_EXPECTS)});
return value.image_url;
}
if (typeof value.type !== "string") throw new TypeErrorCtor(${JSON.stringify(IMAGE_HELPER_EXPECTS)});
if (value.type !== "image") {
throw new TypeErrorCtor('image only accepts MCP image blocks, got "' + value.type + '"');
}
if (typeof value.data !== "string" || value.data === "") throw new TypeErrorCtor("image expected MCP image data");
if (value.data.toLowerCase().startsWith("data:")) return value.data;
const mimeType = typeof value.mimeType === "string" && value.mimeType ? value.mimeType : "application/octet-stream";
return "data:" + mimeType + ";base64," + value.data;
}
function image(value) {
const url = imageUrl(value);
if (url === "") throw new TypeErrorCtor(${JSON.stringify(IMAGE_HELPER_EXPECTS)});
const colon = url.indexOf(":");
const scheme = colon === -1 ? "" : url.slice(0, colon).toLowerCase();
if (scheme === "http" || scheme === "https") {
throw new TypeErrorCtor("remote image URLs are not supported in tool outputs. Pass a base64 data URI instead");
}
const comma = url.indexOf(",");
const header = comma === -1 ? [] : url.slice(colon + 1, comma).split(";");
if (scheme !== "data" || comma === -1 || header.slice(1).every((part) => part.toLowerCase() !== "base64")) {
throw new TypeErrorCtor("invalid image output. Pass a base64 data URI instead");
}
if (!finished) bridge("output", "image", url.slice(comma + 1), header[0] || "application/octet-stream");
}
function exit() {
let writesJson;
try {
writesJson = serializeWrites();
} catch (error) {
done(false, describeError(error));
throw EXIT;
}
done(true, undefined, writesJson);
throw EXIT;
}
const console = {};
for (const level of ["log", "info", "warn", "error", "debug"]) {
console[level] = (...args) => {
if (!finished) bridge("output", "text", args.map(format).join(" "));
};
}
Object.freeze(console);
Object.defineProperty(globalThis, "tools", { value: tools, enumerable: true });
Object.defineProperty(globalThis, "ALL_TOOLS", { value: allTools, enumerable: true });
Object.defineProperty(globalThis, "console", { value: console, enumerable: true });
Object.defineProperty(globalThis, "text", { value: text, enumerable: true });
Object.defineProperty(globalThis, "image", { value: image, enumerable: true });
Object.defineProperty(globalThis, "exit", { value: exit, enumerable: true });
return {
settle(id, ok, payload) {
const entry = pending.get(id);
if (!entry) return;
pending.delete(id);
if (!ok) {
entry.reject(new ErrorCtor(payload));
return;
}
let value;
try {
value = payload === undefined ? undefined : parse(payload);
} catch (error) {
entry.reject(error);
return;
}
entry.resolve(value);
},
run(fn) {
let promise;
try {
promise = fn(tools, console);
} catch (error) {
done(false, describeError(error));
return;
}
promiseThen.call(
promise,
(value) => {
let json;
try {
json = serialize(value);
} catch (error) {
done(false, describeError(error));
return;
}
done(true, json, serializeWrites());
},
(error) => {
done(false, describeError(error));
},
);
},
stalled() {
if (finished || pending.size > 0) return false;
done(
false,
stringify({
name: "Error",
message:
"The script is waiting on a promise that can never settle: no tool call is pending, and timers do not exist here.",
}),
);
return true;
},
};
})`;

View file

@ -0,0 +1,51 @@
import type { CodemodeOutputItem } from "../types.ts";
import type { CodemodeWasmModule } from "../wasm.ts";
/**
* Messages between the host (main thread) and the worker. Tool arguments,
* results, and values cross as JSON strings: the worker passes them into and
* out of the QuickJS VM as strings and never builds structured values itself.
*/
export interface WorkerData {
code: string;
/** `jsName` is the identifier the script uses; `description` is listed in `ALL_TOOLS`. */
tools: { name: string; jsName: string; description: string }[];
globals: { name: string; spread: boolean }[];
/** Compiled `quickjs-wasi` module. Structured clone shares the compiled code with the worker. */
wasm: CodemodeWasmModule;
memoryLimitBytes: number | undefined;
/** Snapshot for `load()`: key to JSON text. */
store: Record<string, string>;
/**
* One Int32 the host sets to non-zero before terminating the worker. The VM's interrupt handler
* polls it, because Bun's `worker.terminate()` cannot stop a thread that is spinning in wasm.
*/
interrupt: SharedArrayBuffer;
}
/** JSON-encoded `{ name?, message, stack? }` of an error thrown by the script. */
export type ScriptErrorJson = string;
export type WorkerToHostMessage =
| { type: "call"; id: number; target: "tool" | "global"; name: string; args: string | undefined }
| { type: "output"; item: CodemodeOutputItem }
/** `writes` is a JSON array of `[key, json]` for `store()` and `[key]` for deletions. */
| { type: "done"; ok: true; value: string | undefined; writes: string }
| { type: "done"; ok: false; error: ScriptErrorJson }
/** The VM failed outside the script's control, for example a wasm trap. */
| { type: "crash"; message: string };
export type HostToWorkerMessage =
/** `payload` is the JSON result when `ok`, otherwise the error message. */
{ type: "result"; id: number; ok: boolean; payload: string | undefined };
export function isWorkerToHostMessage(value: unknown): value is WorkerToHostMessage {
if (typeof value !== "object" || value === null) return false;
const type = (value as { type?: unknown }).type;
return type === "call" || type === "output" || type === "done" || type === "crash";
}
export function isHostToWorkerMessage(value: unknown): value is HostToWorkerMessage {
return typeof value === "object" && value !== null && (value as { type?: unknown }).type === "result";
}

View file

@ -0,0 +1,159 @@
/**
* Worker thread entry. One worker runs one script inside a fresh QuickJS VM
* (a separate wasm instance), relays tool calls and output to the host, and
* reports the result. The host terminates the worker when the script settles,
* times out, or is aborted; the worker exists so that a spinning script never
* blocks the host thread.
*
* Importing this module starts the worker. Hosts that bundle their code (for
* example a Bun compiled executable) add a file that imports
* `@earendil-works/pi-codemode/worker` as a separate entrypoint and pass its URL
* as `workerUrl`.
*/
import { parentPort, workerData } from "node:worker_threads";
import { JSException, type JSValueHandle, MAX_STACK_SIZE, QuickJS } from "quickjs-wasi";
import { PRELUDE_SOURCE } from "./prelude-source.ts";
import { isHostToWorkerMessage, type WorkerData, type WorkerToHostMessage } from "./protocol.ts";
function post(message: WorkerToHostMessage): void {
parentPort?.postMessage(message);
}
function crash(error: unknown): void {
post({ type: "crash", message: error instanceof Error ? `${error.name}: ${error.message}` : String(error) });
}
/**
* QuickJS writes engine diagnostics to fd 1 and 2, which the default shim
* forwards to the host's stdout and stderr. That output belongs to the host
* application (for example a TUI), so it is discarded. Reporting every byte as
* written keeps libc from retrying.
*/
function discardOutput(memory: { readonly buffer: ArrayBufferLike }) {
return {
fd_write(_fd: number, iovsPtr: number, iovsLen: number, nwrittenPtr: number): number {
const view = new DataView(memory.buffer);
let written = 0;
for (let i = 0; i < iovsLen; i++) {
written += view.getUint32(iovsPtr + i * 8 + 4, true);
}
view.setUint32(nwrittenPtr, written, true);
return 0;
},
};
}
function describeException(error: JSException): string {
const head = error.message ? `${error.name}: ${error.message}` : error.name;
const stack = error.stack?.trimEnd();
return JSON.stringify({ name: error.name, message: error.message, stack: stack ? `${head}\n${stack}` : head });
}
async function main(data: WorkerData): Promise<void> {
const interrupt = new Int32Array(data.interrupt);
const vm = await QuickJS.create({
wasm: data.wasm,
memoryLimit: data.memoryLimitBytes,
// Without a guard, deep recursion overflows the wasm stack and traps instead of throwing a
// catchable RangeError.
maxStackSize: MAX_STACK_SIZE,
interruptHandler: () => Atomics.load(interrupt, 0) !== 0,
wasi: discardOutput,
});
// Called from the prelude with primitives only.
const bridge = vm.newFunction("bridge", (kind, a, b, c) => {
switch (kind.toString()) {
case "call":
case "global":
post({
type: "call",
id: a.toNumber(),
target: kind.toString() === "call" ? "tool" : "global",
name: b.toString(),
args: c === undefined || c.isUndefined ? undefined : c.toString(),
});
break;
case "output":
post({
type: "output",
item:
a.toString() === "image"
? { type: "image", data: b.toString(), mimeType: c.toString() }
: { type: "text", text: b.toString() },
});
break;
case "done":
if (a.toBoolean()) {
post({
type: "done",
ok: true,
value: b === undefined || b.isUndefined ? undefined : b.toString(),
writes: c.toString(),
});
} else {
post({ type: "done", ok: false, error: b.toString() });
}
break;
}
return vm.undefined;
});
// The VM lives until the host terminates the worker, so these handles are never disposed.
const api = vm.withScope((scope) =>
scope.escape(
vm.callFunction(
vm.evalCode(PRELUDE_SOURCE, "codemode-prelude.js"),
vm.undefined,
bridge,
vm.newString(JSON.stringify(data.tools)),
vm.newString(JSON.stringify(data.globals)),
vm.newString(JSON.stringify(data.store)),
),
),
);
const settle = api.getProp("settle");
const run = api.getProp("run");
const stalled = api.getProp("stalled");
/** Run queued jobs, then fail a script that waits on nothing that can ever resume it. */
const drain = () => {
vm.executePendingJobs();
vm.callFunction(stalled, api).dispose();
};
parentPort?.on("message", (message: unknown) => {
if (!isHostToWorkerMessage(message)) return;
try {
vm.withScope(() => {
vm.callFunction(
settle,
api,
vm.newNumber(message.id),
message.ok ? vm.true : vm.false,
message.payload === undefined ? vm.undefined : vm.newString(message.payload),
);
});
drain();
} catch (error) {
crash(error);
}
});
// The prefix shares the first line with the script so reported line numbers
// match the script as written.
let fn: JSValueHandle;
try {
fn = vm.evalCode(`(async (tools, console) => {${data.code}\n})`, "codemode.js");
} catch (error) {
if (!(error instanceof JSException)) throw error;
post({ type: "done", ok: false, error: describeException(error) });
return;
}
vm.callFunction(run, api, fn).dispose();
fn.dispose();
drain();
}
if (parentPort) {
main(workerData as WorkerData).catch(crash);
}