feat(core): Implement granular memory loaders for JIT architecture (#12195)

2026-07-15 12:30:32 -07:00 · 2025-10-30 00:09:12 -04:00
parent d5f4008625
commit 5f9a74bb7c
2 changed files with 468 additions and 2 deletions
@@ -323,6 +323,143 @@ function concatenateInstructions(
    .join('\n\n');
 }

+export interface MemoryLoadResult {
+  files: Array<{ path: string; content: string }>;
+}
+
+export async function loadGlobalMemory(
+  debugMode: boolean = false,
+): Promise<MemoryLoadResult> {
+  const userHome = homedir();
+  const geminiMdFilenames = getAllGeminiMdFilenames();
+
+  const accessChecks = geminiMdFilenames.map(async (filename) => {
+    const globalPath = path.join(userHome, GEMINI_DIR, filename);
+    try {
+      await fs.access(globalPath, fsSync.constants.R_OK);
+      if (debugMode) {
+        logger.debug(`Found global memory file: ${globalPath}`);
+      }
+      return globalPath;
+    } catch {
+      debugLogger.debug('A global memory file was not found.');
+      return null;
+    }
+  });
+
+  const foundPaths = (await Promise.all(accessChecks)).filter(
+    (p): p is string => p !== null,
+  );
+
+  const contents = await readGeminiMdFiles(foundPaths, debugMode, 'tree');
+
+  return {
+    files: contents
+      .filter((item) => item.content !== null)
+      .map((item) => ({
+        path: item.filePath,
+        content: item.content as string,
+      })),
+  };
+}
+
+/**
+ * Traverses upward from startDir to stopDir, finding all GEMINI.md variants.
+ *
+ * Files are ordered by directory level (root to leaf), with all filename
+ * variants grouped together per directory.
+ */
+async function findUpwardGeminiFiles(
+  startDir: string,
+  stopDir: string,
+  debugMode: boolean,
+): Promise<string[]> {
+  const upwardPaths: string[] = [];
+  let currentDir = path.resolve(startDir);
+  const resolvedStopDir = path.resolve(stopDir);
+  const geminiMdFilenames = getAllGeminiMdFilenames();
+  const globalGeminiDir = path.join(homedir(), GEMINI_DIR);
+
+  if (debugMode) {
+    logger.debug(
+      `Starting upward search from ${currentDir} stopping at ${resolvedStopDir}`,
+    );
+  }
+
+  while (true) {
+    if (currentDir === globalGeminiDir) {
+      break;
+    }
+
+    // Parallelize checks for all filename variants in the current directory
+    const accessChecks = geminiMdFilenames.map(async (filename) => {
+      const potentialPath = path.join(currentDir, filename);
+      try {
+        await fs.access(potentialPath, fsSync.constants.R_OK);
+        return potentialPath;
+      } catch {
+        return null;
+      }
+    });
+
+    const foundPathsInDir = (await Promise.all(accessChecks)).filter(
+      (p): p is string => p !== null,
+    );
+
+    upwardPaths.unshift(...foundPathsInDir);
+
+    if (
+      currentDir === resolvedStopDir ||
+      currentDir === path.dirname(currentDir)
+    ) {
+      break;
+    }
+    currentDir = path.dirname(currentDir);
+  }
+  return upwardPaths;
+}
+
+export async function loadEnvironmentMemory(
+  trustedRoots: string[],
+  extensionLoader: ExtensionLoader,
+  debugMode: boolean = false,
+): Promise<MemoryLoadResult> {
+  const allPaths = new Set<string>();
+
+  // Trusted Roots Upward Traversal (Parallelized)
+  const traversalPromises = trustedRoots.map(async (root) => {
+    const resolvedRoot = path.resolve(root);
+    if (debugMode) {
+      logger.debug(
+        `Loading environment memory for trusted root: ${resolvedRoot} (Stopping exactly here)`,
+      );
+    }
+    return await findUpwardGeminiFiles(resolvedRoot, resolvedRoot, debugMode);
+  });
+
+  const pathArrays = await Promise.all(traversalPromises);
+  pathArrays.flat().forEach((p) => allPaths.add(p));
+
+  // Extensions
+  const extensionPaths = extensionLoader
+    .getExtensions()
+    .filter((ext) => ext.isActive)
+    .flatMap((ext) => ext.contextFiles);
+  extensionPaths.forEach((p) => allPaths.add(p));
+
+  const sortedPaths = Array.from(allPaths).sort();
+  const contents = await readGeminiMdFiles(sortedPaths, debugMode, 'tree');
+
+  return {
+    files: contents
+      .filter((item) => item.content !== null)
+      .map((item) => ({
+        path: item.filePath,
+        content: item.content as string,
+      })),
+  };
+}
+
 export interface LoadServerHierarchicalMemoryResponse {
  memoryContent: string;
  fileCount: number;
@@ -400,3 +537,68 @@ export async function loadServerHierarchicalMemory(
    filePaths,
  };
 }
+
+export async function loadJitSubdirectoryMemory(
+  targetPath: string,
+  trustedRoots: string[],
+  alreadyLoadedPaths: Set<string>,
+  debugMode: boolean = false,
+): Promise<MemoryLoadResult> {
+  const resolvedTarget = path.resolve(targetPath);
+  let bestRoot: string | null = null;
+
+  // Find the deepest trusted root that contains the target path
+  for (const root of trustedRoots) {
+    const resolvedRoot = path.resolve(root);
+    if (
+      resolvedTarget.startsWith(resolvedRoot) &&
+      (!bestRoot || resolvedRoot.length > bestRoot.length)
+    ) {
+      bestRoot = resolvedRoot;
+    }
+  }
+
+  if (!bestRoot) {
+    if (debugMode) {
+      logger.debug(
+        `JIT memory skipped: ${resolvedTarget} is not in any trusted root.`,
+      );
+    }
+    return { files: [] };
+  }
+
+  if (debugMode) {
+    logger.debug(
+      `Loading JIT memory for ${resolvedTarget} (Trusted root: ${bestRoot})`,
+    );
+  }
+
+  // Traverse from target up to the trusted root
+  const potentialPaths = await findUpwardGeminiFiles(
+    resolvedTarget,
+    bestRoot,
+    debugMode,
+  );
+
+  // Filter out already loaded paths
+  const newPaths = potentialPaths.filter((p) => !alreadyLoadedPaths.has(p));
+
+  if (newPaths.length === 0) {
+    return { files: [] };
+  }
+
+  if (debugMode) {
+    logger.debug(`Found new JIT memory files: ${JSON.stringify(newPaths)}`);
+  }
+
+  const contents = await readGeminiMdFiles(newPaths, debugMode, 'tree');
+
+  return {
+    files: contents
+      .filter((item) => item.content !== null)
+      .map((item) => ({
+        path: item.filePath,
+        content: item.content as string,
+      })),
+  };
+}