From 9364e7e5b25b2cc7b8daf1cb146464ec19dbc67f Mon Sep 17 00:00:00 2001
From: Lann Martin <lann.martin@akamai.com>
Date: Thu, 10 Sep 2026 21:51:13 -0400
Subject: [PATCH] runtime: settle async results before suspended producer exit

---
 runtime/src/exec/boundary.ts                  |  17 +-
 .../jspi/fixtures/task-return-settlement.wasm | Bin 0 -> 1233 bytes
 .../jspi/fixtures/task-return-settlement.wat  | 105 +++++++++
 .../tests/jspi/task_return_settlement_test.ts | 202 ++++++++++++++++++
 runtime/tests/lift_result_driver_test.ts      |  78 +++++++
 5 files changed, 395 insertions(+), 7 deletions(-)
 create mode 100644 runtime/tests/jspi/fixtures/task-return-settlement.wasm
 create mode 100644 runtime/tests/jspi/fixtures/task-return-settlement.wat
 create mode 100644 runtime/tests/jspi/task_return_settlement_test.ts
 create mode 100644 runtime/tests/lift_result_driver_test.ts

diff --git a/runtime/src/exec/boundary.ts b/runtime/src/exec/boundary.ts
index 4650fcd..6d908dc 100644
--- a/runtime/src/exec/boundary.ts
+++ b/runtime/src/exec/boundary.ts
@@ -1283,15 +1283,18 @@ export function createLiftedFunction(input: {
 
     let outcome: DriveExit | Promise<DriveExit>;
     try {
-      // Driver completion is not thread exhaustion. Require a captured result,
-      // no awaiting wasm call of this task, and no entry hop anywhere in the
-      // store. Other tasks' genuine SuspensionPoint parks may remain background
-      // work, but any engine hop this driver started still needs servicing.
-      // Callback tasks can retain waiting threads after task.return; awaiting
-      // those threads' final exit would prevent long-lived producers returning.
       const midWasmCall = () => task.threads.some((t) => store.awaiting.has(t));
       const hopParked = () => entryHopThreads(store).length > 0;
-      const driveDone = () => resolvedSeen && !midWasmCall() && !hopParked();
+      // Driver completion is not thread exhaustion. CONTRACT: an async result
+      // is independent of producer lifetime
+      // (embedder-api.md:180-185; definitions.py:521-526,2360-2369). A
+      // genuine SuspensionPoint may therefore be handed to the settlement
+      // pump once the result exists. Entry hops remain part of result-memory
+      // ordering and must complete first. Sync lifts retain full activation
+      // liveness. Callback tasks may retain waiting producer threads after
+      // task.return; awaiting their final exit would prevent result delivery.
+      const driveDone = () =>
+        resolvedSeen && !hopParked() && (ft.async || !midWasmCall());
       outcome = drive(
         store,
         driveDone,
diff --git a/runtime/tests/jspi/fixtures/task-return-settlement.wasm b/runtime/tests/jspi/fixtures/task-return-settlement.wasm
new file mode 100644
index 0000000000000000000000000000000000000000..a12c96bd03abfb319cc06d1a314f14f0c94afb41
GIT binary patch
literal 1233
zcma)6v2xQu5PiFMI>}OO6azC1Gk`K_7-sB>4~QlO75`w_IYIzSL6Y&Lasg!;S_;09
zk6`Z<kEKA7!gsf;z1_X{_8CR99)l1DbS_(*5|X9;P}O#JS+w@>VMNg0LlS!CLL-9u
zSX~qE9TUQ6Rh8{kxv>|d_U<WRd|g*dJF{hLYx$(<6jGM}1d2dmNNR1@3?#ZN))rd5
z7gM`{=tVTbjno4f33UsLrR1ma>eNVIET&wK0?vRj-RgaEdTAE4!veHh01Mv9Sr)O>
z&Tagg5`WcYOxmLPJg;rLsY|KbN5R!{Z+CU&#nQmxonj1MCysPmT(!mWi=8#LotJhi
z7sq?w{8M#RO8>`uhYxawJm{#|EPWprwbX&qiUJCsDOGPu#75THYKbK9c;YDjN$NY1
z+NE@Y4Edw4Ila?oNaU=)olt5{IHQcyFykC)^gN?<l)p)ls+1pe<1Sy0*Nm*ooohMg
z+|`kMsYap%FoQJ6!vwd~kr?Q686&!k2hdU<SoW|g2dX|m8Vmt7X4^OrL$meRMWNUZ
zuh9T4eoC=I*tA96y73bvnYT@EK=>CiC4SOq$acRcP{%cUs@0(_=wV+7{^i-+qdEK!
z%h8~1jAW=Wb*efCZp;%=0kP*>_B+NSnR3Kz!X)sjlDNbvg1bS8L=-rJNggxYO$tO<
zeniLse`vC**4I@j@|!t%(GQnx$jWg~^z)**DOVi3k?+nb5{Bl^t0mW+!o3|#2h#Sw
n2NFgT9815Vbm}H$f~nZ<k3g^7<h+%4^qJ|1@x1hCxyQ<1e(x9U

literal 0
HcmV?d00001

diff --git a/runtime/tests/jspi/fixtures/task-return-settlement.wat b/runtime/tests/jspi/fixtures/task-return-settlement.wat
new file mode 100644
index 0000000..eadda09
--- /dev/null
+++ b/runtime/tests/jspi/fixtures/task-return-settlement.wat
@@ -0,0 +1,105 @@
+;; An async export publishes its result, then remains in the same activation
+;; while waitable-set.wait JSPI-suspends on a pending async host import.
+(component
+  (import "before-gate" (func $before-gate async))
+  (import "gate" (func $gate async))
+  (import "hop" (func $hop))
+  (import "continued" (func $continued))
+  (import "probe-entered" (func $probe-entered))
+
+  (core module $Mem (memory (export "mem") 1))
+  (core instance $mem (instantiate $Mem))
+
+  (canon task.return (result u32) (core func $task.return))
+  (canon lower (func $before-gate) async (core func $before-gate.async))
+  (canon lower (func $gate) async (core func $gate.async))
+  (canon lower (func $hop) (core func $hop.sync))
+  (canon lower (func $continued) (core func $continued.sync))
+  (canon lower (func $probe-entered) (core func $probe-entered.sync))
+  (canon waitable-set.new (core func $ws.new))
+  (canon waitable.join (core func $w.join))
+  (canon waitable-set.wait (memory (core memory $mem "mem"))
+    (core func $ws.wait))
+  (canon subtask.drop (core func $subtask.drop))
+
+  (core module $Core
+    (import "" "mem" (memory 1))
+    (import "" "task.return" (func $task.return (param i32)))
+    (import "" "before-gate" (func $before-gate (result i32)))
+    (import "" "gate" (func $gate (result i32)))
+    (import "" "hop" (func $hop))
+    (import "" "continued" (func $continued))
+    (import "" "probe-entered" (func $probe-entered))
+    (import "" "waitable-set.new" (func $ws.new (result i32)))
+    (import "" "waitable.join" (func $w.join (param i32 i32)))
+    (import "" "waitable-set.wait"
+      (func $ws.wait (param i32 i32) (result i32)))
+    (import "" "subtask.drop" (func $subtask.drop (param i32)))
+
+    (func (export "run") (result i32)
+      (local $started i32) (local $sub i32) (local $ws i32)
+      ;; First suspend before the result exists. Reaching task.return after
+      ;; this resumption exercises settlement beyond the initial activation.
+      (local.set $started (call $before-gate))
+      (if (i32.ne
+            (i32.and (local.get $started) (i32.const 15))
+            (i32.const 1))
+        (then unreachable))
+      (local.set $sub (i32.shr_u (local.get $started) (i32.const 4)))
+      (local.set $ws (call $ws.new))
+      (call $w.join (local.get $sub) (local.get $ws))
+      (drop (call $ws.wait (local.get $ws) (i32.const 0)))
+      (call $subtask.drop (local.get $sub))
+
+      (call $task.return (i32.const 42))
+      ;; A nonblocking Suspending import adds an engine entry hop between
+      ;; task.return and the genuine second suspension.
+      (call $hop)
+      (local.set $started (call $gate))
+      (if (i32.ne
+            (i32.and (local.get $started) (i32.const 15))
+            (i32.const 1))
+        (then unreachable))
+      (local.set $sub (i32.shr_u (local.get $started) (i32.const 4)))
+      (local.set $ws (call $ws.new))
+      (call $w.join (local.get $sub) (local.get $ws))
+      (if (i32.ne
+            (call $ws.wait (local.get $ws) (i32.const 0))
+            (i32.const 1))
+        (then unreachable))
+      (if (i32.ne (i32.load (i32.const 0)) (local.get $sub))
+        (then unreachable))
+      (if (i32.ne (i32.load (i32.const 4)) (i32.const 2))
+        (then unreachable))
+      (call $subtask.drop (local.get $sub))
+      (call $continued)
+      ;; The result is already delivered. This later producer failure must be
+      ;; recorded for subsequent entry, not turn into an unhandled rejection.
+      unreachable)
+
+    (func (export "run-cb") (param i32 i32 i32) (result i32)
+      unreachable)
+
+    (func (export "probe") (result i32)
+      (call $probe-entered)
+      (i32.const 7)))
+
+  (core instance $core (instantiate $Core (with "" (instance
+    (export "mem" (memory $mem "mem"))
+    (export "task.return" (func $task.return))
+    (export "before-gate" (func $before-gate.async))
+    (export "gate" (func $gate.async))
+    (export "hop" (func $hop.sync))
+    (export "continued" (func $continued.sync))
+    (export "probe-entered" (func $probe-entered.sync))
+    (export "waitable-set.new" (func $ws.new))
+    (export "waitable.join" (func $w.join))
+    (export "waitable-set.wait" (func $ws.wait))
+    (export "subtask.drop" (func $subtask.drop))))))
+
+  (func (export "run") async (result u32)
+    (canon lift (core func $core "run") async
+      (callback (core func $core "run-cb"))))
+  (func (export "probe") (result u32)
+    (canon lift (core func $core "probe")))
+)
diff --git a/runtime/tests/jspi/task_return_settlement_test.ts b/runtime/tests/jspi/task_return_settlement_test.ts
new file mode 100644
index 0000000..06fdf5f
--- /dev/null
+++ b/runtime/tests/jspi/task_return_settlement_test.ts
@@ -0,0 +1,202 @@
+import { instantiate } from "../../src/embedder/mod.ts";
+import { Translator } from "../../src/shim/mod.ts";
+import { isTrap, suspending } from "@polyengine/protocol";
+import { assert, assertEquals, assertRejects } from "./asserts.ts";
+
+const root = new URL("../../../", import.meta.url);
+
+async function readIfPresent(rel: string): Promise<Uint8Array | null> {
+  try {
+    return await Deno.readFile(new URL(rel, root));
+  } catch {
+    return null;
+  }
+}
+
+const shimWasm = await readIfPresent(
+  "target/wasm32-unknown-unknown/release/translator_shim.wasm",
+);
+if (shimWasm === null) {
+  console.warn(
+    "SKIP task.return settlement: missing translator_shim.wasm " +
+      "(cargo build -p translator-shim --release --target wasm32-unknown-unknown)",
+  );
+}
+const componentBytes = await Deno.readFile(
+  new URL("./fixtures/task-return-settlement.wasm", import.meta.url),
+);
+
+Deno.test({
+  name:
+    "#323: task.return settles before a producer's genuine JSPI suspension ends",
+  ignore: shimWasm === null,
+  fn: async () => {
+    const gate = Promise.withResolvers<void>();
+    const entered = Promise.withResolvers<void>();
+    const continued = Promise.withResolvers<void>();
+    const translator = await Translator.create(shimWasm!);
+    const translated = translator.translate(componentBytes);
+    let probeEntered = false;
+    const component = await instantiate({
+      componentBytes,
+      ...translated,
+    }, {
+      "before-gate": () => Promise.resolve(),
+      hop: suspending(() => undefined),
+      gate: () => {
+        entered.resolve();
+        return gate.promise;
+      },
+      continued: () => continued.resolve(),
+      "probe-entered": () => void (probeEntered = true),
+    });
+
+    const result = component.exports.run() as Promise<number>;
+    try {
+      await entered.promise;
+      const early = await Promise.race([
+        result.then((value) => ({ state: "resolved", value })),
+        new Promise<{ state: "pending" }>((resolve) =>
+          setTimeout(() => resolve({ state: "pending" }), 50)
+        ),
+      ]);
+      assertEquals(
+        JSON.stringify(early),
+        JSON.stringify({ state: "resolved", value: 42 }),
+        "task.return result was withheld behind the producer gate",
+      );
+    } finally {
+      // Never strand the real JSPI activation when an assertion fails.
+      gate.resolve();
+    }
+
+    // Observe work after the suspension independently of the settled export.
+    await continued.promise;
+    assertEquals(await result, 42);
+
+    // The producer traps after continuation. The export result stays delivered,
+    // while the failure is retained and surfaced by subsequent entry.
+    const later = await assertRejects(
+      () => component.exports.probe() as Promise<unknown>,
+      "post-delivery producer failure must reach a later entry",
+    );
+    assertEquals(
+      probeEntered,
+      false,
+      "retained poison must precede probe entry",
+    );
+    assert(
+      isTrap(later) && String(later).includes("unreachable"),
+      `expected the original branded producer trap, got ${later}`,
+    );
+  },
+});
+
+Deno.test({
+  name: "#323: a late host rejection remains observable after result delivery",
+  ignore: shimWasm === null,
+  fn: async () => {
+    const gate = Promise.withResolvers<void>();
+    const entered = Promise.withResolvers<void>();
+    const boom = new Error("late gate rejection");
+    let probeEntered = false;
+    const translator = await Translator.create(shimWasm!);
+    const component = await instantiate({
+      componentBytes,
+      ...translator.translate(componentBytes),
+    }, {
+      "before-gate": () => Promise.resolve(),
+      hop: suspending(() => undefined),
+      gate: () => {
+        entered.resolve();
+        return gate.promise;
+      },
+      continued: () => {
+        throw new Error("continued after rejected gate");
+      },
+      "probe-entered": () => void (probeEntered = true),
+    });
+
+    const result = component.exports.run() as Promise<number>;
+    try {
+      await entered.promise;
+      const early = await Promise.race([
+        result.then((value) => ({ state: "resolved", value })),
+        new Promise<{ state: "pending" }>((resolve) =>
+          setTimeout(() => resolve({ state: "pending" }), 50)
+        ),
+      ]);
+      assertEquals(
+        JSON.stringify(early),
+        JSON.stringify({ state: "resolved", value: 42 }),
+      );
+    } finally {
+      gate.reject(boom);
+    }
+    await new Promise((resolve) => setTimeout(resolve, 0));
+
+    const later = await assertRejects(
+      () => component.exports.probe() as Promise<unknown>,
+      "late host rejection must reach a later entry",
+    );
+    assert(
+      isTrap(later) && String(later).includes(boom.message),
+      `expected the branded host rejection with its original cause, got ${later}`,
+    );
+    assertEquals(
+      probeEntered,
+      true,
+      "the probe entered before the resumed producer's rejection surfaced",
+    );
+  },
+});
+
+Deno.test({
+  name: "#323: task.return after one suspension settles at the next suspension",
+  ignore: shimWasm === null,
+  fn: async () => {
+    const before = Promise.withResolvers<void>();
+    const after = Promise.withResolvers<void>();
+    const enteredBefore = Promise.withResolvers<void>();
+    const enteredAfter = Promise.withResolvers<void>();
+    const continued = Promise.withResolvers<void>();
+    const translator = await Translator.create(shimWasm!);
+    const component = await instantiate({
+      componentBytes,
+      ...translator.translate(componentBytes),
+    }, {
+      "before-gate": () => {
+        enteredBefore.resolve();
+        return before.promise;
+      },
+      hop: suspending(() => undefined),
+      gate: () => {
+        enteredAfter.resolve();
+        return after.promise;
+      },
+      continued: () => continued.resolve(),
+      "probe-entered": () => {},
+    });
+
+    const result = component.exports.run() as Promise<number>;
+    try {
+      await enteredBefore.promise;
+      before.resolve();
+      await enteredAfter.promise;
+      const early = await Promise.race([
+        result.then((value) => ({ state: "resolved", value })),
+        new Promise<{ state: "pending" }>((resolve) =>
+          setTimeout(() => resolve({ state: "pending" }), 50)
+        ),
+      ]);
+      assertEquals(
+        JSON.stringify(early),
+        JSON.stringify({ state: "resolved", value: 42 }),
+      );
+    } finally {
+      before.resolve();
+      after.resolve();
+    }
+    await continued.promise;
+  },
+});
diff --git a/runtime/tests/lift_result_driver_test.ts b/runtime/tests/lift_result_driver_test.ts
new file mode 100644
index 0000000..13978b5
--- /dev/null
+++ b/runtime/tests/lift_result_driver_test.ts
@@ -0,0 +1,78 @@
+import { assertEq } from "./support/asserts.ts";
+import {
+  createLiftedFunction,
+  newStats,
+  registerHostCall,
+  type ResolvedOptions,
+} from "../src/exec/mod.ts";
+import {
+  ComponentInstanceState,
+  currentThread,
+  Store,
+} from "../src/task/mod.ts";
+import type { FuncType } from "../src/cabi/types.ts";
+
+const FT: FuncType = { params: [], results: [], async: true };
+
+Deno.test("#323: a foreign driver failure does not unwind the live producer", async () => {
+  const store = new Store();
+  const inst = new ComponentInstanceState(0, store);
+  const opts: ResolvedOptions = {
+    stringEncoding: "utf8",
+    memory: null,
+    realloc: null,
+    postReturn: null,
+    callback: () => (() => 0) as never,
+    async: true,
+    cancellable: false,
+    coreType: { params: [], results: ["i32"] },
+    instance: inst,
+  };
+  const gate = Promise.withResolvers<number>();
+  let releases = 0;
+  let producerThread!: {
+    syncCallStack: { releaseLenders(): void }[];
+  };
+  const ownerWait = {
+    owner: undefined as unknown,
+    ready: () => false,
+    waiting: () => true,
+    resume: () => {},
+  };
+  const producer = createLiftedFunction({
+    name: "producer",
+    ft: FT,
+    opts,
+    core: () => {
+      producerThread = currentThread() as typeof producerThread;
+      producerThread.syncCallStack.push({
+        releaseLenders: () => releases++,
+      });
+      const task = (currentThread() as unknown as {
+        task: { return_(result: never[]): void };
+      }).task;
+      task.return_([]);
+      ownerWait.owner = currentThread();
+      store.startWaiting(ownerWait as never);
+      return gate.promise;
+    },
+    stats: newStats(),
+  });
+
+  const result = producer() as Promise<void>;
+  assertEq(await result, undefined);
+
+  const boom = new Error("foreign driver failure");
+  let wake!: () => void;
+  const host = new Promise<void>((resolve) => (wake = resolve));
+  registerHostCall(store, host);
+  store.hostFailure = boom;
+  wake();
+  await new Promise((resolve) => setTimeout(resolve, 0));
+  assertEq(releases, 0, "foreign failure must not unwind producer scopes");
+
+  store.stopWaiting(ownerWait as never);
+  producerThread.syncCallStack.pop();
+  gate.resolve(0);
+  await new Promise((resolve) => setTimeout(resolve, 0));
+});