feat: add $withType<TFn>() to restrict isomorphicFn

docs/start/framework/react/guide/environment-functions.md

@@ -83,6 +83,28 @@ A no-op (short for "no operation") is a function that does nothing when executed
 function noop() {}
 ```
 
+### Restricting the type of the implementations
+
+You can use `.$withType<TFn>()` to enforce the type of both the server and client implementations matches the same function signature.
+
+```tsx
+import { createIsomorphicFn } from '@tanstack/react-start'
+import type { Environment } from '@/types/environment'
+
+const getEnv = createIsomorphicFn()
+  .$withType<() => Environment>()
+  .server(() => 'server')
+  .client(() => 'client')
+
+const env = getEnv()
+//    ^? Environment
+```
+
+You will get a type error if either implementation mismatches from the function signature passed to `$withType`.
+
+> [!NOTE]
+> When using `$withType()`, TypeScript will enforce that you define both the server and client implementations.
+
 ---
 
 ## `env`Only Functions

packages/start-client-core/src/createIsomorphicFn.ts

@@ -21,7 +21,39 @@ export interface ClientOnlyFn<TArgs extends Array<any>, TClient>
   ) => IsomorphicFn<TArgs, TServer, TClient>
 }
 
+type AnyFn = (...args: Array<any>) => any
+export interface ServerOnlyFnWithType<TArgs extends Array<any>, TReturnType>
+  extends IsomorphicFn<TArgs, TReturnType> {
+  client: (
+    clientImpl: (...args: TArgs) => TReturnType,
+  ) => IsomorphicFn<TArgs, TReturnType, TReturnType>
+}
+
+export interface ClientOnlyFnWithType<TArgs extends Array<any>, TReturnType>
+  extends IsomorphicFn<TArgs, TReturnType, TReturnType> {
+  server: (
+    serverImpl: (...args: TArgs) => TReturnType,
+  ) => IsomorphicFn<TArgs, TReturnType, TReturnType>
+}
+
+export interface IsomorphicFnWithType<TArgs extends Array<any>, TReturnType> {
+  server: (serverImpl: (...args: TArgs) => TReturnType) => {
+    client: (
+      clientImpl: (...args: TArgs) => TReturnType,
+    ) => IsomorphicFn<TArgs, TReturnType, TReturnType>
+  }
+  client: (clientImpl: (...args: TArgs) => TReturnType) => {
+    server: (
+      serverImpl: (...args: TArgs) => TReturnType,
+    ) => IsomorphicFn<TArgs, TReturnType, TReturnType>
+  }
+}
+
 export interface IsomorphicFnBase extends IsomorphicFn {
+  $withType: <TFn extends AnyFn>() => IsomorphicFnWithType<
+    Parameters<TFn>,
+    ReturnType<TFn>
+  >
   server: <TArgs extends Array<any>, TServer>(
     serverImpl: (...args: TArgs) => TServer,
   ) => ServerOnlyFn<TArgs, TServer>
@@ -34,8 +66,12 @@ export interface IsomorphicFnBase extends IsomorphicFn {
 // if we use `createIsomorphicFn` in this library itself, vite tries to execute it before the transformer runs
 // therefore we must return a dummy function that allows calling `server` and `client` method chains.
 export function createIsomorphicFn(): IsomorphicFnBase {
-  return {
+  const baseFns = {
     server: () => ({ client: () => () => {} }),
     client: () => ({ server: () => () => {} }),
+  }
+  return {
+    $withType: () => baseFns,
+    ...baseFns,
   } as any
 }

packages/start-client-core/src/tests/createIsomorphicFn.test-d.ts

@@ -70,3 +70,56 @@ test('createIsomorphicFn with arguments', () => {
   expectTypeOf(fn2).toBeCallableWith(1, 'a')
   expectTypeOf(fn2).returns.toEqualTypeOf<string | number>()
 })
+
+test('createIsomorphicFn with type', () => {
+  const fn1 = createIsomorphicFn()
+    .$withType<(a: string) => string>()
+    .server((a) => {
+      expectTypeOf(a).toEqualTypeOf<string>()
+      return 'data'
+    })
+    .client((a) => {
+      expectTypeOf(a).toEqualTypeOf<string>()
+      return 'data'
+    })
+  expectTypeOf(fn1).toBeCallableWith('foo')
+  expectTypeOf(fn1).returns.toEqualTypeOf<string>()
+
+  const fn2 = createIsomorphicFn()
+    .$withType<(a: string, b: number) => boolean>()
+    .client((a, b) => {
+      expectTypeOf(a).toEqualTypeOf<string>()
+      expectTypeOf(b).toEqualTypeOf<number>()
+      return true as const
+    })
+    .server((a, b) => {
+      expectTypeOf(a).toEqualTypeOf<string>()
+      expectTypeOf(b).toEqualTypeOf<number>()
+      return false as const
+    })
+  expectTypeOf(fn2).toBeCallableWith('foo', 1)
+  expectTypeOf(fn2).returns.toEqualTypeOf<boolean>()
+
+  const noServerImpl = createIsomorphicFn()
+    .$withType<(a: string) => string>()
+    .client((a) => 'data')
+  expectTypeOf(noServerImpl).not.toBeFunction()
+
+  // Missing server implementation
+  const noClientImpl = createIsomorphicFn()
+    .$withType<(a: string) => string>()
+    .server((a) => 'data')
+  expectTypeOf(noClientImpl).not.toBeFunction()
+
+  const _invalidFnReturn = createIsomorphicFn()
+    .$withType<(a: string) => string>()
+    .server(() => '1')
+    // @ts-expect-error - invalid return type
+    .client(() => 2)
+
+  const _invalidFnArgs = createIsomorphicFn()
+    .$withType<(a: string) => string>()
+    .server((a: string) => 'data')
+    // @ts-expect-error - invalid argument type
+    .client((a: number) => 'data')
+})

packages/start-plugin-core/tests/createIsomorphicFn/snapshots/client/createIsomorphicFnDestructured.tsx

@@ -13,4 +13,5 @@ function abstractedClientFn() {
 }
 const clientOnlyFnAbstracted = abstractedClientFn;
 const serverThenClientFnAbstracted = abstractedClientFn;
-const clientThenServerFnAbstracted = abstractedClientFn;
+const clientThenServerFnAbstracted = abstractedClientFn;
+const withTypeRestriction = () => 'client';

packages/start-plugin-core/tests/createIsomorphicFn/snapshots/server/createIsomorphicFnDestructured.tsx

@@ -13,4 +13,5 @@ function abstractedClientFn() {
 }
 const clientOnlyFnAbstracted = () => {};
 const serverThenClientFnAbstracted = abstractedServerFn;
-const clientThenServerFnAbstracted = abstractedServerFn;
+const clientThenServerFnAbstracted = abstractedServerFn;
+const withTypeRestriction = () => 'server';

packages/start-plugin-core/tests/createIsomorphicFn/test-files/createIsomorphicFnDestructured.tsx

@@ -33,3 +33,8 @@ const serverThenClientFnAbstracted = createIsomorphicFn()
 const clientThenServerFnAbstracted = createIsomorphicFn()
   .client(abstractedClientFn)
   .server(abstractedServerFn)
+
+const withTypeRestriction = createIsomorphicFn()
+  .$withType<() => string>()
+  .server(() => 'server')
+  .client(() => 'client')