feat(react-router): Automatically flush on serverless for loaders/actions (#17234)

Author: chargome

packages/react-router/src/server/wrapServerAction.ts

@@ -1,6 +1,7 @@
 import { SEMATTRS_HTTP_TARGET } from '@opentelemetry/semantic-conventions';
 import type { SpanAttributes } from '@sentry/core';
 import {
+  flushIfServerless,
   getActiveSpan,
   getRootSpan,
   SEMANTIC_ATTRIBUTE_SENTRY_OP,
@@ -59,17 +60,21 @@ export function wrapServerAction<T>(options: SpanOptions = {}, actionFn: (args:
       }
     }
 
-    return startSpan(
-      {
-        name,
-        ...options,
-        attributes: {
-          [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.http.react-router.action',
-          [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'function.react-router.action',
-          ...options.attributes,
+    try {
+      return await startSpan(
+        {
+          name,
+          ...options,
+          attributes: {
+            [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.http.react-router.action',
+            [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'function.react-router.action',
+            ...options.attributes,
+          },
         },
-      },
-      () => actionFn(args),
-    );
+        () => actionFn(args),
+      );
+    } finally {
+      await flushIfServerless();
+    }
   };
 }

packages/react-router/src/server/wrapServerLoader.ts

@@ -1,6 +1,7 @@
 import { SEMATTRS_HTTP_TARGET } from '@opentelemetry/semantic-conventions';
 import type { SpanAttributes } from '@sentry/core';
 import {
+  flushIfServerless,
   getActiveSpan,
   getRootSpan,
   SEMANTIC_ATTRIBUTE_SENTRY_OP,
@@ -59,17 +60,21 @@ export function wrapServerLoader<T>(options: SpanOptions = {}, loaderFn: (args:
         }
       }
     }
-    return startSpan(
-      {
-        name,
-        ...options,
-        attributes: {
-          [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.http.react-router.loader',
-          [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'function.react-router.loader',
-          ...options.attributes,
+    try {
+      return await startSpan(
+        {
+          name,
+          ...options,
+          attributes: {
+            [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.http.react-router.loader',
+            [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'function.react-router.loader',
+            ...options.attributes,
+          },
         },
-      },
-      () => loaderFn(args),
-    );
+        () => loaderFn(args),
+      );
+    } finally {
+      await flushIfServerless();
+    }
   };
 }

packages/react-router/test/server/wrapSentryHandleRequest.test.ts

@@ -176,27 +176,6 @@ describe('wrapSentryHandleRequest', () => {
   });
 });
 
-test('should not set span attributes when parameterized path does not exist', async () => {
-  const mockActiveSpan = {};
-  const mockRootSpan = { setAttributes: vi.fn() };
-
-  (getActiveSpan as unknown as ReturnType<typeof vi.fn>).mockReturnValue(mockActiveSpan);
-  (getRootSpan as unknown as ReturnType<typeof vi.fn>).mockReturnValue(mockRootSpan);
-
-  const originalHandler = vi.fn().mockResolvedValue('test');
-  const wrappedHandler = wrapSentryHandleRequest(originalHandler);
-
-  const routerContext = {
-    staticHandlerContext: {
-      matches: [],
-    },
-  } as any;
-
-  await wrappedHandler(new Request('https://guapo.chulo'), 200, new Headers(), routerContext, {} as any);
-
-  expect(mockRootSpan.setAttributes).not.toHaveBeenCalled();
-});
-
 describe('getMetaTagTransformer', () => {
   beforeEach(() => {
     vi.clearAllMocks();

packages/react-router/test/server/wrapServerAction.test.ts

@@ -3,6 +3,15 @@ import type { ActionFunctionArgs } from 'react-router';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import { wrapServerAction } from '../../src/server/wrapServerAction';
 
+vi.mock('@sentry/core', async () => {
+  const actual = await vi.importActual('@sentry/core');
+  return {
+    ...actual,
+    startSpan: vi.fn(),
+    flushIfServerless: vi.fn(),
+  };
+});
+
 describe('wrapServerAction', () => {
   beforeEach(() => {
     vi.clearAllMocks();
@@ -12,11 +21,12 @@ describe('wrapServerAction', () => {
     const mockActionFn = vi.fn().mockResolvedValue('result');
     const mockArgs = { request: new Request('http://test.com') } as ActionFunctionArgs;
 
-    const spy = vi.spyOn(core, 'startSpan');
+    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn());
+
     const wrappedAction = wrapServerAction({}, mockActionFn);
     await wrappedAction(mockArgs);
 
-    expect(spy).toHaveBeenCalledWith(
+    expect(core.startSpan).toHaveBeenCalledWith(
       {
         name: 'Executing Server Action',
         attributes: {
@@ -27,6 +37,7 @@ describe('wrapServerAction', () => {
       expect.any(Function),
     );
     expect(mockActionFn).toHaveBeenCalledWith(mockArgs);
+    expect(core.flushIfServerless).toHaveBeenCalled();
   });
 
   it('should wrap an action function with custom options', async () => {
@@ -40,11 +51,12 @@ describe('wrapServerAction', () => {
     const mockActionFn = vi.fn().mockResolvedValue('result');
     const mockArgs = { request: new Request('http://test.com') } as ActionFunctionArgs;
 
-    const spy = vi.spyOn(core, 'startSpan');
+    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn());
+
     const wrappedAction = wrapServerAction(customOptions, mockActionFn);
     await wrappedAction(mockArgs);
 
-    expect(spy).toHaveBeenCalledWith(
+    expect(core.startSpan).toHaveBeenCalledWith(
       {
         name: 'Custom Action',
         attributes: {
@@ -56,5 +68,43 @@ describe('wrapServerAction', () => {
       expect.any(Function),
     );
     expect(mockActionFn).toHaveBeenCalledWith(mockArgs);
+    expect(core.flushIfServerless).toHaveBeenCalled();
+  });
+
+  it('should call flushIfServerless on successful execution', async () => {
+    const mockActionFn = vi.fn().mockResolvedValue('result');
+    const mockArgs = { request: new Request('http://test.com') } as ActionFunctionArgs;
+
+    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn());
+
+    const wrappedAction = wrapServerAction({}, mockActionFn);
+    await wrappedAction(mockArgs);
+
+    expect(core.flushIfServerless).toHaveBeenCalled();
+  });
+
+  it('should call flushIfServerless even when action throws an error', async () => {
+    const mockError = new Error('Action failed');
+    const mockActionFn = vi.fn().mockRejectedValue(mockError);
+    const mockArgs = { request: new Request('http://test.com') } as ActionFunctionArgs;
+
+    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn());
+
+    const wrappedAction = wrapServerAction({}, mockActionFn);
+
+    await expect(wrappedAction(mockArgs)).rejects.toThrow('Action failed');
+    expect(core.flushIfServerless).toHaveBeenCalled();
+  });
+
+  it('should propagate errors from action function', async () => {
+    const mockError = new Error('Test error');
+    const mockActionFn = vi.fn().mockRejectedValue(mockError);
+    const mockArgs = { request: new Request('http://test.com') } as ActionFunctionArgs;
+
+    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn());
+
+    const wrappedAction = wrapServerAction({}, mockActionFn);
+
+    await expect(wrappedAction(mockArgs)).rejects.toBe(mockError);
   });
 });

packages/react-router/test/server/wrapServerLoader.test.ts

@@ -3,6 +3,15 @@ import type { LoaderFunctionArgs } from 'react-router';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import { wrapServerLoader } from '../../src/server/wrapServerLoader';
 
+vi.mock('@sentry/core', async () => {
+  const actual = await vi.importActual('@sentry/core');
+  return {
+    ...actual,
+    startSpan: vi.fn(),
+    flushIfServerless: vi.fn(),
+  };
+});
+
 describe('wrapServerLoader', () => {
   beforeEach(() => {
     vi.clearAllMocks();
@@ -12,11 +21,12 @@ describe('wrapServerLoader', () => {
     const mockLoaderFn = vi.fn().mockResolvedValue('result');
     const mockArgs = { request: new Request('http://test.com') } as LoaderFunctionArgs;
 
-    const spy = vi.spyOn(core, 'startSpan');
+    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn());
+
     const wrappedLoader = wrapServerLoader({}, mockLoaderFn);
     await wrappedLoader(mockArgs);
 
-    expect(spy).toHaveBeenCalledWith(
+    expect(core.startSpan).toHaveBeenCalledWith(
       {
         name: 'Executing Server Loader',
         attributes: {
@@ -27,6 +37,7 @@ describe('wrapServerLoader', () => {
       expect.any(Function),
     );
     expect(mockLoaderFn).toHaveBeenCalledWith(mockArgs);
+    expect(core.flushIfServerless).toHaveBeenCalled();
   });
 
   it('should wrap a loader function with custom options', async () => {
@@ -40,11 +51,12 @@ describe('wrapServerLoader', () => {
     const mockLoaderFn = vi.fn().mockResolvedValue('result');
     const mockArgs = { request: new Request('http://test.com') } as LoaderFunctionArgs;
 
-    const spy = vi.spyOn(core, 'startSpan');
+    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn());
+
     const wrappedLoader = wrapServerLoader(customOptions, mockLoaderFn);
     await wrappedLoader(mockArgs);
 
-    expect(spy).toHaveBeenCalledWith(
+    expect(core.startSpan).toHaveBeenCalledWith(
       {
         name: 'Custom Loader',
         attributes: {
@@ -56,5 +68,43 @@ describe('wrapServerLoader', () => {
       expect.any(Function),
     );
     expect(mockLoaderFn).toHaveBeenCalledWith(mockArgs);
+    expect(core.flushIfServerless).toHaveBeenCalled();
+  });
+
+  it('should call flushIfServerless on successful execution', async () => {
+    const mockLoaderFn = vi.fn().mockResolvedValue('result');
+    const mockArgs = { request: new Request('http://test.com') } as LoaderFunctionArgs;
+
+    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn());
+
+    const wrappedLoader = wrapServerLoader({}, mockLoaderFn);
+    await wrappedLoader(mockArgs);
+
+    expect(core.flushIfServerless).toHaveBeenCalled();
+  });
+
+  it('should call flushIfServerless even when loader throws an error', async () => {
+    const mockError = new Error('Loader failed');
+    const mockLoaderFn = vi.fn().mockRejectedValue(mockError);
+    const mockArgs = { request: new Request('http://test.com') } as LoaderFunctionArgs;
+
+    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn());
+
+    const wrappedLoader = wrapServerLoader({}, mockLoaderFn);
+
+    await expect(wrappedLoader(mockArgs)).rejects.toThrow('Loader failed');
+    expect(core.flushIfServerless).toHaveBeenCalled();
+  });
+
+  it('should propagate errors from loader function', async () => {
+    const mockError = new Error('Test error');
+    const mockLoaderFn = vi.fn().mockRejectedValue(mockError);
+    const mockArgs = { request: new Request('http://test.com') } as LoaderFunctionArgs;
+
+    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn());
+
+    const wrappedLoader = wrapServerLoader({}, mockLoaderFn);
+
+    await expect(wrappedLoader(mockArgs)).rejects.toBe(mockError);
   });
 });