fix(tanstackstart-react): Flush events in server entry point for serverless environments

packages/tanstackstart-react/src/server/wrapFetchWithSentry.ts

@@ -1,3 +1,4 @@
+import { flushIfServerless } from '@sentry/core';
 import { SEMANTIC_ATTRIBUTE_SENTRY_OP, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, startSpan } from '@sentry/node';
 import { extractServerFunctionSha256 } from './utils';
 
@@ -32,35 +33,41 @@ export type ServerEntry = {
 export function wrapFetchWithSentry(serverEntry: ServerEntry): ServerEntry {
   if (serverEntry.fetch) {
     serverEntry.fetch = new Proxy<typeof serverEntry.fetch>(serverEntry.fetch, {
-      apply: (target, thisArg, args) => {
-        const request: Request = args[0];
-        const url = new URL(request.url);
-        const method = request.method || 'GET';
+      async apply(target, thisArg, args) {
+        try {
+          const request: Request = args[0];
+          const url = new URL(request.url);
+          const method = request.method || 'GET';
 
-        // instrument server functions
-        if (url.pathname.includes('_serverFn') || url.pathname.includes('createServerFn')) {
-          const functionSha256 = extractServerFunctionSha256(url.pathname);
-          const op = 'function.tanstackstart';
+          // instrument server functions
+          if (url.pathname.includes('_serverFn') || url.pathname.includes('createServerFn')) {
+            const functionSha256 = extractServerFunctionSha256(url.pathname);
+            const op = 'function.tanstackstart';
 
-          const serverFunctionSpanAttributes = {
-            [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.tanstackstart.server',
-            [SEMANTIC_ATTRIBUTE_SENTRY_OP]: op,
-            'tanstackstart.function.hash.sha256': functionSha256,
-          };
+            const serverFunctionSpanAttributes = {
+              [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.tanstackstart.server',
+              [SEMANTIC_ATTRIBUTE_SENTRY_OP]: op,
+              'tanstackstart.function.hash.sha256': functionSha256,
+            };
 
-          return startSpan(
-            {
-              op: op,
-              name: `${method} ${url.pathname}`,
-              attributes: serverFunctionSpanAttributes,
-            },
-            () => {
-              return target.apply(thisArg, args);
-            },
-          );
-        }
+            // eslint-disable-next-line no-return-await
+            return await startSpan(
+              {
+                op: op,
+                name: `${method} ${url.pathname}`,
+                attributes: serverFunctionSpanAttributes,
+              },
+              async () => {
+                return target.apply(thisArg, args);
+              },
+            );
+          }
 
-        return target.apply(thisArg, args);
+          // eslint-disable-next-line no-return-await
+          return await target.apply(thisArg, args);
+        } finally {
+          await flushIfServerless();
+        }
       },
     });
   }

packages/tanstackstart-react/test/server/wrapFetchWithSentry.test.ts

@@ -0,0 +1,66 @@
+import { afterEach, describe, expect, it, vi } from 'vitest';
+
+const startSpanSpy = vi.fn((_, callback) => callback());
+const flushIfServerlessSpy = vi.fn().mockResolvedValue(undefined);
+
+vi.mock('@sentry/node', async importOriginal => {
+  const original = await importOriginal();
+  return {
+    ...original,
+    startSpan: (...args: unknown[]) => startSpanSpy(...args),
+  };
+});
+
+vi.mock('@sentry/core', async importOriginal => {
+  const original = await importOriginal();
+  return {
+    ...original,
+    flushIfServerless: (...args: unknown[]) => flushIfServerlessSpy(...args),
+  };
+});
+
+// Import after mocks are set up
+const { wrapFetchWithSentry } = await import('../../src/server/wrapFetchWithSentry');
+
+describe('wrapFetchWithSentry', () => {
+  afterEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it('calls flushIfServerless after a regular request', async () => {
+    const mockResponse = new Response('ok');
+    const fetchFn = vi.fn().mockResolvedValue(mockResponse);
+
+    const serverEntry = wrapFetchWithSentry({ fetch: fetchFn });
+    const request = new Request('http://localhost:3000/page');
+
+    await serverEntry.fetch(request);
+
+    expect(fetchFn).toHaveBeenCalled();
+    expect(flushIfServerlessSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it('calls flushIfServerless after a server function request', async () => {
+    const mockResponse = new Response('ok');
+    const fetchFn = vi.fn().mockResolvedValue(mockResponse);
+
+    const serverEntry = wrapFetchWithSentry({ fetch: fetchFn });
+    const request = new Request('http://localhost:3000/_serverFn/abc123');
+
+    await serverEntry.fetch(request);
+
+    expect(startSpanSpy).toHaveBeenCalled();
+    expect(flushIfServerlessSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it('calls flushIfServerless even if the handler throws', async () => {
+    const fetchFn = vi.fn().mockRejectedValue(new Error('handler error'));
+
+    const serverEntry = wrapFetchWithSentry({ fetch: fetchFn });
+    const request = new Request('http://localhost:3000/page');
+
+    await expect(serverEntry.fetch(request)).rejects.toThrow('handler error');
+
+    expect(flushIfServerlessSpy).toHaveBeenCalledTimes(1);
+  });
+});