Repository: TanStack/router (ts)

Issue:
serverOnly return type should consider parameters of the passed function
### Which project does this relate to?

Start

### Describe the bug

serverOnly does not return the correct type.

As of v1.82.4, serverOnly is declared like this.
```ts
export declare function serverOnly<T>(value: () => T): () => T | undefined;
```

This means that when I have a serverOnly function like this:
```ts
import { serverOnly } from "@tanstack/start";

export const getPokemon = serverOnly(async (pokemon: string) => {
  return await fetch(`https://pokeapi.co/api/v2/pokemon/${pokemon}`).then(res => res.json());
});
```

getPokemon is inferred to have the type `() => Promise<any> | undefined`, which means passing the `pokemon` parameter leads to a typescript error
```ts
export const Route = createRootRoute({
  beforeLoad: async () => {
    const pokemon = await getPokemon("ditto");
                                     // ^ Expected 0 arguments, but got 1. (2554)
    return {pokemon}
  },
  component: RootComponent,
});
```

## Alternative method
You can wrap the anonymous function inside another anonymous function to get the correct return type
```ts
import { serverOnly } from "@tanstack/start";

// Notice () => async () => {}
export const getPokemon = serverOnly(() => async (pokemon: string) => {
  return await fetch(`https://pokeapi.co/api/v2/pokemon/${pokemon}`).then(res => res.json());
});
```

But you will have to "unwrap" once
```ts
export const Route = createRootRoute({
  beforeLoad: async () => {
    const pokemon = await getPokemon?.()("ditto");
    return { pokemon };
  },
  component: RootComponent,
});
```

### Your Example Website or App

https://stackblitz.com/edit/github-avg9mj-9r1rad?file=app%2Froutes%2F__root.tsx

### Steps to Reproduce the Bug or Issue

1. Go to __root.tsx
2. Observe that I can't pass "ditto" as an argument to getPokemon, even if the function passed to serverOnly has a parameter of type string

### Expected behavior

getPokemon should be inferred as `(pokemon: string) => Promise<any> | undefined`

### Screenshots or Videos

_No response_

### Platform

- OS: Windows
- Browser: Microsoft Edge
- Version: Version 131.0.2903.51 (Official build) (64-bit)


### Additional context

_No response_
`serverOnly` uses wrong types
### Which project does this relate to?

Start

### Describe the bug

the `serverOnly` function gets swapped for client builds with
```ts
() => {
  invariant(false, "serverOnly() functions can only be called on the server!");
}
```
while the type signature is
```ts
declare function serverOnly<T>(value: () => T): () => T | undefined;
```
and the implementation is
```ts
const serverOnly = _ => _
```
This is a problem for multiple reasons:
1. `invariant` is not a valid ambient, and an import is not patched into the file, so the actual error seen is `Uncaught ReferenceError: invariant is not defined`
2. according to typescript semantics, the type must stay `() => T`, there is no scenario for `undefined` because we throw. This just promotes needlessly defensive programming in otherwise valid states
3. because of the `undefined` leak, there's the assumption that in clientside environments, it counts as a `noop`, but in reality we shouldn't expect any fallback behaviors because we don't return `undefined`, we throw.
4. The Implementation is much broader than a `() => T`, it's any `T`.

### Your Example Website or App

https://github.com/datner/router/tree/datner%2FserverOnly

### Steps to Reproduce the Bug or Issue

```sh
pnpm i && pnpm build:all && pnpm -F tanstack-start-example-basic dev
```
1. load page
2. trigger `beforeLoad` clientside (like move a page)


### Expected behavior

As a user, I expected to get `undefined` but U an seeing an error

### Screenshots or Videos

![Image](https://github.com/user-attachments/assets/4fdc4a4c-44f4-4a9e-ae79-6fae80fee8be)


### Platform

- OS: macOS
- Browser: Arc
- Version: 1.85.7


### Additional context

_No response_

Pull request:
fix(start): fix serverOnly implementation/types, and add clientOnly counterpart
fixes #2938 
fixes #2828

This makes the following changes to serverOnly:

- Changes types to allow passing any function, not just ones that don't accept any arguments (and fixes return type, as the replacement function throws instead of no-op'ing)
- Changes replacement function to actually throw an error instead of referencing a missing `invariant` function
- Extracts out function instead of passing to serverOnly in runtime

This also adds `clientOnly`, which does the same as `serverOnly` but for the client.

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
-}


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>()
+})