From 93ace5548199e4c3fcc4dc07c2e18df8330bb1d8 Mon Sep 17 00:00:00 2001
From: Pooya Parsa <pyapar@gmail.com>
Date: Thu, 5 Oct 2023 16:46:53 +0200
Subject: [PATCH] feat(nuxt): native async-context in vue's `withAsyncContext`
 (#23526)

---
 .../nuxt/src/app/composables/asyncContext.ts  |  9 +++++
 .../nuxt/src/core/plugins/async-context.ts    | 37 +++----------------
 2 files changed, 15 insertions(+), 31 deletions(-)
 create mode 100644 packages/nuxt/src/app/composables/asyncContext.ts

diff --git a/packages/nuxt/src/app/composables/asyncContext.ts b/packages/nuxt/src/app/composables/asyncContext.ts
new file mode 100644
index 0000000000..5fa88df529
--- /dev/null
+++ b/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()
+  })
+}
diff --git a/packages/nuxt/src/core/plugins/async-context.ts b/packages/nuxt/src/core/plugins/async-context.ts
index 0ef43e66ff..fd997070dd 100644
--- a/packages/nuxt/src/core/plugins/async-context.ts
+++ b/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)
-
-      let importName: string
-
-      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, ') }')
-                  }
-                }
-              }
-            })
-          }
-        }
-      })
-
+      s.prepend('import { withAsyncContext as _withAsyncContext } from "#app/composables/asyncContext";\n')
+      s.replace(/withAsyncContext as _withAsyncContext,?/, '')
       if (s.hasChanged()) {
         return {
           code: s.toString(),