error: recover async stack frames dropped under AsyncLocalStorage

src/jsc/bindings/AsyncStackTrace.cpp

@@ -16,9 +16,69 @@
 #include <JavaScriptCore/Options.h>
 #include <JavaScriptCore/StackFrame.h>
 #include <JavaScriptCore/UnlinkedCodeBlock.h>
+#include <JavaScriptCore/VMEntryRecord.h>
 
 using namespace JSC;
 
+// dynamicDowncast(JSValue)'s isCell() is true for the empty value on JSVALUE64
+// (null cell deref); asyncStackTraceContext() and reaction getters return empty.
+template<typename T>
+static inline T* cellAs(JSValue v)
+{
+    return (v && v.isCell()) ? dynamicDowncast<T>(v.asCell()) : nullptr;
+}
+
+template<typename T>
+static inline bool dynamicCastValue(JSValue v, T** out)
+{
+    *out = cellAs<T>(v);
+    return *out != nullptr;
+}
+
+static BytecodeIndex computeGeneratorBytecodeIndex(CodeBlock* codeBlock, JSAsyncFunctionGenerator* generator)
+{
+    BytecodeIndex bytecodeIndex(0);
+    JSValue stateValue = generator->internalField(JSAsyncFunctionGenerator::Field::State).get();
+    if (stateValue.isInt32()) {
+        int32_t state = stateValue.asInt32();
+        size_t numberOfJumpTables = codeBlock->numberOfUnlinkedSwitchJumpTables();
+        if (state > 0 && numberOfJumpTables > 0) {
+            size_t lastTableIndex = numberOfJumpTables - 1;
+            const UnlinkedSimpleJumpTable& jumpTable = codeBlock->unlinkedSwitchJumpTable(lastTableIndex);
+            int32_t offset = jumpTable.offsetForValue(state);
+            if (offset)
+                bytecodeIndex = BytecodeIndex(offset);
+        }
+    }
+    return bytecodeIndex;
+}
+
+static bool appendGeneratorFrame(VM& vm, JSCell* owner, JSAsyncFunctionGenerator* generator, WTF::Vector<StackFrame>& results)
+{
+    auto* asyncFunction = cellAs<JSFunction>(generator->next());
+    if (!asyncFunction || asyncFunction->isHostOrPrivateBuiltinFunction())
+        return false;
+    FunctionExecutable* executable = asyncFunction->jsExecutable();
+    if (!executable)
+        return false;
+    if (CodeBlock* codeBlock = executable->codeBlockForCall()) {
+        BytecodeIndex bytecodeIndex = computeGeneratorBytecodeIndex(codeBlock, generator);
+        results.append(StackFrame(vm, owner, asyncFunction, codeBlock, bytecodeIndex, /* isAsyncFrame */ true));
+    } else {
+        results.append(StackFrame(vm, owner, asyncFunction, /* isAsyncFrame */ true));
+    }
+    return true;
+}
+
+// resolveWithInternalMicrotaskForAsyncAwait wraps await contexts as
+// InternalFieldTuple(context, asyncContext) when ALS is active; field 0 is the real cell.
+static inline JSValue unwrapAsyncContextTuple(JSValue v)
+{
+    if (auto* tuple = cellAs<InternalFieldTuple>(v))
+        return tuple->getInternalField(0);
+    return v;
+}
+
 // Walk a promise's reaction chain to find the async generators awaiting it,
 // and collect them as async StackFrames. Used when an error is created from
 // native code at the top of the event loop (e.g. run_from_js_thread in node_fs.rs)
@@ -35,20 +95,8 @@ static void collectAsyncStackFramesFromPromise(JSC::VM& vm, JSC::JSCell* owner,
 
     JSC::AssertNoGC assertNoGC;
 
-    auto dynamicCastValue = []<typename T>(JSC::JSValue v, T** out) -> bool {
-        if (!v || !v.isCell())
-            return false;
-        *out = dynamicDowncast<T>(v.asCell());
-        return *out != nullptr;
-    };
-
     auto unwrapGeneratorFromContext = [&](JSC::JSValue context) -> JSC::JSAsyncFunctionGenerator* {
-        JSC::InternalFieldTuple* tuple = nullptr;
-        if (dynamicCastValue(context, &tuple))
-            context = tuple->getInternalField(0);
-        JSC::JSAsyncFunctionGenerator* generator = nullptr;
-        dynamicCastValue(context, &generator);
-        return generator;
+        return cellAs<JSC::JSAsyncFunctionGenerator>(unwrapAsyncContextTuple(context));
     };
 
     // Walk reaction->context → generator. If context is not a generator (e.g.
@@ -104,43 +152,9 @@ static void collectAsyncStackFramesFromPromise(JSC::VM& vm, JSC::JSCell* owner,
         return nullptr;
     };
 
-    auto computeBytecodeIndex = [&](JSC::CodeBlock* codeBlock, JSC::JSAsyncFunctionGenerator* generator) -> JSC::BytecodeIndex {
-        JSC::BytecodeIndex bytecodeIndex(0);
-        JSC::JSValue stateValue = generator->internalField(JSC::JSAsyncFunctionGenerator::Field::State).get();
-        if (stateValue.isInt32()) {
-            int32_t state = stateValue.asInt32();
-            size_t numberOfJumpTables = codeBlock->numberOfUnlinkedSwitchJumpTables();
-            if (state > 0 && numberOfJumpTables > 0) {
-                size_t lastTableIndex = numberOfJumpTables - 1;
-                const JSC::UnlinkedSimpleJumpTable& jumpTable = codeBlock->unlinkedSwitchJumpTable(lastTableIndex);
-                int32_t offset = jumpTable.offsetForValue(state);
-                if (offset)
-                    bytecodeIndex = JSC::BytecodeIndex(offset);
-            }
-        }
-        return bytecodeIndex;
-    };
-
-    auto appendFrame = [&](JSC::JSAsyncFunctionGenerator* generator) {
-        JSC::JSFunction* asyncFunction = nullptr;
-        if (!dynamicCastValue(generator->next(), &asyncFunction))
-            return;
-        if (asyncFunction->isHostOrPrivateBuiltinFunction())
-            return;
-        JSC::FunctionExecutable* executable = asyncFunction->jsExecutable();
-        if (!executable)
-            return;
-        if (JSC::CodeBlock* codeBlock = executable->codeBlockForCall()) {
-            JSC::BytecodeIndex bytecodeIndex = computeBytecodeIndex(codeBlock, generator);
-            results.append(JSC::StackFrame(vm, owner, asyncFunction, codeBlock, bytecodeIndex, /* isAsyncFrame */ true));
-        } else {
-            results.append(JSC::StackFrame(vm, owner, asyncFunction, /* isAsyncFrame */ true));
-        }
-    };
-
     JSC::JSAsyncFunctionGenerator* gen = getAwaitingGenerator(promise);
     while (gen && results.size() < maxStackSize) {
-        appendFrame(gen);
+        appendGeneratorFrame(vm, owner, gen, results);
         JSC::JSPromise* returnPromise = nullptr;
         if (!dynamicCastValue(gen->context(), &returnPromise))
             break;
@@ -176,3 +190,72 @@ extern "C" void Bun__attachAsyncStackFromPromise(JSC::JSGlobalObject* globalObje
 
     instance->setStackFrames(vm, WTF::move(frames));
 }
+
+// VM::onAppendStackTrace hook, installed lazily by jsSetAsyncHooksEnabled so the
+// no-ALS path is untouched. Interpreter::getAsyncStackTrace's getParentGenerator
+// casts asyncStackTraceContext() to JSAsyncFunctionGenerator and fails on the
+// InternalFieldTuple ALS puts there, dropping every `at async` frame. Run the same
+// walk with and without the unwrap in lockstep and append only the frames the
+// un-unwrapped walk misses, so hops where ALS was inactive aren't doubled. See
+// oven-sh/WebKit#347 for the in-tree unwrap that will obsolete this hook.
+void Bun::appendAsyncLocalStorageStackFrames(VM& vm, JSCell* owner, Vector<StackFrame>& results, size_t maxToAppend)
+{
+    if (!maxToAppend || !Options::useAsyncStackTrace())
+        return;
+
+    AssertNoGC assertNoGC;
+
+    JSAsyncFunctionGenerator* origin = nullptr;
+    for (EntryFrame* entryFrame = vm.topEntryFrame; entryFrame;) {
+        VMEntryRecord* record = vmEntryRecord(entryFrame);
+        if (auto* generator = dynamicDowncast<JSAsyncFunctionGenerator>(record->m_context)) {
+            origin = generator;
+            break;
+        }
+        entryFrame = record->prevTopEntryFrame();
+    }
+    if (!origin)
+        return;
+
+    auto promiseContext = [](JSPromise* p, bool unwrap) -> JSValue {
+        if (!p)
+            return {};
+        JSValue context = p->asyncStackTraceContext();
+        return unwrap ? unwrapAsyncContextTuple(context) : context;
+    };
+
+    // Mirrors Interpreter::getAsyncStackTrace's getParentGenerator for direct
+    // await and Promise.race. Promise.all/allSettled/any use JSPromiseCombinatorsGlobalContext
+    // (private header, not forwarded) so the chain still ends at those hops under ALS.
+    auto getParentGenerator = [&](JSAsyncFunctionGenerator* gen, bool unwrap) -> JSAsyncFunctionGenerator* {
+        auto* returnPromise = cellAs<JSPromise>(gen->internalField(JSAsyncFunctionGenerator::Field::Context).get());
+        JSValue context = promiseContext(returnPromise, unwrap);
+        if (!context)
+            return nullptr;
+        if (auto* generator = cellAs<JSAsyncFunctionGenerator>(context))
+            return generator;
+        if (auto* promise = cellAs<JSPromise>(context))
+            return cellAs<JSAsyncFunctionGenerator>(promiseContext(promise, unwrap));
+        return nullptr;
+    };
+
+    size_t appended = 0;
+    bool jscStopped = false;
+    JSAsyncFunctionGenerator* current = origin;
+    for (unsigned hops = 0; hops < 256 && appended < maxToAppend; hops++) {
+        JSAsyncFunctionGenerator* next = getParentGenerator(current, /* unwrap */ true);
+        if (!next)
+            break;
+        if (!jscStopped) {
+            if (getParentGenerator(current, /* unwrap */ false))
+                current = next;
+            else
+                jscStopped = true;
+        }
+        if (jscStopped) {
+            if (appendGeneratorFrame(vm, owner, next, results))
+                appended++;
+            current = next;
+        }
+    }
+}

src/jsc/bindings/AsyncStackTrace.h

@@ -11,6 +11,10 @@
 // ErrorInstance with no stack of its own; no-op otherwise. Never throws.
 extern "C" void Bun__attachAsyncStackFromPromise(JSC::JSGlobalObject*, JSC::EncodedJSValue errorValue, JSC::JSPromise*);
 
+namespace JSC {
+class StackFrame;
+}
+
 namespace Bun {
 
 // C++ convenience wrapper over Bun__attachAsyncStackFromPromise.
@@ -19,4 +23,8 @@ inline void attachAsyncStackFromPromise(JSC::JSGlobalObject* globalObject, JSC::
     Bun__attachAsyncStackFromPromise(globalObject, JSC::JSValue::encode(error), promise);
 }
 
+// VM::onAppendStackTrace hook: appends the `at async` frames JSC's own walk drops when
+// AsyncLocalStorage wraps await contexts in InternalFieldTuple. See AsyncStackTrace.cpp.
+void appendAsyncLocalStorageStackFrames(JSC::VM&, JSC::JSCell* owner, WTF::Vector<JSC::StackFrame>&, size_t maxToAppend);
+
 } // namespace Bun

src/jsc/bindings/NodeAsyncHooks.cpp

@@ -4,6 +4,7 @@
 #include "JavaScriptCore/ObjectConstructor.h"
 #include "JavaScriptCore/ArrayConstructor.h"
 
+#include "AsyncStackTrace.h"
 #include "ZigGlobalObject.h"
 
 namespace Bun {
@@ -27,7 +28,11 @@ JSC_DEFINE_HOST_FUNCTION(jsCleanupLater, (JSC::JSGlobalObject * globalObject, JS
 JSC_DEFINE_HOST_FUNCTION(jsSetAsyncHooksEnabled, (JSC::JSGlobalObject * globalObject, JSC::CallFrame* callFrame))
 {
     ASSERT(callFrame->argumentCount() == 1);
-    globalObject->setAsyncContextTrackingEnabled(callFrame->uncheckedArgument(0).toBoolean(globalObject));
+    bool enabled = callFrame->uncheckedArgument(0).toBoolean(globalObject);
+    globalObject->setAsyncContextTrackingEnabled(enabled);
+    auto& vm = JSC::getVM(globalObject);
+    if (enabled && !vm.onAppendStackTrace())
+        vm.setOnAppendStackTrace(Bun::appendAsyncLocalStorageStackFrames);
     return JSC::JSValue::encode(JSC::jsUndefined());
 }
 

test/regression/issue/24003.test.ts

@@ -0,0 +1,133 @@
+// https://github.com/oven-sh/bun/issues/24003
+// Async stack traces were missing their `at async <fn>` frames when an
+// AsyncLocalStorage store was active at the point of `await`.
+import { describe, expect, test } from "bun:test";
+import { bunEnv, bunExe } from "harness";
+
+async function run(source: string): Promise<string> {
+  await using proc = Bun.spawn({
+    cmd: [bunExe(), "-e", source],
+    env: bunEnv,
+    stdout: "pipe",
+    stderr: "pipe",
+  });
+  const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+  expect(stderr).toBe("");
+  expect(exitCode).toBe(0);
+  return stdout;
+}
+
+function asyncFrames(stack: string): string[] {
+  return stack
+    .split("\n")
+    .filter(l => l.includes("at async "))
+    .map(l => l.trim().replace(/ \(.*\)$/, ""));
+}
+
+describe.concurrent("issue #24003", () => {
+  test("async stack frames under AsyncLocalStorage.run()", async () => {
+    const stdout = await run(`
+    const { AsyncLocalStorage } = require("node:async_hooks");
+    const als = new AsyncLocalStorage();
+    async function fn3() { await 0; throw new Error("boom"); }
+    async function fn2() { await fn3(); }
+    async function fn1() { await fn2(); }
+    async function main() {
+      try { await als.run({}, () => fn1()); }
+      catch (e) { console.log(e.stack); }
+    }
+    main();
+  `);
+    expect(stdout).toContain("at fn3");
+    expect(asyncFrames(stdout)).toEqual(["at async fn2", "at async fn1", "at async main"]);
+  });
+
+  test("async stack frames under AsyncLocalStorage.enterWith()", async () => {
+    const stdout = await run(`
+    const { AsyncLocalStorage } = require("node:async_hooks");
+    const als = new AsyncLocalStorage();
+    als.enterWith({ x: 1 });
+    async function inner() { await 0; console.log(new Error("boom").stack); }
+    async function outer() { await inner(); }
+    async function main() { await outer(); }
+    main();
+  `);
+    expect(stdout).toContain("at inner");
+    expect(asyncFrames(stdout)).toEqual(["at async outer", "at async main"]);
+  });
+
+  test("async stack frames through Promise.race under AsyncLocalStorage", async () => {
+    const stdout = await run(`
+    const { AsyncLocalStorage } = require("node:async_hooks");
+    const als = new AsyncLocalStorage();
+    async function leaf() { await 0; throw new Error("boom"); }
+    async function viaRace() { await Promise.race([leaf()]); }
+    async function caller() { await viaRace(); }
+    async function main() {
+      try { await als.run({}, () => caller()); }
+      catch (e) { console.log(e.stack); }
+    }
+    main();
+  `);
+    expect(stdout).toContain("at leaf");
+    expect(asyncFrames(stdout)).toEqual(["at async viaRace", "at async caller", "at async main"]);
+  });
+
+  // The hook is installed lazily from the AsyncLocalStorage constructor; construct one
+  // without entering a store so the lockstep walk runs but JSC succeeds at every hop.
+  test("async stack frames are not duplicated when no AsyncLocalStorage store is active", async () => {
+    const stdout = await run(`
+    new (require("node:async_hooks").AsyncLocalStorage)();
+    async function fn3() { await 0; throw new Error("boom"); }
+    async function fn2() { await fn3(); }
+    async function fn1() { await fn2(); }
+    async function main() {
+      try { await fn1(); }
+      catch (e) { console.log(e.stack); }
+    }
+    main();
+  `);
+    expect(stdout).toContain("at fn3");
+    expect(asyncFrames(stdout)).toEqual(["at async fn2", "at async fn1", "at async main"]);
+  });
+
+  test("async stack frames through Promise.race are not duplicated when no AsyncLocalStorage store is active", async () => {
+    const stdout = await run(`
+    new (require("node:async_hooks").AsyncLocalStorage)();
+    async function leaf() { await 0; throw new Error("boom"); }
+    async function viaRace() { await Promise.race([leaf()]); }
+    async function caller() { await viaRace(); }
+    async function main() {
+      try { await caller(); }
+      catch (e) { console.log(e.stack); }
+    }
+    main();
+  `);
+    expect(stdout).toContain("at leaf");
+    expect(asyncFrames(stdout)).toEqual(["at async viaRace", "at async caller", "at async main"]);
+  });
+
+  test("async stack frames when AsyncLocalStorage is entered mid-chain", async () => {
+    // fn2 awaits fn3 before ALS is entered, so JSC's own walk finds fn2; fn1
+    // awaits fn2 after enterWith, so JSC stops there and the hook must supply
+    // fn1/main without duplicating fn2.
+    const stdout = await run(`
+    const { AsyncLocalStorage } = require("node:async_hooks");
+    const als = new AsyncLocalStorage();
+    async function fn3() { await 0; throw new Error("boom"); }
+    async function fn2() { await fn3(); }
+    async function fn1() {
+      const p = fn2();
+      als.enterWith({});
+      await p;
+    }
+    async function main() {
+      try { await fn1(); }
+      catch (e) { console.log(e.stack); }
+    }
+    main();
+  `);
+    expect(stdout).toContain("at fn3");
+    expect(asyncFrames(stdout)).toEqual(["at async fn2", "at async fn1", "at async main"]);
+  });
+});