cellb: lower await on in-browser sites, so AsyncLocalStorage keeps concurrent stores

packages/cellb/docs/browser-host.md

@@ -22,9 +22,10 @@ and WebKit.
   rollback journal keeps its pool slot between commits instead of being created and deleted on each
   one. The database file is the same as with the default `DELETE` (only WAL marks it).
 - **`AsyncLocalStorage` needs help.** Browsers have no async context yet: for a nodejs_compat
-  bundle, `browserProjectRuntime` makes promise and timer continuations carry their stores, and
-  `bundleWorkerForBrowser` lowers `await` so that it goes through them. An in-browser site keeps
-  run-at-a-time stores ([gaps.md](gaps.md)).
+  project, `browserProjectRuntime` makes promise and timer continuations carry their stores, and
+  the code's `await` is lowered so that it goes through them: by `bundleWorkerForBrowser` in a
+  bundle, by `createProjectSite` in a site's modules ([gaps.md](gaps.md)). A browser with TC39
+  AsyncContext needs neither: `node:async_hooks` then runs on it, and a site lowers nothing.
 - **A ServiceWorker cannot spawn workers** and is stopped when idle, dropping every port it held.
   It can route requests, hold Web Locks and pass MessagePorts, so it routes and a page owns the
   workers.
@@ -61,8 +62,9 @@ code and builds the runtime from the wrangler configuration. The code is either:
   `@statewalker/webrun-site-host`, which compiles the project's sources on request
   (`@statewalker/webrun-modules`). Loaded Workers' code is uploaded to it and imported from URLs.
   The project's `cloudflare:workers` has its own `env` and `exports` (Workers loaded into its
-  realm get empty ones; in isolates, their own), and `cloudflare:email` loads; `await` is not
-  lowered, so `AsyncLocalStorage` keeps run-at-a-time stores there (see below).
+  realm get empty ones; in isolates, their own), and `cloudflare:email` loads. With `nodeCompat`
+  the site lowers `await` in the modules it serves, so `AsyncLocalStorage` answers as in a bundle;
+  the site's `nodeCompat` must match the project's flag (see below).
 
 A loaded Worker's `cloudflare:workers` has its own `env` and `exports` (its `WorkerCode.env`, and
 stubs of its own entrypoints), except on a site without `isolateScript`.
@@ -274,9 +276,10 @@ tabs or workers with no server. Its conditional writes are atomic only within on
   outran the pool's growth (`browserProjectRuntime` starts with 64 files). With `sahpoolOpener`,
   pass a larger `initialCapacity` for workloads that open many cells together.
 - **`AsyncLocalStorage.getStore()` returns another cell's store, or `undefined` after an
-  `await`**, in a browser. The code was not bundled by `bundleWorkerForBrowser` with nodejs_compat:
-  an in-browser site keeps run-at-a-time stores (shared between concurrent runs), and a bundle of
-  your own whose `await` is not lowered loses its store at the first `await`.
+  `await`**, in a browser. The code's `await` was not lowered: a bundle of your own, not made by
+  `bundleWorkerForBrowser` with nodejs_compat, or a site made by `createProjectSite` without
+  `nodeCompat` for a nodejs_compat project. Code in an async generator, or (on a site) in an async
+  method that uses `super`, sees no store after its `await`s.
 - **Every S3 request fails with a CORS error** in the console, or `PUT` succeeds but ownership never
   works. The bucket lacks a CORS rule, or does not expose `etag`: conditional writes need it.
 - **`SignatureDoesNotMatch` through a proxy.** `signAs` is missing or differs from the host the proxy

packages/cellb/docs/gaps.md

@@ -89,12 +89,17 @@ Dedicated Workers all work in Firefox 157, and the whole suite passes there. Wha
 - **`AsyncLocalStorage` in a browser carries its stores only through code it can see.** The
   `node:async_hooks` polyfill keeps each `run()` to its synchronous part and makes continuations
   carry their stores (`Promise.prototype.then`, timers, `queueMicrotask`); a native `await` bypasses
-  `then`, so `bundleWorkerForBrowser` lowers `await` to generators in nodejs_compat bundles. A
-  bundle then answers as on Node (test/async-context-worker.ts, 16 objects at once). An in-browser
-  site's modules are compiled without lowering, so there the polyfill keeps its run-at-a-time
-  stores: a `run()`'s store stays set until its promise settles, which is right for one run at a
-  time and shared between concurrent ones (two cells answering at once can read each other's).
-  Event listeners (`addEventListener`, MessagePort `onmessage`) do not carry stores, on Node either.
+  `then`, so with nodejs_compat `await` is lowered to generators: by `bundleWorkerForBrowser`
+  (esbuild) in a bundle, by `createProjectSite` (`lowerAwait`, after webrun-modules) in a site's
+  modules. Both then answer as on Node (test/async-context-worker.ts, 16 objects at once). Not
+  lowered: async generators' own awaits, async functions that use `super` or `new.target` (on a
+  site; esbuild lowers those too), and top-level await; code there resumes in no store after an
+  `await`, rather than in another event's. In a browser with TC39 AsyncContext (Stage 2, shipped
+  by none yet) the polyfill is a layer over `AsyncContext.Variable` and `Snapshot` instead, sites
+  lower nothing (`bundleWorkerForBrowser` takes `lowerAwait: false`), and the engine carries the
+  stores through native awaits as workerd and celld do; `enterWith()` throws there
+  (test/p18 runs it on an AsyncContext made of Node's own). Event listeners (`addEventListener`, MessagePort
+  `onmessage`) do not carry stores, on Node either.
 - **An object's events start in the async context the runtime was made in.** A call to an object
   is a new event, as on Cloudflare: the callee (and its constructor) does not see the caller's
   `AsyncLocalStorage` stores. Where Cloudflare has none, cellb has the stores that were current

packages/cellb/e2e-browser/site-projects.spec.ts

@@ -8,6 +8,11 @@ import {
   goldenDir,
   scenarios,
 } from "../conformance/scenarios.ts";
+import {
+  ASYNC_CONTEXT_CONFIG,
+  ASYNC_CONTEXT_WORKER,
+  expectedSeen,
+} from "../test/async-context-worker.ts";
 import { startReferenceServer } from "../test/ws-reference-server.ts";
 
 // The celld projects served by an IN-BROWSER HTTP SERVER: the page mounts a
@@ -127,3 +132,31 @@ for (const noIsolate of [false, true])
       noIsolate ? { env: [], exports: [] } : { env: ["NAME"], exports: ["Echo"] },
     );
   });
+
+// node:async_hooks on a site: its modules' awaits are lowered as they are served, so 16 objects at
+// once keep their own stores, as in a bundle (agents-sdk-host.spec) and on Node (test/p13).
+test("AsyncLocalStorage on a site keeps each concurrent object's store, as on Node", async ({
+  page,
+}) => {
+  const b64 = (s: string) => Buffer.from(s).toString("base64");
+  await page.goto("/site-projects/");
+  await page.waitForFunction(() => document.body.dataset.ready === "true");
+  const r = await page.evaluate((job) => (window as any).__runOnSite(job), {
+    name: "site-als",
+    scenario: "",
+    config: ASYNC_CONTEXT_CONFIG,
+    files: [
+      ["/index.js", b64(ASYNC_CONTEXT_WORKER)],
+      ["/wrangler.json", b64(JSON.stringify(ASYNC_CONTEXT_CONFIG))],
+    ],
+    assets: [],
+    paths: ["/", "/"],
+  });
+  expect(r.ok, r.error).toBe(true);
+  for (const answer of r.answers) {
+    const run = JSON.parse(answer);
+    expect(run.outside).toBeNull();
+    for (const [name, seen] of Object.entries(run.results))
+      expect(seen).toEqual(expectedSeen(name));
+  }
+});

packages/cellb/package.json

@@ -34,19 +34,20 @@
     "@noble/hashes": "^1.8.0",
     "@sqlite.org/sqlite-wasm": "^3.53.4-build1",
     "@statewalker/webrun-files": "^0.10.0",
+    "@statewalker/webrun-http-browser": "catalog:",
     "@statewalker/webrun-http-streams": "^0.2.2",
     "@statewalker/webrun-rpc": "^0.4.0",
-    "@statewalker/webrun-streams": "^0.2.0",
-    "worker-timers": "^8.0.0",
-    "sqlite-vec": "^0.1.7",
     "@statewalker/webrun-site-builder": "catalog:",
     "@statewalker/webrun-site-host": "catalog:",
-    "@statewalker/webrun-http-browser": "catalog:",
+    "@statewalker/webrun-streams": "^0.2.0",
+    "acorn": "^8.19.0",
+    "assert": "^2.1.0",
+    "buffer": "^6.0.3",
     "events": "^3.3.0",
     "path-browserify": "^1.0.1",
+    "sqlite-vec": "^0.1.7",
     "util": "^0.12.5",
-    "buffer": "^6.0.3",
-    "assert": "^2.1.0"
+    "worker-timers": "^8.0.0"
   },
   "devDependencies": {
     "@playwright/test": "^1.63.0",
@@ -56,16 +57,16 @@
     "@statewalker/webrun-http-browser": "^0.6.0",
     "@statewalker/webrun-modules": "catalog:",
     "@types/node": "catalog:",
+    "@types/ws": "catalog:",
     "esbuild": "^0.28.0",
+    "html-rewriter-wasm": "^0.4.1",
+    "miniflare": "4.20260730.0",
+    "sqlite-vec-wasm-demo": "0.1.6",
     "typescript": "catalog:",
     "vite": "catalog:",
     "vitest": "catalog:",
     "ws": "catalog:",
-    "@types/ws": "catalog:",
-    "zod": "^3.25.0",
-    "sqlite-vec-wasm-demo": "0.1.6",
-    "html-rewriter-wasm": "^0.4.1",
-    "miniflare": "4.20260730.0"
+    "zod": "^3.25.0"
   },
   "private": true
 }

packages/cellb/src/core/async-context.ts

@@ -3,12 +3,42 @@ export type AsyncRoot = <R>(fn: () => R) => R;
 
 const POLYFILL_STATE = Symbol.for("cellb.asyncContext.state");
 
+// biome-ignore lint/suspicious/noExplicitAny: any function, called with its own arguments
+type Fn<R = unknown> = (...args: any[]) => R;
+
+/** TC39 AsyncContext (Stage 2), the part cellb uses. */
+export interface NativeAsyncContext {
+  Variable: new <T>(options?: {
+    name?: string;
+    defaultValue?: T;
+  }) => {
+    get(): T | undefined;
+    run<R>(value: T, fn: Fn<R>, ...args: unknown[]): R;
+  };
+  Snapshot: {
+    new (): { run<R>(fn: Fn<R>, ...args: unknown[]): R };
+    wrap<F extends Fn>(fn: F): F;
+  };
+}
+
+/**
+ * The engine's AsyncContext, when the realm has one: async context then follows every `await`
+ * natively, as in workerd, celld and Node, and code needs no lowering.
+ */
+export function nativeAsyncContext(): NativeAsyncContext | undefined {
+  const ac = (globalThis as { AsyncContext?: Partial<NativeAsyncContext> }).AsyncContext;
+  return typeof ac?.Variable === "function" && typeof ac.Snapshot === "function"
+    ? (ac as NativeAsyncContext)
+    : undefined;
+}
+
 /**
  * The async context current now, to start events in: on Cloudflare a call to a Durable Object is
  * a new event, which sees no `AsyncLocalStorage` store its caller set. On Node it is every store
- * (`AsyncLocalStorage.snapshot()`); in a browser, the `node:async_hooks` polyfill's frame (a
- * realm's polyfill copies share it), or its root frame if the polyfill loads later. Elsewhere
- * events run in their caller's context, as before.
+ * (`AsyncLocalStorage.snapshot()`); in a browser with AsyncContext, an `AsyncContext.Snapshot`;
+ * in other browsers, the `node:async_hooks` polyfill's frame (a realm's polyfill copies share
+ * it), or its root frame if the polyfill loads later. Elsewhere events run in their caller's
+ * context, as before.
  */
 export function asyncRoot(): AsyncRoot {
   const hooks = (
@@ -28,6 +58,11 @@ export function asyncRoot(): AsyncRoot {
   // AsyncContextFrame) stays its own instead of reaching the next event started in the snapshot.
   if (snapshot && Resource)
     return (fn) => snapshot(() => new Resource("cellb.event").runInAsyncScope(fn));
+  const native = nativeAsyncContext();
+  if (native) {
+    const captured = new native.Snapshot();
+    return (fn) => captured.run(fn);
+  }
   type State = { root: unknown; current: unknown };
   const realm = globalThis as Record<symbol, State | undefined>;
   const captured = realm[POLYFILL_STATE]?.current;

packages/cellb/src/node-compat/async_hooks.ts

@@ -1,25 +1,48 @@
 /**
- * `node:async_hooks` AsyncLocalStorage for browsers (nodejs_compat). Browsers have no async
- * context propagation yet (TC39 AsyncContext), so this is the polyfill's design:
+ * `node:async_hooks` AsyncLocalStorage for browsers (nodejs_compat).
+ *
+ * In a realm with TC39 AsyncContext (Stage 2; no browser ships it yet), each AsyncLocalStorage
+ * is an `AsyncContext.Variable`: the engine carries it through every `await`, as workerd, celld
+ * and Node do, and `installAsyncContext()` patches nothing. Code needs no lowering there (lowered
+ * code works too). `enterWith()` has no AsyncContext counterpart and throws.
+ *
+ * Elsewhere this module simulates it:
  *
  * - a context is an immutable frame (store per AsyncLocalStorage); `run()` makes its frame
  *   current for the synchronous part of the call only;
  * - `installAsyncContext()` makes continuations carry the frame that was current when they
  *   were scheduled: `Promise.prototype.then` (so `catch`, `finally` and `await` on a thenable),
  *   `setTimeout`, `setInterval` and `queueMicrotask`;
  * - a native `await` does not call `then`, so code that relies on propagation is compiled with
- *   `async`/`await` lowered to generators (`bundleWorkerForBrowser` does it with nodejs_compat):
- *   each `await` then goes through `Promise.prototype.then`.
+ *   `async`/`await` lowered to generators (with nodejs_compat, `bundleWorkerForBrowser` does it
+ *   with esbuild, an in-browser site with `lowerAwait`): each `await` then goes through
+ *   `Promise.prototype.then`.
  *
  * Code that is not lowered (cellb itself) resumes after a native `await` in the root frame: it
  * neither sees nor leaks another event's stores.
  *
- * Without `installAsyncContext()` (an in-browser site, whose modules are compiled without
- * lowering), `run()` keeps its store set until its promise settles instead: one run at a time
- * sees its store after `await`, concurrent runs share whichever was set last.
+ * Without `installAsyncContext()` (code compiled without lowering), `run()` keeps its store set
+ * until its promise settles instead: one run at a time sees its store after `await`, concurrent
+ * runs share whichever was set last.
  */
+import { type NativeAsyncContext, nativeAsyncContext } from "../core/async-context.js";
+
 type Frame = ReadonlyMap<AsyncLocalStorage<unknown>, unknown>;
 
+const native = nativeAsyncContext();
+type NativeVariable = InstanceType<NativeAsyncContext["Variable"]>;
+const variables = new WeakMap<AsyncLocalStorage<unknown>, NativeVariable>();
+/** Stores whose `disable()` holds until their next `run()` (native backend). */
+const disabled = new WeakSet<AsyncLocalStorage<unknown>>();
+const variableOf = (als: AsyncLocalStorage<unknown>, ac: NativeAsyncContext) => {
+  let v = variables.get(als);
+  if (!v) {
+    v = new ac.Variable();
+    variables.set(als, v);
+  }
+  return v;
+};
+
 // One state per realm: a browser bundle carries its own copy of this module (node:async_hooks
 // bundled in), and its stores must be the frames that cellb's patched `then` restores.
 const STATE = Symbol.for("cellb.asyncContext.state");
@@ -58,9 +81,15 @@ function bindToCurrent<F extends (...args: any[]) => unknown>(fn: F): F {
 
 export class AsyncLocalStorage<T> {
   getStore(): T | undefined {
-    return state.current.get(this as AsyncLocalStorage<unknown>) as T | undefined;
+    const self = this as AsyncLocalStorage<unknown>;
+    if (native) return disabled.has(self) ? undefined : (variableOf(self, native).get() as T);
+    return state.current.get(self) as T | undefined;
   }
   run<R>(store: T, fn: (...args: unknown[]) => R, ...args: unknown[]): R {
+    if (native) {
+      disabled.delete(this as AsyncLocalStorage<unknown>);
+      return variableOf(this as AsyncLocalStorage<unknown>, native).run(store, fn, ...args);
+    }
     const frame = new Map(state.current).set(this as AsyncLocalStorage<unknown>, store);
     if (state.propagate) return runIn(frame, fn, undefined, args);
     // Without installAsyncContext() (code whose `await` is not lowered): the store stays set
@@ -84,14 +113,28 @@ export class AsyncLocalStorage<T> {
   }
   /** The store for the rest of the current synchronous part and what it schedules. */
   enterWith(store: T): void {
+    if (native)
+      throw new Error(
+        "AsyncLocalStorage.enterWith() is not available with AsyncContext: use run() instead",
+      );
     state.current = new Map(state.current).set(this as AsyncLocalStorage<unknown>, store);
   }
   exit<R>(fn: (...args: unknown[]) => R, ...args: unknown[]): R {
+    if (native)
+      return variableOf(this as AsyncLocalStorage<unknown>, native).run(
+        undefined,
+        fn,
+        ...args,
+      );
     const frame = new Map(state.current);
     frame.delete(this as AsyncLocalStorage<unknown>);
     return runIn(frame, fn, undefined, args);
   }
   disable(): void {
+    if (native) {
+      disabled.add(this as AsyncLocalStorage<unknown>);
+      return;
+    }
     if (!state.current.has(this as AsyncLocalStorage<unknown>)) return;
     const frame = new Map(state.current);
     frame.delete(this as AsyncLocalStorage<unknown>);
@@ -100,10 +143,15 @@ export class AsyncLocalStorage<T> {
   /** `fn` bound to the current stores. */
   // biome-ignore lint/suspicious/noExplicitAny: any function, called with its own arguments
   static bind<F extends (...args: any[]) => unknown>(fn: F): F {
+    if (native) return native.Snapshot.wrap(fn);
     return bindToCurrent(fn);
   }
   /** Runs a function later in the stores current now. */
   static snapshot(): <R>(fn: (...args: unknown[]) => R, ...args: unknown[]) => R {
+    if (native) {
+      const snapshot = new native.Snapshot();
+      return (fn, ...args) => snapshot.run(fn, ...args);
+    }
     const frame = state.current;
     return (fn, ...args) => runIn(frame, fn, undefined, args);
   }
@@ -113,9 +161,10 @@ const INSTALLED = Symbol.for("cellb.asyncContext");
 
 /**
  * Makes continuations carry the current stores, in the realm of `g` (a browser worker that runs
- * Worker code). Idempotent.
+ * Worker code). Idempotent. With AsyncContext the engine carries them: nothing to install.
  */
 export function installAsyncContext(g: typeof globalThis = globalThis): void {
+  if (native) return;
   const proto = g.Promise.prototype as Promise<unknown> & { [INSTALLED]?: true };
   if (proto[INSTALLED]) return;
   Object.defineProperty(proto, INSTALLED, { value: true });

packages/cellb/src/project/browser-bundle.ts

@@ -53,7 +53,16 @@ const bareBuiltin = new Set(builtinModules.filter((m) => !m.startsWith("_") && !
 /** Bundles the Worker module at `entry` into one ES module for `browserProjectRuntime`. */
 export async function bundleWorkerForBrowser(
   entry: string,
-  opts: { nodeCompat?: boolean } = {},
+  opts: {
+    nodeCompat?: boolean;
+    /**
+     * Lower `await` to generators (default: with `nodeCompat`), so that the node:async_hooks
+     * polyfill carries AsyncLocalStorage stores through it. `false` only for browsers with TC39
+     * AsyncContext, whose engine carries them through native awaits; a lowered bundle runs there
+     * too.
+     */
+    lowerAwait?: boolean;
+  } = {},
 ): Promise<string> {
   const out = await build({
     entryPoints: [entry],
@@ -67,7 +76,7 @@ export async function bundleWorkerForBrowser(
     keepNames: true,
     // nodejs_compat: `await` lowered to generators, so that each one resumes through
     // Promise.prototype.then and carries its AsyncLocalStorage stores (node-compat/async_hooks.ts).
-    ...(opts.nodeCompat
+    ...((opts.lowerAwait ?? opts.nodeCompat)
       ? { supported: { "async-await": false, "async-generator": false, "for-await": false } }
       : {}),
     logLevel: "silent",

packages/cellb/src/project/browser-host.ts

@@ -105,10 +105,11 @@ async function importWorkerCode(
   const nodeCompat = flags.some((f) => f.startsWith("nodejs_compat"));
   if (nodeCompat) {
     installNodeGlobals();
-    // node:async_hooks: continuations carry their AsyncLocalStorage stores. Only for a bundle:
-    // bundleWorkerForBrowser lowers its awaits so that they go through them; a site's modules
-    // keep native awaits and the polyfill's run-at-a-time stores.
-    if ("bundle" in code) installAsyncContext();
+    // node:async_hooks: continuations carry their AsyncLocalStorage stores (with AsyncContext
+    // the engine does, and this installs nothing). The code's awaits are lowered so that they go
+    // through them: by bundleWorkerForBrowser for a bundle, by createProjectSite (nodeCompat) for
+    // a site's modules.
+    installAsyncContext();
   }
   const globals = {
     ...workersGlobals,