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

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

@@ -1,9 +1,11 @@
 import completionCall from './completionCall';
 import generateEntityScenarioDataGenerator from './promptGenerators/generateEntityScenarioDataGenerator';
+import generateMissingKeysPrompt from './promptGenerators/generateMissingKeysPrompt';
 import { saveLlmCall } from '~codeyam/aws/dynamodb';
 import type {
   Analysis,
   Entity,
+  ExecutionFlow,
   Scenario,
   ScenarioData,
   ScenariosDataStructure,
@@ -11,6 +13,109 @@ import type {
 import validateJson from './validateJson';
 import { awsLog, awsLogDebugLevel } from '~codeyam/utils';
 import { AI, parseJsonSafe } from '~codeyam/ai';
+import convertNullToUndefinedBySchema from './dataStructure/helpers/convertNullToUndefinedBySchema';
+
+/**
+ * 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';
 
@@ -20,6 +125,7 @@ interface GenerateEntityScenarioDataArgs {
   entity: Pick<Entity, 'name' | 'filePath' | 'metadata'>;
   structure: ScenariosDataStructure;
   scenarios: Scenario[];
+  executionFlows?: ExecutionFlow[];
   incompleteResponse?: string;
   analysis: Analysis;
   model?: AI.Model;
@@ -33,10 +139,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,
@@ -47,6 +239,7 @@ export async function generateDataForScenario({
   const prompt = generateEntityScenarioDataGenerator(
     structure,
     scenario,
+    executionFlows,
     defaultScenarioData,
     incompleteResponse,
   );
@@ -91,7 +284,7 @@ export async function generateDataForScenario({
     ...response.stats,
   });
 
-  const { completion, finishReason } = response;
+  let { completion, finishReason } = response;
 
   if (!completion) {
     console.log(
@@ -100,42 +293,80 @@ 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 = {};
+      }
+
+      // 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 = {
     ...scenarioDataWithoutDescription,
@@ -144,8 +375,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] = {};
@@ -166,6 +398,38 @@ export async function generateDataForScenario({
     }
   }
 
+  // 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,
+    );
+  }
+  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
+  if (isDefault) {
+    await fillMissingMockDataKeys({
+      structure,
+      scenario,
+      executionFlows,
+      fullScenarioData,
+      model,
+    });
+  }
+
   return {
     scenarioData: fullScenarioData,
     llmCall: { name: scenario.name, id: llmCall.id },
@@ -176,6 +440,7 @@ export default async function generateEntityScenarioData({
   entity,
   structure,
   scenarios,
+  executionFlows,
   incompleteResponse,
   analysis,
   model,
@@ -196,6 +461,7 @@ export default async function generateEntityScenarioData({
       entity,
       structure,
       scenario: defaultScenario,
+      executionFlows,
       incompleteResponse,
       analysis,
       model,
@@ -218,6 +484,7 @@ export default async function generateEntityScenarioData({
           entity,
           structure,
           scenario,
+          executionFlows,
           defaultScenarioData,
           incompleteResponse,
           analysis,
@@ -238,37 +505,77 @@ 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 };
+      }
 
-    console.log('CodeYam Debug: generateEntityScenarioData results', {
-      filePath: entity.filePath,
-      entityName: entity.name,
-      totalScenarios: scenarios.length,
-      defaultScenarioGenerated: !!defaultScenarioResult,
-      otherScenariosRequested: scenarios.length - 1,
-      otherScenariosResults: results.length,
-      nullResults: results.filter((r) => r === null).length,
-      validResults: validResults.length,
-      finalScenarioDatasCount: scenarioDatas.length,
+      // Merge argumentsData with default argumentsData
+      if (
+        defaultScenarioData.data?.argumentsData &&
+        Array.isArray(defaultScenarioData.data.argumentsData) &&
+        scenarioData.data?.argumentsData &&
+        Array.isArray(scenarioData.data.argumentsData)
+      ) {
+        for (
+          let i = 0;
+          i < defaultScenarioData.data.argumentsData.length;
+          i++
+        ) {
+          const scenarioArg = scenarioData.data.argumentsData[i];
+          const defaultArg = defaultScenarioData.data.argumentsData[i];
+
+          // Only merge if both are objects (LLM sometimes returns primitives)
+          if (
+            scenarioArg &&
+            typeof scenarioArg === 'object' &&
+            !Array.isArray(scenarioArg) &&
+            defaultArg &&
+            typeof defaultArg === 'object' &&
+            !Array.isArray(defaultArg)
+          ) {
+            scenarioData.data.argumentsData[i] = deepMergeScenarioData(
+              defaultArg,
+              scenarioArg,
+            );
+          } else if (scenarioArg !== undefined) {
+            // Keep the scenario value as-is (even if primitive)
+            scenarioData.data.argumentsData[i] = scenarioArg;
+          } else if (defaultArg && typeof defaultArg === 'object') {
+            // Use default if scenario is undefined and default is an object
+            scenarioData.data.argumentsData[i] = { ...defaultArg };
+          } else {
+            scenarioData.data.argumentsData[i] = defaultArg;
+          }
+        }
+      } else if (
+        defaultScenarioData.data?.argumentsData &&
+        !scenarioData.data?.argumentsData
+      ) {
+        // Use default argumentsData if scenario has none
+        scenarioData.data.argumentsData =
+          defaultScenarioData.data.argumentsData.map((arg) => ({ ...arg }));
+      }
+
+      return scenarioData;
     });
 
-    awsLog(
-      'CodeYam: scenarioDatas :>> ',
-      JSON.stringify(
-        {
-          filePath: entity.filePath,
-          entityName: entity.name,
-          scenarioDatas: scenarioDatas.map((sd) => ({
-            scenarioName: sd.scenarioName,
-            hasData: !!sd.data,
-            dataKeys: sd.data ? Object.keys(sd.data) : [],
-          })),
-        },
-        null,
-        2,
-      ),
-    );
+    scenarioDatas.push(...mergedScenarioDatas);
+    llmCalls.push(...validResults.map((result) => result.llmCall));
 
     return { scenarioDatas, llmCalls };
   } catch (error) {
@@ -280,176 +587,129 @@ export default async function generateEntityScenarioData({
 export const generateSystemMessage = (
   scenarioName: string,
   defaultScenario: boolean,
-) => `## Overview
-- You will be provided with a structure for scenario data, the "mockData structure", and for the arguments of the function, the "argumentsData structure", and a scenario describing a new state of data for this function.
-
-- Your goal is to generate data to fulfill both the data structures and the scenario, to be returned in a JSON object.
-
-## Provided Information
-
-- The "mockData structure" and "argumentsData structure" describe all of the data required for the simulation to work. Copy the data structures exactly (if an attribute is a function it will look like "myFunction()" and that exact key "myFunction()" should be in the return value structure) and replace any primitives (e.g. "string", "number", "boolean", "date") with realistic data. 
-
-${defaultScenario ? '' : '- The "Default Scenario Data" provides the data that has already been generated for the default scenario. This data will be merged into the data for each scenario. Your goal is to modify this data as needed to satisfy the requirements of this scenario. As such you only need to specify data that is different or should be removed from the default scenario data. To remove data set some portion of the data structure to null.'}
-
-- Follow the "keyAttributeInstructions" for the scenario carefully.${defaultScenario ? ' As the Default Scenario ensure that the entire data structure is filled in with appropriate values even if they are not explicitly described in the "keyAttributeInstructions".' : ''}
-
-## Instructions
-${
-  defaultScenario
-    ? `
-- We need data for the default scenario.
-
-- The default scenario requires robust and complete data that fulfills all the "mockData" and "argumentsData" data structures if they exist, except for error attributes.
-
-- The one exception to this is error data. Sometimes the data structure will contain structures for both success and error data. In this case you should only fill in the success data and leave out the error data completely.
-
-- This data should cover for every object, attribute, and nested attribute described in the data structures (other than error data). If an item in a data structure is an array then provide multiple objects in the array. Fill in every nested attribute with realistic data. This should provide a robust and complete "database" for the simulation. If the type is described as "any" then create a realistic example of that type based on its name. If the type is "date" or "timestamp" then use the instructions below to generate a dynamic date or timestamp. 
-
-- Remember to fill in each and every attribute and nested attribute in the data structure for the default scenario (other than error data). This will make the simulation more realistic and valuable. Do not add any attributes that are not in the data structure. If an attribute is not in the data structure then it should not be included in the data.
-`
-    : `
-- This scenario is a modification of the default scenario. Figure out what needs to change about the Default Scenario in order to satisfy this scenario. If data must be removed then mark is as null at the appropriate level of the structure that the deletion should occur. Otherwise you can change any data that needs to be modified. Do not include any aspect of the structure that should remain the same as the Default Scenario.
-
-- The data generated will be merged into the Default Scenario Data in order to create the full data for this scenario.
-`
-}
-
-## Example Response
-
+  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)
+- Do not skip any keys, even simple or small entries
+- Arrays should have 3-5 items to provide realistic test data variety
+- Don't skip nested attributes unless the execution flow requirements specify a parent attribute should be null or undefined
+- This provides the baseline data for all other scenarios`
+    : `## Non-Default Scenario
+Generate ONLY the differences from the default scenario.
+- Include only fields that need to change
+- Set to \`null\` to remove data
+- Omit unchanged fields—they merge from default`;
+
+  // Only include the "NO ERROR DATA" instruction when the scenario doesn't require error data
+  const noErrorDataInstruction = requiresErrorData
+    ? `## IMPORTANT: ERROR DATA REQUIRED
+This scenario tests error handling. You MUST include "error" fields with realistic error messages.
+- Set error fields to truthy values (e.g., "Error: Operation failed" or "Something went wrong")
+- The error data is REQUIRED to trigger the correct error UI state`
+    : `## CRITICAL: NO ERROR DATA
+NEVER include "error" fields in responses. Skip them entirely.
+- If structure has \`{ data: {...}, error: {...} }\`, only fill \`data\`
+- Leave out any attribute named "error"—do not set to null, omit entirely`;
+
+  return `You are a test data generator. Create mock data matching a data structure and scenario requirements.
+
+## Execution Flow Requirements
+Each scenario has \`coveredFlows\` which lists the execution flows (distinct outcomes/behaviors) this scenario should demonstrate.
+Each flow has \`requiredValues\` - the attribute values that MUST be set to produce that outcome.
+
+**Your job**: Generate mock data that satisfies ALL the requiredValues from ALL coveredFlows.
+
+For example, if a flow requires:
+\`\`\`json
 {
-  "scenarioData": {
-    "scenarioName": ${scenarioName},
-    "data': {
-      "mockData": {
-        "useUser()": {
-          "user": {
-            "name": "John Doe",
-            "metadata": {
-              "hobbies": ["fishing", "hiking", "camping"]
-              "favoriteColor": "blue",
-            },
-            "joinedAt": {
-              "~~codeyam-code~~": "new Date(Date.now() - 3 * 24 * 60 * 60 * 1000)"         
-            }
-          },
-        },
-        "loadProject()": {
-          "project": {
-            "name": "New Project",
-            "slug": "new-project",
-            "files": [
-              {
-                "name": "File 1",
-                "slug": "file-1",
-                "content": "This is the content of file 1"
-              },
-              {
-                "name": "File 2",
-                "slug": "file-2",
-                "content": "This is the content of file 2"
-              },
-              {
-                "name": "File 3",
-                "slug": "file-3",
-                "content": "This is the content of file 3"
-              }
-            ]
-          }
-        }
-      },
-      "argumentsData": [
-        {
-          "open": "true",
-          "children: {
-            "~~codeyam-jsx~~": "<div>Hi!</div>"         
-          }      
-        }    
-      ]  
-    }
-  }
+  "attributePath": "signature[0].isLoading",
+  "value": "false",
+  "comparison": "equals",
+  "valueType": "boolean"
 }
+\`\`\`
+Then set \`isLoading: false\` in the mockData.
 
-## scenarioData
-
-You will return a JSON object called scenarioData containing the name and data for this scenario.
-
-If you encounter a "date" then you should use "~~codeyam-code~~" along with code in string format. That string will be evaled to create the date or timestamp needed.
+## 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 example the "timestamp" above should probably not be a hard-coded timestamp but should be a function that returns a timestamp. If you want to return a timestamp that is 24 hours in the future, you would return this object:
+- 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
 {
-  "scenarioData": {
-    "scenarioName": ${scenarioName},
-    "data": {
-      "mockData": {
-        "dataLoaders": {
-          "fetchUser()": {
-            "user": {
-              "name": "John Doe",
-              "metadata": {
-                "hobbies": ["fishing", "hiking", "camping"],
-                "favoriteColor": "blue"
-              },
-              "joinedAt": {
-                "~~codeyam-code~~": "new Date(Date.now() - 3 * 24 * 60 * 60 * 1000).toISOString()"
-              }  
-            }
-          }
-        }
-      }
-    }
-  }
+  "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.
 
-Keep in mind the above pattern is also valid for argumentsData.
+${scenarioType}
 
-All code should be written in JavaScript. It will run in a node environment (not the browser) with no reference to external libraries. It will essentially be run in a sandboxed environment. Do not reference any external libraries in the code.
+${noErrorDataInstruction}
 
-With argumentsData you may also need to return a "children" argument. "children" can either be a string or a jsx element.
+## Special Markers
 
-If you want to return a jsx element you would return something like this:
+### 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
+- In general we want robust data, not minimal data unless specified otherwise
+
+## CRITICAL: Preserve Exact Structure
+Your response MUST mirror the EXACT nested structure provided in mockData Structure.
+- Do NOT reorganize, split, or create duplicate keys
+- The hierarchy of nested objects must match exactly what was provided unless overridden by scenario rules
+- Only change the leaf VALUES (replacing type descriptions like "string" with actual data like "hello")
+- Copy the key strings EXACTLY from the structure
+- Do NOT modify type parameters, arguments, or any part of the key
+- The keys preserve the exact function call as written in the original code
+
+## Response Format
+\`\`\`json
 {
   "scenarioData": {
-    "scenarioName": ${scenarioName},
+    "scenarioName": "${scenarioName}",
     "data": {
-      "argumentsData": [
-        {
-          "children": {
-            "~~codeyam-jsx~~": "<div>Click Me!</div>"
-          }
-        }
-      ]
+      "mockData": { ... },
+      "argumentsData": [ ... ]
     }
   }
 }
 \`\`\`
 
-The "~~codeyam-jsx~~" key in a "children" prop will be rendered directly to the code as a jsx element, not a string. Only return basic elements like "<div>".
-
-Do not import or reference any libraries in this code. No external libraries will be imported.
-
-If either mockData or argumentsData is empty you can return an empty object for mockData ("{}") or empty array for argumentsData ("[]"). 
-
-If an object inside of mockData or argumentsData is intended to not exist you can set it to null.
-
-Never use "~~codeyam-code~~" or "~~codeyam-jsx~~" as a partial key in the response. Even if the key looks like it should have a relative date or code. "~~codeyam-code~~" and "~~codeyam-jsx~~" are only used as independent keys used to mark values that are code.
-
-## Common Mistakes
-
-- Do not use the term undefined in the response. If something is missing leave it out or use null.
-- Leave out error data (e.g. a data attribute labeled as "error"), especially when it is in a data structure that has success data as well. Only include the success data and simply leave out the error data. Error data is usually labeled as "error" in the structure. Leave it out entirely.
-  - Please double check this. Ensure no error data is included in the response unless specifically requested. ${defaultScenario ? `This includes the ${`DEFAULT_SCENARIO_NAME`}.` : ''} Leave out any error attributes from the response completely!
-- Do not reference data in the response (e.g. if there was a "posts" field you can not say "posts[0]" elsewhere in the response - you need to duplicate the data). Unfortunately all data must be written out directly in the response.
-- Ensure each scenario in the response matches the full scenario name, commas and all, in the scenarios provided.
-- Do not throw errors from this data in any manner. Simple leave out any error data.
-- If the data structure contains a type that is a "date" then use ~~codeyam-code~~ to generate a dynamic timestamp or date. Do not hard-code dates unless it needs to be a very specific date. Otherwise it should be a relative, dynamic date.
-- The response must be valid JSON. Do not include raw code in the response.
-${defaultScenario ? `- As the ${DEFAULT_SCENARIO_NAME} be very robust in generating data that reflects common real-world usage. All arrays and attributes, except error attributes, should be filled in with high quality data so that default simulation provides a good default of how the function behaves` : ''}
+## 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)
+- Scenario name must match exactly: "${scenarioName}"
+- Empty mockData: \`{}\`, empty argumentsData: \`[]\`
+
+## IMPORTANT: Avoid Identifier Collisions
+When generating identifier values (SHA hashes, entity IDs, etc.):
+- Use DISTINCT values for different identifier fields
+- Avoid matching: if \`entity.sha\` is "abc123", arrays like \`jobs[].entityShas\` or \`currentlyExecuting.entityShas\` should NOT contain "abc123"
+- This prevents accidental blocking of UI conditionals that check if IDs are in/not-in arrays
 `;
+};
 
 export const generateIncompleteSystemMessage = (
   scenarioName: string,
@@ -463,3 +723,56 @@ 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.
+
+## Special Markers
+
+### Dynamic Dates (\`~~codeyam-code~~\`)
+\`\`\`json
+{ "createdAt": { "~~codeyam-code~~": "new Date(Date.now() - 24*60*60*1000)" } }
+\`\`\`
+Use for relative dates. Code runs in Node (no browser APIs, no external libraries).
+
+### JSX Children (\`~~codeyam-jsx~~\`)
+\`\`\`json
+{ "children": { "~~codeyam-jsx~~": "<div>Hello</div>" } }
+\`\`\`
+Use simple elements only (\`<div>\`, \`<span>\`). No custom components.
+
+### Arrays
+- Arrays should have many items (at least 4) unless specified otherwise
+- Each item must follow the exact structure provided
+
+## CRITICAL: Preserve Exact Structure
+Your response MUST mirror the EXACT nested structure provided for the missing keys.
+- Only change the leaf VALUES (replacing type descriptions like "string" with actual data)
+- Copy the key strings EXACTLY from the structure
+- Do NOT modify type parameters, arguments, or any part of the key
+
+## CRITICAL: NO ERROR DATA
+NEVER include "error" fields in responses. Skip them entirely.
+- If structure has \`{ data: {...}, error: {...} }\`, only fill \`data\`
+- Leave out any attribute named "error"—do not set to null, omit entirely
+
+## Response Format
+\`\`\`json
+{
+  "mockData": {
+    // fill in ONLY the missing keys
+  }
+}
+\`\`\`
+
+## Rules
+- Valid JSON only—no raw code outside markers
+- No \`undefined\`—use \`null\` or omit
+- No data references (can't use \`posts[0]\` elsewhere — duplicate the value)
+`;