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

packages/core/src/utils/memoryDiscovery.test.ts

@@ -8,7 +8,12 @@ import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest';
 import * as fsPromises from 'node:fs/promises';
 import * as os from 'node:os';
 import * as path from 'node:path';
-import { loadServerHierarchicalMemory } from './memoryDiscovery.js';
+import {
+  loadServerHierarchicalMemory,
+  loadGlobalMemory,
+  loadEnvironmentMemory,
+  loadJitSubdirectoryMemory,
+} from './memoryDiscovery.js';
 import {
   setGeminiMdFilename,
   DEFAULT_CONTEXT_FILENAME,
@@ -26,7 +31,7 @@ vi.mock('os', async (importOriginal) => {
   };
 });
 
-describe('loadServerHierarchicalMemory', () => {
+describe('memoryDiscovery', () => {
   const DEFAULT_FOLDER_TRUST = true;
   let testRootDir: string;
   let cwd: string;
@@ -612,4 +617,263 @@ included directory memory
     expect(parentOccurrences).toBe(1);
     expect(childOccurrences).toBe(1);
   });
+
+  describe('loadGlobalMemory', () => {
+    it('should load global memory file if it exists', async () => {
+      const globalMemoryFile = await createTestFile(
+        path.join(homedir, GEMINI_DIR, DEFAULT_CONTEXT_FILENAME),
+        'Global memory content',
+      );
+
+      const result = await loadGlobalMemory();
+
+      expect(result.files).toHaveLength(1);
+      expect(result.files[0].path).toBe(globalMemoryFile);
+      expect(result.files[0].content).toBe('Global memory content');
+    });
+
+    it('should return empty content if global memory file does not exist', async () => {
+      const result = await loadGlobalMemory();
+
+      expect(result.files).toHaveLength(0);
+    });
+  });
+
+  describe('loadEnvironmentMemory', () => {
+    it('should load extension memory', async () => {
+      const extFile = await createTestFile(
+        path.join(testRootDir, 'ext', 'GEMINI.md'),
+        'Extension content',
+      );
+      const mockExtensionLoader = new SimpleExtensionLoader([
+        {
+          isActive: true,
+          contextFiles: [extFile],
+        } as GeminiCLIExtension,
+      ]);
+
+      const result = await loadEnvironmentMemory([], mockExtensionLoader);
+
+      expect(result.files).toHaveLength(1);
+      expect(result.files[0].path).toBe(extFile);
+      expect(result.files[0].content).toBe('Extension content');
+    });
+
+    it('should NOT traverse upward beyond trusted root (even with .git)', async () => {
+      // Setup: /temp/parent/repo/.git
+      const parentDir = await createEmptyDir(path.join(testRootDir, 'parent'));
+      const repoDir = await createEmptyDir(path.join(parentDir, 'repo'));
+      await createEmptyDir(path.join(repoDir, '.git'));
+      const srcDir = await createEmptyDir(path.join(repoDir, 'src'));
+
+      await createTestFile(
+        path.join(parentDir, DEFAULT_CONTEXT_FILENAME),
+        'Parent content',
+      );
+      await createTestFile(
+        path.join(repoDir, DEFAULT_CONTEXT_FILENAME),
+        'Repo content',
+      );
+      const srcFile = await createTestFile(
+        path.join(srcDir, DEFAULT_CONTEXT_FILENAME),
+        'Src content',
+      );
+
+      // Trust srcDir. Should ONLY load srcFile.
+      // Repo and Parent are NOT trusted.
+      const result = await loadEnvironmentMemory(
+        [srcDir],
+        new SimpleExtensionLoader([]),
+      );
+
+      expect(result.files).toHaveLength(1);
+      expect(result.files[0].path).toBe(srcFile);
+      expect(result.files[0].content).toBe('Src content');
+    });
+
+    it('should NOT traverse upward beyond trusted root (no .git)', async () => {
+      // Setup: /homedir/docs/notes (no .git anywhere)
+      const docsDir = await createEmptyDir(path.join(homedir, 'docs'));
+      const notesDir = await createEmptyDir(path.join(docsDir, 'notes'));
+
+      await createTestFile(
+        path.join(homedir, DEFAULT_CONTEXT_FILENAME),
+        'Home content',
+      );
+      const docsFile = await createTestFile(
+        path.join(docsDir, DEFAULT_CONTEXT_FILENAME),
+        'Docs content',
+      );
+
+      // Trust notesDir. Should load NOTHING because notesDir has no file,
+      // and we do not traverse up to docsDir.
+      const resultNotes = await loadEnvironmentMemory(
+        [notesDir],
+        new SimpleExtensionLoader([]),
+      );
+      expect(resultNotes.files).toHaveLength(0);
+
+      // Trust docsDir. Should load docsFile, but NOT homeFile.
+      const resultDocs = await loadEnvironmentMemory(
+        [docsDir],
+        new SimpleExtensionLoader([]),
+      );
+      expect(resultDocs.files).toHaveLength(1);
+      expect(resultDocs.files[0].path).toBe(docsFile);
+      expect(resultDocs.files[0].content).toBe('Docs content');
+    });
+
+    it('should deduplicate paths when same root is trusted multiple times', async () => {
+      const repoDir = await createEmptyDir(path.join(testRootDir, 'repo'));
+      await createEmptyDir(path.join(repoDir, '.git'));
+
+      const repoFile = await createTestFile(
+        path.join(repoDir, DEFAULT_CONTEXT_FILENAME),
+        'Repo content',
+      );
+
+      // Trust repoDir twice.
+      const result = await loadEnvironmentMemory(
+        [repoDir, repoDir],
+        new SimpleExtensionLoader([]),
+      );
+
+      expect(result.files).toHaveLength(1);
+      expect(result.files[0].path).toBe(repoFile);
+    });
+
+    it('should keep multiple memory files from the same directory adjacent and in order', async () => {
+      // Configure multiple memory filenames
+      setGeminiMdFilename(['PRIMARY.md', 'SECONDARY.md']);
+
+      const dir = await createEmptyDir(
+        path.join(testRootDir, 'multi_file_dir'),
+      );
+      await createEmptyDir(path.join(dir, '.git'));
+
+      const primaryFile = await createTestFile(
+        path.join(dir, 'PRIMARY.md'),
+        'Primary content',
+      );
+      const secondaryFile = await createTestFile(
+        path.join(dir, 'SECONDARY.md'),
+        'Secondary content',
+      );
+
+      const result = await loadEnvironmentMemory(
+        [dir],
+        new SimpleExtensionLoader([]),
+      );
+
+      expect(result.files).toHaveLength(2);
+      // Verify order: PRIMARY should come before SECONDARY because they are
+      // sorted by path and PRIMARY.md comes before SECONDARY.md alphabetically
+      // if in same dir.
+      expect(result.files[0].path).toBe(primaryFile);
+      expect(result.files[1].path).toBe(secondaryFile);
+      expect(result.files[0].content).toBe('Primary content');
+      expect(result.files[1].content).toBe('Secondary content');
+    });
+  });
+
+  describe('loadJitSubdirectoryMemory', () => {
+    it('should load JIT memory when target is inside a trusted root', async () => {
+      const rootDir = await createEmptyDir(path.join(testRootDir, 'jit_root'));
+      const subDir = await createEmptyDir(path.join(rootDir, 'subdir'));
+      const targetFile = path.join(subDir, 'target.txt');
+
+      const subDirMemory = await createTestFile(
+        path.join(subDir, DEFAULT_CONTEXT_FILENAME),
+        'Subdir JIT content',
+      );
+
+      const result = await loadJitSubdirectoryMemory(
+        targetFile,
+        [rootDir],
+        new Set(),
+      );
+
+      expect(result.files).toHaveLength(1);
+      expect(result.files[0].path).toBe(subDirMemory);
+      expect(result.files[0].content).toBe('Subdir JIT content');
+    });
+
+    it('should skip JIT memory when target is outside trusted roots', async () => {
+      const trustedRoot = await createEmptyDir(
+        path.join(testRootDir, 'trusted'),
+      );
+      const untrustedDir = await createEmptyDir(
+        path.join(testRootDir, 'untrusted'),
+      );
+      const targetFile = path.join(untrustedDir, 'target.txt');
+
+      await createTestFile(
+        path.join(untrustedDir, DEFAULT_CONTEXT_FILENAME),
+        'Untrusted content',
+      );
+
+      const result = await loadJitSubdirectoryMemory(
+        targetFile,
+        [trustedRoot],
+        new Set(),
+      );
+
+      expect(result.files).toHaveLength(0);
+    });
+
+    it('should skip already loaded paths', async () => {
+      const rootDir = await createEmptyDir(path.join(testRootDir, 'jit_root'));
+      const subDir = await createEmptyDir(path.join(rootDir, 'subdir'));
+      const targetFile = path.join(subDir, 'target.txt');
+
+      const rootMemory = await createTestFile(
+        path.join(rootDir, DEFAULT_CONTEXT_FILENAME),
+        'Root content',
+      );
+      const subDirMemory = await createTestFile(
+        path.join(subDir, DEFAULT_CONTEXT_FILENAME),
+        'Subdir content',
+      );
+
+      // Simulate root memory already loaded (e.g., by loadEnvironmentMemory)
+      const alreadyLoaded = new Set([rootMemory]);
+
+      const result = await loadJitSubdirectoryMemory(
+        targetFile,
+        [rootDir],
+        alreadyLoaded,
+      );
+
+      expect(result.files).toHaveLength(1);
+      expect(result.files[0].path).toBe(subDirMemory);
+      expect(result.files[0].content).toBe('Subdir content');
+    });
+
+    it('should use the deepest trusted root when multiple nested roots exist', async () => {
+      const outerRoot = await createEmptyDir(path.join(testRootDir, 'outer'));
+      const innerRoot = await createEmptyDir(path.join(outerRoot, 'inner'));
+      const targetFile = path.join(innerRoot, 'target.txt');
+
+      const outerMemory = await createTestFile(
+        path.join(outerRoot, DEFAULT_CONTEXT_FILENAME),
+        'Outer content',
+      );
+      const innerMemory = await createTestFile(
+        path.join(innerRoot, DEFAULT_CONTEXT_FILENAME),
+        'Inner content',
+      );
+
+      const result = await loadJitSubdirectoryMemory(
+        targetFile,
+        [outerRoot, innerRoot],
+        new Set(),
+      );
+
+      expect(result.files).toHaveLength(1);
+      expect(result.files[0].path).toBe(innerMemory);
+      expect(result.files[0].content).toBe('Inner content');
+      // Ensure outer memory is NOT loaded
+      expect(result.files.find((f) => f.path === outerMemory)).toBeUndefined();
+    });
+  });
 });

packages/core/src/utils/memoryDiscovery.ts

@@ -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,
+      })),
+  };
+}