fix: generic env handling

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

@@ -146,7 +146,7 @@ export function handleCreateServerFn(
             ),
           ),
         ],
-        [t.directive(t.directiveLiteral('use server'))],
+        [t.directive(t.directiveLiteral(opts.directive))],
       ),
     ),
   )

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

@@ -1,8 +1,7 @@
-import { TRANSFORM_ID_REGEX, VITE_ENVIRONMENT_NAMES } from '../constants'
+import { TRANSFORM_ID_REGEX } 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'
 
 function cleanId(id: string): string {
@@ -39,44 +38,29 @@ const getLookupConfigurationsForEnv = (
     return [createServerFnConfig]
   }
 }
+const SERVER_FN_LOOKUP = 'server-fn-module-lookup'
 export function createServerFnPlugin(opts: {
   framework: CompileStartFrameworkOptions
   directive: string
+  environments: Array<{ name: string; type: 'client' | 'server' }>
 }): PluginOption {
-  const SERVER_FN_LOOKUP = 'server-fn-module-lookup'
+  const compilers: Record<string /* envName */, ServerFnCompiler> = {}
 
-  const compilers: Partial<Record<ViteEnvironmentNames, ServerFnCompiler>> = {}
-  return [
-    {
-      name: 'tanstack-start-core:capture-server-fn-module-lookup',
-      // we only need this plugin in dev mode
-      apply: 'serve',
-      applyToEnvironment(env) {
-        return [
-          VITE_ENVIRONMENT_NAMES.client,
-          VITE_ENVIRONMENT_NAMES.server,
-        ].includes(env.name as ViteEnvironmentNames)
-      },
-      transform: {
-        filter: {
-          id: new RegExp(`${SERVER_FN_LOOKUP}$`),
-        },
-        handler(code, id) {
-          const compiler =
-            compilers[this.environment.name as ViteEnvironmentNames]
-          compiler?.ingestModule({ code, id: cleanId(id) })
-        },
-      },
-    },
-    {
-      name: 'tanstack-start-core::server-fn',
-      enforce: 'pre',
+  function perEnvServerFnPlugin(environment: {
+    name: string
+    type: 'client' | 'server'
+  }): PluginOption {
+    // in server environments, we don't transform middleware calls
+    const transformCodeFilter =
+      environment.type === 'client'
+        ? [/\.\s*handler\(/, /\.\s*createMiddleware\(\)/]
+        : [/\.\s*handler\(/]
 
+    return {
+      name: `tanstack-start-core::server-fn:${environment.name}`,
+      enforce: 'pre',
       applyToEnvironment(env) {
-        return [
-          VITE_ENVIRONMENT_NAMES.client,
-          VITE_ENVIRONMENT_NAMES.server,
-        ].includes(env.name as ViteEnvironmentNames)
+        return env.name === environment.name
       },
       transform: {
         filter: {
@@ -85,32 +69,18 @@ export function createServerFnPlugin(opts: {
             include: TRANSFORM_ID_REGEX,
           },
           code: {
-            // 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: [/\.\s*handler\(/, /\.\s*createMiddleware\(\)/],
+            include: transformCodeFilter,
           },
         },
         async handler(code, id) {
-          let compiler =
-            compilers[this.environment.name as ViteEnvironmentNames]
+          let compiler = compilers[this.environment.name]
           if (!compiler) {
-            const env =
-              this.environment.name === VITE_ENVIRONMENT_NAMES.client
-                ? 'client'
-                : this.environment.name === VITE_ENVIRONMENT_NAMES.server
-                  ? 'server'
-                  : (() => {
-                      throw new Error(
-                        `Environment ${this.environment.name} not configured`,
-                      )
-                    })()
-
             compiler = new ServerFnCompiler({
-              env,
+              env: environment.type,
               directive: opts.directive,
-              lookupKinds: LookupKindsPerEnv[env],
+              lookupKinds: LookupKindsPerEnv[environment.type],
               lookupConfigurations: getLookupConfigurationsForEnv(
-                env,
+                environment.type,
                 opts.framework,
               ),
               loadModule: async (id: string) => {
@@ -147,7 +117,7 @@ export function createServerFnPlugin(opts: {
                 return null
               },
             })
-            compilers[this.environment.name as ViteEnvironmentNames] = compiler
+            compilers[this.environment.name] = compiler
           }
 
           id = cleanId(id)
@@ -157,8 +127,7 @@ export function createServerFnPlugin(opts: {
       },
 
       hotUpdate(ctx) {
-        const compiler =
-          compilers[this.environment.name as ViteEnvironmentNames]
+        const compiler = compilers[this.environment.name]
 
         ctx.modules.forEach((m) => {
           if (m.id) {
@@ -173,6 +142,27 @@ export function createServerFnPlugin(opts: {
           }
         })
       },
+    }
+  }
+
+  return [
+    ...opts.environments.map(perEnvServerFnPlugin),
+    {
+      name: 'tanstack-start-core:capture-server-fn-module-lookup',
+      // we only need this plugin in dev mode
+      apply: 'serve',
+      applyToEnvironment(env) {
+        return !!opts.environments.find((e) => e.name === env.name)
+      },
+      transform: {
+        filter: {
+          id: new RegExp(`${SERVER_FN_LOOKUP}$`),
+        },
+        handler(code, id) {
+          const compiler = compilers[this.environment.name]
+          compiler?.ingestModule({ code, id: cleanId(id) })
+        },
+      },
     },
   ]
 }

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

@@ -39,6 +39,7 @@ export interface TanStackStartVitePluginCoreOptions {
     ssr?: {
       getServerFnById?: string
     }
+    providerEnv?: string
   }
 }
 
@@ -105,6 +106,19 @@ export function TanStackStartVitePluginCore(
     return bundle
   }
 
+  const environments: Array<{ name: string; type: 'client' | 'server' }> = [
+    { name: VITE_ENVIRONMENT_NAMES.client, type: 'client' },
+    { name: VITE_ENVIRONMENT_NAMES.server, type: 'server' },
+  ]
+  if (
+    corePluginOpts.serverFn?.providerEnv &&
+    !environments.find((e) => e.name === corePluginOpts.serverFn?.providerEnv)
+  ) {
+    environments.push({
+      name: corePluginOpts.serverFn.providerEnv,
+      type: 'server',
+    })
+  }
   return [
     {
       name: 'tanstack-start-core:config',
@@ -341,8 +355,12 @@ export function TanStackStartVitePluginCore(
     },
     tanStackStartRouter(startPluginOpts, getConfig, corePluginOpts),
     // N.B. TanStackStartCompilerPlugin must be before the TanStackServerFnPlugin
-    startCompilerPlugin(corePluginOpts.framework),
-    createServerFnPlugin({ framework: corePluginOpts.framework, directive }),
+    startCompilerPlugin({ framework: corePluginOpts.framework, environments }),
+    createServerFnPlugin({
+      framework: corePluginOpts.framework,
+      directive,
+      environments,
+    }),
 
     TanStackServerFnPlugin({
       // This is the ID that will be available to look up and import
@@ -371,7 +389,8 @@ export function TanStackStartVitePluginCore(
         getRuntimeCode: () =>
           `import { createServerRpc } from '@tanstack/${corePluginOpts.framework}-start/server-rpc'`,
         replacer: (d) => `createServerRpc('${d.functionId}', ${d.fn})`,
-        envName: VITE_ENVIRONMENT_NAMES.server,
+        envName:
+          corePluginOpts.serverFn?.providerEnv || VITE_ENVIRONMENT_NAMES.server,
       },
     }),
     loadEnvPlugin(),

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

@@ -6,11 +6,10 @@ import { VIRTUAL_MODULES } from '@tanstack/start-server-core'
 import { normalizePath } from 'vite'
 import path from 'pathe'
 import { makeIdFiltersToMatchWithQuery } from '@rolldown/pluginutils'
-import { TRANSFORM_ID_REGEX, VITE_ENVIRONMENT_NAMES } from '../constants'
+import { TRANSFORM_ID_REGEX } from '../constants'
 import { compileStartOutputFactory } from './compilers'
 import { transformFuncs } from './constants'
-import type { ViteEnvironmentNames } from '../constants'
-import type { Plugin } from 'vite'
+import type { Plugin, PluginOption } from 'vite'
 import type { CompileStartFrameworkOptions } from './compilers'
 
 const debug =
@@ -36,82 +35,77 @@ function resolvePackage(packageName: string): string {
   return pkgRoot
 }
 
-export function startCompilerPlugin(
-  framework: CompileStartFrameworkOptions,
-): Plugin {
-  const compileStartOutput = compileStartOutputFactory(framework)
-
-  return {
-    name: 'tanstack-start-core:compiler',
-    enforce: 'pre',
-    applyToEnvironment(env) {
-      return [
-        VITE_ENVIRONMENT_NAMES.client,
-        VITE_ENVIRONMENT_NAMES.server,
-      ].includes(env.name as ViteEnvironmentNames)
-    },
-    transform: {
-      filter: {
-        code: tokenRegex,
-        id: {
-          include: TRANSFORM_ID_REGEX,
-          exclude: [
-            VIRTUAL_MODULES.serverFnManifest,
-            // N.B. the following files either just re-export or provide the runtime implementation of those functions
-            // we do not want to include them in the transformation
-            // however, those packages (especially start-client-core ATM) also USE these functions
-            // (namely `createIsomorphicFn` in `packages/start-client-core/src/getRouterInstance.ts`) and thus need to be transformed
-            ...makeIdFiltersToMatchWithQuery([
-              ...resolveRuntimeFiles({
-                package: '@tanstack/start-client-core',
-                files: [
-                  'index.js',
-                  'createIsomorphicFn.js',
-                  'envOnly.js',
-                  'serverFnFetcher.js',
-                  'createStart.js',
-                  'createMiddleware.js',
-                ],
-              }),
-              ...resolveRuntimeFiles({
-                package: '@tanstack/start-server-core',
-                files: ['index.js', 'server-functions-handler.js'],
-              }),
-            ]),
+const transformFilter = {
+  code: tokenRegex,
+  id: {
+    include: TRANSFORM_ID_REGEX,
+    exclude: [
+      VIRTUAL_MODULES.serverFnManifest,
+      // N.B. the following files either just re-export or provide the runtime implementation of those functions
+      // we do not want to include them in the transformation
+      // however, those packages (especially start-client-core ATM) also USE these functions
+      // (namely `createIsomorphicFn` in `packages/start-client-core/src/getRouterInstance.ts`) and thus need to be transformed
+      ...makeIdFiltersToMatchWithQuery([
+        ...resolveRuntimeFiles({
+          package: '@tanstack/start-client-core',
+          files: [
+            'index.js',
+            'createIsomorphicFn.js',
+            'envOnly.js',
+            'serverFnFetcher.js',
+            'createStart.js',
+            'createMiddleware.js',
           ],
-        },
-      },
-      handler(code, id) {
-        const env =
-          this.environment.name === VITE_ENVIRONMENT_NAMES.client
-            ? 'client'
-            : this.environment.name === VITE_ENVIRONMENT_NAMES.server
-              ? 'server'
-              : (() => {
-                  throw new Error(
-                    `Environment ${this.environment.name} not configured`,
-                  )
-                })()
+        }),
+        ...resolveRuntimeFiles({
+          package: '@tanstack/start-server-core',
+          files: ['index.js', 'server-functions-handler.js'],
+        }),
+      ]),
+    ],
+  },
+}
 
-        const url = pathToFileURL(id)
-        url.searchParams.delete('v')
-        id = fileURLToPath(url).replace(/\\/g, '/')
+export function startCompilerPlugin(opts: {
+  framework: CompileStartFrameworkOptions
+  environments: Array<{ name: string; type: 'client' | 'server' }>
+}): PluginOption {
+  const compileStartOutput = compileStartOutputFactory(opts.framework)
 
-        if (debug) console.info(`${env} Compiling Start: `, id)
+  function perEnvCompilerPlugin(environment: {
+    name: string
+    type: 'client' | 'server'
+  }): Plugin {
+    return {
+      name: `tanstack-start-core:compiler:${environment.name}`,
+      enforce: 'pre',
+      applyToEnvironment(env) {
+        return env.name === environment.name
+      },
+      transform: {
+        filter: transformFilter,
+        handler(code, id) {
+          const url = pathToFileURL(id)
+          url.searchParams.delete('v')
+          id = fileURLToPath(url).replace(/\\/g, '/')
+
+          if (debug) console.info(`${environment.name} Compiling Start: `, id)
 
-        const compiled = compileStartOutput({
-          code,
-          filename: id,
-          env,
-        })
+          const compiled = compileStartOutput({
+            code,
+            filename: id,
+            env: environment.type,
+          })
 
-        if (debug) {
-          logDiff(code, compiled.code)
-          console.log('Output:\n', compiled.code + '\n\n')
-        }
+          if (debug) {
+            logDiff(code, compiled.code)
+            console.log('Output:\n', compiled.code + '\n\n')
+          }
 
-        return compiled
+          return compiled
+        },
       },
-    },
+    }
   }
+  return opts.environments.map(perEnvCompilerPlugin)
 }