From 1a55dbf619c090dd612ff505476aa1cf2002d3b6 Mon Sep 17 00:00:00 2001
From: Dmitry Vasilev <dmitry.vasil@gmail.com>
Date: Mon, 19 Dec 2022 19:05:36 +0800
Subject: [PATCH] wrap promises from external calls

---
 src/eval.js  | 35 ++++++++++++++++++++---------------
 test/test.js | 22 ++++++++++++++++++++++
 2 files changed, 42 insertions(+), 15 deletions(-)

diff --git a/src/eval.js b/src/eval.js
index fecb2b7..c4f08cd 100644
--- a/src/eval.js
+++ b/src/eval.js
@@ -372,6 +372,24 @@ export const eval_modules = (
       }
     }
 
+    const set_promise_status = value => {
+      if(value instanceof Promise.Original) {
+        // record stack for async calls, so expand calls works sync
+        set_record_call()
+        return value
+          .then(v => {
+            value.status = {ok: true, value: v}
+            return v
+          })
+          .catch(e => {
+            value.status = {ok: false, error: e}
+            throw e
+          })
+      } else {
+        return value
+      }
+    }
+
     const __with_restore_children = async value => {
       // children is an array of child calls for current function call. But it
       // can be null to save one empty array allocation in case it has no child
@@ -421,20 +439,7 @@ export const eval_modules = (
         try {
           value = fn(...args)
           ok = true
-          if(value instanceof Promise.Original) {
-            set_record_call()
-            return value
-              .then(v => {
-                value.status = {ok: true, value: v}
-                return v
-              })
-              .catch(e => {
-                value.status = {ok: false, error: e}
-                throw e
-              })
-          } else {
-            return value
-          }
+          return set_promise_status(value)
         } catch(_error) {
           ok = false
           error = _error
@@ -534,7 +539,7 @@ export const eval_modules = (
           value = undefined
         }
         ok = true
-        return value
+        return set_promise_status(value)
       } catch(_error) {
         ok = false
         error = _error
diff --git a/test/test.js b/test/test.js
index 693b60d..c845ad1 100644
--- a/test/test.js
+++ b/test/test.js
@@ -2767,7 +2767,29 @@ const y = x()`
     assert_equal(logs, [1, 2])
   }),
 
+  test('async/await external async fn', async () => {
+    await assert_code_evals_to_async(
+      `
+        const AsyncFunction = 
+          new Function('return (async () => {}).constructor')()
+        const async_fn = new AsyncFunction('return 1')
+        await async_fn()
+      `,
+      1
+    )
+  }),
+
+  // TODO
   /*
+  test('async/await Promise.then creates subcall', async () => {
+    const i = await test_initial_state_async(`
+      await Promise.resolve(1).then(x => {
+      })
+    `)
+    console.log('i', root_calltree_node(i))
+
+  }),
+
   test('async/await bug', async () => {
     const code = `
       const f = async () => {