fix(rsc): show import chain for server-only and client-only import error

Author: hi-ogawa

packages/plugin-rsc/src/plugins/validate-import.ts

@@ -1,4 +1,4 @@
-import type { Plugin } from 'vite'
+import type { DevEnvironment, Plugin, Rollup } from 'vite'
 
 // https://github.com/vercel/next.js/blob/90f564d376153fe0b5808eab7b83665ee5e08aaf/packages/next/src/build/webpack-config.ts#L1249-L1280
 // https://github.com/pcattori/vite-env-only/blob/68a0cc8546b9a37c181c0b0a025eb9b62dbedd09/src/deny-imports.ts
@@ -8,37 +8,131 @@ export function validateImportPlugin(): Plugin {
     name: 'rsc:validate-imports',
     resolveId: {
       order: 'pre',
-      async handler(source, importer, options) {
+      async handler(source, _importer, options) {
         // optimizer is not aware of server/client boudnary so skip
         if ('scan' in options && options.scan) {
           return
         }
 
         // Validate client-only imports in server environments
-        if (source === 'client-only') {
-          if (this.environment.name === 'rsc') {
-            throw new Error(
-              `'client-only' cannot be imported in server build (importer: '${importer ?? 'unknown'}', environment: ${this.environment.name})`,
-            )
+        if (source === 'client-only' || source === 'server-only') {
+          if (
+            (source === 'client-only' && this.environment.name === 'rsc') ||
+            (source === 'server-only' && this.environment.name !== 'rsc')
+          ) {
+            return {
+              id: `\0virtual:vite-rsc/validate-imports/invalid/${source}`,
+              moduleSideEffects: true,
+            }
           }
-          return { id: `\0virtual:vite-rsc/empty`, moduleSideEffects: false }
-        }
-        if (source === 'server-only') {
-          if (this.environment.name !== 'rsc') {
-            throw new Error(
-              `'server-only' cannot be imported in client build (importer: '${importer ?? 'unknown'}', environment: ${this.environment.name})`,
-            )
+          return {
+            id: `\0virtual:vite-rsc/validate-imports/valid/${source}`,
+            moduleSideEffects: false,
           }
-          return { id: `\0virtual:vite-rsc/empty`, moduleSideEffects: false }
         }
 
         return
       },
     },
     load(id) {
-      if (id.startsWith('\0virtual:vite-rsc/empty')) {
+      if (id.startsWith('\0virtual:vite-rsc/validate-imports/invalid/')) {
+        // it should surface as build error but we make a runtime error just in case.
+        const source = id.slice(id.lastIndexOf('/') + 1)
+        return `throw new Error("invalid import of '${source}'")`
+      }
+      if (id.startsWith('\0virtual:vite-rsc/validate-imports/')) {
         return `export {}`
       }
     },
+    // need a different way to probe module graph for dev and build
+    transform: {
+      order: 'post',
+      async handler(_code, id) {
+        if (this.environment.mode === 'dev') {
+          if (id.startsWith(`\0virtual:vite-rsc/validate-imports/invalid/`)) {
+            const chain = getImportChainDev(this.environment, id)
+            const error = formatError(chain, this.environment.name)
+            if (error) {
+              this.error({
+                id: chain[1],
+                message: error,
+              })
+            }
+          }
+        }
+      },
+    },
+    generateBundle() {
+      if (this.environment.mode === 'build') {
+        const serverOnly = getImportChainBuild(
+          this,
+          '\0virtual:vite-rsc/validate-imports/invalid/server-only',
+        )
+        const serverOnlyError = formatError(serverOnly, this.environment.name)
+        if (serverOnlyError) {
+          throw new Error(serverOnlyError)
+        }
+        const clientOnly = getImportChainBuild(
+          this,
+          '\0virtual:vite-rsc/validate-imports/invalid/client-only',
+        )
+        const clientOnlyError = formatError(clientOnly, this.environment.name)
+        if (clientOnlyError) {
+          throw new Error(clientOnlyError)
+        }
+      }
+    },
+  }
+}
+
+function getImportChainDev(environment: DevEnvironment, id: string) {
+  const chain: string[] = []
+  const recurse = (id: string) => {
+    if (chain.includes(id)) return
+    const info = environment.moduleGraph.getModuleById(id)
+    if (!info) return
+    chain.push(id)
+    const next = [...info.importers][0]
+    if (next && next.id) {
+      recurse(next.id)
+    }
+  }
+  recurse(id)
+  return chain
+}
+
+function getImportChainBuild(ctx: Rollup.PluginContext, id: string): string[] {
+  const chain: string[] = []
+  const recurse = (id: string) => {
+    if (chain.includes(id)) return
+    const info = ctx.getModuleInfo(id)
+    if (!info) return
+    chain.push(id)
+    const next = info.importers[0]
+    if (next) {
+      recurse(next)
+    }
+  }
+  recurse(id)
+  return chain
+}
+
+function formatError(
+  chain: string[],
+  environmentName: string,
+): string | undefined {
+  if (chain.length === 0) return
+  const id = chain[0]!
+  const source = id.slice(id.lastIndexOf('/') + 1)
+  let result = `'${source}' cannot be imported in '${environmentName}' environment:\n`
+  result += chain
+    .slice(1, 6)
+    .map(
+      (id, i) => ' '.repeat(i + 1) + `imported by ${id.replaceAll('\0', '')}\n`,
+    )
+    .join('')
+  if (chain.length > 6) {
+    result += ' '.repeat(7) + '...\n'
   }
+  return result
 }