fix(remix): Support merging json responses from root loader functions. (#5548)

Introduces support of `json` responses while injecting `sentryTrace` and `sentryBaggage` into the `root` loader. 
Added a set of tests to ensure that we don't omit user-defined loader responses.

Author: onurtemizkan

packages/remix/package.json

@@ -61,8 +61,9 @@
     "lint:prettier": "prettier --check \"{src,test,scripts}/**/*.ts\"",
     "test": "run-s test:unit",
     "test:integration": "run-s test:integration:prepare test:integration:client test:integration:server",
+    "test:integration:clean": "(cd test/integration && rimraf .cache build node_modules)",
     "test:integration:ci": "run-s test:integration:prepare test:integration:client:ci test:integration:server",
-    "test:integration:prepare": "(cd test/integration && yarn)",
+    "test:integration:prepare": "yarn test:integration:clean && (cd test/integration && yarn)",
     "test:integration:client": "yarn playwright install-deps && yarn playwright test test/integration/test/client/",
     "test:integration:client:ci": "yarn test:integration:client --browser='all' --reporter='line'",
     "test:integration:server": "jest --config=test/integration/jest.config.js test/integration/test/server/",

packages/remix/src/utils/instrumentServer.ts

@@ -100,7 +100,7 @@ export interface RouteMatch<Route> {
 }
 
 // Taken from Remix Implementation
-// https://github.com/remix-run/remix/blob/7688da5c75190a2e29496c78721456d6e12e3abe/packages/remix-server-runtime/responses.ts#L54-L62
+// https://github.com/remix-run/remix/blob/32300ec6e6e8025602cea63e17a2201989589eab/packages/remix-server-runtime/responses.ts#L60-L77
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 function isResponse(value: any): value is Response {
   return (
@@ -114,6 +114,15 @@ function isResponse(value: any): value is Response {
   );
 }
 
+const redirectStatusCodes = new Set([301, 302, 303, 307, 308]);
+function isRedirectResponse(response: Response): boolean {
+  return redirectStatusCodes.has(response.status);
+}
+
+function isCatchResponse(response: Response): boolean {
+  return response.headers.get('X-Remix-Catch') != null;
+}
+
 // Based on Remix Implementation
 // https://github.com/remix-run/remix/blob/7688da5c75190a2e29496c78721456d6e12e3abe/packages/remix-server-runtime/data.ts#L131-L145
 function extractData(response: Response): Promise<unknown> {
@@ -207,6 +216,7 @@ function makeWrappedDataFunction(origFn: DataFunction, name: 'action' | 'loader'
     try {
       const span = activeTransaction.startChild({
         op: `remix.server.${name}`,
+        // TODO: Consider using the `id` of the route this function belongs to
         description: activeTransaction.name,
         tags: {
           name,
@@ -235,8 +245,8 @@ function makeWrappedAction(origAction: DataFunction): DataFunction {
   return makeWrappedDataFunction(origAction, 'action');
 }
 
-function makeWrappedLoader(origAction: DataFunction): DataFunction {
-  return makeWrappedDataFunction(origAction, 'loader');
+function makeWrappedLoader(origLoader: DataFunction): DataFunction {
+  return makeWrappedDataFunction(origLoader, 'loader');
 }
 
 function getTraceAndBaggage(): { sentryTrace?: string; sentryBaggage?: string } {
@@ -262,8 +272,21 @@ function getTraceAndBaggage(): { sentryTrace?: string; sentryBaggage?: string }
 function makeWrappedRootLoader(origLoader: DataFunction): DataFunction {
   return async function (this: unknown, args: DataFunctionArgs): Promise<Response | AppData> {
     const res = await origLoader.call(this, args);
+    const traceAndBaggage = getTraceAndBaggage();
+
+    // Note: `redirect` and `catch` responses do not have bodies to extract
+    if (isResponse(res) && !isRedirectResponse(res) && !isCatchResponse(res)) {
+      const data = await extractData(res);
+
+      if (typeof data === 'object') {
+        return { ...data, ...traceAndBaggage };
+      } else {
+        __DEBUG_BUILD__ && logger.warn('Skipping injection of trace and baggage as the response body is not an object');
+        return data;
+      }
+    }
 
-    return { ...res, ...getTraceAndBaggage() };
+    return { ...res, ...traceAndBaggage };
   };
 }
 

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

@@ -1,4 +1,4 @@
-import type { MetaFunction } from '@remix-run/node';
+import { MetaFunction, LoaderFunction, json, redirect } from '@remix-run/node';
 import { Links, LiveReload, Meta, Outlet, Scripts, ScrollRestoration } from '@remix-run/react';
 import { withSentry } from '@sentry/remix';
 
@@ -10,7 +10,35 @@ export const meta: MetaFunction = ({ data }) => ({
   baggage: data.sentryBaggage,
 });
 
-export function App() {
+export const loader: LoaderFunction = async ({ request }) => {
+  const url = new URL(request.url);
+  const type = url.searchParams.get('type');
+
+  switch (type) {
+    case 'empty':
+      return {};
+    case 'plain':
+      return {
+        data_one: [],
+        data_two: 'a string',
+      };
+    case 'json':
+      return json({ data_one: [], data_two: 'a string' }, { headers: { 'Cache-Control': 'max-age=300' } });
+    case 'null':
+      return null;
+    case 'undefined':
+      return undefined;
+    case 'throwRedirect':
+      throw redirect('/?type=plain');
+    case 'returnRedirect':
+      return redirect('/?type=plain');
+    default: {
+      return {};
+    }
+  }
+};
+
+function App() {
   return (
     <html lang="en">
       <head>

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

@@ -0,0 +1,141 @@
+import { test, expect, Page } from '@playwright/test';
+
+async function getRouteData(page: Page): Promise<any> {
+  return page.evaluate('window.__remixContext.routeData').catch(err => {
+    console.warn(err);
+
+    return {};
+  });
+}
+
+async function extractTraceAndBaggageFromMeta(
+  page: Page,
+): Promise<{ sentryTrace?: string | null; sentryBaggage?: string | null }> {
+  const sentryTraceTag = await page.$('meta[name="sentry-trace"]');
+  const sentryTraceContent = await sentryTraceTag?.getAttribute('content');
+
+  const sentryBaggageTag = await page.$('meta[name="baggage"]');
+  const sentryBaggageContent = await sentryBaggageTag?.getAttribute('content');
+
+  return { sentryTrace: sentryTraceContent, sentryBaggage: sentryBaggageContent };
+}
+
+test('should inject `sentry-trace` and `baggage` into root loader returning an empty object.', async ({ page }) => {
+  await page.goto('/?type=empty');
+
+  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,
+    sentryBaggage,
+  });
+});
+
+test('should inject `sentry-trace` and `baggage` into root loader returning a plain object.', async ({ page }) => {
+  await page.goto('/?type=plain');
+
+  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({
+    data_one: [],
+    data_two: 'a string',
+    sentryTrace: sentryTrace,
+    sentryBaggage: sentryBaggage,
+  });
+});
+
+test('should inject `sentry-trace` and `baggage` into root loader returning a `JSON response`.', async ({ page }) => {
+  await page.goto('/?type=json');
+
+  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({
+    data_one: [],
+    data_two: 'a string',
+    sentryTrace: sentryTrace,
+    sentryBaggage: sentryBaggage,
+  });
+});
+
+test('should inject `sentry-trace` and `baggage` into root loader returning `null`.', async ({ page }) => {
+  await page.goto('/?type=null');
+
+  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 `undefined`.', async ({ page }) => {
+  await page.goto('/?type=undefined');
+
+  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 throwing a redirection to a plain object.', async ({
+  page,
+}) => {
+  await page.goto('/?type=throwRedirect');
+
+  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 a redirection to a plain object', async ({
+  page,
+}) => {
+  await page.goto('/?type=returnRedirect');
+
+  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,
+  });
+});

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

@@ -22,6 +22,13 @@ describe('Remix API Actions', () => {
           description: 'routes/action-json-response/$id',
           op: 'remix.server.action',
         },
+        // TODO: These two spans look exactly the same, but they are not.
+        // One is from the parent route, and the other is from the route we are reaching.
+        // We need to pass the names of the routes as their descriptions while wrapping loaders and actions.
+        {
+          description: 'routes/action-json-response/$id',
+          op: 'remix.server.loader',
+        },
         {
           description: 'routes/action-json-response/$id',
           op: 'remix.server.loader',

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

@@ -9,23 +9,6 @@ import {
 jest.spyOn(console, 'error').mockImplementation();
 
 describe('Remix API Loaders', () => {
-  it('does not add a loader if there is not one defined.', async () => {
-    const baseURL = await runServer();
-    const url = `${baseURL}/`;
-    const envelope = await getEnvelopeRequest(url);
-    const transaction = envelope[2];
-
-    assertSentryTransaction(transaction, {
-      transaction: 'root',
-      spans: [
-        {
-          description: 'root',
-          op: 'remix.server.documentRequest',
-        },
-      ],
-    });
-  });
-
   it('reports an error thrown from the loader', async () => {
     const baseURL = await runServer();
     const url = `${baseURL}/loader-json-response/-2`;
@@ -75,6 +58,13 @@ describe('Remix API Loaders', () => {
         source: 'route',
       },
       spans: [
+        // TODO: These two spans look exactly the same, but they are not.
+        // One is from the parent route, and the other is from the route we are reaching.
+        // We need to pass the names of the routes as their descriptions while wrapping loaders and actions.
+        {
+          description: 'routes/loader-json-response/$id',
+          op: 'remix.server.loader',
+        },
         {
           description: 'routes/loader-json-response/$id',
           op: 'remix.server.loader',