feat(sfc): auto restore current instance after await statements in async setup()

packages/runtime-core/__tests__/apiSetupHelpers.spec.ts

@@ -1,9 +1,13 @@
 import {
+  ComponentInternalInstance,
   defineComponent,
+  getCurrentInstance,
   h,
   nodeOps,
+  onMounted,
   render,
-  SetupContext
+  SetupContext,
+  Suspense
 } from '@vue/runtime-test'
 import {
   defineEmits,
@@ -12,7 +16,8 @@ import {
   withDefaults,
   useAttrs,
   useSlots,
-  mergeDefaults
+  mergeDefaults,
+  withAsyncContext
 } from '../src/apiSetupHelpers'
 
 describe('SFC <script setup> helpers', () => {
@@ -89,4 +94,39 @@ describe('SFC <script setup> helpers', () => {
       `props default key "foo" has no corresponding declaration`
     ).toHaveBeenWarned()
   })
+
+  test('withAsyncContext', async () => {
+    const spy = jest.fn()
+
+    let beforeInstance: ComponentInternalInstance | null = null
+    let afterInstance: ComponentInternalInstance | null = null
+    let resolve: (msg: string) => void
+
+    const Comp = defineComponent({
+      async setup() {
+        beforeInstance = getCurrentInstance()
+        const msg = await withAsyncContext(
+          new Promise(r => {
+            resolve = r
+          })
+        )
+        // register the lifecycle after an await statement
+        onMounted(spy)
+        afterInstance = getCurrentInstance()
+        return () => msg
+      }
+    })
+
+    const root = nodeOps.createElement('div')
+    render(h(() => h(Suspense, () => h(Comp))), root)
+
+    expect(spy).not.toHaveBeenCalled()
+    resolve!('hello')
+    // wait a macro task tick for all micro ticks to resolve
+    await new Promise(r => setTimeout(r))
+    // mount hook should have been called
+    expect(spy).toHaveBeenCalled()
+    // should retain same instance before/after the await call
+    expect(beforeInstance).toBe(afterInstance)
+  })
 })

packages/runtime-core/src/apiSetupHelpers.ts

@@ -1,7 +1,8 @@
 import {
   getCurrentInstance,
   SetupContext,
-  createSetupContext
+  createSetupContext,
+  setCurrentInstance
 } from './component'
 import { EmitFn, EmitsOptions } from './componentEmits'
 import {
@@ -226,3 +227,17 @@ export function mergeDefaults(
   }
   return props
 }
+
+/**
+ * Runtime helper for storing and resuming current instance context in
+ * async setup().
+ * @internal
+ */
+export async function withAsyncContext<T>(
+  awaitable: T | Promise<T>
+): Promise<T> {
+  const ctx = getCurrentInstance()
+  const res = await awaitable
+  setCurrentInstance(ctx)
+  return res
+}

packages/runtime-core/src/index.ts

@@ -48,12 +48,14 @@ export { defineAsyncComponent } from './apiAsyncComponent'
 // <script setup> API ----------------------------------------------------------
 
 export {
+  // macros runtime, for warnings only
   defineProps,
   defineEmits,
   defineExpose,
   withDefaults,
   // internal
   mergeDefaults,
+  withAsyncContext,
   // deprecated
   defineEmit,
   useContext