fix(start): clean up server function handler imports

.changeset/clean-server-fn-imports.md

@@ -0,0 +1,6 @@
+---
+'@tanstack/start-plugin-core': patch
+'@tanstack/router-utils': patch
+---
+
+Fix Start compiler handling for TS-wrapped server/env-only function calls and safely remove unused imported server function handlers from client output.

packages/router-utils/src/compiler-helpers.ts

@@ -1,5 +1,15 @@
 import * as t from '@babel/types'
 
+type CompilerNodePath<TNode extends t.Node = t.Node> = {
+  node: TNode
+  parentPath: CompilerNodePath | null
+  isVariableDeclarator(): boolean
+}
+
+type ReplacePathNode<TPath, TNode extends t.Node> = Omit<TPath, 'node'> & {
+  node: TNode
+}
+
 type IdentifierScopeFrame = {
   kind: 'program' | 'function' | 'block'
   bindings: Set<string>
@@ -23,6 +33,55 @@ export interface ExtractedModuleInfo {
   reExportAllSources: Array<string>
 }
 
+function getTransparentWrapperExpression(node: t.Node): t.Expression | null {
+  if (
+    t.isTSAsExpression(node) ||
+    t.isTSSatisfiesExpression(node) ||
+    t.isTSTypeAssertion(node) ||
+    t.isTSNonNullExpression(node) ||
+    t.isParenthesizedExpression(node)
+  ) {
+    return node.expression
+  }
+
+  return null
+}
+
+export function unwrapExpression(expr: t.Expression): t.Expression {
+  let inner = getTransparentWrapperExpression(expr)
+  while (inner) {
+    expr = inner
+    inner = getTransparentWrapperExpression(expr)
+  }
+
+  return expr
+}
+
+export function getVariableDeclaratorForExpressionPath<
+  TPath extends CompilerNodePath<t.Expression>,
+>(path: TPath): ReplacePathNode<TPath, t.VariableDeclarator> | null {
+  let currentPath: CompilerNodePath = path
+  let parentPath = currentPath.parentPath
+
+  while (
+    parentPath &&
+    getTransparentWrapperExpression(parentPath.node) === currentPath.node
+  ) {
+    currentPath = parentPath
+    parentPath = parentPath.parentPath
+  }
+
+  if (
+    parentPath?.isVariableDeclarator() &&
+    t.isVariableDeclarator(parentPath.node) &&
+    parentPath.node.init === currentPath.node
+  ) {
+    return parentPath as ReplacePathNode<TPath, t.VariableDeclarator>
+  }
+
+  return null
+}
+
 function getModuleExportName(node: t.Identifier | t.StringLiteral) {
   return t.isIdentifier(node) ? node.name : node.value
 }

packages/router-utils/src/index.ts

@@ -23,10 +23,12 @@ export {
   expandSharedDestructuredDeclarators,
   expandTransitively,
   extractModuleInfoFromAst,
+  getVariableDeclaratorForExpressionPath,
   removeBindingsTransitivelyDependingOn,
   removeModuleLevelBindings,
   retainModuleLevelDeclarations,
   stripUnreferencedTopLevelExpressionStatements,
+  unwrapExpression,
   unwrapExportedDeclarations,
 } from './compiler-helpers'
 export type { ExtractedModuleInfo, ModuleInfoBinding } from './compiler-helpers'

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

@@ -5,7 +5,9 @@ import {
   extractModuleInfoFromAst,
   findReferencedIdentifiers,
   generateFromAst,
+  getVariableDeclaratorForExpressionPath,
   parseAst,
+  unwrapExpression,
 } from '@tanstack/router-utils'
 import babel from '@babel/core'
 import { handleCreateServerFn } from './handleCreateServerFn'
@@ -374,6 +376,17 @@ function isTopLevelDirectCallCandidateNode(node: t.CallExpression): boolean {
   return isSimpleDirectCallExpression(node)
 }
 
+function getPotentialCandidateCallExpression(
+  node: t.Expression | null | undefined,
+): t.CallExpression | null {
+  if (!node) {
+    return null
+  }
+
+  const unwrapped = unwrapExpression(node)
+  return t.isCallExpression(unwrapped) ? unwrapped : null
+}
+
 /**
  * Checks if a CallExpression path is a top-level direct-call candidate.
  * Top-level means the call is the init of a VariableDeclarator at program level.
@@ -392,15 +405,24 @@ function isTopLevelDirectCallCandidate(
   }
 
   // Must be top-level: VariableDeclarator -> VariableDeclaration -> Program
-  const parent = path.parent
-  if (!t.isVariableDeclarator(parent) || parent.init !== node) {
+  // or VariableDeclarator -> VariableDeclaration -> ExportNamedDeclaration -> Program.
+  const variableDeclarator = getVariableDeclaratorForExpressionPath(
+    path as babel.NodePath<t.Expression>,
+  )
+  if (!variableDeclarator) {
     return false
   }
-  const grandParent = path.parentPath.parent
-  if (!t.isVariableDeclaration(grandParent)) {
+
+  const variableDeclaration = variableDeclarator.parentPath
+  if (!variableDeclaration.isVariableDeclaration()) {
     return false
   }
-  return t.isProgram(path.parentPath.parentPath?.parent)
+
+  const parent = variableDeclaration.parentPath
+  return (
+    parent.isProgram() ||
+    (parent.isExportNamedDeclaration() && parent.parentPath.isProgram())
+  )
 }
 
 function isDirectCallCandidateForKind(
@@ -1039,11 +1061,11 @@ export class StartCompiler {
 
           if (declarations) {
             for (const decl of declarations) {
-              if (decl.init && t.isCallExpression(decl.init)) {
+              const init = getPotentialCandidateCallExpression(decl.init)
+              if (init) {
                 if (
-                  isMethodChainCandidate(decl.init, fileKinds) ||
-                  (checkDirectCalls &&
-                    isTopLevelDirectCallCandidateNode(decl.init))
+                  isMethodChainCandidate(init, fileKinds) ||
+                  (checkDirectCalls && isTopLevelDirectCallCandidateNode(init))
                 ) {
                   candidateIndices.push(i)
                   break // Only need to mark this statement once
@@ -1774,7 +1796,7 @@ export class StartCompiler {
       getLookupSetup(resolvedKind, this.externalLookupSetup)?.type ===
         'directCall' &&
       binding.init &&
-      t.isCallExpression(binding.init)
+      t.isCallExpression(unwrapExpression(binding.init))
     ) {
       binding.resolvedKind = 'None'
       return 'None'
@@ -1792,14 +1814,7 @@ export class StartCompiler {
       return 'None'
     }
 
-    // Unwrap common TypeScript/parenthesized wrappers first for efficiency
-    while (
-      t.isTSAsExpression(expr) ||
-      t.isTSNonNullExpression(expr) ||
-      t.isParenthesizedExpression(expr)
-    ) {
-      expr = expr.expression
-    }
+    expr = unwrapExpression(expr)
 
     let result: Kind = 'None'
 

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

@@ -1,6 +1,7 @@
 import * as t from '@babel/types'
 import babel from '@babel/core'
 import { hasKeys } from '@tanstack/router-core'
+import { getVariableDeclaratorForExpressionPath } from '@tanstack/router-utils'
 import path from 'pathe'
 import { cleanId, codeFrameError, stripMethodCall } from './utils'
 import type { CompilationContext, RewriteCandidate, ServerFn } from './types'
@@ -222,13 +223,17 @@ export function handleCreateServerFn(
     const { path: candidatePath, methodChain } = candidate
     const { inputValidator, handler } = methodChain
 
+    const candidateVariableDeclarator = getVariableDeclaratorForExpressionPath(
+      candidatePath as babel.NodePath<t.Expression>,
+    )
+
     // Check if the call is assigned to a variable
-    if (!candidatePath.parentPath.isVariableDeclarator()) {
+    if (!candidateVariableDeclarator) {
       throw new Error('createServerFn must be assigned to a variable!')
     }
 
     // Get the identifier name of the variable
-    const variableDeclarator = candidatePath.parentPath.node
+    const variableDeclarator = candidateVariableDeclarator.node
     if (!t.isIdentifier(variableDeclarator.id)) {
       throw codeFrameError(
         context.code,
@@ -361,8 +366,8 @@ export function handleCreateServerFn(
       ])
 
       // Find the variable declaration statement containing our createServerFn
-      const variableDeclaration = candidatePath.parentPath.parentPath
-      if (!variableDeclaration.isVariableDeclaration()) {
+      const variableDeclaration = candidateVariableDeclarator.parentPath
+      if (!variableDeclaration?.isVariableDeclaration()) {
         throw new Error(
           'Expected createServerFn to be in a VariableDeclaration',
         )
@@ -398,15 +403,6 @@ export function handleCreateServerFn(
       // or
       // const myFn = createServerFn().handler(createSsrRpc("id"))
 
-      // If the handler function is an identifier, we need to remove the bound function
-      // from the file since it won't be needed
-      if (t.isIdentifier(handlerFn)) {
-        const binding = handlerFnPath.scope.getBinding(handlerFn.name)
-        if (binding) {
-          binding.path.remove()
-        }
-      }
-
       // Generate the RPC stub using pre-compiled templates
       // Note: Caller files only pass functionId, not the full serverFnMeta
       const rpcStub = generateCallerRpcStub(functionId, envConfig)

packages/start-plugin-core/tests/createServerFn/createServerFn.test.ts

@@ -158,6 +158,91 @@ describe('createServerFn compiles correctly', async () => {
     `)
   })
 
+  test('should remove imports used only as caller handler identifiers', async () => {
+    const code = `
+      import { createServerFn } from '@tanstack/react-start'
+      import { getUsers } from './server'
+
+      export const getUsersFn = createServerFn().handler(getUsers)
+    `
+
+    const compiledResult = await compile({
+      code,
+      env: 'client',
+      isProviderFile: false,
+      mode: 'build',
+    })
+
+    expect(compiledResult!.code).toMatchInlineSnapshot(`
+      "import { createClientRpc } from '@tanstack/react-start/client-rpc';
+      import { createServerFn } from '@tanstack/react-start';
+      export const getUsersFn = createServerFn().handler(createClientRpc(\"a78c63d4bb3c0b10a8b70902c73611fbf0b9229e0807b065c03c939f9c0ce100\"));"
+    `)
+  })
+
+  test('should not remove handler identifier bindings that are still referenced', async () => {
+    const code = `
+      import { createServerFn } from '@tanstack/react-start'
+      import { getUsers } from './server'
+
+      export const getUsersFn = createServerFn().handler(getUsers)
+      console.log(getUsers)
+    `
+
+    const compiledResult = await compile({
+      code,
+      env: 'client',
+      isProviderFile: false,
+      mode: 'build',
+    })
+
+    expect(compiledResult!.code).toMatchInlineSnapshot(`
+      "import { createClientRpc } from '@tanstack/react-start/client-rpc';
+      import { createServerFn } from '@tanstack/react-start';
+      import { getUsers } from './server';
+      export const getUsersFn = createServerFn().handler(createClientRpc(\"a78c63d4bb3c0b10a8b70902c73611fbf0b9229e0807b065c03c939f9c0ce100\"));
+      console.log(getUsers);"
+    `)
+  })
+
+  test('should compile server functions wrapped in TypeScript as expressions', async () => {
+    const code = `
+      import { createServerFn } from '@tanstack/react-start'
+
+      export const getUsersFn = createServerFn().handler(() => 'server') as any
+    `
+
+    const compiledResult = await compile({
+      code,
+      env: 'client',
+      isProviderFile: false,
+      mode: 'build',
+    })
+
+    expect(compiledResult).not.toBeNull()
+    expect(compiledResult!.code).toContain('createClientRpc')
+    expect(compiledResult!.code).not.toContain('server')
+  })
+
+  test('should compile server functions wrapped in TypeScript satisfies expressions', async () => {
+    const code = `
+      import { createServerFn } from '@tanstack/react-start'
+
+      export const getUsersFn = createServerFn().handler(() => 'server') satisfies unknown
+    `
+
+    const compiledResult = await compile({
+      code,
+      env: 'client',
+      isProviderFile: false,
+      mode: 'build',
+    })
+
+    expect(compiledResult).not.toBeNull()
+    expect(compiledResult!.code).toContain('createClientRpc')
+    expect(compiledResult!.code).not.toContain('server')
+  })
+
   test('should use dce by default', async () => {
     const code = `
       import { createServerFn } from '@tanstack/react-start'

packages/start-plugin-core/tests/envOnly/snapshots/client/envOnlyExportedAlias.tsx

@@ -0,0 +1,6 @@
+export const fn = () => {
+  throw new Error("createServerOnlyFn() functions can only be called on the server!");
+};
+export const wrappedFn = (() => {
+  throw new Error("createServerOnlyFn() functions can only be called on the server!");
+}) as () => string;

packages/start-plugin-core/tests/envOnly/snapshots/server/envOnlyExportedAlias.tsx

@@ -0,0 +1,2 @@
+export const fn = () => 'server-only-value';
+export const wrappedFn = (() => 'wrapped-server-only-value') as () => string;

packages/start-plugin-core/tests/envOnly/test-files/envOnlyExportedAlias.tsx

@@ -0,0 +1,4 @@
+import { createServerOnlyFn as so } from '@tanstack/react-start'
+
+export const fn = so(() => 'server-only-value')
+export const wrappedFn = so(() => 'wrapped-server-only-value') as () => string