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

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

@@ -21,7 +21,36 @@ 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>
+  extends IsomorphicFn<TArgs, TReturnType> {
+  server: (
+    serverImpl: (...args: TArgs) => TReturnType,
+  ) => ServerOnlyFnWithType<TArgs, TReturnType>
+  client: (
+    clientImpl: (...args: TArgs) => TReturnType,
+  ) => ClientOnlyFnWithType<TArgs, 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>
@@ -35,7 +64,9 @@ export interface IsomorphicFnBase extends IsomorphicFn {
 // therefore we must return a dummy function that allows calling `server` and `client` method chains.
 export function createIsomorphicFn(): IsomorphicFnBase {
   return {
+    $withType: () => {},
     server: () => ({ client: () => () => {} }),
     client: () => ({ server: () => () => {} }),
   } as any
 }
+

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

@@ -70,3 +70,45 @@ 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 _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')