feat(remix): Support defer() usage in rootAuthLoader()

packages/remix/src/ssr/rootAuthLoader.ts

@@ -1,11 +1,15 @@
 import { sanitizeAuthObject } from '@clerk/backend';
+import type { defer } from '@remix-run/server-runtime';
+import { isDeferredData } from '@remix-run/server-runtime/dist/responses';
 
 import { invalidRootLoaderCallbackReturn } from '../errors';
 import { authenticateRequest } from './authenticateRequest';
 import type { LoaderFunctionArgs, LoaderFunctionReturn, RootAuthLoaderCallback, RootAuthLoaderOptions } from './types';
 import {
   assertValidHandlerResult,
+  getResponseClerkState,
   injectAuthIntoRequest,
+  injectRequestStateIntoDeferredData,
   injectRequestStateIntoResponse,
   interstitialJsonResponse,
   isRedirect,
@@ -64,6 +68,16 @@ export const rootAuthLoader: RootAuthLoader = async (
   const handlerResult = await handler(injectAuthIntoRequest(args, sanitizeAuthObject(requestState.toAuth())));
   assertValidHandlerResult(handlerResult, invalidRootLoaderCallbackReturn);
 
+  // When using defer(), we need to inject the clerk auth state into its internal data object.
+  if (isDeferredData(handlerResult)) {
+    return injectRequestStateIntoDeferredData(
+      // This is necessary because the DeferredData type is not exported from remix.
+      handlerResult as unknown as ReturnType<typeof defer>,
+      requestState,
+      args.context,
+    );
+  }
+
   if (isResponse(handlerResult)) {
     try {
       // respect and pass-through any redirects without modifying them

packages/remix/src/ssr/utils.ts

@@ -1,6 +1,6 @@
 import type { AuthObject, RequestState } from '@clerk/backend';
 import { constants, debugRequestState, loadInterstitialFromLocal } from '@clerk/backend';
-import type { AppLoadContext } from '@remix-run/server-runtime';
+import type { AppLoadContext, defer } from '@remix-run/server-runtime';
 import { json } from '@remix-run/server-runtime';
 import cookie from 'cookie';
 
@@ -108,9 +108,51 @@ export const injectRequestStateIntoResponse = async (
   context: AppLoadContext,
 ) => {
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
-  const { reason, message, isSignedIn, isInterstitial, ...rest } = requestState;
   const clone = response.clone();
   const data = await clone.json();
+
+  const { clerkState, headers } = getResponseClerkState(requestState, context);
+
+  // set the correct content-type header in case the user returned a `Response` directly
+  // without setting the header, instead of using the `json()` helper
+  clone.headers.set(constants.Headers.ContentType, constants.ContentTypes.Json);
+  headers.forEach((value, key) => {
+    clone.headers.set(key, value);
+  });
+
+  return json({ ...(data || {}), ...clerkState }, clone);
+};
+
+export function injectRequestStateIntoDeferredData(
+  data: ReturnType<typeof defer>,
+  requestState: RequestState,
+  context: AppLoadContext,
+) {
+  const { clerkState, headers } = getResponseClerkState(requestState, context);
+
+  // Avoid creating a new object here to retain referential equality.
+  data.data.clerkState = clerkState.clerkState;
+
+  if (typeof data.init !== 'undefined') {
+    data.init.headers = new Headers(data.init.headers);
+
+    headers.forEach((value, key) => {
+      // @ts-expect-error -- We are ensuring headers is defined above
+      data.init.headers.set(key, value);
+    });
+  }
+
+  return data;
+}
+
+/**
+ * Returns the clerk state object and observability headers to be injected into a loader response.
+ *
+ * @internal
+ */
+export function getResponseClerkState(requestState: RequestState, context: AppLoadContext) {
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  const { reason, message, isSignedIn, isInterstitial, ...rest } = requestState;
   const clerkState = wrapWithClerkState({
     __clerk_ssr_state: rest.toAuth(),
     __frontendApi: requestState.frontendApi,
@@ -126,15 +168,14 @@ export const injectRequestStateIntoResponse = async (
     __clerkJSUrl: getEnvVariable('CLERK_JS', context),
     __clerkJSVersion: getEnvVariable('CLERK_JS_VERSION', context),
   });
-  // set the correct content-type header in case the user returned a `Response` directly
-  // without setting the header, instead of using the `json()` helper
-  clone.headers.set(constants.Headers.ContentType, constants.ContentTypes.Json);
-  observabilityHeadersFromRequestState(requestState).forEach((value, key) => {
-    clone.headers.set(key, value);
-  });
 
-  return json({ ...(data || {}), ...clerkState }, clone);
-};
+  const headers = observabilityHeadersFromRequestState(requestState);
+
+  return {
+    clerkState,
+    headers,
+  };
+}
 
 /**
  * Wraps obscured clerk internals with a readable `clerkState` key.

playground/remix-node/app/root.tsx

@@ -1,18 +1,20 @@
-import type { DataFunctionArgs, Headers } from '@remix-run/node';
+import { defer, type DataFunctionArgs, type Headers } from '@remix-run/node';
 import type { MetaFunction } from '@remix-run/react';
-import { Links, LiveReload, Meta, Outlet, Scripts, ScrollRestoration, useLoaderData } from '@remix-run/react';
+import { Await, Links, LiveReload, Meta, Outlet, Scripts, ScrollRestoration, useLoaderData } from '@remix-run/react';
 import { getClerkDebugHeaders, rootAuthLoader } from '@clerk/remix/ssr.server';
 import { ClerkApp, ClerkErrorBoundary } from '@clerk/remix';
+import { Suspense } from 'react';
 
 export const loader = (args: DataFunctionArgs) => {
   return rootAuthLoader(
     args,
     ({ request }) => {
       const { user } = request;
+      const data: Promise<{ foo: string }> = new Promise(r => r({ foo: 'bar' }))
 
       console.log('root User:', user);
 
-      return { user };
+      return defer({ user, data }, { headers: { 'x-bryce': 'cool guy' } })
     },
     { loadUser: true },
   );
@@ -27,10 +29,8 @@ export function headers({
   loaderHeaders: Headers;
   parentHeaders: Headers;
 }) {
-  return {
-    'x-parent-header': parentHeaders.get('x-parent-header'),
-    ...getClerkDebugHeaders(loaderHeaders),
-  };
+  console.log(loaderHeaders)
+  return loaderHeaders 
 }
 
 export const meta: MetaFunction = () => {
@@ -60,6 +60,11 @@ function App() {
         <Links />
       </head>
       <body>
+        <Suspense fallback="Loading...">
+          <Await resolve={loaderData.data}>
+            {val => (<>Hello {val.foo}</>)}
+          </Await>
+        </Suspense>
         <Outlet />
         <ScrollRestoration />
         <Scripts />