add flushing

Author: chargome

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

@@ -2,11 +2,13 @@ import { context } from '@opentelemetry/api';
 import { getRPCMetadata, RPCType } from '@opentelemetry/core';
 import { ATTR_HTTP_ROUTE } from '@opentelemetry/semantic-conventions';
 import {
+  flush,
   getActiveSpan,
   getRootSpan,
   getTraceMetaTags,
   SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
   SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
+  vercelWaitUntil,
 } from '@sentry/core';
 import type { AppLoadContext, EntryContext } from 'react-router';
 import type { PassThrough } from 'stream';
@@ -58,11 +60,16 @@ export function wrapSentryHandleRequest(originalHandle: OriginalHandleRequest):
       });
     }
 
-    return originalHandle(request, responseStatusCode, responseHeaders, routerContext, loadContext);
+    try {
+      return await originalHandle(request, responseStatusCode, responseHeaders, routerContext, loadContext);
+    } finally {
+      vercelWaitUntil(flush());
+    }
   };
 }
 
 /** @deprecated Use `wrapSentryHandleRequest` instead. */
+// todo(v11): remove this
 export const sentryHandleRequest = wrapSentryHandleRequest;
 
 /**

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

@@ -1,11 +1,13 @@
 import { RPCType } from '@opentelemetry/core';
 import { ATTR_HTTP_ROUTE } from '@opentelemetry/semantic-conventions';
 import {
+  flush,
   getActiveSpan,
   getRootSpan,
   getTraceMetaTags,
   SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
   SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
+  vercelWaitUntil,
 } from '@sentry/core';
 import { PassThrough } from 'stream';
 import { beforeEach, describe, expect, test, vi } from 'vitest';
@@ -22,13 +24,66 @@ vi.mock('@sentry/core', () => ({
   getActiveSpan: vi.fn(),
   getRootSpan: vi.fn(),
   getTraceMetaTags: vi.fn(),
+  flush: vi.fn(),
+  vercelWaitUntil: vi.fn(),
 }));
 
 describe('wrapSentryHandleRequest', () => {
   beforeEach(() => {
     vi.clearAllMocks();
   });
 
+  test('should call flush on successful execution', async () => {
+    const originalHandler = vi.fn().mockResolvedValue('success response');
+    const wrappedHandler = wrapSentryHandleRequest(originalHandler);
+
+    const request = new Request('https://example.com');
+    const responseStatusCode = 200;
+    const responseHeaders = new Headers();
+    const routerContext = { staticHandlerContext: { matches: [] } } as any;
+    const loadContext = {} as any;
+
+    await wrappedHandler(request, responseStatusCode, responseHeaders, routerContext, loadContext);
+
+    expect(vercelWaitUntil).toHaveBeenCalledWith(flush());
+    expect(flush).toHaveBeenCalled();
+  });
+
+  test('should call flush even when original handler throws an error', async () => {
+    const mockError = new Error('Handler failed');
+    const originalHandler = vi.fn().mockRejectedValue(mockError);
+    const wrappedHandler = wrapSentryHandleRequest(originalHandler);
+
+    const request = new Request('https://example.com');
+    const responseStatusCode = 200;
+    const responseHeaders = new Headers();
+    const routerContext = { staticHandlerContext: { matches: [] } } as any;
+    const loadContext = {} as any;
+
+    await expect(
+      wrappedHandler(request, responseStatusCode, responseHeaders, routerContext, loadContext),
+    ).rejects.toThrow('Handler failed');
+
+    expect(vercelWaitUntil).toHaveBeenCalledWith(flush());
+    expect(flush).toHaveBeenCalled();
+  });
+
+  test('should propagate errors from original handler', async () => {
+    const mockError = new Error('Test error');
+    const originalHandler = vi.fn().mockRejectedValue(mockError);
+    const wrappedHandler = wrapSentryHandleRequest(originalHandler);
+
+    const request = new Request('https://example.com');
+    const responseStatusCode = 500;
+    const responseHeaders = new Headers();
+    const routerContext = { staticHandlerContext: { matches: [] } } as any;
+    const loadContext = {} as any;
+
+    await expect(wrappedHandler(request, responseStatusCode, responseHeaders, routerContext, loadContext)).rejects.toBe(
+      mockError,
+    );
+  });
+
   test('should call original handler with same parameters', async () => {
     const originalHandler = vi.fn().mockResolvedValue('original response');
     const wrappedHandler = wrapSentryHandleRequest(originalHandler);
@@ -62,7 +117,8 @@ describe('wrapSentryHandleRequest', () => {
     (getActiveSpan as unknown as ReturnType<typeof vi.fn>).mockReturnValue(mockActiveSpan);
     (getRootSpan as unknown as ReturnType<typeof vi.fn>).mockReturnValue(mockRootSpan);
     const getRPCMetadata = vi.fn().mockReturnValue(mockRpcMetadata);
-    vi.mocked(vi.importActual('@opentelemetry/core')).getRPCMetadata = getRPCMetadata;
+    (vi.importActual('@opentelemetry/core') as unknown as { getRPCMetadata: typeof getRPCMetadata }).getRPCMetadata =
+      getRPCMetadata;
 
     const routerContext = {
       staticHandlerContext: {
@@ -110,7 +166,8 @@ describe('wrapSentryHandleRequest', () => {
     (getActiveSpan as unknown as ReturnType<typeof vi.fn>).mockReturnValue(null);
 
     const getRPCMetadata = vi.fn().mockReturnValue(mockRpcMetadata);
-    vi.mocked(vi.importActual('@opentelemetry/core')).getRPCMetadata = getRPCMetadata;
+    (vi.importActual('@opentelemetry/core') as unknown as { getRPCMetadata: typeof getRPCMetadata }).getRPCMetadata =
+      getRPCMetadata;
 
     const routerContext = {
       staticHandlerContext: {
@@ -132,7 +189,7 @@ describe('getMetaTagTransformer', () => {
     );
   });
 
-  test('should inject meta tags before closing head tag', done => {
+  test('should inject meta tags before closing head tag', () => {
     const outputStream = new PassThrough();
     const bodyStream = new PassThrough();
     const transformer = getMetaTagTransformer(bodyStream);
@@ -145,7 +202,6 @@ describe('getMetaTagTransformer', () => {
     outputStream.on('end', () => {
       expect(outputData).toContain('<meta name="sentry-trace" content="test-trace-id"></head>');
       expect(outputData).not.toContain('</head></head>');
-      done();
     });
 
     transformer.pipe(outputStream);
@@ -154,7 +210,7 @@ describe('getMetaTagTransformer', () => {
     bodyStream.end();
   });
 
-  test('should not modify chunks without head closing tag', done => {
+  test('should not modify chunks without head closing tag', () => {
     const outputStream = new PassThrough();
     const bodyStream = new PassThrough();
     const transformer = getMetaTagTransformer(bodyStream);
@@ -167,7 +223,6 @@ describe('getMetaTagTransformer', () => {
     outputStream.on('end', () => {
       expect(outputData).toBe('<html><body>Test</body></html>');
       expect(getTraceMetaTags).toHaveBeenCalled();
-      done();
     });
 
     transformer.pipe(outputStream);
@@ -176,7 +231,7 @@ describe('getMetaTagTransformer', () => {
     bodyStream.end();
   });
 
-  test('should handle buffer input', done => {
+  test('should handle buffer input', () => {
     const outputStream = new PassThrough();
     const bodyStream = new PassThrough();
     const transformer = getMetaTagTransformer(bodyStream);
@@ -188,7 +243,6 @@ describe('getMetaTagTransformer', () => {
 
     outputStream.on('end', () => {
       expect(outputData).toContain('<meta name="sentry-trace" content="test-trace-id"></head>');
-      done();
     });
 
     transformer.pipe(outputStream);