@codeyam/codeyam-cli 0.1.0-staging.1669d45 → 0.1.0-staging.27d5a59

package/analyzer-template/project/writeScenarioComponents.ts

@@ -987,10 +987,19 @@ function addMockToContent(
           /[.*+?^${}()|[\]\\]/g,
           '\\$&',
         );
-        // Create regex that matches the call (with optional whitespace variations)
+        // Create regex that matches the call (with optional whitespace variations).
+        // TypeScript formatters commonly break type parameters across lines, e.g.:
+        //   useLoaderData<
+        //     typeof loader
+        //   >()
+        // So we allow optional whitespace around < and > delimiters, not just
+        // where whitespace already exists in the call signature string.
         const callRegex = new RegExp(
-          escapedCallSignature.replace(/\s+/g, '\\s*'),
-          'g',
+          escapedCallSignature
+            .replace(/\s+/g, '\\s*')
+            .replace(/</g, '\\s*<\\s*')
+            .replace(/>/g, '\\s*>\\s*'),
+          'gs',
         );
         fileContent = fileContent.replace(callRegex, `${mockFunctionName}()`);
       }
@@ -1012,15 +1021,29 @@ function addMockToContent(
         rootAnalysis.metadata?.scenariosDataStructure?.dataForMocks;
 
       // Find matching call signature key in dataForMocks
+      // IMPORTANT: First try exact match with type parameters (e.g., "useLoaderData<LoaderData>()")
+      // to avoid picking the wrong variant (e.g., "useLoaderData<typeof loader>()" which may
+      // have different properties). Fall back to base name matching only if exact match fails.
       const hookBaseName = importedExport.name.split(/[<(]/)[0];
-      const callSignatureKey = dataForMocks
-        ? Object.keys(dataForMocks).find((key) => {
+      const expectedKey = importedExport.calls?.[0];
+      let callSignatureKey: string | undefined;
+
+      if (dataForMocks) {
+        const keys = Object.keys(dataForMocks);
+
+        // First try exact match with the expected call signature
+        if (expectedKey && keys.includes(expectedKey)) {
+          callSignatureKey = expectedKey;
+        } else {
+          // Fall back to base name matching
+          callSignatureKey = keys.find((key) => {
             // Split on ., <, or ( to get the true base name
             // This handles both "useFlags()" -> "useFlags" and "trpc.useUtils()" -> "trpc"
             const keyBaseName = key.split(/[.<(]/)[0];
             return keyBaseName === hookBaseName;
-          })
-        : undefined;
+          });
+        }
+      }
 
       // Use the call signature if found, otherwise construct it
       const dataKey =
@@ -1032,11 +1055,17 @@ function addMockToContent(
       // use the base name (e.g., "trpc") when calling constructMockCode. This ensures
       // constructMockCode generates a complete nested mock from the schema without
       // referencing __cyOriginal variables.
+      // NOTE: Only check for dots BEFORE the first open paren. Function calls with
+      // dotted arguments like "getSurveyWithMetadata(params.surveyId)" have dots
+      // INSIDE parens — these are NOT method chains and should keep the full signature.
       const dataKeyBaseName = dataKey.split(/[.<(]/)[0];
+      const firstParenIndex = dataKey.indexOf('(');
+      const partBeforeFirstParen =
+        firstParenIndex >= 0 ? dataKey.substring(0, firstParenIndex) : dataKey;
       const isMethodChainDataKey =
         dataKeyBaseName === importedExport.name &&
         dataKey !== importedExport.name &&
-        dataKey.includes('.');
+        partBeforeFirstParen.includes('.');
       const mockNameToUse = isMethodChainDataKey
         ? importedExport.name
         : dataKey;
@@ -1061,18 +1090,36 @@ function addMockToContent(
       // IMPORTANT: We inline scenarios().data() inside the function rather than
       // storing in a const - see comment in constructMockCode.ts for why.
       if (!mockCode) {
-        mockCode = `function ${importedExport.name}(...args) {
-  return scenarios().data()?.["${dataKey}"];
+        mockCode = `// PATCHED: memoize to return stable reference (prevents infinite useEffect re-triggers)
+const _${importedExport.name}Ref = {
+  current: null,
+};
+function ${importedExport.name}(...args) {
+  if (!_${importedExport.name}Ref.current) {
+    _${importedExport.name}Ref.current = scenarios().data()?.["${dataKey}"];
+  }
+  return _${importedExport.name}Ref.current;
 }`;
       }
     } else {
       // Helper to find matching call signature key from dataForMocks
+      // IMPORTANT: First try exact match with type parameters (e.g., "useLoaderData<LoaderData>()")
+      // to avoid picking the wrong variant. Fall back to base name matching only if needed.
       const hookBaseName = importedExport.name.split(/[<(]/)[0];
+      const expectedKey = importedExport.calls?.[0];
       const findMatchingKey = (
         dataForMocks: Record<string, unknown> | undefined,
       ): string | undefined => {
         if (!dataForMocks) return undefined;
-        return Object.keys(dataForMocks).find((key) => {
+        const keys = Object.keys(dataForMocks);
+
+        // First try exact match with the expected call signature
+        if (expectedKey && keys.includes(expectedKey)) {
+          return expectedKey;
+        }
+
+        // Fall back to base name matching
+        return keys.find((key) => {
           // Split on ., <, or ( to get the true base name
           // This handles both "useFlags()" -> "useFlags" and "trpc.useUtils()" -> "trpc"
           const keyBaseName = key.split(/[.<(]/)[0];
@@ -1104,12 +1151,20 @@ function addMockToContent(
       // referencing __cyOriginal variables. The __cyOriginal pattern is only needed
       // for partial mocking where we preserve some original methods, not for complete
       // method-chain mocks where we provide all implementations.
+      // NOTE: Only check for dots BEFORE the first open paren. Function calls with
+      // dotted arguments like "getSurveyWithMetadata(params.surveyId)" have dots
+      // INSIDE parens — these are NOT method chains and should keep the full signature.
       const dataKeyBaseName = dataKey?.split(/[.<(]/)[0];
+      const elseFirstParenIndex = dataKey?.indexOf('(') ?? -1;
+      const elsePartBeforeFirstParen =
+        dataKey && elseFirstParenIndex >= 0
+          ? dataKey.substring(0, elseFirstParenIndex)
+          : dataKey;
       const isMethodChainDataKey =
         dataKey &&
         dataKeyBaseName === importedExport.name &&
         dataKey !== importedExport.name &&
-        dataKey.includes('.');
+        (elsePartBeforeFirstParen?.includes('.') ?? false);
       const mockNameToUse = isMethodChainDataKey
         ? importedExport.name
         : (dataKey ?? importedExport.calls?.[0] ?? `${importedExport.name}()`);
@@ -1131,8 +1186,15 @@ function addMockToContent(
       // IMPORTANT: We inline scenarios().data() inside the function rather than
       // storing in a const - see comment in constructMockCode.ts for why.
       if (!mockCode && dataKey) {
-        mockCode = `function ${importedExport.name}(...args) {
-  return scenarios().data()?.["${dataKey}"];
+        mockCode = `// PATCHED: memoize to return stable reference (prevents infinite useEffect re-triggers)
+const _${importedExport.name}Ref = {
+  current: null,
+};
+function ${importedExport.name}(...args) {
+  if (!_${importedExport.name}Ref.current) {
+    _${importedExport.name}Ref.current = scenarios().data()?.["${dataKey}"];
+  }
+  return _${importedExport.name}Ref.current;
 }`;
       }
     }
@@ -1205,6 +1267,21 @@ function addMockToContent(
         fileContent = fileContent.replace(importRegExp, '$1$2');
       }
 
+      // Also handle namespace imports (import * as foo from '...')
+      // These need to be renamed to foo__cyOriginal when the mock code spreads from the original.
+      // Note: We match any path (not just escapedImportPath) because the import path may have
+      // been rewritten by transitive import handling before this code runs.
+      if (shouldRenameToOriginal) {
+        const namespaceImportRegExp = new RegExp(
+          `(import\\s+\\*\\s+as\\s+)${escapeRegExp(firstPart)}(\\s+from\\s+['"][^'"]*['"])`,
+          'm',
+        );
+        fileContent = fileContent.replace(
+          namespaceImportRegExp,
+          `$1${firstPart}__cyOriginal$2`,
+        );
+      }
+
       // Remove empty imports entirely to avoid partial commenting issues with multiline imports
       // This handles both single-line and multiline empty imports
       fileContent = fileContent.replace(
@@ -1698,9 +1775,16 @@ export default async function writeScenarioComponents({
   });
 
   let importedExportIndex = 0;
+  const loopStartTime = Date.now();
+  console.log(
+    `[WriteScenario] Starting import loop for ${entity.name}: ${sortedImportedExports.length} imports`,
+  );
   for (const importedExport of sortedImportedExports) {
     importedExportIndex++;
     if (importedExportIndex % 5 === 0 || importedExportIndex === 1) {
+      console.log(
+        `[WriteScenario] ${entity.name} import ${importedExportIndex}/${sortedImportedExports.length}: ${importedExport.name} elapsed=${Date.now() - loopStartTime}ms`,
+      );
       debugLog(
         `Processing importedExport ${importedExportIndex}/${sortedImportedExports.length}`,
         {
@@ -1850,6 +1934,17 @@ export default async function writeScenarioComponents({
               dataFileContent = stripServerOnlyImport(dataFileContent);
               dataFileContent = applyServerOnlyMocks(dataFileContent);
 
+              // Process transitive imports: if the data entity file imports from files
+              // that have `import "server-only"`, rewrite those imports to transitive
+              // copies with server-only stripped. Without this, data entity files that
+              // end up in a client bundle (e.g., via a stripped "use server" boundary)
+              // would pull in server-only code and cause build failures.
+              dataFileContent = await processTransitiveImportsRecursively(
+                dataFileContent,
+                fileNotMocked.path,
+                dataScenarioPath,
+              );
+
               // Write the transformed data entity file
               await writeFile(dataScenarioPath, dataFileContent);
               scenarioComponentPaths.push(dataScenarioPath);
@@ -1875,9 +1970,9 @@ export default async function writeScenarioComponents({
               importedExport.name,
             );
 
-            // Don't recurse - data entities don't need their dependencies processed
-            // The import rewriting will happen later in this same loop iteration
-            // (at lines ~1590-1702) to point imports to this transformed file
+            // Don't recurse into full writeScenarioComponents - data entities don't
+            // need mock generation for their dependencies. Transitive server-only
+            // import processing is already handled above via processTransitiveImportsRecursively.
           } else {
             // For visual/library entities, recurse to process their dependencies
 
@@ -1888,6 +1983,10 @@ export default async function writeScenarioComponents({
               importedExport.resolvedIsDefault === true &&
               importedExport.isDefault === false;
 
+            console.log(
+              `[WriteScenario] RECURSE START: ${entity.name} -> ${importedExportEntity.name}`,
+            );
+            const recurseStartTime = Date.now();
             debugLog(
               `Recursing into writeScenarioComponents for ${importedExportEntity.name}`,
               {
@@ -1919,7 +2018,10 @@ export default async function writeScenarioComponents({
                   ? importedExport.name
                   : undefined,
               }),
-              60000, // 1 minute timeout for recursive calls
+              180000, // 3 minute timeout for recursive calls (complex components need more time)
+            );
+            console.log(
+              `[WriteScenario] RECURSE END: ${entity.name} -> ${importedExportEntity.name} took ${Date.now() - recurseStartTime}ms`,
             );
             debugLog(
               `Completed recursive writeScenarioComponents for ${importedExportEntity.name}`,
@@ -1943,20 +2045,25 @@ export default async function writeScenarioComponents({
       // that stubbing would break (e.g., Zod schemas with .superRefine())
       const isDataEntity = entityType === 'data' || entityType === 'type';
 
-      // Heuristic: Zod schemas are often misclassified as 'library' but should be preserved
-      // Detect by: name starts with Z + uppercase letter, AND has Zod method calls
-      const looksLikeZodSchema =
-        entityType === 'library' &&
-        /^Z[A-Z]/.test(importedExport.name) &&
-        importedExport.calls?.some((call: string) =>
-          /\.(superRefine|refine|transform|default|optional|nullable|array|object|string|number|boolean|parse|safeParse)\s*\(/.test(
-            call,
-          ),
-        );
+      // If calls data shows the entity is only accessed via properties/methods
+      // (e.g., formValidator.validate(), schema.superRefine()) and never directly
+      // invoked (e.g., getInitialProps()), it's used as an object and should be
+      // preserved rather than replaced with a Proxy stub.
+      const onlyPropertyAccessed =
+        importedExport.calls?.length > 0 &&
+        !importedExport.calls.some((call: string) => {
+          const afterName = call.slice(importedExport.name.length);
+          return afterName.startsWith('(') || afterName.startsWith('<');
+        });
 
-      // Callable entities can be safely stubbed (but not Zod schemas)
+      // Callable entities can be safely stubbed. Entities that are only
+      // property-accessed should be preserved (their methods need to work).
+      // 'other' entities are unknown types — safer to preserve than stub.
       const isCallable =
-        !isDataEntity && !looksLikeZodSchema && entityType !== undefined;
+        !isDataEntity &&
+        !onlyPropertyAccessed &&
+        entityType !== undefined &&
+        entityType !== 'other';
 
       // Determine what action to take
       const shouldStripAndReplace = hasMock;
@@ -2261,6 +2368,10 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
     }
   }
 
+  // Track post-import-loop timing
+  const postLoopStartTime = Date.now();
+  console.log(`[WriteScenario] POST-LOOP START: ${entity.name}`);
+
   // Collect universal mocks BEFORE processing nodeModuleImports
   // This is needed to check if a node module import is handled by a universal mock
   const universalMocks = project.metadata?.universalMocks ?? [];
@@ -2320,6 +2431,10 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
     );
   }
 
+  console.log(
+    `[WriteScenario] POST-LOOP ${entity.name}: node+universal mocks took ${Date.now() - postLoopStartTime}ms`,
+  );
+
   if (
     rootAnalysis.entitySha === entity.sha &&
     entity.metadata?.notExported &&
@@ -2398,6 +2513,10 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
   );
   debugLog('Completed rewriteRelativeModuleImports');
 
+  console.log(
+    `[WriteScenario] POST-LOOP ${entity.name}: transformations took ${Date.now() - postLoopStartTime}ms`,
+  );
+
   /**
    * Recursively process a file's imports to create transitive copies with server-only stripped.
    * This handles chains of arbitrary depth: A -> B -> C -> D where each needs server-only removed.
@@ -2418,7 +2537,7 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
     startTime: number = Date.now(),
   ): Promise<string> {
     // Global timeout for entire transitive processing
-    const GLOBAL_TIMEOUT_MS = 60000; // 1 minute max for all transitive processing
+    const GLOBAL_TIMEOUT_MS = 180000; // 3 minutes max for all transitive processing (complex components need more time)
     const elapsed = Date.now() - startTime;
     if (elapsed > GLOBAL_TIMEOUT_MS) {
       throw new Error(
@@ -2427,6 +2546,10 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
     }
 
     const importPaths = extractInternalImportPaths(content);
+    // Always log to help debug timeout issues
+    console.log(
+      `[TransitiveImports] depth=${depth} file=${path.basename(sourceFilePath)} imports=${importPaths.length} visited=${visitedPaths.size} elapsed=${Date.now() - startTime}ms`,
+    );
     debugLog(`processTransitiveImportsRecursively depth=${depth}`, {
       sourceFilePath,
       importCount: importPaths.length,
@@ -2481,8 +2604,10 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
       );
       const extension = importFile.name.split('.').pop();
       const isIndex = isIndexPath(importFile.path);
-      const pathHash = safeFileName(importFile.path);
-      const transitiveFilePath = `${PROJECT_RELATIVE_PATH}/${basePath}/${pathHash}_${isIndex ? 'index_' : ''}transitive_${safeFileName(scenario.name)}.${extension}`;
+      // Limit pathHash length to prevent ENAMETOOLONG errors on macOS (255 char limit)
+      const pathHash = safeFileName(importFile.path, { maxLength: 80 });
+      const scenarioSlug = safeFileName(scenario.name, { maxLength: 60 });
+      const transitiveFilePath = `${PROJECT_RELATIVE_PATH}/${basePath}/${pathHash}_${isIndex ? 'index_' : ''}transitive_${scenarioSlug}.${extension}`;
 
       // Check if this is a circular import (we're already processing this file)
       const isCircularImport = visitedPaths.has(resolvedPath);
@@ -2606,6 +2731,10 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
     return modifiedContent;
   }
 
+  console.log(
+    `[WriteScenario] POST-LOOP ${entity.name}: before remaining imports ${Date.now() - postLoopStartTime}ms`,
+  );
+
   // Process remaining internal imports that weren't in importedExports
   // This handles transitive dependencies: when the file content includes code (e.g., from
   // other functions in the same file) that imports from files with "server-only"
@@ -2634,6 +2763,20 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
     debugLog(
       `[REMAINING LOOP] import ${remainingImportIndex}/${remainingImportPaths.length}: ${importPath}`,
     );
+
+    // Skip imports that point to generated CodeYam files (same skip logic as
+    // rewriteRelativeModuleImports). Without this, MockData files from earlier
+    // captures that get discovered by the TypeScript compiler would be treated as
+    // regular imports, creating transitive copies with stale content.
+    const scenarioFilePattern = /[a-f0-9]{64}_\w+_[A-Z]\w*$/;
+    const mockDataPattern = /__codeyamMocks__\//;
+    if (
+      scenarioFilePattern.test(importPath) ||
+      mockDataPattern.test(importPath)
+    ) {
+      continue;
+    }
+
     // Resolve the import path to a project file path
     const resolvedFilePath = resolveImportPath(importPath, file.path, project);
 
@@ -2681,8 +2824,10 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
     );
     const targetFileExtension = targetFile.name.split('.').pop();
     const targetFileIsIndex = isIndexPath(targetFile.path);
-    const filePathHash = safeFileName(targetFile.path);
-    const transformedFilePath = `${PROJECT_RELATIVE_PATH}/${targetFileBasePath}/${filePathHash}_${targetFileIsIndex ? 'index_' : ''}transitive_${safeFileName(scenario.name)}.${targetFileExtension}`;
+    // Limit path hash length to prevent ENAMETOOLONG errors
+    const filePathHash = safeFileName(targetFile.path, { maxLength: 80 });
+    const targetScenarioSlug = safeFileName(scenario.name, { maxLength: 60 });
+    const transformedFilePath = `${PROJECT_RELATIVE_PATH}/${targetFileBasePath}/${filePathHash}_${targetFileIsIndex ? 'index_' : ''}transitive_${targetScenarioSlug}.${targetFileExtension}`;
 
     // Check if we've already processed this file as a transitive copy
     // Note: __data_file_written__ is for entity-specific scenario files with different naming,
@@ -2737,8 +2882,12 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
         );
         const nestedExtension = nestedFile.name.split('.').pop();
         const nestedIsIndex = isIndexPath(nestedFile.path);
-        const nestedPathHash = safeFileName(nestedFile.path);
-        const nestedTransformedPath = `${PROJECT_RELATIVE_PATH}/${nestedBasePath}/${nestedPathHash}_${nestedIsIndex ? 'index_' : ''}transitive_${safeFileName(scenario.name)}.${nestedExtension}`;
+        // Limit path hash length to prevent ENAMETOOLONG errors
+        const nestedPathHash = safeFileName(nestedFile.path, { maxLength: 80 });
+        const nestedScenarioSlug = safeFileName(scenario.name, {
+          maxLength: 60,
+        });
+        const nestedTransformedPath = `${PROJECT_RELATIVE_PATH}/${nestedBasePath}/${nestedPathHash}_${nestedIsIndex ? 'index_' : ''}transitive_${nestedScenarioSlug}.${nestedExtension}`;
 
         // Check if already processed as a transitive file (we can rewrite to point to it)
         // Note: __data_file_written__ is for entity-specific scenario files with different naming,
@@ -2859,6 +3008,10 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
     fileContent = fileContent.replace(importRegex, `$1${safeRelativePath}$2`);
   }
 
+  console.log(
+    `[WriteScenario] POST-LOOP ${entity.name}: remaining imports loop took ${Date.now() - postLoopStartTime}ms`,
+  );
+
   const scenarioComponentComment = `// This file is auto-generated by CodeYam. Do not edit this file manually.
 // This file contains content for a scenario component:
 // Analyses being written: ${JSON.stringify(fileAnalyses?.map((a) => ({ id: a.id, entityName: a.entityName })))}
@@ -2870,6 +3023,83 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
 // Scenario: ${scenario.id} - ${scenario.name}
 `;
 
+  // Final pass: Rename any namespace imports (import * as X from '...') that have
+  // corresponding mock code using ...X__cyOriginal spread pattern.
+  // This handles cases where:
+  // 1. The import path was rewritten by transitive import handling
+  // 2. The import wasn't caught by the earlier renaming logic
+  // We scan for all __cyOriginal references in the mock code and ensure the imports are renamed.
+  const cyOriginalReferences = fileContent.match(/\.\.\.(\w+)__cyOriginal/g);
+  if (cyOriginalReferences) {
+    const uniqueNames = [
+      ...new Set(
+        cyOriginalReferences.map((ref) =>
+          ref.replace('...', '').replace('__cyOriginal', ''),
+        ),
+      ),
+    ];
+    for (const name of uniqueNames) {
+      // Match namespace imports for this name that haven't been renamed yet
+      const namespaceImportRegex = new RegExp(
+        `(import\\s+\\*\\s+as\\s+)${escapeRegExp(name)}(\\s+from\\s+['"][^'"]*['"])`,
+        'g',
+      );
+      fileContent = fileContent.replace(
+        namespaceImportRegex,
+        `$1${name}__cyOriginal$2`,
+      );
+    }
+  }
+
+  // For route components (page.tsx, layout.tsx), the component IS the Next.js page.
+  // There's no wrapper page to inject argumentsData as props (unlike non-route components
+  // which get a scenarioComponent wrapper). We need to wrap the default export so that
+  // the scenario's argumentsData is passed as props to the component.
+  if (
+    isFrameworkRoute(file, entity, framework, file === rootFile) &&
+    rootAnalysis.metadata?.scenariosDataStructure?.arguments?.length > 0
+  ) {
+    const positionalArguments =
+      rootAnalysis.metadata.scenariosDataStructure.arguments;
+    const hasNamedArgs =
+      positionalArguments.length === 1 &&
+      typeof positionalArguments[0] === 'object';
+
+    if (hasNamedArgs) {
+      // Match: export default function Name(
+      // Also: export default async function Name(
+      const defaultExportMatch = fileContent.match(
+        /export\s+default\s+(async\s+)?function\s+(\w+)\s*\(/,
+      );
+
+      if (defaultExportMatch) {
+        const funcName = defaultExportMatch[2];
+
+        // Remove "export default" from the original function declaration
+        fileContent = fileContent.replace(
+          /export\s+default\s+(async\s+)?function\s+(\w+)\s*\(/,
+          '$1function $2(',
+        );
+
+        // Ensure scenarios import is present
+        const mockDataPath = `${relativeMocksDir}/MockData_${safeFileName(scenario.name)}`;
+        if (fileContent.indexOf('import { scenarios } from') === -1) {
+          fileContent = `import { scenarios } from "${mockDataPath}";\n\n${fileContent}`;
+        }
+
+        // Add wrapper default export that injects argumentsData as props.
+        // Filter null values: the LLM uses null to mean "don't pass this prop"
+        // but spreading { prop: null } explicitly passes null, which differs from
+        // undefined (absent) and crashes components that check `prop !== undefined`.
+        fileContent += `\n\nexport default function _CYRouteWrapper(props: any) {
+  const _cyArgs = scenarios().data()?.['arguments']?.[0] ?? {};
+  const _cyCleanArgs = Object.fromEntries(Object.entries(_cyArgs).filter(([_, v]) => v !== null));
+  return <${funcName} {...props} {..._cyCleanArgs} />;
+}\n`;
+      }
+    }
+  }
+
   // Use the directive that was extracted at the beginning of processing
   // This ensures it stays at the very top even after imports are prepended
   // NOTE: We only preserve "use client" directives, NOT "use server" directives.
@@ -2891,5 +3121,9 @@ ${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${
   debugLog('Successfully wrote scenario file');
   scenarioComponentPaths.push(scenarioComponentPath);
 
+  console.log(
+    `[WriteScenario] POST-LOOP ${entity.name}: COMPLETE total=${Date.now() - postLoopStartTime}ms`,
+  );
+
   return { scenarioComponentPaths, writtenScenarioComponents };
 }

package/analyzer-template/project/writeSimpleRoot.ts

@@ -277,16 +277,17 @@ export default async function writeSimpleRoot(
               let mockKey: string | undefined;
 
               if (matchingKeys.length > 1) {
-                // Parse root.tsx to find how this hook is used
-                // Look for patterns like: const { prop1, prop2 } = useLoaderData()
-                // or: const result = useLoaderData()
+                // Parse root.tsx to find how this hook is used, INCLUDING type parameters
+                // Look for patterns like: const { prop1, prop2 } = useLoaderData<LoaderData>()
+                // or: const result = useLoaderData<typeof loader>()
                 const hookUsageRegex = new RegExp(
-                  `(?:const|let|var)\\s+(?:\\{([^}]+)\\}|([a-zA-Z_][a-zA-Z0-9_]*))\\s*=\\s*${entityName}(?:<[^>]*>)?\\s*\\(`,
+                  `(?:const|let|var)\\s+(?:\\{([^}]+)\\}|([a-zA-Z_][a-zA-Z0-9_]*))\\s*=\\s*${entityName}(<[^>]+>)?\\s*\\(`,
                   'g',
                 );
 
                 const matches = [...newContent.matchAll(hookUsageRegex)];
                 const rootPropertyNames: string[] = [];
+                const typeParamsUsed: Set<string> = new Set();
 
                 for (const match of matches) {
                   if (match[1]) {
@@ -300,17 +301,26 @@ export default async function writeSimpleRoot(
                     // Simple assignment: const result = ...
                     rootPropertyNames.push(match[2]);
                   }
+                  // Capture type parameters used (e.g., "<LoaderData>", "<typeof loader>")
+                  if (match[3]) {
+                    typeParamsUsed.add(match[3]);
+                  }
                 }
 
-                // Find a key that matches by call signature
-                // For call signature format (e.g., "useLoaderData()"), just use the first match
-                mockKey = matchingKeys[0];
+                // First, try to find a key that matches the exact type parameters used in the source
+                // e.g., if source has useLoaderData<LoaderData>(), prefer "useLoaderData<LoaderData>()" key
+                if (typeParamsUsed.size > 0) {
+                  for (const typeParam of typeParamsUsed) {
+                    const exactKey = `${entityName}${typeParam}()`;
+                    if (matchingKeys.includes(exactKey)) {
+                      mockKey = exactKey;
+                      break;
+                    }
+                  }
+                }
 
-                // Fallback: if no match found but there are multiple keys,
-                // exclude the route's key (which we can identify from mockedImports variable names)
+                // Fallback: if no exact type parameter match, use the first match
                 if (!mockKey) {
-                  // The route's variable names would be in its own analysis
-                  // For simplicity, just use the first key that's different from the most common pattern
                   mockKey = matchingKeys[0];
                 }
               } else if (matchingKeys.length === 1) {