fix(remix): Extract deferred responses correctly in root loaders. (#8305)

Co-authored-by: Abhijeet Prasad <[email protected]>

Author: onurtemizkan

packages/remix/src/utils/instrumentServer.ts

@@ -27,7 +27,7 @@ import type {
   ServerRoute,
   ServerRouteManifest,
 } from './types';
-import { extractData, getRequestMatch, isResponse, json, matchServerRoutes } from './vendor/response';
+import { extractData, getRequestMatch, isDeferredData, isResponse, json, matchServerRoutes } from './vendor/response';
 import { normalizeRemixRequest } from './web-fetch';
 
 // Flag to track if the core request handler is instrumented.
@@ -229,6 +229,13 @@ function makeWrappedRootLoader(origLoader: DataFunction): DataFunction {
     const res = await origLoader.call(this, args);
     const traceAndBaggage = getTraceAndBaggage();
 
+    if (isDeferredData(res)) {
+      return {
+        ...res.data,
+        ...traceAndBaggage,
+      };
+    }
+
     // Note: `redirect` and `catch` responses do not have bodies to extract
     if (isResponse(res) && !isRedirectResponse(res) && !isCatchResponse(res)) {
       const data = await extractData(res);

packages/remix/src/utils/types.ts

@@ -62,6 +62,15 @@ export interface RouteData {
   [routeId: string]: AppData;
 }
 
+export type DeferredData = {
+  data: Record<string, unknown>;
+  init?: ResponseInit;
+  deferredKeys: string[];
+  subscribe(fn: (aborted: boolean, settledKey?: string) => void): () => boolean;
+  cancel(): void;
+  resolveData(signal: AbortSignal): Promise<boolean>;
+};
+
 export interface MetaFunction {
   (args: { data: AppData; parentsData: RouteData; params: Params; location: Location }): HtmlMetaDescriptor;
 }

packages/remix/src/utils/vendor/response.ts

@@ -6,7 +6,7 @@
 //
 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
-import type { ReactRouterDomPkg, RouteMatch, ServerRoute } from '../types';
+import type { DeferredData, ReactRouterDomPkg, RouteMatch, ServerRoute } from '../types';
 
 /**
  * Based on Remix Implementation
@@ -124,3 +124,18 @@ export function getRequestMatch(url: URL, matches: RouteMatch<ServerRoute>[]): R
 
   return match;
 }
+
+/**
+ * https://github.com/remix-run/remix/blob/3e589152bc717d04e2054c31bea5a1056080d4b9/packages/remix-server-runtime/responses.ts#L75-L85
+ */
+export function isDeferredData(value: any): value is DeferredData {
+  const deferred: DeferredData = value;
+  return (
+    deferred &&
+    typeof deferred === 'object' &&
+    typeof deferred.data === 'object' &&
+    typeof deferred.subscribe === 'function' &&
+    typeof deferred.cancel === 'function' &&
+    typeof deferred.resolveData === 'function'
+  );
+}

packages/remix/test/integration/app/root.tsx

@@ -1,4 +1,4 @@
-import { MetaFunction, LoaderFunction, json, redirect } from '@remix-run/node';
+import { MetaFunction, LoaderFunction, json, defer, redirect } from '@remix-run/node';
 import { Links, LiveReload, Meta, Outlet, Scripts, ScrollRestoration } from '@remix-run/react';
 import { withSentry } from '@sentry/remix';
 
@@ -24,6 +24,8 @@ export const loader: LoaderFunction = async ({ request }) => {
       };
     case 'json':
       return json({ data_one: [], data_two: 'a string' }, { headers: { 'Cache-Control': 'max-age=300' } });
+    case 'defer':
+      return defer({ data_one: [], data_two: 'a string' });
     case 'null':
       return null;
     case 'undefined':

packages/remix/test/integration/app/routes/loader-defer-response/index.tsx

@@ -0,0 +1,20 @@
+import { defer, LoaderFunction } from '@remix-run/node';
+import { useLoaderData } from '@remix-run/react';
+
+type LoaderData = { id: string };
+
+export const loader: LoaderFunction = async ({ params: { id } }) => {
+  return defer({
+    id,
+  });
+};
+
+export default function LoaderJSONResponse() {
+  const data = useLoaderData<LoaderData>();
+
+  return (
+    <div>
+      <h1>{data && data.id ? data.id : 'Not Found'}</h1>
+    </div>
+  );
+}

packages/remix/test/integration/test/client/root-loader.test.ts

@@ -72,6 +72,22 @@ test('should inject `sentry-trace` and `baggage` into root loader returning a `J
   });
 });
 
+test('should inject `sentry-trace` and `baggage` into root loader returning a deferred response', async ({ page }) => {
+  await page.goto('/?type=defer');
+
+  const { sentryTrace, sentryBaggage } = await extractTraceAndBaggageFromMeta(page);
+
+  expect(sentryTrace).toEqual(expect.any(String));
+  expect(sentryBaggage).toEqual(expect.any(String));
+
+  const rootData = (await getRouteData(page))['root'];
+
+  expect(rootData).toMatchObject({
+    sentryTrace: sentryTrace,
+    sentryBaggage: sentryBaggage,
+  });
+});
+
 test('should inject `sentry-trace` and `baggage` into root loader returning `null`.', async ({ page }) => {
   await page.goto('/?type=null');
 

packages/remix/test/integration/test/server/loader.test.ts

@@ -187,4 +187,32 @@ describe.each(['builtin', 'express'])('Remix API Loaders with adapter = %s', ada
       },
     });
   });
+
+  it('correctly instruments a deferred loader', async () => {
+    const env = await RemixTestEnv.init(adapter);
+    const url = `${env.url}/loader-defer-response`;
+    const envelope = await env.getEnvelopeRequest({ url, envelopeType: 'transaction' });
+    const transaction = envelope[2];
+
+    assertSentryTransaction(transaction, {
+      transaction: 'root',
+      transaction_info: {
+        source: 'route',
+      },
+      spans: [
+        {
+          description: 'root',
+          op: 'function.remix.loader',
+        },
+        {
+          description: 'routes/loader-defer-response/index',
+          op: 'function.remix.loader',
+        },
+        {
+          description: 'root',
+          op: 'function.remix.document_request',
+        },
+      ],
+    });
+  });
 });