fix(start): fix serverOnly implementation/types, and add clientOnly c…

…ounterpart (#2954) * Fix type of serverOnly and add clientOnly function * Replace invalid invariant call with throwing an error, and support clientOnly * only build handlers once and reuse * write type test and make linter happy * Add snapshot testing for serverOnly and clientOnly * Write e2e tests * use a single source of truth for tokens to transform and do dead code elimination after
TanStack · Dec 8, 2024 · 402bae6 · 402bae6
1 parent 52e9d4f
commit 402bae6
Show file tree

Hide file tree

Showing 19 changed files with 373 additions and 48 deletions.
diff --git 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
@@ -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
@@ -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/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
@@ -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,