bun test: attribute late expect() calls of an abandoned test to that test instead of the next one

src/js/node/async_hooks.ts

@@ -22,6 +22,9 @@
 // each key is an AsyncLocalStorage object and the value is the associated value. There are a ton of
 // calls to $assert which will verify this invariant (only during bun-debug)
 //
+// The one other key that can appear is bun:test's AsyncContextRef (stored as both key and value, only
+// read from native code; see src/runtime/test_runner/AsyncContextRef.rs). Nothing in here touches it.
+//
 const setAsyncHooksEnabled = $newCppFunction("NodeAsyncHooks.cpp", "jsSetAsyncHooksEnabled", 1);
 const cleanupLater = $newCppFunction("NodeAsyncHooks.cpp", "jsCleanupLater", 0);
 const { validateFunction, validateString, validateObject } = require("internal/validators");
@@ -49,7 +52,7 @@ function assertValidAsyncContextArray(array: unknown): array is ReadonlyArray<an
   $assert(array.length > 0, "AsyncContextData should be undefined if empty, got", Bun.inspect(array, { depth: 1 }));
   for (var i = 0; i < array.length; i += 2) {
     $assert(
-      array[i] instanceof AsyncLocalStorage,
+      array[i] instanceof AsyncLocalStorage || isBunTestAsyncContextRef(array[i]),
       `Odd indexes in AsyncContextData should be an array of AsyncLocalStorage\nIndex %s was %s`,
       i,
       array[i],
@@ -58,12 +61,17 @@ function assertValidAsyncContextArray(array: unknown): array is ReadonlyArray<an
   return true;
 }
 
+// Only run during debug. The generated prototype's toStringTag is the class name.
+function isBunTestAsyncContextRef(key: any) {
+  return typeof key === "object" && key !== null && key[Symbol.toStringTag] === "AsyncContextRef";
+}
+
 // Only run during debug
 function debugFormatContextValue(value: ReadonlyArray<any> | undefined) {
   if (value === undefined) return "undefined";
   let str = "{\n";
   for (var i = 0; i < value.length; i += 2) {
-    str += `  ${value[i].__id__}: typeof = ${typeof value[i + 1]}\n`;
+    str += `  ${isBunTestAsyncContextRef(value[i]) ? "bun:test" : value[i].__id__}: typeof = ${typeof value[i + 1]}\n`;
   }
   str += "}";
   return str;

src/jsc/bindings/AsyncContextFrame.cpp

@@ -1,7 +1,10 @@
 #include "root.h"
 #include "ZigGlobalObject.h"
+#include "ZigGeneratedClasses.h"
 #include "AsyncContextFrame.h"
+#include <JavaScriptCore/ArgList.h>
 #include <JavaScriptCore/InternalFieldTuple.h>
+#include <JavaScriptCore/JSArray.h>
 
 #if ASSERT_ENABLED
 #include <JavaScriptCore/IntegrityInlines.h>
@@ -131,3 +134,155 @@ JSValue AsyncContextFrame::profiledCall(JSGlobalObject* global, JSValue function
     return AsyncContextFrame::call(global, functionObject, thisValue, args);
 }
 #undef ASYNCCONTEXTFRAME_CALL_IMPL
+
+// ── bun:test (src/runtime/test_runner/AsyncContextRef.rs) ──────────────────
+//
+// The context array ([key, value, ...], see node/async_hooks.ts) gets one extra
+// pair per test or hook invocation: (ref, ref), the invocation's AsyncContextRef.
+
+static bool isAsyncContextRef(JSValue value)
+{
+    return dynamicDowncast<WebCore::JSAsyncContextRef>(value) != nullptr;
+}
+
+static JSValue findAsyncContextRef(JSValue context)
+{
+    auto* array = dynamicDowncast<JSArray>(context);
+    if (!array)
+        return jsUndefined();
+    unsigned length = array->length();
+    for (unsigned i = 0; i < length; i += 2) {
+        if (!array->canGetIndexQuickly(i))
+            continue;
+        JSValue key = array->getIndexQuickly(i);
+        if (isAsyncContextRef(key))
+            return key;
+    }
+    return jsUndefined();
+}
+
+// A context the runner alone populated: without the runner there would be none.
+static bool holdsOnlyAsyncContextRefs(JSValue context)
+{
+    auto* array = dynamicDowncast<JSArray>(context);
+    if (!array)
+        return false;
+    unsigned length = array->length();
+    if (length == 0)
+        return false;
+    for (unsigned i = 0; i < length; i += 2) {
+        if (!array->canGetIndexQuickly(i) || !isAsyncContextRef(array->getIndexQuickly(i)))
+            return false;
+    }
+    return true;
+}
+
+// The caller checks for an exception afterwards.
+static void appendPairsWithoutRefs(JSGlobalObject* globalObject, JSValue context, MarkedArgumentBuffer& entries)
+{
+    auto& vm = JSC::getVM(globalObject);
+    auto scope = DECLARE_THROW_SCOPE(vm);
+    auto* array = dynamicDowncast<JSArray>(context);
+    if (!array)
+        return;
+    unsigned length = array->length();
+    entries.ensureCapacity(length + 2);
+    for (unsigned i = 0; i < length; i += 2) {
+        JSValue key = array->getIndex(globalObject, i);
+        RETURN_IF_EXCEPTION(scope, );
+        JSValue value = i + 1 < length ? array->getIndex(globalObject, i + 1) : jsUndefined();
+        RETURN_IF_EXCEPTION(scope, );
+        if (isAsyncContextRef(key))
+            continue;
+        entries.append(key);
+        entries.append(value);
+    }
+}
+
+// node:vm contexts copy this flag when created, so it is set before a test file loads.
+extern "C" [[ZIG_EXPORT(nothrow)]] void Bun__AsyncContextRef__enableTracking(JSC::JSGlobalObject* globalObject)
+{
+    globalObject->setAsyncContextTrackingEnabled(true);
+}
+
+// AsyncContextFrame::withAsyncContextIfNeeded for a callback being registered: the
+// registering invocation's (ref, ref) is not a context to capture. __enter drops it
+// from a context that is.
+extern "C" [[ZIG_EXPORT(nothrow)]] JSC::EncodedJSValue Bun__AsyncContextRef__withAsyncContextIfNeeded(JSC::JSGlobalObject* globalObject, JSC::EncodedJSValue callbackValue)
+{
+    if (holdsOnlyAsyncContextRefs(globalObject->m_asyncContextData.get()->getInternalField(0)))
+        return callbackValue;
+    return JSValue::encode(AsyncContextFrame::withAsyncContextIfNeeded(globalObject, JSValue::decode(callbackValue)));
+}
+
+// Returns what to invoke in place of `callback` so that it runs with its usual
+// context plus (ref, ref). A callback registered under a context (an
+// AsyncContextFrame) gets a new frame, which Bun__JSValue__call installs and
+// restores as usual. Any other callback gets the array installed in place, so
+// that, as before, what it does to the context (als.enterWith()) outlives it;
+// __leave then only removes the ref.
+extern "C" [[ZIG_EXPORT(zero_is_throw)]] JSC::EncodedJSValue Bun__AsyncContextRef__enter(JSC::JSGlobalObject* globalObject, JSC::EncodedJSValue callbackValue, JSC::EncodedJSValue refValue)
+{
+    auto& vm = JSC::getVM(globalObject);
+    auto scope = DECLARE_THROW_SCOPE(vm);
+
+    JSValue callback = JSValue::decode(callbackValue);
+    JSValue ref = JSValue::decode(refValue);
+    ASSERT(isAsyncContextRef(ref));
+    auto* slot = globalObject->m_asyncContextData.get();
+
+    auto* registrationFrame = dynamicDowncast<AsyncContextFrame>(callback);
+    JSValue previousContext = registrationFrame ? registrationFrame->context.get() : slot->getInternalField(0);
+
+    MarkedArgumentBuffer entries;
+    appendPairsWithoutRefs(globalObject, previousContext, entries);
+    RETURN_IF_EXCEPTION(scope, {});
+    entries.append(ref);
+    entries.append(ref);
+    if (entries.hasOverflowed()) [[unlikely]] {
+        throwOutOfMemoryError(globalObject, scope);
+        return {};
+    }
+    JSArray* context = constructArray(globalObject, static_cast<ArrayAllocationProfile*>(nullptr), entries);
+    RETURN_IF_EXCEPTION(scope, {});
+
+    globalObject->setAsyncContextTrackingEnabled(true);
+
+    if (registrationFrame)
+        return JSValue::encode(AsyncContextFrame::create(globalObject, registrationFrame->callback.get(), context));
+
+    slot->putInternalField(vm, 0, context);
+    return JSValue::encode(callback);
+}
+
+// Takes the refs out of whatever the callback left in the slot, keeping the rest
+// in effect as before. Always rebuilt: als.disable() splices the installed array
+// in place. After a frame the call restored the slot itself, so there is no ref.
+extern "C" [[ZIG_EXPORT(check_slow)]] void Bun__AsyncContextRef__leave(JSC::JSGlobalObject* globalObject)
+{
+    auto& vm = JSC::getVM(globalObject);
+    auto scope = DECLARE_THROW_SCOPE(vm);
+
+    auto* slot = globalObject->m_asyncContextData.get();
+    JSValue current = slot->getInternalField(0);
+    if (findAsyncContextRef(current).isUndefined())
+        return;
+    MarkedArgumentBuffer entries;
+    appendPairsWithoutRefs(globalObject, current, entries);
+    RETURN_IF_EXCEPTION(scope, );
+    if (entries.hasOverflowed()) [[unlikely]] {
+        throwOutOfMemoryError(globalObject, scope);
+        return;
+    }
+    JSValue remaining = jsUndefined();
+    if (!entries.isEmpty()) {
+        remaining = constructArray(globalObject, static_cast<ArrayAllocationProfile*>(nullptr), entries);
+        RETURN_IF_EXCEPTION(scope, );
+    }
+    slot->putInternalField(vm, 0, remaining);
+}
+
+extern "C" [[ZIG_EXPORT(nothrow)]] JSC::EncodedJSValue Bun__AsyncContextRef__current(JSC::JSGlobalObject* globalObject)
+{
+    return JSValue::encode(findAsyncContextRef(globalObject->m_asyncContextData.get()->getInternalField(0)));
+}

src/jsc/event_loop.rs

@@ -1138,26 +1138,6 @@ impl EventLoop {
         }
     }
 
-    /// Prefer `runCallbackWithResult` unless you really need to make sure that microtasks are drained.
-    pub fn run_callback_with_result_and_forcefully_drain_microtasks(
-        &mut self,
-        callback: JSValue,
-        global_object: &JSGlobalObject,
-        this_value: JSValue,
-        arguments: &[JSValue],
-    ) -> JsResult<JSValue> {
-        // Same gate as `run_callback`.
-        if global_object.has_exception() {
-            return Ok(JSValue::UNDEFINED);
-        }
-        let result = callback.call(global_object, this_value, arguments)?;
-        result.ensure_still_alive();
-        let jsc_vm = global_object.bun_vm().jsc_vm();
-        self.drain_microtasks_with_global(global_object, jsc_vm)
-            .map_err(|stopped| stopped.throw(global_object))?;
-        Ok(result)
-    }
-
     /// Keep one poll registered with the loop so `us_loop_run_bun_tick` parks
     /// instead of returning immediately on `num_polls == 0`.
     #[cfg(not(windows))]

src/jsc/generated_classes_list.rs

@@ -33,6 +33,7 @@ pub mod Classes {
     pub use crate::crypto::CryptoHasher;
     pub use crate::image as Image;
     pub use crate::shell::Interpreter as ShellInterpreter;
+    pub use crate::test_runner::async_context_ref::AsyncContextRef;
     pub use crate::test_runner::done_callback::DoneCallback;
     pub use crate::test_runner::expect::Expect;
     pub use crate::test_runner::expect::ExpectAny;

src/runtime/cli/test_command.rs

@@ -3276,6 +3276,8 @@ impl TestCommand {
             }
             // need to wake up so autoTick() doesn't wait for 16-100ms after loading the entrypoint
             vm.wakeup();
+            // Before the file loads: node:vm contexts copy this flag when they are created.
+            jsc::cpp::Bun__AsyncContextRef__enableTracking(vm.global());
             let promise = vm.load_entry_point_for_test_runner(file_path)?;
             // Only count the file once, not once per repeat
             if repeat_index == 0 {

src/runtime/test_runner/AsyncContextRef.rs

@@ -0,0 +1,64 @@
+//! Put into the async context for the duration of one test or hook invocation,
+//! so that the invocation's continuations can still be traced back to it.
+//!
+//! `expect()`, the snapshot matchers, `expect.assertions()` and hook registration
+//! belong to whatever entry the runner is executing when they are called. Once
+//! the runner has abandoned an invocation that is still running (timeout, or an
+//! unhandled error while it was awaiting), that would be the next entry;
+//! `RefData::abandoned` makes `bun_test::caller_ref` attribute such late calls
+//! to the abandoned invocation instead, which rejects them. Invocations that
+//! completed are never consulted, so callbacks they registered keep belonging to
+//! whichever entry runs them. Uncaught errors are not covered: by the time one is
+//! reported (`jest::on_unhandled_rejection`) the context it was thrown in is gone.
+//!
+//! The slot manipulation lives next to `AsyncContextFrame` (AsyncContextFrame.cpp).
+
+use bun_jsc::{JSGlobalObject, JSValue, JsClass as _, JsResult};
+
+use crate::test_runner::bun_test::RefDataPtr;
+
+#[bun_jsc::JsClass(no_construct, no_constructor)] // codegen wires to_js / from_js
+pub struct AsyncContextRef {
+    /// Owned `+1`, released in `finalize`.
+    r#ref: RefDataPtr,
+}
+
+impl AsyncContextRef {
+    // Codegen calls `finalize(Box<Self>)`; clippy::boxed_local is a false positive.
+    #[allow(clippy::boxed_local)]
+    pub fn finalize(self: Box<Self>) {
+        self.r#ref.deref(); // `RefPtr` has no `Drop`
+    }
+
+    /// Puts `refdata` (a `+1`, consumed) into the context `callback` is about to
+    /// run with, and returns what to invoke in its place. Pair with [`Self::leave`].
+    pub(crate) fn enter(global: &JSGlobalObject, callback: JSValue, refdata: RefDataPtr) -> JsResult<JSValue> {
+        let ref_js = AsyncContextRef { r#ref: refdata }.to_js(global);
+        let callable = bun_jsc::cpp::Bun__AsyncContextRef__enter(global, callback, ref_js);
+        ref_js.ensure_still_alive();
+        callable
+    }
+
+    /// Call once the callback returned, before its microtasks run.
+    pub(crate) fn leave(global: &JSGlobalObject) -> JsResult<()> {
+        bun_jsc::cpp::Bun__AsyncContextRef__leave(global)
+    }
+
+    /// A `+1` to the abandoned invocation the running JS descends from, if any.
+    pub(crate) fn abandoned_caller(global: &JSGlobalObject) -> Option<RefDataPtr> {
+        Self::abandoned_in_context(global).map(RefDataPtr::dupe_ref)
+    }
+
+    pub(crate) fn caller_is_abandoned(global: &JSGlobalObject) -> bool {
+        Self::abandoned_in_context(global).is_some()
+    }
+
+    /// To use right away: the borrow is of a wrapper that the context array
+    /// installed at this moment keeps alive.
+    fn abandoned_in_context(global: &JSGlobalObject) -> Option<&RefDataPtr> {
+        let this = Self::from_js(bun_jsc::cpp::Bun__AsyncContextRef__current(global))?;
+        // SAFETY: live payload of the wrapper `__current` just found in the context array.
+        let refdata: &RefDataPtr = unsafe { &(*this).r#ref };
+        refdata.abandoned.get().then_some(refdata)
+    }
+}

src/runtime/test_runner/Collection.rs

@@ -141,7 +141,7 @@ impl Collection {
             // `self.active_scope` in `step()` and mutated through).
             self.current_scope_callback_queue.push(QueuedDescribe {
                 active_scope: self.active_scope,
-                callback: DeprecatedStrong::init(cb),
+                callback: DeprecatedStrong::init(bun_test::keep_registration_async_context(cb)),
                 new_scope: NonNull::from(new_scope),
             });
         }
@@ -238,6 +238,7 @@ impl Collection {
                 buntest_strong,
                 global_this,
                 callback.get(),
+                None,
                 false,
                 RefDataValue::Collection { active_scope: previous_scope },
                 &Timespec::EPOCH,