flush in server actions/loaders

chargome · chargome · commit 9af345593b8b · 2025-07-30T18:54:58.000+02:00
diff --git a/packages/react-router/src/server/wrapServerAction.ts b/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();
+    }
   };
 }
diff --git a/packages/react-router/src/server/wrapServerLoader.ts b/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();
+    }
   };
 }
diff --git a/packages/react-router/test/server/wrapServerAction.test.ts b/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);
   });
 });
diff --git a/packages/react-router/test/server/wrapServerLoader.test.ts b/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);
   });
 });
