@codeyam/codeyam-cli 0.1.0-staging.1669d45 → 0.1.0-staging.1a2737b

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

@@ -1,7 +1,9 @@
 import completionCall from './completionCall';
 import generateEntityScenarioDataGenerator from './promptGenerators/generateEntityScenarioDataGenerator';
 import generateMissingKeysPrompt from './promptGenerators/generateMissingKeysPrompt';
+import generateChunkPrompt from './promptGenerators/generateChunkPrompt';
 import { saveLlmCall } from '~codeyam/aws/dynamodb';
+import { trackDataSnapshot } from './e2eDataTracking';
 import type {
   Analysis,
   Entity,
@@ -14,6 +16,15 @@ import validateJson from './validateJson';
 import { awsLog, awsLogDebugLevel } from '~codeyam/utils';
 import { AI, parseJsonSafe } from '~codeyam/ai';
 import convertNullToUndefinedBySchema from './dataStructure/helpers/convertNullToUndefinedBySchema';
+import convertTypeAnnotationsToValues from './dataStructure/helpers/convertTypeAnnotationsToValues';
+import fixNullIdsBySchema from './dataStructure/helpers/fixNullIdsBySchema';
+import coerceObjectsToPrimitivesBySchema from './dataStructure/helpers/coerceObjectsToPrimitivesBySchema';
+import { JsonTypeDefinition } from '~codeyam/types';
+import { deepMerge } from '~codeyam/generate';
+import {
+  chunkDataStructure,
+  getRequiredValuesForChunk,
+} from './dataStructureChunking';
 
 /**
  * Check if any of the scenario's covered flows require error data.
@@ -119,6 +130,528 @@ function deepMergeScenarioData(
 
 const DEFAULT_SCENARIO_NAME = 'Default Scenario';
 
+/**
+ * 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: string,
+  obj: JsonTypeDefinition,
+  currentPath: string[] = [],
+): string[] | null {
+  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 as JsonTypeDefinition, [
+        ...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: Record<string, unknown>,
+  dataForMocks: JsonTypeDefinition,
+  currentPathForLogging: string[] = [],
+): void {
+  if (typeof dataForMocks !== 'object' || dataForMocks === null) {
+    return;
+  }
+
+  const keysInSchema = Object.keys(dataForMocks);
+  const keysToRelocate: { key: string; path: string[] }[] = [];
+
+  // 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: Record<string, unknown> = mockData;
+    for (const pathKey of path) {
+      if (current[pathKey] === undefined) {
+        current[pathKey] = {};
+      }
+      current = current[pathKey] as Record<string, unknown>;
+    }
+
+    // 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] as Record<string, unknown>,
+        value as Record<string, unknown>,
+      );
+    } 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 as Record<string, unknown>,
+        schemaValue as JsonTypeDefinition,
+        [...currentPathForLogging, key],
+      );
+    }
+  }
+}
+
+/**
+ * Generate default mock data for a schema type.
+ * Returns reasonable default values based on the schema type string.
+ */
+function generateDefaultForSchemaType(schemaType: unknown): unknown {
+  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 (Array.isArray(schemaType)) {
+    if (schemaType.length === 0) return [];
+    // Generate a single default element based on the first element's schema
+    const elementDefault = generateDefaultForSchemaType(schemaType[0]);
+    return elementDefault !== undefined ? [elementDefault] : [];
+  }
+  if (typeof schemaType === 'object' && schemaType !== null) {
+    // Recursively generate defaults for nested objects
+    const result: Record<string, unknown> = {};
+    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: string,
+  key: string,
+): string[] | null {
+  // 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: Record<string, unknown>,
+): void {
+  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 as Record<string, unknown>);
+    } 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 as Record<string, unknown>);
+        }
+      }
+    }
+  }
+}
+
+/**
+ * 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: Record<string, unknown>,
+  dataForMocks: JsonTypeDefinition,
+  pathPrefix: string = '',
+): void {
+  if (typeof dataForMocks !== 'object' || dataForMocks === null) {
+    return;
+  }
+
+  const missingKeys: string[] = [];
+
+  for (const key of Object.keys(dataForMocks)) {
+    if (key === '_nullable') continue; // Internal marker, not a data key
+
+    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 as Record<string, unknown>,
+          schemaValue as JsonTypeDefinition,
+          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 ? '...' : ''}`,
+    );
+  }
+}
+
+/**
+ * Enforce execution flow requiredValues by setting falsy paths to null.
+ *
+ * The LLM doesn't reliably generate null for `comparison: 'falsy'` requirements.
+ * For example, a flow like "diffView: falsy" should hide a modal, but the LLM
+ * might generate a truthy object, causing the modal to show in all screenshots.
+ *
+ * This function:
+ * 1. Gets requiredValues from covered flows
+ * 2. For 'falsy' comparisons: sets the value to null
+ * 3. For 'truthy' comparisons with falsy values: generates a default truthy value
+ */
+function enforceRequiredValues(
+  mockData: Record<string, unknown>,
+  coveredFlowIds: string[],
+  executionFlows: ExecutionFlow[],
+): void {
+  if (!coveredFlowIds.length || !executionFlows.length) {
+    return;
+  }
+
+  // Get all requiredValues from covered flows
+  const coveredFlows = executionFlows.filter((flow) =>
+    coveredFlowIds.includes(flow.id),
+  );
+
+  for (const flow of coveredFlows) {
+    if (!flow.requiredValues) continue;
+
+    for (const rv of flow.requiredValues) {
+      if (!rv.attributePath) continue;
+
+      // Find the value in mockData - the path could be nested
+      // e.g., attributePath: "diffView" could be at mockData['useDiffModal()'].diffView
+      const result = findAndSetValueInMockData(
+        mockData,
+        rv.attributePath,
+        rv.comparison,
+        rv.valueType,
+      );
+
+      if (result.found) {
+        awsLog(
+          `CodeYam: Enforced ${rv.comparison} for ${rv.attributePath} (set to ${result.newValue === null ? 'null' : typeof result.newValue})`,
+        );
+      }
+    }
+  }
+}
+
+/**
+ * Find a value in mockData by attributePath and enforce the comparison.
+ * The attributePath could be a simple key or a nested path.
+ *
+ * Returns { found: boolean, newValue: unknown }
+ */
+function findAndSetValueInMockData(
+  mockData: Record<string, unknown>,
+  attributePath: string,
+  comparison: string,
+  valueType?: string,
+): { found: boolean; newValue?: unknown } {
+  // Try to find the path at various nesting levels
+  // The attributePath might be "diffView" but the actual location is
+  // mockData['useDiffModal()'].diffView
+
+  // Strategy 1: Direct path (e.g., mockData[attributePath])
+  if (attributePath in mockData) {
+    const currentValue = mockData[attributePath];
+    const { shouldChange, newValue } = getEnforcedValue(
+      currentValue,
+      comparison,
+      valueType,
+    );
+    if (shouldChange) {
+      mockData[attributePath] = newValue;
+      return { found: true, newValue };
+    }
+    return { found: true, newValue: currentValue };
+  }
+
+  // Strategy 2: Search in nested objects
+  for (const [key, value] of Object.entries(mockData)) {
+    if (value && typeof value === 'object' && !Array.isArray(value)) {
+      const nestedObj = value as Record<string, unknown>;
+
+      // Check if attributePath exists in this nested object
+      if (attributePath in nestedObj) {
+        const currentValue = nestedObj[attributePath];
+        const { shouldChange, newValue } = getEnforcedValue(
+          currentValue,
+          comparison,
+          valueType,
+        );
+        if (shouldChange) {
+          nestedObj[attributePath] = newValue;
+          return { found: true, newValue };
+        }
+        return { found: true, newValue: currentValue };
+      }
+
+      // Also check dot-notation paths (e.g., "diffView.type")
+      if (attributePath.includes('.')) {
+        const parts = attributePath.split('.');
+        const firstPart = parts[0];
+        if (firstPart in nestedObj) {
+          // Recurse with the rest of the path
+          const result = findAndSetValueInMockData(
+            nestedObj,
+            attributePath,
+            comparison,
+            valueType,
+          );
+          if (result.found) return result;
+        }
+      }
+
+      // Recursively search deeper
+      const result = findAndSetValueInMockData(
+        nestedObj,
+        attributePath,
+        comparison,
+        valueType,
+      );
+      if (result.found) return result;
+    }
+  }
+
+  return { found: false };
+}
+
+/**
+ * Determine if a value should be changed to match a comparison requirement.
+ *
+ * For 'falsy' comparison: truthy values should become null
+ * For 'truthy' comparison: falsy values should become a default truthy value
+ */
+function getEnforcedValue(
+  currentValue: unknown,
+  comparison: string,
+  valueType?: string,
+): { shouldChange: boolean; newValue: unknown } {
+  const isTruthy = Boolean(currentValue);
+
+  if (comparison === 'falsy') {
+    // Value should be falsy
+    if (isTruthy) {
+      return { shouldChange: true, newValue: null };
+    }
+    return { shouldChange: false, newValue: currentValue };
+  }
+
+  if (comparison === 'truthy') {
+    // Value should be truthy
+    if (!isTruthy) {
+      // Generate a default truthy value based on valueType
+      const defaultValue = generateDefaultTruthyValue(valueType);
+      return { shouldChange: true, newValue: defaultValue };
+    }
+    return { shouldChange: false, newValue: currentValue };
+  }
+
+  // For other comparisons (equals, exists, etc.), don't auto-enforce
+  return { shouldChange: false, newValue: currentValue };
+}
+
+/**
+ * Generate a default truthy value for a given type.
+ */
+function generateDefaultTruthyValue(valueType?: string): unknown {
+  if (!valueType) return { _placeholder: true };
+
+  switch (valueType.toLowerCase()) {
+    case 'string':
+      return 'default-value';
+    case 'number':
+      return 1;
+    case 'boolean':
+      return true;
+    case 'array':
+      return [{ _placeholder: true }];
+    case 'object':
+    default:
+      return { _placeholder: true };
+  }
+}
+
 type ScenarioDataWithoutDescription = Omit<ScenarioData, 'scenarioDescription'>;
 
 interface GenerateEntityScenarioDataArgs {
@@ -236,12 +769,122 @@ export async function generateDataForScenario({
 }: GenerateDataForScenarioArgs) {
   awsLogDebugLevel(1, `Generating data for ${entity.name}: ${scenario.name}`);
 
+  // Check if we should chunk the data structure for focused processing
+  let chunkedMockData: Record<string, unknown> | undefined;
+  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 as any).mockData;
+            if (chunkMockData && typeof chunkMockData === 'object') {
+              Object.assign(chunkedMockData, chunkMockData);
+              awsLog(
+                `Chunk ${i + 1} generated data for: ${Object.keys(chunkMockData).join(', ')}`,
+              );
+            }
+          }
+        }
+      }
+
+      // Detect keys that were lost from failed or partial chunk responses
+      // and fill them with schema-based defaults so they aren't permanently lost.
+      const allChunkedKeys = chunks.flatMap((c) => Object.keys(c || {}));
+      const returnedKeys = new Set(Object.keys(chunkedMockData));
+      const missingChunkKeys = allChunkedKeys.filter(
+        (k) => !returnedKeys.has(k),
+      );
+
+      if (missingChunkKeys.length > 0) {
+        awsLog(
+          `Chunked processing: ${missingChunkKeys.length} key(s) missing from chunk results, filling with defaults: ${missingChunkKeys.join(', ')}`,
+        );
+        const dataForMocksRecord = structure.dataForMocks as Record<
+          string,
+          unknown
+        >;
+        for (const key of missingChunkKeys) {
+          chunkedMockData[key] = generateDefaultForSchemaType(
+            dataForMocksRecord[key],
+          );
+        }
+      }
+
+      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;
@@ -390,14 +1033,47 @@ export async function generateDataForScenario({
     }
   }
 
-  if (structure.dataForMocks && !fullScenarioData.data.mockData) {
-    fullScenarioData.data.mockData = {};
-    for (const propKey of Object.keys(structure.arguments)) {
-      const dataAsAny = fullScenarioData.data as any;
-      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) {
+    fullScenarioData.data.mockData ||= {};
+    const dataAsAny = fullScenarioData.data as any;
+    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
@@ -408,6 +1084,59 @@ export async function generateDataForScenario({
       structure.dataForMocks,
     );
   }
+
+  // Coerce objects/arrays to primitives when the schema expects a primitive type.
+  // The LLM sometimes generates an object where the schema expects "string",
+  // e.g., { body: { "env": "production" } } instead of { body: "some string" }.
+  // This causes runtime errors like "TypeError: body.match is not a function".
+  // Must run BEFORE convertCommaSeparatedStringsToArrays, which intentionally
+  // overrides schema types for array-like field names.
+  if (structure.dataForMocks && fullScenarioData.data.mockData) {
+    coerceObjectsToPrimitivesBySchema(
+      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);
+  }
+
+  // Convert type annotation strings that appear as values to actual values.
+  // The LLM sometimes echoes the schema type annotation as the value.
+  // For example, if the schema says { filePath: "string | undefined" },
+  // the LLM might return { filePath: "string | undefined" } instead of
+  // generating an actual value. This converts those type strings to
+  // appropriate default values (e.g., "string | undefined" → undefined).
+  if (fullScenarioData.data.mockData) {
+    convertTypeAnnotationsToValues(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);
+  }
+
+  // Enforce execution flow requiredValues by setting falsy paths to null.
+  // The LLM doesn't reliably generate null for falsy requirements (e.g., diffView: falsy
+  // to hide a modal). This post-processing ensures that scenarios match their
+  // covered flows' requiredValues.
+  if (fullScenarioData.data.mockData && executionFlows) {
+    enforceRequiredValues(
+      fullScenarioData.data.mockData,
+      scenario.metadata?.coveredFlows || [],
+      executionFlows,
+    );
+  }
+
   if (structure.arguments && fullScenarioData.data.argumentsData) {
     for (let i = 0; i < fullScenarioData.data.argumentsData.length; i++) {
       if (structure.arguments[i]) {
@@ -419,7 +1148,8 @@ export async function generateDataForScenario({
     }
   }
 
-  // For Default Scenario only: check for missing keys and make follow-up call if needed
+  // 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,
@@ -430,6 +1160,31 @@ export async function generateDataForScenario({
     });
   }
 
+  // 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 },
@@ -571,6 +1326,60 @@ export default async function generateEntityScenarioData({
           defaultScenarioData.data.argumentsData.map((arg) => ({ ...arg }));
       }
 
+      // Enforce requiredValues AFTER merge for non-default scenarios
+      // This ensures that if a non-default scenario doesn't cover a certain flow,
+      // it inherits the enforcement from the default scenario
+      if (scenarioData.data?.mockData && executionFlows) {
+        // Get the flows covered by this scenario
+        const scenarioCoveredFlows =
+          nonDefaultScenarios.find(
+            (s) => s.name === result.scenarioData.scenarioDescription,
+          )?.metadata?.coveredFlows || [];
+
+        // Get the paths that the scenario's flows affect
+        const scenarioAffectedPaths = new Set<string>();
+        for (const flowId of scenarioCoveredFlows) {
+          const flow = executionFlows.find((f) => f.id === flowId);
+          if (flow?.requiredValues) {
+            for (const rv of flow.requiredValues) {
+              if (rv.attributePath) {
+                // Extract base path (e.g., "diffView" from "diffView.type")
+                const basePath = rv.attributePath.split('.')[0];
+                scenarioAffectedPaths.add(basePath);
+              }
+            }
+          }
+        }
+
+        // Get the default scenario's flows
+        const defaultCoveredFlows =
+          defaultScenario.metadata?.coveredFlows || [];
+
+        // For paths NOT affected by the scenario, apply default's enforcement
+        const defaultFlowsForUnaffectedPaths = defaultCoveredFlows.filter(
+          (flowId) => {
+            const flow = executionFlows.find((f) => f.id === flowId);
+            if (!flow?.requiredValues) return false;
+
+            // Check if this flow affects a path that the scenario doesn't cover
+            return flow.requiredValues.some((rv) => {
+              if (!rv.attributePath) return false;
+              const basePath = rv.attributePath.split('.')[0];
+              return !scenarioAffectedPaths.has(basePath);
+            });
+          },
+        );
+
+        // Apply enforcement from default scenario for unaffected paths
+        if (defaultFlowsForUnaffectedPaths.length > 0) {
+          enforceRequiredValues(
+            scenarioData.data.mockData,
+            defaultFlowsForUnaffectedPaths,
+            executionFlows,
+          );
+        }
+      }
+
       return scenarioData;
     });
 
@@ -600,7 +1409,8 @@ Generate COMPLETE, robust data for the entire data structure.
     : `## Non-Default Scenario
 Generate ONLY the differences from the default scenario.
 - Include only fields that need to change
-- Set to \`null\` to remove data
+- For object/scalar fields: set to \`null\` to remove/unset the data
+- For array fields: use \`[]\` for empty arrays (not \`null\`) unless the schema type explicitly includes \`| null\`
 - Omit unchanged fields—they merge from default`;
 
   // Only include the "NO ERROR DATA" instruction when the scenario doesn't require error data
@@ -633,6 +1443,17 @@ For example, if a flow requires:
 \`\`\`
 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:
@@ -673,6 +1494,7 @@ Use simple elements only (\`<div>\`, \`<span>\`). No custom components.
 - 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
+- For empty arrays, use \`[]\` (not \`null\`) unless the schema type explicitly includes \`| null\` and the scenario requires the attribute be removed
 
 ## CRITICAL: Preserve Exact Structure
 Your response MUST mirror the EXACT nested structure provided in mockData Structure.
@@ -733,6 +1555,8 @@ export const generateMissingKeysSystemMessage =
 
 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~~\`)
@@ -750,6 +1574,7 @@ 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
+- For empty arrays, use \`[]\` (not \`null\`) unless the schema type explicitly includes \`| null\` and the scenario requires the attribute be removed
 
 ## CRITICAL: Preserve Exact Structure
 Your response MUST mirror the EXACT nested structure provided for the missing keys.
@@ -776,3 +1601,107 @@ NEVER include "error" fields in responses. Skip them entirely.
 - 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: string,
+) => `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: string,
+) => `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
+- For empty arrays, use \`[]\` (not \`null\`) unless the schema type explicitly includes \`| null\` and the scenario requires the attribute be removed
+`;