diff --git a/packages/core/src/utils/memoryDiscovery.test.ts b/packages/core/src/utils/memoryDiscovery.test.ts index 5f8c4b534c..fba1901f43 100644 --- a/packages/core/src/utils/memoryDiscovery.test.ts +++ b/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(); + }); + }); }); diff --git a/packages/core/src/utils/memoryDiscovery.ts b/packages/core/src/utils/memoryDiscovery.ts index 212dd7f935..bd98f0ab01 100644 --- a/packages/core/src/utils/memoryDiscovery.ts +++ b/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 { + 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 { + 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 { + const allPaths = new Set(); + + // 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, + debugMode: boolean = false, +): Promise { + 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, + })), + }; +}