Capture and enter async context frames in Rust JSG

src/rust/AGENTS.md

@@ -14,7 +14,7 @@ _Snapshot — the set drifts as crates come and go; `bazel query //src/rust/...`
 
 | Crate                   | Purpose                                                                                                                                                                                                                       |
 | ----------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `jsg/`                  | Rust JSG bindings: `Lock`, `Rc<T>`, `Resource`, `Struct`, `Type`, `Realm`, `FeatureFlags`, module registration; V8 handle types including typed arrays, `ArrayBuffer`, `ArrayBufferView`, `SharedArrayBuffer`, `BackingStore`; JS functions from Rust closures (`v8::Function::new`) |
+| `jsg/`                  | Rust JSG bindings: `Lock`, `Rc<T>`, `Resource`, `Struct`, `Type`, `Realm`, `FeatureFlags`, module registration; V8 handle types including typed arrays, `ArrayBuffer`, `ArrayBufferView`, `SharedArrayBuffer`, `BackingStore`; JS functions from Rust closures (`v8::Function::new`); async context frames (`AsyncContextFrame`, `TracedAsyncContextFrame`) |
 | `jsg-macros/`           | Proc macros: `#[jsg_struct]`, `#[jsg_method]`, `#[jsg_resource]`, `#[jsg_oneof]`, `#[jsg_static_constant]`, `#[jsg_constructor]`, `#[jsg_property]`, `#[jsg_inspect_property]`, `#[jsg_wildcard_property]`                   |
 | `jsg-test/`             | Test harness (`Harness`) for JSG Rust bindings                                                                                                                                                                                |
 | `api/`                  | Rust-implemented Node.js APIs; registers modules via `register_nodejs_modules()`                                                                                                                                              |

src/rust/jsg-compile-tests/async-context-enter-holds-the-lock.rs

@@ -0,0 +1,11 @@
+// Copyright (c) 2026 Cloudflare, Inc.
+// Licensed under the Apache 2.0 license found in the LICENSE file or at:
+//     https://opensource.org/licenses/Apache-2.0
+
+// `enter` borrows the caller's lock for as long as the frame is current, so the closure can only
+// use the lock it is handed, not the caller's.
+pub fn t(lock: &mut jsg::Lock<'_>, frame: &jsg::AsyncContextFrame) {
+    frame.enter(lock, |_inner| lock.new_object()); //~ ERROR E0501
+
+    //~| ERROR E0500
+}

src/rust/jsg-compile-tests/async-context-enter-lock-cannot-escape.rs

@@ -0,0 +1,11 @@
+// Copyright (c) 2026 Cloudflare, Inc.
+// Licensed under the Apache 2.0 license found in the LICENSE file or at:
+//     https://opensource.org/licenses/Apache-2.0
+
+// The lock that `enter` hands its closure is only borrowed for the closure, while the frame is
+// current, so it cannot be kept for use after the previous frame is restored.
+pub fn t<'s>(lock: &mut jsg::Lock<'s>, frame: &jsg::AsyncContextFrame) {
+    let mut escaped = None;
+    frame.enter(lock, |lock| escaped = Some(lock)); //~ ERROR E0521
+    let _ = escaped.map(|lock: &mut jsg::Lock<'s>| lock.new_object());
+}

src/rust/jsg-compile-tests/async-context-frame-is-not-traced.rs

@@ -0,0 +1,18 @@
+// Copyright (c) 2026 Cloudflare, Inc.
+// Licensed under the Apache 2.0 license found in the LICENSE file or at:
+//     https://opensource.org/licenses/Apache-2.0
+
+// An `AsyncContextFrame` holds a strong `Global`, so a cycle through it is never collected. It is
+// not `Traced`, so it cannot be a field of a resource that GC traces; a resource holds a
+// `TracedAsyncContextFrame`.
+use std::cell::Cell;
+
+use jsg_macros::jsg_resource;
+
+#[jsg_resource] //~ ERROR E0277
+struct Holder {
+    frame: Cell<Option<jsg::AsyncContextFrame>>,
+}
+
+#[jsg_resource]
+impl Holder {}

src/rust/jsg-compile-tests/ok-async-context-frame.rs

@@ -0,0 +1,34 @@
+// Copyright (c) 2026 Cloudflare, Inc.
+// Licensed under the Apache 2.0 license found in the LICENSE file or at:
+//     https://opensource.org/licenses/Apache-2.0
+
+// An `AsyncContextFrame` is `'static`, `Send` and `Sync`, so a future passed to async I/O can hold
+// it, and a `TracedAsyncContextFrame` is `Send` and can be a resource field. `enter` opens no
+// scope, so a handle created in the frame stays valid after it.
+use std::cell::Cell;
+
+use jsg::v8;
+use jsg_macros::jsg_resource;
+
+fn assert_static_send_and_sync<T: Send + Sync + 'static>() {}
+fn assert_static_send<T: Send + 'static>() {}
+
+#[jsg_resource]
+struct Holder {
+    frame: Cell<Option<jsg::TracedAsyncContextFrame>>,
+}
+
+#[jsg_resource]
+impl Holder {}
+
+pub fn t<'s>(lock: &mut jsg::Lock<'s>) -> v8::Local<'s, v8::Object> {
+    assert_static_send_and_sync::<jsg::AsyncContextFrame>();
+    assert_static_send::<jsg::TracedAsyncContextFrame>();
+    let frame = jsg::AsyncContextFrame::current(lock);
+    let object = frame.enter(lock, |lock| lock.new_object());
+    let holder = Holder {
+        frame: Cell::new(Some(frame.into())),
+    };
+    drop(holder);
+    object
+}

src/rust/jsg-compile-tests/traced-async-context-frame-is-not-sync.rs

@@ -0,0 +1,17 @@
+// Copyright (c) 2026 Cloudflare, Inc.
+// Licensed under the Apache 2.0 license found in the LICENSE file or at:
+//     https://opensource.org/licenses/Apache-2.0
+
+// A `TracedAsyncContextFrame` holds a `TracedGlobal`, which can move to another thread but not be
+// shared with one.
+//
+// This compiles only if the type is not Sync; see `traced-global-is-not-sync.rs`.
+trait AmbiguousIfSync<A> {
+    fn some_item() {}
+}
+impl<T: ?Sized> AmbiguousIfSync<()> for T {}
+impl<T: ?Sized + Sync> AmbiguousIfSync<u8> for T {}
+
+pub fn t() {
+    let _ = <jsg::TracedAsyncContextFrame as AmbiguousIfSync<_>>::some_item;
+}

src/rust/jsg-test/ffi.c++

@@ -5,6 +5,7 @@
 #include "ffi.h"
 
 #include <workerd/io/features.h>
+#include <workerd/jsg/async-context.h>
 #include <workerd/jsg/setup.h>
 #include <workerd/rust/jsg-test/lib.rs.h>
 #include <workerd/rust/jsg/ffi-inl.h>
@@ -254,6 +255,39 @@ void hold_across_gc(Isolate* isolate,
   request_gc(isolate, GcType::Full);
 }
 
+namespace {
+// The storage key of the async context frames the tests create, like the key of one
+// AsyncLocalStorage.
+const kj::Arc<::workerd::jsg::AsyncContextFrame::StorageKey>& testStorageKey() {
+  static const auto key = kj::arc<::workerd::jsg::AsyncContextFrame::StorageKey>();
+  return key;
+}
+}  // namespace
+
+void run_in_async_context(Isolate* isolate,
+    ::workerd::rust::jsg::Local value,
+    size_t data,
+    ::rust::Fn<void(size_t)> callback) {
+  auto& js = ::workerd::jsg::Lock::from(isolate);
+  ::workerd::jsg::AsyncContextFrame::StorageScope scope(js, testStorageKey().addRef(),
+      js.v8Ref(::workerd::rust::jsg::local_from_ffi<v8::Value>(kj::mv(value))));
+  callback(data);
+}
+
+::workerd::rust::jsg::MaybeLocal async_context_value(Isolate* isolate) {
+  KJ_IF_SOME(frame, ::workerd::jsg::AsyncContextFrame::current(isolate)) {
+    KJ_IF_SOME(value, frame.get(*testStorageKey())) {
+      return ::workerd::rust::jsg::maybe_local_to_ffi(
+          v8::MaybeLocal<v8::Value>(value.getHandle(isolate)));
+    }
+  }
+  return ::workerd::rust::jsg::maybe_local_to_ffi(v8::MaybeLocal<v8::Value>());
+}
+
+bool cpp_sees_root_frame(Isolate* isolate) {
+  return ::workerd::jsg::AsyncContextFrame::current(isolate) == kj::none;
+}
+
 }  // namespace rust::jsg_test
 
 }  // namespace workerd

src/rust/jsg-test/ffi.h

@@ -111,6 +111,20 @@ void hold_across_gc(Isolate* isolate,
     ::rust::Fn<::workerd::rust::jsg::MaybeLocal(Isolate*)> create,
     ::rust::Fn<void(Isolate*, ::workerd::rust::jsg::Local)> check);
 
+// Calls `callback` in a new C++ jsg::AsyncContextFrame that stores `value` under the tests'
+// storage key, as AsyncLocalStorage.run() does, then makes the previous frame current again.
+void run_in_async_context(Isolate* isolate,
+    ::workerd::rust::jsg::Local value,
+    size_t data,
+    ::rust::Fn<void(size_t)> callback);
+
+// The value the current async context frame stores under the tests' storage key, or empty in the
+// root frame or a frame that stores none.
+::workerd::rust::jsg::MaybeLocal async_context_value(Isolate* isolate);
+
+// Returns whether C++ jsg::AsyncContextFrame::current() finds no frame, i.e. sees the root frame.
+bool cpp_sees_root_frame(Isolate* isolate);
+
 }  // namespace rust::jsg_test
 
 }  // namespace workerd

src/rust/jsg-test/lib.rs

@@ -137,6 +137,22 @@ mod ffi {
             create: unsafe fn(*mut Isolate) -> MaybeLocal,
             check: unsafe fn(*mut Isolate, Local),
         );
+
+        /// Calls `callback` in a new C++ `jsg::AsyncContextFrame` that stores `value` under the
+        /// tests' storage key, as `AsyncLocalStorage.run()` does.
+        pub unsafe fn run_in_async_context(
+            isolate: *mut Isolate,
+            value: Local,
+            data: usize, /* callback */
+            callback: unsafe fn(usize /* callback */),
+        );
+
+        /// The value the current async context frame stores under the tests' storage key, or
+        /// empty if it stores none.
+        pub unsafe fn async_context_value(isolate: *mut Isolate) -> MaybeLocal;
+
+        /// Whether C++ `jsg::AsyncContextFrame::current()` finds no frame.
+        pub unsafe fn cpp_sees_root_frame(isolate: *mut Isolate) -> bool;
     }
 }
 
@@ -498,6 +514,71 @@ impl Harness {
         // SAFETY: the isolate is valid and locked with a context entered (guaranteed by Lock).
         unsafe { ffi::hold_across_gc(lock.isolate().as_ffi(), create, check) };
     }
+
+    /// Runs `f` in a new C++ `jsg::AsyncContextFrame` that stores `value` under the tests'
+    /// storage key, as `AsyncLocalStorage.run()` does, and makes the previous frame current
+    /// again afterwards. [`Harness::async_context_value`] reads the value back.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `f` panics, once the previous frame is current again.
+    pub fn run_in_async_context<'s, F: FnOnce(&mut jsg::Lock<'s>) -> R, R>(
+        lock: &mut jsg::Lock<'s>,
+        value: v8::Local<'s, v8::Value>,
+        f: F,
+    ) -> R {
+        struct Call<'a, 's, F, R> {
+            lock: &'a mut jsg::Lock<'s>,
+            f: Option<F>,
+            result: Option<Result<R, Box<dyn std::any::Any + Send>>>,
+        }
+
+        fn trampoline<'s, F: FnOnce(&mut jsg::Lock<'s>) -> R, R>(data: usize) {
+            // SAFETY: data was cast from &raw mut Call in run_in_async_context below.
+            let call = unsafe { &mut *(data as *mut Call<'_, 's, F, R>) };
+            let Some(f) = call.f.take() else {
+                return;
+            };
+            let lock = &mut *call.lock;
+            call.result = Some(std::panic::catch_unwind(std::panic::AssertUnwindSafe(
+                || f(lock),
+            )));
+        }
+
+        let isolate = lock.isolate().as_ffi();
+        // SAFETY: `value` is valid in the lock's scope.
+        let value = unsafe { value.into_ffi() };
+        let mut call = Call::<F, R> {
+            lock,
+            f: Some(f),
+            result: None,
+        };
+        // SAFETY: the lock keeps the isolate locked with a context entered and a scope open, and
+        // C++ opens no scope around the callback, so the lock stays valid in it; `call` outlives
+        // the call.
+        unsafe {
+            ffi::run_in_async_context(isolate, value, &raw mut call as usize, trampoline::<F, R>);
+        }
+        match call.result.expect("the callback ran") {
+            Ok(result) => result,
+            Err(payload) => std::panic::resume_unwind(payload),
+        }
+    }
+
+    /// The value that the current async context frame stores under the tests' storage key (see
+    /// [`Harness::run_in_async_context`]), or `None` in the root frame.
+    pub fn async_context_value<'s>(lock: &mut jsg::Lock<'s>) -> Option<v8::Local<'s, v8::Value>> {
+        // SAFETY: the lock keeps the isolate locked with a scope open, which the value belongs
+        // to.
+        unsafe { v8::Local::from_ffi_maybe(ffi::async_context_value(lock.isolate().as_ffi())) }
+    }
+
+    /// Whether C++ `jsg::AsyncContextFrame::current()` sees the root frame, to compare with
+    /// Rust's [`jsg::AsyncContextFrame::is_root`].
+    pub fn cpp_sees_root_frame(lock: &mut jsg::Lock<'_>) -> bool {
+        // SAFETY: the lock keeps the isolate locked.
+        unsafe { ffi::cpp_sees_root_frame(lock.isolate().as_ffi()) }
+    }
 }
 
 /// Whether `a` and `b` belong to the same `v8::IsolateGroup`.