From d6c3c2439a4c1e233e0f81e07d1b592c7e0520a6 Mon Sep 17 00:00:00 2001
From: Daniel Roe <daniel@roe.dev>
Date: Mon, 10 Apr 2023 22:57:13 +0100
Subject: [PATCH] feat(nuxt): support async transform of object properties
 (#20182)

---
 docs/1.getting-started/6.data-fetching.md     |  2 +-
 docs/2.guide/1.concepts/1.auto-imports.md     |  2 +-
 .../nuxt/src/app/composables/component.ts     |  7 +--
 packages/nuxt/src/core/nuxt.ts                |  8 +++-
 packages/nuxt/src/core/plugins/unctx.ts       | 44 ++++++++++++-------
 packages/nuxt/src/pages/module.ts             |  6 ++-
 packages/schema/build.config.ts               |  1 +
 packages/schema/package.json                  |  1 +
 packages/schema/src/config/build.ts           | 14 ++++++
 pnpm-lock.yaml                                |  5 ++-
 test/fixtures/basic/pages/chunk-error.vue     |  5 ++-
 .../basic/pages/legacy/async-data.vue         |  5 +++
 .../basic/pages/navigate-to-redirect.vue      |  4 +-
 13 files changed, 73 insertions(+), 31 deletions(-)

diff --git a/docs/1.getting-started/6.data-fetching.md b/docs/1.getting-started/6.data-fetching.md
index 7b091c6679..089f61b40a 100644
--- a/docs/1.getting-started/6.data-fetching.md
+++ b/docs/1.getting-started/6.data-fetching.md
@@ -356,7 +356,7 @@ const { data } = await useFetch('/api/price')
 
 ## Using Async Setup
 
-If you are using `async setup()`, the current component instance will be lost after the first `await`. (This is a Vue 3 limitation.) If you want to use multiple async operations, such as multiple calls to `useFetch`, you will need to use `<script setup>` or await them together at the end of setup.
+If you are using `async setup()`, the current component instance will be lost after the first `await`. (This is a Vue 3 limitation.) If you want to use multiple async operations, such as multiple calls to `useFetch`, you will need to use `<script setup>`, await them together at the end of setup, or alternatively use `defineNuxtComponent` (which applies a custom transform to the setup function).
 
 ::alert{icon=👉}
 Using `<script setup>` is recommended, as it removes the limitation of using top-level await. [Read more](https://vuejs.org/api/sfc-script-setup.html#top-level-await)
diff --git a/docs/2.guide/1.concepts/1.auto-imports.md b/docs/2.guide/1.concepts/1.auto-imports.md
index a7888a3b0b..1e2d1f1d19 100644
--- a/docs/2.guide/1.concepts/1.auto-imports.md
+++ b/docs/2.guide/1.concepts/1.auto-imports.md
@@ -54,7 +54,7 @@ When you are using the built-in Composition API composables provided by Vue and
 
 During a component lifecycle, Vue tracks the temporary instance of the current component (and similarly, Nuxt tracks a temporary instance of `nuxtApp`) via a global variable, and then unsets it in same tick. This is essential when server rendering, both to avoid cross-request state pollution (leaking a shared reference between two users) and to avoid leakage between different components.
 
-That means that (with very few exceptions) you cannot use them outside a Nuxt plugin, Nuxt route middleware or Vue setup function. On top of that, you must use them synchronously - that is, you cannot use `await` before calling a composable, except within `<script setup>` blocks, in `defineNuxtPlugin` or in `defineNuxtRouteMiddleware`, where we perform a transform to keep the synchronous context even after the `await`.
+That means that (with very few exceptions) you cannot use them outside a Nuxt plugin, Nuxt route middleware or Vue setup function. On top of that, you must use them synchronously - that is, you cannot use `await` before calling a composable, except within `<script setup>` blocks, within the setup function of a component declared with `defineNuxtComponent`, in `defineNuxtPlugin` or in `defineNuxtRouteMiddleware`, where we perform a transform to keep the synchronous context even after the `await`.
 
 If you get an error message like `Nuxt instance is unavailable` then it probably means you are calling a Nuxt composable in the wrong place in the Vue or Nuxt lifecycle.
 
diff --git a/packages/nuxt/src/app/composables/component.ts b/packages/nuxt/src/app/composables/component.ts
index 1c401e16cf..9213550011 100644
--- a/packages/nuxt/src/app/composables/component.ts
+++ b/packages/nuxt/src/app/composables/component.ts
@@ -2,7 +2,7 @@ import { getCurrentInstance, reactive, toRefs } from 'vue'
 import type { DefineComponent, defineComponent } from 'vue'
 import { useHead } from '@unhead/vue'
 import type { NuxtApp } from '../nuxt'
-import { useNuxtApp } from '../nuxt'
+import { callWithNuxt, useNuxtApp } from '../nuxt'
 import { useAsyncData } from './asyncData'
 import { useRoute } from './router'
 
@@ -14,7 +14,7 @@ async function runLegacyAsyncData (res: Record<string, any> | Promise<Record<str
   const vm = getCurrentInstance()!
   const { fetchKey } = vm.proxy!.$options
   const key = typeof fetchKey === 'function' ? fetchKey(() => '') : fetchKey || route.fullPath
-  const { data } = await useAsyncData(`options:asyncdata:${key}`, () => fn(nuxt))
+  const { data } = await useAsyncData(`options:asyncdata:${key}`, () => callWithNuxt(nuxt, fn, [nuxt]))
   if (data.value && typeof data.value === 'object') {
     Object.assign(await res, toRefs(reactive(data.value)))
   } else if (process.dev) {
@@ -38,7 +38,8 @@ export const defineNuxtComponent: typeof defineComponent =
       [NuxtComponentIndicator]: true,
       ...options,
       setup (props, ctx) {
-        const res = setup?.(props, ctx) || {}
+        const nuxtApp = useNuxtApp()
+        const res = setup ? Promise.resolve(callWithNuxt(nuxtApp, setup, [props, ctx])).then(r => r || {}) : {}
 
         const promises: Promise<any>[] = []
         if (options.asyncData) {
diff --git a/packages/nuxt/src/core/nuxt.ts b/packages/nuxt/src/core/nuxt.ts
index 45405aa0f9..6659fab06f 100644
--- a/packages/nuxt/src/core/nuxt.ts
+++ b/packages/nuxt/src/core/nuxt.ts
@@ -98,8 +98,12 @@ async function initNuxt (nuxt: Nuxt) {
 
   nuxt.hook('modules:done', () => {
     // Add unctx transform
-    addVitePlugin(UnctxTransformPlugin(nuxt).vite({ sourcemap: nuxt.options.sourcemap.server || nuxt.options.sourcemap.client }))
-    addWebpackPlugin(UnctxTransformPlugin(nuxt).webpack({ sourcemap: nuxt.options.sourcemap.server || nuxt.options.sourcemap.client }))
+    const options = {
+      sourcemap: nuxt.options.sourcemap.server || nuxt.options.sourcemap.client,
+      transformerOptions: nuxt.options.optimization.asyncTransforms
+    }
+    addVitePlugin(UnctxTransformPlugin.vite(options))
+    addWebpackPlugin(UnctxTransformPlugin.webpack(options))
 
     // Add composable tree-shaking optimisations
     const serverTreeShakeOptions: TreeShakeComposablesPluginOptions = {
diff --git a/packages/nuxt/src/core/plugins/unctx.ts b/packages/nuxt/src/core/plugins/unctx.ts
index 5cf7b748a5..042950010c 100644
--- a/packages/nuxt/src/core/plugins/unctx.ts
+++ b/packages/nuxt/src/core/plugins/unctx.ts
@@ -1,30 +1,42 @@
-import { normalize } from 'pathe'
+import { pathToFileURL } from 'node:url'
+import { parseQuery, parseURL } from 'ufo'
+import type { TransformerOptions } from 'unctx/transform'
 import { createTransformer } from 'unctx/transform'
 import { createUnplugin } from 'unplugin'
-import type { Nuxt, NuxtApp } from 'nuxt/schema'
 
 const TRANSFORM_MARKER = '/* _processed_nuxt_unctx_transform */\n'
 
-export const UnctxTransformPlugin = (nuxt: Nuxt) => {
-  const transformer = createTransformer({
-    asyncFunctions: ['defineNuxtPlugin', 'defineNuxtRouteMiddleware']
-  })
+interface UnctxTransformPluginOptions {
+  sourcemap?: boolean
+  transformerOptions: TransformerOptions
+}
 
-  let app: NuxtApp | undefined
-  nuxt.hook('app:resolve', (_app) => { app = _app })
-
-  return createUnplugin((options: { sourcemap?: boolean } = {}) => ({
+export const UnctxTransformPlugin = createUnplugin((options: UnctxTransformPluginOptions) => {
+  const transformer = createTransformer(options.transformerOptions)
+  return {
     name: 'unctx:transform',
     enforce: 'post',
     transformInclude (id) {
-      if (id.includes('macro=true')) { return true }
+      const { pathname, search } = parseURL(decodeURIComponent(pathToFileURL(id).href))
+      const query = parseQuery(search)
 
-      id = normalize(id).replace(/\?.*$/, '')
-      return app?.plugins.some(i => i.src === id) || app?.middleware.some(m => m.path === id)
+      // Vue files
+      if (
+        pathname.endsWith('.vue') ||
+        'macro' in query ||
+        ('vue' in query && (query.type === 'template' || query.type === 'script' || 'setup' in query))
+      ) {
+        return true
+      }
+
+      // JavaScript files
+      if (pathname.match(/\.((c|m)?j|t)sx?$/g)) {
+        return true
+      }
     },
     transform (code, id) {
       // TODO: needed for webpack - update transform in unctx/unplugin?
-      if (code.startsWith(TRANSFORM_MARKER)) { return }
+      if (code.startsWith(TRANSFORM_MARKER) || !transformer.shouldTransform(code)) { return }
       const result = transformer.transform(code)
       if (result) {
         return {
@@ -35,5 +47,5 @@ export const UnctxTransformPlugin = (nuxt: Nuxt) => {
         }
       }
     }
-  }))
-}
+  }
+})
diff --git a/packages/nuxt/src/pages/module.ts b/packages/nuxt/src/pages/module.ts
index 485ca84827..8997954d9a 100644
--- a/packages/nuxt/src/pages/module.ts
+++ b/packages/nuxt/src/pages/module.ts
@@ -156,8 +156,10 @@ export default defineNuxtModule({
         layer => resolve(layer.config.srcDir, layer.config.dir?.pages || 'pages')
       )
     }
-    addVitePlugin(PageMetaPlugin.vite(pageMetaOptions))
-    addWebpackPlugin(PageMetaPlugin.webpack(pageMetaOptions))
+    nuxt.hook('modules:done', () => {
+      addVitePlugin(PageMetaPlugin.vite(pageMetaOptions))
+      addWebpackPlugin(PageMetaPlugin.webpack(pageMetaOptions))
+    })
 
     // Add prefetching support for middleware & layouts
     addPlugin(resolve(runtimeDir, 'plugins/prefetch.client'))
diff --git a/packages/schema/build.config.ts b/packages/schema/build.config.ts
index beb7e14bf3..01f52c4bf6 100644
--- a/packages/schema/build.config.ts
+++ b/packages/schema/build.config.ts
@@ -23,6 +23,7 @@ export default defineBuildConfig({
     'vue-bundle-renderer',
     '@unhead/schema',
     'vue',
+    'unctx',
     'hookable',
     'nitropack',
     'webpack',
diff --git a/packages/schema/package.json b/packages/schema/package.json
index 3ebebe54eb..f9bef6b958 100644
--- a/packages/schema/package.json
+++ b/packages/schema/package.json
@@ -28,6 +28,7 @@
     "@vitejs/plugin-vue-jsx": "^3.0.1",
     "nitropack": "^2.3.2",
     "unbuild": "latest",
+    "unctx": "^2.2.0",
     "vite": "~4.2.1"
   },
   "dependencies": {
diff --git a/packages/schema/src/config/build.ts b/packages/schema/src/config/build.ts
index e4f74a60da..076146b50a 100644
--- a/packages/schema/src/config/build.ts
+++ b/packages/schema/src/config/build.ts
@@ -184,5 +184,19 @@ export default defineUntypedSchema({
         }
       }
     },
+
+    /**
+     * Options passed directly to the transformer from `unctx` that preserves async context
+     * after `await`.
+     *
+     * @type {import('unctx').TransformerOptions}
+     */
+    asyncTransforms: {
+      asyncFunctions: ['defineNuxtPlugin', 'defineNuxtRouteMiddleware'],
+      objectDefinitions: {
+        defineNuxtComponent: ['asyncData', 'setup'],
+        definePageMeta: ['middleware', 'validate']
+      }
+    }
   }
 })
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index bd3297a745..0fe492507b 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -739,6 +739,9 @@ importers:
       unbuild:
         specifier: ^1.2.0
         version: 1.2.0
+      unctx:
+        specifier: ^2.2.0
+        version: 2.2.0
       vite:
         specifier: ^4.2.1
         version: 4.2.1(@types/node@18.15.11)
@@ -5099,7 +5102,6 @@ packages:
     resolution: {integrity: sha512-7RUKfXgSMMkzt6ZuXmqapOurLGPPfgj6l9uRZ7lRGolvk0y2yocc35LdcxKC5PQZdn2DMqioAQ2NoWcrTKmm6g==}
     dependencies:
       '@types/estree': 1.0.0
-    dev: false
 
   /esutils@2.0.3:
     resolution: {integrity: sha512-kVscqXk4OCp68SZ0dkgEKVi6/8ij300KBWTJq32P/dYeWTSwK41WyTxalN1eRmA5Z9UU/LX9D7FWSmV9SAYx6g==}
@@ -8712,7 +8714,6 @@ packages:
       estree-walker: 3.0.3
       magic-string: 0.30.0
       unplugin: 1.3.1
-    dev: false
 
   /undici@5.20.0:
     resolution: {integrity: sha512-J3j60dYzuo6Eevbawwp1sdg16k5Tf768bxYK4TUJRH7cBM4kFCbf3mOnM/0E3vQYXvpxITbbWmBafaDbxLDz3g==}
diff --git a/test/fixtures/basic/pages/chunk-error.vue b/test/fixtures/basic/pages/chunk-error.vue
index f0785e59df..5e7c5bf016 100644
--- a/test/fixtures/basic/pages/chunk-error.vue
+++ b/test/fixtures/basic/pages/chunk-error.vue
@@ -1,12 +1,13 @@
 <script setup lang="ts">
 definePageMeta({
-  middleware: defineNuxtRouteMiddleware(async (to, from) => {
+  async middleware (to, from) {
+    await new Promise(resolve => setTimeout(resolve, 1))
     const nuxtApp = useNuxtApp()
     if (process.client && from !== to && !nuxtApp.isHydrating) {
       // trigger a loading error when navigated to via client-side navigation
       await import(/* webpackIgnore: true */ /* @vite-ignore */ `some-non-exis${''}ting-module`)
     }
-  })
+  }
 })
 const someValue = useState('val', () => 1)
 </script>
diff --git a/test/fixtures/basic/pages/legacy/async-data.vue b/test/fixtures/basic/pages/legacy/async-data.vue
index 3ca56e0798..8b36739185 100644
--- a/test/fixtures/basic/pages/legacy/async-data.vue
+++ b/test/fixtures/basic/pages/legacy/async-data.vue
@@ -6,7 +6,12 @@
 
 <script>
 export default defineNuxtComponent({
+  async setup () {
+    await nextTick()
+    useRuntimeConfig()
+  },
   async asyncData () {
+    await nextTick()
     return {
       hello: await $fetch('/api/hello')
     }
diff --git a/test/fixtures/basic/pages/navigate-to-redirect.vue b/test/fixtures/basic/pages/navigate-to-redirect.vue
index b705612434..9e934988b7 100644
--- a/test/fixtures/basic/pages/navigate-to-redirect.vue
+++ b/test/fixtures/basic/pages/navigate-to-redirect.vue
@@ -4,9 +4,9 @@
 
 <script setup>
 definePageMeta({
-  middleware: defineNuxtRouteMiddleware(async () => {
+  middleware: async () => {
     await new Promise(resolve => setTimeout(resolve, 1))
     return navigateTo({ path: '/' }, { redirectCode: 307 })
-  })
+  }
 })
 </script>