mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-10 22:21:22 -07:00
feat(core): Implement granular memory loaders for JIT architecture (#12195)
This commit is contained in:
Abhi
@@ -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();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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,
|
||||
})),
|
||||
};
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user