poe-code 10.0.1 → 10.0.3

package/packages/safejs/dist/run.js

@@ -44,6 +44,7 @@ export function run(source, options = {}) {
     const lifecycle = { hostCallbackDepth: 0 };
     const dumpController = createDumpController(lifecycle);
     const promiseTracker = createSandboxPromiseRejectionTracker();
+    let completedSnapshot;
     const execute = async () => {
         const promiseReplay = new PromiseReplay(options.snapshot?.promiseReplay);
         return promiseReplayContext.run(promiseReplay, async () => {
@@ -52,8 +53,18 @@ export function run(source, options = {}) {
             let leaveHostReplay;
             let leaveInputReplay;
             let createFailureSnapshot;
-            let lastCheckpoint;
+            let cancellationSnapshot;
+            let cancellationSnapshotError;
             let snapshotScheduler;
+            const captureCancellationSnapshot = () => {
+                try {
+                    cancellationSnapshot = createFailureSnapshot?.();
+                }
+                catch (reason) {
+                    cancellationSnapshotError = { reason };
+                }
+            };
+            options.signal?.addEventListener("abort", captureCancellationSnapshot, { once: true });
             try {
                 const restoredSnapshot = options.snapshot === undefined ? undefined : restore(options.snapshot, { source });
                 if (restoredSnapshot !== undefined) {
@@ -165,6 +176,9 @@ export function run(source, options = {}) {
                     bindings: executionScope.snapshot().bindings,
                     clock: options.clock,
                     hostCalls: hostCalls.snapshot(),
+                    replay: hostCalls.snapshotReplay(),
+                    initialInputs: initialInputs.snapshot,
+                    promiseReplay: promiseReplay.snapshot(),
                     random,
                     sourceHash
                 });
@@ -196,7 +210,6 @@ export function run(source, options = {}) {
                                 randomResumeState: typeof yieldPoint.replayState === "number" ? yieldPoint.replayState : undefined,
                                 sourceHash
                             });
-                            lastCheckpoint = snapshot;
                             return snapshot;
                         };
                         activeSnapshotScheduler.onYield(createSnapshot);
@@ -242,7 +255,6 @@ export function run(source, options = {}) {
                                         : undefined,
                                     sourceHash
                                 });
-                                lastCheckpoint = snapshot;
                                 return snapshot;
                             };
                             activeSnapshotScheduler.onYield(createSnapshot);
@@ -278,6 +290,7 @@ export function run(source, options = {}) {
                 if (replayError !== undefined)
                     snapshot.replayError = replayError;
                 dumpController.finalize(snapshot);
+                completedSnapshot = snapshot;
                 return {
                     ...result,
                     snapshot
@@ -288,9 +301,11 @@ export function run(source, options = {}) {
                 materializeWrappedErrorCause(error);
                 if (createFailureSnapshot !== undefined && snapshotScheduler !== undefined) {
                     try {
-                        const snapshot = lastCheckpoint ?? createFailureSnapshot();
-                        await snapshotScheduler.write(snapshot);
+                        if (cancellationSnapshotError !== undefined)
+                            throw cancellationSnapshotError.reason;
+                        const snapshot = cancellationSnapshot ?? createFailureSnapshot();
                         dumpController.finalize(snapshot);
+                        await snapshotScheduler.write(snapshot);
                     }
                     catch (snapshotError) {
                         if (snapshotError instanceof UnsnapshotableValueError) {
@@ -305,6 +320,7 @@ export function run(source, options = {}) {
                 throw error;
             }
             finally {
+                options.signal?.removeEventListener("abort", captureCancellationSnapshot);
                 leaveInputReplay?.();
                 leaveHostReplay?.();
                 promiseReplay.dispose();
@@ -313,7 +329,18 @@ export function run(source, options = {}) {
             }
         });
     };
-    const result = withRunResources(options.signal, () => withSandboxPromiseRejectionTracker(promiseTracker, execute));
+    const result = withRunResources(options.signal, () => withSandboxPromiseRejectionTracker(promiseTracker, execute)).catch(async (error) => {
+        if (completedSnapshot !== undefined) {
+            try {
+                await createSnapshotScheduler(options).write(completedSnapshot);
+            }
+            catch (snapshotError) {
+                console.warn("Failed to write failure snapshot.", snapshotError);
+            }
+        }
+        dumpController.fail(error);
+        throw error;
+    });
     return attachDumpController(result, dumpController);
 }
 async function throwIfReturnedPromiseRejected(result) {

package/packages/safejs/dist/runner/run-harness.d.ts

@@ -1,7 +1,7 @@
 import { type Diagnostic } from "../lint.js";
 import type { ModuleRegistry } from "../modules/registry.js";
 import type { OtelSink } from "../observability/otel.js";
-import { type RunResult } from "../run.js";
+import { type RunOptions, type RunResult } from "../run.js";
 import type { SnapshotBackend } from "../snapshot/backend.js";
 type HarnessMeta = {
     filepath: string;
@@ -9,6 +9,7 @@ type HarnessMeta = {
     version: unknown;
 };
 export type RunHarnessOptions = {
+    budget?: RunOptions["budget"];
     modulesFor: (frontmatter: Record<string, unknown>, meta: HarnessMeta) => ModuleRegistry;
     otelSink?: OtelSink;
     signal?: AbortSignal;

package/packages/safejs/dist/runner/run-harness.js

@@ -66,6 +66,7 @@ export async function runHarness(filepath, options) {
         throw new LintError(lintErrors);
     }
     return runHarnessSource(executableSource, {
+        budget: options.budget,
         entryPointArgs: hasDefaultExport(executableSource, filepath) ? [] : undefined,
         filename: filepath,
         modules,
@@ -102,6 +103,7 @@ export async function runHarnessPair(filepath, options) {
         throw new LintError(lintErrors);
     }
     return runHarnessSource(executableSource, {
+        budget: options.budget,
         entryPointArgs: [frontmatter],
         filename: pair.scriptPath,
         importMeta: {

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

@@ -1,17 +1,20 @@
 import type { RunResult, RunSnapshot } from "../run.js";
+export type DumpOptions = {
+    onFailure?: "throw" | "checkpoint";
+};
 type DumpController = {
     fail(error: unknown): void;
     finalize(snapshot: RunSnapshot): void;
     onYield(createSnapshot: () => RunSnapshot): void;
     requestCurrentSnapshot(): Promise<string>;
-    requestSnapshot(): Promise<string>;
+    requestSnapshot(options?: DumpOptions): Promise<string>;
 };
 export type RunLifecycle = {
     hostCallbackDepth: number;
 };
 export declare function attachDumpController(result: Promise<RunResult>, controller: DumpController): Promise<RunResult>;
 export declare function createDumpController(lifecycle?: RunLifecycle): DumpController;
-export declare function dump(result: Pick<RunResult, "snapshot"> | PromiseLike<RunResult>): Promise<string>;
+export declare function dump(result: Pick<RunResult, "snapshot"> | PromiseLike<RunResult>, options?: DumpOptions): Promise<string>;
 export declare function dumpCurrent(result: Pick<RunResult, "snapshot"> | PromiseLike<RunResult>): Promise<string>;
 export declare function serializeRunSnapshot(snapshot: RunSnapshot): string;
 export {};

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

@@ -61,13 +61,14 @@ export function createDumpController(lifecycle) {
             }
             return this.requestSnapshot();
         },
-        requestSnapshot() {
+        requestSnapshot(options = {}) {
             assertDumpAllowed();
             if (failed !== undefined) {
-                if (isDataBudgetError(failed.error) &&
-                    (latestSnapshot !== undefined || latestSnapshotFactory !== undefined)) {
+                if ((options.onFailure === "checkpoint" ||
+                    (options.onFailure === undefined && isDataBudgetError(failed.error))) &&
+                    finalSnapshot !== undefined) {
                     try {
-                        return Promise.resolve(serializeRunSnapshot(latestSnapshot ?? latestSnapshotFactory()));
+                        return Promise.resolve(serializeRunSnapshot(finalSnapshot));
                     }
                     catch (error) {
                         return Promise.reject(error);
@@ -134,10 +135,10 @@ function isDataBudgetError(error) {
         "budget" in error &&
         error.budget === "dataSize");
 }
-export function dump(result) {
+export function dump(result, options = {}) {
     const controller = readDumpController(result);
     if (controller !== undefined) {
-        return controller.requestSnapshot();
+        return controller.requestSnapshot(options);
     }
     if (hasSnapshot(result)) {
         try {