@codeyam/codeyam-cli 0.1.0-staging.dd216e0 → 0.1.0-staging.e2d4438

package/analyzer-template/packages/analyze/src/lib/files/scenarios/gatherDataForMocks.ts

@@ -1,8 +1,10 @@
-// import { writeFile } from 'fs/promises';
 import {
   convertDotNotation,
   fillInDirectSchemaGapsAndUnknowns,
+  buildSchemaIndexes,
+  type SchemaIndexes,
   joinParenthesesAndArrays,
+  mergeJsonTypeDefinitions,
   splitOutsideParenthesesAndArrays,
 } from '~codeyam/ai';
 import { DataStructure, JsonTypeDefinition } from '~codeyam/types';
@@ -15,6 +17,65 @@ function cleanFunctionName(name: string): string {
   return name?.split('<')[0];
 }
 
+/**
+ * Escapes special regex characters in a string.
+ */
+function escapeRegExp(str: string): string {
+  return str.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+}
+
+/**
+ * Counts only DIRECT function calls, excluding method chains on return values.
+ *
+ * When a hook returns an object with methods that are later called, the analysis
+ * tracks those as call signatures like 'useAuth().functionCallReturnValue.signOut()'.
+ * These are NOT separate calls to useAuth - they're method invocations on the return value.
+ *
+ * This function filters out such method chains to get the true count of direct calls.
+ *
+ * @example
+ * // One direct call, one method on return value
+ * countDirectCalls(['useAuth()', 'useAuth().functionCallReturnValue.signOut()']) // => 1
+ *
+ * // Two direct calls
+ * countDirectCalls(['useFetcher()', 'useFetcher()']) // => 2
+ */
+function countDirectCalls(calls: string[]): number {
+  return calls.filter((sig) => !sig.includes('.functionCallReturnValue.'))
+    .length;
+}
+
+/**
+ * Checks if all direct calls (non-method-chain calls) are identical.
+ * This happens when destructuring a hook call:
+ *   const { authLoading, signOut } = useAuth();
+ * produces calls: ['useAuth()', 'useAuth()'] (same signature repeated)
+ *
+ * In this case, we should use the full unqualified schema since there's
+ * really only ONE function call, not multiple different calls.
+ *
+ * @example
+ * // Same call repeated (from destructuring)
+ * allDirectCallsIdentical(['useAuth()', 'useAuth()']) // => true
+ *
+ * // Different calls
+ * allDirectCallsIdentical(['useFetcher<User>()', 'useFetcher<Report>()']) // => false
+ *
+ * // Single call
+ * allDirectCallsIdentical(['useAuth()']) // => true
+ *
+ * // With method chain (method chain is ignored)
+ * allDirectCallsIdentical(['useAuth()', 'useAuth().functionCallReturnValue.signOut()']) // => true
+ */
+function allDirectCallsIdentical(calls: string[]): boolean {
+  const directCalls = calls.filter(
+    (sig) => !sig.includes('.functionCallReturnValue.'),
+  );
+  if (directCalls.length === 0) return true;
+  const firstCall = directCalls[0];
+  return directCalls.every((call) => call === firstCall);
+}
+
 /**
  * Preprocesses the schema to:
  * 1. Filter out all signature paths
@@ -68,6 +129,18 @@ function preprocessSchemaForMocks(
   return result;
 }
 
+/**
+ * Options for gatherDataForMocks
+ */
+export interface GatherDataForMocksOptions {
+  /**
+   * When provided, generates canonical keys in format `EntityName::hookName::index`.
+   * This ensures consistent keys across scenarios regardless of variable names
+   * chosen during analysis.
+   */
+  entityName?: string;
+}
+
 /**
  * Processes a single call (with optional variable name) and returns the mock data entry.
  * Returns null if no valid schema is found.
@@ -78,6 +151,8 @@ function preprocessSchemaForMocks(
  * @param variableNameOccurrence - When a variable name is reused, this is the occurrence index
  *                                  (e.g., first 'data' = 0, second 'data' = 1)
  *                                  Used to disambiguate keys when the same variable is reused.
+ * @param options - Additional options including entityName for canonical key generation
+ * @param hookCallIndex - The global index of this hook call for canonical key generation
  */
 function processCall(
   importedExport: {
@@ -89,31 +164,194 @@ function processCall(
   },
   dependencySchemas: DataStructure['dependencySchemas'],
   callIndex: number,
-  variableNameOccurrence?: number,
+  variableNameOccurrence: number | undefined,
+  options: GatherDataForMocksOptions | undefined,
+  hookCallIndex: number,
+  prebuiltIndexes?: SchemaIndexes,
 ): { key: string; value: JsonTypeDefinition } | null {
   const callName = importedExport.calls?.[callIndex] ?? importedExport.name;
-  const variableName = importedExport.callVariableNames?.[callIndex];
+
+  // Determine if we should use variable-qualified schema or base schema.
+  // Use variable-qualified schema when:
+  // 1. variableNameOccurrence is defined (we're processing multiple DISTINCT calls), OR
+  // 2. Single call with single variable name (directCallCount === 1 && varNamesLength === 1)
+  // Use base schema when:
+  // - Multiple calls that are all identical (destructuring case: const { a, b } = useHook())
+  // - directCallCount !== varNamesLength (rest spread or method chain case)
+  // This is important for rest spread destructuring like: const { x, y, ...data } = useLoaderData()
+  // where data.queueState.jobs should come from the base schema, not from a variable-qualified schema.
+  // Also important: method chains like useAuth().functionCallReturnValue.signOut() should NOT count
+  // as separate calls - only direct hook calls count.
+  const directCallCount = countDirectCalls(importedExport.calls ?? []);
+  const varNamesLength = importedExport.callVariableNames?.length ?? 0;
+  const shouldUseVariableQualifiedSchema =
+    variableNameOccurrence !== undefined ||
+    (directCallCount === 1 && varNamesLength === 1);
+  const variableName = shouldUseVariableQualifiedSchema
+    ? importedExport.callVariableNames?.[callIndex]
+    : undefined;
 
   // Build the schema lookup key - use variable-qualified key if available
-  const schemaLookupKey = variableName
-    ? `${importedExport.name}::${variableName}`
-    : importedExport.name;
+  // When the same variable name is reused (variableNameOccurrence > 0), try indexed lookup first
+  // e.g., "useFetcher::fetcher[1]" for the second use of "fetcher"
+  const indexedVariableName =
+    variableName &&
+    variableNameOccurrence !== undefined &&
+    variableNameOccurrence > 0
+      ? `${variableName}[${variableNameOccurrence}]`
+      : variableName;
+
+  // Schema lookup order (most specific to least specific):
+  // 1. Call signature (e.g., 'useLoaderData<typeof loader>()')
+  // 2. Variable-qualified key with index (e.g., 'useFetcher::fetcher[1]')
+  // 3. Variable-qualified key without index (e.g., 'useFetcher::fetcher')
+  // 4. Base name (e.g., 'useLoaderData')
+
+  let rawSchema: Record<string, string> = {};
+
+  // 1. Try call signature first (most specific - new format)
+  if (callName !== importedExport.name) {
+    rawSchema =
+      dependencySchemas?.[importedExport.filePath]?.[callName]
+        ?.returnValueSchema ?? {};
+  }
+
+  // 2. Try variable-qualified key with index
+  if (Object.keys(rawSchema).length === 0 && indexedVariableName) {
+    const indexedKey = `${importedExport.name}::${indexedVariableName}`;
+    rawSchema =
+      dependencySchemas?.[importedExport.filePath]?.[indexedKey]
+        ?.returnValueSchema ?? {};
+  }
 
-  const rawSchema =
-    dependencySchemas?.[importedExport.filePath]?.[schemaLookupKey]
+  // 3. Try variable-qualified key without index (fallback for indexed lookup)
+  if (
+    Object.keys(rawSchema).length === 0 &&
+    variableNameOccurrence !== undefined &&
+    variableNameOccurrence > 0 &&
+    variableName
+  ) {
+    const nonIndexedKey = `${importedExport.name}::${variableName}`;
+    rawSchema =
+      dependencySchemas?.[importedExport.filePath]?.[nonIndexedKey]
+        ?.returnValueSchema ?? {};
+  }
+
+  // 4. Get base name schema and merge with call-specific schema
+  // The base schema often has more complete array element details that the
+  // call-specific schema might be missing. Instead of using as fallback only
+  // when empty, MERGE both schemas so we get the most complete picture.
+  //
+  // HOWEVER: For type-parameterized calls where the base schema contains paths
+  // from MULTIPLE DIFFERENT type parameters, DON'T merge to avoid contamination.
+  // The call-specific schema (rawSchema) is already isolated and complete.
+  const baseSchema =
+    dependencySchemas?.[importedExport.filePath]?.[importedExport.name]
       ?.returnValueSchema ?? {};
 
-  // If no schema found with variable-qualified key, fall back to base name
-  const fallbackSchema =
-    Object.keys(rawSchema).length === 0
-      ? (dependencySchemas?.[importedExport.filePath]?.[importedExport.name]
-          ?.returnValueSchema ?? {})
-      : rawSchema;
+  // Check if this is a type-parameterized call with a call-specific schema
+  const hasTypeParams = callName.includes('<') && callName.includes('>');
+  const hasCallSpecificSchema = Object.keys(rawSchema).length > 0;
 
-  // Preprocess to filter signatures and extract functionCallReturnValue
-  let preprocessedSchema = preprocessSchemaForMocks(fallbackSchema);
+  // Check if the base schema has multiple DIFFERENT type parameter variants
+  // If so, merging would bring in wrong fields from other variants
+  let baseHasMultipleTypeParamVariants = false;
+  if (hasTypeParams && hasCallSpecificSchema) {
+    const baseFunctionName = cleanFunctionName(callName.split('(')[0]);
+    const typeParamVariantsInBase = new Set<string>();
+
+    for (const path of Object.keys(baseSchema)) {
+      // Look for paths like "useFetcher<TypeA>(...)" or "useFetcher<TypeB>(...)"
+      const match = path.match(
+        new RegExp(`^${escapeRegExp(baseFunctionName)}<([^>]+)>\\(`),
+      );
+      if (match) {
+        typeParamVariantsInBase.add(match[1]);
+      }
+    }
+    baseHasMultipleTypeParamVariants = typeParamVariantsInBase.size > 1;
+  }
 
-  const functionNames: string[] = [];
+  // For type-parameterized calls where we have a call-specific schema, we need to
+  // be careful about merging with the base schema. The base schema may contain
+  // merged data from ALL type parameter variants, which would pollute the result.
+  let mergedSchema: Record<string, string>;
+  if (hasCallSpecificSchema && hasTypeParams) {
+    if (baseHasMultipleTypeParamVariants) {
+      // Base schema has paths with different type parameters - filter to only
+      // include paths matching the current type parameter.
+      const typeParamMatch = callName.match(/<([^>]+)>/);
+      const currentTypeParam = typeParamMatch ? typeParamMatch[1] : null;
+      const baseFunctionName = cleanFunctionName(callName.split('(')[0]);
+
+      if (currentTypeParam) {
+        // Filter base schema to only include paths that:
+        // 1. Match current type param (e.g., useLoaderData<LoaderData>())
+        // 2. OR don't have any type param (generic paths like useLoaderData())
+        const filteredBaseSchema: Record<string, string> = {};
+        for (const [path, type] of Object.entries(baseSchema)) {
+          // Check if path has a type parameter for this function
+          const pathTypeParamMatch = path.match(
+            new RegExp(`^${escapeRegExp(baseFunctionName)}<([^>]+)>\\(`),
+          );
+          if (pathTypeParamMatch) {
+            // Path has a type param - only include if it matches current
+            if (pathTypeParamMatch[1] === currentTypeParam) {
+              filteredBaseSchema[path] = type;
+            }
+            // Otherwise skip (it's for a different type param variant)
+          } else {
+            // Path doesn't have a type param for this function - include it
+            // (could be generic paths or paths for other functions)
+            filteredBaseSchema[path] = type;
+          }
+        }
+        mergedSchema = { ...filteredBaseSchema, ...rawSchema };
+      } else {
+        // Couldn't extract type param - fall back to rawSchema only
+        mergedSchema = rawSchema;
+      }
+    } else {
+      // Base schema has zero or one type-parameterized variant.
+      // When there's only one variant, the base schema paths are safe to merge
+      // (no contamination from other variants). This is important for cases like:
+      // - Array element details (gitStatus[].path) that are in base but not call-specific
+      // - Enriched paths that should be preserved
+      // Filter to only include paths matching current type param or generic paths.
+      const typeParamMatch = callName.match(/<([^>]+)>/);
+      const currentTypeParam = typeParamMatch ? typeParamMatch[1] : null;
+      const baseFunctionName = cleanFunctionName(callName.split('(')[0]);
+
+      if (currentTypeParam) {
+        const filteredBaseSchema: Record<string, string> = {};
+        for (const [path, type] of Object.entries(baseSchema)) {
+          const pathTypeParamMatch = path.match(
+            new RegExp(`^${escapeRegExp(baseFunctionName)}<([^>]+)>\\(`),
+          );
+          if (pathTypeParamMatch) {
+            // Path has a type param - only include if it matches current
+            if (pathTypeParamMatch[1] === currentTypeParam) {
+              filteredBaseSchema[path] = type;
+            }
+          } else {
+            // Path doesn't have a type param - include generic paths
+            filteredBaseSchema[path] = type;
+          }
+        }
+        mergedSchema = { ...filteredBaseSchema, ...rawSchema };
+      } else {
+        // No type param to extract - merge normally
+        mergedSchema = { ...baseSchema, ...rawSchema };
+      }
+    }
+  } else if (Object.keys(rawSchema).length === 0) {
+    mergedSchema = baseSchema;
+  } else {
+    mergedSchema = { ...baseSchema, ...rawSchema };
+  }
+
+  // Preprocess to filter signatures and extract functionCallReturnValue
+  let preprocessedSchema = preprocessSchemaForMocks(mergedSchema);
 
   const isFunctionNameWithoutFunction = (
     path: string,
@@ -199,38 +437,66 @@ function processCall(
     }
   }
 
-  const relevantMergedDependencySchema = Object.keys(preprocessedSchema).reduce(
-    (acc, path) => {
-      if (path.endsWith(')')) {
-        const pathParts = splitOutsideParenthesesAndArrays(path);
-        const functionName = joinParenthesesAndArrays([
-          ...pathParts.slice(0, -1),
-          pathParts[pathParts.length - 1].split('(')[0],
-        ]);
-        functionNames.push(functionName);
-        for (const existingPath in acc) {
-          if (isFunctionNameWithoutFunction(existingPath, functionName)) {
-            delete acc[existingPath];
-          }
-        }
-      }
+  // Two-pass function name filtering to avoid O(k²) retroactive deletion.
+  // Pass 1: Discover all unique function names from paths ending with ')'
+  const functionNameSet = new Set<string>();
+  const schemaKeys = Object.keys(preprocessedSchema);
+  for (const path of schemaKeys) {
+    if (path.endsWith(')')) {
+      const pathParts = splitOutsideParenthesesAndArrays(path);
+      const functionName = joinParenthesesAndArrays([
+        ...pathParts.slice(0, -1),
+        pathParts[pathParts.length - 1].split('(')[0],
+      ]);
+      functionNameSet.add(functionName);
+    }
+  }
 
-      if (functionNames.some((fn) => isFunctionNameWithoutFunction(path, fn))) {
-        return acc;
-      }
+  // Index function names by first segment for O(1) bucket lookup instead of
+  // checking all function names for every path. Most paths only need to check
+  // 1-2 function names instead of all F.
+  const fnByFirstSegment = new Map<string, string[]>();
+  for (const fn of functionNameSet) {
+    const dotIdx = fn.indexOf('.');
+    const firstSeg = dotIdx >= 0 ? fn.slice(0, dotIdx) : fn;
+    let bucket = fnByFirstSegment.get(firstSeg);
+    if (!bucket) {
+      bucket = [];
+      fnByFirstSegment.set(firstSeg, bucket);
+    }
+    bucket.push(fn);
+  }
 
-      acc[path] = preprocessedSchema[path];
-      return acc;
-    },
-    {} as Record<string, string>,
-  );
+  // Pass 2: Filter paths that match any discovered function name prefix.
+  // Use indexed lookup: extract path's first segment, check only matching function names.
+  const relevantMergedDependencySchema: Record<string, string> = {};
+  for (const path of schemaKeys) {
+    // Extract first segment of path (before first '.', '(', or '<')
+    const dotIdx = path.indexOf('.');
+    const parenIdx = path.indexOf('(');
+    const angleIdx = path.indexOf('<');
+    let minIdx = path.length;
+    if (dotIdx >= 0 && dotIdx < minIdx) minIdx = dotIdx;
+    if (parenIdx >= 0 && parenIdx < minIdx) minIdx = parenIdx;
+    if (angleIdx >= 0 && angleIdx < minIdx) minIdx = angleIdx;
+    const firstSeg = path.slice(0, minIdx);
+
+    const candidates = fnByFirstSegment.get(firstSeg);
+    if (
+      !candidates ||
+      !candidates.some((fn) => isFunctionNameWithoutFunction(path, fn))
+    ) {
+      relevantMergedDependencySchema[path] = preprocessedSchema[path];
+    }
+  }
 
-  const schema = convertDotNotation(
-    fillInDirectSchemaGapsAndUnknowns({
-      scopeName: importedExport.name,
-      schema: relevantMergedDependencySchema,
-    }),
-  );
+  const filledSchema = fillInDirectSchemaGapsAndUnknowns({
+    scopeName: importedExport.name,
+    schema: relevantMergedDependencySchema,
+    prebuiltIndexes,
+  });
+
+  const schema = convertDotNotation(filledSchema);
 
   // Determine if this is an object-style or hook-style dependency:
   // - Hook-style: calls = ['useAuth()'] - call starts with name + '('
@@ -245,23 +511,11 @@ function processCall(
     ? `${baseSchemaKey}::${variableName}`
     : baseSchemaKey;
 
-  // Build output key - for variable-qualified, use "variableName <- functionName" format
-  // This reads naturally as "variableName receives from functionName"
-  // When a variable name is reused (variableNameOccurrence > 0), disambiguate with index
-  let outputKey: string;
-  if (variableName) {
-    // When the same variable name is used multiple times, append the occurrence index
-    // to create unique keys: "data[0] <- useFetch", "data[1] <- useFetch"
-    const disambiguator =
-      variableNameOccurrence !== undefined && variableNameOccurrence > 0
-        ? `[${variableNameOccurrence}]`
-        : '';
-    outputKey = `${variableName}${disambiguator} <- ${importedExport.name}`;
-  } else if (isObjectStyle) {
-    outputKey = importedExport.name;
-  } else {
-    outputKey = callName;
-  }
+  // Build output key - use the ORIGINAL CALL SIGNATURE from the code
+  // This makes it clear to the LLM exactly which call needs which data.
+  // Example: useFetcher<{ data: UserData }>() vs useFetcher<{ data: ReportData }>()
+  // Example: db.select(usersQuery) vs db.select(postsQuery)
+  const outputKey = callName;
 
   // Try to find return value - check with variable qualifier first, then without
   let returnValue = schema[schemaKey];
@@ -271,6 +525,38 @@ function processCall(
   if (!returnValue) {
     returnValue = schema['returnValue'];
   }
+  // Fallback for function call return values (e.g., useFetcher returns via functionCallReturnValue)
+  if (!returnValue) {
+    returnValue = schema['functionCallReturnValue'];
+  }
+
+  // Handle method-chain style calls (e.g., createClient().from("workouts").select("*") or db.select(query))
+  // After convertDotNotation, the schema is nested: { "createClient()": { "from(...)": { ... } } }
+  // We need to navigate the nested structure using the call signature parts.
+  // For object-style dependencies (isObjectStyle=true), the baseSchemaKey is just the object name (e.g., "db"),
+  // so returnValue might be the parent object, not the actual return value of the method chain.
+  // We should always try method-chain navigation for calls with dots and parens.
+  if (callName.includes('.') && callName.includes('(')) {
+    const callParts = splitOutsideParenthesesAndArrays(callName);
+    let current: unknown = schema;
+    for (const part of callParts) {
+      if (
+        current &&
+        typeof current === 'object' &&
+        !Array.isArray(current) &&
+        part in (current as Record<string, unknown>)
+      ) {
+        current = (current as Record<string, unknown>)[part];
+      } else {
+        current = undefined;
+        break;
+      }
+    }
+    // Use method-chain result if found, otherwise keep previous returnValue
+    if (current !== undefined) {
+      returnValue = current as typeof returnValue;
+    }
+  }
 
   if (!returnValue) {
     return null;
@@ -279,6 +565,27 @@ function processCall(
   return { key: outputKey, value: returnValue as JsonTypeDefinition };
 }
 
+/**
+ * Adds a result to the accumulator, merging schemas if the key already exists.
+ * This handles the case where the same call signature appears multiple times
+ * (e.g., two calls to useFetcher() with the same generic type).
+ */
+function addResultToAccumulator(
+  acc: { [key: string]: JsonTypeDefinition },
+  key: string,
+  value: JsonTypeDefinition,
+): void {
+  if (acc[key]) {
+    // Key already exists - merge the schemas
+    acc[key] = mergeJsonTypeDefinitions(
+      acc[key] as Record<string, unknown>,
+      value as Record<string, unknown>,
+    ) as JsonTypeDefinition;
+  } else {
+    acc[key] = value;
+  }
+}
+
 export default function gatherDataForMocks(
   importedExports: {
     filePath?: string;
@@ -288,19 +595,162 @@ export default function gatherDataForMocks(
     callVariableNames?: string[];
   }[],
   dependencySchemas: DataStructure['dependencySchemas'],
+  options?: GatherDataForMocksOptions,
 ) {
+  // Build schema indexes ONCE from all dependency schemas combined.
+  // This avoids rebuilding indexes (~1-2 seconds for 18k keys) for each processCall.
+  //
+  // IMPORTANT: We must preprocess schemas BEFORE building indexes!
+  // The schemas contain .functionCallReturnValue paths that get stripped during processing.
+  // If we build indexes from original paths, lookups with preprocessed paths will fail.
+  // Example: 'hook().functionCallReturnValue.x.length' → index maps 'hook().functionCallReturnValue.x' → ['length']
+  //          But preprocessed path 'hook().x.length' looks for 'hook().x' which doesn't exist in the index.
+  const allSchemaKeys: Record<string, string> = {};
+  if (dependencySchemas) {
+    for (const filePath in dependencySchemas) {
+      for (const entityName in dependencySchemas[filePath]) {
+        const schema =
+          dependencySchemas[filePath][entityName]?.returnValueSchema;
+        if (schema) {
+          // Preprocess each schema to remove .functionCallReturnValue before collecting
+          const preprocessed = preprocessSchemaForMocks(schema);
+          Object.assign(allSchemaKeys, preprocessed);
+        }
+      }
+    }
+  }
+  const prebuiltIndexes =
+    Object.keys(allSchemaKeys).length > 0
+      ? buildSchemaIndexes(allSchemaKeys)
+      : undefined;
+
+  // Track the global index for each hook name across all imports
+  // This enables canonical keys like EntityName::useLoaderData::0, EntityName::useFetcher::0, etc.
+  const hookCallIndices: Record<string, number> = {};
+
   const result = importedExports.reduce(
     (acc, importedExport) => {
       if (!importedExport.isMocked) return acc;
 
-      // Check if we have multiple calls with variable names
-      const hasMultipleCalls =
-        importedExport.calls &&
-        importedExport.calls.length > 1 &&
+      // Get the clean hook name for indexing
+      const hookName = cleanFunctionName(importedExport.name);
+
+      // Check if we have multiple DISTINCT direct calls with variable names.
+      // Important: Only count direct calls to the function, not method chains on its return value.
+      // For example, ['useAuth()', 'useAuth().functionCallReturnValue.signOut()'] has only 1 direct call.
+      //
+      // ALSO: If all direct calls are IDENTICAL (same signature repeated), this is from
+      // destructuring like `const { a, b } = useHook()` and we should use the FULL unqualified
+      // schema, not iterate through partial variable-qualified schemas.
+      const directCallCount = countDirectCalls(importedExport.calls ?? []);
+      const hasMultipleDistinctCalls =
+        directCallCount > 1 &&
         importedExport.callVariableNames &&
-        importedExport.callVariableNames.length === importedExport.calls.length;
+        importedExport.callVariableNames.length === directCallCount &&
+        !allDirectCallsIdentical(importedExport.calls ?? []);
+
+      // Check for multiple distinct method-chain calls WITHOUT callVariableNames.
+      // This pattern is common with tRPC, Supabase, etc. where a single import
+      // (e.g., `trpc`) is used with multiple different method chains (e.g.,
+      // `trpc.useUtils()`, `trpc.quote.getQuotesByOrg.useQuery()`).
+      // These should each get their own entry in dataForMocks.
+      //
+      // IMPORTANT: This is DIFFERENT from the rest spread pattern where a function
+      // like useLoaderData is called directly, and other calls are method chains
+      // on its return value (e.g., `useLoaderData<typeof loader>()` with
+      // `useLoaderData<typeof loader>().uncommittedEntitiesList.filter(...)`).
+      // In that case, all the additional calls start with the base function call.
+      const directCalls = (importedExport.calls ?? []).filter(
+        (sig) => !sig.includes('.functionCallReturnValue.'),
+      );
+      const uniqueDirectCalls = [...new Set(directCalls)];
+
+      // Check if any call is a direct invocation of the import (e.g., `useLoaderData<typeof loader>()`)
+      // as opposed to all being property accesses (e.g., `trpc.useUtils()`, `trpc.quote.xxx()`).
+      // If there's a direct function call, the other calls are likely method chains on its result.
+      const baseName = importedExport.name;
+      const hasDirectFunctionCall = directCalls.some((sig) => {
+        // Pattern: baseName() or baseName<...>()
+        // This matches "useLoaderData()", "useLoaderData<typeof loader>()", etc.
+        // but NOT "trpc.useUtils()" which has a dot before the parens
+        const baseCallPattern = new RegExp(`^${baseName}(<[^>]*>)?\\(`);
+        return baseCallPattern.test(sig);
+      });
+
+      const hasMultipleDistinctMethodChains =
+        !hasMultipleDistinctCalls &&
+        uniqueDirectCalls.length > 1 &&
+        !allDirectCallsIdentical(directCalls) &&
+        !hasDirectFunctionCall; // Exclude rest spread patterns
+
+      if (hasMultipleDistinctMethodChains) {
+        // For multiple method-chain calls on the same import, create a NESTED structure
+        // under the base import name. This allows constructMockCode to generate a
+        // combined mock object with all method chains.
+        //
+        // Example: For calls ['trpc.useUtils()', 'trpc.quote.getQuotesByOrg.useQuery()'],
+        // create: { "trpc": { "useUtils()": {...}, "quote": { "getQuotesByOrg": { "useQuery()": {...} } } } }
+        const nestedResult: Record<string, unknown> = {};
+
+        for (let i = 0; i < uniqueDirectCalls.length; i++) {
+          // Find the index of this call in the original calls array
+          const callIndex = (importedExport.calls ?? []).indexOf(
+            uniqueDirectCalls[i],
+          );
+          if (callIndex === -1) continue;
+
+          const hookCallIndex = hookCallIndices[hookName] ?? 0;
+          hookCallIndices[hookName] = hookCallIndex + 1;
+
+          const processResult = processCall(
+            importedExport,
+            dependencySchemas,
+            callIndex,
+            undefined,
+            options,
+            hookCallIndex,
+            prebuiltIndexes,
+          );
 
-      if (hasMultipleCalls) {
+          if (processResult) {
+            // The key is like "trpc.useUtils()" or "trpc.quote.getQuotesByOrg.useQuery()"
+            // We need to nest the value at the path AFTER the base name
+            const callSignature = processResult.key;
+
+            // Split the call signature into path parts, handling parens
+            // e.g., "trpc.useUtils()" -> ["trpc", "useUtils()"]
+            // e.g., "trpc.quote.getQuotesByOrg.useQuery()" -> ["trpc", "quote", "getQuotesByOrg", "useQuery()"]
+            const pathParts = splitOutsideParenthesesAndArrays(callSignature);
+
+            // Skip the base name (first part) and nest the value at the remaining path
+            if (pathParts.length > 1 && pathParts[0] === baseName) {
+              const nestedPath = pathParts.slice(1);
+
+              // Set the value at the nested path, merging with existing structure
+              let current = nestedResult;
+              for (let j = 0; j < nestedPath.length - 1; j++) {
+                const part = nestedPath[j];
+                if (!(part in current)) {
+                  current[part] = {};
+                }
+                current = current[part] as Record<string, unknown>;
+              }
+              // Set the leaf value
+              const leafKey = nestedPath[nestedPath.length - 1];
+              current[leafKey] = processResult.value;
+            }
+          }
+        }
+
+        // Add the nested result under the base import name
+        if (Object.keys(nestedResult).length > 0) {
+          addResultToAccumulator(
+            acc,
+            baseName,
+            nestedResult as JsonTypeDefinition,
+          );
+        }
+      } else if (hasMultipleDistinctCalls) {
         // Track how many times each variable name has been seen
         // to disambiguate keys when the same variable is reused
         const variableNameCounts: Record<string, number> = {};
@@ -313,21 +763,167 @@ export default function gatherDataForMocks(
           const occurrence = variableNameCounts[varName] ?? 0;
           variableNameCounts[varName] = occurrence + 1;
 
+          // Get the current hook call index and increment for next call
+          const hookCallIndex = hookCallIndices[hookName] ?? 0;
+          hookCallIndices[hookName] = hookCallIndex + 1;
+
           const processResult = processCall(
             importedExport,
             dependencySchemas,
             i,
             occurrence,
+            options,
+            hookCallIndex,
+            prebuiltIndexes,
           );
           if (processResult) {
-            acc[processResult.key] = processResult.value;
+            addResultToAccumulator(acc, processResult.key, processResult.value);
+          }
+        }
+      } else if (
+        // Handle multiple UNIQUE type-parameterized direct calls when variable count doesn't match.
+        // This happens when a hook is called with different type parameters (useFetcher<TypeA>(), useFetcher<TypeB>()),
+        // but the callVariableNames count doesn't match directCallCount (common with rest spread patterns).
+        // Each unique call signature should get its own entry in dataForMocks.
+        uniqueDirectCalls.length > 1 &&
+        !allDirectCallsIdentical(directCalls) &&
+        hasDirectFunctionCall
+      ) {
+        // Process each UNIQUE direct call signature
+        for (let i = 0; i < uniqueDirectCalls.length; i++) {
+          const callSignature = uniqueDirectCalls[i];
+
+          // Find the index of this call in the original calls array
+          const callIndex = (importedExport.calls ?? []).indexOf(callSignature);
+          if (callIndex === -1) continue;
+
+          const hookCallIndex = hookCallIndices[hookName] ?? 0;
+          hookCallIndices[hookName] = hookCallIndex + 1;
+
+          const processResult = processCall(
+            importedExport,
+            dependencySchemas,
+            callIndex,
+            undefined,
+            options,
+            hookCallIndex,
+            prebuiltIndexes,
+          );
+
+          if (processResult) {
+            addResultToAccumulator(acc, processResult.key, processResult.value);
           }
         }
       } else {
-        // Original behavior for single calls or calls without variable names
-        const processResult = processCall(importedExport, dependencySchemas, 0);
-        if (processResult) {
-          acc[processResult.key] = processResult.value;
+        // Single calls, or multiple identical calls (from destructuring)
+        // Get the current hook call index and increment for next call
+        const hookCallIndex = hookCallIndices[hookName] ?? 0;
+        hookCallIndices[hookName] = hookCallIndex + 1;
+
+        // First, try to get result from unqualified schema
+        const processResult = processCall(
+          importedExport,
+          dependencySchemas,
+          0,
+          undefined,
+          options,
+          hookCallIndex,
+          prebuiltIndexes,
+        );
+
+        // Check if we got a meaningful result (not empty object or null)
+        // Arrays are always meaningful (even empty ones represent valid data structure)
+        const hasUnqualifiedResult =
+          processResult &&
+          processResult.value !== null &&
+          (typeof processResult.value !== 'object' ||
+            Array.isArray(processResult.value) ||
+            Object.keys(processResult.value as object).length > 0);
+
+        // Check if we have multiple callVariableNames that are DIFFERENT
+        // (e.g., from parent + child components using the same hook differently)
+        // In this case, we should also merge in variable-qualified schemas
+        // to capture all destructured properties from both parent and child.
+        const uniqueCallVarNames = new Set(
+          importedExport.callVariableNames ?? [],
+        );
+        const hasDifferentVarNames =
+          importedExport.callVariableNames &&
+          importedExport.callVariableNames.length > 1 &&
+          uniqueCallVarNames.size > 1 &&
+          // Only merge variable-qualified schemas when calls are all identical
+          // (e.g., useFlags() repeated). Don't do this for method chains
+          // or rest spread cases where calls.length !== callVariableNames.length
+          allDirectCallsIdentical(importedExport.calls ?? []);
+
+        if (hasUnqualifiedResult) {
+          let finalValue = processResult.value;
+
+          // If we have different variable names (e.g., from parent and child components),
+          // also merge in variable-qualified schemas to ensure we capture all properties.
+          // This handles the case where parent uses useFlags().enableNewAdminDashboard
+          // and child uses useFlags().flagShowReferenceIdOnQuoteTable.
+          if (hasDifferentVarNames) {
+            let mergedValue = processResult.value as JsonTypeDefinition;
+            for (let i = 0; i < importedExport.callVariableNames!.length; i++) {
+              const varQualifiedResult = processCall(
+                importedExport,
+                dependencySchemas,
+                i,
+                0, // Use 0 to trigger variable-qualified lookup
+                options,
+                hookCallIndex,
+                prebuiltIndexes,
+              );
+              if (varQualifiedResult && varQualifiedResult.value) {
+                mergedValue = mergeJsonTypeDefinitions(
+                  mergedValue as Record<string, unknown>,
+                  varQualifiedResult.value as Record<string, unknown>,
+                ) as JsonTypeDefinition;
+              }
+            }
+            finalValue = mergedValue;
+          }
+
+          addResultToAccumulator(
+            acc,
+            processResult.key,
+            finalValue as JsonTypeDefinition,
+          );
+        } else if (
+          directCallCount > 1 &&
+          allDirectCallsIdentical(importedExport.calls ?? []) &&
+          importedExport.callVariableNames &&
+          importedExport.callVariableNames.length > 0
+        ) {
+          // No unqualified schema exists, but we have multiple identical calls with variable names
+          // Merge all the variable-qualified schemas together
+          const callSignature =
+            importedExport.calls?.[0] ?? `${importedExport.name}()`;
+          let mergedValue: JsonTypeDefinition = {};
+
+          for (let i = 0; i < importedExport.callVariableNames.length; i++) {
+            const varQualifiedResult = processCall(
+              importedExport,
+              dependencySchemas,
+              i,
+              0, // Use 0 to trigger variable-qualified lookup
+              options,
+              hookCallIndex,
+              prebuiltIndexes,
+            );
+            if (varQualifiedResult && varQualifiedResult.value) {
+              // Deep merge the value into the accumulated result
+              mergedValue = mergeJsonTypeDefinitions(
+                mergedValue,
+                varQualifiedResult.value,
+              );
+            }
+          }
+
+          if (Object.keys(mergedValue as object).length > 0) {
+            addResultToAccumulator(acc, callSignature, mergedValue);
+          }
         }
       }