fix(p3-shim): restore task context for deferred fused adapter starts

crates/js-component-bindgen/src/intrinsics/mod.rs

@@ -1724,8 +1724,14 @@ mod tests {
         assert!(source.contains("startSubtask();"));
         assert!(source.contains("registerSubtaskProgress();"));
         assert!(source.contains("driveJspiCallee();"));
-        assert!(source.contains("} else if (enteredSynchronously === null) {"));
+        let delayed_start = source
+            .find("} else if (enteredSynchronously === null) {")
+            .expect("async start intrinsic should contain the delayed-entry branch");
+        let delayed_source = &source[delayed_start..];
         assert!(source.contains("const enterPromise = preparedTask.enter().then(started => {"));
+        assert!(delayed_source.contains("taskID: preparedTask.id(),"));
+        assert!(delayed_source.contains("componentIdx: preparedTask.componentIdx(),"));
+        assert!(delayed_source.contains("fn: startSubtask,"));
         assert!(source.contains("if (subtask.isReturned()) {"));
         assert!(source.contains("callerComponentState.handles.remove(subtask.waitableRep())"));
     }

crates/js-component-bindgen/src/intrinsics/p3/host.rs

@@ -708,7 +708,17 @@ impl HostIntrinsic {
                                         throw preparedTask.isErrored() ?? new Error("task failed to start");
                                     }}
 
-                                    startSubtask();
+                                    // PrepareCall installs the callee task before the eager
+                                    // path runs, but that metadata is gone by the time a
+                                    // contended entry resumes in this timer. Fused adapters
+                                    // can execute context.get while lowering their inputs,
+                                    // so restore the same task context around the delayed
+                                    // start that the eager path observes.
+                                    {with_global_current_task_meta_fn}({{
+                                        taskID: preparedTask.id(),
+                                        componentIdx: preparedTask.componentIdx(),
+                                        fn: startSubtask,
+                                    }});
 
                                     if (callee._jcoMaySuspend === false) {{
                                         driveDirectCallee();

packages/jco-transpile/test/fixtures/components/p3/deferred-subtask-start.wat

@@ -1,9 +1,11 @@
 ;; A sync-lifted, async-typed export makes two async-lowered calls into a
-;; callback-lifted export of another component. The first callee yields, which
-;; suspends it while the second call is prepared, then makes a sync-lowered
-;; call into a third component and returns. The second callee's entry is
-;; deferred (the first holds the component's slice), and must still be
-;; attributed to the right task and never fail outside the call.
+;; callback-lifted export of another component. The first callee raises
+;; backpressure and yields, which suspends it while the second call is
+;; prepared. The second callee is therefore forced through delayed entry.
+;; Once the first resumes it clears backpressure, makes a sync-lowered call
+;; into a third component, and returns. String lowering in the delayed fused
+;; start adapter calls the callee's realloc, which observes context slot 0 and
+;; therefore requires the deferred task metadata to be installed first.
 (component
   (component $leaf
     (core module $m (func (export "leaf")))
@@ -15,61 +17,89 @@
   (component $mid
     (import "leaf" (func $leaf))
 
-    (core module $libc (memory (export "mem") 1))
-    (core instance $libc (instantiate $libc))
+    (core func $get0 (canon context.get i32 0))
+    (core module $libc
+      (import "context" "get0" (func $get0 (result i32)))
+      (memory (export "mem") 1)
+      (func (export "realloc") (param i32 i32 i32 i32) (result i32)
+        (if (call $get0) (then unreachable))
+        (i32.const 64)))
+    (core instance $context (export "get0" (func $get0)))
+    (core instance $libc
+      (instantiate $libc (with "context" (instance $context))))
 
     (core func $leaf (canon lower (func $leaf)))
+    (core func $backpressure-inc (canon backpressure.inc))
+    (core func $backpressure-dec (canon backpressure.dec))
     (core func $yield (canon thread.yield))
     (core func $task-return (canon task.return))
 
     (core module $m
       (import "" "leaf" (func $leaf))
+      (import "" "backpressure.inc" (func $backpressure-inc))
+      (import "" "backpressure.dec" (func $backpressure-dec))
       (import "" "thread.yield" (func $yield (result i32)))
       (import "" "task.return" (func $task-return))
 
-      (func (export "mid") (result i32)
+      (global $completed (mut i32) (i32.const 0))
+
+      (func (export "mid") (param i32 i32) (result i32)
+        (call $backpressure-inc)
         (drop (call $yield))
+        (call $backpressure-dec)
         (call $leaf)
+        (global.set $completed (i32.add (global.get $completed) (i32.const 1)))
         (call $task-return)
         (i32.const 0))
 
       (func (export "callback") (param i32 i32 i32) (result i32)
         (i32.const 0))
+
+      (func (export "completed") (result i32)
+        (global.get $completed))
     )
 
     (core instance $i
       (instantiate $m
         (with "" (instance
           (export "leaf" (func $leaf))
+          (export "backpressure.inc" (func $backpressure-inc))
+          (export "backpressure.dec" (func $backpressure-dec))
           (export "thread.yield" (func $yield))
           (export "task.return" (func $task-return))
         ))
       )
     )
 
-    (func (export "mid") async
+    (func (export "mid") async (param "value" string)
       (canon lift (core func $i "mid") async (callback (core func $i "callback"))
-        (memory (core memory $libc "mem"))))
+        (memory (core memory $libc "mem")) (realloc (core func $libc "realloc"))))
+    (func (export "completed") (result u32)
+      (canon lift (core func $i "completed")))
   )
   (instance $mid (instantiate $mid (with "leaf" (func $leaf "leaf"))))
 
   (component $top
-    (import "mid" (func $mid async))
+    (import "mid" (func $mid async (param "value" string)))
 
-    (core module $libc (memory (export "mem") 1))
+    (core module $libc
+      (memory (export "mem") 1)
+      (data (i32.const 64) "hey"))
     (core instance $libc (instantiate $libc))
 
     ;; async-lowered
     (core func $mid (canon lower (func $mid) async (memory (core memory $libc "mem"))))
 
     (core module $m
-      (import "" "mid" (func $mid (result i32)))
+      (import "" "mid" (func $mid (param i32 i32) (result i32)))
 
       ;; Returns the two calls' packed subtask states, first call in the high byte.
       (func (export "run") (result i32)
         (local $first i32)
-        (local.set $first (call $mid))
-        (i32.or (i32.shl (local.get $first) (i32.const 8)) (call $mid)))
+        (local.set $first (call $mid (i32.const 64) (i32.const 3)))
+        (i32.or
+          (i32.shl (local.get $first) (i32.const 8))
+          (call $mid (i32.const 64) (i32.const 3))))
     )
 
     (core instance $i
@@ -87,4 +117,5 @@
   (instance $top (instantiate $top (with "mid" (func $mid "mid"))))
 
   (export "run" (func $top "run"))
+  (export "completed" (func $mid "completed"))
 )

packages/jco-transpile/test/p3/canonical-abi-traps.ts

@@ -213,7 +213,7 @@ suite.skipIf(typeof WebAssembly.Suspending !== 'function')('Canonical ABI traps'
         },
     );
 
-    test.concurrent('a deferred async-lowered callee start is attributed to its own task', async () => {
+    test.concurrent('a deferred fused adapter start restores the callee task context', async () => {
         const { instance, cleanup } = await setupAsyncTest({
             asyncMode: 'jspi',
             component: {
@@ -222,17 +222,21 @@ suite.skipIf(typeof WebAssembly.Suspending !== 'function')('Canonical ABI traps'
             },
         });
         try {
-            // Each call returns its subtask's state (low nibble) before the
-            // callee has returned: STARTING (0) or STARTED (1).
+            const waitForCompleted = async (expected: number) => {
+                for (let i = 0; i < 50 && instance.completed() !== expected; i++) {
+                    await new Promise<void>((resolve) => setTimeout(resolve, 0));
+                }
+                assert.strictEqual(instance.completed(), expected);
+            };
+            // The first call enters and raises backpressure before yielding;
+            // the second must therefore take the delayed STARTING path.
             const states = (packed: number) => [(packed >>> 8) & 0xf, packed & 0xf];
-            for (const state of states(await instance.run())) {
-                assert.include([0, 1], state);
-            }
+            assert.deepEqual(states(await instance.run()), [1, 0]);
+            await waitForCompleted(2);
             // Neither the first callee's `task.return` nor the second callee's
             // deferred entry failed outside the call: the instance is still usable.
-            for (const state of states(await instance.run())) {
-                assert.include([0, 1], state);
-            }
+            assert.deepEqual(states(await instance.run()), [1, 0]);
+            await waitForCompleted(4);
         } finally {
             await cleanup();
         }