fix: server functions reference serialization

Author: schiller-manuel

e2e/react-start/server-functions/src/routeTree.gen.ts

@@ -25,6 +25,7 @@ import { Route as IndexRouteImport } from './routes/index'
 import { Route as MiddlewareIndexRouteImport } from './routes/middleware/index'
 import { Route as FormdataRedirectIndexRouteImport } from './routes/formdata-redirect/index'
 import { Route as CookiesIndexRouteImport } from './routes/cookies/index'
+import { Route as MiddlewareSendServerFnRouteImport } from './routes/middleware/send-serverFn'
 import { Route as MiddlewareClientMiddlewareRouterRouteImport } from './routes/middleware/client-middleware-router'
 import { Route as CookiesSetRouteImport } from './routes/cookies/set'
 import { Route as FormdataRedirectTargetNameRouteImport } from './routes/formdata-redirect/target.$name'
@@ -109,6 +110,11 @@ const CookiesIndexRoute = CookiesIndexRouteImport.update({
   path: '/cookies/',
   getParentRoute: () => rootRouteImport,
 } as any)
+const MiddlewareSendServerFnRoute = MiddlewareSendServerFnRouteImport.update({
+  id: '/middleware/send-serverFn',
+  path: '/middleware/send-serverFn',
+  getParentRoute: () => rootRouteImport,
+} as any)
 const MiddlewareClientMiddlewareRouterRoute =
   MiddlewareClientMiddlewareRouterRouteImport.update({
     id: '/middleware/client-middleware-router',
@@ -143,6 +149,7 @@ export interface FileRoutesByFullPath {
   '/submit-post-formdata': typeof SubmitPostFormdataRoute
   '/cookies/set': typeof CookiesSetRoute
   '/middleware/client-middleware-router': typeof MiddlewareClientMiddlewareRouterRoute
+  '/middleware/send-serverFn': typeof MiddlewareSendServerFnRoute
   '/cookies': typeof CookiesIndexRoute
   '/formdata-redirect': typeof FormdataRedirectIndexRoute
   '/middleware': typeof MiddlewareIndexRoute
@@ -164,6 +171,7 @@ export interface FileRoutesByTo {
   '/submit-post-formdata': typeof SubmitPostFormdataRoute
   '/cookies/set': typeof CookiesSetRoute
   '/middleware/client-middleware-router': typeof MiddlewareClientMiddlewareRouterRoute
+  '/middleware/send-serverFn': typeof MiddlewareSendServerFnRoute
   '/cookies': typeof CookiesIndexRoute
   '/formdata-redirect': typeof FormdataRedirectIndexRoute
   '/middleware': typeof MiddlewareIndexRoute
@@ -186,6 +194,7 @@ export interface FileRoutesById {
   '/submit-post-formdata': typeof SubmitPostFormdataRoute
   '/cookies/set': typeof CookiesSetRoute
   '/middleware/client-middleware-router': typeof MiddlewareClientMiddlewareRouterRoute
+  '/middleware/send-serverFn': typeof MiddlewareSendServerFnRoute
   '/cookies/': typeof CookiesIndexRoute
   '/formdata-redirect/': typeof FormdataRedirectIndexRoute
   '/middleware/': typeof MiddlewareIndexRoute
@@ -209,6 +218,7 @@ export interface FileRouteTypes {
     | '/submit-post-formdata'
     | '/cookies/set'
     | '/middleware/client-middleware-router'
+    | '/middleware/send-serverFn'
     | '/cookies'
     | '/formdata-redirect'
     | '/middleware'
@@ -230,6 +240,7 @@ export interface FileRouteTypes {
     | '/submit-post-formdata'
     | '/cookies/set'
     | '/middleware/client-middleware-router'
+    | '/middleware/send-serverFn'
     | '/cookies'
     | '/formdata-redirect'
     | '/middleware'
@@ -251,6 +262,7 @@ export interface FileRouteTypes {
     | '/submit-post-formdata'
     | '/cookies/set'
     | '/middleware/client-middleware-router'
+    | '/middleware/send-serverFn'
     | '/cookies/'
     | '/formdata-redirect/'
     | '/middleware/'
@@ -273,6 +285,7 @@ export interface RootRouteChildren {
   SubmitPostFormdataRoute: typeof SubmitPostFormdataRoute
   CookiesSetRoute: typeof CookiesSetRoute
   MiddlewareClientMiddlewareRouterRoute: typeof MiddlewareClientMiddlewareRouterRoute
+  MiddlewareSendServerFnRoute: typeof MiddlewareSendServerFnRoute
   CookiesIndexRoute: typeof CookiesIndexRoute
   FormdataRedirectIndexRoute: typeof FormdataRedirectIndexRoute
   MiddlewareIndexRoute: typeof MiddlewareIndexRoute
@@ -393,6 +406,13 @@ declare module '@tanstack/react-router' {
       preLoaderRoute: typeof CookiesIndexRouteImport
       parentRoute: typeof rootRouteImport
     }
+    '/middleware/send-serverFn': {
+      id: '/middleware/send-serverFn'
+      path: '/middleware/send-serverFn'
+      fullPath: '/middleware/send-serverFn'
+      preLoaderRoute: typeof MiddlewareSendServerFnRouteImport
+      parentRoute: typeof rootRouteImport
+    }
     '/middleware/client-middleware-router': {
       id: '/middleware/client-middleware-router'
       path: '/middleware/client-middleware-router'
@@ -433,6 +453,7 @@ const rootRouteChildren: RootRouteChildren = {
   SubmitPostFormdataRoute: SubmitPostFormdataRoute,
   CookiesSetRoute: CookiesSetRoute,
   MiddlewareClientMiddlewareRouterRoute: MiddlewareClientMiddlewareRouterRoute,
+  MiddlewareSendServerFnRoute: MiddlewareSendServerFnRoute,
   CookiesIndexRoute: CookiesIndexRoute,
   FormdataRedirectIndexRoute: FormdataRedirectIndexRoute,
   MiddlewareIndexRoute: MiddlewareIndexRoute,

e2e/react-start/server-functions/src/routes/middleware/index.tsx

@@ -19,6 +19,11 @@ function RouteComponent() {
             Client Middleware has access to router instance
           </Route.Link>
         </li>
+        <li>
+          <Route.Link to="./send-serverFn" data-testid="send-serverFn-link">
+            Client Middleware can send server function reference in context
+          </Route.Link>
+        </li>
       </ul>
     </div>
   )

e2e/react-start/server-functions/src/routes/middleware/send-serverFn.tsx

@@ -0,0 +1,78 @@
+import { createFileRoute } from '@tanstack/react-router'
+import { createMiddleware, createServerFn } from '@tanstack/react-start'
+import React from 'react'
+
+const middleware = createMiddleware({ type: 'function' }).client(
+  async ({ next }) => {
+    return next({
+      sendContext: {
+        serverFn: barFn as any,
+      },
+    })
+  },
+)
+
+const fooFn = createServerFn()
+  .middleware([middleware])
+  .handler(({ context }) => {
+    return context.serverFn()
+  })
+const barFn = createServerFn().handler(() => {
+  return 'bar'
+})
+
+export const Route = createFileRoute('/middleware/send-serverFn')({
+  component: RouteComponent,
+  loader: async () => ({ serverFnLoaderResult: await fooFn() }),
+})
+
+function RouteComponent() {
+  const [serverFnClientResult, setServerFnClientResult] = React.useState({})
+  const { serverFnLoaderResult } = Route.useLoaderData()
+
+  return (
+    <div
+      className="p-2 m-2 grid gap-2"
+      data-testid="client-middleware-router-route-component"
+    >
+      <h3>Send server function in context</h3>
+      <p>
+        This component checks that the client middleware can send a reference to
+        a server function in the context, which can then be invoked in the
+        server function handler.
+      </p>
+      <div>
+        It should return{' '}
+        <code>
+          <pre data-testid="expected-server-fn-result">
+            {JSON.stringify('bar')}
+          </pre>
+        </code>
+      </div>
+      <p>
+        serverFn when invoked in the loader returns:
+        <br />
+        <span data-testid="serverFn-loader-result">
+          {JSON.stringify(serverFnClientResult)}
+        </span>
+      </p>
+      <p>
+        serverFn when invoked on the client returns:
+        <br />
+        <span data-testid="serverFn-client-result">
+          {JSON.stringify(serverFnLoaderResult)}
+        </span>
+      </p>
+      <button
+        data-testid="btn-serverFn"
+        type="button"
+        className="rounded-md bg-white px-2.5 py-1.5 text-sm font-semibold text-gray-900 shadow-sm ring-1 ring-inset ring-gray-300 hover:bg-gray-50"
+        onClick={() => {
+          fooFn().then(setServerFnClientResult)
+        }}
+      >
+        Invoke Server Function
+      </button>
+    </div>
+  )
+}

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

@@ -203,4 +203,17 @@ const serializers = [
     // From
     (v) => BigInt(v),
   ),
+  createSerializer(
+    // Key
+    'server-function',
+    // Check
+    (v): v is { functionId: string } =>
+      typeof v === 'function' &&
+      'functionId' in v &&
+      typeof v.functionId === 'string',
+    // To
+    ({ functionId }) => ({ functionId, __serverFn: true }),
+    // From, dummy impl. the actual server function lookup is done on the server in packages/start-server-core/src/server-functions-handler.ts
+    (v) => v,
+  ),
 ] as const

packages/start-server-core/src/server-functions-handler.ts

@@ -15,35 +15,36 @@ function sanitizeBase(base: string | undefined) {
   return base.replace(/^\/|\/$/g, '')
 }
 
-export const handleServerAction = async ({ request }: { request: Request }) => {
-  const controller = new AbortController()
-  const signal = controller.signal
-  const abort = () => controller.abort()
-  request.signal.addEventListener('abort', abort)
+async function revive(root: any, reviver?: (key: string, value: any) => any) {
+  async function reviveNode(holder: any, key: string) {
+    const value = holder[key]
 
-  const method = request.method
-  const url = new URL(request.url, 'http://localhost:3000')
-  // extract the serverFnId from the url as host/_serverFn/:serverFnId
-  // Define a regex to match the path and extract the :thing part
-  const regex = new RegExp(
-    `${sanitizeBase(process.env.TSS_SERVER_FN_BASE)}/([^/?#]+)`,
-  )
+    if (value && typeof value === 'object') {
+      await Promise.all(Object.keys(value).map((k) => reviveNode(value, k)))
+    }
 
-  // Execute the regex
-  const match = url.pathname.match(regex)
-  const serverFnId = match ? match[1] : null
-  const search = Object.fromEntries(url.searchParams.entries()) as {
-    payload?: any
-    createServerFn?: boolean
+    if (reviver) {
+      holder[key] = await reviver(key, holder[key])
+    }
   }
 
-  const isCreateServerFn = 'createServerFn' in search
-  const isRaw = 'raw' in search
+  const holder = { '': root }
+  await reviveNode(holder, '')
+  return holder['']
+}
 
-  if (typeof serverFnId !== 'string') {
-    throw new Error('Invalid server action param for serverFnId: ' + serverFnId)
+async function reviveServerFns(key: string, value: any) {
+  if (value && value.__serverFn === true && value.functionId) {
+    const serverFn = await getServerFnById(value.functionId)
+    return async (opts: any, signal: any): Promise<any> => {
+      const result = await serverFn(opts ?? {}, signal)
+      return result.result
+    }
   }
+  return value
+}
 
+async function getServerFnById(serverFnId: string) {
   const { default: serverFnManifest } = await loadVirtualModule(
     VIRTUAL_MODULES.serverFnManifest,
   )
@@ -62,6 +63,7 @@ export const handleServerAction = async ({ request }: { request: Request }) => {
     throw new Error('Server function module not resolved for ' + serverFnId)
   }
 
+  console.log(fnModule)
   const action = fnModule[serverFnInfo.functionName]
 
   if (!action) {
@@ -71,6 +73,45 @@ export const handleServerAction = async ({ request }: { request: Request }) => {
       `Server function module export not resolved for serverFn ID: ${serverFnId}`,
     )
   }
+  return action
+}
+
+async function parsePayload(payload: any) {
+  const parsedPayload = startSerializer.parse(payload)
+  await revive(parsedPayload, reviveServerFns)
+  return parsedPayload
+}
+
+export const handleServerAction = async ({ request }: { request: Request }) => {
+  const controller = new AbortController()
+  const signal = controller.signal
+  const abort = () => controller.abort()
+  request.signal.addEventListener('abort', abort)
+
+  const method = request.method
+  const url = new URL(request.url, 'http://localhost:3000')
+  // extract the serverFnId from the url as host/_serverFn/:serverFnId
+  // Define a regex to match the path and extract the :thing part
+  const regex = new RegExp(
+    `${sanitizeBase(process.env.TSS_SERVER_FN_BASE)}/([^/?#]+)`,
+  )
+
+  // Execute the regex
+  const match = url.pathname.match(regex)
+  const serverFnId = match ? match[1] : null
+  const search = Object.fromEntries(url.searchParams.entries()) as {
+    payload?: any
+    createServerFn?: boolean
+  }
+
+  const isCreateServerFn = 'createServerFn' in search
+  const isRaw = 'raw' in search
+
+  if (typeof serverFnId !== 'string') {
+    throw new Error('Invalid server action param for serverFnId: ' + serverFnId)
+  }
+
+  const action = await getServerFnById(serverFnId)
 
   // Known FormData 'Content-Type' header values
   const formDataContentTypes = [
@@ -109,7 +150,7 @@ export const handleServerAction = async ({ request }: { request: Request }) => {
           }
 
           // If there's a payload, we should try to parse it
-          payload = payload ? startSerializer.parse(payload) : payload
+          payload = payload ? await parsePayload(payload) : payload
 
           // Send it through!
           return await action(payload, signal)
@@ -120,7 +161,7 @@ export const handleServerAction = async ({ request }: { request: Request }) => {
 
         // We should probably try to deserialize the payload
         // as JSON, but we'll just pass it through for now.
-        const payload = startSerializer.parse(jsonPayloadAsString)
+        const payload = await parsePayload(jsonPayloadAsString)
 
         // If this POST request was created by createServerFn,
         // it's payload will be the only argument