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

package/analyzer-template/project/orchestrateCapture.ts

@@ -32,7 +32,10 @@ import type {
 } from './orchestrateCapture/taskRunner';
 import KyselyAnalysisLoader from './orchestrateCapture/KyselyAnalysisLoader';
 
-const MAX_BATCH_SIZE = 6;
+// Temporarily set to 1 (effectively disabling batching) to work around reliability
+// issues where multi-analysis batches cause cascading failures. Revisit once we
+// reassess the SQS protocol and worker error handling.
+const MAX_BATCH_SIZE = 1;
 const POLLING_INTERVAL = 30 * 1000; // 30 seconds
 const TASK_HEALTH_CHECK_INTERVAL = 60 * 1000; // 1 minute
 const ANALYSIS_TIMEOUT = 20 * 60 * 1000;
@@ -50,6 +53,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 +68,7 @@ export default async function orchestrateCapture({
   taskRunner,
   analysisLoader,
   pollingInterval = POLLING_INTERVAL,
+  metadataCommit,
 }: OrchestrateCaptureArgs) {
   if (!full && !analysisIds && !commit && !entityShas) {
     throw new Error(
@@ -95,15 +102,20 @@ export default async function orchestrateCapture({
   const attemptCount = new Map<string, number>();
   const maxAttempts = 2;
 
+  // Track which analyses had a worker report success — only these can be "captured"
+  const confirmedSuccessAnalyses = new Set<string>();
+
   // Register batch completion callback with task runner
-  taskRunner.onBatchCompleted((analysisIds: string[]) => {
+  taskRunner.onBatchCompleted((analysisIds: string[], success: boolean) => {
     console.log(
-      `CodeYam Capture: Received batch completion for ${analysisIds.length} analyses: ${analysisIds.join(', ')}`,
+      `CodeYam Capture: Received batch completion (success=${success}) for ${analysisIds.length} analyses: ${analysisIds.join(', ')}`,
     );
-    // Remove from inFlight and increment attempt count for each analysis
     for (const analysisId of analysisIds) {
       inFlight.delete(analysisId);
       attemptCount.set(analysisId, (attemptCount.get(analysisId) || 0) + 1);
+      if (success) {
+        confirmedSuccessAnalyses.add(analysisId);
+      }
     }
   });
 
@@ -191,9 +203,13 @@ export default async function orchestrateCapture({
         knownAnalyses.add(id);
       }
 
-      // an analysis we've seen, but is no longer "ready to be captured", is captured!
+      // An analysis is "captured" only if:
+      // 1. A worker reported success for its batch (in confirmedSuccessAnalyses)
+      // 2. It no longer appears in the ready-to-be-captured list (DB confirms)
+      // This prevents transient empty DB results from inflating the captured count.
       const capturedAnalysisIds = Array.from(knownAnalyses).filter(
-        (id) => !allReadyAnalysisIds.includes(id),
+        (id) =>
+          confirmedSuccessAnalyses.has(id) && !allReadyAnalysisIds.includes(id),
       );
 
       // Remove captured analyses from inFlight tracking
@@ -318,9 +334,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,
@@ -348,6 +367,32 @@ 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 if analysis complete but no entities ready for capture
         // Only check this after at least one poll (not firstLoop) to ensure we've
@@ -364,7 +409,9 @@ export default async function orchestrateCapture({
           // check if all enqueued analyses have been either captured or abandoned
           const allAnalysesSettled = Array.from(enqueuedAnalyses).every(
             (id) => {
-              const captured = !allReadyAnalysisIds.includes(id);
+              const captured =
+                confirmedSuccessAnalyses.has(id) &&
+                !allReadyAnalysisIds.includes(id);
               const abandoned = (attemptCount.get(id) || 0) >= maxAttempts;
               return captured || abandoned;
             },
@@ -404,5 +451,33 @@ export default async function orchestrateCapture({
   } finally {
     // Always clean up task runner resources in finally block
     await taskRunner.cleanup();
+
+    const trackingCommit = metadataCommit ?? commit;
+    if (trackingCommit?.id) {
+      const captureCompletedAt = new Date().toISOString();
+
+      await updateCommitMetadata({
+        commitId: trackingCommit.id,
+        runStatusUpdate: {
+          captureCompletedAt,
+          capturePid: undefined,
+        },
+        // completedAt = max(analysisCompletedAt, captureCompletedAt)
+        // Normally capture finishes after analysis, but using max means illogical
+        // orderings leave a detectable signal (completedAt !== captureCompletedAt).
+        updateCallback(metadata) {
+          const run = metadata.currentRun!;
+          run.completedAt =
+            run.analysisCompletedAt &&
+            run.analysisCompletedAt > captureCompletedAt
+              ? run.analysisCompletedAt
+              : captureCompletedAt;
+        },
+      }).catch((error) => {
+        console.log(
+          `Warning: Failed to update capture completion metadata: ${error}`,
+        );
+      });
+    }
   }
 }

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>,
@@ -20,13 +26,38 @@ export function reconcileMockDataKeys(
         ? reconcileMockDataKeys(value as Record<string, unknown>, expectedKeys)
         : value;
 
-    // Check if key matches an expected key exactly
+    // Check if this is a flat dotted key (method chain) that should be nested.
+    // e.g., "ZId.safeParse(params.surveyId)" → { ZId: { "safeParse(params.surveyId)": ... } }
+    // e.g., "trpc.fastener.getAttributeValues.useQuery({ ... })" → nested structure
+    // constructMockCode generates nested lookups like data["ZId"]["safeParse(...)"]
+    // for method chains, so the data must be nested to match.
+    const flatKeyParts = parseFlatDottedKey(key);
+    if (flatKeyParts && flatKeyParts.length > 1) {
+      // Nest when:
+      // 1. The full key matches an expected key exactly (e.g., "ZId.safeParse(params.surveyId)")
+      // 2. OR individual parts match expected keys (e.g., "trpc" is an expected key)
+      //    This handles LLM-generated flat keys that should be nested (Issue 8)
+      const shouldNest =
+        expectedKeys.includes(key) ||
+        flatKeyParts.some((part) => expectedKeys.includes(part));
+      if (shouldNest) {
+        deepMergeIntoResult(
+          result,
+          flatKeyParts,
+          reconciledValue,
+          expectedKeys,
+        );
+        continue;
+      }
+    }
+
+    // Check if key matches an expected key exactly (for non-dotted keys)
     if (expectedKeys.includes(key)) {
       result[key] = reconciledValue;
       continue;
     }
 
-    // Try to find closest match
+    // Try to find closest match using existing logic
     const closestKey = findClosestKey(key, expectedKeys);
     if (closestKey) {
       result[closestKey] = reconciledValue;
@@ -39,6 +70,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 +166,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 +198,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

package/analyzer-template/project/runAnalysis.ts

@@ -215,6 +215,11 @@ export default async function runAnalysis({
       entityShas,
       branchName: branch?.name,
     });
+
+    // Propagate the fallback branch to branchCommit so downstream code can use it
+    if (branchCommit && branch) {
+      branchCommit.branch = branch;
+    }
   }
 
   // After this block, 'commit' is the union of a regular or baseline commit
@@ -308,6 +313,10 @@ export default async function runAnalysis({
   // Route to appropriate analyzer
   try {
     if (runReadyToBeCapturedPath) {
+      // Signal that analysis can produce captures immediately since analyses are already ready.
+      // This is needed for recapture to work - orchestration waits for the 'analysisProducing' event.
+      controller?.setAnalysisProducingPhase();
+
       return await analyzeReadyToBeCaptured({
         project,
         projectAnalyzer,
@@ -330,6 +339,7 @@ export default async function runAnalysis({
         desiredAnalyses,
         force,
         fileStore,
+        onAnalysisProducing: () => controller?.setAnalysisProducingPhase(),
       });
     }
 
@@ -350,6 +360,7 @@ export default async function runAnalysis({
       backgroundJob,
       desiredAnalyses,
       fileStore,
+      onAnalysisProducing: () => controller?.setAnalysisProducingPhase(),
     });
   } catch (e) {
     awsLog('CodeYam Error: Analysis failed', {

package/analyzer-template/project/serverOnlyModules.ts

@@ -213,6 +213,108 @@ export function applyServerOnlyMocks(fileContent: string): string {
   return applyServerOnlyMocksAst(fileContent);
 }
 
+/**
+ * Check if a file exports an "env" variable that's derived from a node module import.
+ * This is a generic pattern for env configuration packages (like @t3-oss/env-*, envalid, etc.)
+ * that create server-only environment objects.
+ *
+ * Pattern detected:
+ *   import { someFunc } from "some-package";
+ *   export const env = someFunc({...});
+ *
+ * Returns the name of the imported function if detected, null otherwise.
+ */
+function detectEnvStoreFromNodeModule(sourceFile: ts.SourceFile): {
+  importedFunctionName: string;
+  modulePath: string;
+} | null {
+  // First, collect all imported function names from node modules
+  const nodeModuleImports = new Map<string, string>(); // functionName -> modulePath
+
+  for (const statement of sourceFile.statements) {
+    if (!ts.isImportDeclaration(statement)) continue;
+
+    const moduleSpecifier = statement.moduleSpecifier;
+    if (!ts.isStringLiteral(moduleSpecifier)) continue;
+
+    const modulePath = moduleSpecifier.text;
+    // Skip relative imports - only look at node modules
+    if (modulePath.startsWith('.') || modulePath.startsWith('/')) continue;
+
+    if (
+      statement.importClause?.namedBindings &&
+      ts.isNamedImports(statement.importClause.namedBindings)
+    ) {
+      for (const element of statement.importClause.namedBindings.elements) {
+        nodeModuleImports.set(element.name.text, modulePath);
+      }
+    }
+    if (statement.importClause?.name) {
+      nodeModuleImports.set(statement.importClause.name.text, modulePath);
+    }
+  }
+
+  // Now look for: export const env = importedFunc(...)
+  for (const statement of sourceFile.statements) {
+    if (!ts.isVariableStatement(statement)) continue;
+
+    // Check if it's exported
+    const isExported = statement.modifiers?.some(
+      (m) => m.kind === ts.SyntaxKind.ExportKeyword,
+    );
+    if (!isExported) continue;
+
+    for (const decl of statement.declarationList.declarations) {
+      // Check if variable is named "env"
+      if (!ts.isIdentifier(decl.name) || decl.name.text !== 'env') continue;
+
+      // Check if initializer is a call expression
+      if (!decl.initializer || !ts.isCallExpression(decl.initializer)) continue;
+
+      // Get the function being called
+      const callExpr = decl.initializer.expression;
+      let funcName: string | null = null;
+
+      if (ts.isIdentifier(callExpr)) {
+        funcName = callExpr.text;
+      } else if (
+        ts.isPropertyAccessExpression(callExpr) &&
+        ts.isIdentifier(callExpr.name)
+      ) {
+        funcName = callExpr.name.text;
+      }
+
+      if (funcName && nodeModuleImports.has(funcName)) {
+        return {
+          importedFunctionName: funcName,
+          modulePath: nodeModuleImports.get(funcName)!,
+        };
+      }
+    }
+  }
+
+  return null;
+}
+
+/**
+ * Generate mock code for env store patterns.
+ * Replaces the imported function with one that returns a proxy.
+ */
+function generateEnvStoreMock(functionName: string): string {
+  return `// Mock for env store function (server-only env vars)
+const ${functionName} = (config: any) => {
+  // Return a proxy that returns undefined for any property access
+  // This allows the code to run on the client without throwing
+  return new Proxy({} as any, {
+    get: (target, prop) => {
+      if (prop === 'then' || prop === Symbol.toStringTag) return undefined;
+      return undefined;
+    },
+  });
+};
+`;
+}
+
 /**
  * AST-based implementation of applyServerOnlyMocks for large files.
  * Uses TypeScript parser instead of regex to avoid catastrophic backtracking.
@@ -229,6 +331,10 @@ function applyServerOnlyMocksAst(fileContent: string): string {
     const replacements: { start: number; end: number; replacement: string }[] =
       [];
 
+    // Check if this file exports an "env" variable from a node module function call
+    // This is a generic pattern for env configuration packages
+    const envStoreInfo = detectEnvStoreFromNodeModule(sourceFile);
+
     for (const statement of sourceFile.statements) {
       if (!ts.isImportDeclaration(statement)) continue;
 
@@ -236,8 +342,6 @@ function applyServerOnlyMocksAst(fileContent: string): string {
       if (!ts.isStringLiteral(moduleSpecifier)) continue;
 
       const modulePath = moduleSpecifier.text;
-      const config = findServerOnlyModule(modulePath);
-      if (!config) continue;
 
       // Extract import names
       const importNames: string[] = [];
@@ -265,6 +369,27 @@ function applyServerOnlyMocksAst(fileContent: string): string {
         }
       }
 
+      // Check if this import provides the function used to create the env store
+      if (
+        envStoreInfo &&
+        modulePath === envStoreInfo.modulePath &&
+        importNames.includes(envStoreInfo.importedFunctionName)
+      ) {
+        const mockCode = generateEnvStoreMock(
+          envStoreInfo.importedFunctionName,
+        );
+        replacements.push({
+          start: statement.getStart(sourceFile),
+          end: statement.getEnd(),
+          replacement: mockCode,
+        });
+        continue;
+      }
+
+      // Check for known server-only modules (from explicit list)
+      const config = findServerOnlyModule(modulePath);
+      if (!config) continue;
+
       const mockCode = config.getMockCode(importNames, defaultImportName);
       replacements.push({
         start: statement.getStart(sourceFile),