add persisted pi workflow extension
This commit is contained in:
@@ -0,0 +1,339 @@
|
||||
import { randomBytes } from "node:crypto";
|
||||
import { spawn, type ChildProcess } from "node:child_process";
|
||||
import * as path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import { safeStringify, toSerializable } from "../_workflows/serialization.ts";
|
||||
|
||||
const MAX_SOURCE_BYTES = 512 * 1024;
|
||||
const MAX_ARGS_BYTES = 256 * 1024;
|
||||
const MAX_RESULT_BYTES = 1024 * 1024;
|
||||
const MAX_ASSIGNMENT_MESSAGE_BYTES = 512 * 1024;
|
||||
const MAX_ASSIGNMENT_REQUESTS = 32;
|
||||
const MAX_STAGE_UPDATES = 64;
|
||||
const DEFAULT_SCRIPT_RUNTIME_MS = 10 * 60 * 1000;
|
||||
|
||||
export interface ScriptAssignmentResult {
|
||||
ok: boolean;
|
||||
output: string;
|
||||
structured?: unknown;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
export interface RunScriptSandboxOptions {
|
||||
source: string;
|
||||
args: unknown;
|
||||
cwd: string;
|
||||
signal: AbortSignal;
|
||||
runtimeTimeoutMs?: number;
|
||||
onAssign: (
|
||||
role: string,
|
||||
prompt: string,
|
||||
signal: AbortSignal,
|
||||
requestId: number,
|
||||
) => Promise<ScriptAssignmentResult>;
|
||||
onStage: (title: string) => void;
|
||||
}
|
||||
|
||||
function byteLength(value: string) {
|
||||
return Buffer.byteLength(value, "utf8");
|
||||
}
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === "object" && value !== null && !Array.isArray(value);
|
||||
}
|
||||
|
||||
function errorText(error: unknown) {
|
||||
return error instanceof Error ? error.message : String(error);
|
||||
}
|
||||
|
||||
function terminateChild(child: ChildProcess) {
|
||||
if (child.exitCode !== null || child.signalCode !== null) return;
|
||||
child.kill("SIGTERM");
|
||||
const force = setTimeout(() => {
|
||||
if (child.exitCode === null && child.signalCode === null)
|
||||
child.kill("SIGKILL");
|
||||
}, 1_000);
|
||||
force.unref?.();
|
||||
}
|
||||
|
||||
export function validateScriptIpcEnvelope(raw: unknown, token: string) {
|
||||
if (!isRecord(raw) || raw.token !== token || typeof raw.kind !== "string") {
|
||||
throw new Error("Script sandbox sent an invalid IPC message");
|
||||
}
|
||||
return raw;
|
||||
}
|
||||
|
||||
export function parseScriptAssignmentPayload(payloadJson: unknown) {
|
||||
if (
|
||||
typeof payloadJson !== "string" ||
|
||||
byteLength(payloadJson) > MAX_ASSIGNMENT_MESSAGE_BYTES
|
||||
) {
|
||||
throw new Error("Script sandbox sent an oversized assignment request");
|
||||
}
|
||||
let payload: unknown;
|
||||
try {
|
||||
payload = JSON.parse(payloadJson);
|
||||
} catch {
|
||||
throw new Error("Script sandbox sent malformed assignment JSON");
|
||||
}
|
||||
if (
|
||||
!isRecord(payload) ||
|
||||
!Number.isSafeInteger(payload.id) ||
|
||||
typeof payload.id !== "number" ||
|
||||
payload.id < 1 ||
|
||||
typeof payload.role !== "string" ||
|
||||
payload.role.trim().length === 0 ||
|
||||
payload.role.length > 80 ||
|
||||
typeof payload.prompt !== "string" ||
|
||||
payload.prompt.length > 100_000
|
||||
) {
|
||||
throw new Error("Script sandbox sent an invalid assignment request");
|
||||
}
|
||||
return {
|
||||
id: payload.id,
|
||||
role: payload.role.trim(),
|
||||
prompt: payload.prompt,
|
||||
};
|
||||
}
|
||||
|
||||
export class ScriptRequestBudget {
|
||||
private readonly ids = new Set<number>();
|
||||
private count = 0;
|
||||
|
||||
accept(id: number) {
|
||||
if (this.ids.has(id)) {
|
||||
throw new Error("Script sandbox sent a duplicate assignment request");
|
||||
}
|
||||
if (++this.count > MAX_ASSIGNMENT_REQUESTS) {
|
||||
throw new Error("Script sandbox exceeded its assignment request budget");
|
||||
}
|
||||
this.ids.add(id);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute orchestration code in a separate, permission-restricted Node process.
|
||||
* The child can only invoke the narrow agent/phase IPC protocol and is always
|
||||
* terminated on completion, cancellation, or protocol failure. The workflow
|
||||
* itself and its assignment requests have no wall-clock deadline. Active requests
|
||||
* are aborted only when the workflow is cancelled or the sandbox is cleaned up.
|
||||
*/
|
||||
export function runScriptSandbox(options: RunScriptSandboxOptions) {
|
||||
if (!process.allowedNodeEnvironmentFlags.has("--permission")) {
|
||||
return Promise.reject(
|
||||
new Error("This Node runtime cannot enforce script child permissions"),
|
||||
);
|
||||
}
|
||||
if (byteLength(options.source) > MAX_SOURCE_BYTES) {
|
||||
return Promise.reject(
|
||||
new Error(`Script source exceeds the ${MAX_SOURCE_BYTES} byte limit`),
|
||||
);
|
||||
}
|
||||
|
||||
try {
|
||||
const exactArgsJson = JSON.stringify({
|
||||
defined: options.args !== undefined,
|
||||
value: options.args,
|
||||
});
|
||||
if (exactArgsJson && byteLength(exactArgsJson) > MAX_ARGS_BYTES) {
|
||||
return Promise.reject(new Error("Script args exceed the IPC limit"));
|
||||
}
|
||||
} catch {
|
||||
// The bounded serializer below handles cycles and hostile values.
|
||||
}
|
||||
|
||||
const argsJson = safeStringify(
|
||||
{ defined: options.args !== undefined, value: options.args },
|
||||
{ maxBytes: MAX_ARGS_BYTES + 1, maxDepth: 16, maxNodes: 10_000 },
|
||||
);
|
||||
if (byteLength(argsJson) > MAX_ARGS_BYTES) {
|
||||
return Promise.reject(new Error("Script args exceed the IPC limit"));
|
||||
}
|
||||
|
||||
return new Promise<unknown>((resolve, reject) => {
|
||||
const workerPath = fileURLToPath(
|
||||
new URL("./sandbox-child.cjs", import.meta.url),
|
||||
);
|
||||
const child = spawn(
|
||||
process.execPath,
|
||||
[
|
||||
"--permission",
|
||||
`--allow-fs-read=${path.dirname(workerPath)}`,
|
||||
"--max-old-space-size=128",
|
||||
"--stack-size=2048",
|
||||
workerPath,
|
||||
],
|
||||
{
|
||||
cwd: options.cwd,
|
||||
env: {
|
||||
PATH: process.env.PATH ?? "",
|
||||
NODE_NO_WARNINGS: "1",
|
||||
},
|
||||
stdio: ["ignore", "ignore", "ignore", "ipc"],
|
||||
},
|
||||
);
|
||||
const token = randomBytes(24).toString("hex");
|
||||
const requestBudget = new ScriptRequestBudget();
|
||||
const activeAssignmentRequests = new Map<number, AbortController>();
|
||||
let stageCount = 0;
|
||||
let finished = false;
|
||||
const runtimeTimer = setTimeout(
|
||||
() => finish(new Error("Script exceeded its wall-clock time limit")),
|
||||
options.runtimeTimeoutMs ?? DEFAULT_SCRIPT_RUNTIME_MS,
|
||||
);
|
||||
runtimeTimer.unref?.();
|
||||
|
||||
const cleanup = () => {
|
||||
for (const abortController of activeAssignmentRequests.values()) {
|
||||
abortController.abort(new Error("Script stopped"));
|
||||
}
|
||||
activeAssignmentRequests.clear();
|
||||
clearTimeout(runtimeTimer);
|
||||
options.signal.removeEventListener("abort", onAbort);
|
||||
child.removeAllListeners("message");
|
||||
child.removeAllListeners("error");
|
||||
child.removeAllListeners("exit");
|
||||
terminateChild(child);
|
||||
};
|
||||
const finish = (error?: Error, value?: unknown) => {
|
||||
if (finished) return;
|
||||
finished = true;
|
||||
cleanup();
|
||||
if (error) reject(error);
|
||||
else resolve(value);
|
||||
};
|
||||
const onAbort = () => finish(new Error("Script was aborted"));
|
||||
|
||||
options.signal.addEventListener("abort", onAbort, { once: true });
|
||||
if (options.signal.aborted) {
|
||||
onAbort();
|
||||
return;
|
||||
}
|
||||
|
||||
child.on("error", (error) => finish(error));
|
||||
child.on("exit", (code, exitSignal) => {
|
||||
if (!finished) {
|
||||
finish(
|
||||
new Error(
|
||||
`Script sandbox exited before completion (${exitSignal ?? code ?? "unknown"})`,
|
||||
),
|
||||
);
|
||||
}
|
||||
});
|
||||
child.on("message", (message: unknown) => {
|
||||
let raw: Record<string, unknown>;
|
||||
try {
|
||||
raw = validateScriptIpcEnvelope(message, token);
|
||||
} catch (error) {
|
||||
finish(error instanceof Error ? error : new Error(String(error)));
|
||||
return;
|
||||
}
|
||||
if (raw.kind === "stage") {
|
||||
if (++stageCount > MAX_STAGE_UPDATES) {
|
||||
finish(new Error("Script sandbox exceeded its stage update budget"));
|
||||
return;
|
||||
}
|
||||
if (
|
||||
typeof raw.payloadJson !== "string" ||
|
||||
byteLength(raw.payloadJson) > 4096
|
||||
) {
|
||||
finish(new Error("Script sandbox sent an invalid stage update"));
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const payload: unknown = JSON.parse(raw.payloadJson);
|
||||
if (!isRecord(payload) || typeof payload.title !== "string") {
|
||||
throw new Error("invalid title");
|
||||
}
|
||||
options.onStage(payload.title.slice(0, 160));
|
||||
} catch {
|
||||
finish(new Error("Script sandbox sent an invalid stage update"));
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (raw.kind === "assign") {
|
||||
let payload: ReturnType<typeof parseScriptAssignmentPayload>;
|
||||
try {
|
||||
payload = parseScriptAssignmentPayload(raw.payloadJson);
|
||||
} catch (error) {
|
||||
finish(error instanceof Error ? error : new Error(String(error)));
|
||||
return;
|
||||
}
|
||||
try {
|
||||
requestBudget.accept(payload.id);
|
||||
} catch (error) {
|
||||
finish(error instanceof Error ? error : new Error(String(error)));
|
||||
return;
|
||||
}
|
||||
const id = payload.id;
|
||||
const abortController = new AbortController();
|
||||
const sendResult = (result: ScriptAssignmentResult) => {
|
||||
if (!activeAssignmentRequests.delete(id)) return;
|
||||
if (finished || !child.connected) return;
|
||||
const normalized = toSerializable(result, {
|
||||
maxDepth: 16,
|
||||
maxNodes: 10_000,
|
||||
maxStringBytes: 128 * 1024,
|
||||
});
|
||||
let resultJson = JSON.stringify(normalized);
|
||||
if (byteLength(resultJson) > MAX_ASSIGNMENT_MESSAGE_BYTES) {
|
||||
resultJson = JSON.stringify({
|
||||
ok: false,
|
||||
output: "",
|
||||
error: "Assignment result exceeded the script IPC output limit",
|
||||
});
|
||||
}
|
||||
child.send({ token, kind: "assignmentResult", id, resultJson });
|
||||
};
|
||||
activeAssignmentRequests.set(id, abortController);
|
||||
void options
|
||||
.onAssign(
|
||||
payload.role.trim(),
|
||||
payload.prompt,
|
||||
abortController.signal,
|
||||
payload.id,
|
||||
)
|
||||
.then(sendResult)
|
||||
.catch((error) =>
|
||||
sendResult({ ok: false, output: "", error: errorText(error) }),
|
||||
);
|
||||
return;
|
||||
}
|
||||
if (raw.kind === "result") {
|
||||
if (
|
||||
typeof raw.resultJson !== "string" ||
|
||||
byteLength(raw.resultJson) > MAX_RESULT_BYTES
|
||||
) {
|
||||
finish(new Error("Script result exceeded the IPC limit"));
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const normalized = toSerializable(JSON.parse(raw.resultJson));
|
||||
finish(undefined, JSON.parse(JSON.stringify(normalized)));
|
||||
} catch (error) {
|
||||
finish(
|
||||
new Error(`Script returned invalid JSON: ${errorText(error)}`),
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (raw.kind === "error" && typeof raw.error === "string") {
|
||||
finish(new Error(raw.error.slice(0, 16 * 1024)));
|
||||
return;
|
||||
}
|
||||
finish(new Error("Script sandbox sent an unknown IPC message"));
|
||||
});
|
||||
|
||||
child.send(
|
||||
{
|
||||
kind: "init",
|
||||
token,
|
||||
source: options.source,
|
||||
argsJson,
|
||||
},
|
||||
(error) => {
|
||||
if (error) finish(error);
|
||||
},
|
||||
);
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user