2 files changed, 63 insertions, 8 deletions

diff --git a/src/bun.js/async_hooks.exports.js b/src/bun.js/async_hooks.exports.js
index 21792e295..9dfff7ef1 100644
--- a/src/bun.js/async_hooks.exports.js
+++ b/src/bun.js/async_hooks.exports.js
@@ -1,3 +1,8 @@
+var drainMicrotasks = () => {
+  ({ drainMicrotasks } = import.meta.require("bun:jsc"));
+  drainMicrotasks();
+};
+
 const warnOnce = fn => {
   let warned = false;
   return (...args) => {
@@ -16,6 +21,7 @@ class AsyncLocalStorage {
 
   constructor() {
     this._enabled = false;
+    this.#store = null;
   }
 
   enterWith(store) {
@@ -28,19 +34,28 @@ class AsyncLocalStorage {
   exit(cb, ...args) {
     this.#store = null;
     notImplemented();
-    cb(...args);
+    typeof cb === "function" && cb(...args);
   }
 
   run(store, callback, ...args) {
     if (typeof callback !== "function") throw new TypeError("ERR_INVALID_CALLBACK");
     const prev = this.#store;
-    this.enterWith(store);
-
-    try {
-      return callback(...args);
-    } finally {
-      this.#store = prev;
+    var result, err;
+    process.nextTick(() => {
+      this.enterWith(store);
+      try {
+        result = callback(...args);
+      } catch (e) {
+        err = e;
+      } finally {
+        this.#store = prev;
+      }
+    });
+    drainMicrotasks();
+    if (typeof err !== "undefined") {
+      throw err;
     }
+    return result;
   }
 
   getStore() {
@@ -153,7 +168,17 @@ class AsyncResource {
 
   runInAsyncScope(fn, ...args) {
     notImplemented();
-    process.nextTick(fn, ...args);
+    var result, err;
+    process.nextTick(fn => {
+      try {
+        result = fn(...args);
+      } catch (err2) {
+        err = err2;
+      }
+    }, fn);
+    drainMicrotasks();
+    if (err) throw err;
+    return result;
   }
 
   asyncId() {
diff --git a/test/bun.js/disabled-module.test.js b/test/bun.js/disabled-module.test.js
index 43b08731d..c12676959 100644
--- a/test/bun.js/disabled-module.test.js
+++ b/test/bun.js/disabled-module.test.js
@@ -6,3 +6,33 @@ test("not implemented yet module masquerades as undefined and throws an error",
   expect(typeof worker_threads).toBe("undefined");
   expect(typeof worker_threads.getEnvironmentData).toBe("undefined");
 });
+
+test("AsyncLocalStorage polyfill", () => {
+  const { AsyncLocalStorage } = import.meta.require("async_hooks");
+
+  const store = new AsyncLocalStorage();
+  var called = false;
+  expect(store.getStore()).toBe(null);
+  store.run({ foo: "bar" }, () => {
+    expect(store.getStore()).toEqual({ foo: "bar" });
+    called = true;
+  });
+  expect(store.getStore()).toBe(null);
+  expect(called).toBe(true);
+});
+
+test("AsyncResource polyfill", () => {
+  const { AsyncResource } = import.meta.require("async_hooks");
+
+  const resource = new AsyncResource("test");
+  var called = false;
+  resource.runInAsyncScope(
+    () => {
+      called = true;
+    },
+    null,
+    "foo",
+    "bar",
+  );
+  expect(called).toBe(true);
+});