@codeyam/codeyam-cli 0.1.0-staging.bbe4da9 → 0.1.0-staging.c90f8c9

package/analyzer-template/project/orchestrateCapture/SequentialCaptureTaskRunner.ts

@@ -37,7 +37,10 @@ import {
 export interface SequentialCaptureTaskRunnerContext {
   absoluteProjectRootPath: string;
   absoluteTempPath?: string;
-  startArgs: StartArgs;
+  startArgs: Omit<
+    StartArgs,
+    'absoluteProjectRootPath' | 'absoluteTempPath' | 'noServer' | 'fast'
+  >;
   /** Universal mocks to apply to the capturer's project directory */
   universalMocks?: UniversalMock[];
   /** Environment variables to pass to the spawned processes */
@@ -72,8 +75,6 @@ export default class SequentialCaptureTaskRunner implements CaptureTaskRunner {
     config: TaskRunnerConfig,
     context: SequentialCaptureTaskRunnerContext,
   ): Promise<SequentialCaptureTaskRunner> {
-    console.log('CodeYam DEBUG: Creating SequentialCaptureTaskRunner...');
-
     // Validate that cwd and project path share parent directory
     // Use realpathSync to resolve symlinks (e.g., /tmp -> /private/tmp on macOS)
     const parentOfCodeyam = fs.realpathSync(path.dirname(process.cwd()));
@@ -90,56 +91,65 @@ export default class SequentialCaptureTaskRunner implements CaptureTaskRunner {
 
     // Derive paths
     const parentDir = parentOfCodeyam;
-    const capturerPath = path.normalize(path.join(parentDir, 'capturer-1'));
 
-    console.log(`CodeYam: Sequential runner using capturer at ${capturerPath}`);
+    // In CAPTURE_ONLY mode, run directly from the main project folder (skip rsync)
+    // This preserves any manual edits made to generated files
+    const captureOnlyEnvValue = process.env.CAPTURE_ONLY;
+    const captureOnly = captureOnlyEnvValue === 'true';
 
-    // Create capturer base directory
-    if (!fs.existsSync(capturerPath)) {
-      console.log(
-        `CodeYam DEBUG: Creating capturer directory at ${capturerPath}`,
-      );
-      fs.mkdirSync(capturerPath, { recursive: true });
-    }
-
-    // Copy both codeyam/ and project/ into capturer-1/
-    // start.ts waits for repo checkout to complete before creating this runner
-    console.log('CodeYam: Copying codeyam/ and project/ to capturer-1/...');
-    await rsyncCopy({
-      sourcePath: parentDir,
-      destinationPath: capturerPath,
-      excludes: ['codeyam/log.txt', 'capturer-*'],
-      silent: false,
-    });
-    console.log('CodeYam DEBUG: Directory copy completed');
-
-    // Sanitize env files in capturer project (replaces real values with mock placeholders)
-    const capturerProjectPath = path.join(capturerPath, 'project');
-    console.log('CodeYam: Sanitizing env files in capturer project...');
-    console.log(`CodeYam DEBUG: capturerProjectPath = ${capturerProjectPath}`);
-    const sanitizeResult = sanitizeEnvFiles({
-      projectPath: capturerProjectPath,
-      log: (msg) => console.log(`CodeYam Sanitize: ${msg}`),
-    });
     console.log(
-      `CodeYam DEBUG: Env files sanitized - created: ${sanitizeResult.createdFiles.length}, sanitized: ${sanitizeResult.sanitizedFiles.length}`,
+      `CodeYam: CAPTURE_ONLY env = ${captureOnlyEnvValue === undefined ? 'undefined' : `"${captureOnlyEnvValue}"`}`,
     );
-    if (sanitizeResult.sanitizedFiles.length > 0) {
-      console.log(
-        `CodeYam DEBUG: Sanitized files: ${sanitizeResult.sanitizedFiles.join(', ')}`,
-      );
-    }
 
-    // Apply universal mocks to capturer project
-    if (context.universalMocks && context.universalMocks.length > 0) {
+    const capturerPath = captureOnly
+      ? parentDir
+      : path.normalize(path.join(parentDir, 'capturer-1'));
+
+    console.log(`CodeYam: Sequential runner using capturer at ${capturerPath}`);
+
+    if (captureOnly) {
       console.log(
-        `CodeYam: Applying ${context.universalMocks.length} universal mock(s) to capturer project...`,
+        'CodeYam: CAPTURE_ONLY mode - running from main project folder, skipping rsync',
       );
-      await applyUniversalMocks({
-        mocks: context.universalMocks,
+    } else {
+      // Create capturer base directory
+      if (!fs.existsSync(capturerPath)) {
+        fs.mkdirSync(capturerPath, { recursive: true });
+      }
+
+      // Copy both codeyam/ and project/ into capturer-1/
+      // start.ts waits for repo checkout to complete before creating this runner
+      console.log('CodeYam: Copying codeyam/ and project/ to capturer-1/...');
+      await rsyncCopy({
+        sourcePath: parentDir,
+        destinationPath: capturerPath,
+        excludes: ['codeyam/log.txt', 'capturer-*'],
+        silent: false,
+      });
+
+      // Sanitize env files in capturer project (replaces real values with mock placeholders)
+      const capturerProjectPath = path.join(capturerPath, 'project');
+      console.log('CodeYam: Sanitizing env files in capturer project...');
+      const sanitizeResult = sanitizeEnvFiles({
         projectPath: capturerProjectPath,
+        log: (msg) => console.log(`CodeYam Sanitize: ${msg}`),
       });
-      console.log('CodeYam DEBUG: Universal mocks applied');
+      if (sanitizeResult.sanitizedFiles.length > 0) {
+        console.log(
+          `CodeYam: Env files sanitized: ${sanitizeResult.sanitizedFiles.join(', ')}`,
+        );
+      }
+
+      // Apply universal mocks to capturer project
+      if (context.universalMocks && context.universalMocks.length > 0) {
+        console.log(
+          `CodeYam: Applying ${context.universalMocks.length} universal mock(s) to capturer project...`,
+        );
+        await applyUniversalMocks({
+          mocks: context.universalMocks,
+          projectPath: capturerProjectPath,
+        });
+      }
     }
 
     const runner = new SequentialCaptureTaskRunner(
@@ -150,12 +160,8 @@ export default class SequentialCaptureTaskRunner implements CaptureTaskRunner {
 
     // Start background processing loop
     // Each batch will spawn its own paired start+capture processes
-    console.log('CodeYam DEBUG: Starting processing loop...');
     runner.startProcessingLoop();
 
-    console.log(
-      'CodeYam DEBUG: SequentialCaptureTaskRunner created successfully',
-    );
     return runner;
   }
 
@@ -179,23 +185,18 @@ export default class SequentialCaptureTaskRunner implements CaptureTaskRunner {
   }
 
   private async runProcessingLoop(): Promise<void> {
-    console.log('CodeYam DEBUG: Processing loop started');
     while (!this.isCleanedUp) {
       if (this.queue.length === 0) {
         await new Promise((resolve) => setTimeout(resolve, 1000));
         continue;
       }
 
-      console.log(
-        `CodeYam DEBUG: Processing queue - ${this.queue.length} batches waiting`,
-      );
       this.isProcessing = true;
       const batch = this.queue.shift()!;
 
       try {
         await this.processBatch(batch);
 
-        console.log(`CodeYam DEBUG: Batch completed successfully`);
         // Notify callback of completion
         if (this.batchCompletedCallback) {
           this.batchCompletedCallback(batch);
@@ -210,7 +211,6 @@ export default class SequentialCaptureTaskRunner implements CaptureTaskRunner {
         this.isProcessing = false;
       }
     }
-    console.log('CodeYam DEBUG: Processing loop exited');
   }
 
   private async processBatch(analysisIds: string[]): Promise<void> {
@@ -265,6 +265,10 @@ export default class SequentialCaptureTaskRunner implements CaptureTaskRunner {
         READY_TO_BE_CAPTURED: 'true',
         PROJECT_SLUG: this.config.projectSlug,
         CODEYAM_LOG_PATH: capturerLogPath,
+        // Pass CAPTURE_ONLY through from parent process to skip file regeneration
+        ...(process.env.CAPTURE_ONLY && {
+          CAPTURE_ONLY: process.env.CAPTURE_ONLY,
+        }),
       },
       workingDir: capturerCodeyamPath,
       processName: 'project',
@@ -298,14 +302,38 @@ export default class SequentialCaptureTaskRunner implements CaptureTaskRunner {
       }, timeoutMs);
     });
 
+    // Track which process exits first and when
+    const batchStartTime = Date.now();
+    let startExitCode: number | null = null;
+    let captureExitCode: number | null = null;
+    let firstToExit: 'start' | 'capture' | 'timeout' | null = null;
+
+    const startPromiseWithLogging = startResult.promise.then((code) => {
+      startExitCode = code;
+      if (!firstToExit) {
+        firstToExit = 'start';
+      }
+      return code;
+    });
+
+    const capturePromiseWithLogging = captureResult.promise.then((code) => {
+      captureExitCode = code;
+      if (!firstToExit) {
+        firstToExit = 'capture';
+      }
+      return code;
+    });
+
     try {
-      // Race: whichever finishes first (capture exits when done, or timeout)
+      // Wait for EITHER process to complete (or timeout)
+      // If start.js exits, we should investigate why - it shouldn't exit before capture is done
       await Promise.race([
-        captureResult.promise,
-        startResult.promise,
+        capturePromiseWithLogging,
+        startPromiseWithLogging,
         timeoutPromise,
       ]);
-      console.log('CodeYam: Batch processing completed successfully');
+      const elapsed = ((Date.now() - batchStartTime) / 1000).toFixed(1);
+      console.log(`CodeYam: Batch processing completed after ${elapsed}s`);
     } finally {
       // Clean up timeout handle
       if (timeoutHandle) {
@@ -330,9 +358,6 @@ export default class SequentialCaptureTaskRunner implements CaptureTaskRunner {
   }
 
   enqueueWork(analysisIds: string[]): Promise<void> {
-    console.log(
-      `CodeYam DEBUG: Enqueueing batch of ${analysisIds.length} analyses`,
-    );
     this.queue.push(analysisIds);
     return Promise.resolve();
   }

package/analyzer-template/project/orchestrateCapture/SupabaseAnalysisLoader.ts

@@ -1,5 +1,5 @@
 import type { Commit, Project } from '~codeyam/types';
-import { getSupabase } from '~codeyam/supabase';
+import { getSupabase } from '~codeyam/database';
 import type { AnalysisLoader, ReadyAnalysisData } from './taskRunner';
 
 interface SupabaseAnalysisLoaderConfig {
@@ -48,7 +48,7 @@ export default class SupabaseAnalysisLoader implements AnalysisLoader {
     let pageOffset = 0;
 
     while (true) {
-      const query = getSupabase()
+      const { data, error } = await getSupabase()
         .rpc('get_ready_to_be_captured_analyses', args)
         .select(
           'id, commit_id, branch_commit_sha, status->processingForCaptureAt',
@@ -56,7 +56,7 @@ export default class SupabaseAnalysisLoader implements AnalysisLoader {
         .range(pageOffset, pageOffset + pageSize - 1)
         .returns<ReadyAnalysisData[]>();
 
-      const { data, error } = await query;
+      const rows = data as ReadyAnalysisData[] | null;
 
       if (error) {
         console.error(
@@ -67,13 +67,13 @@ export default class SupabaseAnalysisLoader implements AnalysisLoader {
         return null;
       }
 
-      if (!data || data.length === 0) {
+      if (!rows || rows.length === 0) {
         break;
       }
 
-      allData.push(...data);
+      allData.push(...rows);
 
-      if (data.length < pageSize) {
+      if (rows.length < pageSize) {
         break;
       }
 

package/analyzer-template/project/orchestrateCapture.ts

@@ -23,7 +23,7 @@
  */
 
 import type { Commit, Project } from '~codeyam/types';
-import { updateCommitMetadata } from '~codeyam/supabase';
+import { updateCommitMetadata } from '~codeyam/database';
 
 import type {
   AnalysisLoader,
@@ -50,6 +50,9 @@ interface OrchestrateCaptureArgs {
   taskRunner: CaptureTaskRunner;
   analysisLoader?: AnalysisLoader;
   pollingInterval?: number;
+  // Optional commit for metadata/progress tracking only (not filtering).
+  // Use when full=true but you still want capture progress updates on a commit.
+  metadataCommit?: Commit | null;
 }
 
 export default async function orchestrateCapture({
@@ -62,6 +65,7 @@ export default async function orchestrateCapture({
   taskRunner,
   analysisLoader,
   pollingInterval = POLLING_INTERVAL,
+  metadataCommit,
 }: OrchestrateCaptureArgs) {
   if (!full && !analysisIds && !commit && !entityShas) {
     throw new Error(
@@ -318,9 +322,12 @@ export default async function orchestrateCapture({
         // console.log(
         //   `Codeyam Debug: All known analyses: ${Array.from(knownAnalyses).join(', ')}`,
         // );
-        if (commit) {
-          commit.metadata = await updateCommitMetadata({
-            commitId: commit.id,
+        // Use metadataCommit if provided, otherwise fall back to commit.
+        // This allows progress tracking even when commit is null (e.g., full=true mode).
+        const trackingCommit = metadataCommit ?? commit;
+        if (trackingCommit) {
+          trackingCommit.metadata = await updateCommitMetadata({
+            commitId: trackingCommit.id,
             runStatusUpdate: {
               readyToBeCaptured: knownAnalyses.size,
               capturesAttempted: capturedAnalysisIds.length + abandoned.length,
@@ -331,7 +338,9 @@ export default async function orchestrateCapture({
       }
 
       // Early exit: if we've processed work and the runner has naturally stopped
-      if (knownAnalyses.size > 0 && enqueuedAnalyses.size > 0) {
+      // IMPORTANT: Only exit early if we're draining (analysis phase complete).
+      // Without isDraining check, we might exit before slow-running analyses finish.
+      if (isDraining && knownAnalyses.size > 0 && enqueuedAnalyses.size > 0) {
         const isStillRunning = await taskRunner.isRunning();
         if (
           !isStillRunning &&
@@ -346,9 +355,37 @@ export default async function orchestrateCapture({
         }
       }
 
+      // Early exit when all enqueued work is settled (captured or abandoned)
+      // This avoids waiting for the 20-minute ANALYSIS_TIMEOUT when capture is complete
+      // IMPORTANT: Only check this when isDraining (analysis phase complete) to avoid
+      // exiting while slow-running analyses are still in progress
+      if (
+        isDraining &&
+        !firstLoop &&
+        enqueuedAnalyses.size > 0 &&
+        inFlight.size === 0 &&
+        remaining.length === 0 &&
+        toRetry.length === 0
+      ) {
+        const allAnalysesSettled = Array.from(enqueuedAnalyses).every((id) => {
+          const captured = !allReadyAnalysisIds.includes(id);
+          const abandoned = (attemptCount.get(id) || 0) >= maxAttempts;
+          return captured || abandoned;
+        });
+
+        if (allAnalysesSettled) {
+          console.log(
+            'CodeYam Capture: All enqueued analyses settled (captured or abandoned), exiting early',
+          );
+          break;
+        }
+      }
+
       if (isDraining) {
-        // Exit immediately if analysis complete but no entities ready for capture
-        if (knownAnalyses.size === 0) {
+        // Exit if analysis complete but no entities ready for capture
+        // Only check this after at least one poll (not firstLoop) to ensure we've
+        // had a chance to query the database for ready analyses
+        if (knownAnalyses.size === 0 && !firstLoop) {
           console.log(
             'CodeYam Capture: Analysis complete but no entities ready for capture, exiting',
           );

package/analyzer-template/project/prepareRepo.ts

@@ -39,7 +39,7 @@ export default async function prepareRepo(args: PrepareRepoArgs) {
       throw new Error(`Local project path does not exist: ${localProjectPath}`);
     }
 
-    awsLogDebugLevel(1, 'Rsyncing repo...');
+    awsLogDebugLevel(1, 'Syncing repo...');
 
     // Rsync project files respecting .gitignore to avoid copying secrets
     // (.env files, credentials, private keys, etc.) while preserving node_modules

package/analyzer-template/project/reconcileMockDataKeys.ts

@@ -5,7 +5,13 @@
  * (e.g., "eq('user_id', 'user-12345')") when the mock code expects variable
  * references (e.g., "eq('user_id', user.id)").
  *
- * This function remaps mismatched keys to their expected equivalents.
+ * Keys are now original call signatures (e.g., "useFetcher<User>()", "db.select(query)").
+ *
+ * Also handles flat dotted keys (e.g., "trpc.fastener.getAttributeValues.useQuery(...)")
+ * by converting them to nested structure.
+ *
+ * @param mockData The mock data object to reconcile
+ * @param expectedKeys The expected keys from the structure
  */
 export function reconcileMockDataKeys(
   mockData: Record<string, unknown>,
@@ -26,7 +32,27 @@ export function reconcileMockDataKeys(
       continue;
     }
 
-    // Try to find closest match
+    // Check if this is a flat dotted key that should be converted to nested structure
+    // e.g., "trpc.fastener.getAttributeValues.useQuery({ ... })"
+    const flatKeyParts = parseFlatDottedKey(key);
+    if (flatKeyParts && flatKeyParts.length > 1) {
+      // Check if the parts match expected keys (indicating this should be nested)
+      const partsMatchExpected = flatKeyParts.some((part) =>
+        expectedKeys.includes(part),
+      );
+      if (partsMatchExpected) {
+        // Convert flat key to nested structure and merge
+        deepMergeIntoResult(
+          result,
+          flatKeyParts,
+          reconciledValue,
+          expectedKeys,
+        );
+        continue;
+      }
+    }
+
+    // Try to find closest match using existing logic
     const closestKey = findClosestKey(key, expectedKeys);
     if (closestKey) {
       result[closestKey] = reconciledValue;
@@ -39,6 +65,94 @@ export function reconcileMockDataKeys(
   return result;
 }
 
+/**
+ * Parse a flat dotted key into path segments.
+ * Handles keys like "trpc.fastener.getAttributeValues.useQuery({ ... })"
+ * where the last part might contain parentheses with function arguments.
+ *
+ * Returns null if the key doesn't look like a dotted path.
+ */
+function parseFlatDottedKey(key: string): string[] | null {
+  // Don't treat keys that start with a function call as dotted paths
+  // e.g., "useQuery({ id: 123 })" should not be parsed
+  if (key.match(/^[a-zA-Z_][a-zA-Z0-9_]*\s*\(/)) {
+    // Check if there's a dot BEFORE the first parenthesis
+    const parenIndex = key.indexOf('(');
+    const dotBeforeParen = key.slice(0, parenIndex).includes('.');
+    if (!dotBeforeParen) {
+      return null;
+    }
+  }
+
+  // Must contain a dot to be a dotted path
+  if (!key.includes('.')) {
+    return null;
+  }
+
+  const parts: string[] = [];
+  let current = '';
+  let depth = 0;
+
+  for (let i = 0; i < key.length; i++) {
+    const char = key[i];
+
+    if (char === '(' || char === '{' || char === '[') {
+      depth++;
+      current += char;
+    } else if (char === ')' || char === '}' || char === ']') {
+      depth--;
+      current += char;
+    } else if (char === '.' && depth === 0) {
+      if (current) {
+        parts.push(current);
+      }
+      current = '';
+    } else {
+      current += char;
+    }
+  }
+
+  if (current) {
+    parts.push(current);
+  }
+
+  // Need at least 2 parts for it to be a meaningful dotted path
+  return parts.length >= 2 ? parts : null;
+}
+
+/**
+ * Deep merge a value into the result object following the path segments.
+ * Reconciles the leaf key against expected keys.
+ */
+function deepMergeIntoResult(
+  result: Record<string, unknown>,
+  parts: string[],
+  value: unknown,
+  expectedKeys: string[],
+): void {
+  let current = result;
+
+  // Navigate/create path for all but the last part
+  for (let i = 0; i < parts.length - 1; i++) {
+    const part = parts[i];
+    if (!(part in current)) {
+      current[part] = {};
+    }
+    const next = current[part];
+    if (next && typeof next === 'object' && !Array.isArray(next)) {
+      current = next as Record<string, unknown>;
+    } else {
+      // Can't traverse further, structure mismatch
+      return;
+    }
+  }
+
+  // Handle the leaf key
+  const leafKey = parts[parts.length - 1];
+  const closestLeafKey = findClosestKey(leafKey, expectedKeys) || leafKey;
+  current[closestLeafKey] = value;
+}
+
 /**
  * Find the closest matching key from expectedKeys for a given key.
  *
@@ -47,9 +161,15 @@ export function reconcileMockDataKeys(
  * - Ignore the last argument value (could be literal or variable)
  * - Match based on the common prefix
  * - If key has no parentheses, try matching with "()" appended
+ * - If key is "funcName()" (empty args), match expected key "funcName({...})" by function name
+ * - If key has object arguments (like useQuery({...})), match by function name
+ * - If key has single argument with literal value, match by function name
  *
  * e.g., "eq('user_id', 'user-12345')" matches "eq('user_id', user.id)"
  * e.g., "getAll" matches "getAll()"
+ * e.g., "getSurveysAction()" matches "getSurveysAction({ environmentId, ... })"
+ * e.g., "useQuery({ id: 123 })" matches "useQuery({ id: Number(id) })"
+ * e.g., "useCustomSizes('example-slug')" matches "useCustomSizes(projectSlug)"
  */
 function findClosestKey(key: string, expectedKeys: string[]): string | null {
   const keyBase = extractKeyBase(key);
@@ -73,9 +193,132 @@ function findClosestKey(key: string, expectedKeys: string[]): string | null {
     }
   }
 
+  // If key is "funcName()" (empty args), try matching expected key "funcName({...})"
+  // This handles the case where LLM generates "getSurveysAction()" but expected key
+  // is "getSurveysAction({ environmentId, limit, ... })" with an object argument
+  const emptyArgsMatch = key.match(/^([a-zA-Z_][a-zA-Z0-9_]*)\(\)$/);
+  if (emptyArgsMatch) {
+    const funcName = emptyArgsMatch[1];
+    for (const expected of expectedKeys) {
+      // Match if expected key starts with the same function name and has args
+      const expectedMatch = expected.match(
+        new RegExp(`^${escapeRegExp(funcName)}\\([^)]+\\)$`),
+      );
+      if (expectedMatch) {
+        return expected;
+      }
+    }
+  }
+
+  // If key has object literal arguments (like "useQuery({ id: 123, ... })"),
+  // try matching expected keys with the same function path AND first object property.
+  // This handles cases where LLM generates literal values but expected keys have variable references.
+  // e.g., "useQuery({ attributes: ['a', 'b'] })" matches "useQuery({ attributes: items || [] })"
+  // e.g., "trpc.users.getById.useQuery({ userId: 'user-123' })" matches "trpc.users.getById.useQuery({ userId })"
+  // Uses first property matching when multiple candidates exist for disambiguation.
+  const objectArgMatch = key.match(/^([a-zA-Z_][a-zA-Z0-9_.]*)\s*\(\s*\{/);
+  if (objectArgMatch) {
+    const funcPath = objectArgMatch[1]; // Can be "funcName" or "path.to.funcName"
+    // Extract the first property name from the object argument
+    const firstPropMatch = key.match(
+      /\(\s*\{\s*\n?\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*:/,
+    );
+    const firstProp = firstPropMatch ? firstPropMatch[1] : null;
+
+    // Collect all matching candidates
+    const candidates: string[] = [];
+    for (const expected of expectedKeys) {
+      // Match if expected key has the same function path and starts with object argument
+      const expectedObjArgMatch = expected.match(
+        new RegExp(`^${escapeRegExp(funcPath)}\\s*\\(\\s*\\{`),
+      );
+      if (expectedObjArgMatch) {
+        candidates.push(expected);
+      }
+    }
+
+    // If we have a first property, prefer matching by property name
+    if (firstProp && candidates.length > 1) {
+      // Look for a candidate with the same first property
+      const propRegex = new RegExp(
+        `\\(\\s*\\{[\\s\\n]*${escapeRegExp(firstProp)}\\s*:`,
+      );
+      const propMatch = candidates.find((c) => propRegex.test(c));
+      if (propMatch) {
+        return propMatch;
+      }
+    }
+
+    // Fall back to first candidate if no property match found
+    if (candidates.length > 0) {
+      return candidates[0];
+    }
+  }
+
+  // If key has a single argument (literal value OR variable reference),
+  // try matching expected keys with the same function name but different argument.
+  // This handles cases where:
+  // - LLM generates "useCustomSizes('example-slug')" but expected key is "useCustomSizes(projectSlug)"
+  // - LLM generates "useCustomSizes(someProjectSlug)" but expected key is "useCustomSizes(projectSlug)"
+  //
+  // The second case is important because when LLM generates multiple keys with different
+  // variable names (e.g., "useCustomSizes(someProjectSlug)" and "useCustomSizes(projectSlug)"),
+  // without reconciliation, normalizeMockData merges them into "useCustomSizes()" which
+  // doesn't match the expected key "useCustomSizes(projectSlug)", causing undefined data access.
+  const singleArgMatch = key.match(
+    /^([a-zA-Z_][a-zA-Z0-9_.]*)\s*\(([^,)]+)\)$/,
+  );
+  if (singleArgMatch) {
+    const funcPath = singleArgMatch[1];
+    const arg = singleArgMatch[2].trim();
+    // Apply this matching if the argument is a literal OR a simple variable reference
+    const isLiteralArg = /^['"`].*['"`]$/.test(arg) || /^\d+$/.test(arg);
+    const isVariableRef = /^[a-zA-Z_][a-zA-Z0-9_]*$/.test(arg);
+    if (isLiteralArg || isVariableRef) {
+      for (const expected of expectedKeys) {
+        // Match if expected key has the same function path with a single argument
+        const expectedSingleArgMatch = expected.match(
+          new RegExp(`^${escapeRegExp(funcPath)}\\s*\\([^,)]+\\)$`),
+        );
+        if (expectedSingleArgMatch) {
+          return expected;
+        }
+      }
+    }
+  }
+
+  // If key has multiple arguments where one contains a literal value,
+  // try matching expected keys with the same function name and argument count.
+  // This handles cases like "getScenarioDisplayState(scenario, undefined, undefined, 'entitysha123', undefined)"
+  // matching "getScenarioDisplayState(scenario, undefined, undefined, entitySha, undefined)"
+  const multiArgWithLiteralMatch = key.match(
+    /^([a-zA-Z_][a-zA-Z0-9_.]*)\s*\((.+)\)$/,
+  );
+  if (multiArgWithLiteralMatch) {
+    const funcPath = multiArgWithLiteralMatch[1];
+    const argsStr = multiArgWithLiteralMatch[2];
+    // Check if any argument is a string literal
+    const hasLiteralArg = /'[^']*'|"[^"]*"/.test(argsStr);
+    if (hasLiteralArg) {
+      for (const expected of expectedKeys) {
+        // Match if expected key starts with the same function path
+        if (expected.startsWith(funcPath + '(')) {
+          return expected;
+        }
+      }
+    }
+  }
+
   return null;
 }
 
+/**
+ * Escape special regex characters in a string.
+ */
+function escapeRegExp(str: string): string {
+  return str.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+}
+
 /**
  * Normalize quotes in a string: convert double quotes to single quotes.
  * This handles cases where the LLM generates keys with double quotes