Preserve AsyncLocalStorage context in unhandledRejection handlers

scripts/build/deps/webkit.ts

@@ -7,7 +7,13 @@
 // -lto variants built with ThinLTO (per-module summaries for cross-language
 // importing), and the Windows ICU data table filtered + per-item zstd
 // compressed (lazily decompressed via bun_icu_decompress.cpp).
-export const WEBKIT_VERSION = "4895f45dfbd0d1226c4d41799887bc0ecb9f341b";
+//
+// Preview build of oven-sh/WebKit#268 (PromiseFinallyAwaitJob carries the
+// async context across), on top of WebKit main a8d15c1c — which already has
+// the AsyncFunctionResume settle-ordering fix (#295) and the
+// AsyncContextSwapScope RAII helper (#301) that #268 now uses. Re-pin to the
+// autobuild tag of its merge commit once it lands on WebKit main.
+export const WEBKIT_VERSION = "autobuild-preview-pr-268-5f70edce";
 
 /**
  * WebKit (JavaScriptCore) — the JS engine.

src/jsc/JSGlobalObject.rs

@@ -1138,6 +1138,17 @@ impl JSGlobalObject {
         });
     }
 
+    /// Installs `context` as the current async context, returning the previous one.
+    pub fn exchange_async_context(&self, context: JSValue) -> JSValue {
+        unsafe extern "C" {
+            safe fn AsyncContextFrame__exchangeAsyncContext(
+                global: &JSGlobalObject,
+                context: JSValue,
+            ) -> JSValue;
+        }
+        AsyncContextFrame__exchangeAsyncContext(self, context)
+    }
+
     pub fn readable_stream_to_array_buffer(&self, value: JSValue) -> JSValue {
         ZigGlobalObject__readableStreamToArrayBuffer(self, value)
     }
@@ -1728,6 +1739,31 @@ unsafe extern "C" {
     safe fn ScriptExecutionContextIdentifier__forGlobalObject(global: &JSGlobalObject) -> u32;
 }
 
+/// Clears the current async context for the guard's lifetime, restoring it on drop.
+///
+/// Top-level microtask drains and GC must not observe an installed context: the
+/// propagation machinery assumes the ambient context is undefined there, which is
+/// why `JSNextTickQueue` resets the slot after draining.
+pub struct ClearedAsyncContextScope<'a> {
+    global: &'a JSGlobalObject,
+    previous: JSValue,
+}
+
+impl<'a> ClearedAsyncContextScope<'a> {
+    pub fn new(global: &'a JSGlobalObject) -> Self {
+        Self {
+            global,
+            previous: global.exchange_async_context(JSValue::UNDEFINED),
+        }
+    }
+}
+
+impl Drop for ClearedAsyncContextScope<'_> {
+    fn drop(&mut self) {
+        self.global.exchange_async_context(self.previous);
+    }
+}
+
 impl ScriptExecutionContextIdentifier {
     /// Returns `None` if the context referred to by `self` no longer exists.
     pub fn global_object(self) -> Option<GlobalRef> {

src/jsc/VirtualMachine.rs

@@ -3351,12 +3351,19 @@ impl VirtualMachine {
 
         if isBunTest.load(core::sync::atomic::Ordering::Relaxed) {
             self.unhandled_error_counter += 1;
+            // The test runner's handler may drive the next user callback, and
+            // those are registered unwrapped — don't let one inherit the
+            // promise's context. `handleRejectedPromises` restores afterwards.
+            let _scope = jsc::ClearedAsyncContextScope::new(global_object);
             (self.on_unhandled_rejection)(self, global_object, reason);
             return;
         }
 
         // Each arm drains microtasks on exit — hoisted into a closure.
+        // `handleRejectedPromises` runs this dispatch with the rejected promise's
+        // async context installed; a top-level drain must not see it.
         let drain = |this: &mut Self| {
+            let _scope = jsc::ClearedAsyncContextScope::new(global_object);
             let _ = this.event_loop_mut().drain_microtasks();
         };
         // Wrapper over the `Bun__handleUnhandledRejection` FFI call (returns
@@ -3430,7 +3437,11 @@ impl VirtualMachine {
                 // continue to default handler — but RETURN if this drain
                 // errors (the VM is dead; don't bump the counter or invoke the
                 // handler).
-                if self.event_loop_mut().drain_microtasks().is_err() {
+                let drained = {
+                    let _scope = jsc::ClearedAsyncContextScope::new(global_object);
+                    self.event_loop_mut().drain_microtasks()
+                };
+                if drained.is_err() {
                     return;
                 }
             }

src/jsc/bindings/AsyncContextFrame.cpp

@@ -96,6 +96,15 @@ extern "C" JSC::EncodedJSValue AsyncContextFrame__withAsyncContextIfNeeded(JSGlo
     return JSValue::encode(AsyncContextFrame::withAsyncContextIfNeeded(globalObject, JSValue::decode(callback)));
 }
 
+// Installs `context` as the current async context, returning the previous one.
+extern "C" JSC::EncodedJSValue AsyncContextFrame__exchangeAsyncContext(JSGlobalObject* globalObject, JSC::EncodedJSValue context)
+{
+    auto* asyncContextData = globalObject->m_asyncContextData.get();
+    JSValue previous = asyncContextData->getInternalField(0);
+    asyncContextData->putInternalField(JSC::getVM(globalObject), 0, JSValue::decode(context));
+    return JSValue::encode(previous);
+}
+
 #define ASYNCCONTEXTFRAME_CALL_IMPL(...)                                            \
     if (!functionObject.isCell())                                                   \
         return jsUndefined();                                                       \

src/jsc/bindings/BunProcess.cpp

@@ -1354,7 +1354,22 @@ extern "C" int Bun__handleUnhandledRejection(JSC::JSGlobalObject* lexicalGlobalO
         MarkedArgumentBuffer args;
         args.append(reason);
         args.append(promise);
-        wrapped.emit(eventType, args);
+        // The caller emits this event with the promise's async context installed.
+        // A listener that throws is reported as an uncaught exception; let the
+        // throw out of emit and clear the slot before reporting it, instead of
+        // letting EventEmitter report it inside the installed window. Node's
+        // uncaughtException handler observes undefined here too (even with a
+        // persistent enterWith), so "clear" — not "restore the drain's ambient" —
+        // is the semantic the dual-runtime test pins.
+        WTF::NakedPtr<JSC::Exception> listenerException;
+        wrapped.emit(eventType, args, listenerException);
+        if (listenerException) [[unlikely]] {
+            auto* asyncContextData = globalObject->m_asyncContextData.get();
+            JSC::JSValue saved = asyncContextData->getInternalField(0);
+            asyncContextData->putInternalField(JSC::getVM(globalObject), 0, JSC::jsUndefined());
+            Bun__reportUnhandledError(globalObject, JSC::JSValue::encode(JSC::JSValue(listenerException.get())));
+            asyncContextData->putInternalField(JSC::getVM(globalObject), 0, saved);
+        }
         return true;
     }
 

src/jsc/bindings/ZigGlobalObject.cpp

@@ -1059,18 +1059,39 @@ void GlobalObject::reportUncaughtExceptionAtEventLoop(JSGlobalObject* globalObje
 
 extern "C" void Bun__handleHandledPromise(Zig::GlobalObject* JSGlobalObject, JSC::JSPromise* promise);
 
+// Entries in m_aboutToBeNotifiedRejectedPromises are either the rejected
+// JSPromise itself, or an AsyncContextFrame pairing it (as `callback`) with
+// the async context that was active when it was rejected.
+static JSC::JSPromise* rejectedPromiseFromEntry(JSC::JSCell* entry, JSC::JSValue* asyncContext = nullptr)
+{
+    if (auto* frame = dynamicDowncast<AsyncContextFrame>(entry)) {
+        if (asyncContext)
+            *asyncContext = frame->context.get();
+        return uncheckedDowncast<JSC::JSPromise>(frame->callback.get());
+    }
+    return uncheckedDowncast<JSC::JSPromise>(entry);
+}
+
 void GlobalObject::promiseRejectionTracker(JSGlobalObject* obj, JSC::JSPromise* promise,
     JSC::JSPromiseRejectionOperation operation)
 {
     auto* globalObj = static_cast<GlobalObject*>(obj);
 
     switch (operation) {
-    case JSPromiseRejectionOperation::Reject:
-        globalObj->m_aboutToBeNotifiedRejectedPromises.append(obj->vm(), globalObj, promise);
+    case JSPromiseRejectionOperation::Reject: {
+        // Snapshot the async context at rejection time, so that
+        // "unhandledRejection" listeners observe the rejected promise's
+        // AsyncLocalStorage context, like Node.js. The event itself is only
+        // emitted later (at the end of the tick, from handleRejectedPromises),
+        // by which point the context that was live here has been unwound.
+        // Wraps the promise in an AsyncContextFrame only when a context is active.
+        JSC::JSCell* entry = AsyncContextFrame::withAsyncContextIfNeeded(obj, promise).asCell();
+        globalObj->m_aboutToBeNotifiedRejectedPromises.append(obj->vm(), globalObj, entry);
         break;
-    case JSPromiseRejectionOperation::Handle:
-        bool removed = globalObj->m_aboutToBeNotifiedRejectedPromises.removeFirstMatching(globalObj, [&](JSC::WriteBarrier<JSC::JSPromise>& unhandledPromise) {
-            return unhandledPromise.get() == promise;
+    }
+    case JSPromiseRejectionOperation::Handle: {
+        bool removed = globalObj->m_aboutToBeNotifiedRejectedPromises.removeFirstMatching(globalObj, [&](JSC::WriteBarrier<JSC::JSCell>& entry) {
+            return rejectedPromiseFromEntry(entry.get()) == promise;
         });
         if (removed) break;
         // handleRejectedPromises() drains the list into a local buffer before
@@ -1081,14 +1102,15 @@ void GlobalObject::promiseRejectionTracker(JSGlobalObject* obj, JSC::JSPromise*
         // handleRejectedPromises(), so there may be more than one).
         for (auto* inflight = globalObj->m_rejectedPromisesBeingProcessed; inflight; inflight = inflight->outer) {
             for (size_t i = inflight->index, n = inflight->buffer->size(); i < n; ++i) {
-                if (inflight->buffer->at(i).asCell() == promise)
+                if (rejectedPromiseFromEntry(inflight->buffer->at(i).asCell()) == promise)
                     return;
             }
         }
         // The promise rejection has already been notified, now we need to queue it for the rejectionHandled event
         Bun__handleHandledPromise(globalObj, promise);
         break;
     }
+    }
 }
 
 void GlobalObject::setConsole(void* console)
@@ -3414,12 +3436,34 @@ void GlobalObject::handleRejectedPromises()
         InFlightRejections inflight { &promises, 0, m_rejectedPromisesBeingProcessed };
         WTF::SetForScope inflightScope(m_rejectedPromisesBeingProcessed, &inflight);
         for (size_t i = 0, size = promises.size(); i < size; ++i) {
-            auto* promise = static_cast<JSC::JSPromise*>(promises.at(i).asCell());
+            JSC::JSValue asyncContext;
+            auto* promise = rejectedPromiseFromEntry(promises.at(i).asCell(), &asyncContext);
             if (promise->isHandled())
                 continue;
             inflight.index = i + 1;
 
+            // Emit the event in the context the promise was rejected in, the way
+            // Node exchanges the context frame around this dispatch. Entries
+            // rejected with no context replay `undefined` rather than inherit the
+            // drain's: it can be re-entered synchronously from user code.
+            InternalFieldTuple* asyncContextData = nullptr;
+            JSC::JSValue restoreAsyncContext;
+            if (asyncContext || isAsyncContextTrackingEnabled()) {
+                if (!asyncContext)
+                    asyncContext = JSC::jsUndefined();
+                asyncContextData = m_asyncContextData.get();
+                restoreAsyncContext = asyncContextData->getInternalField(0);
+                asyncContextData->putInternalField(virtual_machine, 0, asyncContext);
+            }
+
             Bun__handleRejectedPromise(this, promise);
+
+            // Any uncaught exception that leaks out of the dispatch is reported
+            // after restoring, like Node's processPromiseRejections does in its
+            // finally before the throw reaches triggerUncaughtException.
+            if (asyncContextData)
+                asyncContextData->putInternalField(virtual_machine, 0, restoreAsyncContext);
+
             if (auto ex = scope.exception()) {
                 if (virtual_machine.isTerminationException(ex)) [[unlikely]]
                     return;

src/jsc/bindings/ZigGlobalObject.h

@@ -797,7 +797,11 @@ class GlobalObject : public Bun::GlobalScope {
     DOMGuardedObjectSet m_guardedObjects WTF_GUARDED_BY_LOCK(m_gcLock);
     WebCore::SubtleCrypto* m_subtleCrypto = nullptr;
 
-    Bun::WriteBarrierList<JSC::JSPromise> m_aboutToBeNotifiedRejectedPromises;
+    // Rejected promises waiting for the "unhandledRejection" notification at
+    // the end of the tick. Each entry is either the JSPromise itself, or an
+    // AsyncContextFrame wrapping the JSPromise (as `callback`) together with
+    // the async context that was active when the promise was rejected.
+    Bun::WriteBarrierList<JSC::JSCell> m_aboutToBeNotifiedRejectedPromises;
 
 public:
     // While handleRejectedPromises() is iterating its drained snapshot, this

src/jsc/bindings/webcore/EventEmitter.cpp

@@ -124,6 +124,23 @@ bool EventEmitter::emit(const Identifier& eventType, const MarkedArgumentBuffer&
     return fireEventListeners(eventType, arguments);
 }
 
+bool EventEmitter::emit(const Identifier& eventType, const MarkedArgumentBuffer& arguments, WTF::NakedPtr<JSC::Exception>& returnedException)
+{
+    auto* data = eventTargetData();
+    if (!data)
+        return false;
+
+    auto* listenersVector = data->eventListenerMap.find(eventType);
+    if (!listenersVector) [[unlikely]]
+        return false;
+
+    bool prevFiringEventListeners = data->isFiringEventListeners;
+    data->isFiringEventListeners = true;
+    auto fired = innerInvokeEventListeners(eventType, *listenersVector, arguments, &returnedException);
+    data->isFiringEventListeners = prevFiringEventListeners;
+    return fired;
+}
+
 void EventEmitter::uncaughtExceptionInEventHandler()
 {
 }
@@ -206,7 +223,7 @@ bool EventEmitter::fireEventListeners(const Identifier& eventType, const MarkedA
 // Intentionally creates a copy of the listeners vector to avoid event listeners added after this point from being run.
 // Note that removal still has an effect due to the removed field in RegisteredEventListener.
 // https://dom.spec.whatwg.org/#concept-event-listener-inner-invoke
-bool EventEmitter::innerInvokeEventListeners(const Identifier& eventType, SimpleEventListenerVector listeners, const MarkedArgumentBuffer& arguments)
+bool EventEmitter::innerInvokeEventListeners(const Identifier& eventType, SimpleEventListenerVector listeners, const MarkedArgumentBuffer& arguments, WTF::NakedPtr<JSC::Exception>* returnedException)
 {
     Ref<EventEmitter> protectedThis(*this);
     ASSERT(!listeners.isEmpty());
@@ -252,6 +269,10 @@ bool EventEmitter::innerInvokeEventListeners(const Identifier& eventType, Simple
         auto* exception = exceptionPtr.get();
 
         if (exception) [[unlikely]] {
+            if (returnedException) {
+                *returnedException = exception;
+                return fired;
+            }
             auto errorIdentifier = vm.propertyNames->error;
             auto hasErrorListener = this->hasActiveEventListeners(errorIdentifier);
             if (!hasErrorListener || eventType == errorIdentifier) {

src/jsc/bindings/webcore/EventEmitter.h

@@ -56,6 +56,7 @@ class EventEmitter final : public ScriptWrappable, public CanMakeWeakPtr<EventEm
     WEBCORE_EXPORT bool removeAllListeners(const Identifier& eventType);
 
     WEBCORE_EXPORT bool emit(const Identifier&, const MarkedArgumentBuffer&);
+    WEBCORE_EXPORT bool emit(const Identifier&, const MarkedArgumentBuffer&, WTF::NakedPtr<JSC::Exception>& returnedException);
     WEBCORE_EXPORT void uncaughtExceptionInEventHandler();
 
     WEBCORE_EXPORT Vector<Identifier> getEventNames();
@@ -107,7 +108,7 @@ class EventEmitter final : public ScriptWrappable, public CanMakeWeakPtr<EventEm
     {
     }
 
-    bool innerInvokeEventListeners(const Identifier&, SimpleEventListenerVector, const MarkedArgumentBuffer& arguments);
+    bool innerInvokeEventListeners(const Identifier&, SimpleEventListenerVector, const MarkedArgumentBuffer& arguments, WTF::NakedPtr<JSC::Exception>* returnedException = nullptr);
     void invalidateEventListenerRegions();
 
     EventEmitterData m_eventTargetData;

src/jsc/lib.rs

@@ -997,7 +997,7 @@ pub use abort_signal::{AbortSignal, AbortSignalRef};
 // re-exported here so `crate::VM` and `crate::vm::VM` name the same nominal
 // type (and likewise for `JSGlobalObject`). Both structs carry `UnsafeCell`
 // so `&T → *mut T` for FFI is sound under Stacked Borrows.
-pub use self::js_global_object::{GlobalRef, JSGlobalObject};
+pub use self::js_global_object::{ClearedAsyncContextScope, GlobalRef, JSGlobalObject};
 pub use self::vm::{HeapType, Lock as ApiLock, VM};
 
 /// Options for `JSGlobalObject::validate_integer_range` / `validate_bigint_range`.

src/jsc/virtual_machine_exports.rs

@@ -154,6 +154,9 @@ pub fn handle_rejected_promise(global: &JSGlobalObject, promise: &mut JSPromise)
     }
 
     jsc_vm.unhandled_rejection(global, result, promise.to_js());
+    // The caller emits this event with the promise's async context installed;
+    // GC (and the finalizers it may run) must not observe it.
+    let _scope = crate::ClearedAsyncContextScope::new(global);
     jsc_vm.auto_garbage_collect();
 }