fix(devtools): rework timeline function wrapping to run after key injection

packages/devtools/package.json

@@ -42,7 +42,8 @@
     "dev:playground": "pnpm build && nuxi dev playground",
     "dev:prepare": "pnpm run stub && nuxi prepare client",
     "prepare": "tsx scripts/prepare.ts",
-    "prepack": "pnpm build"
+    "prepack": "pnpm build",
+    "test": "vitest run"
   },
   "peerDependencies": {
     "vite": ">=6.0"
@@ -58,13 +59,16 @@
     "consola": "catalog:prod",
     "destr": "catalog:prod",
     "error-stack-parser-es": "catalog:frontend",
+    "escape-string-regexp": "catalog:prod",
     "fast-npm-meta": "catalog:prod",
     "get-port-please": "catalog:prod",
     "hookable": "catalog:frontend",
     "image-meta": "catalog:prod",
     "launch-editor": "catalog:prod",
     "local-pkg": "catalog:prod",
+    "magic-string": "catalog:prod",
     "magicast": "catalog:prod",
+    "mlly": "catalog:prod",
     "ohash": "catalog:frontend",
     "pathe": "catalog:frontend",
     "perfect-debounce": "catalog:frontend",
@@ -74,6 +78,7 @@
     "structured-clone-es": "catalog:frontend",
     "tinyexec": "catalog:prod",
     "tinyglobby": "catalog:prod",
+    "unplugin": "catalog:prod",
     "unstorage": "catalog:prod",
     "vite-plugin-inspect": "catalog:prod",
     "vite-plugin-vue-tracer": "catalog:prod",
@@ -91,6 +96,7 @@
     "@iconify-json/ri": "catalog:icons",
     "@iconify-json/simple-icons": "catalog:icons",
     "@iconify-json/tabler": "catalog:icons",
+    "@nuxt/schema": "catalog:types",
     "@nuxt/test-utils": "catalog:cli",
     "@parcel/watcher": "catalog:buildtools",
     "@types/markdown-it-link-attributes": "catalog:types",
@@ -131,6 +137,7 @@
     "vanilla-jsoneditor": "catalog:frontend",
     "vis-data": "catalog:frontend",
     "vis-network": "catalog:frontend",
+    "vitest": "catalog:cli",
     "vue-tsc": "catalog:cli",
     "vue-virtual-scroller": "catalog:frontend"
   }

packages/devtools/src/integrations/timeline-wrap.ts

@@ -0,0 +1,82 @@
+import type { SourceMap } from 'magic-string'
+import MagicString from 'magic-string'
+import { findStaticImports, parseStaticImport } from 'mlly'
+
+export const HELPER_NAME = '__nuxtTimelineWrap'
+export const HELPER_NAME_RE = /__nuxtTimelineWrap/
+const RENAME_PREFIX = '_$__'
+
+export const TRANSFORM_ID_INCLUDE: RegExp[] = [
+  /^[^?]*\.(?:m?[jt]sx?|vue)(?:$|\?)/,
+]
+export const TRANSFORM_ID_EXCLUDE: RegExp[] = [
+  /^\0/,
+  /[\\/]node_modules[\\/]/,
+  /[?&]macro=true/,
+  /[?&]type=(?:style|template|custom)\b/,
+]
+
+export function wrapTimelineImports(
+  code: string,
+  isWrappable: (source: string, name: string) => boolean,
+  helperPath: string,
+): { code: string, map: SourceMap } | undefined {
+  if (!code.includes('import') || code.includes(HELPER_NAME))
+    return
+
+  const staticImports = findStaticImports(code)
+  if (!staticImports.length)
+    return
+
+  const s = new MagicString(code)
+  const wrappers: string[] = []
+
+  for (const imp of staticImports) {
+    const { specifier, defaultImport, namespacedImport, namedImports } = parseStaticImport(imp)
+    let wrapped = false
+
+    // returns the local binding for an imported name, renamed when it gets wrapped
+    const bind = (name: string, local: string): string => {
+      if (!isWrappable(specifier, name))
+        return local
+      wrapped = true
+      // a default import has no meaningful exported name, record it under its local name
+      wrappers.push(`const ${local} = ${HELPER_NAME}(${JSON.stringify(name === 'default' ? local : name)}, ${RENAME_PREFIX}${local});`)
+      return RENAME_PREFIX + local
+    }
+
+    const clause: string[] = []
+    if (defaultImport)
+      clause.push(bind('default', defaultImport))
+    if (namespacedImport)
+      clause.push(`* as ${namespacedImport}`)
+    const named = Object.entries(namedImports ?? {}).map(([name, local]) => {
+      const bound = bind(name, local)
+      return name === bound ? name : `${name} as ${bound}`
+    })
+
+    if (named.length)
+      clause.push(`{ ${named.join(', ')} }`)
+
+    if (!wrapped)
+      continue
+
+    // rebuild only the clause between `import` and `from`, keeping the specifier untouched
+    const clauseStart = imp.start + imp.code.indexOf(imp.imports)
+    s.overwrite(clauseStart, clauseStart + imp.imports.length, `${clause.join(', ')} `)
+  }
+
+  if (!wrappers.length)
+    return
+
+  s.prepend([
+    `import { ${HELPER_NAME} } from ${JSON.stringify(helperPath)};`,
+    ...wrappers,
+    '',
+  ].join('\n'))
+
+  return {
+    code: s.toString(),
+    map: s.generateMap({ hires: 'boundary' }),
+  }
+}

packages/devtools/src/integrations/timeline.ts

@@ -1,10 +1,16 @@
+import type { NuxtHooks, NuxtOptions } from '@nuxt/schema'
 import type { Import } from 'unimport'
 import type { NuxtDevtoolsServerContext } from '../types'
+import { addBuildPlugin } from '@nuxt/kit'
+import escapeRE from 'escape-string-regexp'
 import { resolve } from 'pathe'
-import semver from 'semver'
+import { createUnplugin } from 'unplugin'
 import { runtimeDir } from '../dirs'
+import { HELPER_NAME_RE, TRANSFORM_ID_EXCLUDE, TRANSFORM_ID_INCLUDE, wrapTimelineImports } from './timeline-wrap'
 
 const DEFINE_UPPER_RE = /^define[A-Z]/
+const RUNTIME_DIR_RE = new RegExp(`^${escapeRE(runtimeDir)}`)
+const IMPORT_RE = /\bimport\b/
 
 export function setup({ nuxt, options }: NuxtDevtoolsServerContext) {
   const helperPath = resolve(runtimeDir, 'function-metrics-helpers')
@@ -34,50 +40,100 @@ export function setup({ nuxt, options }: NuxtDevtoolsServerContext) {
     return true
   }
 
-  nuxt.hook('imports:context', (unimport) => {
-    const ctx = unimport.getInternalContext()
+  let unimport: Parameters<NuxtHooks['imports:context']>[0] | undefined
+  nuxt.hook('imports:context', (ctx) => {
+    unimport = ctx
+  })
+
+  // wrappable auto-imports keyed by `${source}\0${name}`, cached until the import list changes
+  const wrappableCache = new WeakMap<Import[], Promise<Map<string, Import>>>()
+  const resolvedFrom = new Map<string, Promise<string | undefined>>()
 
-    if (!ctx.version || !semver.gte(ctx.version, '3.1.0'))
-      throw new Error(`[Nuxt DevTools] The timeline feature requires \`unimport\` >= v3.1.0, but got \`${ctx.version || '(unknown)'}\`. Please upgrade using \`nuxi upgrade --force\`.`)
+  async function getWrappable(): Promise<Map<string, Import>> {
+    if (!unimport)
+      return new Map()
+    const imports = await unimport.getImports()
+    let cached = wrappableCache.get(imports)
+    if (!cached) {
+      cached = buildWrappable(imports)
+      wrappableCache.set(imports, cached)
+    }
+    return cached
+  }
 
-    ctx.addons.push(
-      {
-        injectImportsResolved(imports, _code, id) {
-          if (id?.includes('?macro=true'))
-            return
-          return imports.map((i) => {
-            if (!filter(i))
-              return i
+  async function buildWrappable(imports: Import[]): Promise<Map<string, Import>> {
+    const map = new Map<string, Import>()
+    // never wrap keyed-function factory macros (`createUseFetch`, etc.)
+    // the option exists since Nuxt 4.4 and is not yet part of the schema types we build against (+ we need it as optional for backwards compat)
+    const { keyedComposableFactories } = nuxt.options.optimization as NuxtOptions['optimization'] & {
+      keyedComposableFactories?: { name: string }[]
+    }
+    const factoryNames = new Set(keyedComposableFactories?.map(f => f.name) ?? [])
+    const ctx = unimport?.getInternalContext()
 
-            const name = i.as || i.name
+    // group by source so that each unique `from` is only resolved once
+    const bySource = new Map<string, Import[]>()
+    for (const i of imports) {
+      if (!filter(i) || factoryNames.has(i.name))
+        continue
+      map.set(`${i.from}\0${i.name}`, i)
+      const group = bySource.get(i.from)
+      if (group)
+        group.push(i)
+      else
+        bySource.set(i.from, [i])
+    }
 
-            return {
-              ...i,
-              meta: {
-                wrapperOriginalAs: name,
-              },
-              as: `_$__${name}`,
-            }
-          })
-        },
-        injectImportsStringified(str, imports, s, id) {
-          if (id?.includes('?macro=true'))
-            return
-          const code = s.toString()
-          const injected = imports.filter(i => i.meta?.wrapperOriginalAs)
-          if (injected.length) {
-            const result = [
-              str,
-              code.includes('__nuxtTimelineWrap')
-                ? ''
-                : `import { __nuxtTimelineWrap } from ${JSON.stringify(helperPath)}`,
-              ...injected.map(i => `const ${i.meta!.wrapperOriginalAs} = __nuxtTimelineWrap(${JSON.stringify(i.name)}, ${i.as})`),
-              '',
-            ].join(';')
-            return result
-          }
+    if (!ctx?.resolveId)
+      return map
+
+    // unimport injects imports with the resolved specifier, which may differ from `i.from`
+    await Promise.all(Array.from(bySource, async ([from, items]) => {
+      let resolving = resolvedFrom.get(from)
+      if (!resolving) {
+        resolving = Promise.resolve(ctx.resolveId(from)).then(r => r || undefined, () => undefined)
+        resolvedFrom.set(from, resolving)
+      }
+      const resolved = await resolving
+      if (resolved && resolved !== from) {
+        for (const i of items)
+          map.set(`${resolved}\0${i.name}`, i)
+      }
+    }))
+    return map
+  }
+
+  // the wrapping must happen AFTER Nuxt's key injection and keyed function factory macro rewriting
+
+  // Nuxt registers these plugins in the `build:before` hook during core module setup, so we register ours
+  // from `modules:done` to be queued after them
+  nuxt.hook('modules:done', () => {
+    nuxt.hook('build:before', () => {
+      addBuildPlugin(createUnplugin(() => ({
+        name: 'nuxt:devtools:timeline-function-wrap',
+        enforce: 'post',
+        transform: {
+          filter: {
+            id: {
+              include: TRANSFORM_ID_INCLUDE,
+              exclude: [
+                ...TRANSFORM_ID_EXCLUDE,
+                RUNTIME_DIR_RE,
+              ],
+            },
+            code: {
+              include: IMPORT_RE,
+              exclude: HELPER_NAME_RE,
+            },
+          },
+          async handler(code) {
+            const wrappable = await getWrappable()
+            if (!wrappable.size)
+              return
+            return wrapTimelineImports(code, (source, name) => wrappable.has(`${source}\0${name}`), helperPath)
+          },
         },
-      },
-    )
+      })))
+    })
   })
 }