From 24b399a47cda56596fa9e4c38bd689bdb459df10 Mon Sep 17 00:00:00 2001 From: Alex Nahas Date: Sun, 23 Aug 2026 02:23:47 -0700 Subject: [PATCH] Resume transformed awaits at the continuation boundary The previous helper published actor identity before resolving the awaited promise and cleared it one microtask later. Nested native promise jobs could resume after that window while the actor lock was still held, leaving the next transformed await unable to identify its owner. Wrap each transformed result and restore its captured context at the first post-await instruction. Keep that identity until the same checkpoint boundary that releases the input lock. --- .changeset/quiet-awaits-resume.md | 5 +++ src/gate.test.ts | 45 ++++++++++++++++++- src/gate.ts | 73 ++++++++++++++++++++++++++----- src/vite.test.ts | 17 +++---- src/vite.ts | 6 ++- 5 files changed, 125 insertions(+), 21 deletions(-) create mode 100644 .changeset/quiet-awaits-resume.md diff --git a/.changeset/quiet-awaits-resume.md b/.changeset/quiet-awaits-resume.md new file mode 100644 index 0000000..d4e6f42 --- /dev/null +++ b/.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. diff --git a/src/gate.test.ts b/src/gate.test.ts index 28d24a2..7a8b0db 100644 --- a/src/gate.test.ts +++ b/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; diff --git a/src/gate.ts b/src/gate.ts index 2798367..ca9fadc 100644 --- a/src/gate.ts +++ b/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 = | { readonly ok: true; readonly value: T } | { readonly ok: false; readonly exception: unknown }; +const TRANSFORMED_AWAIT = Symbol("@mcp-b/do-runtime/transformed-await"); + +type TransformedAwait = { + readonly [TRANSFORMED_AWAIT]: true; + readonly context: IoContext; + readonly outcome: Outcome; +}; + function isThenable(value: unknown): value is PromiseLike { return ( (typeof value === "object" && value !== null) || @@ -33,20 +41,46 @@ export function __gate(value: T): T | Promise> { return resumeWithContext(context, Promise.resolve(value)); } -function resumeWithContext(context: IoContext, promise: Promise): Promise { - return new Promise((resolve, reject) => { +/** Capture an actor await without publishing its context before the continuation runs. */ +export function __gateAwait(value: T): T | Promise>> { + 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(value: T | TransformedAwait): 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(value: T | TransformedAwait): value is TransformedAwait { + 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( + context: IoContext, + promise: Promise, + finish: (outcome: Outcome) => Result, +): Promise { + return new Promise((resolve, reject) => { const publish = context.makeTransformReentryCallback((outcome: Outcome) => { 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(context: IoContext, promise: Promise): Promise< }); } +function resumeAwaitWithContext( + context: IoContext, + promise: Promise, +): Promise> { + return publishOutcome(context, promise, (outcome) => ({ + [TRANSFORMED_AWAIT]: true, + context, + outcome, + })); +} + +function resumeWithContext(context: IoContext, promise: Promise): Promise { + 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 diff --git a/src/vite.test.ts b/src/vite.test.ts index 206f64a..4f6b8e0 100644 --- a/src/vite.test.ts +++ b/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); }); diff --git a/src/vite.ts b/src/vite.ts index 0a91b3b..155789c 100644 --- a/src/vite.ts +++ b/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) {