fix(codemode): fail stalled scripts and expand local schema references

- A script awaiting a promise nothing can settle (no pending host call,
  no timers in the VM) now fails instead of hanging until aborted, which
  matters because pi runs scripts without a default deadline
- Declarations resolve #/$defs and #/definitions references, which MCP
  servers built on pydantic or zod emit for nested arguments

packages/codemode/README.md

@@ -42,7 +42,7 @@ await sandbox.close();
 - `store(key, value)` and `load(key)` read and write JSON values synchronously. See [Store](#store).
 - Nothing else: no timers, `fetch`, `process`, `require`, modules, or `WebAssembly`. `eval` and `Function` work but only produce more code inside the same VM.
 
-`timeoutMs: Infinity` disables the deadline; the script then runs until it settles or `signal` aborts it.
+`timeoutMs: Infinity` disables the deadline; the script then runs until it settles or `signal` aborts it. A script that waits on a promise nothing can settle (no tool call pending, and the VM has no timers or I/O) fails right away instead of hanging.
 
 `memoryLimitBytes` caps the VM's heap. Allocations beyond it fail inside the script as `InternalError: out of memory`.
 
@@ -102,7 +102,7 @@ renderDeclarations({ tools: sandbox.tools, globals: sandbox.globals });
 // };
 ```
 
-Schemas only shape the declarations; values are not validated against them.
+Schemas only shape the declarations; values are not validated against them. Local references (`#/$defs/...`, `#/definitions/...`) are expanded; recursive and remote references render as `unknown`.
 
 `execute()` never rejects for script failures. `result.error.kind` is one of:
 

packages/codemode/src/declarations.ts

@@ -13,7 +13,7 @@ export interface RenderDeclarationsOptions {
  * description. Tools become members of `declare const tools`, globals become
  * `declare function` statements, and `ns.member` globals members of `declare const ns`. Descriptions become doc comments; schemas become types
  * (`unknown` where a schema is missing or uses features TypeScript cannot express, such as
- * `$ref`).
+ * recursive `$ref`s).
  *
  * ```ts
  * declare const tools: {
@@ -94,28 +94,61 @@ function union(types: string[]): string {
 
 /**
  * Convert a JSON Schema to a TypeScript type expression. `indent` is the indentation of the line
- * the type starts on; nested object members are indented one level deeper.
+ * the type starts on; nested object members are indented one level deeper. Local references
+ * (`#/$defs/...`, `#/definitions/...`) resolve against `schema` itself; recursive references and
+ * references elsewhere render as `unknown`.
  */
 export function schemaToType(schema: CodemodeJsonSchema, indent = ""): string {
+	return toType(schema, indent, { root: schema, resolving: new Set() });
+}
+
+interface SchemaContext {
+	root: CodemodeJsonSchema;
+	/** References being expanded on the current path, to stop at recursive types. */
+	resolving: Set<string>;
+}
+
+function resolveRef(ref: string, root: CodemodeJsonSchema): CodemodeJsonSchema | undefined {
+	if (ref !== "#" && !ref.startsWith("#/")) return undefined;
+	let current: unknown = root;
+	for (const segment of ref.slice(2).split("/").filter(Boolean)) {
+		const key = decodeURIComponent(segment).replaceAll("~1", "/").replaceAll("~0", "~");
+		if (!isObject(current) || !(key in current)) return undefined;
+		current = current[key];
+	}
+	return typeof current === "boolean" || isObject(current) ? current : undefined;
+}
+
+function toType(schema: CodemodeJsonSchema, indent: string, context: SchemaContext): string {
 	if (schema === true) return "unknown";
 	if (schema === false) return "never";
 	if (!isObject(schema)) return "unknown";
-	if (typeof schema.$ref === "string") return "unknown";
+	if (typeof schema.$ref === "string") {
+		const ref = schema.$ref;
+		const target = context.resolving.has(ref) ? undefined : resolveRef(ref, context.root);
+		if (target === undefined) return "unknown";
+		context.resolving.add(ref);
+		try {
+			return toType(target, indent, context);
+		} finally {
+			context.resolving.delete(ref);
+		}
+	}
 
 	if ("const" in schema) return JSON.stringify(schema.const) ?? "unknown";
 	if (Array.isArray(schema.enum)) return union(schema.enum.map((value) => JSON.stringify(value) ?? "unknown"));
 
 	const variants = Array.isArray(schema.anyOf) ? schema.anyOf : Array.isArray(schema.oneOf) ? schema.oneOf : undefined;
-	if (variants) return union(variants.map((variant) => schemaToType(variant as CodemodeJsonSchema, indent)));
+	if (variants) return union(variants.map((variant) => toType(variant as CodemodeJsonSchema, indent, context)));
 	if (Array.isArray(schema.allOf)) {
-		const parts = schema.allOf.map((part) => schemaToType(part as CodemodeJsonSchema, indent));
+		const parts = schema.allOf.map((part) => toType(part as CodemodeJsonSchema, indent, context));
 		const meaningful = parts.filter((part) => part !== "unknown");
 		return meaningful.length === 0 ? "unknown" : meaningful.map(asElement).join(" & ");
 	}
 
 	const type = schema.type;
 	if (Array.isArray(type)) {
-		return union(type.map((entry) => schemaToType({ ...schema, type: entry }, indent)));
+		return union(type.map((entry) => toType({ ...schema, type: entry }, indent, context)));
 	}
 	switch (type) {
 		case "string":
@@ -128,33 +161,35 @@ export function schemaToType(schema: CodemodeJsonSchema, indent = ""): string {
 		case "null":
 			return "null";
 		case "array":
-			return arrayType(schema, indent);
+			return arrayType(schema, indent, context);
 		case "object":
-			return objectType(schema, indent);
+			return objectType(schema, indent, context);
 		case undefined:
-			if (isObject(schema.properties) || isObject(schema.additionalProperties)) return objectType(schema, indent);
-			if (schema.items !== undefined || schema.prefixItems !== undefined) return arrayType(schema, indent);
+			if (isObject(schema.properties) || isObject(schema.additionalProperties)) {
+				return objectType(schema, indent, context);
+			}
+			if (schema.items !== undefined || schema.prefixItems !== undefined) return arrayType(schema, indent, context);
 			return "unknown";
 		default:
 			return "unknown";
 	}
 }
 
-function arrayType(schema: Record<string, unknown>, indent: string): string {
+function arrayType(schema: Record<string, unknown>, indent: string, context: SchemaContext): string {
 	const tuple = Array.isArray(schema.prefixItems)
 		? schema.prefixItems
 		: Array.isArray(schema.items)
 			? schema.items
 			: [];
 	if (tuple.length > 0) {
-		return `[${tuple.map((item) => schemaToType(item as CodemodeJsonSchema, indent)).join(", ")}]`;
+		return `[${tuple.map((item) => toType(item as CodemodeJsonSchema, indent, context)).join(", ")}]`;
 	}
 	const items = schema.items;
 	if (items === undefined || Array.isArray(items)) return "unknown[]";
-	return `${asElement(schemaToType(items as CodemodeJsonSchema, indent))}[]`;
+	return `${asElement(toType(items as CodemodeJsonSchema, indent, context))}[]`;
 }
 
-function objectType(schema: Record<string, unknown>, indent: string): string {
+function objectType(schema: Record<string, unknown>, indent: string, context: SchemaContext): string {
 	const properties = isObject(schema.properties) ? schema.properties : {};
 	const required = new Set(Array.isArray(schema.required) ? schema.required : []);
 	const additional = schema.additionalProperties;
@@ -163,11 +198,11 @@ function objectType(schema: Record<string, unknown>, indent: string): string {
 	for (const [name, property] of Object.entries(properties)) {
 		const description = isObject(property) && typeof property.description === "string" ? property.description : "";
 		const optional = required.has(name) ? "" : "?";
-		const type = schemaToType(property as CodemodeJsonSchema, inner);
+		const type = toType(property as CodemodeJsonSchema, inner, context);
 		members.push(`${docComment(description, inner)}${inner}${propertyKey(name)}${optional}: ${type};`);
 	}
 	if (additional !== undefined && additional !== false) {
-		const type = additional === true ? "unknown" : schemaToType(additional as CodemodeJsonSchema, inner);
+		const type = additional === true ? "unknown" : toType(additional as CodemodeJsonSchema, inner, context);
 		members.push(`${inner}[key: string]: ${type};`);
 	} else if (members.length === 0 && additional === undefined) {
 		return "Record<string, unknown>";

packages/codemode/src/runtime/prelude-source.ts

@@ -10,7 +10,9 @@
  * JSON text passed in as `storeJson`, and the keys the script wrote are reported
  * with a successful "done".
  *
- * Evaluates to a function `(bridge, toolNamesJson, globalsJson, storeJson) => { settle, run }`.
+ * Evaluates to a function `(bridge, toolNamesJson, globalsJson, storeJson) => { settle, run, stalled }`.
+ * `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),
  * "log" (level, message) or "done" (ok, valueJsonOrErrorJson, writesJson).
@@ -26,6 +28,13 @@ export const PRELUDE_SOURCE: string = `(function (bridge, toolNamesJson, globals
 	const ErrorCtor = Error;
 	const pending = new Map();
 	let nextId = 1;
+	let finished = false;
+
+	function done(ok, payload, writes) {
+		if (finished) return;
+		finished = true;
+		bridge("done", ok, payload, writes);
+	}
 
 	function serialize(value) {
 		return value === undefined ? undefined : stringify(value);
@@ -187,7 +196,7 @@ export const PRELUDE_SOURCE: string = `(function (bridge, toolNamesJson, globals
 			try {
 				promise = fn(tools, console);
 			} catch (error) {
-				bridge("done", false, describeError(error));
+				done(false, describeError(error));
 				return;
 			}
 			promiseThen.call(
@@ -197,15 +206,27 @@ export const PRELUDE_SOURCE: string = `(function (bridge, toolNamesJson, globals
 					try {
 						json = serialize(value);
 					} catch (error) {
-						bridge("done", false, describeError(error));
+						done(false, describeError(error));
 						return;
 					}
-					bridge("done", true, json, serializeWrites());
+					done(true, json, serializeWrites());
 				},
 				(error) => {
-					bridge("done", false, describeError(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;
+		},
 	};
 })`;

packages/codemode/src/runtime/worker.ts

@@ -108,6 +108,12 @@ async function main(data: WorkerData): Promise<void> {
 	);
 	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;
@@ -121,7 +127,7 @@ async function main(data: WorkerData): Promise<void> {
 					message.payload === undefined ? vm.undefined : vm.newString(message.payload),
 				);
 			});
-			vm.executePendingJobs();
+			drain();
 		} catch (error) {
 			crash(error);
 		}
@@ -139,7 +145,7 @@ async function main(data: WorkerData): Promise<void> {
 	}
 	vm.callFunction(run, api, fn).dispose();
 	fn.dispose();
-	vm.executePendingJobs();
+	drain();
 }
 
 if (parentPort) {

packages/codemode/test/declarations.test.ts

@@ -17,6 +17,29 @@ describe("schemaToType", () => {
 		expect(schemaToType(false)).toBe("never");
 	});
 
+	it("resolves local references and stops at recursive ones", () => {
+		const schema = {
+			type: "object",
+			properties: {
+				item: { $ref: "#/$defs/Item" },
+				legacy: { $ref: "#/definitions/Legacy" },
+				remote: { $ref: "https://example.com/schema.json" },
+			},
+			required: ["item"],
+			$defs: {
+				Item: {
+					type: "object",
+					properties: { id: { type: "string" }, parent: { $ref: "#/$defs/Item" } },
+					required: ["id"],
+				},
+			},
+			definitions: { Legacy: { enum: ["a", "b"] } },
+		};
+		expect(schemaToType(schema)).toBe(
+			'{\n  item: {\n    id: string;\n    parent?: unknown;\n  };\n  legacy?: "a" | "b";\n  remote?: unknown;\n}',
+		);
+	});
+
 	it("renders arrays and tuples", () => {
 		expect(schemaToType({ type: "array", items: { type: "string" } })).toBe("string[]");
 		expect(schemaToType({ type: "array", items: { anyOf: [{ type: "string" }, { type: "number" }] } })).toBe(

packages/codemode/test/sandbox.test.ts

@@ -354,6 +354,20 @@ describe("limits and lifetime", () => {
 		expect(result).toMatchObject({ ok: true, value: "late" });
 	});
 
+	it("fails a script that waits on a promise nothing can settle", async () => {
+		const sandbox = createSandbox([{ name: "echo", execute: (args) => args }]);
+		const result = await sandbox.execute("await tools.echo(1); await new Promise(() => {}); return 'never'", {
+			timeoutMs: Number.POSITIVE_INFINITY,
+		});
+		expect(result).toMatchObject({
+			ok: false,
+			error: { kind: "script", message: expect.stringContaining("can never settle") },
+		});
+		// Returning while a call is still pending is not a stall.
+		const returned = await sandbox.execute("tools.echo(2); return 'early'", { timeoutMs: Number.POSITIVE_INFINITY });
+		expect(returned).toMatchObject({ ok: true, value: "early" });
+	});
+
 	it("terminates a microtask-spinning loop on timeout", async () => {
 		const sandbox = createSandbox();
 		const result = await sandbox.execute("while (true) await null", { timeoutMs: 200 });