@codeyam/codeyam-cli 0.1.0-staging.28f73cf → 0.1.0-staging.323686

package/packages/ai/src/lib/generateEntityScenarioData.js

@@ -1,13 +1,489 @@
 import completionCall from "./completionCall.js";
 import generateEntityScenarioDataGenerator from "./promptGenerators/generateEntityScenarioDataGenerator.js";
+import generateMissingKeysPrompt from "./promptGenerators/generateMissingKeysPrompt.js";
+import generateChunkPrompt from "./promptGenerators/generateChunkPrompt.js";
 import { saveLlmCall } from "../../../../packages/aws/dynamodb/index.js";
+import { trackDataSnapshot } from "./e2eDataTracking.js";
 import validateJson from "./validateJson.js";
 import { awsLog, awsLogDebugLevel } from "../../../../packages/utils/index.js";
 import { parseJsonSafe } from "../../../../packages/ai/index.js";
+import convertNullToUndefinedBySchema from "./dataStructure/helpers/convertNullToUndefinedBySchema.js";
+import fixNullIdsBySchema from "./dataStructure/helpers/fixNullIdsBySchema.js";
+import { deepMerge } from "../../../../packages/generate/index.js";
+import { chunkDataStructure, getRequiredValuesForChunk, } from "./dataStructureChunking.js";
+/**
+ * Check if any of the scenario's covered flows require error data.
+ * Returns true if any requiredValue has an error path with truthy comparison.
+ */
+function scenarioRequiresErrorData(scenario, executionFlows) {
+    const coveredFlowIds = scenario.metadata?.coveredFlows || [];
+    for (const flowId of coveredFlowIds) {
+        const flow = executionFlows?.find((f) => f.id === flowId);
+        if (!flow?.requiredValues)
+            continue;
+        for (const rv of flow.requiredValues) {
+            // Check if any requiredValue has an error path and requires it to be truthy
+            if (rv.attributePath?.toLowerCase().includes('.error') &&
+                rv.comparison === 'truthy') {
+                return true;
+            }
+        }
+    }
+    return false;
+}
+/**
+ * Deep merge scenario data with default scenario data.
+ * The scenario-specific data takes precedence, with default filling in missing fields.
+ *
+ * IMPORTANT: null values are PRESERVED (not removed) in the result.
+ * This is critical because writeMockDataTsx.ts does another deepMerge with default data,
+ * and it needs null values to prevent defaults from being filled back in.
+ * If we removed null here, the second merge would restore the defaults,
+ * making scenarios identical to the default scenario.
+ */
+function deepMergeScenarioData(defaultData, scenarioData) {
+    // Guard against non-object inputs (LLM sometimes returns primitives)
+    if (typeof scenarioData !== 'object' ||
+        scenarioData === null ||
+        Array.isArray(scenarioData)) {
+        // Return scenario value directly if it's not a mergeable object
+        return scenarioData;
+    }
+    if (typeof defaultData !== 'object' ||
+        defaultData === null ||
+        Array.isArray(defaultData)) {
+        // Return scenario value if default isn't mergeable
+        return scenarioData;
+    }
+    const result = {};
+    // Start with all keys from default
+    for (const key of Object.keys(defaultData)) {
+        if (key in scenarioData) {
+            const scenarioValue = scenarioData[key];
+            const defaultValue = defaultData[key];
+            // null means explicitly override with null (falsy value)
+            // IMPORTANT: We preserve null instead of removing the key
+            // This ensures writeMockDataTsx's deepMerge won't fill in defaults
+            if (scenarioValue === null) {
+                result[key] = null;
+                continue;
+            }
+            // Deep merge objects (but not arrays)
+            if (typeof scenarioValue === 'object' &&
+                !Array.isArray(scenarioValue) &&
+                typeof defaultValue === 'object' &&
+                !Array.isArray(defaultValue) &&
+                defaultValue !== null) {
+                result[key] = deepMergeScenarioData(defaultValue, scenarioValue);
+            }
+            else {
+                // Use scenario value (overrides default)
+                result[key] = scenarioValue;
+            }
+        }
+        else {
+            // Key not in scenario, use default
+            result[key] = defaultData[key];
+        }
+    }
+    // Add any keys that are only in scenario data (including null values)
+    for (const key of Object.keys(scenarioData)) {
+        if (!(key in defaultData)) {
+            result[key] = scenarioData[key];
+        }
+    }
+    return result;
+}
 const DEFAULT_SCENARIO_NAME = 'Default Scenario';
-export async function generateDataForScenario({ entity, structure, scenario, defaultScenarioData, incompleteResponse, analysis, model, }) {
+/**
+ * Find the path to a key within a nested dataForMocks structure.
+ * Returns the path as an array of keys, or null if not found.
+ *
+ * @example
+ * // dataForMocks = { trpc: { fastener: { "useMutation()": { isLoading: "boolean" } } } }
+ * // findKeyPath("fastener", dataForMocks) returns ["trpc"]
+ */
+function findKeyPath(targetKey, obj, currentPath = []) {
+    if (typeof obj !== 'object' || obj === null || Array.isArray(obj)) {
+        return null;
+    }
+    for (const key of Object.keys(obj)) {
+        if (key === targetKey) {
+            return currentPath;
+        }
+        // Recursively search in nested objects
+        const nested = obj[key];
+        if (typeof nested === 'object' &&
+            nested !== null &&
+            !Array.isArray(nested)) {
+            const result = findKeyPath(targetKey, nested, [
+                ...currentPath,
+                key,
+            ]);
+            if (result !== null) {
+                return result;
+            }
+        }
+    }
+    return null;
+}
+/**
+ * Relocate misplaced nested keys in mockData to their correct position
+ * based on the dataForMocks structure.
+ *
+ * When the LLM returns mockData with keys at the wrong nesting level
+ * (e.g., { trpc: { quote: {...} }, fastener: {...} } when fastener should
+ * be inside trpc), this function moves them to the correct position.
+ *
+ * This function works recursively to handle nested misplacements, not just
+ * root-level ones. For example, if getQuote is at trpc.getQuote instead of
+ * trpc.quote.getQuote, it will be relocated.
+ *
+ * @example
+ * // dataForMocks: { trpc: { quote: {...}, fastener: {...} } }
+ * // mockData: { trpc: { quote: {...} }, fastener: {...} }
+ * // After: mockData: { trpc: { quote: {...}, fastener: {...} } }
+ *
+ * @example (nested case)
+ * // dataForMocks: { trpc: { quote: { getQuote: {...} } } }
+ * // mockData: { trpc: { quote: {...}, getQuote: {...} } }
+ * // After: mockData: { trpc: { quote: { getQuote: {...} } } }
+ */
+function relocateMisplacedNestedKeys(mockData, dataForMocks, currentPathForLogging = []) {
+    if (typeof dataForMocks !== 'object' || dataForMocks === null) {
+        return;
+    }
+    const keysInSchema = Object.keys(dataForMocks);
+    const keysToRelocate = [];
+    // Find keys in mockData that are NOT at this level in dataForMocks
+    // but DO exist somewhere nested in dataForMocks
+    for (const key of Object.keys(mockData)) {
+        if (!keysInSchema.includes(key)) {
+            // This key is at this level in mockData but not at this level in dataForMocks
+            // Check if it exists somewhere nested in dataForMocks
+            const path = findKeyPath(key, dataForMocks);
+            if (path !== null && path.length > 0) {
+                keysToRelocate.push({ key, path });
+            }
+        }
+    }
+    // Relocate each misplaced key to its correct nested position
+    for (const { key, path } of keysToRelocate) {
+        const value = mockData[key];
+        // Navigate to the correct parent in mockData, creating nested objects if needed
+        let current = mockData;
+        for (const pathKey of path) {
+            if (current[pathKey] === undefined) {
+                current[pathKey] = {};
+            }
+            current = current[pathKey];
+        }
+        // Deep merge the value into the correct location
+        // Use deep merge to preserve existing data at that location
+        if (current[key] !== undefined && typeof current[key] === 'object') {
+            current[key] = deepMerge(current[key], value);
+        }
+        else {
+            current[key] = value;
+        }
+        // Remove the key from its current (wrong) level
+        delete mockData[key];
+        const fullPath = [...currentPathForLogging, ...path].join('.');
+        awsLog(`CodeYam: Relocated misplaced key "${key}" from [${currentPathForLogging.join('.')}] to [${fullPath}]`);
+    }
+    // Recursively process nested objects to handle deeply nested misplacements
+    for (const key of Object.keys(mockData)) {
+        const mockValue = mockData[key];
+        const schemaValue = dataForMocks[key];
+        // Only recurse if both mockData and schema have nested objects at this key
+        if (typeof mockValue === 'object' &&
+            mockValue !== null &&
+            !Array.isArray(mockValue) &&
+            typeof schemaValue === 'object' &&
+            schemaValue !== null &&
+            !Array.isArray(schemaValue)) {
+            relocateMisplacedNestedKeys(mockValue, schemaValue, [...currentPathForLogging, key]);
+        }
+    }
+}
+/**
+ * Generate default mock data for a schema type.
+ * Returns reasonable default values based on the schema type string.
+ */
+function generateDefaultForSchemaType(schemaType) {
+    if (typeof schemaType === 'string') {
+        // Handle common type strings
+        if (schemaType === 'function')
+            return () => { };
+        if (schemaType === 'promise')
+            return Promise.resolve();
+        if (schemaType === 'boolean')
+            return false;
+        if (schemaType === 'string')
+            return '';
+        if (schemaType === 'number')
+            return 0;
+        if (schemaType.includes('number | undefined'))
+            return undefined;
+        if (schemaType.includes('string | undefined'))
+            return undefined;
+        if (schemaType.includes('boolean | undefined'))
+            return undefined;
+        if (schemaType.includes('| undefined'))
+            return undefined;
+        if (schemaType.includes('| null'))
+            return null;
+        return schemaType; // Return the type as a string placeholder
+    }
+    if (typeof schemaType === 'object' &&
+        schemaType !== null &&
+        !Array.isArray(schemaType)) {
+        // Recursively generate defaults for nested objects
+        const result = {};
+        for (const [key, value] of Object.entries(schemaType)) {
+            result[key] = generateDefaultForSchemaType(value);
+        }
+        return result;
+    }
+    return undefined;
+}
+/**
+ * Detect if a string should be converted to an array.
+ * Returns the array if the field appears to be an array field, or null if it should remain a string.
+ *
+ * This handles two cases:
+ * 1. Comma-separated values: "color,size" -> ["color", "size"]
+ * 2. Single values for array-named fields: "Finish" -> ["Finish"]
+ */
+function parseCommaSeparatedStringAsArray(value, key) {
+    // Heuristic: if the key name suggests it's an array field, convert it
+    // Common patterns: *_attributes, *_ids, *_items, *_tags, *_values, plural names
+    // Check this FIRST because array-named fields should be converted regardless
+    // of whether they contain commas (single values become single-element arrays).
+    const arrayFieldPatterns = [
+        /_attributes$/i,
+        /_ids$/i,
+        /_items$/i,
+        /_tags$/i,
+        /_values$/i,
+        /_types$/i,
+        /_names$/i,
+        /_keys$/i,
+        /^attributes$/i,
+        /^items$/i,
+        /^tags$/i,
+        /^values$/i,
+    ];
+    const looksLikeArrayField = arrayFieldPatterns.some((pattern) => pattern.test(key));
+    if (looksLikeArrayField) {
+        // Skip newlines check - multiline values shouldn't be split
+        if (value.includes('\n')) {
+            return null;
+        }
+        // Split by comma and trim whitespace
+        const parts = value.split(',').map((s) => s.trim());
+        // Filter out empty strings - this handles both "Finish" -> ["Finish"]
+        // and "" -> []
+        return parts.filter((s) => s.length > 0);
+    }
+    // For non-array-named fields, only convert if there are commas
+    if (!value.includes(',')) {
+        return null;
+    }
+    // For non-array-named fields, apply stricter sentence detection
+    // Skip if it looks like a sentence (comma followed by space and lowercase)
+    if (/,\s+[a-z]/.test(value)) {
+        return null;
+    }
+    // Skip if it contains newlines (likely formatted text)
+    if (value.includes('\n')) {
+        return null;
+    }
+    return null;
+}
+/**
+ * Convert comma-separated string values to arrays when they look like array data.
+ * This handles cases where the LLM generates strings like "color,size" instead
+ * of arrays like ["color", "size"] due to schema type misdetection.
+ */
+function convertCommaSeparatedStringsToArrays(mockData) {
+    for (const [key, value] of Object.entries(mockData)) {
+        if (typeof value === 'string') {
+            const asArray = parseCommaSeparatedStringAsArray(value, key);
+            if (asArray !== null) {
+                mockData[key] = asArray;
+                awsLog(`CodeYam: Converted comma-separated string to array for key "${key}": "${value}" -> [${asArray.map((s) => `"${s}"`).join(', ')}]`);
+            }
+        }
+        else if (value !== null &&
+            typeof value === 'object' &&
+            !Array.isArray(value)) {
+            // Recursively process nested objects
+            convertCommaSeparatedStringsToArrays(value);
+        }
+        else if (Array.isArray(value)) {
+            // Recursively process arrays (each element could be an object)
+            for (const item of value) {
+                if (item !== null && typeof item === 'object' && !Array.isArray(item)) {
+                    convertCommaSeparatedStringsToArrays(item);
+                }
+            }
+        }
+    }
+}
+/**
+ * Ensure all keys from dataForMocks have corresponding data in mockData.
+ * For missing keys, generate default values based on the schema.
+ * Recursively checks nested objects to fill in any missing nested fields.
+ */
+function fillMissingMockDataKeysWithDefaults(mockData, dataForMocks, pathPrefix = '') {
+    if (typeof dataForMocks !== 'object' || dataForMocks === null) {
+        return;
+    }
+    const missingKeys = [];
+    for (const key of Object.keys(dataForMocks)) {
+        const fullPath = pathPrefix ? `${pathPrefix}.${key}` : key;
+        if (mockData[key] === undefined) {
+            missingKeys.push(fullPath);
+            // Generate default data based on schema
+            const schemaForKey = dataForMocks[key];
+            mockData[key] = generateDefaultForSchemaType(schemaForKey);
+        }
+        else {
+            // Key exists, but if both are objects, recursively check for missing nested keys
+            const schemaValue = dataForMocks[key];
+            const mockValue = mockData[key];
+            if (typeof schemaValue === 'object' &&
+                schemaValue !== null &&
+                !Array.isArray(schemaValue) &&
+                typeof mockValue === 'object' &&
+                mockValue !== null &&
+                !Array.isArray(mockValue)) {
+                fillMissingMockDataKeysWithDefaults(mockValue, schemaValue, fullPath);
+            }
+        }
+    }
+    if (missingKeys.length > 0) {
+        awsLog(`CodeYam: Generated default mock data for ${missingKeys.length} missing key(s): ${missingKeys.slice(0, 10).join(', ')}${missingKeys.length > 10 ? '...' : ''}`);
+    }
+}
+/**
+ * For Default Scenario only: detect missing mockData keys and make a follow-up
+ * LLM call to fill them in. This handles cases where the LLM completes normally
+ * but misses some keys (often small/simple ones when the schema is large).
+ */
+async function fillMissingMockDataKeys({ structure, scenario, executionFlows, fullScenarioData, model, }) {
+    if (!structure.dataForMocks ||
+        typeof structure.dataForMocks !== 'object' ||
+        Array.isArray(structure.dataForMocks)) {
+        return;
+    }
+    const expectedKeys = Object.keys(structure.dataForMocks);
+    const generatedKeys = Object.keys(fullScenarioData.data.mockData || {});
+    const missingKeys = expectedKeys.filter((k) => !generatedKeys.includes(k));
+    if (missingKeys.length === 0) {
+        return;
+    }
+    awsLog(`Default Scenario missing ${missingKeys.length} keys, making follow-up call`, { missingKeys });
+    // Build subset schema with only missing keys
+    const missingSchema = {};
+    for (const key of missingKeys) {
+        missingSchema[key] = structure.dataForMocks[key];
+    }
+    const followUpPrompt = generateMissingKeysPrompt({
+        scenario,
+        executionFlows,
+        generatedMockData: fullScenarioData.data.mockData || {},
+        missingSchema,
+    });
+    const followUpResponse = await completionCall({
+        type: 'generateMissingMockData',
+        systemMessage: generateMissingKeysSystemMessage(),
+        prompt: followUpPrompt,
+        model,
+    });
+    if (!followUpResponse.completion) {
+        return;
+    }
+    const followUpJson = validateJson(followUpResponse.completion);
+    const followUpParsed = parseJsonSafe(followUpJson);
+    if (!followUpParsed ||
+        typeof followUpParsed !== 'object' ||
+        !('mockData' in followUpParsed)) {
+        return;
+    }
+    const followUpMockData = followUpParsed.mockData;
+    if (followUpMockData && typeof followUpMockData === 'object') {
+        fullScenarioData.data.mockData = {
+            ...fullScenarioData.data.mockData,
+            ...followUpMockData,
+        };
+    }
+}
+export async function generateDataForScenario({ entity, structure, scenario, executionFlows, defaultScenarioData, incompleteResponse, analysis, model, }) {
+    var _a;
     awsLogDebugLevel(1, `Generating data for ${entity.name}: ${scenario.name}`);
-    const prompt = generateEntityScenarioDataGenerator(structure, scenario, defaultScenarioData, incompleteResponse);
+    // Check if we should chunk the data structure for focused processing
+    let chunkedMockData;
+    const coveredFlowIds = scenario.metadata?.coveredFlows || [];
+    if (structure.dataForMocks &&
+        !incompleteResponse // Don't do chunked calls on continuation
+    ) {
+        const chunks = chunkDataStructure(structure.dataForMocks);
+        // If we have multiple chunks, process each one with a focused call
+        if (chunks.length > 1) {
+            awsLog(`Data structure has ${Object.keys(structure.dataForMocks).length} keys, splitting into ${chunks.length} chunks for focused processing`);
+            chunkedMockData = {};
+            for (let i = 0; i < chunks.length; i++) {
+                const chunk = chunks[i];
+                const chunkKeys = Object.keys(chunk || {});
+                // Get relevant requiredValues for this chunk
+                const relevantRequiredValues = getRequiredValuesForChunk(chunk, executionFlows || [], coveredFlowIds);
+                awsLog(`Processing chunk ${i + 1}/${chunks.length}: ${chunkKeys.join(', ')}`);
+                const chunkPrompt = generateChunkPrompt({
+                    scenario,
+                    chunk,
+                    chunkIndex: i,
+                    totalChunks: chunks.length,
+                    relevantRequiredValues,
+                });
+                const chunkResponse = await completionCall({
+                    type: 'generateChunkMockData',
+                    systemMessage: generateChunkSystemMessage(scenario.name),
+                    prompt: chunkPrompt,
+                    model,
+                });
+                // Save chunk call to LLM log for replay support
+                await saveLlmCall({
+                    object_type: 'analysis',
+                    object_id: analysis.id,
+                    propsJson: {
+                        entity: { name: entity.name, filePath: entity.filePath },
+                        scenario: { name: scenario.name },
+                        chunkIndex: i,
+                        totalChunks: chunks.length,
+                    },
+                    ...chunkResponse.stats,
+                });
+                if (chunkResponse.completion) {
+                    const validJson = validateJson(chunkResponse.completion);
+                    const parsed = parseJsonSafe(validJson);
+                    if (parsed && typeof parsed === 'object' && 'mockData' in parsed) {
+                        const chunkMockData = parsed.mockData;
+                        if (chunkMockData && typeof chunkMockData === 'object') {
+                            Object.assign(chunkedMockData, chunkMockData);
+                            awsLog(`Chunk ${i + 1} generated data for: ${Object.keys(chunkMockData).join(', ')}`);
+                        }
+                    }
+                }
+            }
+            awsLog(`Chunked processing complete. Generated ${Object.keys(chunkedMockData).length} keys total`);
+        }
+    }
+    // When we have chunked mock data with actual content, tell the main prompt to skip mockData generation
+    // Important: Check for actual keys, not just truthy object, because {} would skip generation incorrectly
+    const hasChunkedData = chunkedMockData && Object.keys(chunkedMockData).length > 0;
+    const prompt = generateEntityScenarioDataGenerator(structure, scenario, executionFlows, defaultScenarioData, incompleteResponse, { mockDataAlreadyGenerated: hasChunkedData });
     const isDefault = scenario.name === DEFAULT_SCENARIO_NAME;
     const response = await completionCall({
         type: 'generateEntityScenarioData',
@@ -45,46 +521,79 @@ export async function generateDataForScenario({ entity, structure, scenario, def
         },
         ...response.stats,
     });
-    const { completion, finishReason } = response;
+    let { completion, finishReason } = response;
     if (!completion) {
         console.log('CodeYam Error: Example data generation failed: No response from AI');
         return null;
     }
-    awsLog(`LLMCall ${llmCall ? llmCall.id : 'N/A'}: ${entity.filePath} ${entity.metadata?.exportAlias ?? entity.name} scenario data completion :>> Finish reason:`, {
-        finishReason,
-        completion,
-    });
+    awsLogDebugLevel(1, `LLMCall ${llmCall ? llmCall.id : 'N/A'}: ${entity.filePath} ${entity.metadata?.exportAlias ?? entity.name} finishReason: ${finishReason}`);
+    // If response was truncated due to token limit, make a continuation call
+    if (finishReason === 'length') {
+        awsLogDebugLevel(1, 'Response truncated, making continuation call');
+        const continuationResponse = await completionCall({
+            type: 'generateEntityScenarioData',
+            systemMessage: generateIncompleteSystemMessage(scenario.name, isDefault),
+            prompt: completion, // Pass the incomplete response as the prompt
+            model,
+        });
+        if (continuationResponse.completion) {
+            completion = completion + continuationResponse.completion;
+            finishReason = continuationResponse.finishReason;
+        }
+    }
     const validJson = validateJson(completion);
     const parsed = parseJsonSafe(validJson);
     if (!parsed || typeof parsed !== 'object' || !('scenarioData' in parsed)) {
-        console.log('CodeYam Debug: generateDataForScenario failed to parse', {
-            entityName: entity.name,
-            scenarioName: scenario.name,
-            hasParsed: !!parsed,
-            parsedType: typeof parsed,
-            hasScenarioData: parsed && 'scenarioData' in parsed,
-            completionPreview: completion.substring(0, 200),
-        });
+        awsLog(`Failed to parse scenario data for ${entity.name}/${scenario.name}`);
         return null;
     }
-    const { scenarioData: scenarioDataWithoutDescription } = parsed;
-    console.log('CodeYam Debug: generateDataForScenario parsed successfully', {
-        entityName: entity.name,
-        scenarioName: scenario.name,
-        parsedScenarioName: scenarioDataWithoutDescription.scenarioName,
-        hasData: !!scenarioDataWithoutDescription.data,
-        dataKeys: scenarioDataWithoutDescription.data
-            ? Object.keys(scenarioDataWithoutDescription.data)
-            : [],
-    });
+    let { scenarioData: scenarioDataWithoutDescription } = parsed;
+    // FIX: LLM sometimes puts mock data keys directly under scenarioData instead of
+    // under scenarioData.data.mockData. Detect and fix this structural issue.
+    if (structure.dataForMocks) {
+        const scenarioDataAsAny = scenarioDataWithoutDescription;
+        const reservedKeys = new Set([
+            'scenarioName',
+            'data',
+            'scenarioDescription',
+        ]);
+        const misplacedKeys = [];
+        // Find keys that are directly under scenarioData but should be in mockData
+        for (const key of Object.keys(scenarioDataAsAny)) {
+            if (reservedKeys.has(key))
+                continue;
+            // If this key exists in the dataForMocks schema, it's misplaced
+            if (key in structure.dataForMocks) {
+                misplacedKeys.push(key);
+            }
+        }
+        if (misplacedKeys.length > 0) {
+            // Ensure data.mockData exists
+            if (!scenarioDataAsAny.data) {
+                scenarioDataAsAny.data = {};
+            }
+            if (!scenarioDataAsAny.data.mockData) {
+                scenarioDataAsAny.data.mockData = {};
+            }
+            // Move misplaced keys to mockData
+            for (const key of misplacedKeys) {
+                scenarioDataAsAny.data.mockData[key] = scenarioDataAsAny[key];
+                delete scenarioDataAsAny[key];
+            }
+            // Update the reference
+            scenarioDataWithoutDescription =
+                scenarioDataAsAny;
+        }
+    }
     const fullScenarioData = {
         ...scenarioDataWithoutDescription,
         scenarioDescription: scenario?.description ?? '',
     };
     if (structure.dataForMocks && !fullScenarioData.data.argumentsData) {
         fullScenarioData.data.argumentsData = [];
-        if (structure.arguments && !fullScenarioData.data.argumentsData) {
-            fullScenarioData.data.argumentsData = [];
+        // Populate argumentsData from structure.arguments using top-level data values
+        // Bug fix: removed redundant !argumentsData check that was always false after setting to []
+        if (structure.arguments) {
             for (let i = 0; i < structure.arguments.length; ++i) {
                 if (!fullScenarioData.data.argumentsData[i]) {
                     fullScenarioData.data.argumentsData[i] = {};
@@ -96,19 +605,98 @@ export async function generateDataForScenario({ entity, structure, scenario, def
             }
         }
     }
-    if (structure.dataForMocks && !fullScenarioData.data.mockData) {
-        fullScenarioData.data.mockData = {};
-        for (const propKey of Object.keys(structure.arguments)) {
-            const dataAsAny = fullScenarioData.data;
-            fullScenarioData.data.mockData[propKey] = dataAsAny[propKey];
+    // Merge flat-level mock data into mockData.
+    // Sometimes the LLM returns some data inside data.mockData but other data at the flat
+    // data level (e.g., data.useRouter() instead of data.mockData.useRouter()).
+    // This code ensures all dataForMocks keys end up in mockData.
+    if (structure.dataForMocks) {
+        (_a = fullScenarioData.data).mockData || (_a.mockData = {});
+        const dataAsAny = fullScenarioData.data;
+        for (const propKey of Object.keys(structure.dataForMocks)) {
+            // Only copy if it exists at flat level and not already in mockData
+            if (dataAsAny[propKey] !== undefined &&
+                !fullScenarioData.data.mockData[propKey]) {
+                fullScenarioData.data.mockData[propKey] = dataAsAny[propKey];
+            }
+        }
+    }
+    // Merge chunked mock data from focused calls (takes priority over main call's data)
+    // This ensures keys processed with focused attention are correctly generated.
+    if (chunkedMockData && fullScenarioData.data.mockData) {
+        for (const [key, value] of Object.entries(chunkedMockData)) {
+            // Chunked data takes priority - overwrite main call's potentially wrong data
+            fullScenarioData.data.mockData[key] = value;
+        }
+        awsLog(`Merged chunked mock data for keys: ${Object.keys(chunkedMockData).join(', ')}`);
+    }
+    // Relocate misplaced nested keys to their correct position.
+    // The LLM sometimes places nested keys at root level instead of inside their
+    // parent object (e.g., 'fastener' at root instead of inside 'trpc').
+    // This ensures the mockData structure matches the dataForMocks schema.
+    if (structure.dataForMocks && fullScenarioData.data.mockData) {
+        relocateMisplacedNestedKeys(fullScenarioData.data.mockData, structure.dataForMocks);
+    }
+    // Convert null values to undefined based on schema type constraints.
+    // LLM uses null for "no value" (JSON doesn't support undefined), but TypeScript
+    // types like "string | undefined" don't accept null. This converts null→undefined
+    // for fields typed as "T | undefined" (but preserves null for "T | null").
+    if (structure.dataForMocks && fullScenarioData.data.mockData) {
+        convertNullToUndefinedBySchema(fullScenarioData.data.mockData, structure.dataForMocks);
+    }
+    // Convert comma-separated strings to arrays when appropriate.
+    // The LLM sometimes generates strings like "color,size" instead of arrays
+    // like ["color", "size"] when the schema type is incorrectly inferred as
+    // 'string' instead of 'string[]'. This causes runtime errors when code
+    // calls array methods like .map() on the value.
+    if (fullScenarioData.data.mockData) {
+        convertCommaSeparatedStringsToArrays(fullScenarioData.data.mockData);
+    }
+    // Fix null values for ID fields when the schema indicates they should be non-null.
+    // The LLM sometimes generates `null` for ID fields (e.g., `"id": null`) when
+    // the schema type is `"number"`. This causes runtime issues when code checks
+    // `if (!data?.id)` expecting a truthy value.
+    if (structure.dataForMocks && fullScenarioData.data.mockData) {
+        fixNullIdsBySchema(fullScenarioData.data.mockData, structure.dataForMocks);
+    }
+    if (structure.arguments && fullScenarioData.data.argumentsData) {
+        for (let i = 0; i < fullScenarioData.data.argumentsData.length; i++) {
+            if (structure.arguments[i]) {
+                convertNullToUndefinedBySchema(fullScenarioData.data.argumentsData[i], structure.arguments[i]);
+            }
         }
     }
+    // For Default Scenario only: check for missing keys and make follow-up call if needed.
+    // This tries to get better-quality data via LLM before falling back to defaults.
+    if (isDefault) {
+        await fillMissingMockDataKeys({
+            structure,
+            scenario,
+            executionFlows,
+            fullScenarioData,
+            model,
+        });
+    }
+    // Fill in missing mock data keys with default values (after trying LLM follow-up).
+    // The LLM sometimes doesn't generate data for all keys in large schemas.
+    // This ensures all dataForMocks keys have corresponding data to prevent
+    // runtime errors like "Cannot read properties of undefined".
+    // Only run for Default Scenario - non-default scenarios will get missing
+    // data filled in from the merge with default scenario data.
+    if (isDefault && structure.dataForMocks && fullScenarioData.data.mockData) {
+        fillMissingMockDataKeysWithDefaults(fullScenarioData.data.mockData, structure.dataForMocks);
+    }
+    // Track the final scenario data for E2E debugging
+    trackDataSnapshot('generateDataForScenario_result', {
+        scenarioName: scenario.name,
+        mockData: fullScenarioData.data.mockData,
+        argumentsData: fullScenarioData.data.argumentsData,
+    }, entity.name, scenario.name);
     return {
         scenarioData: fullScenarioData,
         llmCall: { name: scenario.name, id: llmCall.id },
     };
 }
-export default async function generateEntityScenarioData({ entity, structure, scenarios, incompleteResponse, analysis, model, }) {
+export default async function generateEntityScenarioData({ entity, structure, scenarios, executionFlows, incompleteResponse, analysis, model, }) {
     if (scenarios.length === 0) {
         return { scenarioDatas: [], llmCalls: [] };
     }
@@ -120,6 +708,7 @@ export default async function generateEntityScenarioData({ entity, structure, sc
             entity,
             structure,
             scenario: defaultScenario,
+            executionFlows,
             incompleteResponse,
             analysis,
             model,
@@ -136,6 +725,7 @@ export default async function generateEntityScenarioData({ entity, structure, sc
                 entity,
                 structure,
                 scenario,
+                executionFlows,
                 defaultScenarioData,
                 incompleteResponse,
                 analysis,
@@ -148,28 +738,59 @@ export default async function generateEntityScenarioData({ entity, structure, sc
         if (nullCount > 0) {
             awsLog(`⚠️ Warning: ${nullCount} of ${results.length} non-default scenarios failed to generate data for ${entity.name}`);
         }
-        scenarioDatas.push(...validResults.map((result) => result.scenarioData));
-        llmCalls.push(...validResults.map((result) => result.llmCall));
-        console.log('CodeYam Debug: generateEntityScenarioData results', {
-            filePath: entity.filePath,
-            entityName: entity.name,
-            totalScenarios: scenarios.length,
-            defaultScenarioGenerated: !!defaultScenarioResult,
-            otherScenariosRequested: scenarios.length - 1,
-            otherScenariosResults: results.length,
-            nullResults: results.filter((r) => r === null).length,
-            validResults: validResults.length,
-            finalScenarioDatasCount: scenarioDatas.length,
+        // Merge non-default scenario data with default scenario data
+        // The LLM generates partial data (only differences), we need to merge with default
+        const mergedScenarioDatas = validResults.map((result) => {
+            const scenarioData = result.scenarioData;
+            // Merge mockData with default mockData
+            if (defaultScenarioData.data?.mockData && scenarioData.data?.mockData) {
+                scenarioData.data.mockData = deepMergeScenarioData(defaultScenarioData.data.mockData, scenarioData.data.mockData);
+            }
+            else if (defaultScenarioData.data?.mockData &&
+                !scenarioData.data?.mockData) {
+                // Use default mockData if scenario has none
+                scenarioData.data.mockData = { ...defaultScenarioData.data.mockData };
+            }
+            // Merge argumentsData with default argumentsData
+            if (defaultScenarioData.data?.argumentsData &&
+                Array.isArray(defaultScenarioData.data.argumentsData) &&
+                scenarioData.data?.argumentsData &&
+                Array.isArray(scenarioData.data.argumentsData)) {
+                for (let i = 0; i < defaultScenarioData.data.argumentsData.length; i++) {
+                    const scenarioArg = scenarioData.data.argumentsData[i];
+                    const defaultArg = defaultScenarioData.data.argumentsData[i];
+                    // Only merge if both are objects (LLM sometimes returns primitives)
+                    if (scenarioArg &&
+                        typeof scenarioArg === 'object' &&
+                        !Array.isArray(scenarioArg) &&
+                        defaultArg &&
+                        typeof defaultArg === 'object' &&
+                        !Array.isArray(defaultArg)) {
+                        scenarioData.data.argumentsData[i] = deepMergeScenarioData(defaultArg, scenarioArg);
+                    }
+                    else if (scenarioArg !== undefined) {
+                        // Keep the scenario value as-is (even if primitive)
+                        scenarioData.data.argumentsData[i] = scenarioArg;
+                    }
+                    else if (defaultArg && typeof defaultArg === 'object') {
+                        // Use default if scenario is undefined and default is an object
+                        scenarioData.data.argumentsData[i] = { ...defaultArg };
+                    }
+                    else {
+                        scenarioData.data.argumentsData[i] = defaultArg;
+                    }
+                }
+            }
+            else if (defaultScenarioData.data?.argumentsData &&
+                !scenarioData.data?.argumentsData) {
+                // Use default argumentsData if scenario has none
+                scenarioData.data.argumentsData =
+                    defaultScenarioData.data.argumentsData.map((arg) => ({ ...arg }));
+            }
+            return scenarioData;
         });
-        awsLog('CodeYam: scenarioDatas :>> ', JSON.stringify({
-            filePath: entity.filePath,
-            entityName: entity.name,
-            scenarioDatas: scenarioDatas.map((sd) => ({
-                scenarioName: sd.scenarioName,
-                hasData: !!sd.data,
-                dataKeys: sd.data ? Object.keys(sd.data) : [],
-            })),
-        }, null, 2));
+        scenarioDatas.push(...mergedScenarioDatas);
+        llmCalls.push(...validResults.map((result) => result.llmCall));
         return { scenarioDatas, llmCalls };
     }
     catch (error) {
@@ -177,27 +798,81 @@ export default async function generateEntityScenarioData({ entity, structure, sc
         throw error;
     }
 }
-export const generateSystemMessage = (scenarioName, defaultScenario) => {
+export const generateSystemMessage = (scenarioName, defaultScenario, requiresErrorData = false) => {
     const scenarioType = defaultScenario
         ? `## Default Scenario
 Generate COMPLETE, robust data for the entire data structure.
 - Fill ALL fields with realistic values (except error attributes)
-- Arrays should have 2-3 items
-- Don't skip nested attributes
+- Do not skip any keys, even simple or small entries
+- Arrays should have 3-5 items to provide realistic test data variety
+- Don't skip nested attributes unless the execution flow requirements 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.
 - Include only fields that need to change
 - Set to \`null\` to remove data
 - Omit unchanged fields—they merge from default`;
+    // Only include the "NO ERROR DATA" instruction when the scenario doesn't require error data
+    const noErrorDataInstruction = requiresErrorData
+        ? `## IMPORTANT: ERROR DATA REQUIRED
+This scenario tests error handling. You MUST include "error" fields with realistic error messages.
+- Set error fields to truthy values (e.g., "Error: Operation failed" or "Something went wrong")
+- The error data is REQUIRED to trigger the correct error UI state`
+        : `## CRITICAL: NO ERROR DATA
+NEVER include "error" fields in responses. Skip them entirely.
+- If structure has \`{ data: {...}, error: {...} }\`, only fill \`data\`
+- Leave out any attribute named "error"—do not set to null, omit entirely`;
     return `You are a test data generator. Create mock data matching a data structure and scenario requirements.
 
+## Execution Flow Requirements
+Each scenario has \`coveredFlows\` which lists the execution flows (distinct outcomes/behaviors) this scenario should demonstrate.
+Each flow has \`requiredValues\` - the attribute values that MUST be set to produce that outcome.
+
+**Your job**: Generate mock data that satisfies ALL the requiredValues from ALL coveredFlows.
+
+For example, if a flow requires:
+\`\`\`json
+{
+  "attributePath": "signature[0].isLoading",
+  "value": "false",
+  "comparison": "equals",
+  "valueType": "boolean"
+}
+\`\`\`
+Then set \`isLoading: false\` in the mockData.
+
+### Array Length Requirements (length< and length>)
+For array size variation flows:
+- \`comparison: "length<"\` with \`value: "0"\` → generate EMPTY array \`[]\`
+- \`comparison: "length<"\` with \`value: "3"\` → generate 1-2 items (few items)
+- \`comparison: "length>"\` with \`value: "10"\` → generate 12+ items (many items)
+
+### String Length Requirements (normal vs long)
+For text length variation flows:
+- \`value: "normal"\` with \`valueType: "string"\` → generate normal length text (10-50 chars)
+- \`value: "long"\` with \`valueType: "string"\` → generate LONG text (200+ chars) to test overflow/truncation
+
+## CRITICAL: Blocking Flows to Avoid
+If the scenario includes \`blockingFlowsToAvoid\`, these are flows (like modals, overlays) that would BLOCK the expected UI.
+You MUST generate mock data that PREVENTS these flows from triggering:
+
+- For \`comparison: "truthy"\` requirements → set the value to \`false\`, \`null\`, \`undefined\`, or \`0\`
+- For \`comparison: "exists"\` requirements → set the value to \`null\` or omit it entirely
+- For \`comparison: "equals"\` requirements → set a DIFFERENT value than what's required
+
+For example, if a blocking flow has:
+\`\`\`json
+{
+  "attributePath": "useFetcher().data.success",
+  "value": "true",
+  "comparison": "truthy"
+}
+\`\`\`
+Then you MUST set \`useFetcher().data\` to \`null\` or \`{ success: false }\` to prevent the modal from appearing.
+
 ${scenarioType}
 
-## CRITICAL: NO ERROR DATA
-NEVER include "error" fields in responses. Skip them entirely.
-- If structure has \`{ data: {...}, error: {...} }\`, only fill \`data\`
-- Leave out any attribute named "error"—do not set to null, omit entirely
+${noErrorDataInstruction}
 
 ## Special Markers
 
@@ -213,16 +888,19 @@ 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:
-\`\`\`json
-{
-  "mockData": {
-    "useUser()": { "user": { "name": "John" } },
-    "from().select()": [{ "id": "1" }]
-  }
-}
-\`\`\`
+### Arrays
+- Arrays should have many items (at least 4) unless specified otherwise
+- Each item must follow the exact structure provided
+- In general we want robust data, not minimal data unless specified otherwise
+
+## CRITICAL: Preserve Exact Structure
+Your response MUST mirror the EXACT nested structure provided in mockData Structure.
+- Do NOT reorganize, split, or create duplicate keys
+- The hierarchy of nested objects must match exactly what was provided unless overridden by scenario rules
+- Only change the leaf VALUES (replacing type descriptions like "string" with actual data like "hello")
+- Copy the key strings EXACTLY from the structure
+- Do NOT modify type parameters, arguments, or any part of the key
+- The keys preserve the exact function call as written in the original code
 
 ## Response Format
 \`\`\`json
@@ -240,9 +918,15 @@ Preserve function call syntax exactly as written in the structure:
 ## Rules
 - Valid JSON only—no raw code outside markers
 - No \`undefined\`—use \`null\` or omit
-- No data references (can't use \`posts[0]\` elsewhere—duplicate the value)
+- No data references (can't use \`posts[0]\` elsewhere — duplicate the value)
 - Scenario name must match exactly: "${scenarioName}"
 - Empty mockData: \`{}\`, empty argumentsData: \`[]\`
+
+## IMPORTANT: Avoid Identifier Collisions
+When generating identifier values (SHA hashes, entity IDs, etc.):
+- Use DISTINCT values for different identifier fields
+- Avoid matching: if \`entity.sha\` is "abc123", arrays like \`jobs[].entityShas\` or \`currentlyExecuting.entityShas\` should NOT contain "abc123"
+- This prevents accidental blocking of UI conditionals that check if IDs are in/not-in arrays
 `;
 };
 export const generateIncompleteSystemMessage = (scenarioName, isDefault) => `Your previous response provided us with an incomplete json object.
@@ -254,4 +938,154 @@ Here is the original system message as well:
 ${generateSystemMessage(scenarioName, isDefault)}
 \`\`\`
 `;
+/**
+ * System message for follow-up calls to generate missing mockData keys.
+ * Includes the same rules as the main system message but with a simpler response format.
+ */
+export const generateMissingKeysSystemMessage = () => `You are completing mock data generation for the Default Scenario. The initial response was missing some keys.
+
+Generate data ONLY for the missing keys provided in the prompt. Do not skip any of them.
+
+- Scenario name must match exactly: "Default Scenario"
+
+## Special Markers
+
+### Dynamic Dates (\`~~codeyam-code~~\`)
+\`\`\`json
+{ "createdAt": { "~~codeyam-code~~": "new Date(Date.now() - 24*60*60*1000)" } }
+\`\`\`
+Use for relative dates. Code runs in Node (no browser APIs, no external libraries).
+
+### JSX Children (\`~~codeyam-jsx~~\`)
+\`\`\`json
+{ "children": { "~~codeyam-jsx~~": "<div>Hello</div>" } }
+\`\`\`
+Use simple elements only (\`<div>\`, \`<span>\`). No custom components.
+
+### Arrays
+- Arrays should have many items (at least 4) unless specified otherwise
+- Each item must follow the exact structure provided
+
+## CRITICAL: Preserve Exact Structure
+Your response MUST mirror the EXACT nested structure provided for the missing keys.
+- Only change the leaf VALUES (replacing type descriptions like "string" with actual data)
+- Copy the key strings EXACTLY from the structure
+- Do NOT modify type parameters, arguments, or any part of the key
+
+## CRITICAL: NO ERROR DATA
+NEVER include "error" fields in responses. Skip them entirely.
+- If structure has \`{ data: {...}, error: {...} }\`, only fill \`data\`
+- Leave out any attribute named "error"—do not set to null, omit entirely
+
+## Response Format
+\`\`\`json
+{
+  "mockData": {
+    // fill in ONLY the missing keys
+  }
+}
+\`\`\`
+
+## Rules
+- Valid JSON only—no raw code outside markers
+- No \`undefined\`—use \`null\` or omit
+- No data references (can't use \`posts[0]\` elsewhere — duplicate the value)
+`;
+/**
+ * System message for focused calls to generate critical mockData keys.
+ * These are keys referenced by the scenario's execution flow requiredValues.
+ */
+export const generateCriticalKeysSystemMessage = (scenarioName) => `You are generating mock data for CRITICAL keys that control scenario behavior.
+
+These keys are referenced by the execution flow's requiredValues - they directly determine
+what the component renders. Pay EXTRA attention to matching the exact structure and values.
+
+- Scenario name must match exactly: "${scenarioName}"
+
+## CRITICAL: Special Characters in Keys
+Keys like \`*\` are LITERAL string keys, NOT wildcards or patterns.
+- If the schema shows \`{ "*": "string" }\`, generate \`{ "*": "some value" }\`
+- Do NOT interpret \`*\` as "any key" - use it as an actual key name
+
+## CRITICAL: Preserve Exact Structure
+Your response MUST mirror the EXACT nested structure provided.
+- Copy key strings EXACTLY as shown (including special characters)
+- Only change leaf VALUES (replacing type descriptions with actual data)
+- Do NOT modify keys, type parameters, or add extra keys
+
+## Matching requiredValues
+When the prompt shows requiredValues like:
+- \`attributePath: "useParams().functionCallReturnValue.*"\`
+- \`value: "scenarios"\`
+
+This means set the \`*\` key to include "scenarios". For URL paths split by \`/\`,
+generate a path like \`"scenarios/id/mode"\` where segments match requirements.
+
+## Response Format
+\`\`\`json
+{
+  "scenarioData": {
+    "scenarioName": "${scenarioName}",
+    "data": {
+      "mockData": {
+        // generate data for ONLY the critical keys
+      }
+    }
+  }
+}
+\`\`\`
+
+## Rules
+- Valid JSON only
+- No \`undefined\`—use \`null\` or omit
+- Match the exact schema structure provided
+`;
+/**
+ * System message for focused calls to generate mock data for a chunk of keys.
+ * Used when data structures are large and need to be processed in smaller pieces.
+ */
+export const generateChunkSystemMessage = (scenarioName) => `You are generating mock data for a SUBSET of keys from a larger data structure.
+
+This chunk contains fewer keys so you can focus on generating HIGH QUALITY data for each one.
+Pay EXTRA attention to matching the exact structure and values for each key.
+
+- Scenario name must match exactly: "${scenarioName}"
+
+## CRITICAL: Special Characters in Keys
+Keys like \`*\` are LITERAL string keys, NOT wildcards or patterns.
+- If the schema shows \`{ "*": "string" }\`, generate \`{ "*": "some value" }\`
+- Do NOT interpret \`*\` as "any key" - use it as an actual key name
+
+## CRITICAL: Preserve Exact Structure
+Your response MUST mirror the EXACT nested structure provided.
+- Copy key strings EXACTLY as shown (including special characters)
+- Only change leaf VALUES (replacing type descriptions with actual data)
+- Do NOT modify keys, type parameters, or add extra keys
+
+## Matching requiredValues
+If the prompt includes requiredValues, these are specific values that MUST be set:
+- For \`attributePath: "useParams().functionCallReturnValue.*"\` with \`value: "scenarios"\`
+  → Set the \`*\` key to include "scenarios" (e.g., "scenarios/id/mode")
+- For URL paths, generate realistic paths that satisfy the requirements
+
+## CRITICAL: NO ERROR DATA
+NEVER include "error" fields in responses. Skip them entirely.
+- If structure has \`{ data: {...}, error: {...} }\`, only fill \`data\`
+- Leave out any attribute named "error"—do not set to null, omit entirely
+
+## Response Format
+\`\`\`json
+{
+  "mockData": {
+    // generate data for ONLY the keys in this chunk
+  }
+}
+\`\`\`
+
+## Rules
+- Valid JSON only
+- No \`undefined\`—use \`null\` or omit
+- Generate data for ALL keys in the chunk (don't skip any)
+- Arrays should have many items (at least 4) unless specified otherwise
+`;
 //# sourceMappingURL=generateEntityScenarioData.js.map