feat(nuxt): support custom keyed composables (#19490)

packages/schema/src/config/build.ts

@@ -132,6 +132,27 @@ export default defineUntypedSchema({
    * Build time optimization configuration.
    */
   optimization: {
+    /**
+     * Functions to inject a key for.
+     *
+     * As long as the number of arguments passed to the function is less than `argumentLength`, an
+     * additional magic string will be injected that can be used to deduplicate requests between server
+     * and client. You will need to take steps to handle this additional key.
+     *
+     * The key will be unique based on the location of the function being invoked within the file.
+     *
+     * @type {Array<{ name: string, argumentLength: number }>}
+     */
+    keyedComposables: {
+      $resolve: (val) => [
+        { name: 'useState', argumentLength: 2 },
+        { name: 'useFetch', argumentLength: 3 },
+        { name: 'useAsyncData', argumentLength: 3 },
+        { name: 'useLazyAsyncData', argumentLength: 3 },
+        { name: 'useLazyFetch', argumentLength: 3 },
+      ].concat(val).filter(Boolean)
+    },
+
     /**
      * Tree shake code from specific builds.
      */

packages/vite/src/plugins/composable-keys.ts

@@ -7,20 +7,24 @@ import MagicString from 'magic-string'
 import { hash } from 'ohash'
 import type { CallExpression } from 'estree'
 import { parseQuery, parseURL } from 'ufo'
+import escapeRE from 'escape-string-regexp'
 import { findStaticImports, parseStaticImport } from 'mlly'
 
 export interface ComposableKeysOptions {
   sourcemap: boolean
   rootDir: string
+  composables: Array<{ name: string, argumentLength: number }>
 }
 
-const keyedFunctions = [
-  'useState', 'useFetch', 'useAsyncData', 'useLazyAsyncData', 'useLazyFetch'
-]
-const KEYED_FUNCTIONS_RE = new RegExp(`(${keyedFunctions.join('|')})`)
 const stringTypes = ['Literal', 'TemplateLiteral']
 
 export const composableKeysPlugin = createUnplugin((options: ComposableKeysOptions) => {
+  const composableMeta = Object.fromEntries(options.composables.map(({ name, ...meta }) => [name, meta]))
+
+  const maxLength = Math.max(...options.composables.map(({ argumentLength }) => argumentLength))
+  const keyedFunctions = new Set(options.composables.map(({ name }) => name))
+  const KEYED_FUNCTIONS_RE = new RegExp(`\\b(${[...keyedFunctions].map(f => escapeRE(f)).join('|')})\\b`)
+
   return {
     name: 'nuxt:composable-keys',
     enforce: 'post',
@@ -44,24 +48,28 @@ export const composableKeysPlugin = createUnplugin((options: ComposableKeysOptio
           if (_node.type !== 'CallExpression' || (_node as CallExpression).callee.type !== 'Identifier') { return }
           const node: CallExpression = _node as CallExpression
           const name = 'name' in node.callee && node.callee.name
-          if (!name || !keyedFunctions.includes(name) || node.arguments.length >= 4) { return }
+          if (!name || !keyedFunctions.has(name) || node.arguments.length >= maxLength) { return }
 
           imports = imports || detectImportNames(script)
           if (imports.has(name)) { return }
 
+          const meta = composableMeta[name]
+
+          if (node.arguments.length >= meta.argumentLength) { return }
+
           switch (name) {
             case 'useState':
-              if (node.arguments.length >= 2 || stringTypes.includes(node.arguments[0]?.type)) { return }
+              if (stringTypes.includes(node.arguments[0]?.type)) { return }
               break
 
             case 'useFetch':
             case 'useLazyFetch':
-              if (node.arguments.length >= 3 || stringTypes.includes(node.arguments[1]?.type)) { return }
+              if (stringTypes.includes(node.arguments[1]?.type)) { return }
               break
 
             case 'useAsyncData':
             case 'useLazyAsyncData':
-              if (node.arguments.length >= 3 || stringTypes.includes(node.arguments[0]?.type) || stringTypes.includes(node.arguments[node.arguments.length - 1]?.type)) { return }
+              if (stringTypes.includes(node.arguments[0]?.type) || stringTypes.includes(node.arguments[node.arguments.length - 1]?.type)) { return }
               break
           }
 

packages/vite/src/vite.ts

@@ -80,7 +80,11 @@ export async function bundle (nuxt: Nuxt) {
           }
         },
         plugins: [
-          composableKeysPlugin.vite({ sourcemap: nuxt.options.sourcemap.server || nuxt.options.sourcemap.client, rootDir: nuxt.options.rootDir }),
+          composableKeysPlugin.vite({
+            sourcemap: nuxt.options.sourcemap.server || nuxt.options.sourcemap.client,
+            rootDir: nuxt.options.rootDir,
+            composables: nuxt.options.optimization.keyedComposables
+          }),
           replace({
             ...Object.fromEntries([';', '(', '{', '}', ' ', '\t', '\n'].map(d => [`${d}global.`, `${d}globalThis.`])),
             preventAssignment: true

packages/webpack/src/webpack.ts

@@ -46,7 +46,8 @@ export async function bundle (nuxt: Nuxt) {
     }
     config.plugins!.push(composableKeysPlugin.webpack({
       sourcemap: nuxt.options.sourcemap[config.name as 'client' | 'server'],
-      rootDir: nuxt.options.rootDir
+      rootDir: nuxt.options.rootDir,
+      composables: nuxt.options.optimization.keyedComposables
     }))
 
     // Create compiler

test/fixtures/basic/nuxt.config.ts

@@ -49,6 +49,14 @@ export default defineNuxtConfig({
       ]
     }
   },
+  optimization: {
+    keyedComposables: [
+      {
+        name: 'useKeyedComposable',
+        argumentLength: 1
+      }
+    ]
+  },
   runtimeConfig: {
     baseURL: '',
     baseAPIToken: '',

test/fixtures/basic/pages/keyed-composables.vue

@@ -32,6 +32,11 @@ const { data: useFetchTest2 } = await useLocalFetch()
 const useLocalLazyFetch = () => useLazyFetch(() => '/api/counter')
 const { data: useLazyFetchTest1 } = await useLocalLazyFetch()
 const { data: useLazyFetchTest2 } = await useLocalLazyFetch()
+
+const useKeyedComposable = (arg?: string) => arg
+const useLocalKeyedComposable = () => useKeyedComposable()
+const useMyAsyncDataTest1 = useLocalKeyedComposable()
+const useMyAsyncDataTest2 = useLocalKeyedComposable()
 </script>
 
 <template>
@@ -41,6 +46,7 @@ const { data: useLazyFetchTest2 } = await useLocalLazyFetch()
     {{ useLazyAsyncDataTest1 === useLazyAsyncDataTest2 }}
     {{ useFetchTest1 === useFetchTest2 }}
     {{ useLazyFetchTest1 === useLazyFetchTest2 }}
+    {{ !!useMyAsyncDataTest1 && useMyAsyncDataTest1 === useMyAsyncDataTest2 }}
   </div>
 </template>