poe-code 8.0.0 → 9.0.0

package/packages/safejs/dist/modules/mcp-transport.js

@@ -0,0 +1,259 @@
+import { spawn } from "node:child_process";
+import path from "node:path";
+import { HttpTransport, McpClient, StdioTransport } from "tiny-mcp-client";
+export function parseMcpConfig(source, directory) {
+    const parsed = JSON.parse(source);
+    const options = normalizeMcpOptions(parsed, [
+        "servers",
+        "requestTimeoutMs",
+        "closeTimeoutMs",
+        "maxToolPages"
+    ]);
+    for (const server of Object.values(options.servers)) {
+        if ("command" in server) {
+            server.cwd = path.resolve(directory, server.cwd ?? ".");
+            if (server.command.includes("/") || server.command.includes("\\")) {
+                server.command = path.resolve(directory, server.command);
+            }
+        }
+    }
+    return options;
+}
+export function normalizeMcpOptions(value, keys = [
+    "servers",
+    "requestTimeoutMs",
+    "closeTimeoutMs",
+    "maxToolPages",
+    "signal",
+    "fetch",
+    "spawn"
+]) {
+    const options = readRecord(value, "MCP options", keys);
+    const servers = readRecord(options.servers, "MCP servers");
+    const normalized = Object.create(null);
+    for (const [name, input] of Object.entries(servers)) {
+        if (name.length === 0 || name.trim() !== name)
+            throw new TypeError("MCP server names must be non-empty without surrounding whitespace.");
+        const server = readRecord(input, `MCP server ${name}`);
+        if (Object.hasOwn(server, "command")) {
+            readRecord(server, `MCP server ${name}`, ["command", "args", "cwd", "env"]);
+            const command = readText(server.command, "MCP command");
+            if (server.args !== undefined &&
+                (!Array.isArray(server.args) ||
+                    server.args.some((argument) => typeof argument !== "string")))
+                throw new TypeError("MCP args must be an array of strings.");
+            normalized[name] = {
+                command,
+                ...(server.args === undefined ? {} : { args: [...server.args] }),
+                ...(server.cwd === undefined ? {} : { cwd: readText(server.cwd, "MCP cwd") }),
+                env: readStrings(server.env ?? {}, "MCP environment")
+            };
+        }
+        else {
+            readRecord(server, `MCP server ${name}`, ["url", "headers"]);
+            const url = new URL(readText(server.url, "MCP URL"));
+            if (!["http:", "https:"].includes(url.protocol) || url.username || url.password || url.hash)
+                throw new TypeError("MCP URL must use HTTP or HTTPS without credentials or a fragment.");
+            const headers = readStrings(server.headers ?? {}, "MCP headers");
+            new Headers(headers);
+            normalized[name] = { url: url.href, headers };
+        }
+    }
+    for (const name of ["fetch", "spawn"])
+        if (options[name] !== undefined && typeof options[name] !== "function")
+            throw new TypeError(`MCP ${name} must be a function.`);
+    if (options.signal !== undefined && !(options.signal instanceof AbortSignal))
+        throw new TypeError("MCP signal must be an AbortSignal.");
+    return {
+        servers: normalized,
+        requestTimeoutMs: readTimeout(options.requestTimeoutMs, 30_000),
+        closeTimeoutMs: readTimeout(options.closeTimeoutMs, 1_000),
+        maxToolPages: readTimeout(options.maxToolPages, 100),
+        ...(options.signal === undefined ? {} : { signal: options.signal }),
+        ...(options.fetch === undefined ? {} : { fetch: options.fetch }),
+        ...(options.spawn === undefined ? {} : { spawn: options.spawn })
+    };
+}
+function readRecord(value, label, keys) {
+    if (typeof value !== "object" ||
+        value === null ||
+        Array.isArray(value) ||
+        ![Object.prototype, null].includes(Object.getPrototypeOf(value)))
+        throw new TypeError(`${label} must be a plain object.`);
+    const result = Object.create(null);
+    for (const [key, descriptor] of Object.entries(Object.getOwnPropertyDescriptors(value))) {
+        if (!("value" in descriptor) || (keys !== undefined && !keys.includes(key)))
+            throw new TypeError(`${label} contains unsupported field ${key}.`);
+        result[key] = descriptor.value;
+    }
+    return result;
+}
+function readText(value, label) {
+    if (typeof value !== "string" || value.trim().length === 0 || value.includes("\0"))
+        throw new TypeError(`${label} must be a non-empty string without NUL characters.`);
+    return value;
+}
+function readStrings(value, label) {
+    const entries = readRecord(value, label);
+    for (const [name, entry] of Object.entries(entries))
+        if (typeof entry !== "string" || name.includes("\0") || entry.includes("\0"))
+            throw new TypeError(`${label} must contain strings without NUL characters.`);
+    return entries;
+}
+function readTimeout(value, fallback) {
+    if (value === undefined)
+        return fallback;
+    if (typeof value !== "number" ||
+        !Number.isSafeInteger(value) ||
+        value <= 0 ||
+        value > 2_147_483_647)
+        throw new TypeError("MCP timeouts must be positive 32-bit integers.");
+    return value;
+}
+export function connectMcpTransport(server, options, signal) {
+    signal?.throwIfAborted();
+    options.signal?.throwIfAborted();
+    const client = new McpClient({
+        clientInfo: { name: "safejs", version: "1" },
+        requestTimeoutMs: options.requestTimeoutMs
+    });
+    let child;
+    const activeRequests = new Set();
+    const transport = "command" in server
+        ? new StdioTransport({
+            ...server,
+            spawn(command, args, spawnOptions) {
+                child =
+                    options.spawn === undefined
+                        ? spawn(command, args, spawnOptions)
+                        : options.spawn(command, args, spawnOptions);
+                return child;
+            }
+        })
+        : new HttpTransport({
+            ...server,
+            async fetch(input, init) {
+                const controller = new AbortController();
+                activeRequests.add(controller);
+                const cancel = () => controller.abort(init?.signal?.reason);
+                init?.signal?.addEventListener("abort", cancel, { once: true });
+                if (init?.signal?.aborted)
+                    cancel();
+                const timeout = setTimeout(() => controller.abort(new Error("MCP HTTP request timed out.")), init?.method === "DELETE" ? options.closeTimeoutMs : options.requestTimeoutMs);
+                const finish = () => {
+                    clearTimeout(timeout);
+                    activeRequests.delete(controller);
+                    init?.signal?.removeEventListener("abort", cancel);
+                };
+                try {
+                    const response = await (options.fetch ?? fetch)(input, {
+                        ...init,
+                        signal: controller.signal,
+                        redirect: "error"
+                    });
+                    if (response.body === null) {
+                        finish();
+                        return response;
+                    }
+                    const reader = response.body.getReader();
+                    let stopReading;
+                    const release = () => {
+                        finish();
+                        if (stopReading !== undefined)
+                            controller.signal.removeEventListener("abort", stopReading);
+                    };
+                    const body = new ReadableStream({
+                        start(stream) {
+                            stopReading = () => {
+                                void reader.cancel(controller.signal.reason).catch(() => undefined);
+                                stream.error(controller.signal.reason);
+                                release();
+                            };
+                            controller.signal.addEventListener("abort", stopReading, { once: true });
+                            if (controller.signal.aborted)
+                                stopReading();
+                        },
+                        async pull(stream) {
+                            try {
+                                const result = await reader.read();
+                                if (controller.signal.aborted)
+                                    return;
+                                if (result.done) {
+                                    release();
+                                    stream.close();
+                                }
+                                else
+                                    stream.enqueue(result.value);
+                            }
+                            catch (error) {
+                                release();
+                                stream.error(error);
+                            }
+                        },
+                        async cancel(reason) {
+                            release();
+                            await reader.cancel(reason);
+                        }
+                    });
+                    return new Response(body, {
+                        status: response.status,
+                        statusText: response.statusText,
+                        headers: response.headers
+                    });
+                }
+                catch (error) {
+                    finish();
+                    throw error;
+                }
+            }
+        });
+    let closing;
+    const cancel = () => {
+        void close().catch(() => undefined);
+    };
+    function close() {
+        closing ??= (async () => {
+            signal?.removeEventListener("abort", cancel);
+            options.signal?.removeEventListener("abort", cancel);
+            for (const request of activeRequests)
+                request.abort(new Error("MCP connection closed."));
+            await client.close();
+            transport.dispose();
+            await new Promise((resolve, reject) => {
+                let forced;
+                const timeout = setTimeout(() => {
+                    if (child !== undefined && child.exitCode === null && child.signalCode === null)
+                        child.kill("SIGKILL");
+                    forced = setTimeout(() => reject(new Error("MCP transport did not close.")), options.closeTimeoutMs);
+                }, options.closeTimeoutMs);
+                void transport.closed.then(() => {
+                    clearTimeout(timeout);
+                    clearTimeout(forced);
+                    resolve();
+                }, (error) => {
+                    clearTimeout(timeout);
+                    clearTimeout(forced);
+                    reject(error);
+                });
+            });
+        })();
+        return closing;
+    }
+    signal?.addEventListener("abort", cancel, { once: true });
+    options.signal?.addEventListener("abort", cancel, { once: true });
+    const aborted = signal?.aborted || options.signal?.aborted;
+    if (aborted)
+        cancel();
+    const connecting = aborted
+        ? Promise.reject(signal?.aborted ? signal.reason : options.signal?.reason)
+        : client.connect(transport);
+    const ready = connecting.then(() => client, async (error) => {
+        await close().catch(() => undefined);
+        throw signal?.aborted
+            ? signal.reason
+            : options.signal?.aborted
+                ? options.signal.reason
+                : error;
+    });
+    return { ready, close };
+}

package/packages/safejs/dist/modules/mcp.d.ts

@@ -1,3 +1,4 @@
+import { type McpModuleOptions } from "./mcp-transport.js";
 export type McpModuleServerHandle = {
     command: string;
     args?: string[];
@@ -42,8 +43,21 @@ type McpConnection = {
     }>): Promise<McpModuleToolBatchResult[]>;
 };
 export type ConnectMcp = (server: McpModuleServerHandle) => Promise<McpConnection> | McpConnection;
-export declare function makeMcpModule(connectMcp: ConnectMcp): {
+export type McpNamedServerHandle = {
+    name: string;
+};
+export type ManagedMcpClient = McpModuleClient & {
+    close(): Promise<void>;
+};
+export type ManagedMcpModule = {
+    servers: Record<string, ManagedMcpClient>;
+    server(name: string): McpNamedServerHandle;
+    client(handle: McpNamedServerHandle | string): Promise<ManagedMcpClient>;
+};
+type CustomMcpModule = {
     server(server: McpModuleServerHandle): McpModuleServerHandle;
     client(handle: McpModuleServerHandle): Promise<McpModuleClient>;
 };
+export declare function makeMcpModule(options: McpModuleOptions): ManagedMcpModule;
+export declare function makeMcpModule(connectMcp: ConnectMcp): CustomMcpModule;
 export {};

package/packages/safejs/dist/modules/mcp.js

@@ -1,5 +1,112 @@
+import { McpError } from "tiny-mcp-client";
+import { hostErrorData } from "../error/shape.js";
+import { declareHostOperation } from "../interp/host-bridge.js";
+import { runResources } from "../interp/resources.js";
+import { connectMcpTransport, normalizeMcpOptions } from "./mcp-transport.js";
 const DEFAULT_TOOL_BATCH_CONCURRENCY = 4;
 export function makeMcpModule(connectMcp) {
+    if (typeof connectMcp !== "function") {
+        const options = normalizeMcpOptions(connectMcp);
+        const servers = Object.create(null);
+        for (const [name, config] of Object.entries(options.servers)) {
+            const connections = new WeakMap();
+            const directScope = {};
+            const connection = () => {
+                const scope = runResources.getStore();
+                scope?.signal.throwIfAborted();
+                options.signal?.throwIfAborted();
+                const owner = scope ?? directScope;
+                let current = connections.get(owner);
+                if (current === undefined) {
+                    current = connectMcpTransport(config, options, scope?.signal);
+                    connections.set(owner, current);
+                    scope?.add(current.close);
+                }
+                return current.ready;
+            };
+            const invoke = async (operation) => {
+                try {
+                    const result = await operation(await connection());
+                    runResources.getStore()?.signal.throwIfAborted();
+                    options.signal?.throwIfAborted();
+                    return result;
+                }
+                catch (error) {
+                    runResources.getStore()?.signal.throwIfAborted();
+                    options.signal?.throwIfAborted();
+                    if (error instanceof McpError) {
+                        hostErrorData.set(error, {
+                            code: error.code,
+                            ...(error.data === undefined ? {} : { data: error.data })
+                        });
+                    }
+                    throw error;
+                }
+            };
+            servers[name] = {
+                tools: declareHostOperation(async () => invoke(async (client) => {
+                    const tools = [];
+                    const cursors = new Set();
+                    let pages = 0;
+                    let cursor;
+                    do {
+                        const result = await client.listTools(cursor === undefined ? undefined : { cursor });
+                        pages += 1;
+                        tools.push(...normalizeToolsResult(result));
+                        cursor = result.nextCursor;
+                        if (cursor !== undefined) {
+                            if (typeof cursor !== "string" || cursors.has(cursor))
+                                throw new Error("MCP tools pagination did not advance.");
+                            cursors.add(cursor);
+                            if (pages >= options.maxToolPages)
+                                throw new Error("MCP tools page limit exceeded.");
+                        }
+                    } while (cursor !== undefined);
+                    return tools;
+                }), "re-issue"),
+                tool: declareHostOperation(async (toolName, args) => {
+                    const request = createNullRecord({
+                        name: readNonEmptyTrimmedString(toolName, "MCP tool name"),
+                        ...(args === undefined ? {} : { arguments: readToolArguments(args) })
+                    });
+                    return invoke((client) => client.callTool(request));
+                }, "read-side-effect"),
+                toolBatch: declareHostOperation(async (calls) => {
+                    const normalized = readToolBatchCalls(calls);
+                    if (normalized.length === 0)
+                        return [];
+                    return invoke((client) => executeToolBatch(client, normalized));
+                }, "read-side-effect"),
+                close: declareHostOperation(async () => {
+                    const owner = runResources.getStore() ?? directScope;
+                    const current = connections.get(owner);
+                    if (current !== undefined) {
+                        await current.close();
+                        connections.delete(owner);
+                    }
+                }, "re-issue")
+            };
+        }
+        const readName = (value) => {
+            const name = readNonEmptyString(value, "MCP server name");
+            if (!Object.hasOwn(servers, name))
+                throw new Error(`MCP server '${name}' is not configured.`);
+            return name;
+        };
+        return {
+            servers,
+            server(name) {
+                return createNullRecord({ name: readName(name) });
+            },
+            async client(handle) {
+                if (typeof handle === "string")
+                    return servers[readName(handle)];
+                if (!isRecord(handle) || Object.keys(handle).length !== 1 || !Object.hasOwn(handle, "name"))
+                    throw new TypeError("MCP client requires a named server handle.");
+                return servers[readName(getOwnProperty(handle, "name"))];
+            }
+        };
+    }
     return {
         server(server) {
             return normalizeServerHandle(server);

package/packages/safejs/dist/run.js

@@ -1,4 +1,5 @@
 import { hashSource } from "./parse/hash.js";
+import { withRunResources } from "./interp/resources.js";
 import { createReplayableRandom } from "./random.js";
 import { attachErrorSpan, describeThrownValue, materializeWrappedErrorCause, replaceErrorStack } from "./error/shape.js";
 import { parseExecutableModule } from "./parse/parser.js";
@@ -43,7 +44,7 @@ export function run(source, options = {}) {
     const lifecycle = { hostCallbackDepth: 0 };
     const dumpController = createDumpController(lifecycle);
     const promiseTracker = createSandboxPromiseRejectionTracker();
-    const result = withSandboxPromiseRejectionTracker(promiseTracker, async () => {
+    const execute = async () => {
         const promiseReplay = new PromiseReplay(options.snapshot?.promiseReplay);
         return promiseReplayContext.run(promiseReplay, async () => {
             const deactivateOtelSink = activateOtelSink(options.otelSink);
@@ -140,7 +141,8 @@ export function run(source, options = {}) {
                     if (!isSandboxClosure(binding))
                         throw new TypeError("Invalid initial promise operation.");
                     return binding.call([]);
-                });
+                }, hostCalls.rebindHostCapability.bind(hostCalls));
+                hostCalls.validateHostCapabilities();
                 budget.setRetainedDataUsage(initialInputs, measureSandboxData([
                     ...Object.values(initialInputs.values.bindings),
                     ...Object.values(initialInputs.values.imports),
@@ -310,7 +312,8 @@ export function run(source, options = {}) {
                 deactivateOtelSink();
             }
         });
-    });
+    };
+    const result = withRunResources(options.signal, () => withSandboxPromiseRejectionTracker(promiseTracker, execute));
     return attachDumpController(result, dumpController);
 }
 async function throwIfReturnedPromiseRejected(result) {

package/packages/safejs/dist/snapshot/dump-format.d.ts

@@ -1,5 +1,5 @@
 export declare const DUMP_FORMAT_VERSION = 1;
-export declare const EXECUTION_SEMANTICS = "jobs-v4";
+export declare const EXECUTION_SEMANTICS = "jobs-v5";
 export type DumpableSnapshot = {
     sourceHash: string;
     [key: string]: unknown;

package/packages/safejs/dist/snapshot/dump-format.js

@@ -1,5 +1,5 @@
 export const DUMP_FORMAT_VERSION = 1;
-export const EXECUTION_SEMANTICS = "jobs-v4";
+export const EXECUTION_SEMANTICS = "jobs-v5";
 import { assertSnapshotGraphDepth } from "../graph-depth.js";
 import { sandboxErrorTypes } from "../error/shape.js";
 import { getSandboxArgumentEntries, isSandboxArguments } from "../interp/arguments.js";

package/packages/safejs/dist/snapshot/replay-data.d.ts

@@ -61,5 +61,6 @@ export declare function encodeReplayData(value: SandboxValue, options?: {
 export declare function decodeReplayData(input: unknown, options?: {
     resolveCapability?: (id: string) => SandboxClosure | undefined;
     resolvePromise?: (id: string) => SandboxPromise | undefined;
+    onCapabilityRestored?: (original: SandboxClosure, restored: SandboxClosure) => void;
 }): SandboxValue;
 export {};

package/packages/safejs/dist/snapshot/replay-data.js

@@ -185,7 +185,7 @@ export function decodeReplayData(input, options = {}) {
                 restored.set(id, capability);
                 return capability;
             }
-            return createSandboxClosure({
+            const copy = createSandboxClosure({
                 ...capability,
                 boundTarget: capability.boundTarget,
                 cancellationSignal: capability.cancellationSignal,
@@ -202,6 +202,8 @@ export function decodeReplayData(input, options = {}) {
                     return properties;
                 }
             });
+            options.onCapabilityRestored?.(capability, copy);
+            return copy;
         }
         if (kind === "map") {
             const result = createSandboxMap();

package/packages/safejs/dist/snapshot/replay-inputs.d.ts

@@ -1,4 +1,4 @@
-import { type SandboxPromise, type SandboxValue } from "../interp/values.js";
+import { type SandboxClosure, type SandboxPromise, type SandboxValue } from "../interp/values.js";
 import { type ReplayData } from "./replay-data.js";
 export type ReplayInputs = {
     bindings: Record<string, SandboxValue>;
@@ -6,7 +6,7 @@ export type ReplayInputs = {
     entryPointArgs: SandboxValue[] | undefined;
     importMeta: SandboxValue;
 };
-export declare function prepareReplayInputs(current: ReplayInputs, saved?: unknown, preparePromise?: (promise: SandboxPromise | undefined, id: string) => SandboxPromise): {
+export declare function prepareReplayInputs(current: ReplayInputs, saved?: unknown, preparePromise?: (promise: SandboxPromise | undefined, id: string) => SandboxPromise, onCapabilityRestored?: (original: SandboxClosure, restored: SandboxClosure) => void): {
     values: ReplayInputs;
     snapshot: ReplayData;
 };

package/packages/safejs/dist/snapshot/replay-inputs.js

@@ -1,7 +1,7 @@
 import { createSandboxPromise, isSandboxClosure, isSandboxMap, isSandboxSet, isSandboxPromise } from "../interp/values.js";
 import { decodeReplayData, encodeReplayData } from "./replay-data.js";
 const validationPromise = createSandboxPromise(Promise.resolve(undefined));
-export function prepareReplayInputs(current, saved, preparePromise) {
+export function prepareReplayInputs(current, saved, preparePromise, onCapabilityRestored) {
     const identities = new WeakMap();
     const capabilities = new Map();
     const promises = new Map();
@@ -81,6 +81,7 @@ export function prepareReplayInputs(current, saved, preparePromise) {
     assertReplayInputShape(validated);
     const restored = decodeReplayData(saved, {
         resolveCapability,
+        onCapabilityRestored,
         resolvePromise: (id) => {
             const value = readCapability(id);
             if (!promises.has(id) && preparePromise !== undefined)