@codeyam/codeyam-cli 0.1.0-staging.ad31e3e → 0.1.0-staging.b8a55ba

package/analyzer-template/packages/ai/src/lib/gatherRelevantDependentKeyAttributes.ts

@@ -1,4 +1,9 @@
-import { Entity, ReadonlyAnalysis, ReadonlyAnalysisMap } from '~codeyam/types';
+import {
+  DataStructure,
+  Entity,
+  ReadonlyAnalysis,
+  ReadonlyAnalysisMap,
+} from '~codeyam/types';
 import {
   joinParenthesesAndArrays,
   splitOutsideParenthesesAndArrays,
@@ -7,6 +12,7 @@ import { gatherAttributesMap } from './promptGenerators/gatherAttributesMap';
 
 export default function gatherRelevantDependentKeyAttributes(
   entity: Pick<Entity, 'name' | 'metadata'>,
+  mergedDataStructure: Omit<DataStructure, 'equivalentSignatureVariables'>,
   dependentAnalyses: ReadonlyAnalysisMap,
 ) {
   // console.info(
@@ -52,10 +58,16 @@ export default function gatherRelevantDependentKeyAttributes(
   //   ),
   // );
 
-  const { associationMap } = gatherAttributesMap(entity);
+  const isolatedDataStructure = entity.metadata?.isolatedDataStructure;
+  const equivalentSignatureVariables =
+    isolatedDataStructure?.equivalentSignatureVariables ?? {};
+  const { associationMap } = gatherAttributesMap(
+    isolatedDataStructure ?? mergedDataStructure,
+    mergedDataStructure,
+    equivalentSignatureVariables,
+  );
 
-  const dependencySchemas =
-    entity.metadata?.isolatedDataStructure?.dependencySchemas;
+  const dependencySchemas = mergedDataStructure?.dependencySchemas;
   const relevantDependentKeyAttributes: ReadonlyAnalysis['metadata']['keyAttributes'][0][] =
     [];
   Object.keys(dependentAnalyses).forEach((path) => {
@@ -70,37 +82,178 @@ export default function gatherRelevantDependentKeyAttributes(
       const signatureSchema = dependencySchema?.signatureSchema;
       if (!signatureSchema) return;
 
-      const sourceEquivalencies = dependencySchema?.sourceEquivalencies ?? {};
+      // Get sourceEquivalencies from the dependent analysis's mergedDataStructure,
+      // not from dependencySchemas (which doesn't store them)
+      const sourceEquivalencies =
+        dependentAnalysis?.metadata?.mergedDataStructure?.sourceEquivalencies ??
+        {};
+
+      // Build a fallback mapping from parent's usageEquivalencies
+      // This helps when the child's sourceEquivalencies don't have a mapping
+      // For example, if parent passes `entity` to child's `signature[0]`,
+      // usageEquivalencies might have: entity -> [{schemaPath: 'signature[0]', scopeNodeName: entityName}]
+      const usageEquivalencies = mergedDataStructure?.usageEquivalencies ?? {};
+      const signatureToParentPath: Record<string, string> = {};
+
+      for (const parentPath in usageEquivalencies) {
+        const usages = usageEquivalencies[parentPath];
+        for (const usage of usages) {
+          if (usage.scopeNodeName !== entityName) continue;
+
+          // Extract signature path from schemaPath
+          // Format can be:
+          // 1. "functionName(args).signature[N]" or "functionName(args).signature[N].property" (real data)
+          // 2. "signature[N]" or "signature[N].property" (simplified/test format)
+          let signaturePath: string;
+          const signatureIndex = usage.schemaPath.indexOf('.signature[');
+          if (signatureIndex !== -1) {
+            // Full format: extract "signature[N]" or "signature[N].property" after the dot
+            signaturePath = usage.schemaPath.slice(signatureIndex + 1);
+          } else if (usage.schemaPath.startsWith('signature[')) {
+            // Simplified format: already starts with signature[
+            signaturePath = usage.schemaPath;
+          } else {
+            continue;
+          }
+
+          // Map child's signature path to parent's path
+          // e.g., 'signature[0]' -> 'entity'
+          signatureToParentPath[signaturePath] = parentPath;
+        }
+      }
 
       // If sourceEquivalencies is empty but the child has mocked imports,
-      // propagate key attributes directly since they affect the parent's visual output
+      // only propagate key attributes that come from mocked data sources
+      // (not props that the parent passes to the child)
       if (Object.keys(sourceEquivalencies).length === 0) {
-        // Check both importedExports and nodeModuleImports for mocked imports
+        // Collect names of mocked imports
         const importedExports =
           dependentAnalysis.entity?.metadata?.importedExports ?? [];
         const nodeModuleImports =
           dependentAnalysis.entity?.metadata?.nodeModuleImports ?? {};
 
-        const hasMockedImportedExports = importedExports.some(
-          (imp: { isMocked?: boolean }) => imp.isMocked,
-        );
-        const hasMockedNodeModuleImports = Object.values(
-          nodeModuleImports,
-        ).some((imports: any[]) =>
-          imports?.some((imp: { isMocked?: boolean }) => imp.isMocked),
-        );
+        const mockedImportNames = new Set<string>();
+
+        // Collect mocked names from importedExports
+        for (const imp of importedExports) {
+          if ((imp as { isMocked?: boolean }).isMocked) {
+            mockedImportNames.add((imp as { name: string }).name);
+          }
+        }
+
+        // Collect mocked names from nodeModuleImports
+        for (const imports of Object.values(nodeModuleImports)) {
+          for (const imp of imports as any[]) {
+            if ((imp as { isMocked?: boolean }).isMocked) {
+              mockedImportNames.add((imp as { name: string }).name);
+            }
+          }
+        }
+
+        if (mockedImportNames.size > 0) {
+          for (const dependentKeyAttribute of dependentKeyAtttributes) {
+            // Include key attributes that come from controllable data sources:
+            // 1. Mocked data sources (like getScenarioDisplayState())
+            // 2. Arguments/props passed to the child (signature[N].*)
+            // Exclude attributes that are clearly just the prop name itself (like "size")
+            const externalPath = dependentKeyAttribute.externalPath || '';
+            const dataStructurePath =
+              dependentKeyAttribute.dataStructurePath || '';
 
-        const childHasMockedImports =
-          hasMockedImportedExports || hasMockedNodeModuleImports;
+            const comesFromMockedSource = Array.from(mockedImportNames).some(
+              (mockedName) =>
+                externalPath.includes(`${mockedName}(`) ||
+                externalPath.includes(`${mockedName}.`) ||
+                dataStructurePath.includes(`${mockedName}(`) ||
+                dataStructurePath.includes(`${mockedName}.`) ||
+                dataStructurePath.startsWith(mockedName),
+            );
+
+            // Check if it comes from an argument (signature[N].something.deeper)
+            // but not just a simple prop like "size" or "isOutdated"
+            const comesFromArgument =
+              externalPath.startsWith('signature[') &&
+              // Must have nested path (at least one dot after the prop name)
+              // e.g., signature[0].scenario.metadata.* is OK
+              // but signature[0].size is not (it's just a simple prop)
+              (externalPath.match(/\./g) || []).length >= 2;
+
+            if (comesFromMockedSource || comesFromArgument) {
+              relevantDependentKeyAttributes.push({
+                ...dependentKeyAttribute,
+                dependentEntityName: entityName,
+                originalEntityName: entityName,
+                filePath: path,
+              } as any);
+            }
+          }
+          return;
+        }
 
-        if (childHasMockedImports) {
+        // No mocked imports, try usageEquivalencies fallback for direct argument usage
+        if (Object.keys(signatureToParentPath).length > 0) {
           for (const dependentKeyAttribute of dependentKeyAtttributes) {
-            relevantDependentKeyAttributes.push({
-              ...dependentKeyAttribute,
-              dependentEntityName: entityName,
-              originalEntityName: entityName,
-              filePath: path,
-            } as any);
+            if (!dependentKeyAttribute.externalPath?.startsWith('signature['))
+              continue;
+
+            const dependentKeyAttributeExternalPathParts =
+              splitOutsideParenthesesAndArrays(
+                dependentKeyAttribute.externalPath,
+              );
+
+            // Try to find a matching signature path
+            for (
+              let i = dependentKeyAttributeExternalPathParts.length;
+              i > 0;
+              i--
+            ) {
+              const subExternalPath = joinParenthesesAndArrays(
+                dependentKeyAttributeExternalPathParts.slice(0, i),
+              );
+
+              if (signatureToParentPath[subExternalPath]) {
+                const parentPath = signatureToParentPath[subExternalPath];
+
+                // parentPath is already in parent's schema format (e.g., "signature[0].entity" or just "entity")
+                // We need to handle both cases
+                const parentPathStartsWithSignature =
+                  parentPath.startsWith('signature[');
+                let externalPath = parentPathStartsWithSignature
+                  ? parentPath
+                  : `signature[0].${parentPath}`;
+
+                // internalPath is parentPath without the signature[N]. prefix
+                let internalPath: string | undefined =
+                  parentPathStartsWithSignature
+                    ? parentPath.replace(/^signature\[\d+\]\./, '')
+                    : parentPath;
+
+                if (i < dependentKeyAttributeExternalPathParts.length) {
+                  const remainingPath = joinParenthesesAndArrays(
+                    dependentKeyAttributeExternalPathParts.slice(i),
+                  );
+                  externalPath = `${externalPath}.${remainingPath}`;
+                  internalPath = `${internalPath}.${remainingPath}`;
+                }
+
+                // Look up in associationMap for proper internal path
+                if (associationMap[externalPath]) {
+                  internalPath = associationMap[externalPath];
+                }
+
+                relevantDependentKeyAttributes.push({
+                  ...dependentKeyAttribute,
+                  dependentEntityName: entityName,
+                  originalEntityName: entityName,
+                  filePath: path,
+                  internalPath,
+                  externalPath,
+                  dataStructurePath: internalPath,
+                } as any);
+
+                break;
+              }
+            }
           }
         }
         return;
@@ -111,6 +264,8 @@ export default function gatherRelevantDependentKeyAttributes(
           splitOutsideParenthesesAndArrays(dependentKeyAttribute.externalPath);
 
         let matchFound = false;
+
+        // First, try to match using child's sourceEquivalencies
         for (
           let i = dependentKeyAttributeExternalPathParts.length;
           i > 0;
@@ -120,8 +275,8 @@ export default function gatherRelevantDependentKeyAttributes(
             dependentKeyAttributeExternalPathParts.slice(0, i),
           );
 
-          for (const sourceEquivalencyPath in dependencySchema?.sourceEquivalencies ??
-            {}) {
+          // Use the sourceEquivalencies from the child's analysis
+          for (const sourceEquivalencyPath in sourceEquivalencies) {
             const sourceEquivalencyPathParts = splitOutsideParenthesesAndArrays(
               sourceEquivalencyPath,
             );
@@ -133,8 +288,7 @@ export default function gatherRelevantDependentKeyAttributes(
             if (relevantSourceEquivalencyPath !== subExternalPath) continue;
 
             const sourcePath =
-              dependencySchema?.sourceEquivalencies[sourceEquivalencyPath][0]
-                ?.schemaPath;
+              sourceEquivalencies[sourceEquivalencyPath][0]?.schemaPath;
 
             if (!sourcePath) continue;
 
@@ -169,6 +323,8 @@ export default function gatherRelevantDependentKeyAttributes(
               filePath: path, // File path for restructuring into nested format
               internalPath,
               externalPath,
+              // Transform dataStructurePath to parent's path structure
+              dataStructurePath: internalPath,
             } as any);
 
             matchFound = true;
@@ -177,6 +333,74 @@ export default function gatherRelevantDependentKeyAttributes(
 
           if (matchFound) break;
         }
+
+        // Fallback: if no sourceEquivalency match, try using parent's usageEquivalencies
+        // This handles direct argument usage like entity.entityType
+        if (
+          !matchFound &&
+          dependentKeyAttribute.externalPath?.startsWith('signature[')
+        ) {
+          // Try to find a matching signature path in the usageEquivalencies-derived map
+          for (
+            let i = dependentKeyAttributeExternalPathParts.length;
+            i > 0;
+            i--
+          ) {
+            const subExternalPath = joinParenthesesAndArrays(
+              dependentKeyAttributeExternalPathParts.slice(0, i),
+            );
+
+            if (signatureToParentPath[subExternalPath]) {
+              const parentPath = signatureToParentPath[subExternalPath];
+
+              // parentPath is already in parent's schema format (e.g., "signature[0].entity" or just "entity")
+              // We need to handle both cases
+              const parentPathStartsWithSignature =
+                parentPath.startsWith('signature[');
+              let externalPath = parentPathStartsWithSignature
+                ? parentPath
+                : `signature[0].${parentPath}`;
+
+              // internalPath is parentPath without the signature[N]. prefix
+              let internalPath: string | undefined =
+                parentPathStartsWithSignature
+                  ? parentPath.replace(/^signature\[\d+\]\./, '')
+                  : parentPath;
+
+              if (i < dependentKeyAttributeExternalPathParts.length) {
+                const remainingPath = joinParenthesesAndArrays(
+                  dependentKeyAttributeExternalPathParts.slice(i),
+                );
+                externalPath = `${externalPath}.${remainingPath}`;
+                internalPath = `${internalPath}.${remainingPath}`;
+              }
+
+              // Look up the full external path in associationMap for proper internal path
+              const fullExternalPath = externalPath;
+              if (associationMap[fullExternalPath]) {
+                internalPath = associationMap[fullExternalPath];
+              }
+
+              const dependentEntityName =
+                dependentKeyAttribute.dependentEntityName
+                  ? `${dependentKeyAttribute.dependentEntityName} (through ${entityName})`
+                  : entityName;
+
+              relevantDependentKeyAttributes.push({
+                ...dependentKeyAttribute,
+                dependentEntityName,
+                originalEntityName: entityName,
+                filePath: path,
+                internalPath,
+                externalPath,
+                dataStructurePath: internalPath,
+              } as any);
+
+              matchFound = true;
+              break;
+            }
+          }
+        }
       }
     });
   });

package/analyzer-template/packages/ai/src/lib/generateChangesEntityKeyAttributes.ts

@@ -1,11 +1,17 @@
 import completionCall from './completionCall';
 import { saveLlmCall } from '~codeyam/aws/dynamodb';
-import { Analysis, Entity } from '~codeyam/types';
+import type {
+  Analysis,
+  CompoundConditional,
+  ConditionalUsage,
+  Entity,
+} from '~codeyam/types';
 import generateEntityKeyAttributesGenerator from './promptGenerators/generateEntityKeyAttributesGenerator';
 import { awsLog } from '~codeyam/utils';
 import { AI, DEFAULT_LARGER_MODEL } from '~codeyam/ai';
 import isFrontend from './isFrontend';
 import { gatherAttributesMap } from './promptGenerators/gatherAttributesMap';
+import { formatConditionalUsagesForPrompt } from './getConditionalUsagesFromCode';
 
 interface GenerateChangesEntityKeyAttributesArgs {
   entity: Entity;
@@ -29,14 +35,34 @@ export default async function generateChangesEntityKeyAttributes({
     existingKeyAttributes: existingKeyAttributes?.map((a) => a.externalPath),
   });
 
+  const isolatedDataStructure = entity.metadata?.isolatedDataStructure;
+  const mergedDataStructure = analysis.metadata.mergedDataStructure;
+  const equivalentSignatureVariables =
+    isolatedDataStructure?.equivalentSignatureVariables ?? {};
   const { attributesMap, associationMap, dataStructurePathMap } =
-    gatherAttributesMap(entity, true);
+    gatherAttributesMap(
+      isolatedDataStructure ?? mergedDataStructure,
+      mergedDataStructure,
+      equivalentSignatureVariables,
+      true,
+    );
+
+  // Use pre-computed conditional usages from metadata for static analysis context
+  const conditionalUsages: Record<string, ConditionalUsage[]> =
+    isolatedDataStructure?.conditionalUsages ?? {};
+  const compoundConditionals: CompoundConditional[] =
+    isolatedDataStructure?.compoundConditionals ?? [];
+  const staticAnalysisContext = formatConditionalUsagesForPrompt(
+    conditionalUsages,
+    compoundConditionals,
+  );
 
   const prompt = generateEntityKeyAttributesGenerator({
     code: entity.code,
     attributesMap,
     existingKeyAttributes,
     commitDiff,
+    staticAnalysisContext,
   });
 
   const response = await completionCall({
@@ -102,6 +128,9 @@ export default async function generateChangesEntityKeyAttributes({
       description: addedKeyAttribute.description,
       validValueOptions: addedKeyAttribute.validValueOptions,
       errorValueOptions: addedKeyAttribute.errorValueOptions,
+      // Preserve new fields from LLM response
+      valueType: addedKeyAttribute.valueType,
+      dependencies: addedKeyAttribute.dependencies,
     };
     if (index !== -1) {
       keyAttributes.splice(index, 1, fullKeyAttribute);
@@ -167,6 +196,14 @@ Note: A conditional that shows a different UI state (e.g., loading screen, redir
 
 Only return attributes from the "Attributes and Value Types" list—these are the only values we can control. Use the exact path provided. Use \`functionCallReturnValue\` for function return values. Use \`[]\` for array items (e.g., \`items[]\`).
 
+## Dependencies
+
+If this key attribute is only used if one or more other attributes have specific values, include those dependencies in a \`dependencies\` array.
+
+## Value Type
+
+Should reiterate the type of the attribute (e.g., "string", "object", "array", "boolean", etc.) based on the "Attributes and Value Types" list.
+
 ## Response Format
 
 For \`validValueOptions\`, include boundary/extreme values that trigger different behavior—but NOT values that would cause errors.
@@ -177,6 +214,10 @@ For \`validValueOptions\`, include boundary/extreme values that trigger differen
     {
       "path": "user",
       "description": "Controls whether the main form or redirect screen displays",
+      "dependencies": [
+        { "projectId": "must be a valid project ID for 'user' to have effect" }
+      ],
+      "valueType": "object",
       "validValueOptions": ["truthy user object", "null/undefined"]
     }
   ],
@@ -214,6 +255,14 @@ NOT key attributes:
 
 Only return attributes from the "Attributes and Value Types" list—these are the only values we can control. Use the exact path provided. Use \`functionCallReturnValue\` for function return values. Use \`[]\` for array items (e.g., \`items[]\`).
 
+## Dependencies
+
+If this key attribute is only used if one or more other attributes have specific values, include those dependencies in a \`dependencies\` array.
+
+## Value Type
+
+Should reiterate the type of the attribute (e.g., "string", "object", "array", "boolean", etc.) based on the "Attributes and Value Types" list.
+
 ## Response Format
 
 For \`validValueOptions\`, include boundary/extreme values that trigger different behavior—but NOT values that would cause errors.
@@ -224,6 +273,10 @@ For \`validValueOptions\`, include boundary/extreme values that trigger differen
     {
       "path": "config.enabled",
       "description": "Controls whether the main processing logic runs",
+      "dependencies": [
+        { "analysisMode": "must be 'full' for 'enabled' to have effect" }
+      ],
+      "valueType": "boolean",
       "validValueOptions": ["true", "false"]
     }
   ],

package/analyzer-template/packages/ai/src/lib/generateChangesEntityScenarioData.ts

@@ -328,7 +328,7 @@ const generateChangesSystemMessage = (
 Preserve existing data and only modify what's necessary for the new structure.
 - If structure unchanged and data sufficient, return empty data: \`{ "data": {} }\`
 - If new fields added, fill them with realistic values
-- Fill ALL fields with realistic values (except error attributes)
+- Fill ALL fields with realistic values (except error attributes and key attributes set to null or undefined)
 - NEVER include "error" fields—omit entirely`
     : `## Non-Default Scenario (Incremental Update)
 Preserve existing data and only modify what changed.
@@ -365,8 +365,11 @@ Use simple elements only (\`<div>\`, \`<span>\`). No custom components.
 { "onClick": { "~~codeyam-code~~": "() => console.log('clicked')" } }
 \`\`\`
 
-## Function Keys
-Preserve function call syntax exactly as written in the structure.
+## Mock Data Keys
+Preserve keys exactly as written in the structure. There are two formats:
+- Standard function calls: \`"useUser()"\`, \`"from().select()"\`
+- Variable-qualified calls (for multiple calls to the same function): \`"entityDiffFetcher <- useFetcher"\`
+  This reads as "entityDiffFetcher receives from useFetcher". Each variable gets distinct mock data.
 
 ## Response Format
 \`\`\`json

package/analyzer-template/packages/ai/src/lib/generateChangesEntityScenarios.ts

@@ -46,6 +46,7 @@ export default async function generateChangesEntityScenarios({
 }: GenerateChangesEntityScenariosArgs) {
   const relevantDependentKeyAttributes = gatherRelevantDependentKeyAttributes(
     entity,
+    analysis.metadata.mergedDataStructure,
     dependentAnalyses,
   );
   if (
@@ -71,6 +72,7 @@ export default async function generateChangesEntityScenarios({
 
   const prompt = generateChangesEntityScenariosGenerator({
     entity,
+    mergedDataStructure: analysis.metadata.mergedDataStructure,
     keyAttributes,
     addedKeyAttributes,
     removedKeyAttributes,
@@ -286,7 +288,16 @@ Focus scenario updates on demonstrating this change.${scenarioCount ? ` Generate
 
   return `You are updating data scenarios after a code change. Determine which existing scenarios need updating, which should be removed, and if new ones are needed.
 
-${contextSection}## Update Rules
+${contextSection}## Coverage Goals (IMPORTANT)
+Your primary goal is to MAXIMIZE coverage of key attributes across all scenarios:
+
+1. **Use indexed validValueOptions**: Each key attribute has indexed \`validValueOptions\` (with \`index\` and \`value\`). Reference these by index in your valueOptionRef.
+
+2. **Prioritize attributes with dependencies**: Key attributes with a \`dependencies\` array control conditionally-rendered content. Satisfy these dependencies to reveal gated UI/behavior.
+
+3. **Cover conditional branches**: Use \`codeUsages\` to identify all conditional paths that need coverage.
+
+## Update Rules
 
 ### Removed Key Attributes
 - Update scenarios using removed attributes—replace with new attributes or remove if no longer relevant
@@ -295,6 +306,7 @@ ${contextSection}## Update Rules
 
 ### Added Key Attributes
 - Create or update scenarios to demonstrate new behavior
+- Use different validValueOption indices than existing scenarios to maximize coverage
 
 ### Playwright Instructions
 - Validate existing instructions still work after code changes
@@ -306,7 +318,26 @@ ${contextSection}## Update Rules
 ### Default Scenario
 - First scenario must be "${DEFAULT_SCENARIO_NAME}" (exact name)
 - NEVER include playwrightInstructions—captures initial state only
-- Must include all current key attributes with valid values
+- Must include all current key attributes with valid values from their validValueOptions
+
+## keyAttributeInstructions Format
+Use \`dataStructurePath\` as key. Include variable name, instruction, AND a valueOptionRef that references which validValueOption you're using:
+\`\`\`json
+{
+  "useStatus().status": {
+    "localVariable": "status",
+    "instruction": "Set to 'active' mode",
+    "valueOptionRef": { "index": 0, "refType": "exact" }
+  }
+}
+\`\`\`
+
+### valueOptionRef Types
+- \`"exact"\`: Uses the validValueOption at this index exactly
+- \`"derived"\`: Based on this validValueOption but modified (include \`derivation\` field explaining how)
+- \`"custom"\`: Value not from validValueOptions list (use sparingly)
+
+The \`index\` refers to the validValueOptions array position (0 = first option).
 
 ## Response Format
 
@@ -330,7 +361,11 @@ Otherwise:
       "testName": "it(\\"shows empty message when no notifications\\")",
       "description": "User has no notifications.",
       "keyAttributeInstructions": {
-        "notifications[]": { "localVariable": "notifications[]", "instruction": "Empty array" }
+        "useData().notifications[]": {
+          "localVariable": "notifications[]",
+          "instruction": "Empty array [] to show empty state",
+          "valueOptionRef": { "index": 1, "refType": "exact" }
+        }
       }
     }
   ]
@@ -340,7 +375,7 @@ Otherwise:
 ## Rules
 - Scenario names: letters, numbers, spaces only
 - testName: Jest style with specific expectations
-- Use values from \`validValueOptions\` in key attributes list
+- Use values from \`validValueOptions\` in key attributes list (reference by index)
 - Include dependency key attributes if dependencies involved
 - React elements: simple only (\`<div>\`, \`<span>\`)
 - No external library references—plain JS/TS values

package/analyzer-template/packages/ai/src/lib/generateEntityDataStructure.ts

@@ -3,6 +3,7 @@ import isolateScopes from './isolateScopes';
 import analyzeScope from './analyzeScope';
 import { FileAnalyzer } from '~codeyam/analyze';
 import { AI, SerializableDataStructure } from '~codeyam/ai';
+import { resetScopeDataStructureMetrics } from './dataStructure/ScopeDataStructure';
 
 // import { awsLog } from '~codeyam/utils';
 
@@ -17,15 +18,23 @@ export default async function generateEntityDataStructure({
   fileAnalyzer,
   model,
 }: GenerateEntityDataStructureArgs): Promise<SerializableDataStructure> {
+  // Reset metrics counters to prevent false "infinite loop" detection when
+  // processing multiple entities in the same run
+  resetScopeDataStructureMetrics();
+
   // awsLog('CodeYam: Generating entity data structure', {
   //   filePath: entity.filePath,
   //   entityName: entity.name,
   //   // code: entity.code.length,
   // });
 
-  const fileScope = isolateScopes(fileAnalyzer.sourceFile.getText(), {
+  // Use getFullText() to include leading trivia (comments/whitespace) so that
+  // positions from isolateScopes match AST node positions which include trivia
+  const fileText = fileAnalyzer.sourceFile.getFullText();
+  const fileScope = isolateScopes(fileText, {
     entityNames: [entity.name],
   });
+
   let entityScope =
     fileScope.name === entity.name
       ? fileScope
@@ -39,12 +48,10 @@ export default async function generateEntityDataStructure({
   }
 
   if (!entityScope) {
-    if (entity.code === fileAnalyzer.sourceFile.getText()) {
+    if (entity.code === fileText) {
       entityScope = fileScope;
     } else {
-      const entityStartIndex = fileAnalyzer.sourceFile
-        .getText()
-        .indexOf(entity.code);
+      const entityStartIndex = fileText.indexOf(entity.code);
       const entityIndex = entityStartIndex + entity.code.length / 2;
       entityScope = Object.values(fileScope.childScopes).find(
         (s) => entityIndex > s.start && entityIndex < s.end,