refactor

Author: nicohrubec

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

@@ -42,6 +42,24 @@ export function wrapGlobalMiddleware(code: string, id: string, debug: boolean):
   return { code: transformed, didWrap, skipped };
 }
 
+/**
+ * Adds the wrapMiddlewaresWithSentry import to the code.
+ * Handles 'use client' and 'use server' directives by inserting the import after them.
+ */
+export function addSentryImport(code: string): string {
+  const sentryImport = "import { wrapMiddlewaresWithSentry } from '@sentry/tanstackstart-react';\n";
+
+  // Check for 'use server' or 'use client' directives, these need to be before any imports
+  const directiveMatch = code.match(/^(['"])use (client|server)\1;?\s*\n?/);
+
+  if (!directiveMatch) {
+    return sentryImport + code;
+  }
+
+  const directive = directiveMatch[0];
+  return directive + sentryImport + code.slice(directive.length);
+}
+
 /**
  * Wraps route middleware arrays in createFileRoute() files.
  */
@@ -140,17 +158,7 @@ export function makeAutoInstrumentMiddlewarePlugin(options: AutoInstrumentMiddle
         return null;
       }
 
-      const sentryImport = "import { wrapMiddlewaresWithSentry } from '@sentry/tanstackstart-react';\n";
-
-      // Check for 'use server' or 'use client' directives, these need to be before any imports
-      const directiveMatch = transformed.match(/^(['"])use (client|server)\1;?\s*\n?/);
-      if (directiveMatch) {
-        // Insert import after the directive
-        const directive = directiveMatch[0];
-        transformed = directive + sentryImport + transformed.slice(directive.length);
-      } else {
-        transformed = sentryImport + transformed;
-      }
+      transformed = addSentryImport(transformed);
 
       return { code: transformed, map: null };
     },

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

@@ -1,6 +1,7 @@
 import type { Plugin } from 'vite';
 import { describe, expect, it, vi } from 'vitest';
 import {
+  addSentryImport,
   arrayToObjectShorthand,
   makeAutoInstrumentMiddlewarePlugin,
   wrapGlobalMiddleware,
@@ -61,30 +62,6 @@ createStart(() => ({ requestMiddleware: wrapMiddlewaresWithSentry({ myMiddleware
     expect(result).toBeNull();
   });
 
-  it('handles files with use server directive', () => {
-    const plugin = makeAutoInstrumentMiddlewarePlugin() as PluginWithTransform;
-    const code = `'use server';
-import { createStart } from '@tanstack/react-start';
-createStart(() => ({ requestMiddleware: [authMiddleware] }));
-`;
-    const result = plugin.transform(code, '/app/start.ts');
-
-    expect(result).not.toBeNull();
-    expect(result!.code).toMatch(/^'use server';\s*\nimport \{ wrapMiddlewaresWithSentry \}/);
-  });
-
-  it('handles files with use client directive', () => {
-    const plugin = makeAutoInstrumentMiddlewarePlugin() as PluginWithTransform;
-    const code = `"use client";
-import { createStart } from '@tanstack/react-start';
-createStart(() => ({ requestMiddleware: [authMiddleware] }));
-`;
-    const result = plugin.transform(code, '/app/start.ts');
-
-    expect(result).not.toBeNull();
-    expect(result!.code).toMatch(/^"use client";\s*\nimport \{ wrapMiddlewaresWithSentry \}/);
-  });
-
   it('adds import statement when wrapping middlewares', () => {
     const plugin = makeAutoInstrumentMiddlewarePlugin() as PluginWithTransform;
     const result = plugin.transform(createStartFile, '/app/start.ts');
@@ -352,6 +329,33 @@ export const Route = createFileRoute('/foo')({
   });
 });
 
+describe('addSentryImport', () => {
+  it('prepends import to code without directives', () => {
+    const code = 'const foo = 1;';
+    const result = addSentryImport(code);
+
+    expect(result).toBe(
+      "import { wrapMiddlewaresWithSentry } from '@sentry/tanstackstart-react';\nconst foo = 1;",
+    );
+  });
+
+  it('inserts import after use server directive', () => {
+    const code = "'use server';\nconst foo = 1;";
+    const result = addSentryImport(code);
+
+    expect(result).toMatch(/^'use server';\nimport \{ wrapMiddlewaresWithSentry \}/);
+    expect(result).toContain('const foo = 1;');
+  });
+
+  it('inserts import after use client directive', () => {
+    const code = '"use client";\nconst foo = 1;';
+    const result = addSentryImport(code);
+
+    expect(result).toMatch(/^"use client";\nimport \{ wrapMiddlewaresWithSentry \}/);
+    expect(result).toContain('const foo = 1;');
+  });
+});
+
 describe('arrayToObjectShorthand', () => {
   it('converts single identifier', () => {
     expect(arrayToObjectShorthand('foo')).toBe('{ foo }');