feat(nuxt): experimental native async context support (#20918)

Co-authored-by: Daniel Roe <daniel@roe.dev>

packages/nuxt/src/app/nuxt.ts

@@ -17,7 +17,9 @@ import type { RouteMiddleware } from '../../app'
 import type { NuxtError } from '../app/composables/error'
 import type { AsyncDataRequestStatus } from '../app/composables/asyncData'
 
-const nuxtAppCtx = /* #__PURE__ */ getContext<NuxtApp>('nuxt-app')
+const nuxtAppCtx = /* #__PURE__ */ getContext<NuxtApp>('nuxt-app', {
+  asyncContext: !!process.env.NUXT_ASYNC_CONTEXT && process.server
+})
 
 type HookResult = Promise<void> | void
 

packages/nuxt/src/core/nitro.ts

@@ -45,7 +45,8 @@ export async function initNitro (nuxt: Nuxt & { _nitro?: Nitro }) {
     buildDir: nuxt.options.buildDir,
     experimental: {
       // @ts-expect-error `typescriptBundlerResolution` coming in next nitro version
-      typescriptBundlerResolution: nuxt.options.experimental.typescriptBundlerResolution || nuxt.options.typescript?.tsConfig?.compilerOptions?.moduleResolution?.toLowerCase() === 'bundler' || _nitroConfig.typescript?.tsConfig?.compilerOptions?.moduleResolution?.toLowerCase() === 'bundler'
+      typescriptBundlerResolution: nuxt.options.experimental.typescriptBundlerResolution || nuxt.options.typescript?.tsConfig?.compilerOptions?.moduleResolution?.toLowerCase() === 'bundler' || _nitroConfig.typescript?.tsConfig?.compilerOptions?.moduleResolution?.toLowerCase() === 'bundler',
+      asyncContext: nuxt.options.experimental.asyncContext
     },
     imports: {
       autoImport: nuxt.options.imports.autoImport as boolean,
@@ -191,6 +192,7 @@ export async function initNitro (nuxt: Nuxt & { _nitro?: Nitro }) {
       'process.env.NUXT_INLINE_STYLES': !!nuxt.options.experimental.inlineSSRStyles,
       'process.env.NUXT_JSON_PAYLOADS': !!nuxt.options.experimental.renderJsonPayloads,
       'process.env.NUXT_COMPONENT_ISLANDS': !!nuxt.options.experimental.componentIslands,
+      'process.env.NUXT_ASYNC_CONTEXT': !!nuxt.options.experimental.asyncContext,
       'process.dev': nuxt.options.dev,
       __VUE_PROD_DEVTOOLS__: false
     },

packages/nuxt/src/core/nuxt.ts

@@ -25,6 +25,7 @@ import { addModuleTranspiles } from './modules'
 import { initNitro } from './nitro'
 import schemaModule from './schema'
 import { RemovePluginMetadataPlugin } from './plugins/plugin-metadata'
+import { AsyncContextInjectionPlugin } from './plugins/async-context'
 import { resolveDeepImportsPlugin } from './plugins/resolve-deep-imports'
 
 export function createNuxt (options: NuxtOptions): Nuxt {
@@ -143,6 +144,11 @@ async function initNuxt (nuxt: Nuxt) {
     addWebpackPlugin(() => DevOnlyPlugin.webpack({ sourcemap: nuxt.options.sourcemap.server || nuxt.options.sourcemap.client }))
   }
 
+  // Transform initial composable call within `<script setup>` to preserve context
+  if (nuxt.options.experimental.asyncContext) {
+    addBuildPlugin(AsyncContextInjectionPlugin(nuxt))
+  }
+
   // TODO: [Experimental] Avoid emitting assets when flag is enabled
   if (nuxt.options.experimental.noScripts && !nuxt.options.dev) {
     nuxt.hook('build:manifest', async (manifest) => {

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

@@ -0,0 +1,54 @@
+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',
+    transformInclude (id) {
+      return isVue(id, { type: ['template', 'script'] })
+    },
+    transform (code) {
+      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, ') }')
+                  }
+                }
+              }
+            })
+          }
+        }
+      })
+
+      if (s.hasChanged()) {
+        return {
+          code: s.toString(),
+          map: nuxt.options.sourcemap.client || nuxt.options.sourcemap.server
+            ? s.generateMap({ hires: true })
+            : undefined
+        }
+      }
+    }
+  }
+})

packages/nuxt/src/core/runtime/nitro/renderer.ts

@@ -1,3 +1,4 @@
+import { AsyncLocalStorage } from 'node:async_hooks'
 import {
   createRenderer,
   getPrefetchLinks,
@@ -34,6 +35,11 @@ globalThis.__buildAssetsURL = buildAssetsURL
 // @ts-expect-error private property consumed by vite-generated url helpers
 globalThis.__publicAssetsURL = publicAssetsURL
 
+// Polyfill for unctx (https://github.com/unjs/unctx#native-async-context)
+if (process.env.NUXT_ASYNC_CONTEXT && !('AsyncLocalStorage' in globalThis)) {
+  (globalThis as any).AsyncLocalStorage = AsyncLocalStorage
+}
+
 export interface NuxtRenderHTMLContext {
   island?: boolean
   htmlAttrs: string[]

packages/schema/src/config/experimental.ts

@@ -220,6 +220,13 @@ export default defineUntypedSchema({
      */
     watcher: 'chokidar-granular',
 
+    /**
+     * Enable native async context to be accessable for nested composables
+     *
+     * @see https://github.com/nuxt/nuxt/pull/20918
+     */
+    asyncContext: false,
+
     /**
      * Add the capo.js head plugin in order to render tags in of the head in a more performant way.
      *

packages/vite/src/vite.ts

@@ -82,7 +82,10 @@ export async function bundle (nuxt: Nuxt) {
           exclude: ['nuxt/app']
         },
         css: resolveCSSOptions(nuxt),
-        define: { __NUXT_VERSION__: JSON.stringify(nuxt._version) },
+        define: {
+          __NUXT_VERSION__: JSON.stringify(nuxt._version),
+          'process.env.NUXT_ASYNC_CONTEXT': nuxt.options.experimental.asyncContext
+        },
         build: {
           copyPublicDir: false,
           rollupOptions: {

packages/webpack/src/presets/base.ts

@@ -218,6 +218,7 @@ function getEnv (ctx: WebpackConfigContext) {
     'process.env.NODE_ENV': JSON.stringify(ctx.config.mode),
     __NUXT_VERSION__: JSON.stringify(ctx.nuxt._version),
     'process.env.VUE_ENV': JSON.stringify(ctx.name),
+    'process.env.NUXT_ASYNC_CONTEXT': ctx.options.experimental.asyncContext,
     'process.dev': ctx.options.dev,
     'process.test': isTest,
     'process.browser': ctx.isClient,

test/basic.test.ts

@@ -1901,6 +1901,12 @@ describe.skipIf(isDev() || isWindows || !isRenderingJson)('payload rendering', (
   })
 })
 
+describe('Async context', () => {
+  it('should be available', async () => {
+    expect(await $fetch('/async-context')).toContain('"hasApp": true')
+  })
+})
+
 describe.skipIf(isWindows)('useAsyncData', () => {
   it('single request resolves', async () => {
     await expectNoClientErrors('/useAsyncData/single')

test/fixtures/basic/composables/async-context.ts

@@ -0,0 +1,19 @@
+const delay = () => new Promise(resolve => setTimeout(resolve, 10))
+
+export async function nestedAsyncComposable () {
+  await delay()
+  return await fn1()
+}
+
+async function fn1 () {
+  await delay()
+  return await fn2()
+}
+
+async function fn2 () {
+  await delay()
+  const app = useNuxtApp()
+  return {
+    hasApp: !!app
+  }
+}

test/fixtures/basic/nuxt.config.ts

@@ -188,6 +188,7 @@ export default defineNuxtConfig({
     reactivityTransform: true,
     treeshakeClientOnly: true,
     payloadExtraction: true,
+    asyncContext: true,
     headCapoPlugin: true
   },
   appConfig: {

test/fixtures/basic/pages/async-context.vue

@@ -0,0 +1,12 @@
+<template>
+  <div>
+    <Head>
+      <Title>Native Async Context</Title>
+    </Head>
+    {{ data }}
+  </div>
+</template>
+
+<script setup lang="ts">
+const data = await nestedAsyncComposable()
+</script>

vitest.config.ts

@@ -11,6 +11,9 @@ export default defineConfig({
       '@nuxt/test-utils': resolve('./packages/test-utils/src/index.ts')
     }
   },
+  define: {
+    'process.env.NUXT_ASYNC_CONTEXT': 'false'
+  },
   test: {
     globalSetup: './test/setup.ts',
     testTimeout: isWindows ? 60000 : 10000,