@codeyam/codeyam-cli 0.1.0-staging.e38f7bd → 0.1.0-staging.eb21b2f

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

@@ -1,9 +1,13 @@
 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,
+  ExecutionFlow,
   Scenario,
   ScenarioData,
   ScenariosDataStructure,
@@ -11,16 +15,650 @@ import type {
 import validateJson from './validateJson';
 import { awsLog, awsLogDebugLevel } from '~codeyam/utils';
 import { AI, parseJsonSafe } from '~codeyam/ai';
-import transformMockDataToMatchSchema from './transformMockDataToMatchSchema';
+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.
+ * Returns true if any requiredValue has an error path with truthy comparison.
+ */
+function scenarioRequiresErrorData(
+  scenario: Scenario,
+  executionFlows: ExecutionFlow[] | undefined,
+): boolean {
+  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: Record<string, any>,
+  scenarioData: Record<string, any>,
+): Record<string, any> {
+  // 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: Record<string, any> = {};
+
+  // 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';
 
+/**
+ * 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 {
   entity: Pick<Entity, 'name' | 'filePath' | 'metadata'>;
   structure: ScenariosDataStructure;
   scenarios: Scenario[];
+  executionFlows?: ExecutionFlow[];
   incompleteResponse?: string;
   analysis: Analysis;
   model?: AI.Model;
@@ -34,10 +672,96 @@ type GenerateDataForScenarioArgs = Omit<
   defaultScenarioData?: ScenarioData;
 };
 
+interface FillMissingMockDataKeysArgs {
+  structure: ScenariosDataStructure;
+  scenario: Scenario;
+  executionFlows: ExecutionFlow[] | undefined;
+  fullScenarioData: ScenarioData;
+  model: AI.Model | undefined;
+}
+
+/**
+ * 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,
+}: FillMissingMockDataKeysArgs): Promise<void> {
+  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: Record<string, any> = {};
+  for (const key of missingKeys) {
+    missingSchema[key] = (structure.dataForMocks as Record<string, any>)[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 as any).mockData;
+  if (followUpMockData && typeof followUpMockData === 'object') {
+    fullScenarioData.data.mockData = {
+      ...fullScenarioData.data.mockData,
+      ...followUpMockData,
+    };
+  }
+}
+
 export async function generateDataForScenario({
   entity,
   structure,
   scenario,
+  executionFlows,
   defaultScenarioData,
   incompleteResponse,
   analysis,
@@ -45,11 +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;
@@ -92,7 +927,7 @@ export async function generateDataForScenario({
     ...response.stats,
   });
 
-  const { completion, finishReason } = response;
+  let { completion, finishReason } = response;
 
   if (!completion) {
     console.log(
@@ -101,51 +936,79 @@ export async function generateDataForScenario({
     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 as {
+  let { scenarioData: scenarioDataWithoutDescription } = parsed as {
     scenarioData: ScenarioDataWithoutDescription;
   };
 
-  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)
-      : [],
-  });
+  // 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 as any;
+    const reservedKeys = new Set([
+      'scenarioName',
+      'data',
+      'scenarioDescription',
+    ]);
+    const misplacedKeys: string[] = [];
+
+    // 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 = {};
+      }
 
-  // Transform mock data values to match expected schema types
-  // (e.g., LLM might generate array for route params that should be strings)
-  if (scenarioDataWithoutDescription.data?.mockData && structure.dataForMocks) {
-    scenarioDataWithoutDescription.data.mockData =
-      transformMockDataToMatchSchema(
-        scenarioDataWithoutDescription.data.mockData as Record<string, unknown>,
-        structure.dataForMocks as Record<string, unknown>,
-      ) as typeof scenarioDataWithoutDescription.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 as ScenarioDataWithoutDescription;
+    }
   }
 
   const fullScenarioData: ScenarioData = {
@@ -155,8 +1018,9 @@ export async function generateDataForScenario({
 
   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] = {};
@@ -169,14 +1033,158 @@ 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
+  // for fields typed as "T | undefined" (but preserves null for "T | null").
+  if (structure.dataForMocks && fullScenarioData.data.mockData) {
+    convertNullToUndefinedBySchema(
+      fullScenarioData.data.mockData,
+      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]) {
+        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 },
@@ -187,6 +1195,7 @@ export default async function generateEntityScenarioData({
   entity,
   structure,
   scenarios,
+  executionFlows,
   incompleteResponse,
   analysis,
   model,
@@ -207,6 +1216,7 @@ export default async function generateEntityScenarioData({
       entity,
       structure,
       scenario: defaultScenario,
+      executionFlows,
       incompleteResponse,
       analysis,
       model,
@@ -229,6 +1239,7 @@ export default async function generateEntityScenarioData({
           entity,
           structure,
           scenario,
+          executionFlows,
           defaultScenarioData,
           incompleteResponse,
           analysis,
@@ -249,37 +1260,131 @@ export default async function generateEntityScenarioData({
       );
     }
 
-    scenarioDatas.push(...validResults.map((result) => result.scenarioData));
-    llmCalls.push(...validResults.map((result) => result.llmCall));
+    // 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 }));
+      }
+
+      // 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,
+          );
+        }
+      }
 
-    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,
+      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) {
@@ -291,28 +1396,85 @@ export default async function generateEntityScenarioData({
 export const generateSystemMessage = (
   scenarioName: string,
   defaultScenario: boolean,
+  requiresErrorData: boolean = false,
 ) => {
   const scenarioType = defaultScenario
     ? `## Default Scenario
 Generate COMPLETE, robust data for the entire data structure.
-- 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 unless the key attributes specify a parent attribute should be null or undefined
+- Fill ALL fields with realistic values (except error 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
+- 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
+  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
 
@@ -328,30 +1490,20 @@ Use for relative dates. Code runs in Node (no browser APIs, no external librarie
 \`\`\`
 Use simple elements only (\`<div>\`, \`<span>\`). No custom components.
 
-## Mock Data Keys
-Preserve keys exactly as written in the structure. There are two formats:
+### 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
+- For empty arrays, use \`[]\` (not \`null\`) unless the schema type explicitly includes \`| null\` and the scenario requires the attribute be removed
 
-### Standard function calls
-\`\`\`json
-{
-  "mockData": {
-    "useUser()": { "user": { "name": "John" } },
-    "from().select()": [{ "id": "1" }]
-  }
-}
-\`\`\`
-
-### 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.
+## 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
@@ -369,9 +1521,15 @@ This reads as "entityDiffFetcher receives from useFetcher". Each variable gets i
 ## 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
 `;
 };
 
@@ -387,3 +1545,163 @@ 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
+- 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.
+- 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: 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
+`;