fix(sfc): fix <script setup> async context preservation logic

fix #4050

packages/compiler-sfc/__tests__/compileScript.spec.ts

@@ -847,6 +847,9 @@ const emit = defineEmits(['a', 'b'])
       const { content } = compile(`<script setup>${code}</script>`, {
         refSugar: true
       })
+      if (shouldAsync) {
+        expect(content).toMatch(`let __temp, __restore`)
+      }
       expect(content).toMatch(`${shouldAsync ? `async ` : ``}setup(`)
       if (typeof expected === 'string') {
         expect(content).toMatch(expected)
@@ -856,28 +859,35 @@ const emit = defineEmits(['a', 'b'])
     }
 
     test('expression statement', () => {
-      assertAwaitDetection(`await foo`, `await _withAsyncContext(foo)`)
+      assertAwaitDetection(
+        `await foo`,
+        `;(([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore())`
+      )
     })
 
     test('variable', () => {
       assertAwaitDetection(
         `const a = 1 + (await foo)`,
-        `1 + (await _withAsyncContext(foo))`
+        `1 + ((([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore(),__temp))`
       )
     })
 
     test('ref', () => {
       assertAwaitDetection(
         `ref: a = 1 + (await foo)`,
-        `1 + (await _withAsyncContext(foo))`
+        `1 + ((([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore(),__temp))`
       )
     })
 
     test('nested statements', () => {
       assertAwaitDetection(`if (ok) { await foo } else { await bar }`, code => {
         return (
-          code.includes(`await _withAsyncContext(foo)`) &&
-          code.includes(`await _withAsyncContext(bar)`)
+          code.includes(
+            `;(([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore())`
+          ) &&
+          code.includes(
+            `;(([__temp,__restore]=_withAsyncContext(()=>(bar))),__temp=await __temp,__restore())`
+          )
         )
       })
     })

packages/compiler-sfc/src/compileScript.ts

@@ -32,7 +32,8 @@ import {
   LabeledStatement,
   CallExpression,
   RestElement,
-  TSInterfaceBody
+  TSInterfaceBody,
+  AwaitExpression
 } from '@babel/types'
 import { walk } from 'estree-walker'
 import { RawSourceMap } from 'source-map'
@@ -487,6 +488,25 @@ export function compileScript(
     })
   }
 
+  /**
+   * await foo()
+   * -->
+   * (([__temp, __restore] = withAsyncContext(() => foo())),__temp=await __temp,__restore(),__temp)
+   */
+  function processAwait(node: AwaitExpression, isStatement: boolean) {
+    s.overwrite(
+      node.start! + startOffset,
+      node.argument.start! + startOffset,
+      `${isStatement ? `;` : ``}(([__temp,__restore]=${helper(
+        `withAsyncContext`
+      )}(()=>(`
+    )
+    s.appendLeft(
+      node.end! + startOffset,
+      `))),__temp=await __temp,__restore()${isStatement ? `` : `,__temp`})`
+    )
+  }
+
   function processRefExpression(exp: Expression, statement: LabeledStatement) {
     if (exp.type === 'AssignmentExpression') {
       const { left, right } = exp
@@ -949,17 +969,13 @@ export function compileScript(
       node.type.endsWith('Statement')
     ) {
       ;(walk as any)(node, {
-        enter(node: Node) {
-          if (isFunction(node)) {
+        enter(child: Node, parent: Node) {
+          if (isFunction(child)) {
             this.skip()
           }
-          if (node.type === 'AwaitExpression') {
+          if (child.type === 'AwaitExpression') {
             hasAwait = true
-            s.prependRight(
-              node.argument.start! + startOffset,
-              helper(`withAsyncContext`) + `(`
-            )
-            s.appendLeft(node.argument.end! + startOffset, `)`)
+            processAwait(child, parent.type === 'ExpressionStatement')
           }
         }
       })
@@ -1151,6 +1167,11 @@ export function compileScript(
   if (propsIdentifier) {
     s.prependRight(startOffset, `\nconst ${propsIdentifier} = __props`)
   }
+  // inject temp variables for async context preservation
+  if (hasAwait) {
+    const any = isTS ? `:any` : ``
+    s.prependRight(startOffset, `\nlet __temp${any}, __restore${any}\n`)
+  }
 
   const destructureElements =
     hasDefineExposeCall || !options.inlineTemplate ? [`expose`] : []