fix: Make serverFn union response type distributive (#5720)

Author: brhx

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

@@ -301,9 +301,11 @@ export type RscStream<T> = {
 export type Method = 'GET' | 'POST'
 
 export type ServerFnReturnType<TRegister, TResponse> =
-  Awaited<TResponse> extends Response
-    ? TResponse
-    : ValidateSerializableInput<TRegister, TResponse>
+  TResponse extends PromiseLike<infer U>
+    ? Promise<ServerFnReturnType<TRegister, U>>
+    : TResponse extends Response
+      ? TResponse
+      : ValidateSerializableInput<TRegister, TResponse>
 
 export type ServerFn<
   TRegister,

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

@@ -6,9 +6,13 @@ import type {
   Constrain,
   Register,
   TsrSerializable,
+  ValidateSerializableInput,
   Validator,
 } from '@tanstack/router-core'
-import type { ConstrainValidator } from '../createServerFn'
+import type {
+  ConstrainValidator,
+  ServerFnReturnType,
+} from '../createServerFn'
 
 test('createServerFn method with autocomplete', () => {
   createServerFn().handler((options) => {
@@ -308,7 +312,10 @@ test('createServerFn returns Date', () => {
     dates: [new Date(), new Date()] as const,
   }))
 
-  expectTypeOf(fn()).toEqualTypeOf<Promise<{ dates: readonly [Date, Date] }>>()
+  expectTypeOf<ReturnType<typeof fn>>().toMatchTypeOf<Promise<unknown>>()
+  expectTypeOf<Awaited<ReturnType<typeof fn>>>().toMatchTypeOf<
+    ValidateSerializableInput<Register, { dates: readonly [Date, Date] }>
+  >()
 })
 
 test('createServerFn returns undefined', () => {
@@ -379,6 +386,23 @@ describe('response', () => {
 
     expectTypeOf(fn()).toEqualTypeOf<Promise<Response>>()
   })
+
+  test(`client receives union when handler may return Response or string`, () => {
+    const fn = createServerFn().handler(() => {
+      const result: Response | 'Hello World' =
+        Math.random() > 0.5 ? new Response('Hello World') : 'Hello World'
+
+      return result
+    })
+
+    expectTypeOf(fn()).toEqualTypeOf<Promise<Response | 'Hello World'>>()
+  })
+})
+
+test('ServerFnReturnType distributes Response union', () => {
+  expectTypeOf<
+    ServerFnReturnType<Register, Response | 'Hello World'>
+  >().toEqualTypeOf<Response | 'Hello World'>()
 })
 
 test('createServerFn can be used as a mutation function', () => {