fix: compile createMiddleware using start-compiler-plugin

packages/start-plugin-core/src/create-server-fn-plugin/compiler.ts

@@ -29,10 +29,18 @@ type ExportEntry =
 
 type Kind = 'None' | `Root` | `Builder` | LookupKind
 
-type LookupKind = 'ServerFn' | 'Middleware'
+export type LookupKind = 'ServerFn' | 'Middleware'
+
+const LookupSetup: Record<
+  LookupKind,
+  { candidateCallIdentifier: Set<string> }
+> = {
+  ServerFn: { candidateCallIdentifier: new Set(['handler']) },
+  Middleware: {
+    candidateCallIdentifier: new Set(['server', 'client', 'createMiddlewares']),
+  },
+}
 
-const validLookupKinds: Array<LookupKind> = ['ServerFn', 'Middleware']
-const candidateCallIdentifier = ['handler', 'server', 'client']
 export type LookupConfig = {
   libName: string
   rootExport: string
@@ -48,14 +56,18 @@ interface ModuleInfo {
 export class ServerFnCompiler {
   private moduleCache = new Map<string, ModuleInfo>()
   private initialized = false
+  private validLookupKinds: Set<LookupKind>
   constructor(
     private options: {
       env: 'client' | 'server'
       lookupConfigurations: Array<LookupConfig>
+      lookupKinds: Set<LookupKind>
       loadModule: (id: string) => Promise<void>
       resolveId: (id: string, importer?: string) => Promise<string | null>
     },
-  ) {}
+  ) {
+    this.validLookupKinds = options.lookupKinds
+  }
 
   private async init(id: string) {
     await Promise.all(
@@ -207,7 +219,7 @@ export class ServerFnCompiler {
     }> = []
     for (const handler of candidates) {
       const kind = await this.resolveExprKind(handler, id)
-      if (validLookupKinds.includes(kind as LookupKind)) {
+      if (this.validLookupKinds.has(kind as LookupKind)) {
         toRewrite.push({ callExpression: handler, kind: kind as LookupKind })
       }
     }
@@ -261,7 +273,10 @@ export class ServerFnCompiler {
 
     for (const binding of bindings.values()) {
       if (binding.type === 'var') {
-        const handler = isCandidateCallExpression(binding.init)
+        const handler = isCandidateCallExpression(
+          binding.init,
+          this.validLookupKinds,
+        )
         if (handler) {
           candidates.push(handler)
         }
@@ -368,7 +383,7 @@ export class ServerFnCompiler {
         if (calleeKind === `Root` || calleeKind === `Builder`) {
           return `Builder`
         }
-        for (const kind of validLookupKinds) {
+        for (const kind of this.validLookupKinds) {
           if (calleeKind === kind) {
             return kind
           }
@@ -407,16 +422,18 @@ export class ServerFnCompiler {
     if (t.isMemberExpression(callee) && t.isIdentifier(callee.property)) {
       const prop = callee.property.name
 
-      if (prop === 'handler') {
+      if (
+        this.validLookupKinds.has('ServerFn') &&
+        LookupSetup['ServerFn'].candidateCallIdentifier.has(prop)
+      ) {
         const base = await this.resolveExprKind(callee.object, fileId, visited)
         if (base === 'Root' || base === 'Builder') {
           return 'ServerFn'
         }
         return 'None'
       } else if (
-        prop === 'client' ||
-        prop === 'server' ||
-        prop === 'createMiddleware'
+        this.validLookupKinds.has('Middleware') &&
+        LookupSetup['Middleware'].candidateCallIdentifier.has(prop)
       ) {
         const base = await this.resolveExprKind(callee.object, fileId, visited)
         if (base === 'Root' || base === 'Builder' || base === 'Middleware') {
@@ -483,16 +500,19 @@ export class ServerFnCompiler {
 
 function isCandidateCallExpression(
   node: t.Node | null | undefined,
+  lookupKinds: Set<LookupKind>,
 ): undefined | t.CallExpression {
   if (!t.isCallExpression(node)) return undefined
 
   const callee = node.callee
   if (!t.isMemberExpression(callee) || !t.isIdentifier(callee.property)) {
     return undefined
   }
-  if (!candidateCallIdentifier.includes(callee.property.name)) {
-    return undefined
+  for (const kind of lookupKinds) {
+    if (LookupSetup[kind].candidateCallIdentifier.has(callee.property.name)) {
+      return node
+    }
   }
 
-  return node
+  return undefined
 }

packages/start-plugin-core/src/create-server-fn-plugin/plugin.ts

@@ -1,5 +1,6 @@
 import { VITE_ENVIRONMENT_NAMES } from '../constants'
 import { ServerFnCompiler } from './compiler'
+import type { LookupConfig, LookupKind } from './compiler'
 import type { CompileStartFrameworkOptions } from '../start-compiler-plugin/compilers'
 import type { ViteEnvironmentNames } from '../constants'
 import type { PluginOption } from 'vite'
@@ -8,6 +9,36 @@ function cleanId(id: string): string {
   return id.split('?')[0]!
 }
 
+const LookupKindsPerEnv: Record<'client' | 'server', Set<LookupKind>> = {
+  client: new Set(['Middleware', 'ServerFn'] as const),
+  server: new Set(['ServerFn'] as const),
+}
+
+const getLookupConfigurationsForEnv = (
+  env: 'client' | 'server',
+  framework: CompileStartFrameworkOptions,
+): Array<LookupConfig> => {
+  const createServerFnConfig: LookupConfig = {
+    libName: `@tanstack/${framework}-start`,
+    rootExport: 'createServerFn',
+  }
+  if (env === 'client') {
+    return [
+      {
+        libName: `@tanstack/${framework}-start`,
+        rootExport: 'createMiddleware',
+      },
+      {
+        libName: `@tanstack/${framework}-start`,
+        rootExport: 'createStart',
+      },
+
+      createServerFnConfig,
+    ]
+  } else {
+    return [createServerFnConfig]
+  }
+}
 export function createServerFnPlugin(
   framework: CompileStartFrameworkOptions,
 ): PluginOption {
@@ -52,8 +83,9 @@ export function createServerFnPlugin(
             exclude: new RegExp(`${SERVER_FN_LOOKUP}$`),
           },
           code: {
-            // only scan files that mention `.handler(` | `.server(` | `.client(`
-            include: [/\.handler\(/, /\.server\(/, /\.client\(/],
+            // TODO apply this plugin with a different filter per environment so that .createMiddleware() calls are not scanned in server env
+            // only scan files that mention `.handler(` | `.createMiddleware()`
+            include: [/\.handler\(/, /.createMiddleware\(\)/],
           },
         },
         async handler(code, id) {
@@ -73,21 +105,11 @@ export function createServerFnPlugin(
 
             compiler = new ServerFnCompiler({
               env,
-              lookupConfigurations: [
-                {
-                  libName: `@tanstack/${framework}-start`,
-                  rootExport: 'createMiddleware',
-                },
-
-                {
-                  libName: `@tanstack/${framework}-start`,
-                  rootExport: 'createServerFn',
-                },
-                {
-                  libName: `@tanstack/${framework}-start`,
-                  rootExport: 'createStart',
-                },
-              ],
+              lookupKinds: LookupKindsPerEnv[env],
+              lookupConfigurations: getLookupConfigurationsForEnv(
+                env,
+                framework,
+              ),
               loadModule: async (id: string) => {
                 if (this.environment.mode === 'build') {
                   const loaded = await this.load({ id })

packages/start-plugin-core/src/start-compiler-plugin/compilers.ts

@@ -6,6 +6,7 @@ import {
   findReferencedIdentifiers,
 } from 'babel-dead-code-elimination'
 import { generateFromAst, parseAst } from '@tanstack/router-utils'
+import { handleCreateMiddleware } from '../create-server-fn-plugin/handleCreateMiddleware'
 import { transformFuncs } from './constants'
 import { handleCreateIsomorphicFnCallExpression } from './isomorphicFn'
 import {
@@ -38,6 +39,11 @@ export function compileStartOutputFactory(
         handleCallExpression: handleCreateIsomorphicFnCallExpression,
         paths: [],
       },
+      createMiddleware: {
+        name: 'createMiddleware',
+        handleCallExpression: handleCreateMiddleware,
+        paths: [],
+      },
     }
 
     const ast = parseAst(opts)

packages/start-plugin-core/src/start-compiler-plugin/constants.ts

@@ -2,4 +2,5 @@ export const transformFuncs = [
   'createServerOnlyFn',
   'createClientOnlyFn',
   'createIsomorphicFn',
+  'createMiddleware',
 ] as const

packages/start-plugin-core/tests/createMiddleware/createMiddleware.test.ts → packages/start-plugin-core/tests/createMiddleware-create-server-fn-plugin/createMiddleware.test.ts

@@ -17,16 +17,12 @@ async function compile(opts: {
     loadModule: async (id) => {
       // do nothing in test
     },
+    lookupKinds: new Set(['Middleware']),
     lookupConfigurations: [
       {
         libName: `@tanstack/react-start`,
         rootExport: 'createMiddleware',
       },
-
-      {
-        libName: `@tanstack/react-start`,
-        rootExport: 'createServerFn',
-      },
       {
         libName: `@tanstack/react-start`,
         rootExport: 'createStart',

packages/start-plugin-core/tests/createMiddleware-start-compiler-plugin/createMiddleware.test.ts

@@ -0,0 +1,37 @@
+import { readFile, readdir } from 'node:fs/promises'
+import path from 'node:path'
+import { describe, expect, test } from 'vitest'
+import { compileStartOutputFactory } from '../../src/start-compiler-plugin/compilers'
+
+const compileStartOutput = compileStartOutputFactory('react')
+
+async function getFilenames() {
+  return await readdir(path.resolve(import.meta.dirname, './test-files'))
+}
+
+describe('createMiddleware compiles correctly', async () => {
+  const filenames = await getFilenames()
+
+  describe.each(filenames)('should handle "%s"', async (filename) => {
+    const file = await readFile(
+      path.resolve(import.meta.dirname, `./test-files/${filename}`),
+    )
+    const code = file.toString()
+
+    test.each(['client', 'server'] as const)(
+      `should compile for ${filename} %s`,
+      async (env) => {
+        const compiledResult = compileStartOutput({
+          env,
+          code,
+          filename,
+          dce: false,
+        })
+
+        await expect(compiledResult.code).toMatchFileSnapshot(
+          `./snapshots/${env}/${filename}`,
+        )
+      },
+    )
+  })
+})

packages/start-plugin-core/tests/createMiddleware/snapshots/client/create-function-middleware.ts → packages/start-plugin-core/tests/createMiddleware-start-compiler-plugin/snapshots/client/create-function-middleware.ts

@@ -1,4 +1,5 @@
 import { createMiddleware } from '@tanstack/react-start';
+import { foo } from '@some/lib';
 export const fnMw = createMiddleware({
   type: 'function'
 }).client(() => {

packages/start-plugin-core/tests/createMiddleware-start-compiler-plugin/snapshots/client/createMiddlewareDestructured.tsx

@@ -0,0 +1,24 @@
+import { createMiddleware } from '@tanstack/react-start';
+import { z } from 'zod';
+export const withUseServer = createMiddleware({
+  id: 'test'
+});
+export const withoutUseServer = createMiddleware({
+  id: 'test'
+});
+export const withVariable = createMiddleware({
+  id: 'test'
+});
+async function abstractedFunction() {
+  console.info('Fetching posts...');
+  await new Promise(r => setTimeout(r, 500));
+  return axios.get<Array<PostType>>('https://jsonplaceholder.typicode.com/posts').then(r => r.data.slice(0, 10));
+}
+function zodValidator<TSchema extends z.ZodSchema, TResult>(schema: TSchema, fn: (input: z.output<TSchema>) => TResult) {
+  return async (input: unknown) => {
+    return fn(schema.parse(input));
+  };
+}
+export const withZodValidator = createMiddleware({
+  id: 'test'
+});

packages/start-plugin-core/tests/createMiddleware-start-compiler-plugin/snapshots/client/createMiddlewareDestructuredRename.tsx

@@ -0,0 +1,24 @@
+import { createMiddleware as middlewareFn } from '@tanstack/react-start';
+import { z } from 'zod';
+export const withUseServer = middlewareFn({
+  id: 'test'
+});
+export const withoutUseServer = middlewareFn({
+  id: 'test'
+});
+export const withVariable = middlewareFn({
+  id: 'test'
+});
+async function abstractedFunction() {
+  console.info('Fetching posts...');
+  await new Promise(r => setTimeout(r, 500));
+  return axios.get<Array<PostType>>('https://jsonplaceholder.typicode.com/posts').then(r => r.data.slice(0, 10));
+}
+function zodValidator<TSchema extends z.ZodSchema, TResult>(schema: TSchema, fn: (input: z.output<TSchema>) => TResult) {
+  return async (input: unknown) => {
+    return fn(schema.parse(input));
+  };
+}
+export const withZodValidator = middlewareFn({
+  id: 'test'
+});

packages/start-plugin-core/tests/createMiddleware-start-compiler-plugin/snapshots/client/createMiddlewareStarImport.tsx

@@ -0,0 +1,24 @@
+import * as TanStackStart from '@tanstack/react-start';
+import { z } from 'zod';
+export const withUseServer = TanStackStart.createMiddleware({
+  id: 'test'
+});
+export const withoutUseServer = TanStackStart.createMiddleware({
+  id: 'test'
+});
+export const withVariable = TanStackStart.createMiddleware({
+  id: 'test'
+});
+async function abstractedFunction() {
+  console.info('Fetching posts...');
+  await new Promise(r => setTimeout(r, 500));
+  return axios.get<Array<PostType>>('https://jsonplaceholder.typicode.com/posts').then(r => r.data.slice(0, 10));
+}
+function zodValidator<TSchema extends z.ZodSchema, TResult>(schema: TSchema, fn: (input: z.output<TSchema>) => TResult) {
+  return async (input: unknown) => {
+    return fn(schema.parse(input));
+  };
+}
+export const withZodValidator = TanStackStart.createMiddleware({
+  id: 'test'
+});

packages/start-plugin-core/tests/createMiddleware/snapshots/client/createMiddlewareValidator.tsx → packages/start-plugin-core/tests/createMiddleware-start-compiler-plugin/snapshots/client/createMiddlewareValidator.tsx

@@ -1,4 +1,5 @@
 import { createMiddleware } from '@tanstack/react-start';
+import { z } from 'zod';
 export const withUseServer = createMiddleware({
   id: 'test'
 });

packages/start-plugin-core/tests/createMiddleware-start-compiler-plugin/snapshots/client/factory.ts

@@ -0,0 +1,15 @@
+import { createMiddleware } from '@tanstack/react-start';
+import { getCookie } from '@tanstack/react-start/server';
+interface AuthMiddlewareOptions {
+  allowUnauthenticated?: boolean;
+}
+interface AuthContext {
+  session: {
+    id: string;
+  } | null;
+}
+export const createAuthMiddleware = (opts: AuthMiddlewareOptions = {
+  allowUnauthenticated: false
+}) => createMiddleware({
+  type: 'function'
+});