Resume transformed awaits at the continuation boundary

.changeset/quiet-awaits-resume.md

@@ -0,0 +1,5 @@
+---
+"@mcp-b/do-runtime": patch
+---
+
+Restore actor context at the first instruction after each transformed await so delayed continuations cannot outlive their captured input lock.

src/gate.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "vitest";
-import { __gate, __gateAsyncIterable } from "./gate";
+import { __gate, __gateAsyncIterable, __gateAwait, __resumeAwait } from "./gate";
 import { InputGate, OutputGate } from "./io/io-gate";
 import { IoContext, requireInputLock, type Actor, type Timer } from "./io/io-context";
 
@@ -193,6 +193,49 @@ describe("__gate", () => {
   });
 });
 
+describe("transformed await resume", () => {
+  test("restores context at the first instruction after fulfillment", async () => {
+    const context = newContext();
+
+    const held = await context.run(async () => {
+      __resumeAwait(await __gateAwait(portHop()));
+      requireInputLock(context, "explicit transformed continuation");
+      return true;
+    });
+
+    expect(held).toBe(true);
+  });
+
+  test("carries context across nested native promise continuations", async () => {
+    const context = newContext();
+
+    const held = await context.run(async () => {
+      __resumeAwait(await __gateAwait(portHop()));
+      await Promise.resolve();
+      __resumeAwait(await __gateAwait(portHop()));
+      requireInputLock(context, "nested native continuation");
+      return true;
+    });
+
+    expect(held).toBe(true);
+  });
+
+  test("restores context before throwing a rejection", async () => {
+    const context = newContext();
+
+    await context.run(async () => {
+      let caught: unknown;
+      try {
+        __resumeAwait(await __gateAwait(Promise.reject(new Error("expected"))));
+      } catch (error) {
+        caught = error;
+      }
+      expect(caught).toEqual(new Error("expected"));
+      requireInputLock(context, "rejected transformed continuation");
+    });
+  });
+});
+
 describe("__gateAsyncIterable", () => {
   test("forwards early return to the underlying iterator", async () => {
     let canceled = false;

src/gate.ts

@@ -1,6 +1,6 @@
 /* @do-runtime-gated */
 
-import { tryCurrentSlice, type IoContext } from "./io/io-context";
+import { atCheckpointEnd, tryCurrentSlice, type IoContext } from "./io/io-context";
 
 type ContinuationContext = {
   readonly context: IoContext;
@@ -18,6 +18,14 @@ type Outcome<T> =
   | { readonly ok: true; readonly value: T }
   | { readonly ok: false; readonly exception: unknown };
 
+const TRANSFORMED_AWAIT = Symbol("@mcp-b/do-runtime/transformed-await");
+
+type TransformedAwait<T> = {
+  readonly [TRANSFORMED_AWAIT]: true;
+  readonly context: IoContext;
+  readonly outcome: Outcome<T>;
+};
+
 function isThenable(value: unknown): value is PromiseLike<unknown> {
   return (
     (typeof value === "object" && value !== null) ||
@@ -33,20 +41,46 @@ export function __gate<T>(value: T): T | Promise<Awaited<T>> {
   return resumeWithContext(context, Promise.resolve(value));
 }
 
-function resumeWithContext<T>(context: IoContext, promise: Promise<T>): Promise<T> {
-  return new Promise<T>((resolve, reject) => {
+/** Capture an actor await without publishing its context before the continuation runs. */
+export function __gateAwait<T>(value: T): T | Promise<TransformedAwait<Awaited<T>>> {
+  const context = tryCurrentSlice() ?? continuationContext?.context;
+  if (context === undefined) return value;
+  return resumeAwaitWithContext(context, Promise.resolve(value));
+}
+
+/** Restore the captured actor at the first instruction after a transformed await. */
+export function __resumeAwait<T>(value: T | TransformedAwait<T>): T {
+  if (!isTransformedAwait(value)) return value as T;
+
+  restoreContinuation(value.context);
+  if (value.outcome.ok) return value.outcome.value;
+  throw value.outcome.exception;
+}
+
+function isTransformedAwait<T>(value: T | TransformedAwait<T>): value is TransformedAwait<T> {
+  return Reflect.get(Object(value), TRANSFORMED_AWAIT) === true;
+}
+
+function restoreContinuation(context: IoContext): void {
+  const token = {};
+  continuationContext = { context, token };
+  atCheckpointEnd(() => {
+    if (continuationContext?.token === token) continuationContext = undefined;
+  });
+}
+
+function publishOutcome<T, Result>(
+  context: IoContext,
+  promise: Promise<T>,
+  finish: (outcome: Outcome<T>) => Result,
+): Promise<Result> {
+  return new Promise<Result>((resolve, reject) => {
     const publish = context.makeTransformReentryCallback((outcome: Outcome<T>) => {
       if (continuationContext !== undefined) {
         schedulePublication({ publish: () => publish(outcome), reject });
         return;
       }
-      const token = {};
-      continuationContext = { context, token };
-      if (outcome.ok) resolve(outcome.value);
-      else reject(outcome.exception);
-      queueMicrotask(() => {
-        if (continuationContext?.token === token) continuationContext = undefined;
-      });
+      resolve(finish(outcome));
     });
     void promise.then(
       (value) => {
@@ -59,6 +93,25 @@ function resumeWithContext<T>(context: IoContext, promise: Promise<T>): Promise<
   });
 }
 
+function resumeAwaitWithContext<T>(
+  context: IoContext,
+  promise: Promise<T>,
+): Promise<TransformedAwait<T>> {
+  return publishOutcome(context, promise, (outcome) => ({
+    [TRANSFORMED_AWAIT]: true,
+    context,
+    outcome,
+  }));
+}
+
+function resumeWithContext<T>(context: IoContext, promise: Promise<T>): Promise<T> {
+  return publishOutcome(context, promise, (outcome) => {
+    restoreContinuation(context);
+    if (outcome.ok) return outcome.value;
+    throw outcome.exception;
+  });
+}
+
 /**
  * Resolve one transformed await per task, inside a fresh actor slice. Admission
  * attempts are independent so a blocked actor cannot stall the actor that will

src/vite.test.ts

@@ -3,7 +3,7 @@ import { describe, expect, test } from "vitest";
 import { doRuntimeAwaitTransform } from "./vite";
 
 const HEADER =
-  '/* @do-runtime-gated */\nimport { __gate, __gateAsyncIterable } from "@mcp-b/do-runtime/gate";\n';
+  '/* @do-runtime-gated */\nimport { __gateAsyncIterable, __gateAwait, __resumeAwait } from "@mcp-b/do-runtime/gate";\n';
 
 async function transform(
   code: string,
@@ -31,7 +31,7 @@ describe("doRuntimeAwaitTransform", () => {
     const source = "async function run() { return await task; }\n";
 
     await expect(transform(source)).resolves.toBe(
-      `${HEADER}async function run() { return await __gate((task)); }\n`,
+      `${HEADER}async function run() { return __resumeAwait((await __gateAwait((task)))); }\n`,
     );
   });
 
@@ -47,37 +47,38 @@ describe("doRuntimeAwaitTransform", () => {
     {
       name: "nested awaits",
       source: "const value = await outer(await inner);\n",
-      expected: "const value = await __gate((outer(await __gate((inner)))));\n",
+      expected:
+        "const value = __resumeAwait((await __gateAwait((outer(__resumeAwait((await __gateAwait((inner)))))))));\n",
     },
     {
       name: "arrow, object method, and class method bodies",
       source:
         "const arrow = async () => await one;\nconst object = { async method() { await two; } };\nclass Example { async method() { await three; } }\n",
       expected:
-        "const arrow = async () => await __gate((one));\nconst object = { async method() { await __gate((two)); } };\nclass Example { async method() { await __gate((three)); } }\n",
+        "const arrow = async () => __resumeAwait((await __gateAwait((one))));\nconst object = { async method() { __resumeAwait((await __gateAwait((two)))); } };\nclass Example { async method() { __resumeAwait((await __gateAwait((three)))); } }\n",
     },
     {
       name: "async generators",
       source: "async function* values() { yield await item; }\n",
-      expected: "async function* values() { yield await __gate((item)); }\n",
+      expected: "async function* values() { yield __resumeAwait((await __gateAwait((item)))); }\n",
     },
     {
       name: "top-level await",
       source: "const value = await task;\n",
-      expected: "const value = await __gate((task));\n",
+      expected: "const value = __resumeAwait((await __gateAwait((task))));\n",
     },
     {
       name: "await precedence",
       source: "const first = await a ?? b;\nconst second = await (a, b);\n",
       expected:
-        "const first = await __gate((a)) ?? b;\nconst second = await __gate(((a, b)));\n",
+        "const first = __resumeAwait((await __gateAwait((a)))) ?? b;\nconst second = __resumeAwait((await __gateAwait(((a, b)))));\n",
     },
   ])("preserves $name", async ({ source, expected }) => {
     await expect(transform(source)).resolves.toBe(`${HEADER}${expected}`);
   });
 
   test("is idempotent once the marker is present", async () => {
-    const source = `${HEADER}const value = await __gate((task));\n`;
+    const source = `${HEADER}const value = __resumeAwait((await __gateAwait((task))));\n`;
 
     await expect(transform(source)).resolves.toBe(source);
   });

src/vite.ts

@@ -8,7 +8,7 @@ import {
 
 const MARKER = "/* @do-runtime-gated */";
 const IMPORT =
-  'import { __gate, __gateAsyncIterable } from "@mcp-b/do-runtime/gate";';
+  'import { __gateAsyncIterable, __gateAwait, __resumeAwait } from "@mcp-b/do-runtime/gate";';
 
 export interface DoRuntimeAwaitTransformOptions {
   include?: FilterPattern;
@@ -39,8 +39,10 @@ export function doRuntimeAwaitTransform(options?: DoRuntimeAwaitTransformOptions
       const program = this.parse(code);
       new Visitor({
         AwaitExpression(node) {
-          source.prependLeft(node.argument.start, "__gate((");
+          source.prependLeft(node.start, "__resumeAwait((");
+          source.prependLeft(node.argument.start, "__gateAwait((");
           source.appendRight(node.argument.end, "))");
+          source.appendRight(node.end, "))");
           transformed = true;
         },
         ForOfStatement(node) {