feat(tanstackstart-react): Auto-instrument server function middleware

Author: nicohrubec

dev-packages/e2e-tests/test-applications/tanstackstart-react/src/middleware.ts

@@ -1,5 +1,4 @@
 import { createMiddleware } from '@tanstack/react-start';
-import { wrapMiddlewaresWithSentry } from '@sentry/tanstackstart-react';
 
 // Global request middleware - runs on every request
 // NOTE: This is exported unwrapped to test auto-instrumentation via the Vite plugin
@@ -15,8 +14,8 @@ export const globalFunctionMiddleware = createMiddleware({ type: 'function' }).s
   return next();
 });
 
-// Server function middleware
-const serverFnMiddleware = createMiddleware({ type: 'function' }).server(async ({ next }) => {
+// Server function middleware - exported unwrapped for auto-instrumentation via Vite plugin
+export const serverFnMiddleware = createMiddleware({ type: 'function' }).server(async ({ next }) => {
   console.log('Server function middleware executed');
   return next();
 });
@@ -28,21 +27,15 @@ export const serverRouteRequestMiddleware = createMiddleware().server(async ({ n
 });
 
 // Early return middleware - returns without calling next()
-const earlyReturnMiddleware = createMiddleware({ type: 'function' }).server(async () => {
+// Exported unwrapped for auto-instrumentation via Vite plugin
+export const earlyReturnMiddleware = createMiddleware({ type: 'function' }).server(async () => {
   console.log('Early return middleware executed - not calling next()');
   return { earlyReturn: true, message: 'Middleware returned early without calling next()' };
 });
 
 // Error middleware - throws an exception
-const errorMiddleware = createMiddleware({ type: 'function' }).server(async () => {
+// Exported unwrapped for auto-instrumentation via Vite plugin
+export const errorMiddleware = createMiddleware({ type: 'function' }).server(async () => {
   console.log('Error middleware executed - throwing error');
   throw new Error('Middleware Error Test');
 });
-
-// Manually wrap middlewares with Sentry (for middlewares that won't be auto-instrumented)
-export const [wrappedServerFnMiddleware, wrappedEarlyReturnMiddleware, wrappedErrorMiddleware] =
-  wrapMiddlewaresWithSentry({
-    serverFnMiddleware,
-    earlyReturnMiddleware,
-    errorMiddleware,
-  });

dev-packages/e2e-tests/test-applications/tanstackstart-react/src/routes/test-middleware.tsx

@@ -1,10 +1,10 @@
 import { createFileRoute } from '@tanstack/react-router';
 import { createServerFn } from '@tanstack/react-start';
-import { wrappedServerFnMiddleware, wrappedEarlyReturnMiddleware, wrappedErrorMiddleware } from '../middleware';
+import { serverFnMiddleware, earlyReturnMiddleware, errorMiddleware } from '../middleware';
 
 // Server function with specific middleware (also gets global function middleware)
 const serverFnWithMiddleware = createServerFn()
-  .middleware([wrappedServerFnMiddleware])
+  .middleware([serverFnMiddleware])
   .handler(async () => {
     console.log('Server function with specific middleware executed');
     return { message: 'Server function middleware test' };
@@ -18,15 +18,15 @@ const serverFnWithoutMiddleware = createServerFn().handler(async () => {
 
 // Server function with early return middleware (middleware returns without calling next)
 const serverFnWithEarlyReturnMiddleware = createServerFn()
-  .middleware([wrappedEarlyReturnMiddleware])
+  .middleware([earlyReturnMiddleware])
   .handler(async () => {
     console.log('This should not be executed - middleware returned early');
     return { message: 'This should not be returned' };
   });
 
 // Server function with error middleware (middleware throws an error)
 const serverFnWithErrorMiddleware = createServerFn()
-  .middleware([wrappedErrorMiddleware])
+  .middleware([errorMiddleware])
   .handler(async () => {
     console.log('This should not be executed - middleware threw error');
     return { message: 'This should not be returned' };

packages/tanstackstart-react/src/vite/autoInstrumentMiddleware.ts

@@ -26,6 +26,10 @@ function wrapMiddlewareArrays(code: string, id: string, debug: boolean, regex: R
         // eslint-disable-next-line no-console
         console.log(`[Sentry] Auto-wrapping ${key} in ${id}`);
       }
+      // Handle method call syntax like `.middleware([...])` vs object property syntax like `middleware: [...]`
+      if (key.endsWith('(')) {
+        return `${key}wrapMiddlewaresWithSentry(${objContents}))`;
+      }
       return `${key}: wrapMiddlewaresWithSentry(${objContents})`;
     }
     // Track middlewares that couldn't be auto-wrapped
@@ -53,6 +57,13 @@ export function wrapRouteMiddleware(code: string, id: string, debug: boolean): W
   return wrapMiddlewareArrays(code, id, debug, /(middleware)\s*:\s*\[([^\]]*)\]/g);
 }
 
+/**
+ * Wraps middleware arrays in createServerFn().middleware([...]) calls.
+ */
+export function wrapServerFnMiddleware(code: string, id: string, debug: boolean): WrapResult {
+  return wrapMiddlewareArrays(code, id, debug, /(\.middleware\s*\()\s*\[([^\]]*)\]\s*\)/g);
+}
+
 /**
  * A Vite plugin that automatically instruments TanStack Start middlewares:
  * - `requestMiddleware` and `functionMiddleware` arrays in `createStart()`
@@ -78,8 +89,9 @@ export function makeAutoInstrumentMiddlewarePlugin(options: AutoInstrumentMiddle
       // Detect file types that should be instrumented
       const isStartFile = id.includes('start') && code.includes('createStart(');
       const isRouteFile = code.includes('createFileRoute(') && /middleware\s*:\s*\[/.test(code);
+      const isServerFnFile = code.includes('createServerFn') && /\.middleware\s*\(\s*\[/.test(code);
 
-      if (!isStartFile && !isRouteFile) {
+      if (!isStartFile && !isRouteFile && !isServerFnFile) {
         return null;
       }
 
@@ -101,12 +113,20 @@ export function makeAutoInstrumentMiddlewarePlugin(options: AutoInstrumentMiddle
           skippedMiddlewares.push(...result.skipped);
           break;
         }
-        // route middleware
-        case isRouteFile: {
-          const result = wrapRouteMiddleware(transformed, id, debug);
-          transformed = result.code;
-          needsImport = needsImport || result.didWrap;
-          skippedMiddlewares.push(...result.skipped);
+        // non-global middleware
+        case (isRouteFile || isServerFnFile): {
+          if (isRouteFile) {
+            const result = wrapRouteMiddleware(transformed, id, debug);
+            transformed = result.code;
+            needsImport = needsImport || result.didWrap;
+            skippedMiddlewares.push(...result.skipped);
+          }
+          if (isServerFnFile) {
+            const result = wrapServerFnMiddleware(transformed, id, debug);
+            transformed = result.code;
+            needsImport = needsImport || result.didWrap;
+            skippedMiddlewares.push(...result.skipped);
+          }
           break;
         }
         default:

packages/tanstackstart-react/test/vite/autoInstrumentMiddleware.test.ts

@@ -6,6 +6,7 @@ import {
   makeAutoInstrumentMiddlewarePlugin,
   wrapGlobalMiddleware,
   wrapRouteMiddleware,
+  wrapServerFnMiddleware,
 } from '../../src/vite/autoInstrumentMiddleware';
 
 type PluginWithTransform = Plugin & {
@@ -329,6 +330,86 @@ export const Route = createFileRoute('/foo')({
   });
 });
 
+describe('wrapServerFnMiddleware', () => {
+  it('wraps single middleware in createServerFn().middleware()', () => {
+    const code = `
+const serverFn = createServerFn()
+  .middleware([authMiddleware])
+  .handler(async () => ({}));
+`;
+    const result = wrapServerFnMiddleware(code, '/app/routes/foo.ts', false);
+
+    expect(result.didWrap).toBe(true);
+    expect(result.code).toContain('.middleware(wrapMiddlewaresWithSentry({ authMiddleware }))');
+    expect(result.skipped).toHaveLength(0);
+  });
+
+  it('wraps multiple middlewares in createServerFn().middleware()', () => {
+    const code = `
+const serverFn = createServerFn()
+  .middleware([authMiddleware, loggingMiddleware])
+  .handler(async () => ({}));
+`;
+    const result = wrapServerFnMiddleware(code, '/app/routes/foo.ts', false);
+
+    expect(result.didWrap).toBe(true);
+    expect(result.code).toContain('.middleware(wrapMiddlewaresWithSentry({ authMiddleware, loggingMiddleware }))');
+  });
+
+  it('does not wrap empty middleware arrays', () => {
+    const code = `
+const serverFn = createServerFn()
+  .middleware([])
+  .handler(async () => ({}));
+`;
+    const result = wrapServerFnMiddleware(code, '/app/routes/foo.ts', false);
+
+    expect(result.didWrap).toBe(false);
+    expect(result.skipped).toHaveLength(0);
+  });
+
+  it('does not wrap middleware containing function calls', () => {
+    const code = `
+const serverFn = createServerFn()
+  .middleware([createMiddleware()])
+  .handler(async () => ({}));
+`;
+    const result = wrapServerFnMiddleware(code, '/app/routes/foo.ts', false);
+
+    expect(result.didWrap).toBe(false);
+    expect(result.skipped).toContain('.middleware(');
+  });
+
+  it('handles multiple server functions in same file', () => {
+    const code = `
+const serverFn1 = createServerFn()
+  .middleware([authMiddleware])
+  .handler(async () => ({}));
+
+const serverFn2 = createServerFn()
+  .middleware([loggingMiddleware])
+  .handler(async () => ({}));
+`;
+    const result = wrapServerFnMiddleware(code, '/app/routes/foo.ts', false);
+
+    expect(result.didWrap).toBe(true);
+    expect(result.code).toContain('.middleware(wrapMiddlewaresWithSentry({ authMiddleware }))');
+    expect(result.code).toContain('.middleware(wrapMiddlewaresWithSentry({ loggingMiddleware }))');
+  });
+
+  it('handles trailing commas in middleware arrays', () => {
+    const code = `
+const serverFn = createServerFn()
+  .middleware([authMiddleware,])
+  .handler(async () => ({}));
+`;
+    const result = wrapServerFnMiddleware(code, '/app/routes/foo.ts', false);
+
+    expect(result.didWrap).toBe(true);
+    expect(result.code).toContain('.middleware(wrapMiddlewaresWithSentry({ authMiddleware }))');
+  });
+});
+
 describe('addSentryImport', () => {
   it('prepends import to code without directives', () => {
     const code = 'const foo = 1;';