Retain async context frames across scopes

src/workerd/jsg/async-context-test.c++

@@ -0,0 +1,59 @@
+// 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
+
+#include "async-context.h"
+#include "jsg-test.h"
+
+namespace workerd::jsg::test {
+namespace {
+
+V8System v8System({"--expose-gc"_kj});
+
+struct AsyncContextTestContext: public Object, public ContextGlobal {
+  JSG_RESOURCE_TYPE(AsyncContextTestContext) {}
+};
+JSG_DECLARE_ISOLATE_TYPE(AsyncContextTestIsolate, AsyncContextTestContext);
+
+void collectGarbage(Lock& js) {
+  auto script = check(v8::Script::Compile(js.v8Context(), js.str("gc(); gc();"_kj)));
+  check(script->Run(js.v8Context()));
+}
+
+KJ_TEST("AsyncContextFrame::Scope retains its prior frame") {
+  Evaluator<AsyncContextTestContext, AsyncContextTestIsolate> evaluator(v8System);
+  evaluator.run([&](auto& js) {
+    auto key = kj::arc<AsyncContextFrame::StorageKey>();
+    kj::Maybe<Ref<AsyncContextFrame>> prior = AsyncContextFrame::create(js,
+        AsyncContextFrame::StorageEntry(
+            kj::mv(key), js.v8Ref(v8Str(js.v8Isolate, "prior"_kj).template As<v8::Value>())));
+    auto weakPrior = KJ_ASSERT_NONNULL(prior).getWeakRef(js);
+
+    auto priorScope = [&]() {
+      v8::HandleScope handleScope(js.v8Isolate);
+      return kj::heap<AsyncContextFrame::Scope>(js, *KJ_ASSERT_NONNULL(prior).get());
+    }();
+    {
+      auto rootScope = [&]() {
+        v8::HandleScope handleScope(js.v8Isolate);
+        return kj::heap<AsyncContextFrame::Scope>(js, kj::none);
+      }();
+      prior = kj::none;
+      collectGarbage(js);
+      KJ_EXPECT(weakPrior.isAlive(), "scope did not retain the prior async context frame");
+    }
+
+    KJ_IF_SOME(current, AsyncContextFrame::current(js)) {
+      KJ_IF_SOME(restored, weakPrior.tryGet()) {
+        KJ_EXPECT(&current == &restored);
+      } else {
+        KJ_FAIL_EXPECT("restored async context frame was destroyed");
+      }
+    } else {
+      KJ_FAIL_EXPECT("prior async context frame was not restored");
+    }
+  });
+}
+
+}  // namespace
+}  // namespace workerd::jsg::test

src/workerd/jsg/async-context.c++

@@ -57,8 +57,11 @@ kj::Maybe<AsyncContextFrame&> AsyncContextFrame::current(Lock& js) {
 }
 
 kj::Maybe<Ref<AsyncContextFrame>> AsyncContextFrame::currentRef(Lock& js) {
-  return jsg::AsyncContextFrame::current(js).map(
-      [](jsg::AsyncContextFrame& frame) { return frame.addRef(); });
+  return currentRef(js.v8Isolate);
+}
+
+kj::Maybe<Ref<AsyncContextFrame>> AsyncContextFrame::currentRef(v8::Isolate* isolate) {
+  return current(isolate).map([](AsyncContextFrame& frame) { return frame.addRef(); });
 }
 
 kj::Maybe<AsyncContextFrame&> AsyncContextFrame::current(v8::Isolate* isolate) {
@@ -158,7 +161,7 @@ AsyncContextFrame::Scope::Scope(Lock& js, kj::Maybe<AsyncContextFrame&> resource
 
 AsyncContextFrame::Scope::Scope(v8::Isolate* ptr, kj::Maybe<AsyncContextFrame&> maybeFrame)
     : isolate(ptr),
-      prior(AsyncContextFrame::current(ptr)) {
+      prior(AsyncContextFrame::currentRef(ptr)) {
   maybeSetV8ContinuationContext(isolate, maybeFrame);
 }
 
@@ -168,7 +171,8 @@ AsyncContextFrame::Scope::Scope(Lock& js, kj::Maybe<Ref<AsyncContextFrame>>& res
       })) {}
 
 AsyncContextFrame::Scope::~Scope() noexcept(false) {
-  maybeSetV8ContinuationContext(isolate, prior);
+  maybeSetV8ContinuationContext(isolate,
+      prior.map([](Ref<AsyncContextFrame>& frame) -> AsyncContextFrame& { return *frame; }));
 }
 
 AsyncContextFrame::StorageScope::StorageScope(Lock& js, kj::Arc<StorageKey> key, Value store)

src/workerd/jsg/async-context.h

@@ -135,6 +135,7 @@ class AsyncContextFrame final: public Wrappable {
   // Convenience variation on current() that returns the result wrapped in a Ref for when we
   // need to make sure the frame stays alive.
   static kj::Maybe<Ref<AsyncContextFrame>> currentRef(Lock& js);
+  static kj::Maybe<Ref<AsyncContextFrame>> currentRef(v8::Isolate* isolate);
 
   // Create a new AsyncContextFrame. The new frame inherits the storage context of the current
   // frame (if any) and the given StorageEntry is added.
@@ -172,7 +173,7 @@ class AsyncContextFrame final: public Wrappable {
   // stack until the scope is destroyed.
   struct Scope {
     v8::Isolate* isolate;
-    kj::Maybe<AsyncContextFrame&> prior;
+    kj::Maybe<Ref<AsyncContextFrame>> prior;
     // If frame is nullptr, the root frame is assumed.
     Scope(Lock& js, kj::Maybe<AsyncContextFrame&> frame = kj::none);
     // If frame is nullptr, the root frame is assumed.