@codeyam/codeyam-cli 0.1.0-staging.b8a55ba → 0.1.0-staging.b8f4f94

package/packages/analyze/src/lib/files/scenarios/gatherDataForMocks.js

@@ -1,5 +1,4 @@
-// import { writeFile } from 'fs/promises';
-import { convertDotNotation, fillInDirectSchemaGapsAndUnknowns, joinParenthesesAndArrays, splitOutsideParenthesesAndArrays, } from "../../../../../../packages/ai/index.js";
+import { convertDotNotation, fillInDirectSchemaGapsAndUnknowns, buildSchemaIndexes, joinParenthesesAndArrays, mergeJsonTypeDefinitions, splitOutsideParenthesesAndArrays, } from "../../../../../../packages/ai/index.js";
 /**
  * Strips generic type parameters from a function name.
  * e.g., "useFetcher<BranchEntityDiffResult>" -> "useFetcher"
@@ -7,6 +6,61 @@ import { convertDotNotation, fillInDirectSchemaGapsAndUnknowns, joinParenthesesA
 function cleanFunctionName(name) {
     return name?.split('<')[0];
 }
+/**
+ * Escapes special regex characters in a string.
+ */
+function escapeRegExp(str) {
+    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) {
+    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) {
+    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
@@ -62,19 +116,23 @@ function preprocessSchemaForMocks(schema) {
  * @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, dependencySchemas, callIndex, variableNameOccurrence, options, hookCallIndex) {
+function processCall(importedExport, dependencySchemas, callIndex, variableNameOccurrence, options, hookCallIndex, prebuiltIndexes) {
     const callName = importedExport.calls?.[callIndex] ?? importedExport.name;
     // Determine if we should use variable-qualified schema or base schema.
     // Use variable-qualified schema when:
-    // 1. variableNameOccurrence is defined (we're processing multiple calls with matching variable names), OR
-    // 2. calls.length equals callVariableNames.length (single call with single variable name)
+    // 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:
-    // - calls.length !== callVariableNames.length (rest spread or method chain case)
+    // - 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.
-    const callsLength = importedExport.calls?.length ?? 0;
+    // 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 || callsLength === varNamesLength;
+    const shouldUseVariableQualifiedSchema = variableNameOccurrence !== undefined ||
+        (directCallCount === 1 && varNamesLength === 1);
     const variableName = shouldUseVariableQualifiedSchema
         ? importedExport.callVariableNames?.[callIndex]
         : undefined;
@@ -86,12 +144,26 @@ function processCall(importedExport, dependencySchemas, callIndex, variableNameO
         variableNameOccurrence > 0
         ? `${variableName}[${variableNameOccurrence}]`
         : variableName;
-    const schemaLookupKey = indexedVariableName
-        ? `${importedExport.name}::${indexedVariableName}`
-        : importedExport.name;
-    let rawSchema = dependencySchemas?.[importedExport.filePath]?.[schemaLookupKey]
-        ?.returnValueSchema ?? {};
-    // If no schema found with indexed key, try without the index as fallback
+    // 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 = {};
+    // 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 ?? {};
+    }
+    // 3. Try variable-qualified key without index (fallback for indexed lookup)
     if (Object.keys(rawSchema).length === 0 &&
         variableNameOccurrence !== undefined &&
         variableNameOccurrence > 0 &&
@@ -101,14 +173,114 @@ function processCall(importedExport, dependencySchemas, callIndex, variableNameO
             dependencySchemas?.[importedExport.filePath]?.[nonIndexedKey]
                 ?.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;
+    // 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 ?? {};
+    // 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;
+    // 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();
+        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;
+    }
+    // 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;
+    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 = {};
+                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 = {};
+                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(fallbackSchema);
-    const functionNames = [];
+    let preprocessedSchema = preprocessSchemaForMocks(mergedSchema);
     const isFunctionNameWithoutFunction = (path, functionName) => {
         // Check if path starts with functionName but isn't a function call
         // Must also handle generics: useLoaderData<typeof loader>() should NOT
@@ -178,30 +350,62 @@ function processCall(importedExport, dependencySchemas, callIndex, variableNameO
             preprocessedSchema = mergedSchema;
         }
     }
-    const relevantMergedDependencySchema = Object.keys(preprocessedSchema).reduce((acc, path) => {
+    // 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();
+    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],
             ]);
-            functionNames.push(functionName);
-            for (const existingPath in acc) {
-                if (isFunctionNameWithoutFunction(existingPath, functionName)) {
-                    delete acc[existingPath];
-                }
-            }
+            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();
+    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);
         }
-        acc[path] = preprocessedSchema[path];
-        return acc;
-    }, {});
-    const schema = convertDotNotation(fillInDirectSchemaGapsAndUnknowns({
+        bucket.push(fn);
+    }
+    // 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 = {};
+    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 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 + '('
     // - Object-style: calls = ['supabase.from(...)'] - call starts with name + '.'
@@ -212,31 +416,11 @@ function processCall(importedExport, dependencySchemas, callIndex, variableNameO
     const schemaKey = variableName
         ? `${baseSchemaKey}::${variableName}`
         : baseSchemaKey;
-    // Build output key
-    // When entityName is provided, use canonical key format for consistency across scenarios
-    // Otherwise, use the variable-based format for backward compatibility
-    let outputKey;
-    if (options?.entityName) {
-        // Canonical key format: EntityName::hookName::index
-        // This ensures consistent keys regardless of variable names chosen during analysis
-        const hookName = cleanFunctionName(importedExport.name);
-        outputKey = `${options.entityName}::${hookName}::${hookCallIndex}`;
-    }
-    else if (variableName) {
-        // Legacy format: "variableName <- functionName"
-        // This reads naturally as "variableName receives from functionName"
-        // When a variable name is reused (variableNameOccurrence > 0), disambiguate with index
-        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];
     if (!returnValue && variableName) {
@@ -249,12 +433,76 @@ function processCall(importedExport, dependencySchemas, callIndex, variableNameO
     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 = schema;
+        for (const part of callParts) {
+            if (current &&
+                typeof current === 'object' &&
+                !Array.isArray(current) &&
+                part in current) {
+                current = current[part];
+            }
+            else {
+                current = undefined;
+                break;
+            }
+        }
+        // Use method-chain result if found, otherwise keep previous returnValue
+        if (current !== undefined) {
+            returnValue = current;
+        }
+    }
     if (!returnValue) {
         return null;
     }
     return { key: outputKey, value: returnValue };
 }
+/**
+ * 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, value) {
+    if (acc[key]) {
+        // Key already exists - merge the schemas
+        acc[key] = mergeJsonTypeDefinitions(acc[key], value);
+    }
+    else {
+        acc[key] = value;
+    }
+}
 export default function gatherDataForMocks(importedExports, dependencySchemas, options) {
+    // 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 = {};
+    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 = {};
@@ -263,12 +511,94 @@ export default function gatherDataForMocks(importedExports, dependencySchemas, o
             return acc;
         // Get the clean hook name for indexing
         const hookName = cleanFunctionName(importedExport.name);
-        // Check if we have multiple calls with variable names
-        const hasMultipleCalls = importedExport.calls &&
-            importedExport.calls.length > 1 &&
+        // 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;
-        if (hasMultipleCalls) {
+            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 = {};
+            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 (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];
+                        }
+                        // 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);
+            }
+        }
+        else if (hasMultipleDistinctCalls) {
             // Track how many times each variable name has been seen
             // to disambiguate keys when the same variable is reused
             const variableNameCounts = {};
@@ -281,20 +611,99 @@ export default function gatherDataForMocks(importedExports, dependencySchemas, o
                 // 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);
+                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
+            // 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;
-            const processResult = processCall(importedExport, dependencySchemas, 0, undefined, options, hookCallIndex);
-            if (processResult) {
-                acc[processResult.key] = processResult.value;
+            // 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).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;
+                    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, varQualifiedResult.value);
+                        }
+                    }
+                    finalValue = mergedValue;
+                }
+                addResultToAccumulator(acc, processResult.key, finalValue);
+            }
+            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 = {};
+                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).length > 0) {
+                    addResultToAccumulator(acc, callSignature, mergedValue);
+                }
             }
         }
         return acc;