fix: decouple server functions caller and provider

packages/directive-functions-plugin/src/compilers.ts

@@ -41,34 +41,27 @@ export type ReplacerFn = (opts: {
 
 export type CompileDirectivesOpts = ParseAstOptions & {
   directive: string
-  directiveLabel: string
   getRuntimeCode?: (opts: {
     directiveFnsById: Record<string, DirectiveFn>
   }) => string
   generateFunctionId: GenerateFunctionIdFn
   replacer: ReplacerFn
   filename: string
   root: string
-}
-
-function buildDirectiveSplitParam(opts: CompileDirectivesOpts) {
-  return `tsr-directive-${opts.directive.replace(/[^a-zA-Z0-9]/g, '-')}`
+  isDirectiveSplitParam: boolean
+  directiveSplitParam: string
 }
 
 export function compileDirectives(opts: CompileDirectivesOpts): {
   compiledResult: GeneratorResult
   directiveFnsById: Record<string, DirectiveFn>
-  isDirectiveSplitParam: boolean
 } {
-  const directiveSplitParam = buildDirectiveSplitParam(opts)
-  const isDirectiveSplitParam = opts.filename.includes(directiveSplitParam)
-
   const ast = parseAst(opts)
   const refIdents = findReferencedIdentifiers(ast)
   const directiveFnsById = findDirectives(ast, {
     ...opts,
-    directiveSplitParam,
-    isDirectiveSplitParam,
+    directiveSplitParam: opts.directiveSplitParam,
+    isDirectiveSplitParam: opts.isDirectiveSplitParam,
   })
 
   // Add runtime code if there are directives
@@ -84,7 +77,7 @@ export function compileDirectives(opts: CompileDirectivesOpts): {
   // If we are in the source file, we need to remove all exports
   // then make sure that all of our functions are exported under their
   // directive name
-  if (isDirectiveSplitParam) {
+  if (opts.isDirectiveSplitParam) {
     safeRemoveExports(ast)
 
     // Export a single object with all of the functions
@@ -113,13 +106,12 @@ export function compileDirectives(opts: CompileDirectivesOpts): {
   return {
     compiledResult,
     directiveFnsById,
-    isDirectiveSplitParam,
   }
 }
 
 function findNearestVariableName(
   path: babel.NodePath,
-  directiveLabel: string,
+  directive: string,
 ): string {
   let currentPath: babel.NodePath | null = path
   const nameParts: Array<string> = []
@@ -189,7 +181,7 @@ function findNearestVariableName(
         babel.types.isObjectMethod(currentPath.node)
       ) {
         throw new Error(
-          `${directiveLabel} in ClassMethod or ObjectMethod not supported`,
+          `"${directive}" in ClassMethod or ObjectMethod not supported`,
         )
       }
 
@@ -219,7 +211,6 @@ export function findDirectives(
   ast: babel.types.File,
   opts: ParseAstOptions & {
     directive: string
-    directiveLabel: string
     replacer?: ReplacerFn
     generateFunctionId: GenerateFunctionIdFn
     directiveSplitParam: string
@@ -320,7 +311,7 @@ export function findDirectives(
             throw codeFrameError(
               opts.code,
               nearestBlock.node.loc,
-              `${opts.directiveLabel}s cannot be nested in other blocks or functions`,
+              `"${opts.directive}" cannot be nested in other blocks or functions`,
             )
           }
 
@@ -335,7 +326,7 @@ export function findDirectives(
             throw codeFrameError(
               opts.code,
               directiveFn.node.loc,
-              `${opts.directiveLabel}s must be function declarations or function expressions`,
+              `"${opts.directive}" must be function declarations or function expressions`,
             )
           }
 
@@ -397,7 +388,7 @@ export function findDirectives(
     }
 
     // Find the nearest variable name
-    let functionName = findNearestVariableName(directiveFn, opts.directiveLabel)
+    let functionName = findNearestVariableName(directiveFn, opts.directive)
 
     const incrementFunctionNameVersion = (functionName: string) => {
       const [realReferenceName, count] = functionName.split(/_(\d+)$/)

packages/directive-functions-plugin/src/index.ts

@@ -23,136 +23,104 @@ export type {
 export type DirectiveFunctionsViteEnvOptions = Pick<
   CompileDirectivesOpts,
   'getRuntimeCode' | 'replacer'
-> & {
-  envLabel: string
+>
+export type DirectiveFunctionsViteOptions = DirectiveFunctionsViteEnvOptions & {
+  directive: string
+  onDirectiveFnsById?: (directiveFnsById: Record<string, DirectiveFn>) => void
+  generateFunctionId: GenerateFunctionIdFn
 }
 
-export type DirectiveFunctionsViteOptions = Pick<
-  CompileDirectivesOpts,
-  'directive' | 'directiveLabel'
-> &
-  DirectiveFunctionsViteEnvOptions & {
-    onDirectiveFnsById?: (directiveFnsById: Record<string, DirectiveFn>) => void
-    generateFunctionId: GenerateFunctionIdFn
-  }
-
 const createDirectiveRx = (directive: string) =>
   new RegExp(`"${directive}"|'${directive}'`, 'gm')
 
-export type DirectiveFunctionsVitePluginEnvOptions = Pick<
-  CompileDirectivesOpts,
-  'directive' | 'directiveLabel'
-> & {
-  environments: {
-    client: DirectiveFunctionsViteEnvOptions & { envName?: string }
-    server: DirectiveFunctionsViteEnvOptions & { envName?: string }
-  }
+export type DirectiveFunctionsVitePluginEnvOptions = {
+  directive: string
+  callers: Array<DirectiveFunctionsViteEnvOptions & { envName: string }>
+  provider: DirectiveFunctionsViteEnvOptions & { envName: string }
   onDirectiveFnsById?: (directiveFnsById: Record<string, DirectiveFn>) => void
   generateFunctionId: GenerateFunctionIdFn
 }
 
+function buildDirectiveSplitParam(directive: string) {
+  return `tsr-directive-${directive.replace(/[^a-zA-Z0-9]/g, '-')}`
+}
+
 export function TanStackDirectiveFunctionsPluginEnv(
   opts: DirectiveFunctionsVitePluginEnvOptions,
 ): Plugin {
-  opts = {
-    ...opts,
-    environments: {
-      client: {
-        envName: 'client',
-        ...opts.environments.client,
-      },
-      server: {
-        envName: 'server',
-        ...opts.environments.server,
-      },
-    },
-  }
-
   let root: string = process.cwd()
 
   const directiveRx = createDirectiveRx(opts.directive)
 
+  const appliedEnvironments = new Set([
+    ...opts.callers.map((c) => c.envName),
+    opts.provider.envName,
+  ])
+
+  const directiveSplitParam = buildDirectiveSplitParam(opts.directive)
+
   return {
     name: 'tanstack-start-directive-vite-plugin',
     enforce: 'pre',
     buildStart() {
       root = this.environment.config.root
     },
     applyToEnvironment(env) {
-      return [
-        opts.environments.client.envName,
-        opts.environments.server.envName,
-      ].includes(env.name)
+      return appliedEnvironments.has(env.name)
     },
     transform: {
       filter: {
         code: directiveRx,
       },
       handler(code, id) {
-        const envOptions = [
-          opts.environments.client,
-          opts.environments.server,
-        ].find((e) => e.envName === this.environment.name)
-
-        if (!envOptions) {
-          throw new Error(`Environment ${this.environment.name} not found`)
+        const url = pathToFileURL(id)
+        url.searchParams.delete('v')
+        id = fileURLToPath(url).replace(/\\/g, '/')
+
+        const isDirectiveSplitParam = id.includes(directiveSplitParam)
+
+        let envOptions: DirectiveFunctionsViteEnvOptions & { envName: string }
+        if (isDirectiveSplitParam) {
+          envOptions = opts.provider
+          if (debug)
+            console.info(
+              `Compiling Directives for provider in environment ${envOptions.envName}: `,
+              id,
+            )
+        } else {
+          envOptions = opts.callers.find(
+            (e) => e.envName === this.environment.name,
+          )!
+          if (debug)
+            console.info(
+              `Compiling Directives for caller in environment ${envOptions.envName}: `,
+              id,
+            )
         }
-
-        return transformCode({
-          ...opts,
-          ...envOptions,
+        const { compiledResult, directiveFnsById } = compileDirectives({
+          directive: opts.directive,
+          getRuntimeCode: envOptions.getRuntimeCode,
+          generateFunctionId: opts.generateFunctionId,
+          replacer: envOptions.replacer,
           code,
-          id,
           root,
+          filename: id,
+          directiveSplitParam,
+          isDirectiveSplitParam,
         })
-      },
-    },
-  }
-}
+        // when we process a file with a directive split param, we have already encountered the directives in that file
+        // (otherwise we wouldn't have gotten here)
+        if (!isDirectiveSplitParam) {
+          opts.onDirectiveFnsById?.(directiveFnsById)
+        }
 
-function transformCode({
-  code,
-  id,
-  envLabel,
-  directive,
-  directiveLabel,
-  getRuntimeCode,
-  generateFunctionId,
-  replacer,
-  onDirectiveFnsById,
-  root,
-}: DirectiveFunctionsViteOptions & {
-  code: string
-  id: string
-  root: string
-}) {
-  const url = pathToFileURL(id)
-  url.searchParams.delete('v')
-  id = fileURLToPath(url).replace(/\\/g, '/')
-
-  if (debug) console.info(`${envLabel}: Compiling Directives: `, id)
-
-  const { compiledResult, directiveFnsById, isDirectiveSplitParam } =
-    compileDirectives({
-      directive,
-      directiveLabel,
-      getRuntimeCode,
-      generateFunctionId,
-      replacer,
-      code,
-      root,
-      filename: id,
-    })
-  // when we process a file with a directive split param, we have already encountered the directives in that file
-  // (otherwise we wouldn't have gotten here)
-  if (!isDirectiveSplitParam) {
-    onDirectiveFnsById?.(directiveFnsById)
-  }
+        if (debug) {
+          logDiff(code, compiledResult.code)
+          console.log('Output:\n', compiledResult.code + '\n\n')
+        }
 
-  if (debug) {
-    logDiff(code, compiledResult.code)
-    console.log('Output:\n', compiledResult.code + '\n\n')
+        return compiledResult
+      },
+    },
   }
-
-  return compiledResult
 }

packages/directive-functions-plugin/tests/compiler.test.ts

@@ -19,27 +19,28 @@ const generateFunctionId: CompileDirectivesOpts['generateFunctionId'] = (
 
 const clientConfig: Omit<CompileDirectivesOpts, 'code'> = {
   directive: 'use server',
-  directiveLabel: 'Server function',
   root: './test-files',
   filename: './test-files/test.ts',
   getRuntimeCode: () => 'import { createClientRpc } from "my-rpc-lib-client"',
   generateFunctionId,
   replacer: (opts) => `createClientRpc(${JSON.stringify(opts.functionId)})`,
+  directiveSplitParam: 'tsr-directive-use-server',
+  isDirectiveSplitParam: false,
 }
 
 const ssrConfig: Omit<CompileDirectivesOpts, 'code'> = {
   directive: 'use server',
-  directiveLabel: 'Server function',
   root: './test-files',
   filename: './test-files/test.ts',
   getRuntimeCode: () => 'import { createSsrRpc } from "my-rpc-lib-server"',
   generateFunctionId,
   replacer: (opts) => `createSsrRpc(${JSON.stringify(opts.functionId)})`,
+  directiveSplitParam: 'tsr-directive-use-server',
+  isDirectiveSplitParam: false,
 }
 
 const serverConfig: Omit<CompileDirectivesOpts, 'code'> = {
   directive: 'use server',
-  directiveLabel: 'Server function',
   root: './test-files',
   filename: './test-files/test.ts',
   getRuntimeCode: () => 'import { createServerRpc } from "my-rpc-lib-server"',
@@ -52,6 +53,8 @@ const serverConfig: Omit<CompileDirectivesOpts, 'code'> = {
     // For any other server functions the split function may reference,
     // we use the splitImportFn which is a dynamic import of the split file.
     `createServerRpc(${JSON.stringify(opts.functionId)}, ${opts.fn})`,
+  directiveSplitParam: 'tsr-directive-use-server',
+  isDirectiveSplitParam: true,
 }
 
 describe('server function compilation', () => {

packages/react-start/package.json

@@ -62,6 +62,12 @@
         "default": "./dist/esm/server-rpc.js"
       }
     },
+    "./ssr-rpc": {
+      "import": {
+        "types": "./dist/esm/ssr-rpc.d.ts",
+        "default": "./dist/esm/ssr-rpc.js"
+      }
+    },
     "./plugin/vite": {
       "import": {
         "types": "./dist/esm/plugin/vite.d.ts",

packages/react-start/src/ssr-rpc.ts

@@ -0,0 +1 @@
+export { createSsrRpc } from '@tanstack/start-server-core/createSsrRpc'

packages/react-start/vite.config.ts

@@ -29,6 +29,7 @@ export default mergeConfig(
       './src/client-rpc.ts',
       './src/server.tsx',
       './src/server-rpc.ts',
+      './src/ssr-rpc.ts',
       './src/plugin/vite.ts',
     ],
     externalDeps: [