Use startSpanManual in wrapServerFunction to preserve redirect span status

Author: onurtemizkan

packages/react-router/src/server/rsc/wrapServerFunction.ts

@@ -9,7 +9,7 @@ import {
   SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
   SPAN_STATUS_ERROR,
   SPAN_STATUS_OK,
-  startSpan,
+  startSpanManual,
 } from '@sentry/core';
 import { DEBUG_BUILD } from '../../common/debug-build';
 import { isNotFoundResponse, isRedirectResponse } from './responseUtils';
@@ -58,7 +58,9 @@ export function wrapServerFunction<T extends (...args: any[]) => Promise<any>>(
 
       const hasActiveSpan = !!getActiveSpan();
 
-      return startSpan(
+      // startSpanManual is used instead of startSpan because startSpan's error handler
+      // would overwrite the intentional SPAN_STATUS_OK set for redirect responses.
+      return startSpanManual(
         {
           name: spanName,
           forceTransaction: !hasActiveSpan,
@@ -73,15 +75,18 @@ export function wrapServerFunction<T extends (...args: any[]) => Promise<any>>(
         async span => {
           try {
             const result = await originalFunction.apply(thisArg, args);
+            span.end();
             return result;
           } catch (error) {
             if (isRedirectResponse(error)) {
               span.setStatus({ code: SPAN_STATUS_OK });
+              span.end();
               throw error;
             }
 
             if (isNotFoundResponse(error)) {
               span.setStatus({ code: SPAN_STATUS_ERROR, message: 'not_found' });
+              span.end();
               throw error;
             }
 
@@ -97,6 +102,7 @@ export function wrapServerFunction<T extends (...args: any[]) => Promise<any>>(
                 },
               },
             });
+            span.end();
             throw error;
           } finally {
             await flushIfServerless();

packages/react-router/test/server/rsc/wrapServerFunction.test.ts

@@ -6,7 +6,7 @@ vi.mock('@sentry/core', async () => {
   const actual = await vi.importActual('@sentry/core');
   return {
     ...actual,
-    startSpan: vi.fn(),
+    startSpanManual: vi.fn(),
     getIsolationScope: vi.fn(),
     captureException: vi.fn(),
     flushIfServerless: vi.fn().mockResolvedValue(undefined),
@@ -25,7 +25,7 @@ describe('wrapServerFunction', () => {
     const mockSetTransactionName = vi.fn();
 
     (core.getIsolationScope as any).mockReturnValue({ setTransactionName: mockSetTransactionName });
-    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn({ setStatus: vi.fn() }));
+    (core.startSpanManual as any).mockImplementation((_: any, fn: any) => fn({ setStatus: vi.fn(), end: vi.fn() }));
 
     const wrappedFn = wrapServerFunction('testFunction', mockServerFn);
     const result = await wrappedFn('arg1', 'arg2');
@@ -34,7 +34,7 @@ describe('wrapServerFunction', () => {
     expect(mockServerFn).toHaveBeenCalledWith('arg1', 'arg2');
     expect(core.getIsolationScope).toHaveBeenCalled();
     expect(mockSetTransactionName).toHaveBeenCalledWith('serverFunction/testFunction');
-    expect(core.startSpan).toHaveBeenCalledWith(
+    expect(core.startSpanManual).toHaveBeenCalledWith(
       expect.objectContaining({
         name: 'serverFunction/testFunction',
         forceTransaction: true,
@@ -56,12 +56,12 @@ describe('wrapServerFunction', () => {
 
     (core.getIsolationScope as any).mockReturnValue({ setTransactionName: mockSetTransactionName });
     (core.getActiveSpan as any).mockReturnValue({ spanId: 'existing-span' });
-    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn({ setStatus: vi.fn() }));
+    (core.startSpanManual as any).mockImplementation((_: any, fn: any) => fn({ setStatus: vi.fn(), end: vi.fn() }));
 
     const wrappedFn = wrapServerFunction('testFunction', mockServerFn);
     await wrappedFn();
 
-    expect(core.startSpan).toHaveBeenCalledWith(
+    expect(core.startSpanManual).toHaveBeenCalledWith(
       expect.objectContaining({
         forceTransaction: false,
       }),
@@ -74,15 +74,15 @@ describe('wrapServerFunction', () => {
     const mockSetTransactionName = vi.fn();
 
     (core.getIsolationScope as any).mockReturnValue({ setTransactionName: mockSetTransactionName });
-    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn({ setStatus: vi.fn() }));
+    (core.startSpanManual as any).mockImplementation((_: any, fn: any) => fn({ setStatus: vi.fn(), end: vi.fn() }));
 
     const wrappedFn = wrapServerFunction('testFunction', mockServerFn, {
       name: 'Custom Span Name',
     });
     await wrappedFn();
 
     expect(mockSetTransactionName).toHaveBeenCalledWith('Custom Span Name');
-    expect(core.startSpan).toHaveBeenCalledWith(
+    expect(core.startSpanManual).toHaveBeenCalledWith(
       expect.objectContaining({
         name: 'Custom Span Name',
       }),
@@ -95,14 +95,14 @@ describe('wrapServerFunction', () => {
     const mockSetTransactionName = vi.fn();
 
     (core.getIsolationScope as any).mockReturnValue({ setTransactionName: mockSetTransactionName });
-    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn({ setStatus: vi.fn() }));
+    (core.startSpanManual as any).mockImplementation((_: any, fn: any) => fn({ setStatus: vi.fn(), end: vi.fn() }));
 
     const wrappedFn = wrapServerFunction('testFunction', mockServerFn, {
       attributes: { 'custom.attr': 'value' },
     });
     await wrappedFn();
 
-    expect(core.startSpan).toHaveBeenCalledWith(
+    expect(core.startSpanManual).toHaveBeenCalledWith(
       expect.objectContaining({
         attributes: expect.objectContaining({
           [core.SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'function.rsc.server_function',
@@ -117,10 +117,13 @@ describe('wrapServerFunction', () => {
     const mockError = new Error('Server function failed');
     const mockServerFn = vi.fn().mockRejectedValue(mockError);
     const mockSetStatus = vi.fn();
+    const mockEnd = vi.fn();
     const mockSetTransactionName = vi.fn();
 
     (core.getIsolationScope as any).mockReturnValue({ setTransactionName: mockSetTransactionName });
-    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn({ setStatus: mockSetStatus }));
+    (core.startSpanManual as any).mockImplementation((_: any, fn: any) =>
+      fn({ setStatus: mockSetStatus, end: mockEnd }),
+    );
 
     const wrappedFn = wrapServerFunction('testFunction', mockServerFn);
 
@@ -136,6 +139,7 @@ describe('wrapServerFunction', () => {
         },
       },
     });
+    expect(mockEnd).toHaveBeenCalled();
     expect(core.flushIfServerless).toHaveBeenCalled();
   });
 
@@ -146,31 +150,39 @@ describe('wrapServerFunction', () => {
     });
     const mockServerFn = vi.fn().mockRejectedValue(redirectResponse);
     const mockSetStatus = vi.fn();
+    const mockEnd = vi.fn();
     const mockSetTransactionName = vi.fn();
 
     (core.getIsolationScope as any).mockReturnValue({ setTransactionName: mockSetTransactionName });
-    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn({ setStatus: mockSetStatus }));
+    (core.startSpanManual as any).mockImplementation((_: any, fn: any) =>
+      fn({ setStatus: mockSetStatus, end: mockEnd }),
+    );
 
     const wrappedFn = wrapServerFunction('testFunction', mockServerFn);
 
     await expect(wrappedFn()).rejects.toBe(redirectResponse);
     expect(mockSetStatus).toHaveBeenCalledWith({ code: core.SPAN_STATUS_OK });
+    expect(mockEnd).toHaveBeenCalled();
     expect(core.captureException).not.toHaveBeenCalled();
   });
 
   it('should not capture not-found errors as exceptions', async () => {
     const notFoundResponse = new Response(null, { status: 404 });
     const mockServerFn = vi.fn().mockRejectedValue(notFoundResponse);
     const mockSetStatus = vi.fn();
+    const mockEnd = vi.fn();
     const mockSetTransactionName = vi.fn();
 
     (core.getIsolationScope as any).mockReturnValue({ setTransactionName: mockSetTransactionName });
-    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn({ setStatus: mockSetStatus }));
+    (core.startSpanManual as any).mockImplementation((_: any, fn: any) =>
+      fn({ setStatus: mockSetStatus, end: mockEnd }),
+    );
 
     const wrappedFn = wrapServerFunction('testFunction', mockServerFn);
 
     await expect(wrappedFn()).rejects.toBe(notFoundResponse);
     expect(mockSetStatus).toHaveBeenCalledWith({ code: core.SPAN_STATUS_ERROR, message: 'not_found' });
+    expect(mockEnd).toHaveBeenCalled();
     expect(core.captureException).not.toHaveBeenCalled();
   });
 
@@ -180,7 +192,7 @@ describe('wrapServerFunction', () => {
     const mockSetTransactionName = vi.fn();
 
     (core.getIsolationScope as any).mockReturnValue({ setTransactionName: mockSetTransactionName });
-    (core.startSpan as any).mockImplementation((_: any, fn: any) => fn({ setStatus: vi.fn() }));
+    (core.startSpanManual as any).mockImplementation((_: any, fn: any) => fn({ setStatus: vi.fn(), end: vi.fn() }));
 
     const wrappedFn = wrapServerFunction('testFunction', mockServerFn);