@codeyam/codeyam-cli 0.1.0-staging.09652b8 → 0.1.0-staging.15d0f46

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,6 +82,42 @@ export default function gatherRelevantDependentKeyAttributes(
       const signatureSchema = dependencySchema?.signatureSchema;
       if (!signatureSchema) return;
 
+      const sourceEquivalencies = dependencySchema?.sourceEquivalencies ?? {};
+
+      // If sourceEquivalencies is empty but the child has mocked imports,
+      // propagate key attributes directly since they affect the parent's visual output
+      if (Object.keys(sourceEquivalencies).length === 0) {
+        // Check both importedExports and nodeModuleImports for 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 childHasMockedImports =
+          hasMockedImportedExports || hasMockedNodeModuleImports;
+
+        if (childHasMockedImports) {
+          for (const dependentKeyAttribute of dependentKeyAtttributes) {
+            relevantDependentKeyAttributes.push({
+              ...dependentKeyAttribute,
+              dependentEntityName: entityName,
+              originalEntityName: entityName,
+              filePath: path,
+            } as any);
+          }
+        }
+        return;
+      }
+
       for (const dependentKeyAttribute of dependentKeyAtttributes) {
         const dependentKeyAttributeExternalPathParts =
           splitOutsideParenthesesAndArrays(dependentKeyAttribute.externalPath);

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

@@ -29,8 +29,17 @@ 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,
+    );
 
   const prompt = generateEntityKeyAttributesGenerator({
     code: entity.code,
@@ -102,6 +111,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 +179,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 +197,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 +238,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 +256,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,

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

@@ -23,9 +23,13 @@ export default async function generateEntityDataStructure({
   //   // 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 +43,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,

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

@@ -1,6 +1,6 @@
 import completionCall from './completionCall';
 import { saveLlmCall } from '~codeyam/aws/dynamodb';
-import { Entity } from '~codeyam/types';
+import { DataStructure, Entity } from '~codeyam/types';
 import generateEntityKeyAttributesGenerator from './promptGenerators/generateEntityKeyAttributesGenerator';
 import { awsLog } from '~codeyam/utils';
 import { DEFAULT_LARGER_MODEL } from './aiConfig';
@@ -13,11 +13,13 @@ import type { ConditionalUsage } from '~codeyam/types';
 
 interface GenerateEntityKeyAttributesArgs {
   entity: Pick<Entity, 'sha' | 'name' | 'filePath' | 'code' | 'metadata'>;
+  mergedDataStructure: Omit<DataStructure, 'equivalentSignatureVariables'>;
   model?: AI.Model;
 }
 
 export default async function generateEntityKeyAttributes({
   entity,
+  mergedDataStructure,
   model,
 }: GenerateEntityKeyAttributesArgs) {
   awsLog(`CodeYam: Generating entity key attributes`, {
@@ -25,8 +27,16 @@ export default async function generateEntityKeyAttributes({
     entityName: entity.name,
   });
 
+  const isolatedDataStructure = entity.metadata?.isolatedDataStructure;
+  const equivalentSignatureVariables =
+    isolatedDataStructure?.equivalentSignatureVariables ?? {};
   const { attributesMap, associationMap, dataStructurePathMap } =
-    gatherAttributesMap(entity, true);
+    gatherAttributesMap(
+      isolatedDataStructure ?? mergedDataStructure,
+      mergedDataStructure,
+      equivalentSignatureVariables,
+      true,
+    );
 
   if (Object.keys(attributesMap).length === 0) {
     console.log(
@@ -57,8 +67,8 @@ export default async function generateEntityKeyAttributes({
     staticAnalysisContext,
   });
 
-  console.info('Generated prompt for generateEntityKeyAttributes:');
-  console.info(prompt);
+  console.log('Generated prompt for generateEntityKeyAttributes:');
+  console.log(prompt);
 
   const response = await completionCall({
     type: `generateEntityKeyAttributes`,
@@ -142,6 +152,9 @@ export default async function generateEntityKeyAttributes({
         description: keyAttribute.description,
         validValueOptions: keyAttribute.validValueOptions,
         errorValueOptions: keyAttribute.errorValueOptions,
+        // Preserve new fields from LLM response
+        valueType: keyAttribute.valueType,
+        dependencies: keyAttribute.dependencies,
       };
     });
 
@@ -167,15 +180,37 @@ NOT key attributes:
 - Values that would cause errors if changed
 - Guard clauses that render literally nothing
 
-Note: A conditional that shows a different UI state (e.g., loading screen, redirect message) IS significant—include it.
+## Conditional Key Attributes
+
+We've provided a list of variables that are used in conditional statements in the code. These are very likely key attributes.
+
+Each conditional lists it's source and how it is used. This should help you identify which attributes to include and what values to suggest.
 
 ## Path Format
 
-Only return attributes from the "Attributes and Value Types" list—these are the only values we can control. Use the exact \`fullPath\` provided. Use \`functionCallReturnValue\` for function return values. Use \`[]\` for array items (e.g., \`items[]\`).
+Only return attributes from the "Attributes and Value Types" list — these are the only values we can control. Use the exact \`fullPath\` 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.
+
+All validValueOptions must be of this type.
+
+## Valid Value Options
+
+- For \`validValueOptions\`, include boundary/extreme values that trigger different behavior
+- Each of the \`validValueOptions\` must be the exact type that we provide in "Attributes and Value Types".
+- Be very specific with values. Do not make them ambiguous in any way
+— Do NOT include values that would cause unexpected/unhandled errors.
+- Ensure all values or value ranges that will trigger different rendering behavior are included
 
 ## Response Format
 
-For \`validValueOptions\`, include boundary/extreme values that trigger different behavior—but NOT values that would cause errors.
+List all key attributes in order from most impactful to least impactful regarding how dramatically they will change the rendering of the component using the following format:
 
 \`\`\`json
 {
@@ -183,7 +218,11 @@ For \`validValueOptions\`, include boundary/extreme values that trigger differen
     {
       "fullPath": "signature[0].user",
       "description": "Controls whether the main form or redirect screen displays",
-      "validValueOptions": ["truthy user object", "null/undefined"]
+      "dependencies": [
+        "signature[0].projectId": "must be a valid project ID for 'user' to have effect"
+      ],
+      "valueType": "object",
+      "validValueOptions": ["an empty object {}", "null/undefined"]
     }
   ]
 }
@@ -205,13 +244,37 @@ NOT key attributes:
 - Values that would cause errors if changed
 - Guard clauses that return nothing/minimal data
 
+## Conditional Key Attributes
+
+We've provided a list of variables that are used in conditional statements in the code. These are very likely key attributes.
+
+Each conditional lists it's source and how it is used. This should help you identify which attributes to include and what values to suggest.
+
 ## Path Format
 
-Only return attributes from the "Attributes and Value Types" list—these are the only values we can control. Use the exact \`fullPath\` provided. Use \`functionCallReturnValue\` for function return values. Use \`[]\` for array items (e.g., \`items[]\`).
+Only return attributes from the "Attributes and Value Types" list — these are the only values we can control. Use the exact \`fullPath\` 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.
+
+All validValueOptions must be of this type.
+
+## Valid Value Options
+
+- For \`validValueOptions\`, include boundary/extreme values that trigger different behavior
+- Each of the \`validValueOptions\` must be the exact type that we provide in "Attributes and Value Types".
+- Be very specific with values. Do not make them ambiguous in any way
+— Do NOT include values that would cause unexpected/unhandled errors.
+- Ensure all values or value ranges that will trigger different rendering behavior are included
 
 ## Response Format
 
-For \`validValueOptions\`, include boundary/extreme values that trigger different behavior—but NOT values that would cause errors.
+List all key attributes in order from most impactful to least impactful using the following format:
 
 \`\`\`json
 {
@@ -219,6 +282,10 @@ For \`validValueOptions\`, include boundary/extreme values that trigger differen
     {
       "fullPath": "signature[0].config.enabled",
       "description": "Controls whether the main processing logic runs",
+      "dependencies": [
+        "signature[0].analysisMode": "must be 'full' for 'enabled' to have effect"
+      ],
+      "valueType": "boolean",
       "validValueOptions": ["true", "false"]
     }
   ]

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

@@ -284,9 +284,9 @@ export const generateSystemMessage = (
   const scenarioType = defaultScenario
     ? `## Default Scenario
 Generate COMPLETE, robust data for the entire data structure.
-- 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)
 - Arrays should have 2-3 items
-- Don't skip nested attributes
+- Don't skip nested attributes unless the key attributes specify a parent attribute should be null or undefined
 - This provides the baseline data for all other scenarios`
     : `## Non-Default Scenario
 Generate ONLY the differences from the default scenario.
@@ -317,8 +317,10 @@ Use for relative dates. Code runs in Node (no browser APIs, no external librarie
 \`\`\`
 Use simple elements only (\`<div>\`, \`<span>\`). No custom components.
 
-## 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
 \`\`\`json
 {
   "mockData": {
@@ -328,6 +330,18 @@ Preserve function call syntax exactly as written in the structure:
 }
 \`\`\`
 
+### Variable-qualified calls (for multiple calls to the same function)
+When the same function is called multiple times with results stored in different variables, keys use the format \`variableName <- functionName\`:
+\`\`\`json
+{
+  "mockData": {
+    "entityDiffFetcher <- useFetcher": { "data": null, "state": "idle" },
+    "reportFetcher <- useFetcher": { "data": { "reportId": "abc123" }, "state": "idle" }
+  }
+}
+\`\`\`
+This reads as "entityDiffFetcher receives from useFetcher". Each variable gets its own distinct mock data.
+
 ## Response Format
 \`\`\`json
 {

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

@@ -4,7 +4,7 @@ import { saveLlmCall } from '~codeyam/aws/dynamodb';
 import generateEntityPropsStructureGenerator from './promptGenerators/generateEntityPropsStructureGenerator';
 import identifyReserved, { ReservedFeature } from './identifyReserved';
 import splitOutsideParentheses from './splitOutsideParentheses';
-import { generateSha, loadStatement, saveStatement } from '~codeyam/supabase';
+import { generateSha, loadStatement, saveStatement } from '~codeyam/database';
 import { ScopeInfo } from './dataStructure/ScopeDataStructure';
 import { AI, parseJsonSafe } from '~codeyam/ai';
 

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

@@ -259,8 +259,16 @@ function getLiteralValue(expr: ts.Expression): string | undefined {
 /**
  * Format conditional usages into a string for the LLM prompt
  */
+/**
+ * Interface for conditional usages that may have source tracing information.
+ * This is a superset of ConditionalUsage that includes optional sourceDataPath.
+ */
+interface EnrichedConditionalUsageInput extends ConditionalUsage {
+  sourceDataPath?: string;
+}
+
 export function formatConditionalUsagesForPrompt(
-  conditionalUsages: Record<string, ConditionalUsage[]>,
+  conditionalUsages: Record<string, EnrichedConditionalUsageInput[]>,
 ): string {
   const entries = Object.entries(conditionalUsages);
   if (entries.length === 0) {
@@ -270,7 +278,7 @@ export function formatConditionalUsagesForPrompt(
   const lines: string[] = [
     '\n## Static Analysis: Attributes Used in Conditionals',
     '',
-    'The following attributes were found to be used in conditional statements, making them potential key attributes:',
+    'The following attributes were found to be used in conditional statements, making them very likely key attributes:',
     '',
   ];
 
@@ -281,20 +289,30 @@ export function formatConditionalUsagesForPrompt(
       .flatMap((u) => u.comparedValues || []);
     const uniqueValues = [...new Set(allComparedValues)];
 
-    let description = `- "${path}": `;
+    // Get the source data path if available (traced back to data source)
+    const sourceDataPath = usages.find((u) => u.sourceDataPath)?.sourceDataPath;
+
+    let description = `- "${path}"`;
+
+    // If we have a source data path, show that the local variable derives from a data source
+    if (sourceDataPath) {
+      description += `\n    - Source: "${sourceDataPath}"`;
+    }
+
+    description += '\n    - ';
 
     if (uniqueTypes.includes('switch')) {
-      description += 'used as switch discriminant';
+      description += 'Used as switch discriminant';
       if (uniqueValues.length > 0) {
         description += ` with cases: ${uniqueValues.map((v) => `"${v}"`).join(', ')}`;
       }
     } else if (uniqueTypes.includes('comparison')) {
-      description += 'used in comparison';
+      description += 'Used in comparison';
       if (uniqueValues.length > 0) {
         description += ` against values: ${uniqueValues.map((v) => `"${v}"`).join(', ')}`;
       }
     } else {
-      description += 'used in truthiness check (presence/absence matters)';
+      description += 'Used in truthiness check (presence/absence matters)';
     }
 
     lines.push(description);
@@ -304,6 +322,9 @@ export function formatConditionalUsagesForPrompt(
   lines.push(
     'Consider these attributes as candidates for key attributes, especially those with specific comparison values.',
   );
+  lines.push(
+    'When a local variable derives from a data source, the data source is the key attribute.',
+  );
 
   return lines.join('\n');
 }

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

@@ -12,6 +12,53 @@ export interface Statement {
 }
 
 export default function mergeStatements(statements: Statement[]) {
+  // Debug: Log input statements to understand what we're merging
+  const hasStateStructure = statements.some((s) =>
+    Object.keys(s.structure || {}).some(
+      (k) => k === 'state' || k.includes('state'),
+    ),
+  );
+  const hasStateEquiv = statements.some((s) =>
+    Object.keys(s.equivalentVariables || {}).some(
+      (k) => k === 'state' || k.includes('state'),
+    ),
+  );
+  if (hasStateStructure || hasStateEquiv) {
+    console.log(
+      `[UNION-TYPE] mergeStatements INPUT: ${statements.length} statements`,
+    );
+    statements.forEach((s, i) => {
+      const stateInStructure = Object.entries(s.structure || {}).filter(
+        ([k]) => k === 'state' || k.includes('state'),
+      );
+      const stateInEquiv = Object.entries(s.equivalentVariables || {}).filter(
+        ([k]) => k === 'state' || k.includes('state'),
+      );
+      const sigInEquiv = Object.entries(s.equivalentVariables || {}).filter(
+        ([k, v]) => k === 'state' || v === 'signature[0]',
+      );
+      if (
+        stateInStructure.length > 0 ||
+        stateInEquiv.length > 0 ||
+        sigInEquiv.length > 0
+      ) {
+        console.log(`[UNION-TYPE]   Statement ${i}:`);
+        if (stateInStructure.length > 0)
+          console.log(
+            `[UNION-TYPE]     structure: ${JSON.stringify(Object.fromEntries(stateInStructure))}`,
+          );
+        if (stateInEquiv.length > 0)
+          console.log(
+            `[UNION-TYPE]     equivalentVariables (state): ${JSON.stringify(Object.fromEntries(stateInEquiv))}`,
+          );
+        if (sigInEquiv.length > 0)
+          console.log(
+            `[UNION-TYPE]     equivalentVariables (sig): ${JSON.stringify(Object.fromEntries(sigInEquiv))}`,
+          );
+      }
+    });
+  }
+
   // This should be handled by the AST Analyzer but it currently has a bug
   const ensureFunctionCallReturnValue = (path: string, value: string) => {
     if (path.endsWith(')') && value !== 'function') {
@@ -77,6 +124,15 @@ export default function mergeStatements(statements: Statement[]) {
   ) => {
     if (!mapping[key] || mapping[key] === 'unknown') {
       mapping[key] = value;
+      return;
+    }
+
+    // Prefer union types (containing ' | ') over generic types like "string"
+    const existingIsGeneric =
+      mapping[key] === 'string' || mapping[key] === 'unknown';
+    const newIsUnion = value.includes("'") && value.includes(' | ');
+    if (existingIsGeneric && newIsUnion) {
+      mapping[key] = value;
     }
   };
 
@@ -157,5 +213,53 @@ export default function mergeStatements(statements: Statement[]) {
 
   const llmCalls = statements.map((result) => result.llmCall).filter(Boolean);
 
+  // Propagate more specific types through equivalentVariables
+  // If we have equivalentVariables[A] = B, and structure[A] is more specific than structure[B],
+  // then update structure[B] to match structure[A]
+  // This handles cases like: "state" -> "signature[0]" where "state" has a union type
+
+  // Debug: Log union types in merged structure
+  const unionTypesInStructure = Object.entries(structure).filter(
+    ([_, v]) => typeof v === 'string' && v.includes("'") && v.includes(' | '),
+  );
+  if (unionTypesInStructure.length > 0) {
+    console.log(
+      `[UNION-TYPE] mergeStatements: Found ${unionTypesInStructure.length} union types in merged structure`,
+    );
+    unionTypesInStructure.forEach(([k, v]) =>
+      console.log(`[UNION-TYPE]   structure["${k}"] = ${v}`),
+    );
+    console.log(
+      `[UNION-TYPE] mergeStatements: equivalentVariables = ${JSON.stringify(equivalentVariables)}`,
+    );
+  }
+
+  for (const [source, target] of Object.entries(equivalentVariables) as [
+    string,
+    string,
+  ][]) {
+    const sourceType = structure[source];
+    const targetType = structure[target];
+
+    // Check if source has a more specific type (union type vs generic type)
+    if (sourceType) {
+      const sourceIsUnion =
+        sourceType.includes("'") && sourceType.includes(' | ');
+
+      if (sourceIsUnion) {
+        // Target doesn't exist or is generic - propagate the union type
+        const targetIsGeneric =
+          !targetType || targetType === 'string' || targetType === 'unknown';
+
+        if (targetIsGeneric) {
+          console.log(
+            `[UNION-TYPE] mergeStatements: Propagating "${source}" (${sourceType}) -> "${target}"`,
+          );
+          structure[target] = sourceType;
+        }
+      }
+    }
+  }
+
   return { structure, equivalentVariables, llmCalls };
 }