@codeyam/codeyam-cli 0.1.0-staging.15d0f46 → 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,13 +587,15 @@ 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.
@@ -294,14 +603,57 @@ Generate ONLY the differences from the default scenario.
 - Set to \`null\` to remove data
 - Omit unchanged fields—they merge from default`;
 
+  // Only include the "NO ERROR DATA" instruction when the scenario doesn't require error data
+  const noErrorDataInstruction = requiresErrorData
+    ? `## IMPORTANT: ERROR DATA REQUIRED
+This scenario tests error handling. You MUST include "error" fields with realistic error messages.
+- Set error fields to truthy values (e.g., "Error: Operation failed" or "Something went wrong")
+- The error data is REQUIRED to trigger the correct error UI state`
+    : `## CRITICAL: NO ERROR DATA
+NEVER include "error" fields in responses. Skip them entirely.
+- If structure has \`{ data: {...}, error: {...} }\`, only fill \`data\`
+- Leave out any attribute named "error"—do not set to null, omit entirely`;
+
   return `You are a test data generator. Create mock data matching a data structure and scenario requirements.
 
+## Execution Flow Requirements
+Each scenario has \`coveredFlows\` which lists the execution flows (distinct outcomes/behaviors) this scenario should demonstrate.
+Each flow has \`requiredValues\` - the attribute values that MUST be set to produce that outcome.
+
+**Your job**: Generate mock data that satisfies ALL the requiredValues from ALL coveredFlows.
+
+For example, if a flow requires:
+\`\`\`json
+{
+  "attributePath": "signature[0].isLoading",
+  "value": "false",
+  "comparison": "equals",
+  "valueType": "boolean"
+}
+\`\`\`
+Then set \`isLoading: false\` in the mockData.
+
+## 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
 
@@ -317,30 +669,19 @@ 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
 
-### 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
@@ -358,9 +699,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
 `;
 };
 
@@ -376,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)
+`;