Repository: TanStack/router (ts)

Issue:
Missing allowed types in start's defineConfig "react.babel.plugins"
### Which project does this relate to?

Start

### Describe the bug

The type definition and schema validation of some fields in `@tanstack/start/config`'s exported `defineConfig` is seemingly unnecessarily narrowed down.

```TS
defineConfig({
  react: {
    babel: {
      plugins: (string | [string, any])[] | undefined
    }
  }
})
```

The definition exported in `@vite/plugin-react` is `PluginItem[] | null | undefined`.
`PluginItem` is defined by the following:
```TS
export type PluginItem =
    | ConfigItem
    | PluginObj<any>
    | PluginTarget
    | [PluginTarget, PluginOptions]
    | [PluginTarget, PluginOptions, string | undefined];
```

This does match the zod schema used by Start, exported by `defineConfig`, narrowing down the allowed type definition to a subset of `PluginItem`, but restricts the use of a custom object imported inside `app.config.ts`, as it gets rejected by both the type definition and by Zod at runtime.

This is the only config field I needed, in order to make use of [Jotai's fast refresh plugin with Vite](https://jotai.org/docs/guides/vite). I haven't checked other config fields, there might be other ones that are also narrowed down in such a way that it restricts some use of a custom Vite config.

### Your Example Website or App

https://github.com/busybox11/ts-start-repro-babel/

### Steps to Reproduce the Bug or Issue

1. Install the modules
2. Open `app.config.ts` in your editor => **You should have type errors on the `plugins` array** (line 15, field annotated)
3. Run the dev server => **You should see Zod runtime errors in the console** (also annotated)

### Expected behavior

`defineConfig` should allow the use of custom objects as babel plugins.

### Screenshots or Videos

![Image](https://github.com/user-attachments/assets/9ec4071a-86a0-44ea-ac36-5d5fb48c8fe3)


### Platform

- OS: Linux
- Browser: Zen (Firefox), Chrome
- Version: 1.0.1-t.19 (Firefox 132.0.2), 131.0.6778.108


### Additional context

_No response_

Pull request description:
fix(start): pass through vite react plugin config options
fixes #2952

Patch:
diff --git a/packages/start-vite-plugin/src/compilers.ts b/packages/start-vite-plugin/src/compilers.ts
index a2117d649..13dde6e1b 100644
--- a/packages/start-vite-plugin/src/compilers.ts
+++ b/packages/start-vite-plugin/src/compilers.ts
@@ -32,6 +32,12 @@ export function compileEliminateDeadCode(opts: ParseAstOptions) {
 
 const debug = process.env.TSR_VITE_DEBUG === 'true'
 
+// build these once and reuse them
+const handleServerOnlyCallExpression =
+  buildEnvOnlyCallExpressionHandler('server')
+const handleClientOnlyCallExpression =
+  buildEnvOnlyCallExpressionHandler('client')
+
 type IdentifierConfig = {
   name: string
   type: 'ImportSpecifier' | 'ImportNamespaceSpecifier'
@@ -59,6 +65,7 @@ export function compileStartOutput(opts: ParseAstOptions) {
           createServerFn: IdentifierConfig
           createMiddleware: IdentifierConfig
           serverOnly: IdentifierConfig
+          clientOnly: IdentifierConfig
           createIsomorphicFn: IdentifierConfig
         } = {
           createServerFn: {
@@ -82,6 +89,13 @@ export function compileStartOutput(opts: ParseAstOptions) {
             handleCallExpression: handleServerOnlyCallExpression,
             paths: [],
           },
+          clientOnly: {
+            name: 'clientOnly',
+            type: 'ImportSpecifier',
+            namespaceId: '',
+            handleCallExpression: handleClientOnlyCallExpression,
+            paths: [],
+          },
           createIsomorphicFn: {
             name: 'createIsomorphicFn',
             type: 'ImportSpecifier',
@@ -548,36 +562,49 @@ function handleCreateMiddlewareCallExpression(
   }
 }
 
-function handleServerOnlyCallExpression(
-  path: babel.NodePath<t.CallExpression>,
-  opts: ParseAstOptions,
-) {
-  if (debug)
-    console.info('Handling serverOnly call expression:', path.toString())
+function buildEnvOnlyCallExpressionHandler(env: 'client' | 'server') {
+  return function envOnlyCallExpressionHandler(
+    path: babel.NodePath<t.CallExpression>,
+    opts: ParseAstOptions,
+  ) {
+    if (debug)
+      console.info(`Handling ${env}Only call expression:`, path.toString())
 
-  if (opts.env === 'server') {
-    // Do nothing on the server.
-    return
-  }
+    if (!path.parentPath.isVariableDeclarator()) {
+      throw new Error(`${env}Only() functions must be assigned to a variable!`)
+    }
 
-  // If we're on the client, replace the call expression with a function
-  // that has a single always-triggering invariant.
+    if (opts.env === env) {
+      // extract the inner function from the call expression
+      const innerInputExpression = path.node.arguments[0]
 
-  path.replaceWith(
-    t.arrowFunctionExpression(
-      [],
-      t.blockStatement([
-        t.expressionStatement(
-          t.callExpression(t.identifier('invariant'), [
-            t.booleanLiteral(false),
-            t.stringLiteral(
-              'serverOnly() functions can only be called on the server!',
-            ),
-          ]),
-        ),
-      ]),
-    ),
-  )
+      if (!t.isExpression(innerInputExpression)) {
+        throw new Error(
+          `${env}Only() functions must be called with a function!`,
+        )
+      }
+
+      path.replaceWith(innerInputExpression)
+      return
+    }
+
+    // If we're on the wrong environment, replace the call expression
+    // with a function that always throws an error.
+    path.replaceWith(
+      t.arrowFunctionExpression(
+        [],
+        t.blockStatement([
+          t.throwStatement(
+            t.newExpression(t.identifier('Error'), [
+              t.stringLiteral(
+                `${env}Only() functions can only be called on the ${env}!`,
+              ),
+            ]),
+          ),
+        ]),
+      ),
+    )
+  }
 }
 
 function handleCreateIsomorphicFnCallExpression(
diff --git a/packages/start-vite-plugin/src/index.ts b/packages/start-vite-plugin/src/index.ts
index 55e0b2231..20b377d24 100644
--- a/packages/start-vite-plugin/src/index.ts
+++ b/packages/start-vite-plugin/src/index.ts
@@ -10,6 +10,18 @@ export type TanStackStartViteOptions = {
   env: 'server' | 'client'
 }
 
+const transformFuncs = [
+  'createServerFn',
+  'createMiddleware',
+  'serverOnly',
+  'clientOnly',
+  'createIsomorphicFn',
+]
+const tokenRegex = new RegExp(transformFuncs.join('|'))
+const eitherFuncRegex = new RegExp(
+  `(function ${transformFuncs.join('|function ')})`,
+)
+
 export function TanStackStartViteServerFn(
   opts: TanStackStartViteOptions,
 ): Plugin {
@@ -26,14 +38,8 @@ export function TanStackStartViteServerFn(
       url.searchParams.delete('v')
       id = fileURLToPath(url).replace(/\\/g, '/')
 
-      const includesToken =
-        /createServerFn|createMiddleware|serverOnly|createIsomorphicFn/.test(
-          code,
-        )
-      const includesEitherFunc =
-        /(function createServerFn|function createMiddleware|function serverOnly|function createIsomorphicFn)/.test(
-          code,
-        )
+      const includesToken = tokenRegex.test(code)
+      const includesEitherFunc = eitherFuncRegex.test(code)
 
       if (
         !includesToken ||
@@ -92,12 +98,7 @@ export function TanStackStartViteDeadCodeElimination(
       url.searchParams.delete('v')
       id = fileURLToPath(url).replace(/\\/g, '/')
 
-      if (
-        code.includes('createServerFn') ||
-        code.includes('createMiddleware') ||
-        code.includes('serverOnly') ||
-        code.includes('createIsomorphicFn')
-      ) {
+      if (transformFuncs.some((fn) => code.includes(fn))) {
         const compiled = compileEliminateDeadCode({
           code,
           root: ROOT,
diff --git a/packages/start/src/client/envOnly.ts b/packages/start/src/client/envOnly.ts
new file mode 100644
index 000000000..55a6f3333
--- /dev/null
+++ b/packages/start/src/client/envOnly.ts
@@ -0,0 +1,17 @@
+// A function that will only be available in the server build
+// If called on the client, it will throw an error
+
+export function serverOnly<TFn extends (...args: Array<any>) => any>(
+  fn: TFn,
+): TFn {
+  return fn
+}
+
+// A function that will only be available in the client build
+// If called on the server, it will throw an error
+
+export function clientOnly<TFn extends (...args: Array<any>) => any>(
+  fn: TFn,
+): TFn {
+  return fn
+}
diff --git a/packages/start/src/client/index.tsx b/packages/start/src/client/index.tsx
index ea73a9641..08f8b6f2d 100644
--- a/packages/start/src/client/index.tsx
+++ b/packages/start/src/client/index.tsx
@@ -41,7 +41,7 @@ export {
   registerGlobalMiddleware,
   globalMiddleware,
 } from './registerGlobalMiddleware'
-export { serverOnly } from './serverOnly'
+export { serverOnly, clientOnly } from './envOnly'
 export { DehydrateRouter } from './DehydrateRouter'
 export { json } from './json'
 export { Meta } from './Meta'
diff --git a/packages/start/src/client/serverOnly.ts b/packages/start/src/client/serverOnly.ts
deleted file mode 100644
index 9bd1efb8d..000000000
--- a/packages/start/src/client/serverOnly.ts
+++ /dev/null
@@ -1,6 +0,0 @@
-// A function that will have its contents removed on the client build and only be available on the server build
-// it can accept any value, a function that will return a value, or an async function that will return a value
-
-export function serverOnly<T>(value: () => T): () => T | undefined {
-  return value
-}
diff --git a/packages/start/src/config/index.ts b/packages/start/src/config/index.ts
index 7efef1fb3..f371997c8 100644
--- a/packages/start/src/config/index.ts
+++ b/packages/start/src/config/index.ts
@@ -2,7 +2,7 @@ import path from 'node:path'
 import { existsSync, readFileSync } from 'node:fs'
 import { readFile } from 'node:fs/promises'
 import { fileURLToPath } from 'node:url'
-import reactRefresh from '@vitejs/plugin-react'
+import viteReact from '@vitejs/plugin-react'
 import { resolve } from 'import-meta-resolve'
 import { TanStackRouterVite } from '@tanstack/router-plugin/vite'
 import {
@@ -173,11 +173,7 @@ export function defineConfig(
             serverFunctions.client({
               runtime: '@tanstack/start/client-runtime',
             }),
-            reactRefresh({
-              babel: opts.react?.babel,
-              exclude: opts.react?.exclude,
-              include: opts.react?.include,
-            }),
+            viteReact(opts.react),
             // TODO: RSCS - enable this
             // serverComponents.client(),
           ]
diff --git a/packages/start/src/config/schema.ts b/packages/start/src/config/schema.ts
index af62c095b..dd38a5789 100644
--- a/packages/start/src/config/schema.ts
+++ b/packages/start/src/config/schema.ts
@@ -3,7 +3,7 @@ import { z } from 'zod'
 import type { PluginOption } from 'vite'
 import type { AppOptions as VinxiAppOptions } from 'vinxi'
 import type { NitroOptions } from 'nitropack'
-
+import type { Options as ViteReactOptions } from '@vitejs/plugin-react'
 import type { CustomizableConfig } from 'vinxi/dist/types/lib/vite-dev'
 
 type StartUserViteConfig = CustomizableConfig | (() => CustomizableConfig)
@@ -138,17 +138,7 @@ export const serverSchema = z
 
 const viteSchema = z.custom<StartUserViteConfig>()
 
-const babelSchema = z.object({
-  plugins: z
-    .array(z.union([z.tuple([z.string(), z.any()]), z.string()]))
-    .optional(),
-})
-
-const reactSchema = z.object({
-  babel: babelSchema.optional(),
-  exclude: z.array(z.instanceof(RegExp)).optional(),
-  include: z.array(z.instanceof(RegExp)).optional(),
-})
+const viteReactSchema = z.custom<ViteReactOptions>()
 
 const routersSchema = z.object({
   ssr: z
@@ -192,7 +182,7 @@ const tsrConfig = configSchema.partial().extend({
 })
 
 export const inlineConfigSchema = z.object({
-  react: reactSchema.optional(),
+  react: viteReactSchema.optional(),
   vite: viteSchema.optional(),
   tsr: tsrConfig.optional(),
   routers: routersSchema.optional(),


Test patch:
diff --git a/e2e/start/basic/app/routeTree.gen.ts b/e2e/start/basic/app/routeTree.gen.ts
index 2ab05c9e3..5f29d2300 100644
--- a/e2e/start/basic/app/routeTree.gen.ts
+++ b/e2e/start/basic/app/routeTree.gen.ts
@@ -18,6 +18,7 @@ import { Route as SearchParamsImport } from './routes/search-params'
 import { Route as RedirectImport } from './routes/redirect'
 import { Route as PostsImport } from './routes/posts'
 import { Route as IsomorphicFnsImport } from './routes/isomorphic-fns'
+import { Route as EnvOnlyImport } from './routes/env-only'
 import { Route as DeferredImport } from './routes/deferred'
 import { Route as LayoutImport } from './routes/_layout'
 import { Route as IndexImport } from './routes/index'
@@ -74,6 +75,12 @@ const IsomorphicFnsRoute = IsomorphicFnsImport.update({
   getParentRoute: () => rootRoute,
 } as any)
 
+const EnvOnlyRoute = EnvOnlyImport.update({
+  id: '/env-only',
+  path: '/env-only',
+  getParentRoute: () => rootRoute,
+} as any)
+
 const DeferredRoute = DeferredImport.update({
   id: '/deferred',
   path: '/deferred',
@@ -163,6 +170,13 @@ declare module '@tanstack/react-router' {
       preLoaderRoute: typeof DeferredImport
       parentRoute: typeof rootRoute
     }
+    '/env-only': {
+      id: '/env-only'
+      path: '/env-only'
+      fullPath: '/env-only'
+      preLoaderRoute: typeof EnvOnlyImport
+      parentRoute: typeof rootRoute
+    }
     '/isomorphic-fns': {
       id: '/isomorphic-fns'
       path: '/isomorphic-fns'
@@ -326,6 +340,7 @@ export interface FileRoutesByFullPath {
   '/': typeof IndexRoute
   '': typeof LayoutLayout2RouteWithChildren
   '/deferred': typeof DeferredRoute
+  '/env-only': typeof EnvOnlyRoute
   '/isomorphic-fns': typeof IsomorphicFnsRoute
   '/posts': typeof PostsRouteWithChildren
   '/redirect': typeof RedirectRoute
@@ -346,6 +361,7 @@ export interface FileRoutesByTo {
   '/': typeof IndexRoute
   '': typeof LayoutLayout2RouteWithChildren
   '/deferred': typeof DeferredRoute
+  '/env-only': typeof EnvOnlyRoute
   '/isomorphic-fns': typeof IsomorphicFnsRoute
   '/redirect': typeof RedirectRoute
   '/search-params': typeof SearchParamsRoute
@@ -365,6 +381,7 @@ export interface FileRoutesById {
   '/': typeof IndexRoute
   '/_layout': typeof LayoutRouteWithChildren
   '/deferred': typeof DeferredRoute
+  '/env-only': typeof EnvOnlyRoute
   '/isomorphic-fns': typeof IsomorphicFnsRoute
   '/posts': typeof PostsRouteWithChildren
   '/redirect': typeof RedirectRoute
@@ -388,6 +405,7 @@ export interface FileRouteTypes {
     | '/'
     | ''
     | '/deferred'
+    | '/env-only'
     | '/isomorphic-fns'
     | '/posts'
     | '/redirect'
@@ -407,6 +425,7 @@ export interface FileRouteTypes {
     | '/'
     | ''
     | '/deferred'
+    | '/env-only'
     | '/isomorphic-fns'
     | '/redirect'
     | '/search-params'
@@ -424,6 +443,7 @@ export interface FileRouteTypes {
     | '/'
     | '/_layout'
     | '/deferred'
+    | '/env-only'
     | '/isomorphic-fns'
     | '/posts'
     | '/redirect'
@@ -446,6 +466,7 @@ export interface RootRouteChildren {
   IndexRoute: typeof IndexRoute
   LayoutRoute: typeof LayoutRouteWithChildren
   DeferredRoute: typeof DeferredRoute
+  EnvOnlyRoute: typeof EnvOnlyRoute
   IsomorphicFnsRoute: typeof IsomorphicFnsRoute
   PostsRoute: typeof PostsRouteWithChildren
   RedirectRoute: typeof RedirectRoute
@@ -460,6 +481,7 @@ const rootRouteChildren: RootRouteChildren = {
   IndexRoute: IndexRoute,
   LayoutRoute: LayoutRouteWithChildren,
   DeferredRoute: DeferredRoute,
+  EnvOnlyRoute: EnvOnlyRoute,
   IsomorphicFnsRoute: IsomorphicFnsRoute,
   PostsRoute: PostsRouteWithChildren,
   RedirectRoute: RedirectRoute,
@@ -483,6 +505,7 @@ export const routeTree = rootRoute
         "/",
         "/_layout",
         "/deferred",
+        "/env-only",
         "/isomorphic-fns",
         "/posts",
         "/redirect",
@@ -505,6 +528,9 @@ export const routeTree = rootRoute
     "/deferred": {
       "filePath": "deferred.tsx"
     },
+    "/env-only": {
+      "filePath": "env-only.tsx"
+    },
     "/isomorphic-fns": {
       "filePath": "isomorphic-fns.tsx"
     },
diff --git a/e2e/start/basic/app/routes/env-only.tsx b/e2e/start/basic/app/routes/env-only.tsx
new file mode 100644
index 000000000..50e4da87e
--- /dev/null
+++ b/e2e/start/basic/app/routes/env-only.tsx
@@ -0,0 +1,76 @@
+import { createFileRoute } from '@tanstack/react-router'
+import { clientOnly, createServerFn, serverOnly } from '@tanstack/start'
+import { useState } from 'react'
+
+const serverEcho = serverOnly((input: string) => 'server got: ' + input)
+const clientEcho = clientOnly((input: string) => 'client got: ' + input)
+
+const testOnServer = createServerFn().handler(() => {
+  const serverOnServer = serverEcho('hello')
+  let clientOnServer: string
+  try {
+    clientOnServer = clientEcho('hello')
+  } catch (e) {
+    clientOnServer =
+      'clientEcho threw an error: ' +
+      (e instanceof Error ? e.message : String(e))
+  }
+  return { serverOnServer, clientOnServer }
+})
+
+export const Route = createFileRoute('/env-only')({
+  component: RouteComponent,
+})
+
+function RouteComponent() {
+  const [results, setResults] = useState<Partial<Record<string, string>>>()
+
+  async function handleClick() {
+    const { serverOnServer, clientOnServer } = await testOnServer()
+    const clientOnClient = clientEcho('hello')
+    let serverOnClient: string
+    try {
+      serverOnClient = serverEcho('hello')
+    } catch (e) {
+      serverOnClient =
+        'serverEcho threw an error: ' +
+        (e instanceof Error ? e.message : String(e))
+    }
+    setResults({
+      serverOnServer,
+      clientOnServer,
+      clientOnClient,
+      serverOnClient,
+    })
+  }
+
+  const { serverOnServer, clientOnServer, clientOnClient, serverOnClient } =
+    results || {}
+
+  return (
+    <div>
+      <button onClick={handleClick} data-testid="test-env-only-results-btn">
+        Run
+      </button>
+      {!!results && (
+        <div>
+          <h1>
+            <code>serverEcho</code>
+          </h1>
+          When we called the function on the server:
+          <pre data-testid="server-on-server">{serverOnServer}</pre>
+          When we called the function on the client:
+          <pre data-testid="server-on-client">{serverOnClient}</pre>
+          <br />
+          <h1>
+            <code>clientEcho</code>
+          </h1>
+          When we called the function on the server:
+          <pre data-testid="client-on-server">{clientOnServer}</pre>
+          When we called the function on the client:
+          <pre data-testid="client-on-client">{clientOnClient}</pre>
+        </div>
+      )}
+    </div>
+  )
+}
diff --git a/e2e/start/basic/tests/base.spec.ts b/e2e/start/basic/tests/base.spec.ts
index d1e7ca996..2f2414bcc 100644
--- a/e2e/start/basic/tests/base.spec.ts
+++ b/e2e/start/basic/tests/base.spec.ts
@@ -167,3 +167,28 @@ test('isomorphic functions can have different implementations on client and serv
     'client received hello',
   )
 })
+
+test('env-only functions can only be called on the server or client respectively', async ({
+  page,
+}) => {
+  await page.goto('/env-only')
+
+  await page.waitForLoadState('networkidle')
+
+  await page.getByTestId('test-env-only-results-btn').click()
+  await page.waitForLoadState('networkidle')
+
+  await expect(page.getByTestId('server-on-server')).toContainText(
+    'server got: hello',
+  )
+  await expect(page.getByTestId('server-on-client')).toContainText(
+    'serverEcho threw an error: serverOnly() functions can only be called on the server!',
+  )
+
+  await expect(page.getByTestId('client-on-server')).toContainText(
+    'clientEcho threw an error: clientOnly() functions can only be called on the client!',
+  )
+  await expect(page.getByTestId('client-on-client')).toContainText(
+    'client got: hello',
+  )
+})
diff --git a/packages/start-vite-plugin/tests/envOnly/envOnly.test.ts b/packages/start-vite-plugin/tests/envOnly/envOnly.test.ts
new file mode 100644
index 000000000..a7b4c4a0a
--- /dev/null
+++ b/packages/start-vite-plugin/tests/envOnly/envOnly.test.ts
@@ -0,0 +1,80 @@
+import { readFile, readdir } from 'node:fs/promises'
+import path from 'node:path'
+import { describe, expect, test } from 'vitest'
+
+import { compileStartOutput } from '../../src/compilers'
+
+async function getFilenames() {
+  return await readdir(path.resolve(import.meta.dirname, './test-files'))
+}
+
+describe('envOnly functions compile 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,
+          root: './test-files',
+          filename,
+        })
+
+        await expect(compiledResult.code).toMatchFileSnapshot(
+          `./snapshots/${env}/${filename}`,
+        )
+      },
+    )
+  })
+  test('should error if implementation not provided', () => {
+    expect(() => {
+      compileStartOutput({
+        env: 'client',
+        code: `
+        import { clientOnly } from '@tanstack/start'
+        const fn = clientOnly()`,
+        root: './test-files',
+        filename: 'no-fn.ts',
+      })
+    }).toThrowError()
+    expect(() => {
+      compileStartOutput({
+        env: 'server',
+        code: `
+        import { serverOnly } from '@tanstack/start'
+        const fn = serverOnly()`,
+        root: './test-files',
+        filename: 'no-fn.ts',
+      })
+    }).toThrowError()
+  })
+  test('should be assigned to a variable', () => {
+    expect(() => {
+      compileStartOutput({
+        env: 'server',
+        code: `
+        import { serverOnly } from '@tanstack/start'
+        serverOnly()`,
+        root: './test-files',
+        filename: 'no-fn.ts',
+      })
+    }).toThrowError()
+    expect(() => {
+      compileStartOutput({
+        env: 'client',
+        code: `
+        import { clientOnly } from '@tanstack/start'
+        clientOnly()`,
+        root: './test-files',
+        filename: 'no-fn.ts',
+      })
+    }).toThrowError()
+  })
+})
diff --git a/packages/start-vite-plugin/tests/envOnly/snapshots/client/envOnlyDestructured.tsx b/packages/start-vite-plugin/tests/envOnly/snapshots/client/envOnlyDestructured.tsx
new file mode 100644
index 000000000..3dc714f69
--- /dev/null
+++ b/packages/start-vite-plugin/tests/envOnly/snapshots/client/envOnlyDestructured.tsx
@@ -0,0 +1,5 @@
+import { serverOnly, clientOnly } from '@tanstack/start';
+const serverFunc = () => {
+  throw new Error("serverOnly() functions can only be called on the server!");
+};
+const clientFunc = () => 'client';
\ No newline at end of file
diff --git a/packages/start-vite-plugin/tests/envOnly/snapshots/client/envOnlyDestructuredRename.tsx b/packages/start-vite-plugin/tests/envOnly/snapshots/client/envOnlyDestructuredRename.tsx
new file mode 100644
index 000000000..3dd661c86
--- /dev/null
+++ b/packages/start-vite-plugin/tests/envOnly/snapshots/client/envOnlyDestructuredRename.tsx
@@ -0,0 +1,5 @@
+import { serverOnly as serverFn, clientOnly as clientFn } from '@tanstack/start';
+const serverFunc = () => {
+  throw new Error("serverOnly() functions can only be called on the server!");
+};
+const clientFunc = () => 'client';
\ No newline at end of file
diff --git a/packages/start-vite-plugin/tests/envOnly/snapshots/client/envOnlyStarImport.tsx b/packages/start-vite-plugin/tests/envOnly/snapshots/client/envOnlyStarImport.tsx
new file mode 100644
index 000000000..61f067795
--- /dev/null
+++ b/packages/start-vite-plugin/tests/envOnly/snapshots/client/envOnlyStarImport.tsx
@@ -0,0 +1,5 @@
+import * as TanstackStart from '@tanstack/start';
+const serverFunc = () => {
+  throw new Error("serverOnly() functions can only be called on the server!");
+};
+const clientFunc = () => 'client';
\ No newline at end of file
diff --git a/packages/start-vite-plugin/tests/envOnly/snapshots/server/envOnlyDestructured.tsx b/packages/start-vite-plugin/tests/envOnly/snapshots/server/envOnlyDestructured.tsx
new file mode 100644
index 000000000..24f736eac
--- /dev/null
+++ b/packages/start-vite-plugin/tests/envOnly/snapshots/server/envOnlyDestructured.tsx
@@ -0,0 +1,5 @@
+import { serverOnly, clientOnly } from '@tanstack/start';
+const serverFunc = () => 'server';
+const clientFunc = () => {
+  throw new Error("clientOnly() functions can only be called on the client!");
+};
\ No newline at end of file
diff --git a/packages/start-vite-plugin/tests/envOnly/snapshots/server/envOnlyDestructuredRename.tsx b/packages/start-vite-plugin/tests/envOnly/snapshots/server/envOnlyDestructuredRename.tsx
new file mode 100644
index 000000000..790e1328d
--- /dev/null
+++ b/packages/start-vite-plugin/tests/envOnly/snapshots/server/envOnlyDestructuredRename.tsx
@@ -0,0 +1,5 @@
+import { serverOnly as serverFn, clientOnly as clientFn } from '@tanstack/start';
+const serverFunc = () => 'server';
+const clientFunc = () => {
+  throw new Error("clientOnly() functions can only be called on the client!");
+};
\ No newline at end of file
diff --git a/packages/start-vite-plugin/tests/envOnly/snapshots/server/envOnlyStarImport.tsx b/packages/start-vite-plugin/tests/envOnly/snapshots/server/envOnlyStarImport.tsx
new file mode 100644
index 000000000..51e320d4a
--- /dev/null
+++ b/packages/start-vite-plugin/tests/envOnly/snapshots/server/envOnlyStarImport.tsx
@@ -0,0 +1,5 @@
+import * as TanstackStart from '@tanstack/start';
+const serverFunc = () => 'server';
+const clientFunc = () => {
+  throw new Error("clientOnly() functions can only be called on the client!");
+};
\ No newline at end of file
diff --git a/packages/start-vite-plugin/tests/envOnly/test-files/envOnlyDestructured.tsx b/packages/start-vite-plugin/tests/envOnly/test-files/envOnlyDestructured.tsx
new file mode 100644
index 000000000..1389624d3
--- /dev/null
+++ b/packages/start-vite-plugin/tests/envOnly/test-files/envOnlyDestructured.tsx
@@ -0,0 +1,5 @@
+import { serverOnly, clientOnly } from '@tanstack/start'
+
+const serverFunc = serverOnly(() => 'server')
+
+const clientFunc = clientOnly(() => 'client')
diff --git a/packages/start-vite-plugin/tests/envOnly/test-files/envOnlyDestructuredRename.tsx b/packages/start-vite-plugin/tests/envOnly/test-files/envOnlyDestructuredRename.tsx
new file mode 100644
index 000000000..ac34f0613
--- /dev/null
+++ b/packages/start-vite-plugin/tests/envOnly/test-files/envOnlyDestructuredRename.tsx
@@ -0,0 +1,5 @@
+import { serverOnly as serverFn, clientOnly as clientFn } from '@tanstack/start'
+
+const serverFunc = serverFn(() => 'server')
+
+const clientFunc = clientFn(() => 'client')
diff --git a/packages/start-vite-plugin/tests/envOnly/test-files/envOnlyStarImport.tsx b/packages/start-vite-plugin/tests/envOnly/test-files/envOnlyStarImport.tsx
new file mode 100644
index 000000000..cd262b54a
--- /dev/null
+++ b/packages/start-vite-plugin/tests/envOnly/test-files/envOnlyStarImport.tsx
@@ -0,0 +1,5 @@
+import * as TanstackStart from '@tanstack/start'
+
+const serverFunc = TanstackStart.serverOnly(() => 'server')
+
+const clientFunc = TanstackStart.clientOnly(() => 'client')
diff --git a/packages/start/src/client/tests/envOnly.test-d.ts b/packages/start/src/client/tests/envOnly.test-d.ts
new file mode 100644
index 000000000..ce8ac72c7
--- /dev/null
+++ b/packages/start/src/client/tests/envOnly.test-d.ts
@@ -0,0 +1,34 @@
+import { expectTypeOf, test } from 'vitest'
+import { clientOnly, serverOnly } from '../envOnly'
+
+const inputFn = () => 'output'
+
+const genericInputFn = <T>(input: T) => input
+
+function overloadedFn(input: string): string
+function overloadedFn(input: number): number
+function overloadedFn(input: any) {
+  return input
+}
+
+test("clientOnly returns the function it's given", () => {
+  const outputFn = clientOnly(inputFn)
+  expectTypeOf(outputFn).toEqualTypeOf<typeof inputFn>()
+
+  const genericOutputFn = clientOnly(genericInputFn)
+  expectTypeOf(genericOutputFn).toEqualTypeOf<typeof genericInputFn>()
+
+  const overloadedOutputFn = clientOnly(overloadedFn)
+  expectTypeOf(overloadedOutputFn).toEqualTypeOf<typeof overloadedFn>()
+})
+
+test("serverOnly returns the function it's given", () => {
+  const outputFn = serverOnly(inputFn)
+  expectTypeOf(outputFn).toEqualTypeOf<typeof inputFn>()
+
+  const genericOutputFn = serverOnly(genericInputFn)
+  expectTypeOf(genericOutputFn).toEqualTypeOf<typeof genericInputFn>()
+
+  const overloadedOutputFn = serverOnly(overloadedFn)
+  expectTypeOf(overloadedOutputFn).toEqualTypeOf<typeof overloadedFn>()
+})