feat(nuxt): native async-context in vue's withAsyncContext (#23526)

packages/nuxt/src/app/composables/asyncContext.ts

@@ -0,0 +1,9 @@
+// @ts-expect-error withAsyncContext is internal API
+import { getCurrentInstance, withAsyncContext as withVueAsyncContext } from 'vue'
+
+export function withAsyncContext (fn: () => PromiseLike<unknown>) {
+  return withVueAsyncContext(() => {
+    const nuxtApp = getCurrentInstance()?.appContext.app.$nuxt
+    return nuxtApp ? nuxtApp.runWithContext(fn) : fn()
+  })
+}

packages/nuxt/src/core/plugins/async-context.ts

@@ -1,46 +1,21 @@
 import { createUnplugin } from 'unplugin'
 import MagicString from 'magic-string'
 import type { Nuxt } from '@nuxt/schema'
-import type { Node } from 'estree-walker'
-import { walk } from 'estree-walker'
-import type { BlockStatement } from 'estree'
 import { isVue } from '../utils'
 
 export const AsyncContextInjectionPlugin = (nuxt: Nuxt) => createUnplugin(() => {
   return {
-    name: 'nuxt:async-context-injection',
+    name: 'nuxt:vue-async-context',
     transformInclude (id) {
       return isVue(id, { type: ['template', 'script'] })
     },
     transform (code) {
+      if (!code.includes('_withAsyncContext')) {
+        return
+      }
       const s = new MagicString(code)
-
+      s.prepend('import { withAsyncContext as _withAsyncContext } from "#app/composables/asyncContext";\n')
-      let importName: string
+      s.replace(/withAsyncContext as _withAsyncContext,?/, '')
-
-      walk(this.parse(code) as Node, {
-        enter (node) {
-          // only interested in calls of defineComponent function
-          if (node.type === 'ImportDeclaration' && node.source.value === 'vue') {
-            importName = importName ?? node.specifiers.find(s => s.type === 'ImportSpecifier' && s.imported.name === 'defineComponent')?.local.name
-          }
-          // we only want to transform `async setup()` functions
-          if (node.type === 'CallExpression' && node.callee.type === 'Identifier' && node.callee.name === importName) {
-            walk(node, {
-              enter (setup) {
-                if (setup.type === 'Property' && setup.key.type === 'Identifier' && setup.key.name === 'setup') {
-                  if (setup.value.type === 'FunctionExpression' && setup.value.async) {
-                    const body: BlockStatement = setup.value.body
-                    const { start, end } = body as BlockStatement & { start: number, end: number }
-                    s.appendLeft(start, '{ return useNuxtApp().runWithContext(async () => ')
-                    s.appendRight(end, ') }')
-                  }
-                }
-              }
-            })
-          }
-        }
-      })
-
       if (s.hasChanged()) {
         return {
           code: s.toString(),