@tanstack/start-plugin-core 1.142.12 → 1.143.0

package/src/create-server-fn-plugin/compiler.ts

@@ -10,6 +10,7 @@ import { handleCreateServerFn } from './handleCreateServerFn'
 import { handleCreateMiddleware } from './handleCreateMiddleware'
 import { handleCreateIsomorphicFn } from './handleCreateIsomorphicFn'
 import { handleEnvOnlyFn } from './handleEnvOnly'
+import { handleClientOnlyJSX } from './handleClientOnlyJSX'
 import type { MethodChainPaths, RewriteCandidate } from './types'
 
 type Binding =
@@ -38,6 +39,7 @@ export type LookupKind =
   | 'IsomorphicFn'
   | 'ServerOnlyFn'
   | 'ClientOnlyFn'
+  | 'ClientOnlyJSX'
 
 // Detection strategy for each kind
 type MethodChainSetup = {
@@ -49,8 +51,12 @@ type MethodChainSetup = {
   allowRootAsCandidate?: boolean
 }
 type DirectCallSetup = { type: 'directCall' }
+type JSXSetup = { type: 'jsx'; componentName: string }
 
-const LookupSetup: Record<LookupKind, MethodChainSetup | DirectCallSetup> = {
+const LookupSetup: Record<
+  LookupKind,
+  MethodChainSetup | DirectCallSetup | JSXSetup
+> = {
   ServerFn: {
     type: 'methodChain',
     candidateCallIdentifier: new Set(['handler']),
@@ -66,6 +72,7 @@ const LookupSetup: Record<LookupKind, MethodChainSetup | DirectCallSetup> = {
   },
   ServerOnlyFn: { type: 'directCall' },
   ClientOnlyFn: { type: 'directCall' },
+  ClientOnlyJSX: { type: 'jsx', componentName: 'ClientOnly' },
 }
 
 // Single source of truth for detecting which kinds are present in code
@@ -78,6 +85,7 @@ export const KindDetectionPatterns: Record<LookupKind, RegExp> = {
   IsomorphicFn: /createIsomorphicFn/,
   ServerOnlyFn: /createServerOnlyFn/,
   ClientOnlyFn: /createClientOnlyFn/,
+  ClientOnlyJSX: /<ClientOnly|import\s*\{[^}]*\bClientOnly\b/,
 }
 
 // Which kinds are valid for each environment
@@ -94,6 +102,7 @@ export const LookupKindsPerEnv: Record<'client' | 'server', Set<LookupKind>> = {
     'IsomorphicFn',
     'ServerOnlyFn',
     'ClientOnlyFn',
+    'ClientOnlyJSX', // Only transform on server to remove children
   ] as const),
 }
 
@@ -169,13 +178,53 @@ interface ModuleInfo {
 function needsDirectCallDetection(kinds: Set<LookupKind>): boolean {
   for (const kind of kinds) {
     const setup = LookupSetup[kind]
-    if (setup.type === 'directCall' || setup.allowRootAsCandidate) {
+    if (
+      setup.type === 'directCall' ||
+      (setup.type === 'methodChain' && setup.allowRootAsCandidate)
+    ) {
       return true
     }
   }
   return false
 }
 
+/**
+ * Checks if all kinds in the set are guaranteed to be top-level only.
+ * Only ServerFn is always declared at module level (must be assigned to a variable).
+ * Middleware, IsomorphicFn, ServerOnlyFn, ClientOnlyFn can be nested inside functions.
+ * When all kinds are top-level-only, we can use a fast scan instead of full traversal.
+ */
+function areAllKindsTopLevelOnly(kinds: Set<LookupKind>): boolean {
+  return kinds.size === 1 && kinds.has('ServerFn')
+}
+
+/**
+ * Checks if we need to detect JSX elements (e.g., <ClientOnly>).
+ */
+function needsJSXDetection(kinds: Set<LookupKind>): boolean {
+  for (const kind of kinds) {
+    const setup = LookupSetup[kind]
+    if (setup.type === 'jsx') {
+      return true
+    }
+  }
+  return false
+}
+
+/**
+ * Gets the set of JSX component names to detect.
+ */
+function getJSXComponentNames(kinds: Set<LookupKind>): Set<string> {
+  const names = new Set<string>()
+  for (const kind of kinds) {
+    const setup = LookupSetup[kind]
+    if (setup.type === 'jsx') {
+      names.add(setup.componentName)
+    }
+  }
+  return names
+}
+
 /**
  * Checks if a CallExpression is a direct-call candidate for NESTED detection.
  * Returns true if the callee is a known factory function name.
@@ -234,6 +283,11 @@ export class ServerFnCompiler {
   private moduleCache = new Map<string, ModuleInfo>()
   private initialized = false
   private validLookupKinds: Set<LookupKind>
+  private resolveIdCache = new Map<string, string | null>()
+  private exportResolutionCache = new Map<
+    string,
+    Map<string, { moduleInfo: ModuleInfo; binding: Binding } | null>
+  >()
   // Fast lookup for direct imports from known libraries (e.g., '@tanstack/react-start')
   // Maps: libName → (exportName → Kind)
   // This allows O(1) resolution for the common case without async resolveId calls
@@ -246,11 +300,44 @@ export class ServerFnCompiler {
       lookupKinds: Set<LookupKind>
       loadModule: (id: string) => Promise<void>
       resolveId: (id: string, importer?: string) => Promise<string | null>
+      /**
+       * In 'build' mode, resolution results are cached for performance.
+       * In 'dev' mode (default), caching is disabled to avoid invalidation complexity with HMR.
+       */
+      mode?: 'dev' | 'build'
     },
   ) {
     this.validLookupKinds = options.lookupKinds
   }
 
+  private get mode(): 'dev' | 'build' {
+    return this.options.mode ?? 'dev'
+  }
+
+  private async resolveIdCached(id: string, importer?: string) {
+    if (this.mode === 'dev') {
+      return this.options.resolveId(id, importer)
+    }
+
+    const cacheKey = importer ? `${importer}::${id}` : id
+    const cached = this.resolveIdCache.get(cacheKey)
+    if (cached !== undefined) {
+      return cached
+    }
+    const resolved = await this.options.resolveId(id, importer)
+    this.resolveIdCache.set(cacheKey, resolved)
+    return resolved
+  }
+
+  private getExportResolutionCache(moduleId: string) {
+    let cache = this.exportResolutionCache.get(moduleId)
+    if (!cache) {
+      cache = new Map()
+      this.exportResolutionCache.set(moduleId, cache)
+    }
+    return cache
+  }
+
   private async init() {
     // Register internal stub package exports for recognition.
     // These don't need module resolution - only the knownRootImports fast path.
@@ -274,7 +361,18 @@ export class ServerFnCompiler {
         }
         libExports.set(config.rootExport, config.kind)
 
-        const libId = await this.options.resolveId(config.libName)
+        // For JSX lookups (e.g., ClientOnlyJSX), we only need the knownRootImports
+        // fast path to verify imports. Skip module resolution which may fail if
+        // the package isn't a direct dependency (e.g., @tanstack/react-router from
+        // within start-plugin-core).
+        if (config.kind !== 'Root') {
+          const setup = LookupSetup[config.kind]
+          if (setup.type === 'jsx') {
+            return
+          }
+        }
+
+        const libId = await this.resolveIdCached(config.libName)
         if (!libId) {
           throw new Error(`could not resolve "${config.libName}"`)
         }
@@ -311,9 +409,14 @@ export class ServerFnCompiler {
     this.initialized = true
   }
 
-  public ingestModule({ code, id }: { code: string; id: string }) {
-    const ast = parseAst({ code })
-
+  /**
+   * Extracts bindings and exports from an already-parsed AST.
+   * This is the core logic shared by ingestModule and ingestModuleFromAst.
+   */
+  private extractModuleInfo(
+    ast: ReturnType<typeof parseAst>,
+    id: string,
+  ): ModuleInfo {
     const bindings = new Map<string, Binding>()
     const exports = new Map<string, ExportEntry>()
     const reExportAllSources: Array<string> = []
@@ -414,10 +517,19 @@ export class ServerFnCompiler {
       reExportAllSources,
     }
     this.moduleCache.set(id, info)
+    return info
+  }
+
+  public ingestModule({ code, id }: { code: string; id: string }) {
+    const ast = parseAst({ code })
+    const info = this.extractModuleInfo(ast, id)
     return { info, ast }
   }
 
   public invalidateModule(id: string) {
+    // Note: Resolution caches (resolveIdCache, exportResolutionCache) are only
+    // used in build mode where there's no HMR. In dev mode, caching is disabled,
+    // so we only need to invalidate the moduleCache here.
     return this.moduleCache.delete(id)
   }
 
@@ -448,7 +560,13 @@ export class ServerFnCompiler {
     }
 
     const checkDirectCalls = needsDirectCallDetection(fileKinds)
+    // Optimization: ServerFn is always a top-level declaration (must be assigned to a variable).
+    // If the file only has ServerFn, we can skip full AST traversal and only visit
+    // the specific top-level declarations that have candidates.
+    const canUseFastPath = areAllKindsTopLevelOnly(fileKinds)
 
+    // Always parse and extract module info upfront.
+    // This ensures the module is cached for import resolution even if no candidates are found.
     const { ast } = this.ingestModule({ code, id })
 
     // Single-pass traversal to:
@@ -462,40 +580,145 @@ export class ServerFnCompiler {
       babel.NodePath<t.CallExpression>
     >()
 
-    babel.traverse(ast, {
-      CallExpression: (path) => {
-        const node = path.node
-        const parent = path.parent
-
-        // Check if this call is part of a larger chain (inner call)
-        // If so, store it for method chain lookup but don't treat as candidate
-        if (
-          t.isMemberExpression(parent) &&
-          t.isCallExpression(path.parentPath.parent)
-        ) {
-          // This is an inner call in a chain - store for later lookup
-          chainCallPaths.set(node, path)
-          return
+    // JSX candidates (e.g., <ClientOnly>)
+    const jsxCandidatePaths: Array<babel.NodePath<t.JSXElement>> = []
+    const checkJSX = needsJSXDetection(fileKinds)
+    // Get target component names from JSX setup (e.g., 'ClientOnly')
+    const jsxTargetComponentNames = checkJSX
+      ? getJSXComponentNames(fileKinds)
+      : null
+    // Get module info that was just cached by ingestModule
+    const moduleInfo = this.moduleCache.get(id)!
+
+    if (canUseFastPath) {
+      // Fast path: only visit top-level statements that have potential candidates
+
+      // Collect indices of top-level statements that contain candidates
+      const candidateIndices: Array<number> = []
+      for (let i = 0; i < ast.program.body.length; i++) {
+        const node = ast.program.body[i]!
+        let declarations: Array<t.VariableDeclarator> | undefined
+
+        if (t.isVariableDeclaration(node)) {
+          declarations = node.declarations
+        } else if (t.isExportNamedDeclaration(node) && node.declaration) {
+          if (t.isVariableDeclaration(node.declaration)) {
+            declarations = node.declaration.declarations
+          }
         }
 
-        // Pattern 1: Method chain pattern (.handler(), .server(), .client(), etc.)
-        if (isMethodChainCandidate(node, fileKinds)) {
-          candidatePaths.push(path)
-          return
+        if (declarations) {
+          for (const decl of declarations) {
+            if (decl.init && t.isCallExpression(decl.init)) {
+              if (isMethodChainCandidate(decl.init, fileKinds)) {
+                candidateIndices.push(i)
+                break // Only need to mark this statement once
+              }
+            }
+          }
         }
+      }
 
-        // Pattern 2: Direct call pattern
-        if (checkDirectCalls) {
-          if (isTopLevelDirectCallCandidate(path)) {
-            candidatePaths.push(path)
-          } else if (isNestedDirectCallCandidate(node)) {
+      // Early exit: no potential candidates found at top level
+      if (candidateIndices.length === 0) {
+        return null
+      }
+
+      // Targeted traversal: only visit the specific statements that have candidates
+      // This is much faster than traversing the entire AST
+      babel.traverse(ast, {
+        Program(programPath) {
+          const bodyPaths = programPath.get('body')
+          for (const idx of candidateIndices) {
+            const stmtPath = bodyPaths[idx]
+            if (!stmtPath) continue
+
+            // Traverse only this statement's subtree
+            stmtPath.traverse({
+              CallExpression(path) {
+                const node = path.node
+                const parent = path.parent
+
+                // Check if this call is part of a larger chain (inner call)
+                if (
+                  t.isMemberExpression(parent) &&
+                  t.isCallExpression(path.parentPath.parent)
+                ) {
+                  chainCallPaths.set(node, path)
+                  return
+                }
+
+                // Method chain pattern
+                if (isMethodChainCandidate(node, fileKinds)) {
+                  candidatePaths.push(path)
+                }
+              },
+            })
+          }
+          // Stop traversal after processing Program
+          programPath.stop()
+        },
+      })
+    } else {
+      // Normal path: full traversal for non-fast-path kinds
+      babel.traverse(ast, {
+        CallExpression: (path) => {
+          const node = path.node
+          const parent = path.parent
+
+          // Check if this call is part of a larger chain (inner call)
+          // If so, store it for method chain lookup but don't treat as candidate
+          if (
+            t.isMemberExpression(parent) &&
+            t.isCallExpression(path.parentPath.parent)
+          ) {
+            // This is an inner call in a chain - store for later lookup
+            chainCallPaths.set(node, path)
+            return
+          }
+
+          // Pattern 1: Method chain pattern (.handler(), .server(), .client(), etc.)
+          if (isMethodChainCandidate(node, fileKinds)) {
             candidatePaths.push(path)
+            return
           }
-        }
-      },
-    })
 
-    if (candidatePaths.length === 0) {
+          // Pattern 2: Direct call pattern
+          if (checkDirectCalls) {
+            if (isTopLevelDirectCallCandidate(path)) {
+              candidatePaths.push(path)
+            } else if (isNestedDirectCallCandidate(node)) {
+              candidatePaths.push(path)
+            }
+          }
+        },
+        // Pattern 3: JSX element pattern (e.g., <ClientOnly>)
+        // Collect JSX elements where the component name matches a known import
+        // that resolves to a target component (e.g., ClientOnly from @tanstack/react-router)
+        JSXElement: (path) => {
+          if (!checkJSX || !jsxTargetComponentNames) return
+
+          const openingElement = path.node.openingElement
+          const nameNode = openingElement.name
+
+          // Only handle simple identifier names (not namespaced or member expressions)
+          if (!t.isJSXIdentifier(nameNode)) return
+
+          const componentName = nameNode.name
+          const binding = moduleInfo.bindings.get(componentName)
+
+          // Must be an import binding
+          if (!binding || binding.type !== 'import') return
+
+          // Check if the original import name matches a target component
+          if (jsxTargetComponentNames.has(binding.importedName)) {
+            jsxCandidatePaths.push(path)
+          }
+        },
+      })
+    }
+
+    if (candidatePaths.length === 0 && jsxCandidatePaths.length === 0) {
       return null
     }
 
@@ -512,7 +735,7 @@ export class ServerFnCompiler {
       this.validLookupKinds.has(kind as LookupKind),
     ) as Array<{ path: babel.NodePath<t.CallExpression>; kind: LookupKind }>
 
-    if (validCandidates.length === 0) {
+    if (validCandidates.length === 0 && jsxCandidatePaths.length === 0) {
       return null
     }
 
@@ -605,6 +828,29 @@ export class ServerFnCompiler {
       }
     }
 
+    // Handle JSX candidates (e.g., <ClientOnly>)
+    // Note: We only reach here on the server (ClientOnlyJSX is only in LookupKindsPerEnv.server)
+    // Verify import source using knownRootImports (same as function call resolution)
+    for (const jsxPath of jsxCandidatePaths) {
+      const openingElement = jsxPath.node.openingElement
+      const nameNode = openingElement.name
+      if (!t.isJSXIdentifier(nameNode)) continue
+
+      const componentName = nameNode.name
+      const binding = moduleInfo.bindings.get(componentName)
+      if (!binding || binding.type !== 'import') continue
+
+      // Verify the import source is a known TanStack router package
+      const knownExports = this.knownRootImports.get(binding.source)
+      if (!knownExports) continue
+
+      // Verify the imported name resolves to ClientOnlyJSX kind
+      const kind = knownExports.get(binding.importedName)
+      if (kind !== 'ClientOnlyJSX') continue
+
+      handleClientOnlyJSX(jsxPath, { env: 'server' })
+    }
+
     deadCodeElimination(ast, refIdents)
 
     return generateFromAst(ast, {
@@ -651,6 +897,19 @@ export class ServerFnCompiler {
     exportName: string,
     visitedModules = new Set<string>(),
   ): Promise<{ moduleInfo: ModuleInfo; binding: Binding } | undefined> {
+    const isBuildMode = this.mode === 'build'
+
+    // Check cache first (only for top-level calls in build mode)
+    if (isBuildMode && visitedModules.size === 0) {
+      const moduleCache = this.exportResolutionCache.get(moduleInfo.id)
+      if (moduleCache) {
+        const cached = moduleCache.get(exportName)
+        if (cached !== undefined) {
+          return cached ?? undefined
+        }
+      }
+    }
+
     // Prevent infinite loops in circular re-exports
     if (visitedModules.has(moduleInfo.id)) {
       return undefined
@@ -662,7 +921,12 @@ export class ServerFnCompiler {
     if (directExport) {
       const binding = moduleInfo.bindings.get(directExport.name)
       if (binding) {
-        return { moduleInfo, binding }
+        const result = { moduleInfo, binding }
+        // Cache the result (build mode only)
+        if (isBuildMode) {
+          this.getExportResolutionCache(moduleInfo.id).set(exportName, result)
+        }
+        return result
       }
     }
 
@@ -671,10 +935,11 @@ export class ServerFnCompiler {
     if (moduleInfo.reExportAllSources.length > 0) {
       const results = await Promise.all(
         moduleInfo.reExportAllSources.map(async (reExportSource) => {
-          const reExportTarget = await this.options.resolveId(
+          const reExportTarget = await this.resolveIdCached(
             reExportSource,
             moduleInfo.id,
           )
+
           if (reExportTarget) {
             const reExportModule = await this.getModuleInfo(reExportTarget)
             return this.findExportInModule(
@@ -689,11 +954,19 @@ export class ServerFnCompiler {
       // Return the first valid result
       for (const result of results) {
         if (result) {
+          // Cache the result (build mode only)
+          if (isBuildMode) {
+            this.getExportResolutionCache(moduleInfo.id).set(exportName, result)
+          }
           return result
         }
       }
     }
 
+    // Cache negative result (build mode only)
+    if (isBuildMode) {
+      this.getExportResolutionCache(moduleInfo.id).set(exportName, null)
+    }
     return undefined
   }
 
@@ -719,7 +992,7 @@ export class ServerFnCompiler {
       }
 
       // Slow path: resolve through the module graph
-      const target = await this.options.resolveId(binding.source, fileId)
+      const target = await this.resolveIdCached(binding.source, fileId)
       if (!target) {
         return 'None'
       }
@@ -863,7 +1136,7 @@ export class ServerFnCompiler {
           binding.importedName === '*'
         ) {
           // resolve the property from the target module
-          const targetModuleId = await this.options.resolveId(
+          const targetModuleId = await this.resolveIdCached(
             binding.source,
             fileId,
           )

package/src/create-server-fn-plugin/handleClientOnlyJSX.ts

@@ -0,0 +1,32 @@
+import type * as t from '@babel/types'
+import type * as babel from '@babel/core'
+
+/**
+ * Handles <ClientOnly> JSX elements on the server side.
+ *
+ * On the server, the children of <ClientOnly> should be removed since they
+ * are client-only code. Only the fallback prop (if present) will be rendered.
+ *
+ * Transform:
+ *   <ClientOnly fallback={<Loading />}>{clientOnlyContent}</ClientOnly>
+ * Into:
+ *   <ClientOnly fallback={<Loading />} />
+ *
+ * Or if no fallback:
+ *   <ClientOnly>{clientOnlyContent}</ClientOnly>
+ * Into:
+ *   <ClientOnly />
+ */
+export function handleClientOnlyJSX(
+  path: babel.NodePath<t.JSXElement>,
+  _opts: { env: 'server' },
+): void {
+  const element = path.node
+
+  // Remove all children - they are client-only code
+  element.children = []
+
+  // Make it a self-closing element since there are no children
+  element.openingElement.selfClosing = true
+  element.closingElement = null
+}

package/src/create-server-fn-plugin/plugin.ts

@@ -75,7 +75,15 @@ const getLookupConfigurationsForEnv = (
       ...commonConfigs,
     ]
   } else {
-    return commonConfigs
+    // Server-only: add ClientOnly JSX component lookup
+    return [
+      ...commonConfigs,
+      {
+        libName: `@tanstack/${framework}-router`,
+        rootExport: 'ClientOnly',
+        kind: 'ClientOnlyJSX',
+      },
+    ]
   }
 }
 const SERVER_FN_LOOKUP = 'server-fn-module-lookup'
@@ -118,6 +126,9 @@ export function createServerFnPlugin(opts: {
         async handler(code, id) {
           let compiler = compilers[this.environment.name]
           if (!compiler) {
+            // Default to 'dev' mode for unknown environments (conservative: no caching)
+            const mode =
+              this.environment.mode === 'build' ? 'build' : ('dev' as const)
             compiler = new ServerFnCompiler({
               env: environment.type,
               directive: opts.directive,
@@ -126,6 +137,7 @@ export function createServerFnPlugin(opts: {
                 environment.type,
                 opts.framework,
               ),
+              mode,
               loadModule: async (id: string) => {
                 if (this.environment.mode === 'build') {
                   const loaded = await this.load({ id })