vite-plugin-react-server 1.2.0-alpha.0 → 1.2.0

package/plugin/loader/transformModuleIfNeeded.ts

@@ -1,7 +1,7 @@
-import { transformWithAcornLoose } from "./transformWithAcornLoose.js";
 import type { RawSourceMap } from "source-map";
+import { createTransformer } from "./createTransformer.js";
 import type { TransformOptions } from "./types.js";
-
+import { parse } from "./parse.js";
 
 // --- React RSC Directive Handling ---
 //
@@ -36,30 +36,27 @@ import type { TransformOptions } from "./types.js";
 /**
  * Transforms a module and returns the transformed code with source map attached.
  * This is used by the loader to transform modules and attach source maps.
- * 
+ *
  * @returns The transformed code with source map attached as a URL comment
  */
-export async function transformModuleIfNeeded(
+export type TransformModuleIfNeededFn = (
   source: string,
   moduleId: string,
   options: TransformOptions
-): Promise<{ code: string; map: RawSourceMap | null }> {
-  try {
-    // Get transformed code and source map from acorn-loose transformer
-    const result = await transformWithAcornLoose(
-      source,
-      moduleId,
-      options
-    );
+) => Promise<{ code: string; map: RawSourceMap | null }>;
 
-    if(options.verbose) {
-      console.log("[transformModuleIfNeeded] Transformed module:", { code: result.code, map: result.map });
-    }
-    // Return the transformed code with source map attached as a URL comment
-    return result;
-  } catch (error) {
-    // Log the error and rethrow
-    console.error(`[transformModuleIfNeeded] Error transforming module ${moduleId}:`, error);
-    throw error;
-  }
-}
+/**
+ * Transforms a module using acorn-loose for parsing.
+ * @returns Object containing the transformed code and its source map
+ */
+export const transformModuleIfNeeded: TransformModuleIfNeededFn =
+  async function _transformModuleIfNeeded(source, moduleId, options) {
+    const transformer = createTransformer({
+      parseFn:
+        typeof options.loader?.parse === "function"
+          ? options.loader.parse
+          : parse,
+      options,
+    });
+    return transformer(source, moduleId);
+  };

package/plugin/loader/transformNonServerEnvironment.ts

@@ -0,0 +1,103 @@
+import type { ParseResult } from "./directives/types.js";
+import { createSourceMap } from "./sourceMap.js";
+import type { LoaderConfig, TransformResult } from "./types.js";
+import { removeDirectives } from "./removeDirectives.js";
+import * as acorn from "acorn";
+import { getNodeEnv } from "../config/getNodeEnv.js";
+import { DEFAULT_CONFIG } from "../config/defaults.js";
+/**
+ * Transforms a server module by:
+ * 1. Parsing it to get exports and directives
+ * 2. Collecting function-level server directives
+ * 3. Registering server functions
+ */
+export async function transformNonServerEnvironment(
+  source: string,
+  moduleId: string,
+  transformedModuleId: string,
+  parseResult: ParseResult,
+  loader: Pick<LoaderConfig, "parse" | "verbose" | "logger"> = DEFAULT_CONFIG.RSC_LOADER[getNodeEnv()],
+): Promise<TransformResult> {
+  if (parseResult.type !== "success") {
+    return { code: "", map: null };
+  }
+
+  // Parse the source with Acorn to get accurate directive locations
+  let ast;
+  if (loader.parse) {
+    const parsed = loader.parse(source);
+    if (parsed instanceof Promise) {
+      ast = await parsed;
+    } else {
+      ast = parsed;
+    }
+  }
+  if (typeof ast === "object" && "ast" in ast) {
+    ast = ast.ast;
+  }
+  if (!ast) {
+    ast = acorn.parse(source, {
+      ecmaVersion: "latest",
+      sourceType: "module",
+      locations: true,
+      ranges: true,
+      allowAwaitOutsideFunction: true,
+      allowReturnOutsideFunction: true,
+    });
+  }
+
+  // check if body is iterable
+  if (!(ast.body && typeof ast.body === "object" && "forEach" in ast.body)) {
+    throw new Error(
+      `[transformServerModule] Failed to parse ${moduleId} -> ${transformedModuleId} with loader.parse`
+    );
+  }
+
+  // Collect all directive ranges (file-level and function-level)
+  const rangesToRemove: { start: number; end: number }[] = [];
+
+  // File-level: top-level ExpressionStatement with directive
+  for (const node of ast.body) {
+    if (node.type === "ExpressionStatement") {
+      if (node.directive === "use client") {
+        // Remove "use client" directives in client environment
+        rangesToRemove.push({ start: node.start, end: node.end });
+      } else if (node.directive === "use server") {
+        // In SSR/non-react-server environment, we need to handle server actions
+        // For now, just remove the directive and let the functions execute directly
+        // This works because SSR is still server-side (Node.js), not browser
+        // The server actions will work as regular async functions
+        rangesToRemove.push({ start: node.start, end: node.end });
+        if (loader.verbose) {
+          loader.logger?.info(
+            `[transformNonServerEnvironment] Allowing "use server" in SSR mode: ${moduleId}`
+          );
+        }
+      }
+    }
+    // Stop at first non-directive
+    if (node.type !== "ExpressionStatement" || !node.directive) break;
+  }
+
+  if (loader.verbose) {
+    loader.logger?.info(
+      `[transformNonServerEnvironment] Ranges to remove: ${rangesToRemove
+        .map((r) => `${r.start} -> ${r.end}`)
+        .join(", ")}`
+    );
+  }
+
+  // Remove directives from the source code using the shared utility
+  const transformedCode = removeDirectives(source, rangesToRemove);
+
+  // Only add the import and registrations if there are any registrations to make
+  const finalCode = transformedCode;
+
+  // Create source map based on the final transformed code
+  const map = createSourceMap(finalCode, source, transformedModuleId, []);
+
+  return {
+    code: finalCode,
+    map,
+  };
+}

package/plugin/loader/transformServerModule.ts

@@ -1,12 +1,15 @@
 import type { ParseResult, Program } from "./directives/types.js";
-import type {  ResolvedUserOptions } from "../types.js";
 import { createSourceMap } from "./sourceMap.js";
-import type { TransformResult } from "./types.js";
+import type { LoaderConfig, TransformResult } from "./types.js";
 import { removeDirectives } from "./removeDirectives.js";
 import * as acorn from "acorn";
-import { getNodeEnv } from "../getNodeEnv.js";
+import { getNodeEnv } from "../config/getNodeEnv.js";
 import { DEFAULT_CONFIG } from "../config/defaults.js";
-import type { ArrowFunctionExpression, FunctionDeclaration, FunctionExpression } from "acorn";
+import type {
+  ArrowFunctionExpression,
+  FunctionDeclaration,
+  FunctionExpression,
+} from "acorn";
 
 /**
  * Transforms a server module by:
@@ -17,12 +20,20 @@ import type { ArrowFunctionExpression, FunctionDeclaration, FunctionExpression }
 export async function transformServerModule(
   source: string,
   moduleId: string,
+  transformedModuleId: string,
   parseResult: ParseResult,
   loader: Pick<
-    ResolvedUserOptions["loader"],
-    "registerServerReferenceName" | "importServerPath" 
+    LoaderConfig,
+    | "registerServerReferenceName"
+    | "importServerPath"
+    | "parse"
+    | "isClientComponentCode"
+    | "isClientComponentByName"
+    | "registerClientReferenceName"
+    | "importClientPath"
+    | "verbose"
+    | "logger"
   > = DEFAULT_CONFIG.RSC_LOADER[getNodeEnv()],
-  verbose = false
 ): Promise<TransformResult> {
   if (!loader) {
     loader = DEFAULT_CONFIG.RSC_LOADER[getNodeEnv()];
@@ -31,15 +42,37 @@ export async function transformServerModule(
     return { code: "", map: null };
   }
 
-  // Parse the source with Acorn to get accurate directive locations
-  const ast = acorn.parse(source, {
-    ecmaVersion: "latest",
-    sourceType: "module",
-    locations: true,
-    ranges: true,
-    allowAwaitOutsideFunction: true,
-    allowReturnOutsideFunction: true,
-  });
+  // Parse the source using the loader's parse function or fallback to Acorn
+  let ast;
+  if (typeof loader.parse === "function") {
+    const parsed = loader.parse(source);
+
+    if (parsed instanceof Promise) {
+      ast = await parsed;
+    } else {
+      ast = parsed;
+    }
+  }
+  if (typeof ast === "object" && "ast" in ast) {
+    ast = ast.ast;
+  }
+  if (!ast) {
+    ast = acorn.parse(source, {
+      ecmaVersion: "latest",
+      sourceType: "module",
+      locations: true,
+      ranges: true,
+      allowAwaitOutsideFunction: true,
+      allowReturnOutsideFunction: true,
+    });
+  }
+
+  // check if body is iterable
+  if (!(ast.body && typeof ast.body === "object" && "forEach" in ast.body)) {
+    throw new Error(
+      `[transformServerModule] Failed to parse ${moduleId} -> ${transformedModuleId} with loader.parse`
+    );
+  }
 
   // Collect all directive ranges (file-level and function-level)
   const rangesToRemove: { start: number; end: number }[] = [];
@@ -57,7 +90,13 @@ export async function transformServerModule(
   }
 
   // Function-level: walk all functions and check first statement in body
-  function walkFunctionDirectives(node: FunctionDeclaration | FunctionExpression | ArrowFunctionExpression | Program) {
+  function walkFunctionDirectives(
+    node:
+      | FunctionDeclaration
+      | FunctionExpression
+      | ArrowFunctionExpression
+      | Program
+  ) {
     if (
       (node.type === "FunctionDeclaration" ||
         node.type === "FunctionExpression" ||
@@ -80,54 +119,110 @@ export async function transformServerModule(
       const value = node[key as keyof typeof node];
       if (Array.isArray(value)) {
         value.forEach((item) => {
-          if (item && typeof item === 'object' && 'type' in item) {
+          if (item && typeof item === "object" && "type" in item) {
             walkFunctionDirectives(item as any);
           }
         });
-      } else if (value && typeof value === 'object' && 'type' in value) {
+      } else if (value && typeof value === "object" && "type" in value) {
         walkFunctionDirectives(value as any);
       }
     }
   }
   walkFunctionDirectives(ast);
 
-  if (verbose) {
-    console.log(`[transformServerModule] Ranges to remove:`, rangesToRemove);
+  if (loader.verbose) {
+    for (const range of rangesToRemove) {
+      loader.logger?.info(
+        `[transformServerModule] Ranges to remove: start: ${
+          range.start
+        }, end: ${range.end}, source: ${source.slice(range.start, range.end)}`
+      );
+    }
   }
 
   // Remove directives from the source code using the shared utility
   const transformedCode = removeDirectives(source, rangesToRemove);
 
   // Register all exports as server references
+  // If this function is called, it means the caller determined this is a server module
+  // (either via AST directive detection or regex-based fast path)
   const registrations = [];
+  const hasFileLevelServerDirective = parseResult.directiveInfo.fileLevel?.type === "server";
+  
   for (const exp of parseResult.exports.exports.values()) {
-    // Check if any exports have server directives
-    const hasServerDirective = parseResult.directiveInfo.functionLevel.some(
+    // Check if any exports have server directives at function level
+    const hasFunctionLevelDirective = parseResult.directiveInfo.functionLevel.some(
       (d) => d.type === "server" && d.name === exp.localName
     );
-    // Register if it has its own server directive or if there's a file-level directive
-    if (
-      hasServerDirective ||
-      parseResult.directiveInfo.fileLevel?.type === "server"
-    ) {
+    // Register if:
+    // 1. Has function-level server directive, OR
+    // 2. Has file-level server directive, OR  
+    // 3. No file-level directive detected (AST might miss pre-compiled directives)
+    //    In this case, we trust the caller's decision to call transformServerModule
+    const shouldRegister = hasFunctionLevelDirective || 
+                          hasFileLevelServerDirective || 
+                          !parseResult.directiveInfo.fileLevel;
+    
+    if (shouldRegister) {
       // Use original module ID for re-exports, current module ID for local exports
-      const targetModuleId = exp.originalModuleId || moduleId;
+      const targetModuleId = exp.originalModuleId || transformedModuleId;
       registrations.push(
         `${loader?.registerServerReferenceName}(${exp.localName}, "${targetModuleId}", "${exp.exportName}");`
       );
     }
   }
 
-  // Only add the import and registrations if there are any registrations to make
-  const finalCode =
-    registrations.length > 0
-      ? `\n      import { ${loader?.registerServerReferenceName} } from "${
-          loader?.importServerPath
-        }";\n      ${transformedCode}\n      ${registrations.join("\n")}\n    `
-      : transformedCode;
+  // Also handle client components in server environment
+  // Check if this is a client component by checking for client directive
+  const hasClientDirective =
+    parseResult.directiveInfo?.fileLevel?.type === "client";
+
+  let finalCode = transformedCode;
+  let imports = [];
+
+  // Add server reference imports and registrations if needed
+  if (registrations.length > 0) {
+    imports.push(
+      `import { ${loader?.registerServerReferenceName} } from "${loader?.importServerPath}";`
+    );
+    finalCode = `${finalCode}\n${registrations.join("\n")}`;
+  }
+
+  // Add client reference imports and registrations if this is a client component
+  if (hasClientDirective) {
+    imports.push(
+      `import { ${loader?.registerClientReferenceName} } from "${loader?.importClientPath}";`
+    );
+
+    
+    if (loader.verbose) {
+      loader.logger?.info(
+        `[transformServerModule] moduleID function called: ${moduleId} -> ${transformedModuleId}`
+      );
+    }
+
+    const clientRegistrations = [];
+    for (const exp of parseResult.exports.exports.values()) {
+      if (exp.exportName === "default") {
+        clientRegistrations.push(
+          `export default ${loader?.registerClientReferenceName}(function() { throw new Error("Attempted to call default() on the client"); }, "${transformedModuleId}", "default");`
+        );
+      } else {
+        clientRegistrations.push(
+          `export const ${exp.exportName} = ${loader?.registerClientReferenceName}(function() { throw new Error("Attempted to call ${exp.exportName}() on the client"); }, "${transformedModuleId}", "${exp.exportName}");`
+        );
+      }
+    }
+    finalCode = `${finalCode}\n${clientRegistrations.join("\n")}`;
+  }
+
+  // Add imports at the top if any
+  if (imports.length > 0) {
+    finalCode = `${imports.join("\n")}\n${finalCode}`;
+  }
 
   // Create source map based on the final transformed code
-  const map = createSourceMap(finalCode, source, moduleId, []);
+  const map = createSourceMap(finalCode, source, transformedModuleId, []);
 
   return {
     code: finalCode,

package/plugin/loader/types.ts

@@ -1,32 +1,38 @@
 import type { RawSourceMap } from "source-map";
-import type { DirectiveWarning, ParseResult, Program, AllowedDirectives } from "./directives/types.js";
+import type { DirectiveWarning, ParseResult, Program, AllowedDirectives, ParsedExports } from "./directives/types.js";
+import type { Logger } from "vite";
+import type { PanicThreshold } from "../types.js";
 
 
 
 export type LoaderConfig = {
   serverDirective: RegExp;
   clientDirective: RegExp;
+  directivePattern: RegExp;
   allowedDirectives: AllowedDirectives;
   getDirectiveType: (
     directive: string,
     moduleId?: string
   ) => "client" | "server" | undefined;
-  mode: "development" | "production" | "test";
-  importServerPath: string;
-  importClientPath: string;
-  registerClientReferenceName: string;
-  registerServerReferenceName: string;
-  isServerFunctionCode: (code: string, moduleId?: string) => boolean;
-  isClientComponentCode: (code: string, moduleId?: string) => boolean;
-  parse: (source: string) => Promise<{
-    ast: Program;
-    code: string;
-    map?: { url: string; start: number; end: number; lines: number } | null;
-  }>;
+  mode?: "development" | "production" | "test" | undefined;
+  importServerPath?: string;
+  importClientPath?: string;
+  registerClientReferenceName?: string;
+  registerServerReferenceName?: string;
+  isServerFunctionCode: (code: string, moduleId?: string, transformedModuleId?: string) => boolean;
+  isClientComponentCode: (code: string, moduleId?: string, transformedModuleId?: string) => boolean;
+  isClientComponentByCode: (code: string) => boolean;
+  isClientComponentByName: (moduleId: string, transformedModuleId?: string) => boolean;
+  moduleID: (moduleId: string, sourceContent?: string) => string;
+  parse: ParseFn;
+  verbose?: boolean;
+  logger?: Logger;
 };
 
 export type RscLoader = Pick<LoaderConfig, 'importServerPath' | 'importClientPath' | 'registerClientReferenceName' | 'registerServerReferenceName'>;
 
+export type ParseFn = (source: string) => Program | { ast: Program; code?: string; map?: any; exports?: ParsedExports } | Promise<Program | { ast: Program; code?: string; map?: any; exports?: ParsedExports }>
+
 export type TransformOptions = {
   forceServerFunction?: boolean;
   forceClientComponent?: boolean;
@@ -38,8 +44,12 @@ export type TransformOptions = {
   // based on warning, add warning directive index to avoid warning in development
   addDirectives?: number[];
   verbose?: boolean;
-  panicThreshold?: 'none' | 'critical_errors' | 'all_errors';
+  panicThreshold?: PanicThreshold;
   mode?: "development" | "production" | "test";
+  logger?: Logger;
+  moduleBase?: string;
+  // moduleID is on .loader.moduleID
+  // moduleID?: (moduleId: string) => string;
 };
 
 export type TransformResult = {
@@ -50,21 +60,19 @@ export type TransformResult = {
 export type TransformFunction = (
   source: string,
   moduleId: string,
+  transformedModuleId: string,
   parseResult: ParseResult,
-  options: TransformOptions
+  options: TransformOptions,
 ) => Promise<TransformResult>;
 
 export type TransformerFactory = (options: {
-  parseFn?: (source: string) => Promise<{ ast: Program; code: string; map?: {
-    url: string;
-    start: number;
-    end: number;
-    lines: number;
-  } | null }>;
-  options: Pick<TransformOptions, 'verbose' | 'loader' | 'panicThreshold'>;
+  parseFn?: ParseFn;
+  options: Pick<TransformOptions, 'verbose' | 'loader' | 'panicThreshold' | 'logger' | 'moduleBase'>;
   forceServerFunction?: boolean | undefined;
   forceClientComponent?: boolean | undefined;
   isServerEnvironment?: boolean;
-}) => (source: string, moduleId: string) => Promise<TransformResult>;
+  ssr?: boolean;
+  originalModuleId?: string;
+}) => (source: string, moduleId: string, transformedModuleId?: string) => Promise<TransformResult>;
 
 export type { Program };

package/plugin/metrics/collectStreamMetrics.ts

@@ -1,5 +1,4 @@
 import { Transform } from "node:stream";
-import type { PipeableStream } from "react-dom/server";
 import { createStreamMetrics } from "./createStreamMetrics.js";
 import type { StreamMetrics } from "./types.js";
 
@@ -8,7 +7,7 @@ import type { StreamMetrics } from "./types.js";
  * @param stream The stream to collect metrics from
  * @returns A promise that resolves with the metrics when the stream ends
  */
-export function collectStreamMetrics(stream: PipeableStream | NodeJS.ReadableStream): Promise<StreamMetrics> {
+export function collectStreamMetrics(stream: any | NodeJS.ReadableStream): Promise<StreamMetrics> {
   const metrics = createStreamMetrics();
   const startTime = performance.now()
 

package/plugin/metrics/convertStreamMetrics.ts

@@ -0,0 +1,28 @@
+import type { StreamMetrics, RenderMetrics } from "../types.js";
+import { createRenderMetrics } from "./createRenderMetrics.js";
+
+export function convertStreamMetricsToRenderMetrics(
+  streamMetrics: StreamMetrics,
+  route: string,
+  type: "rsc-full" | "rsc-headless" | "html",
+  fromMainThread: boolean = false,
+  fromRscWorker: boolean = false,
+  fromHtmlWorker: boolean = false,
+  processingTime?: number
+): RenderMetrics<typeof type> {
+  const elapsedTime = processingTime || streamMetrics.duration;
+  
+  return createRenderMetrics({
+    route,
+    type,
+    fromMainThread,
+    fromRscWorker,
+    fromHtmlWorker,
+    processingTime: elapsedTime,
+    chunks: streamMetrics.chunks,
+    chunkRate: streamMetrics.chunks / (elapsedTime / 1000),
+    memoryUsage: process.memoryUsage(),
+    streamMetrics,
+    fileSize: streamMetrics.bytes,
+  }) as RenderMetrics<typeof type>;
+}

package/plugin/metrics/createModuleResolutionMetrics.ts

@@ -0,0 +1,38 @@
+import { isMainThread } from "node:worker_threads";
+import type { ModuleResolutionMetrics } from "./types.js";
+
+export function createModuleResolutionMetrics({
+  route,
+  workerType,
+  resolutionTime,
+  fromMainThread = isMainThread,
+  fromRscWorker = false,
+  fromHtmlWorker = false,
+  memoryUsage = process.memoryUsage(),
+  description,
+}: {
+  route: string;
+  workerType: "rsc" | "html" | "mainThread";
+  resolutionTime: number;
+  fromMainThread?: boolean;
+  fromRscWorker?: boolean;
+  fromHtmlWorker?: boolean;
+  memoryUsage?: NodeJS.MemoryUsage;
+  description?: string;
+}): ModuleResolutionMetrics {
+  return {
+    route,
+    type: "module-resolution",
+    workerType,
+    resolutionTime,
+    fromMainThread,
+    fromRscWorker,
+    fromHtmlWorker,
+    memoryUsage,
+    description,
+  };
+}
+
+
+
+

package/plugin/metrics/createRenderMetrics.ts

@@ -1,17 +1,50 @@
+import { isMainThread } from "node:worker_threads";
 import { createStreamMetrics } from "./createStreamMetrics.js";
-import type { RenderMetrics } from "./types.js";
+import type { BaseRenderMetrics, CreateRenderMetricsFn, RenderMetrics } from "./types.js";
 
-export function createRenderMetrics(route: string): RenderMetrics {
-  return {
+export const createRenderMetrics: CreateRenderMetricsFn = function _createRenderMetrics({
+  route,
+  type,
+  fromMainThread = isMainThread,
+  fromRscWorker = false,
+  fromHtmlWorker = false,
+  processingTime = 0,
+  chunks = 0,
+  chunkRate = 0,
+  memoryUsage = process.memoryUsage(),
+  streamMetrics = createStreamMetrics(),
+  fileSize = 0,
+  fileName,
+  outputPath,
+  baseDir,
+  routePath,
+}) {
+  const base = {
     route,
-    htmlSize: 0,
-    rscSize: 0,
-    processingTime: 0,
-    chunks: 0,
-    chunkRate: 0,
-    memoryUsage: process.memoryUsage(),
-    streamMetrics: createStreamMetrics(),
-    htmlSizes: new Map(),
-    rscSizes: new Map(),
-  };
+    fromMainThread,
+    fromRscWorker,
+    fromHtmlWorker,
+    processingTime,
+    chunks,
+    chunkRate,
+    memoryUsage,
+    streamMetrics,
+  } as BaseRenderMetrics;
+
+  if (type === "html" || type === "rsc-headless") {
+    return {
+      ...base,
+      type,
+      fileSize,
+      fileName,
+      outputPath,
+      baseDir,
+      routePath,
+    } as unknown as RenderMetrics<typeof type>;
+  } else {
+    return {
+      ...base,
+      type,
+    } as unknown as RenderMetrics<typeof type>;
+  }
 }

package/plugin/metrics/createStreamMetrics.ts

@@ -1,14 +1,24 @@
-import type { StreamMetrics } from "./types.js";
+import type { CreateStreamMetricsFn } from "./types.js";
+import { performance } from "node:perf_hooks";
 
-export function createStreamMetrics(): StreamMetrics {
+export const createStreamMetrics: CreateStreamMetricsFn =
+  function _createStreamMetrics(metrics) {
+    if (metrics != null) {
+      return {
+        chunks: metrics.chunks ?? 0,
+        bytes: metrics.bytes ?? 0,
+        backpressureCount: metrics.backpressureCount ?? 0,
+        errorCount: metrics.errorCount ?? 0,
+        duration: metrics.duration ?? 0,
+        startTime: metrics.startTime ?? performance.now(),
+      };
+    }
     return {
       chunks: 0,
       bytes: 0,
       backpressureCount: 0,
-      drainCount: 0,
       errorCount: 0,
       duration: 0,
-      startTime: performance.now()
+      startTime: performance.now(),
     };
-  }
-  
+  };